moe_timing: route on device in the campaign's replay, and gate on it

[qichao-arlo-wang]

qwen3_trace_replay is the only decoder-level program the timing campaign measures. It read topk_indices out of the trace and preloaded them into INT SRAM, so V_TOPK never appeared in a measured program. router_topk was a name in MOE_STAGES that nothing emitted, and every cycle count the campaign produced excluded selection.

	before	after
V_TOPK in the measured program	0	one per token
router_topk in the stage profile	stage never appeared	8 instructions for 2 tokens
functional gate on routing	none	expert ids must equal the trace's, exactly
CI coverage of this program	none	built + run on every push

The trace has no logits

generate_true_routing_with_weights runs the real router, takes topk + softmax, and keeps only those. Regenerating logits means a Qwen3-30B-A3B forward pass, which is what replaying a trace exists to avoid.

router_logits rebuilds a BF16 row that reproduces the recorded top-k instead:

• log-weights in the recorded order — log inverts softmax up to the additive constant the shift supplies
• a strict one-ulp ladder where BF16 flattens a tie. build_route_traces substitutes uniform weights for timing-only runs, which is entirely ties, and V_TOPK breaks ties by low index — so [9, 2, 5] would come back [2, 5, 9] without it
• an unselected floor, and the whole row shifted positive. prestage_bf16_vram_matrix pads with zeros and raw log-weights are all ≤ 0, so a row laid out one column short would have a zero pad outrank every real expert, silently and with plausible numbers
• a refusal, naming the token, when the recorded order and the recorded weights cannot both hold

What it does not rebuild: the 120 unselected logits get a floor, not their true values. Selection and weighting are exercised; the tail's spread is not.

The gate is the point

The existing gate compares an all-zero accumulator against an all-zero golden — which holds for any routing, because the expert weights are dummy zeros. Injecting three router faults:

injection	new router gate	pre-existing smoke gate
per-token INT base dropped	fail (16 mismatches)	pass
every token selected from row 0	fail (8 mismatches)	pass
router not emitted at all — main today	fail (16 mismatches)	pass
control	pass	pass

Expert ids are integers, so that half of the gate has no tolerance in it. The weights are checked against softmax over the BF16 logits actually staged, not against the trace's — folding the reconstruction's own bound in here as well would test the reconstruction twice and the device not at all.

First CI coverage for this program

It had none: every real trace needs the 30B checkpoint. synthetic_trace draws Gaussian logits and hands them to the same trace_from_record and the same schema validation a captured trace goes through, so the shape is real even though the routing distribution is not. Gaussian rather than uniform on purpose — it puts near-ties in the top-k tail, which is where the BF16 reconstruction is hardest.

Verification

• 21 unit guards. Each was injected with the defect it names and confirmed to fire — 8 injections on the reconstruction, 4 on the layout, plus the 3 end-to-end above
• end to end at MLEN 128 (the campaign's shape, one row per token) and at MLEN 64, where 128 experts span two rows per token and the fold is not the identity. expert_ids_match: true on both; max weight error 3.3e-4
• just test-timing-gates still exits 0, all four gates pass
• ruff format --check / ruff check clean across 103 files; the pytest-only CI job 23 passed

test_every_guard_file_is_wired_into_ci caught that the new test file was torch-free and therefore being pulled into the pytest-only job, where it cannot run. That was a real hole rather than a guard to work around: the MLEN 64 fold had no coverage at all, so the guards that close it are the ones that make the torch dependency honest.

Not in scope

The router GEMM. The trace kept the top-k, not the logits, so hidden -> num_experts cannot be replayed faithfully — emitting it with dummy weights would produce arbitrary selection and destroy the trace's routing distribution, which is the whole point of trace replay. router_gemm_included: false is now in the manifest, and the routing-tax note in export_pilot_results says so instead of claiming top-k is excluded too.

🤖 Generated with Claude Code