feat: per-shape probe — measure real decode M-buckets for graph-hidden kernels - #384
Draft
yueliu14 wants to merge 23 commits into
Draft
feat: per-shape probe — measure real decode M-buckets for graph-hidden kernels#384yueliu14 wants to merge 23 commits into
yueliu14 wants to merge 23 commits into
Conversation
…kets) - probe_postprocess: per-tensor dtype (parallel to dims, no dedup) + arg_labels - probe_to_mbuckets: extract measured decode/prefill m_buckets from probe output - kernel_extractor.md: add step 2b — graph-hidden kernels measure m_buckets via enforce-eager probe instead of the M≈conc guess (inference kept as fallback) - shape_capture.md: add Step 1b (measured graph-hidden shapes) + probe hard pitfalls - vllm.sh: optional PROFILE_ACTIVE_ITERS to widen profiler window - capture_shapes_probe.py: per-tensor dtype/arg_labels
…-run robustness End-to-end verification on a short bench (OSL=64) exposed that the 0.1% default over-kept long-tail prefill chunk shapes when the total call count is small (the tail's relative share rises). 1% keeps the traffic-dominant buckets (decode M + the few prefill chunk sizes) and drops the tail. Verified: decode[64,256] + prefill[8192,32768] on gpt-oss-120b.
Remove the hard-coded gpt-oss TARGET_MAP from probe_postprocess.py; kernels are now auto-discovered from the probe output files themselves, so a new model just needs PROBE_TARGETS pointed at its hot kernels with no code change. - probe_postprocess: iterate discovered probe_*.json targets instead of a fixed map; derive label + profile-match substrings from the target string - workload (isl/osl/conc) is now a CLI param (was hard-coded 1024/1024/64), defaults to null when unset - --profile-topn stays optional but downgraded to "strongly recommended": missing profile prints a loud WARNING and marks %GPU unknown instead of silently producing possibly-misleading approx data - note_hipblaslt -> generic note_coverage (no gpt-oss #4 Cijk specifics) - kernel_extractor.md / shape_capture.md: add per-model PROBE_TARGETS derivation playbook; mark gpt-oss values (targets, M=64/256) as examples only Regression verified on gpt-oss-120b (both level-1 postprocess replay on two existing baselines and level-2 end-to-end vLLM run): decode_m_buckets=[64,256], prefill_m_buckets=[8192,32768], shape/count unchanged.
…e probe by default
…ured (warn-but-continue) Extend EXTRACT_OP_SCHEMA with decode/prefill m_buckets + m_buckets_source so the sandboxed workflow script (no fs access) can confirm the probe ran from agent-returned fields. Add verifyProbeMBuckets guard at all three extract_op call sites; warn but continue on synthesized fallback. Mirror scripts/probe_mbuckets_guard.py rule.
…s join profiled kernels Lets a Python launcher like invoke_fused_moe_triton_kernel match the profiled fused_moe_kernel %GPU (Qwen3), not just gpt-oss matmul_ogs. Model-agnostic.
# Conflicts: # e2e_workflow/e2e_workflow.js # e2e_workflow/roles/kernel_extractor.md # e2e_workflow/scripts/adapters/vllm.sh
…er-kernel) Replace the per-kernel probe (N+M enforce-eager workload passes, one inside every extract_op/extract call) with a single global probe over the union of head + milestone candidate kernels, shared by both tracks. - e2e_workflow.js: add PROBE_ALL_SCHEMA + idempotent ensureSharedProbe() (in-memory memo + carryState.shared_probe_json; the probe_all agent skips the server when the on-disk artifact exists). Call it before the head loop and the milestone loop; thread SHARED_PROBE_JSON into all extract_op sites + the milestone extract site; add a Probe entry to meta.phases. - kernel_extractor.md: add PHASE=probe_all (derive union PROBE_TARGETS, one enforce-eager probe, one probe_postprocess -> shared per_shape_probe.json). Rewrite 2b / 2b-op so the DEFAULT path slices the shared artifact via probe_to_mbuckets --kernel-match (no server); the per-kernel probe becomes the FALLBACK for kernels nominated after probe_all ran. Probe scripts unchanged (already support multi-target one-run + --kernel-match slicing). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…pe/weight/baseline) Make the unittest-building stage identical on both tracks; optimization paths unchanged. - kernel_extractor.md: add a single canonical "Unittest contract — shapes / weight / baseline (IDENTICAL for BOTH phases)" section; both PHASE=extract and PHASE=extract_op now point to it (kills doctrine drift between the two parallel prose sections). - e2e_workflow.js: EXTRACT_SCHEMA gains decode_m_buckets/prefill_m_buckets/m_buckets_source (parity with EXTRACT_OP_SCHEMA); call verifyProbeMBuckets after the milestone extract site so the probe-measured-vs-guessed guard now fires on the milestone track too. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The probe_all agent already sets opts.phase:'Probe' for display grouping; the global phase() switch only left the caller's phase dangling after the helper returned. Removed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Probe capture depends on hooking the module attribute the CALL SITE resolves at call time, and C++/asm/CK kernels are reachable via the Python dispatcher above them (not out of scope). Verified end-to-end on MI300X / vLLM v0.25.1. - Fix step 3: a C++/CK kernel (ROCM_ATTN, hipBLASLt GEMM) has no Python entry itself, but hook the module-level Python dispatcher above it; class methods aren't wrappable by the current module:attr probe. - Fix step 5: replace "prefer the definition module" (wrong for module-level imports — early binding misses the wrap) with the call-site binding rule: function-local import -> def module; module-level import -> caller module; 0-calls-after-hooked self-check. - Add worked seams by kernel type: MoE launcher (matmul_ogs), attention dispatcher (rocm_attn:chunked_prefill_paged_decode -> decode M=64/prefill 8192), dense GEMM (torch.nn.functional:linear; but prefer GEMM_SYNTH since dense GEMM is analytic). - probe_all target-derivation points at the seam rule. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The probe mechanism is backend-agnostic — the sglang adapter already propagates OVERLAY_PYTHONPATH/EXTRA_ENV and prepends the sglang source to PYTHONPATH, so no code change is needed. Document the two sglang-specific differences: - enforce-eager = --disable-cuda-graph (via EXTRA_SERVER_ARGS), NOT --enforce-eager. - sglang module:attr seams (from sglang v0.5.12 source; probe-verify per step 7): MoE -> ...moe_runner.triton_utils.fused_moe_triton_kernels:invoke_fused_moe_kernel attn -> ...attention.triton_ops.decode_attention:decode_attention_fwd (+ extend_attention_fwd) GEMM -> torch.nn.functional:linear (prefer synthesis). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
yueliu14
changed the title
The isolated unittest timed the candidate against baseline_src/ -- a second
copy of the source INSIDE the task dir, selected by a meta.baseline_callable
string the generated harness bound by hand. Both trees held byte-identical,
identically-named code, so "which one is the baseline" was a coin flip that
failed SILENTLY: both legs are correct implementations, so correctness still
passed and only the ratio inverted. It also let mode=author time optimized-HIP
against its own naive-HIP seed (15.7x isolated, ~0% e2e).
Both legs now run the SAME leg_runner.py + the SAME cases.py, differing only in
the overlay on PYTHONPATH:
baseline <task>/baseline_overlay frozen CURRENT_OVERLAY = install + every
already-accepted kernel (== the e2e gate's
ref leg)
candidate <task>/_cand_overlay that same stack + exactly ONE entry built
from kernel_src/ (meta.candidate_bind)
Direction becomes a property of the environment rather than of a name, so
speedup = baseline_ms/optimized_ms cannot invert. assert_legs_differ refuses to
measure until both legs provably resolve target_callable to different code AND
the baseline resolves outside the task dir.
Also fixes the ~8% loader gap from comparing a from-file loader against the
installed package: both legs now import through the identical mechanism.
Side effects: the extractor finally receives CURRENT_OVERLAY (every other role
already did), so isolated speedup and e2e_delta_pct share a denominator and the
Amdahl cross-check is meaningful; and two accepted kernels in the same module
COMPOUND instead of the second silently reverting the first, because kernel_src/
now starts from the module as the current stack resolves it.
NOT YET VALIDATED ON GPU -- the torch-dependent paths (measure_legs timing,
oracle .pt roundtrip) have never executed. Verified so far: ast.parse on the
three scripts, and assert_legs_differ against a synthetic fixture (positive plus
both negatives).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Issue #389 bug 2: shape capture for the MiniMax-M3 MoE kernel died at "Server not healthy within 360s", so the most valuable editable target in the run was never authored. The previous fix graded the budget by purpose (OVERLAY_KIND=capture -> 1800s, candidate -> 72*5=360s). That still guesses a constant per model, and it still treats "slow" as a proxy for "wedged". The tight budget exists to reject a wedged candidate overlay -- process alive, /health 503s forever -- before it starves the serving GPU. But wedged shows up as SILENCE, not as elapsed time. Judge on the actual evidence instead: * keep waiting while the server log keeps growing; declare death only after STALL_WINDOW_SEC (600) with no new bytes. Both backends print continuously through dist init, shard load and graph capture, so no-new-bytes IS wedged. This rejects a wedged candidate SOONER than 360s and lets a legitimately slow TP4 cold start run as long as it keeps moving. * the absolute ceiling is now DERIVED from the checkpoint size (600 + 8*GB/TP, clamped 900..7200; 1800 for an HF id) and is only a backstop against a server that spins printing forever. Not a prediction. * retry ONCE, and only for unambiguously transient faults (engine_core_timeout, dist_init_fail). OOM is never retried: lowering MEM_FRACTION would hide "this candidate does not fit", which is a result. * always write OUT_DIR/server_start.json {status, attempts, reason, phase_hint, wait_sec, ceiling_sec, ...} so a failed cold start is a readable reason downstream, not an empty task dir that reads as "authored and found no gain". OVERLAY_KIND / CAPTURE_HEALTH_TRIES / OVERLAY_HEALTH_TRIES no longer select a budget -- OVERLAY_KIND survives as a provenance label only. That also fixes the callers that never set it (the three probe invocations in kernel_extractor.md and director.md's final validation leg all silently got the 360s fail-fast budget). SERVER_STARTUP_TIMEOUT_SEC still overrides the backstop, and the handoff path that exports it is unchanged. Verified without a GPU via a throwaway stub adapter (ADAPTER= override), STALL_WINDOW_SEC=5: prints every 1s, healthy at 25s -> status=ok attempts=1 wait_sec=25 (any fixed budget under 25s kills it) fully silent -> reason=stalled attempts=1 after 5s "EngineCore initialization timed out" -> reason=engine_core_timeout attempts=2 "HIP out of memory" -> reason=oom attempts=1 (no retry) ceiling_sec derived correctly from a 3GB fake checkpoint at TP4 (606 -> 900). NOT yet exercised on a real GPU cold start. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Conflict resolutions: * interface/run_e2e.py, interface/test_run_e2e_alignment.py — took main's side for the divergence code (PR #391's _positive_finite_float / _divergence_pct / _build_baseline_alignment supersede 1585a80's raw_div / same_cfg_div, and main deliberately dropped the legacy baseline_divergence_pct key). Dropped the two 1585a80 tests that asserted those removed aliases. apply_server_timeouts + its tests kept. * e2e_workflow/e2e_workflow.js — kept main's FlyDSL provisioning gate and 1585a80's hasFrozenBaseline comment (baseline_overlay/ + candidate_bind). * kernel_workflow/kernel_workflow.js — main split it into a bakeoff dispatcher + kernel_lane.js worker; took main's dispatcher wholesale and re-applied 1585a80's edits to kernel_lane.js, where the old worker body now lives: overlay wording for the frozen-baseline denominator and the best_patch.diff kernel_src/ pathspec (a bare `git diff` would carry a tampered unittest.py/meta.json into the winner merge). * kernel_workflow/roles/{director,author_engineer}.md — merged both contracts: main's reference_io.pt symlink sharing + optional-golden note, plus 1585a80's baseline_overlay/ + cases.py/leg_runner.py/overlay_setup.py copy set. Both baseline forms are named, since the bakeoff track's oracle_freezer still produces baseline_src/ while the e2e kernel_extractor produces baseline_overlay/. Verified: bash -n bench_e2e.sh; 975 passed / 5 skipped across interface/ + e2e_workflow/scripts/tests/; run_e2e.py --dry-run on ci/fixtures/handoff.dry.json. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ed test
CI run 120 failed the coverage gate (85.77% vs 97%), for two reasons the same
commit introduced.
1. An orphaned test file. `.github/workflows/ci-l0-checks.yml` lists its pytest
files EXPLICITLY (so a GPU test under examples/ can never silently join the
L0 job). test_probe_mbuckets_guard.py was added without touching that list,
so it had never run once — a real bug independent of coverage. Five entries
are added here: the guard test plus the four new ones below.
2. 580 statements landed in the measured trees with 542 uncovered. Tests, not
`omit`: pyproject's source list exists precisely so a new untested script
cannot land without moving the number, and its fail_under comment says
lowering it is the one change that should never be approved.
New tests, each bringing its target to 100%:
test_capture_shapes_probe.py (191 stmts) — the twin file capture_shapes.py is
already at 100% by stubbing torch into sys.modules, and the probe uses the
same lazy _torch() idiom, so the whole thing runs on CPU against fake
tensors exposing only .shape/.dtype. Pins what makes it safe to sit on a hot
serving path: the real kernel runs first and its return value passes through
untouched; a bookkeeping error is swallowed but the kernel's is not; install
never imports the target's module; a triton JITFunction is refused rather
than wrapped (fn[grid](...) would break); kwargs are scanned (vLLM calls
unified_attention by keyword, and a positional-only scan records dims=[],
indistinguishable from graph capture); flush files are per-pid; the daemon
flusher survives a failing snapshot, since atexit does not reliably run on
the EngineCore child.
test_leg_runner.py (82) — the four modes against a real stub task dir on disk,
because leg_runner loads harness_lib and cases.py BY PATH. Pins that
_identity reports an {"error": ...} record instead of raising (a raise would
disarm assert_legs_differ) and that both file lookups are best-effort.
test_probe_postprocess.py (150) — cross-pid sum, the dims+dtypes dedup key,
the calls-weighted avg_us / summed pct_gpu join, and the enforce-eager vs
graph-capture data_semantics switch.
test_probe_to_mbuckets.py (48) — kernel selection by label-or-target with a
busiest-kernel tiebreak, the inclusive M <= conc*mult boundary, a count-share
floor relative to the kernel's OWN total, and (None, reason) rather than
empty-but-plausible bucket lists on the two failure paths.
test_harness_lib.py gains the two-leg measurement section, restoring that file
to 100%: _run_leg (overlay FIRST on PYTHONPATH, no empty entry that would put
CWD — and the task dir's unittest.py — on sys.path, last-JSON-line-wins, raise
on a leg that printed nothing), build_candidate_overlay (stale overlay rebuilt,
both bind kinds), assert_legs_differ (identical identities and a baseline
resolving inside the task dir are both refused before any timing runs), and the
recorded-oracle arm of check_random_vs_baseline, where a missing draw must fail
rather than silently shrink the parity check.
Total: 3949 stmts, 15 miss, 99.62% — the remainder is pre-existing defensive
`except: pass` guards.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… timeout handoff Three unrelated additions get rolled back to their main-branch shape: - probe_mbuckets_guard.py duplicated the fallback rule that e2e_workflow.js already enforces in verifyProbeMBuckets. One copy is enough; the JS comment no longer points at a file that does not exist. - adapters/vllm.sh goes back to a plain --profiler-config. PROFILE_ACTIVE_ITERS had no remaining reader. - run_e2e.py's apply_server_timeouts() plumbed the handoff JSON into SERVER_STARTUP_TIMEOUT_SEC. bench_e2e.sh still reads that variable directly from the environment, so the override contract survives; only the extra indirection goes. Also fixes an order-dependency in the leg_runner PathScrub test: other modules in the same pytest session put SCRIPTS_DIR back on sys.path at import time, so the process-global assertion passed alone and failed in the suite. It now re-runs the scrub and asserts on that. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
# Conflicts: # kernel_workflow/kernel_lane.js
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.