feat: per-shape probe — measure real decode M-buckets for graph-hidden kernels - #384
Draft
yueliu14 wants to merge 18 commits into
Draft
feat: per-shape probe — measure real decode M-buckets for graph-hidden kernels#384yueliu14 wants to merge 18 commits into
yueliu14 wants to merge 18 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>
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.