Fix UT capture losing loader-set dispatch metadata; gate the baseline leg on meta.device_kernel

[yueliu14]

An extracted kernel task can be named after one GPU kernel and measure another, end to end, with every existing gate passing.

The hole

capture_shapes._snapshot recorded a fixed field list (data/dtype/shape/device/contiguous) and nothing from x.__dict__. Several backends keep their dispatch decision there rather than in the data: aiter's fused-MoE gate is use_mxfp4_flydsl = (... and is_shuffled and ...), where is_shuffled is a plain Python attribute the weight loader attaches with setattr. torch.save does not persist it, and .to(device) returns a fresh tensor with an empty __dict__, so the replayed oracle fell into the CK branch while the captured server had run FlyDSL.

Correctness cannot catch this: the golden is frozen from that same baseline leg, so it grades the wrong branch against its own output and agrees with itself. On the Qwen3.8-2.4T MoE head kernel this produced a 1.55x isolated "win" whose e2e was 1.0034x — the speedup was the branch flip, not the patch.

What changes

Preserve the dispatch metadata (capture_shapes, harness_lib)

• _snapshot records JSON-safe loader-set tensor attributes under attrs; reconstruct_captured re-applies them after the device move. Values that cannot round-trip (tensors, modules, callables) are dropped rather than repr'd — a restored repr string still reads truthy to a getattr(w, "...", False) gate, which is worse than an absent attribute.
• meta.capture_env records the dispatch-steering env of the capturing process (tuned-config tables, backend enables, arch pins), for comparison only and never re-exported: replaying with a tuned table the captured server did not have is its own infidelity. Credential-looking names are recorded with the value <redacted> — this runs inside the serving process.

Prove the baseline leg is deployment (harness_lib, leg_runner)

• leg_runner --mode dispatch reports the GPU kernels a leg actually launches, read off a one-shot torch.profiler trace.
• assert_baseline_dispatch, called from measure_legs, refuses to measure unless the baseline leg launches meta.device_kernel. Never the candidate — changing which kernel runs is what a candidate is for.
• A mismatch raises HarnessIncompleteError (exit 3, "regenerate the UT"), never a correctness FAIL: it is a UT-generation defect, not a reason to reject a kernel. The message warns against "fixing" it by exporting a tuned config the captured server did not have.
• It degrades to {"checked": False, "why": ...} — never to a verdict — when evidence cannot be gathered: no device_kernel, kernel_selection.py not vendored, an empty profile, or a dispatch leg that does not run.
• Name matching reuses kernel_selection.kernel_matches, which must not drift from canonicalDeviceKernel in e2e_workflow.js.

Retrofit for oracles already on disk (harness_lib)
The flag is not in existing .pt files and cannot be recovered from them, so without this every pre-existing oracle could only fail the new gate with no remedy short of recapturing from a server that may be gone. apply_declared_attrs applies meta.live_tensor_attrs = {operand: {attr: value}} to a rehydrated bundle, addressing operands by kwarg name or "pos[<i>]". It is also the answer when a rehydrator cannot preserve attrs even from a fresh capture: a packed-fp4 operand is materialized by .view(dt)/.set_(), each returning a fresh tensor with an empty __dict__.

iter_eager_cases_from_oracle/eager_cases_from_oracle take an optional meta so the declaration reaches the frozen correctness cases too — applying it to only the timing legs would have correctness grading CK while timing measured FlyDSL, the very split this closes. A declaration matching no tensor operand raises: a near-identical helper already shipped hand-written in one task's cases.py, returned early on an empty spec, and had therefore never once executed.

Do not blame the candidate for the baseline's own nondeterminism (harness_lib)
Following this PR's remedy on the motivating task turns a green UT red. The fixed baseline runs FlyDSL MoE, which accumulates with atomics and does not reproduce itself: recording it twice at the same seed on gfx950 differs by 1.5e-05 absolute on prefill_M16384, and correct's atol = tol*RMS(ref) inflates that into a 0.5594 "relative error" on near-zero elements — exactly the number the candidate was failed for.

baseline_noise_floor records the baseline a second time at the same seed and scores it with the metric the candidate is judged by; check_random_vs_baseline passes (and labels) a case that misses tol but stays inside noise_margin (2x) of that floor. Absent a floor the gate is exactly as strict as before. Cost: one extra oracle leg, ~15-25 s.

Extractor role (kernel_extractor.md)
Vendor kernel_selection.py; require cases.py to rehydrate via h.reconstruct_captured instead of hand-rolling the walk — a hand-written rehydrator is exactly how the attribute is lost. The exit-code contract now covers both raisers: measure_legs is a second source of HarnessIncompleteError, but the documented main() wrapped only run_correctness, so a dispatch mismatch would have escaped as a traceback and scored exit 1 — the code that drops the head — instead of exit 3.

Verification

Unfixed vs fixed against the motivating task, same candidate:

	isolated speedup	correctness
unfixed (baseline dispatches to CK)	1.5470	PASS
fixed (baseline dispatches to FlyDSL)	1.0605	FAIL, 1 draw of 18

The win collapses, which is the point. The residual FAIL was reproduced with no candidate in the comparison at all — two baseline recordings at the same seed differ by 0.5594 on the same draw, while the CK baseline's own spread is ~8x quieter (max_abs 1.9e-06). Replaying the gate over those two real recordings: the false FAIL goes away with noise_floor on, no other case changes verdict, and an injected uniform 5%-of-RMS error still fails every case.

Report-only sweep, nine existing task dirs (gfx950, sglang v0.5.18, 7-86 s each):

verdict	tasks	what it means
PASS	2 Qwen MoE controls carrying live_tensor_attrs	the fix holds
PASS, but vacuous	1 TileLang DS-V4 task	its device_kernel is main_kernel, which TileLang gives every jitted prim_func
raised	3 Qwen MoE tasks without the declaration	the bug: baseline runs ck::kernel_moe_mxgemm, meta names the FlyDSL kernel. One was not previously known to be affected
raised	2 tasks (Qwen t32x64x256, MiniMax t64x128x256)	same kernel family, different tile: the oracle's shapes never reach the bucket the name came from. Real infidelity the message's causes did not cover — fixed in 614c3c4
unchecked	1 Gemma task	its cases.py cannot call the target under a current image, so the leg does not run and the gate degrades instead of blocking — intended behaviour, on a real failure

Five of nine would exit 3 on their first run after this lands; all five are genuine.

End to end against a live server, not a replayed artifact: Qwen3.8-2.4T MXFP4 on gfx950, TP=8/EP=8, capture overlay on aiter.fused_moe:fused_moe. The server loads, decode CUDA-graph capture completes (258.77 s), the benchmark runs to usable_for_acceptance=true, and the capture flushes 4 cases across prefill and decode with oracle_complete=true. With no hand-written live_tensor_attrs, the oracle carries w1/w2 attrs: {is_shuffled: True} straight off the loader. Replaying record 0 on one GPU:

replay arm	MoE kernel actually launched
as captured	mfma_moe1_silu_mul_afp4_wfp4_bf16_t32x128x256_pm1_async_v32 (+ the matching mfma_moe2_*) — FlyDSL
recorded attrs stripped (pre-fix behaviour)	ck::kernel_moe_mxgemm<...> ×2 — CK

The only difference between the arms is that attribute. _env_snapshot inside the sglang image: 22 vars, all dispatch-relevant, an injected GEAK_KB_STORE_TOKEN came back <redacted>.

Getting that run to exist took two further fixes, both of which are "the overlay cannot observe an MXFP4 server it breaks or silently misses":

• Sub-byte dtypes could not be copied at all (8428dc92). ROCm torch 2.9 has no copy_kernel for float4_e2m1fn_x2, so _snapshot's D2H raised, the blanket except swallowed it, and an MXFP4 capture recorded zero cases while reporting success. Both directions now copy through a uint8 view.
• The overlay perturbed what it measured (609d81b8). install() runs from the overlay's sitecustomize, i.e. during interpreter startup, and used to import the target module there to wrap it — reordering every later import in the process. Pulling aiter.fused_moe in that early makes FlyDSL's JIT abort the server with LLVM ERROR: Do not know how to expand this operator's operand! on a kernel the identical un-overlaid run compiles fine. Reproduced standalone (one GPU, no model, ~90 s) and bisected to that single import. install() now validates eagerly on what needs no import and arms a sys.meta_path shim that binds when the application imports the module; the non-Python-callable refusal moves to bind time and still raises loudly rather than letting a native callable SIGSEGV the server.

Two deliberate limits: a task passes if any case launches the named kernel (tile selection is shape-dependent, so per-case strictness would reject correct tasks) — uncovered cases are reported in unmatched_cases; and a generic device_kernel such as main_kernel makes the check vacuous rather than wrong.

Also in this PR: the head that never reached a measurement

Everything above fixes a head that measured the wrong kernel. Two further commits fix the other half of the same session. In the 20260907 Qwen3.8-2.4T-A95B-Quark-MXFP4 run, four heads were extracted and one reached the optimization lane (63 min); moe1_silu spent 165 min over 6 attempts and gemm 816 min — 13 h 34 m — over 7 attempts, neither producing a candidate.

A mis-transcribed kernel name was diagnosed as a wrong seam (8d681cfa). kernel_selection.py had one failure code for two unrelated defects. A marker that launched no GPU work means the seam is wrong and the extractor should descend. A marker that demonstrably launched kernels, none matching the declared name, means the seam is live and what is wrong is the hand-transcribed device_kernel string — one _HA_S_ typed as _AS_ is invisible to every other check. Both took the same corrective: descend, carrying ATTEMPTED_TARGET_CALLABLES with "Do not return any value again." That ban is the amplifier — it permanently forbids the correct seam from attempt 2 on, leaving the retry loop with no terminating condition. gemm's selection_validation.json shows the seam it re-hunted for 13.6 h was already right on attempt 1.

• device_kernel_not_under_target (descend) is now distinct from device_kernel_name_mismatch and device_kernel_not_in_profile (keep the seam, fix the name).
• A name-only defect no longer bans the seam; the corrective quotes both the kernels the seam launched and the names in PROFILE_TOPN.
• --check-device-kernel runs before capture: the declared name must appear in a profiler row. No GPU, no server, no capture needed to reject a typo. Same matcher as the run-time gate, so there is still one copy.

Extraction had no wall-clock bound at all (e362c900). A retry count does not bound cost. The real envelope is a product never multiplied out: BASELINE_EXTRACT_RETRIES (3, so 4 invocations) × safeAgent's internal tries (3) × AGENT_TIMEOUT_MS (2 h with no global budget) = 24 h on one head — and the global ELAPSED clock is only armed when time_budget_s is passed, so a budget-less run has no guard at all.

• One wall-clock budget per head, armed before the first attempt and spanning every retry including safeAgent's internal ones. Default 120 min via args.capture_budget_s, ~1.9× the only extraction in that session that succeeded.
• Each invocation is capped at what is left (timeoutCapMs); otherwise one invocation could run the full 2 h hung-guard and blow a 2 h budget by itself.
• The budget is checked before paying for an attempt, and capture_min_attempt_s (default 15 min, a measured server boot) refuses to start one it cannot finish.
• A head cut by the budget lands on the normal failure path carrying capture_budget_spent, so "ran out of time" is never read as "unextractable".
• The stopwatch reads elapsed from setTimeout rungs armed at absolute offsets and taken as a monotonic max, because Date.now() / new Date() / Math.random() all throw in a Workflow script.

The three durations are read off the session's logs and selection_validation.json. That 8d681cfa would have cut gemm's 13.6 h, and the ~12.3 h e362c900 would have saved, are inferences over three observation points — neither is a measurement. moe1_silu's 165 min is a third root cause neither commit addresses: 5 of its 6 retries are host OOM materializing a 23 GB recorded oracle. Out of scope here.

Tests

python3 -m unittest discover -s e2e_workflow/scripts/tests: 1255 tests, all pass (one pre-existing collection error from a pytest import in test_e2e_store.py, unrelated). All 8 JS suites pass.

test_selection_name_mismatch.js (11) and test_capture_budget.js (19) extract the real definitions out of e2e_workflow.js — a Workflow script, so it cannot be require()d — and execute them, so editing a test alone cannot satisfy them. The budget tests build the stopwatch over an injected setTimeout and fire the rungs by hand; without that, a stopwatch whose elapsed never updates passes everything else and silently restores the 24 h behaviour.

9e9c6ed0 covers the dispatch-evidence path itself: observed_device_kernels against hand-written chrome traces (only cat: kernel rows are read — cpu_op/ac2g rows carry kernel-ish names and would let any seam certify itself), --mode dispatch profiling the compiled callable when the deployment compiles one, and profile_kernel_names over all four topN document shapes (reading only one yields an empty list, which fails the pre-capture check OPEN).

Migration note for reviewers

assert_baseline_dispatch only fires for tasks whose meta.device_kernel is set and whose dir vendors kernel_selection.py — freshly extracted tasks. Older task dirs degrade to unchecked. Pre-PR oracles that do get regenerated are repaired either by recapturing (preferred — a fresh capture now records attrs on its own) or by declaring meta.live_tensor_attrs with the deployment value; the role doc requires recording why that value is deployment's (a server.log line, a profile), since the gate proves only that the result reaches the right kernel, not that the declared value is right.

Regenerating a UT against the dispatch gate should wire baseline_noise_floor at the same time: the gate's own remedy is what exposes an atomics-nondeterministic baseline, and without a floor that lands as exit 1 on the next author.