From d6e067b05509126253f705c73c866040a5141dfc Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:48:45 -0400 Subject: [PATCH 01/38] docs: scope Phase 5a's agreement claim instead of leaving it an open audit Phase 5a ran on transformers 4.57.6 with finite, non-zero logit diffs, so it is not the RoPE-bug degenerate-output artifact an earlier draft feared. What stands is that 29 positions in one run cannot rule out the near-tie sensitivity 5b found in the same int8 kernel. Resolved as an input to Phase 6: its correctness gate measures agreement at scale on the benchmarked config. --- CHANGELOG.md | 9 ++-- CLAUDE.md | 11 +++-- README.md | 2 +- docs/STATUS.md | 29 +++++++---- .../2026-09-17-phase-5a-quantization-run.md | 11 +++++ ...09-18-phase-5b-speculative-decoding-run.md | 49 ++++++++++++------- 6 files changed, 75 insertions(+), 36 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b5f65e..a852694 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,9 +34,12 @@ once a phase's exit criteria are actually met). 2/4 prompts in one process and 3/4 in another. Root-caused by a batch-width probe and a 10-trial determinism probe to a near-tied logit under the int8 kernel's floating-point precision, not a logic error. - Open, deliberately unaudited: Phase 5a's "perfect top-1/top-k agreement" - used the same kernel and a single-run check that could not see this. - Two findings-doc figures were also corrected after re-verifying against + Phase 5a's "perfect top-1/top-k agreement" used the same kernel but is a + small sample (29 positions, one run, `transformers==4.57.6`, so not + affected by the RoPE bug; its logit differences were finite and + non-zero, so it was not a degenerate-output pass) -- too few positions + to rule this sensitivity out, so Phase 6's correctness gate re-checks + agreement at scale rather than leaning on 5a's number. Two findings-doc figures were also corrected after re-verifying against the results JSONs (see the git history of the findings doc). Cost: **$12.26** across both sessions against a $10 cap that one pod exceeded and that was raised to $20 with disclosure. Also: `docs/findings/` split diff --git a/CLAUDE.md b/CLAUDE.md index 1865dd8..64c891f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -162,10 +162,13 @@ the 8-config sweep draft-model matched in 13 of 16 (prompt, k) combinations and prompt-lookup in 9 of 16, and the same k=4 prompt-lookup config matched 2/4 prompts in one process and 3/4 in another. Root cause, confirmed by two GPU probes: a genuine near-tied logit under the int8 -kernel's floating-point precision, not a logic bug in the loop. **Open -question, deliberately not audited:** Phase 5a's own "perfect top-1/top-k -agreement" claim used this same kernel and a single-run check that could -not see this. Cost: **$12.26** across both sessions, against a $10 cap +kernel's floating-point precision, not a logic bug in the loop. **Scope +of Phase 5a's claim:** its "perfect top-1/top-k agreement" used this same +kernel but is a small sample (29 positions, one run, on +`transformers==4.57.6`, so unaffected by the RoPE bug; finite non-zero +logit diffs, so not a degenerate-output pass) -- too few to rule out +this sensitivity, so Phase 6's correctness gate re-checks agreement at +scale on the benchmarked config instead of leaning on 5a's number. Cost: **$12.26** across both sessions, against a $10 cap that was exceeded on one pod and raised to $20 with disclosure. Full account: `docs/findings/phase-5b/2026-09-18-phase-5b-speculative-decoding-run.md`. diff --git a/README.md b/README.md index 61c7cc1..129b7fa 100644 --- a/README.md +++ b/README.md @@ -165,7 +165,7 @@ including the ones that came back null or mixed: | **2 — vLLM contribution** | Do vLLM's or SGLang's official MoE benchmarks model skewed (zipf) expert load? | Neither did. Upstreamed an opt-in `--expert-load-distribution` flag to vLLM; on one rented RTX 3090, `--tune` under zipf vs. uniform routing picks a **different winning Triton config at 4 of 5 tested batch sizes**. Open: [vllm-project/vllm#57100](https://github.com/vllm-project/vllm/pull/57100) | | **3 — multi-GPU EP** | Does Phase 1's naive-vs-persistent crossover hold under DeepEP's real per-expert token distribution across GPUs? | **It doesn't apply at real scale.** Naive wins at every tested token count (16-2048) on H200, independent of EP. Real DeepEP dispatch across the model's 27 MoE layers produces per-local-expert counts (median 2, max ~20) far below Phase 1's smallest tested point (16) | | **4 — disaggregated prefill/decode** | Does separating prefill and decode across GPU pools relieve contention, holding GPU count fixed at 4? | **Mixed, not a clean result.** Disaggregated TTFT beats co-located at concurrency 4 (0.46s vs 1.12s) but loses at concurrency 8 (0.98s vs 0.70s) — most likely a kernel-warmup confound between independent process launches, reported as genuinely inconclusive rather than forced | -| **5a — int8 quantization** | Does a self-computed, per-channel int8 weight-only kernel cut memory with no model-level correctness cost? | **Yes on both counts, no speedup.** **49.89%** expert-weight memory reduction, perfect top-1/mutual-top-k agreement vs. the naive bf16 kernel at every tested position. Throughput was a near-exact tie with naive (**-1.2%**), as expected — weight-only quantization saves memory bandwidth, not FLOPs. A real bug (quantizing without freeing the original bf16 weights, OOMing a 44GB L40) was found and fixed mid-session | +| **5a — int8 quantization** | Does a self-computed, per-channel int8 weight-only kernel cut memory with no model-level correctness cost? | **Yes on both counts, no speedup.** **49.89%** expert-weight memory reduction, perfect top-1/mutual-top-k agreement vs. the naive bf16 kernel at every tested position (29 positions, one run -- a small sample; see 5b for the near-tie sensitivity it could not rule out). Throughput was a near-exact tie with naive (**-1.2%**), as expected — weight-only quantization saves memory bandwidth, not FLOPs. A real bug (quantizing without freeing the original bf16 weights, OOMing a 44GB L40) was found and fixed mid-session | | **5b — speculative decoding** | Does speculative decoding (a 7B draft model, or model-free prompt-lookup) speed up single-request decode on the int8 target while reproducing plain greedy output exactly? | **Partly, and not exactly.** A40, bf16, unbatched, baseline **25.18 tok/s**: only prompt-lookup beats it (**34.4-50.9 tok/s** across k=1-8, +77.7% at k=4); the 7B draft model is *below* baseline at every k (19.6-23.0) despite 2-7x higher acceptance. Output matched plain greedy in 13 of 16 (prompt, k) combinations for draft-model and 9 of 16 for prompt-lookup; the same k=4 config matched 2/4 prompts in one process and 3/4 in another. Root-caused to near-tied logits under the int8 kernel's floating-point precision, not a logic bug. The first session's numbers were withdrawn by this project's own review (degenerate baseline) and root-caused to a `transformers` bug leaving DeepSeek's RoPE buffers uninitialized | Full method and every number: one folder per phase in diff --git a/docs/STATUS.md b/docs/STATUS.md index 260b3e3..50b36b3 100644 --- a/docs/STATUS.md +++ b/docs/STATUS.md @@ -486,13 +486,23 @@ $10.8848; `5ufi854zxkfwbz` (A40) $0.2022; `d2h1sgmr5wjssv` (A40, where the real baseline/gates/k-sweep ran) $0.49. All four pods stopped or terminated. -**Open question, flagged for the project owner, not audited here**: -this session's root-cause work found concrete evidence of near-tie -floating-point sensitivity in `grouped_matmul_int8`, the same quantized -kernel Phase 5a's own "perfect top-1/mutual-top-k agreement" claim was -measured against. That claim may not be wrong, but the single-run -top-k check it relied on would not have caught this specific failure -mode. Re-auditing Phase 5a was out of this task's scope. +**Scope of Phase 5a's agreement claim (resolved as an input to Phase 6, +2026-09-18)**: this session's root-cause work found concrete evidence of +near-tie floating-point sensitivity in `grouped_matmul_int8`, the same +quantized kernel Phase 5a's "perfect top-1/mutual-top-k agreement" claim +was measured against. Checked against the record, without new GPU spend: +Phase 5a's GPU session ran on `transformers==4.57.6` (a pod-side +override of this repo's `>=5.17.0` floor), the version the RoPE +`inv_freq` bug was never observed on, and its `max_abs_diff` values +(2.125 / 1.906 / 1.344, one per prompt) are finite and differ per +prompt, so it was not a degenerate-output pass. (That 4.57.6 is +unaffected was inferred from those numbers, not tested directly.) What +does stand: the check covered 29 positions in a single run, quantized +vs. naive bf16 -- too small a sample to rule out the 0.1-0.4 logit-gap +near-ties 5b found. Decision: do not re-rent a GPU just to re-audit 5a; +Phase 6's own correctness gate measures agreement at scale on the exact +config being benchmarked, and 5a's agreement figure is not to be quoted +in the head-to-head as more than "true on a small sample". ## Next step @@ -503,5 +513,6 @@ own branch per the one-branch-per-phase convention. Phase 3's PR (#3) also merged. Phase 2's PR is still open and awaiting maintainer review; no further work planned on it beyond responding to review feedback. Phase 6 (final -benchmark vs. vLLM/SGLang) planning can proceed, and should treat the -Phase 5a audit question above as an open input, not a blocker. +benchmark vs. vLLM/SGLang) planning can proceed; its correctness gate +must include an at-scale agreement check on the benchmarked config (see +"Scope of Phase 5a's agreement claim" above). diff --git a/docs/findings/phase-5a/2026-09-17-phase-5a-quantization-run.md b/docs/findings/phase-5a/2026-09-17-phase-5a-quantization-run.md index 29b8d62..a9b1df0 100644 --- a/docs/findings/phase-5a/2026-09-17-phase-5a-quantization-run.md +++ b/docs/findings/phase-5a/2026-09-17-phase-5a-quantization-run.md @@ -276,3 +276,14 @@ abandoned attempts) included. Full records: `docs/findings/phase-5a/2026-09-17-phase-5a-community-attempt1-cost.md`, `2026-09-17-phase-5a-community-attempt2-cost.md`, `2026-09-17-phase-5a-quantization-cost.md`. + +## Addendum, 2026-09-18 (Phase 5b) + +The "perfect top-1/mutual-top-k agreement" above covers 29 positions in a +single run. Phase 5b later measured near-tie floating-point sensitivity +(0.1-0.4 top-2 logit gaps flipping an argmax across process launches) in +this same `grouped_matmul_int8` kernel, which a sample this small cannot +rule out. Read the agreement result as "true on a small sample", not +"shown wrong". It is not a RoPE-bug artifact: this phase's GPU session +ran on `transformers==4.57.6` and its logit diffs are finite and non-zero. See +`docs/findings/phase-5b/2026-09-18-phase-5b-speculative-decoding-run.md`. diff --git a/docs/findings/phase-5b/2026-09-18-phase-5b-speculative-decoding-run.md b/docs/findings/phase-5b/2026-09-18-phase-5b-speculative-decoding-run.md index 09db489..9032886 100644 --- a/docs/findings/phase-5b/2026-09-18-phase-5b-speculative-decoding-run.md +++ b/docs/findings/phase-5b/2026-09-18-phase-5b-speculative-decoding-run.md @@ -295,21 +295,31 @@ from committed files the way the aggregate metrics are. Lesson applied going forward: pull evidence files off a rented pod before any `stop`, not after. -## Open question, explicitly flagged for the project owner, not pursued here - -This session's root-cause work directly strengthens a concern first -raised (but not investigated) earlier in this session: **Phase 5a's own -"perfect model-level top-1/mutual-top-k agreement" claim used this exact -same `grouped_matmul_int8` kernel**, and this session found concrete, -measured evidence that this kernel has genuine near-tie floating-point -sensitivity across separate process launches -- exactly the failure -mode a single-run top-k agreement check cannot see. This is not a claim -that Phase 5a's result is wrong; it is a claim that Phase 5a's own -correctness check was not built to detect this specific failure mode, -and now there is direct measurement (not just a theoretical concern) -that the failure mode is real for this kernel. Auditing Phase 5a is out -of this task's scope and was not attempted -- flagged here for a -decision by the project owner. +## What this means for Phase 5a's agreement claim + +**Phase 5a's "perfect model-level top-1/mutual-top-k agreement" claim +used this exact same `grouped_matmul_int8` kernel**, and this session +found concrete, measured evidence that the kernel has genuine near-tie +floating-point sensitivity across separate process launches. An earlier +draft of this section worried Phase 5a's result might be vacuous (the +RoPE bug producing degenerate output that trivially "agrees"). Checking +the Phase 5a record showed that worry does not hold: + +- Phase 5a's GPU session ran on `transformers==4.57.6`, a pod-side + override of this repo's `>=5.17.0` floor (its own Environment section + and runbook say so). The RoPE `inv_freq` bug was observed only under + `5.17.0`. That 4.57.6 is unaffected is inferred, not tested directly. +- Its `max_abs_diff` values (2.125 / 1.906 / 1.344, one per prompt) are + finite, non-zero and different per prompt; NaN or all-token-0 logits + would not produce that. + +What does stand: that check covered 29 token positions (11 + 11 + 7) in +a single run, comparing quantized against naive bf16. That is too small +a sample to rule out near-ties like the 0.1-0.4 top-2 logit gaps found +here, so Phase 5a's claim is "true on a small sample", not "shown wrong". +No GPU money was spent re-auditing it; instead Phase 6's correctness gate +measures agreement at scale on the exact config it benchmarks, and 5a's +figure should not be quoted in the head-to-head as more than that. ## Summary @@ -346,7 +356,8 @@ runs per config on one pod; the k=4 prompt-lookup repeat (47.21 vs. **Total real GPU cost across both Phase 5b sessions: $12.26 of the $20 cap (61.3%)**, including a disclosed and corrected cost-cap incident -earlier in this session. **One open question is flagged, not resolved, -for the project owner**: whether Phase 5a's own correctness claim should -be re-examined in light of this session's concrete evidence of near-tie -sensitivity in the same quantized kernel. +earlier in this session. **Phase 5a's agreement claim** (same kernel) is +scoped, not overturned: it ran on `transformers==4.57.6`, so it is not +the RoPE-bug artifact an earlier draft feared, but 29 positions in one +run cannot rule out the near-tie sensitivity found here -- Phase 6's +correctness gate re-checks agreement at scale instead. From ae76cf222bddaf75002fa813e2f17d35b3e1f8c2 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:10:23 -0400 Subject: [PATCH 02/38] docs: Phase 6 design -- kernel-level race vs vLLM/SGLang plus engine reference Design doc section 6 assumed a dispatch server for the serving benchmark to target; none exists, so an end-to-end race would measure HF-eager overhead rather than the kernel. Phase 6 instead races the kernels like-for-like on one L40 and reports vLLM/SGLang serving as a labeled reference. Amends design section 6 and records the decision in STATUS.md. --- docs/STATUS.md | 35 ++- .../2026-09-14-dispatch-system-design.md | 9 + .../2026-09-18-phase-6-final-benchmark.md | 266 ++++++++++++++++++ 3 files changed, 306 insertions(+), 4 deletions(-) create mode 100644 docs/design/2026-09-18-phase-6-final-benchmark.md diff --git a/docs/STATUS.md b/docs/STATUS.md index 50b36b3..1c7fd5d 100644 --- a/docs/STATUS.md +++ b/docs/STATUS.md @@ -504,6 +504,31 @@ Phase 6's own correctness gate measures agreement at scale on the exact config being benchmarked, and 5a's agreement figure is not to be quoted in the head-to-head as more than "true on a small sample". +## Phase 6 progress + +**Design approved, 2026-09-18:** +`docs/design/2026-09-18-phase-6-final-benchmark.md`, on branch +`phase-6-final-benchmark`. No implementation plan or code yet, no GPU +spent. Decided in brainstorming: + +- **Kernel-level race plus a labeled engine reference, not an end-to-end + race.** Design doc §6 assumed a dispatch server for the serving + benchmark to target; none exists (the engine is an in-process HF eager + loop), so an end-to-end race would measure that gap, not the kernel. + §6 is amended to say so. A real server belongs to Phase 7. +- **One L40, ~$10 cap** (same GPU class as Phase 1 and 5a's kernel + numbers). Kernel race: dispatch naive/persistent vs. vLLM and SGLang + fused-MoE, uniform and zipf routing, bf16 and weight-only int8. Engine + reference: vLLM and SGLang serving the full model, concurrency 1/4/16/64, + dispatch at concurrency 1 only. +- **Checked live 2026-09-18:** SGLang serves `DeepseekForCausalLM` (V1); + both engines' fused-MoE accept `use_int8_w8a16` + `per_channel_quant`; + SGLang's fused-MoE code was recently reorganized, so engine versions get + pinned and each engine sits behind one adapter. +- **Stage 1 blocks the rest:** the at-scale correctness gate STATUS owes + from 5a/5b, with a large-gap disagreement threshold fixed in the plan + from 5b's near-tie data before any GPU time. + ## Next step Phase 5b is done: root cause fixed, real numbers measured, both @@ -512,7 +537,9 @@ merged (PR #4), as are Phase 5a (PR #5) and Phase 5b (PR #6), each on its own branch per the one-branch-per-phase convention. Phase 3's PR (#3) also merged. Phase 2's PR is still open and awaiting maintainer review; no further -work planned on it beyond responding to review feedback. Phase 6 (final -benchmark vs. vLLM/SGLang) planning can proceed; its correctness gate -must include an at-scale agreement check on the benchmarked config (see -"Scope of Phase 5a's agreement claim" above). +work planned on it beyond responding to review feedback. + +Phase 6's design is approved (see "Phase 6 progress"). Next: write its +implementation plan in `docs/plans/`, fixing the correctness gate's +logit-gap threshold from Phase 5b's recorded data, then execute it task +by task. diff --git a/docs/design/2026-09-14-dispatch-system-design.md b/docs/design/2026-09-14-dispatch-system-design.md index 506eb0f..efaae80 100644 --- a/docs/design/2026-09-14-dispatch-system-design.md +++ b/docs/design/2026-09-14-dispatch-system-design.md @@ -133,6 +133,15 @@ Every comparison runs vLLM/SGLang and this project's server against the **same model, same hardware, same request trace** — an apples-to-apples requirement, not a nice-to-have. +> **Amended 2026-09-18 (Phase 6 design).** "This project's server" was +> never built through Phase 5b: dispatch's engine is an in-process +> Hugging Face eager decode loop with the MoE layer patched, so a +> serving-benchmark race against vLLM/SGLang would measure that gap rather +> than the kernel. Phase 6 instead runs a like-for-like kernel-level race +> plus a labeled engine reference; see +> `docs/design/2026-09-18-phase-6-final-benchmark.md`. A real dispatch +> server belongs to Phase 7 (§8). + ## 7. Phase plan | Phase | What | Cost shape | diff --git a/docs/design/2026-09-18-phase-6-final-benchmark.md b/docs/design/2026-09-18-phase-6-final-benchmark.md new file mode 100644 index 0000000..2850f5c --- /dev/null +++ b/docs/design/2026-09-18-phase-6-final-benchmark.md @@ -0,0 +1,266 @@ +# Phase 6 design: final benchmark vs. vLLM and SGLang + +Per `docs/design/2026-09-14-dispatch-system-design.md` §7 (phase plan, row +6: "Final benchmark vs. vLLM/SGLang, standard tool, same hardware/model"). +Phases 0-5b are complete and merged (Phase 2's upstream PR is open). Every +earlier phase deferred the head-to-head comparison to this one -- see the +"Phase 6 owns that comparison" non-goals in the Phase 4, 5a and 5b designs. + +## 1. Purpose and thesis + +Every number this project has published so far was measured against its +own baseline: stock `moe_infer`, the naive kernel, the bf16 target. None +was measured against a production engine. This phase does that. + +**Thesis: on one L40, at DeepSeek-MoE-16B's real routed-expert shape, how +do dispatch's Triton grouped-GEMM kernels compare with vLLM's and SGLang's +fused-MoE on latency -- across token counts, uniform vs. zipf routing, and +bf16 vs. weight-only int8 -- and where does a full production engine land +on the same model and hardware?** + +The result is a table, not a headline. Nothing is summarized as "beats +vLLM" or "loses to vLLM" unless the per-shape table shows it shape by +shape. This repo has published null and negative results before (Phase 1's +persistent kernel, Phase 3's disproved hypothesis, Phase 4's inconclusive +TTFT) and this phase is designed to tolerate the same: production +fused-MoE kernels are heavily tuned and may simply win at some or all +shapes. + +## 2. Scope, decided during brainstorming + +- **Kernel-level race plus an engine reference, not an end-to-end race.** + Design doc §6 assumed "this project's own server" would exist to point + the standard serving benchmark at. It does not: dispatch's engine is a + Hugging Face eager decode loop with the MoE layer patched -- unbatched, + single-request, no CUDA graphs, no paged KV cache. An end-to-end race + against vLLM/SGLang would measure that gap, not the kernel, which is + the part dispatch actually built. So the like-for-like comparison is at + the kernel/layer level, and the full-engine numbers are a labeled + reference. (Options considered and not chosen: building an + OpenAI-compatible dispatch server for a literal three-way end-to-end + race -- expected to lose on HF-eager overhead rather than kernel + quality; swapping dispatch's kernel into vLLM as a custom MoE backend -- + strongest claim if it works, highest risk, plausibly its own + multi-day integration.) +- **Hardware: one L40, ~$10 cap.** Same GPU class as Phase 1's and Phase + 5a's measured kernel numbers, so this phase extends those tables. 48GB + fits the 16B bf16 model for the engine reference. Rejected: H100 (breaks + continuity, several times the hourly cost) and an L40-plus-H100 split + (the reference would no longer share hardware with the race, which §6 + of the system design forbids). +- **Weight-only int8 gets a like-for-like row.** Checked live 2026-09-18: + both vLLM's `fused_experts_impl` and SGLang's `fused_experts` accept + `use_int8_w8a16` and `per_channel_quant`, the same scheme Phase 5a's + kernel implements (per-output-channel int8 weights, bf16 activations). +- **SGLang serves this model.** Checked live 2026-09-18: SGLang's + `python/sglang/srt/models/deepseek.py` defines `DeepseekForCausalLM`, + the V1 architecture `deepseek-ai/deepseek-moe-16b-base` uses, and routes + its experts through `fused_moe`. +- **Routed experts only.** The timed region matches Phase 1's: the routed + experts' work. DeepSeek's shared experts are a dense MLP separate from + the routed path and are not part of any engine's fused-MoE call. + +## 3. Deliverables, in run order + +Stage 1 gates stage 2; stage 2 gates nothing but is run before stage 3 so +that the cheapest, highest-value evidence is on disk first. + +### 3.1 At-scale correctness gate (blocks the rest) + +`docs/STATUS.md` owes this from Phase 5a and 5b: Phase 5a's "perfect +top-1/mutual-top-k agreement" covered 29 positions in one run, and Phase +5b then found genuine near-tied logits (gaps of 0.1-0.4) under the int8 +kernel's floating-point precision. Phase 5a's figure is not to be quoted +as more than "true on a small sample". + +- Real model (`deepseek-ai/deepseek-moe-16b-base`), the exact configs + being benchmarked: stock `moe_infer` (reference) vs. naive bf16, + persistent bf16, and int8, over hundreds of positions. +- Reports top-1 agreement and mutual top-5 agreement as **rates with + counts**, and splits every disagreement by the reference's top-1/top-2 + logit gap. +- **Gate rule:** any disagreement where the reference's logit gap exceeds + a threshold is a failure (it would mean a bug, not a near-tie). + Disagreements below the threshold are reported, not failed -- they are + the known, explained floating-point sensitivity. +- **The threshold is fixed in the implementation plan, from Phase 5b's + recorded near-tie data, before any GPU time is spent -- not chosen after + seeing results.** +- Reuses `scripts/run_baseline.py`'s reference-capture and comparison + path rather than a new one; extends it to a larger, seeded prompt set. + +### 3.2 Kernel race (`scripts/run_engine_race.py`) + +**Compared:** dispatch naive, dispatch persistent, vLLM `fused_experts`, +SGLang `fused_experts`, each at bf16 and (where the engine supports it) +weight-only int8. + +**Matrix:** total token counts per layer call, spanning decode-sized +batches (1-16) through prefill-sized ones (2048); uniform and +zipf routing (reusing `dispatch.kernels.bench.sample_topk_idx`); bf16 and +int8. The exact grid is fixed in the plan; Phase 1's sweep points are the +default so results extend Phase 1's tables. + +**Same inputs, engine-specific adapters.** One seeded input set is +generated once -- `x`, `topk_ids`, `topk_weights`, and the expert weights +-- and a thin adapter per engine converts it to that engine's layout +(vLLM and SGLang take fused gate+up in `w1` and down in `w2`; dispatch +takes gate, up and down stacked separately). Adapters are the only place +that knows an engine's API, because SGLang's fused-MoE code was recently +reorganized (now under `moe_runner/triton_utils/` and +`sglang/kernels/ops/moe/`) and the vLLM signature has already moved from +per-flag arguments to a `FusedMoEQuantConfig` object at the top level. +Exact engine versions are pinned and recorded in every output JSON. + +**Correctness before timing.** Every engine's output is checked against +the fp32 reference (`assert_matches_reference`, as in Phase 1) before it is +timed. An engine that disagrees with the reference on the benchmark's own +input is refused, not timed -- the same rule `run_kernel_bench.py` already +enforces. + +**Timed region.** The whole routed layer from `(x, topk_ids, +topk_weights)` to output, each engine's own path (grouping/alignment, the +GEMMs, activation, weighted combine), plus one bare-GEMM diagnostic per +engine, matching Phase 1's structure. Timing is `triton.testing.do_bench` +(warmup, repetition and synchronization are its job), reporting mean, p50 +and p99 as Phase 1 does. + +**Tuning fairness.** vLLM and SGLang each get their own tuner run on +this exact L40 -- Phase 2's `--tune` path for vLLM, SGLang's +`tuning_fused_moe_triton.py` -- so a baseline is never handicapped by a +missing config for this GPU. Dispatch gets a block-size sweep of +comparable effort (its `block_m`, plus persistent-vs-naive, which is +already a comparison). **The tuned column is primary.** A stock-default +column (whatever each engine picks out of the box, with no tuned config) +is reported as a secondary sensitivity check, since that is what a user +who never runs the tuner gets. Both columns state exactly what tuning was +done and how long it ran. Tuning under zipf vs. uniform routing is itself +a Phase 2 finding (different winning config at 4 of 5 tested sizes), so +each engine is timed under both distributions, and the table says which +distribution each engine's config was tuned under. + +### 3.3 Engine reference (not a race) + +vLLM and SGLang each serve the full model (bf16) on the same L40, in +**separate virtual environments** (their torch pins conflict), driven by +the standard serving benchmark (vLLM's `vllm bench serve` / +`benchmark_serving.py`, which works against any OpenAI-compatible server), +on one seeded trace, at concurrency 1, 4, 16 and 64. Standard metrics: +TTFT, inter-token latency, end-to-end latency, tokens/sec, requests/sec, +and **cost per million tokens computed from the measured GPU-hour rate** +(design §6), never estimated. + +**Dispatch appears at concurrency 1 only**, replaying the identical +token-ID prompts (a seeded trace file the existing harness will need to +accept) through `scripts/run_baseline.py`'s in-process path. That row +carries an explicit caveat in the table itself: the vLLM/SGLang numbers +include HTTP overhead, CUDA-graph capture and torch.compile warmup +effects that dispatch's in-process loop doesn't, and dispatch has no +batched server. Concurrency 4/16/64 rows for dispatch read "n/a: no +dispatch server", not a blank and not a number. Warmup requests are +discarded and the count is recorded (Phase 4's kernel-warmup confound is +the precedent for why). + +## 4. Architecture + +```text +scripts/run_engine_race.py stage 3.2 driver: build inputs, call adapters, + gate on correctness, time, write JSON +src/dispatch/benchmark/engines/ adapter per engine, one common signature: + base.py (x, topk_ids, topk_weights, weights) -> out + dispatch_kernels.py wraps grouped_moe_routed (naive/persistent, + bf16/int8) + vllm_moe.py wraps vLLM fused_experts, layout conversion + sglang_moe.py wraps SGLang fused_experts, layout conversion +scripts/run_baseline.py stage 3.1 (extended): larger seeded prompt + set, gap-split disagreement report +scripts/gpu/phase6_engine_reference.py + stage 3.3 driver: launch server, run the + standard benchmark, tear down, per engine +docs/findings/phase-6/ outputs (default output dir for the new scripts) +``` + +Adapters import their engine lazily, inside the call, so the module stays +importable -- and unit-testable -- on a machine with neither engine +installed (the same pattern `kernels/bench.py` uses for Triton). + +## 5. Testing + +Follows the repo's existing table (CLAUDE.md, "Testing policy"): + +| Layer | What must be covered | +|---|---| +| Adapters (CPU) | layout conversion is lossless and round-trips (gate+up fused == separate); every adapter accepts identical inputs; a fake engine exercises the refuse-not-time path | +| Race driver (CPU) | an engine that disagrees with the reference is refused, not timed; every output JSON carries full config, engine versions, tuning provenance and hardware; the distribution and token-count grid is what the plan says | +| Gate (CPU) | the gap-split classifier puts near-ties and large-gap disagreements in the right buckets on synthetic logits; a planted large-gap disagreement turns the gate red | +| Race and gate on real hardware (`gpu`) | each adapter meets `assert_matches_reference` at toy, decode- and prefill-shaped dims; a mutation must turn the suite red (Phase 1's rule) | +| Engine reference (`gpu`, paid) | a server that fails to come up, or returns empty completions, refuses to report numbers rather than reporting zeros | + +`gpu`-marked tests stay excluded from CI, with the reason visible in the +test, as elsewhere in the repo. `make check` stays green throughout. + +## 6. Non-goals + +- **A dispatch server.** Phase 7's Rust-router-plus-Python-model-server + architecture (system design §8, ADR-0003) is where a dispatch server + belongs; building a throwaway one here just to have an HTTP endpoint + would measure HF-eager overhead, not the kernel. +- **Swapping dispatch's kernel into vLLM** as a custom MoE backend. +- **H100 or any second hardware class**, and any multi-GPU comparison + (Phase 3's DeepEP result stands on its own; vLLM's and SGLang's own + expert-parallel paths are not benchmarked here). +- **Quantized end-to-end serving comparison.** The engine reference is + bf16 only; int8 is compared at the kernel level, where both engines + expose the same scheme. +- **Beating the engines.** The goal is a correctly measured, honestly + reported table. + +## 7. Risk, cost, and rollout + +**Cost.** One L40 pod, RunPod Secure Cloud, about $0.69/hr as of Phase +5a's rental (checked live at rental time, not assumed). **Budget cap: +$10**, set before the first rental. Three timeboxed stages in the order +above. Cost recorded with `scripts/gpu/provision.py`'s `write_cost_record`. + +**Cost discipline, from CLAUDE.md, applied to this phase:** + +- Pull every evidence file off the pod before `stop`, not after. +- Never leave the pod running across an unbounded wait -- stop first, + restart on the answer. +- Treat `stop` as potentially final for that host's disk (a stopped pod + may be unrestartable); the model download and both engine venvs are + expensive to redo, so the stage order keeps the most valuable evidence + first. +- Develop and debug adapters on CPU against fakes first; rent only for + runs that need the GPU. + +**Risks:** + +- **Production fused-MoE may win at most or all shapes.** A reportable + result, not a failure. What would be a failure is an unfair one: a + baseline running on an untuned or mismatched config, which the tuning + section above exists to prevent. +- **SGLang API drift.** Its fused-MoE code moved recently. Mitigated by + pinned versions and one adapter per engine; if the pinned SGLang version's + API doesn't match what the adapter expects, the adapter fails loudly + rather than falling back to another code path. +- **Disk and install weight.** A 32GB model plus two engine venvs with + conflicting torch pins on one pod. Checked against the pod's disk at + provisioning; models cached on the persistent `/workspace` mount (the + Phase 0 model-cache-on-the-wrong-disk trap). +- **vLLM/SGLang warmup and CUDA-graph confounds** in the serving numbers. + Mitigated by discarded, counted warmup requests, and by stating the + confound in the table where it applies rather than only in prose. +- **Near-tie sensitivity in the correctness gate** (Phase 5b's finding). + Mitigated by the gap-split gate rule and a threshold fixed before the + run. +- **Upstream MoE tuner cost.** vLLM's and SGLang's tuners are slow (Phase + 2's search space was ~1,920 configs per batch size). The plan bounds + tuning to the token counts in the matrix and records how long it ran. + +**Rollout.** Branch `phase-6-final-benchmark` (already created), one commit +per plan task, `docs/STATUS.md` updated in the same commit as each piece +of work, one PR at the end. Findings in `docs/findings/phase-6/`. README, +CLAUDE.md and the story-bank gist refreshed at completion, and design §6 +amended in the same PR to note that the "this project's own server" +assumption was stale and why it wasn't built. From ed9c0a3d4314661c6c87a7418c99d71bca232125 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:57:48 -0400 Subject: [PATCH 03/38] docs: Phase 6 implementation plan, with gate and tuning rules pre-registered 15 tasks: 9 local TDD tasks (code verified in a scratch copy first), then the paid L40 session and wrap-up. Pre-registers the at-scale gate's logit-gap threshold (1.0, from Phase 5b's near-tie data) and the uniform-only tuning rule before any GPU time. Amends the design where live checks corrected it: vllm and sglang share a torch pin so one engines venv works, the per-engine bare-GEMM diagnostic is dropped, and the gate is per config. --- docs/STATUS.md | 38 +- .../2026-09-18-phase-6-final-benchmark.md | 32 +- ...2026-09-19-phase-6-final-benchmark-plan.md | 3974 +++++++++++++++++ 3 files changed, 4031 insertions(+), 13 deletions(-) create mode 100644 docs/plans/2026-09-19-phase-6-final-benchmark-plan.md diff --git a/docs/STATUS.md b/docs/STATUS.md index 1c7fd5d..c1707ec 100644 --- a/docs/STATUS.md +++ b/docs/STATUS.md @@ -529,6 +529,35 @@ spent. Decided in brainstorming: from 5a/5b, with a large-gap disagreement threshold fixed in the plan from 5b's near-tie data before any GPU time. +**Implementation plan written, 2026-09-19:** +`docs/plans/2026-09-19-phase-6-final-benchmark-plan.md`, 15 tasks. Its CPU-testable +code (gap-split classifier, gate prompts, engine adapters, race driver, +summarizer, serving-benchmark helpers) was built and verified in a scratch copy +of the repo before being written into the plan: 236 tests, ruff and +`mypy --strict` clean. Only what the plan says is verified was verified; the real +vLLM and SGLang engines have **not** been exercised (neither installs on this +machine) -- that is Task 10's first job on the pod. + +- **Gate rule pre-registered in the plan, before any GPU time:** a top-1 + disagreement fails only where the reference's top-1/top-2 logit gap exceeds + 1.0 (2.5x Phase 5b's widest measured near-tie of 0.4); minimum 500 compared + positions (the 16 fixed prompts give 1,036 with the real tokenizer); a + stock-vs-stock control measures the gate's own noise floor; a failing config + is excluded from the race rather than halting the session. A test pins the + threshold so it cannot be loosened after seeing a run. +- **Design assumptions the plan's live checks corrected (2026-09-19):** vllm + 0.29.0 and sglang 0.5.20 both pin `torch==2.13.0` with compatible transformers + pins, so one shared engines venv is possible (the design assumed conflicting + pins) and all contestants can share one Triton compiler; SGLang's own tuner + does not list V1's `DeepseekForCausalLM`, so the plan tunes it through a + config-only architecture shim, with a bounded fallback to untuned; the + per-engine bare-GEMM diagnostic is dropped (the fused paths expose no separable + single-GEMM entry point). Design amended accordingly. +- **Tuning rule pre-registered:** every contestant is tuned under uniform routing + only, then timed under both distributions; dispatch's tuned tile size is picked + on uniform results and reused for zipf (a test asserts the choice never looks + at zipf). + ## Next step Phase 5b is done: root cause fixed, real numbers measured, both @@ -539,7 +568,8 @@ merged. Phase 2's PR is still open and awaiting maintainer review; no further work planned on it beyond responding to review feedback. -Phase 6's design is approved (see "Phase 6 progress"). Next: write its -implementation plan in `docs/plans/`, fixing the correctness gate's -logit-gap threshold from Phase 5b's recorded data, then execute it task -by task. +Phase 6's design is approved and its implementation plan is written +(see "Phase 6 progress"). Next: execute +`docs/plans/2026-09-19-phase-6-final-benchmark-plan.md` task by task. +Tasks 1-9 are local, free and TDD; Tasks 10-14 are the paid L40 session and +need the user's explicit go-ahead (live price stated first, $10 cap). diff --git a/docs/design/2026-09-18-phase-6-final-benchmark.md b/docs/design/2026-09-18-phase-6-final-benchmark.md index 2850f5c..86f5246 100644 --- a/docs/design/2026-09-18-phase-6-final-benchmark.md +++ b/docs/design/2026-09-18-phase-6-final-benchmark.md @@ -65,7 +65,7 @@ shapes. Stage 1 gates stage 2; stage 2 gates nothing but is run before stage 3 so that the cheapest, highest-value evidence is on disk first. -### 3.1 At-scale correctness gate (blocks the rest) +### 3.1 At-scale correctness gate (gates each config, per the plan) `docs/STATUS.md` owes this from Phase 5a and 5b: Phase 5a's "perfect top-1/mutual-top-k agreement" covered 29 positions in one run, and Phase @@ -85,7 +85,12 @@ as more than "true on a small sample". the known, explained floating-point sensitivity. - **The threshold is fixed in the implementation plan, from Phase 5b's recorded near-tie data, before any GPU time is spent -- not chosen after - seeing results.** + seeing results.** (Fixed at 1.0 logit; see the plan's "Pre-registered + gate rule".) +- **Gate scope, refined in the plan (2026-09-19):** the gate is per config. + A config that trips it is excluded from the race and reported as failing + its gate, rather than halting the whole paid session; the other configs + proceed. A stock-vs-stock control run measures the gate's own noise floor. - Reuses `scripts/run_baseline.py`'s reference-capture and comparison path rather than a new one; extends it to a larger, seeded prompt set. @@ -101,7 +106,10 @@ zipf routing (reusing `dispatch.kernels.bench.sample_topk_idx`); bf16 and int8. The exact grid is fixed in the plan; Phase 1's sweep points are the default so results extend Phase 1's tables. -**Same inputs, engine-specific adapters.** One seeded input set is +**Same inputs, engine-specific adapters.** All contestants, dispatch's own +kernels included, run inside the one engines venv so they share a torch and +a Triton compiler version (dispatch's `torch>=2.14.0` floor is a +model-loading constraint, not a kernel one). One seeded input set is generated once -- `x`, `topk_ids`, `topk_weights`, and the expert weights -- and a thin adapter per engine converts it to that engine's layout (vLLM and SGLang take fused gate+up in `w1` and down in `w2`; dispatch @@ -120,8 +128,10 @@ enforces. **Timed region.** The whole routed layer from `(x, topk_ids, topk_weights)` to output, each engine's own path (grouping/alignment, the -GEMMs, activation, weighted combine), plus one bare-GEMM diagnostic per -engine, matching Phase 1's structure. Timing is `triton.testing.do_bench` +GEMMs, activation, weighted combine). Phase 1's bare-GEMM diagnostic +stays with dispatch's own kernels only (run_kernel_bench.py already has it): +vLLM's and SGLang's fused-MoE paths expose no separable single-GEMM entry +point, so a per-engine diagnostic isn't measurable like-for-like. Timing is `triton.testing.do_bench` (warmup, repetition and synchronization are its job), reporting mean, p50 and p99 as Phase 1 does. @@ -141,8 +151,12 @@ distribution each engine's config was tuned under. ### 3.3 Engine reference (not a race) -vLLM and SGLang each serve the full model (bf16) on the same L40, in -**separate virtual environments** (their torch pins conflict), driven by +vLLM and SGLang each serve the full model (bf16) on the same L40, from +one shared engines virtual environment (checked live 2026-09-19: vllm 0.29.0 +and sglang 0.5.20 both pin `torch==2.13.0` and their transformers pins are +compatible, so the conflict this design first assumed does not exist at +these versions; if a joint install fails to resolve, fall back to one venv +per engine and record each Triton version in the output), driven by the standard serving benchmark (vLLM's `vllm bench serve` / `benchmark_serving.py`, which works against any OpenAI-compatible server), on one seeded trace, at concurrency 1, 4, 16 and 64. Standard metrics: @@ -151,8 +165,8 @@ and **cost per million tokens computed from the measured GPU-hour rate** (design §6), never estimated. **Dispatch appears at concurrency 1 only**, replaying the identical -token-ID prompts (a seeded trace file the existing harness will need to -accept) through `scripts/run_baseline.py`'s in-process path. That row +prompts (the gate's 16-prompt set, written out as the JSONL trace the +engines read; dispatch runs each prompt twice, one at a time) through `scripts/run_baseline.py`'s in-process path. That row carries an explicit caveat in the table itself: the vLLM/SGLang numbers include HTTP overhead, CUDA-graph capture and torch.compile warmup effects that dispatch's in-process loop doesn't, and dispatch has no diff --git a/docs/plans/2026-09-19-phase-6-final-benchmark-plan.md b/docs/plans/2026-09-19-phase-6-final-benchmark-plan.md new file mode 100644 index 0000000..82497e0 --- /dev/null +++ b/docs/plans/2026-09-19-phase-6-final-benchmark-plan.md @@ -0,0 +1,3974 @@ +# Phase 6 final benchmark vs. vLLM and SGLang -- Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. **Tasks 10-14 rent a GPU and spend real money: never start one without the user's explicit go-ahead for that task, and read "Cost and stop rules" first.** + +**Goal:** Measure, on one rented L40, how dispatch's Triton grouped-GEMM kernels compare with vLLM's and SGLang's fused-MoE on DeepSeek-MoE-16B's real routed-expert shape, gate every config with an at-scale correctness check first, and report where a full production engine lands on the same model as a labeled reference. + +**Architecture:** Three stages on one pod. Stage 1 is a per-config at-scale logit-agreement gate reusing `scripts/run_baseline.py`, extended with a gap-split classifier. Stage 2 is a kernel race: identical seeded inputs and fp32 references are written once to disk, then each engine behind a thin adapter is checked against the reference (refused, not timed, on disagreement) and timed with `triton.testing.do_bench`. Stage 3 launches vLLM and SGLang servers and drives both with `vllm bench serve`; dispatch's own harness supplies the concurrency-1 row. Everything CPU-testable (classifier, adapters, driver, summarizer, serving helpers) is built and verified locally first with fakes. + +**Tech Stack:** Python 3.12+, uv, ruff, `mypy --strict`, pytest, PyTorch, Triton, safetensors; on the pod only: vllm 0.29.0, sglang 0.5.20, torch 2.13.0 (all checked live 2026-09-19). + +**Spec:** `docs/design/2026-09-18-phase-6-final-benchmark.md` (amended 2026-09-19 alongside this plan; the amendments are recorded in its §3 and in system design §6). + +## Global Constraints + +Every task's requirements implicitly include this section. + +- Python 3.12+; `uv`, `ruff` (line length 100), `mypy --strict`, `pytest`. `make check` (lint, typecheck, test) is green before every commit and before any push. +- **Correctness before speed.** No latency number is reported for an engine that has not first matched the fp32 reference (`assert_matches_reference`, rtol 1.6e-2, atol scaled to the reference magnitude). +- **Never quote a benchmark number that wasn't measured** on this hardware, at this config, by this repo. Every reported number states its config: GPU, dtype, precision, token count, routing distribution, tuning label, engine versions. +- `gpu`-marked tests are excluded from CI, with the skip reason visible in the test itself. +- One commit per task. `docs/STATUS.md` is updated in the same commit as the work it describes. +- **Commit messages are plain ASCII: `--`, never an em-dash. No attribution lines in commit messages or PR descriptions** (standing instruction for this session). +- One branch for the phase: `phase-6-final-benchmark` (already created). One PR at the end. Do not push or open the PR without asking the user. +- Model: `deepseek-ai/deepseek-moe-16b-base`. Routed-expert dims (checked live 2026-09-15): hidden 2048, moe_intermediate 1408, 64 routed experts, top-6. +- Hardware: **one L40, RunPod Secure Cloud**. Budget cap **$10**, a ceiling not a target, set before the first rental. +- Token counts (per layer call): `(1, 4, 16, 64, 128, 512, 2048)`; routing `uniform` and `zipf`; precisions `bf16` and weight-only `int8`; seed 0. +- Engine pins: `torch==2.13.0`, `vllm==0.29.0`, `sglang==0.5.20`. A different version is a recorded deviation, never a silent one. +- Findings go to `docs/findings/phase-6/`; the new scripts default their output there. + +## Pre-registered gate rule and decision rules + +Fixed here, **before any Phase 6 GPU time**. Changing any of them after seeing a run's numbers is not allowed; a change needs a written amendment to this section first, with the new evidence that motivates it. + +**The at-scale gate (Stage 1, Task 11).** For each config C in {naive bf16, persistent bf16, int8 (`quantized`)}, run `scripts/run_baseline.py --prompt-set gate` against a same-session stock reference. The gate compares the logits at every prompt position (1,036 positions across 16 prompts with the real tokenizer, measured 2026-09-19; a run under 500 fails). + +- `LARGE_GAP_THRESHOLD = 1.0` logits. A top-1 disagreement where the *reference's* top-1/top-2 gap exceeds this is a **failure**; at or below it is a **near-tie flip**, reported and counted, not failed. +- **Why 1.0:** Phase 5b measured near-tie gaps of 0.1-0.4 on logits of about 20 magnitude (`docs/findings/phase-5b/2026-09-18-phase-5b-speculative-decoding-run.md`), and bf16's own spacing at that magnitude is 0.125. 1.0 is 2.5x the widest measured near-tie and 8 bf16 spacings: loose enough not to fail a correct kernel on floating-point noise, tight enough that a flip at a gap of 1.0+ is not explained by anything measured so far. Phase 5a's `max_abs_diff` (1.3-2.1, taken over every position and every vocabulary entry) is a looser bound on a different quantity and is *not* used: it would classify almost any flip as noise. +- `MIN_GATE_POSITIONS = 500`. Phase 5a's claim rested on 29 positions. +- **Stock-vs-stock control:** one extra run of the unpatched model compared against the first. If the control itself shows a large-gap disagreement, the gate's noise floor is broken (nondeterminism beyond near-ties): **stop the session and investigate before spending on Stages 2-3.** +- **Per-config consequence.** A config that trips the gate is excluded from the race and reported as "failed its correctness gate" with the evidence; the other configs proceed. A trip is a finding to investigate, never a threshold to loosen. + +**The race (Stage 2, Task 12).** + +- Every contestant is tuned under **uniform** routing only (the distribution the engines' own tuners use), then timed under both distributions. dispatch's "tuned" row is the tile size (`block_m` in 16/32/64/128) fastest on uniform at that token count, reused for zipf. Its "default" row is `block_m` 16. vLLM and SGLang get one `default` run (whatever ships for this GPU, recorded) and one `tuned` run. +- An engine that fails the fp32-reference check is refused, not timed. An adapter/API mismatch on the pod gets a **30-minute timebox** to fix; past that, that engine is reported "not measurable at the pinned version" with the failing output kept. +- Tuner timeboxes: **60 minutes per engine per precision.** Token counts a tuner did not finish are labeled `default`, not `tuned`. + +**The engine reference (Stage 3, Task 13).** + +- 30-minute timebox per engine to get a server answering `/health`; past it, the failure log is kept and the engine is reported as unable to serve this checkpoint at the pinned version. +- A bench result with any failed request, or any request that generated other than 64 tokens, is refused (`summarize_bench_result`). +- The dispatch concurrency-1 row is labeled with its non-comparabilities (in-process vs. HTTP, no CUDA graphs); concurrency 4/16/64 read "n/a: no dispatch server". + +## Cost and stop rules + +- **State the live hourly price and get the user's go-ahead before creating the pod.** Estimated cost is *not* a measurement: at the ~$0.7/hr an L40 cost in Phase 5a, the $10 cap is roughly 14 pod-hours; the plan's stages are budgeted at 6-8 hours. Both numbers are planning estimates only. +- **Checkpoints:** at $5 spent, pause and report to the user; at $8, stop starting new work and go to teardown. Never let the total pass $10 without the user raising the cap in writing. +- **Never leave the pod running across an unbounded wait** (a question to the user, a background task with no deadline). Stop it first, restart on the answer. (Phase 5b's $10 cap was exceeded on exactly this.) +- **Pull every evidence file off the pod before `stop` or terminate.** Nothing outside `/workspace` survives a stop, and a stopped pod may be unrestartable ("not enough free GPUs on the host machine", hit on three Phase 5b pods). +- Everything expensive to redo (repo, venvs, model cache, race inputs, tuned configs) lives under `/workspace`. + +## File structure + +| File | Responsibility | Task | +|---|---|---| +| `src/dispatch/benchmark/agreement.py` | gap-split classifier, the pre-registered threshold and position floor | 1 | +| `src/dispatch/benchmark/gate_prompts.py` | the 16 fixed prompts (gate set and engine-reference trace) | 2 | +| `scripts/run_baseline.py` (modify) | `--prompt-set gate`, gap-split output and enforcement, `--ignore-eos` | 2, 3 | +| `src/dispatch/benchmark/harness.py`, `metrics.py` (modify) | `ignore_eos`; public `percentile` | 3 | +| `src/dispatch/benchmark/engines/base.py` | engine contract, seeded case/weight builders, fp32 reference, fused layout | 4 | +| `src/dispatch/benchmark/engines/dispatch_kernels.py` | dispatch's kernels as a contestant | 4 | +| `src/dispatch/benchmark/engines/vllm_moe.py`, `sglang_moe.py`, `registry.py` | the two production engines behind adapters; name -> contestants | 5 | +| `src/dispatch/benchmark/race_summary.py` | tuning rule, rows, markdown table | 6 | +| `scripts/run_engine_race.py` | `prepare` / `run` / `merge` driver | 7 | +| `src/dispatch/benchmark/serving_bench.py`, `scripts/gpu/phase6_engine_reference.py` | serving-benchmark commands, refuse rules, orchestration | 8 | +| `tests/unit/test_engines_real_gpu.py` | real-engine gate (`gpu`), incl. the mutation check | 9 | +| `docs/findings/phase-6/`, `docs/runbooks/phase-6-final-benchmark.md` | evidence and the runbook written from what actually worked | 14, 15 | + +Tasks 1-9 are local, free, and TDD. Tasks 10-14 are the paid pod session. Task 15 is findings, docs and the PR. + +--- + +### Task 1: Gap-split agreement classifier + +**Files:** +- Create: `src/dispatch/benchmark/agreement.py` +- Test: `tests/unit/test_agreement.py` + +**Interfaces:** +- Consumes: nothing from earlier tasks (torch only). +- Produces: `LARGE_GAP_THRESHOLD: float`, `MIN_GATE_POSITIONS: int`, `GapSplitAgreement` (frozen dataclass: `positions`, `disagreements`, `near_tie_disagreements`, `large_gap_disagreements`, `max_disagreement_gap`; property `top1_agreement`; `to_dict()`), `compare_gap_split(actual, reference, *, threshold=LARGE_GAP_THRESHOLD) -> dict[str, GapSplitAgreement]`, `aggregate_gap_split(Iterable[GapSplitAgreement]) -> GapSplitAgreement`. Task 2 consumes all of these. + +- [ ] **Step 1: Write the failing tests** + +Create `tests/unit/test_agreement.py`: + +````python +from __future__ import annotations + +import pytest +import torch + +from dispatch.benchmark.agreement import ( + LARGE_GAP_THRESHOLD, + MIN_GATE_POSITIONS, + GapSplitAgreement, + aggregate_gap_split, + compare_gap_split, +) + + +def _row(top1: float, top2: float, *, winner: int) -> torch.Tensor: + """One 5-token-vocab position whose two largest logits are top1 and top2; + `winner` is which of columns 0/1 holds the larger value.""" + row = torch.zeros(5) + row[winner], row[1 - winner] = top1, top2 + return row + + +def _logits(*rows: torch.Tensor) -> dict[str, torch.Tensor]: + return {"prompt_000_logits": torch.stack(rows)} + + +def test_threshold_and_minimum_positions_are_the_preregistered_values() -> None: + # Pre-registered before any Phase 6 GPU time; changing either needs a + # written plan amendment first, not an edit that makes a red run green. + assert LARGE_GAP_THRESHOLD == 1.0 + assert MIN_GATE_POSITIONS == 500 + + +def test_identical_logits_have_no_disagreements() -> None: + reference = _logits(_row(10.0, 4.0, winner=0), _row(9.0, 8.5, winner=1)) + + result = compare_gap_split(reference, reference)["prompt_000_logits"] + + assert result.positions == 2 + assert result.disagreements == 0 + assert result.top1_agreement == 1.0 + assert result.max_disagreement_gap == 0.0 + + +def test_a_flip_at_a_tiny_gap_is_a_near_tie() -> None: + reference = _logits(_row(20.0, 19.875, winner=0)) # one bf16 spacing at ~20 + actual = _logits(_row(20.0, 19.875, winner=1)) + + result = compare_gap_split(actual, reference)["prompt_000_logits"] + + assert (result.near_tie_disagreements, result.large_gap_disagreements) == (1, 0) + assert result.max_disagreement_gap == pytest.approx(0.125) + + +def test_a_flip_at_a_large_gap_is_reported_separately() -> None: + reference = _logits(_row(10.0, 4.0, winner=0)) + actual = _logits(_row(10.0, 4.0, winner=1)) + + result = compare_gap_split(actual, reference)["prompt_000_logits"] + + assert (result.near_tie_disagreements, result.large_gap_disagreements) == (0, 1) + assert result.max_disagreement_gap == pytest.approx(6.0) + + +def test_a_gap_exactly_at_the_threshold_is_a_near_tie() -> None: + reference = _logits(_row(5.0 + LARGE_GAP_THRESHOLD, 5.0, winner=0)) + actual = _logits(_row(5.0 + LARGE_GAP_THRESHOLD, 5.0, winner=1)) + + result = compare_gap_split(actual, reference)["prompt_000_logits"] + + assert (result.near_tie_disagreements, result.large_gap_disagreements) == (1, 0) + + +def test_agreements_do_not_count_as_disagreements_whatever_their_gap() -> None: + reference = _logits(_row(10.0, 4.0, winner=0), _row(10.0, 9.9, winner=0)) + + result = compare_gap_split(reference, reference)["prompt_000_logits"] + + assert (result.disagreements, result.large_gap_disagreements) == (0, 0) + + +def test_half_precision_inputs_are_compared_in_float32() -> None: + reference = _logits(_row(10.0, 4.0, winner=0)).copy() + reference = {key: value.to(torch.bfloat16) for key, value in reference.items()} + + result = compare_gap_split(reference, reference)["prompt_000_logits"] + + assert result.disagreements == 0 + + +def test_aggregate_sums_counts_and_keeps_the_widest_gap() -> None: + first = GapSplitAgreement( + positions=10, + disagreements=2, + near_tie_disagreements=2, + large_gap_disagreements=0, + max_disagreement_gap=0.3, + ) + second = GapSplitAgreement( + positions=30, + disagreements=1, + near_tie_disagreements=0, + large_gap_disagreements=1, + max_disagreement_gap=4.0, + ) + + total = aggregate_gap_split([first, second]) + + assert (total.positions, total.disagreements) == (40, 3) + assert (total.near_tie_disagreements, total.large_gap_disagreements) == (2, 1) + assert total.max_disagreement_gap == 4.0 + assert total.top1_agreement == pytest.approx(37 / 40) + + +def test_to_dict_carries_the_derived_agreement_rate() -> None: + result = GapSplitAgreement( + positions=4, + disagreements=1, + near_tie_disagreements=1, + large_gap_disagreements=0, + max_disagreement_gap=0.1, + ) + + assert result.to_dict()["top1_agreement"] == 0.75 + + +def test_aggregate_of_nothing_raises() -> None: + with pytest.raises(ValueError, match="zero prompts"): + aggregate_gap_split([]) + + +def test_mismatched_keys_and_shapes_raise() -> None: + reference = _logits(_row(10.0, 4.0, winner=0)) + with pytest.raises(ValueError, match="key mismatch"): + compare_gap_split({"other": reference["prompt_000_logits"]}, reference) + with pytest.raises(ValueError, match="shape mismatch"): + compare_gap_split({"prompt_000_logits": torch.zeros(2, 5)}, reference) +```` + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: `uv run pytest tests/unit/test_agreement.py -v` +Expected: collection FAIL with `ModuleNotFoundError: No module named 'dispatch.benchmark.agreement'`. + +- [ ] **Step 3: Write the implementation** + +Create `src/dispatch/benchmark/agreement.py`: + +````python +"""Gap-split agreement between two logit tensors: every top-1 disagreement is +classified by how close the *reference's* own top-1 and top-2 logits were. + +Phase 5b found genuine near-tied logits (a top-2 gap of 0.1-0.4 out of +~20-magnitude logits; bf16's own spacing at that magnitude is 0.125) that +flip under the int8 kernel's floating-point precision and even between two +processes running the same config. A kernel swap may flip a near-tie. It may +not flip a position the reference was confident about: that is a bug, not +noise. `top1_agreement` alone cannot tell the two apart, and demanding 100% +agreement would fail a correct kernel, so the gate splits them. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import asdict, dataclass + +import torch + +# Pre-registered 2026-09-19, before any Phase 6 GPU time (docs/plans/ +# 2026-09-19-phase-6-final-benchmark-plan.md, "Pre-registered gate rule"). +# 1.0 logit is 2.5x the widest near-tie gap Phase 5b measured (0.4) and 8 +# bf16 spacings at ~20 magnitude. It is NOT to be loosened after seeing a +# run's numbers: a run that trips it is a finding to investigate, and any +# change needs a written amendment to the plan first. A test pins this value. +LARGE_GAP_THRESHOLD = 1.0 + +# Phase 5a's "perfect agreement" covered 29 positions; "at scale" here means +# at least this many compared positions per config. +MIN_GATE_POSITIONS = 500 + + +@dataclass(frozen=True) +class GapSplitAgreement: + positions: int + disagreements: int + near_tie_disagreements: int # reference top1-top2 gap <= threshold + large_gap_disagreements: int # reference top1-top2 gap > threshold + max_disagreement_gap: float # widest reference gap among disagreements; 0.0 if none + + @property + def top1_agreement(self) -> float: + return (self.positions - self.disagreements) / self.positions + + def to_dict(self) -> dict[str, float | int]: + return {**asdict(self), "top1_agreement": self.top1_agreement} + + +def compare_gap_split( + actual: dict[str, torch.Tensor], + reference: dict[str, torch.Tensor], + *, + threshold: float = LARGE_GAP_THRESHOLD, +) -> dict[str, GapSplitAgreement]: + if actual.keys() != reference.keys(): + raise ValueError( + f"key mismatch: actual has {sorted(actual.keys())}, " + f"reference has {sorted(reference.keys())}" + ) + return {key: _gap_split(actual[key], reference[key], threshold) for key in reference} + + +def aggregate_gap_split(per_prompt: Iterable[GapSplitAgreement]) -> GapSplitAgreement: + items = list(per_prompt) + if not items: + raise ValueError("cannot aggregate zero prompts") + return GapSplitAgreement( + positions=sum(item.positions for item in items), + disagreements=sum(item.disagreements for item in items), + near_tie_disagreements=sum(item.near_tie_disagreements for item in items), + large_gap_disagreements=sum(item.large_gap_disagreements for item in items), + max_disagreement_gap=max(item.max_disagreement_gap for item in items), + ) + + +def _gap_split( + actual: torch.Tensor, reference: torch.Tensor, threshold: float +) -> GapSplitAgreement: + if actual.shape != reference.shape: + raise ValueError( + f"shape mismatch: actual {tuple(actual.shape)}, reference {tuple(reference.shape)}" + ) + if reference.shape[-1] < 2: # noqa: PLR2004 -- a top-2 gap needs two logits + raise ValueError("need at least two logits per position to measure a top-2 gap") + reference_f = reference.float() + top2 = reference_f.topk(2, dim=-1).values + gap = top2[..., 0] - top2[..., 1] + disagree = actual.float().argmax(dim=-1) != reference_f.argmax(dim=-1) + large = disagree & (gap > threshold) + near_tie = disagree & ~large + return GapSplitAgreement( + positions=int(disagree.numel()), + disagreements=int(disagree.sum()), + near_tie_disagreements=int(near_tie.sum()), + large_gap_disagreements=int(large.sum()), + max_disagreement_gap=float(gap[disagree].max()) if bool(disagree.any()) else 0.0, + ) +```` + +- [ ] **Step 4: Run the tests to verify they pass** + +Run: `uv run pytest tests/unit/test_agreement.py -v` +Expected: 11 passed. + +- [ ] **Step 5: Mutation check** + +Change `LARGE_GAP_THRESHOLD` to `2.0` in `agreement.py` and re-run. Expected: `test_threshold_and_minimum_positions_are_the_preregistered_values` FAILS. Change it back to `1.0`. (This is the guard against post-hoc loosening.) + +- [ ] **Step 6: Commit** + +Add a bullet at the end of the "Phase 6 progress" section of `docs/STATUS.md`: `- **Task 1 (gap-split classifier)**: dispatch.benchmark.agreement splits every top-1 disagreement by the reference's top1-top2 logit gap; threshold 1.0 and 500-position floor pre-registered in the plan and pinned by a test.` + +```bash +make check +git add src/dispatch/benchmark/agreement.py tests/unit/test_agreement.py docs/STATUS.md +git commit -m "feat: gap-split agreement classifier with pre-registered threshold" +``` + +--- + +### Task 2: Gate prompt set and `run_baseline --prompt-set gate` + +**Files:** +- Create: `src/dispatch/benchmark/gate_prompts.py`, `tests/unit/test_gate_prompts.py` +- Modify: `scripts/run_baseline.py` +- Modify (append tests): `tests/unit/test_run_baseline.py` + +**Interfaces:** +- Consumes: `compare_gap_split`, `aggregate_gap_split`, `MIN_GATE_POSITIONS` from Task 1. +- Produces: `GATE_PROMPTS: tuple[str, ...]` (16 prompts; Task 8 reuses it as the engine-reference trace). `run_baseline.py` gains `--prompt-set {default,gate}`; results JSON gains `prompt_set` and `gap_split` (`None` when no `--compare-reference`, else `GapSplitAgreement.to_dict()`); in gate mode with a reference the run exits non-zero on fewer than 500 positions or on any large-gap disagreement, *after* writing the results and reference files. + +- [ ] **Step 1: Write the failing tests** + +Create `tests/unit/test_gate_prompts.py`: + +````python +from __future__ import annotations + +from dispatch.benchmark.gate_prompts import GATE_PROMPTS + + +def test_gate_prompts_are_sixteen_distinct_nonempty_strings() -> None: + assert len(GATE_PROMPTS) == 16 + assert len(set(GATE_PROMPTS)) == 16 + assert all(prompt.strip() for prompt in GATE_PROMPTS) + + +def test_gate_prompts_are_immutable_so_the_trace_cannot_drift_between_runs() -> None: + assert isinstance(GATE_PROMPTS, tuple) + + +def test_gate_prompts_are_long_enough_to_reach_the_position_floor() -> None: + # A cheap lower bound on the real-tokenizer count (measured 1036 positions + # on 2026-09-19 with deepseek-moe-16b-base's tokenizer; the gate's runtime + # MIN_GATE_POSITIONS check is the authoritative guard). Each prompt is + # several dozen characters per expected token even at 3 chars/token. + assert sum(len(prompt) for prompt in GATE_PROMPTS) / 6 > 500 +```` + +Append to `tests/unit/test_run_baseline.py`, immediately above the first `@pytest.mark.slow` test (it needs no new imports): + +````python +def _gate_logits(rows: int) -> torch.Tensor: + """`rows` positions over a 7-token vocab whose top-1 is column 0 at a 3.0 gap.""" + logits = torch.zeros(rows, 7) + logits[:, 0], logits[:, 1] = 5.0, 2.0 + return logits + + +def _gate_pair( + rows: int, flips: dict[int, float] +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """(reference, actual): identical except that at each row in `flips` the + reference's top-1/top-2 gap is that value and actual has them swapped.""" + reference = _gate_logits(rows) + actual = _gate_logits(rows) + for row, gap in flips.items(): + reference[row, 0], reference[row, 1] = 5.0, 5.0 - gap + actual[row, 0], actual[row, 1] = 5.0 - gap, 5.0 + return {"prompt_000_logits": reference}, {"prompt_000_logits": actual} + + +def _run_gate( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + reference: dict[str, torch.Tensor], + actual: dict[str, torch.Tensor], +) -> None: + reference_path = tmp_path / "stock-reference.safetensors" + save_reference(reference, reference_path) + monkeypatch.setattr(run_baseline_module, "run_baseline", _fake_run_baseline(actual, 27)) + monkeypatch.setattr(run_baseline_module, "summarize", lambda runs: FAKE_SUMMARY) + main( + [ + "--prompt-set", + "gate", + "--moe-kernel", + "quantized", + "--compare-reference", + str(reference_path), + "--output-dir", + str(tmp_path), + "--run-label", + "gate-run", + ] + ) + + +def test_gate_records_the_gap_split_and_passes_a_clean_run( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + reference, actual = _gate_pair(600, {}) + + _run_gate(monkeypatch, tmp_path, reference, actual) + + results = json.loads((tmp_path / "gate-run-results.json").read_text()) + assert results["prompt_set"] == "gate" + assert results["gap_split"]["positions"] == 600 + assert results["gap_split"]["disagreements"] == 0 + assert results["gap_split"]["top1_agreement"] == 1.0 + + +def test_gate_passes_near_tie_flips_but_counts_them( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + reference, actual = _gate_pair(600, {3: 0.125, 40: 0.25, 99: 0.4}) + + _run_gate(monkeypatch, tmp_path, reference, actual) + + split = json.loads((tmp_path / "gate-run-results.json").read_text())["gap_split"] + assert (split["near_tie_disagreements"], split["large_gap_disagreements"]) == (3, 0) + assert split["max_disagreement_gap"] == pytest.approx(0.4) + + +def test_gate_fails_a_large_gap_flip_but_keeps_the_evidence( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + reference, actual = _gate_pair(600, {3: 0.125, 7: 3.0}) + + with pytest.raises(SystemExit, match="a bug, not a near-tie"): + _run_gate(monkeypatch, tmp_path, reference, actual) + + split = json.loads((tmp_path / "gate-run-results.json").read_text())["gap_split"] + assert split["large_gap_disagreements"] == 1 + assert (tmp_path / "gate-run-reference.safetensors").exists() + + +def test_gate_refuses_to_claim_agreement_over_too_few_positions( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + reference, actual = _gate_pair(29, {}) # Phase 5a's sample size + + with pytest.raises(SystemExit, match="fewer than the 500"): + _run_gate(monkeypatch, tmp_path, reference, actual) + + +def test_a_run_without_a_reference_records_no_gap_split( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setattr( + run_baseline_module, "run_baseline", _fake_run_baseline(_gate_pair(600, {})[1], 0) + ) + monkeypatch.setattr(run_baseline_module, "summarize", lambda runs: FAKE_SUMMARY) + + main(["--prompt-set", "gate", "--output-dir", str(tmp_path), "--run-label", "ref-run"]) + + results = json.loads((tmp_path / "ref-run-results.json").read_text()) + assert results["gap_split"] is None + assert (tmp_path / "ref-run-reference.safetensors").exists() +```` + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: `uv run pytest tests/unit/test_gate_prompts.py tests/unit/test_run_baseline.py -v -m "not slow"` +Expected: `test_gate_prompts.py` fails to import; the new `test_run_baseline.py` tests fail with `SystemExit: 2` (argparse: unrecognized `--prompt-set`). + +- [ ] **Step 3: Write the prompt set** + +Create `src/dispatch/benchmark/gate_prompts.py`: + +````python +"""The fixed prompt set for Phase 6's at-scale correctness gate. Real prose, +code, and arithmetic text of mixed style, so the routed experts see varied +token distributions rather than one register. Compared teacher-forced (the +logits at every prompt position), which is why each is a few sentences long: +the gate needs hundreds of positions per config, and the runtime check +against MIN_GATE_POSITIONS enforces that with the real tokenizer. + +Also the request trace for the engine reference (docs/design/ +2026-09-18-phase-6-final-benchmark.md section 3.3), so every engine, and +dispatch's own harness, sees identical text. +""" + +from __future__ import annotations + +GATE_PROMPTS: tuple[str, ...] = ( + "The committee reviewed the proposal for the new harbor bridge over three long " + "sessions. Engineers argued about the cost of steel, while residents worried about " + "traffic noise and the loss of the old fishing pier. In the end they voted to " + "approve a smaller design with a wider footpath.", + "def merge_sorted(left, right):\n result = []\n i = j = 0\n while i < len(left) " + "and j < len(right):\n if left[i] <= right[j]:\n result.append(left[i])\n" + " i += 1\n else:\n result.append(right[j])\n j += 1\n", + "To compute the area of a triangle with base 12 and height 7, multiply the two " + "numbers and divide by two. So 12 times 7 is 84, and half of 84 is 42. The area is " + "therefore 42 square units, which we can check by drawing the triangle on grid paper.", + "When the storm finally reached the coast, the lighthouse keeper lit the lamp an hour " + "early. The waves climbed the rocks below and threw white spray against the glass. He " + "wrote in his logbook that the night was the worst he had seen in thirty years.", + "SELECT customer_id, SUM(amount) AS total_spent FROM orders WHERE order_date >= " + "'2025-01-01' GROUP BY customer_id HAVING SUM(amount) > 1000 ORDER BY total_spent " + "DESC LIMIT 20; This query lists the twenty biggest customers of the year.", + "Photosynthesis converts light energy into chemical energy stored in glucose. Inside " + "the chloroplast, chlorophyll absorbs red and blue light and uses it to split water, " + "releasing oxygen. The resulting energy carriers then drive the fixation of carbon " + "dioxide in the Calvin cycle.", + "Dear Ms. Alvarez, thank you for your patience while we investigated the delayed " + "shipment. Your order left our warehouse on Tuesday and should arrive by Friday. We " + "have refunded the shipping fee and added a discount code to your account.", + "The Roman Republic expanded across the Mediterranean through a mix of alliances and " + "conquest. After the Punic Wars, Carthage was destroyed and Rome controlled the " + "western sea. Wealth flowed into the city, and with it came deep arguments about land " + "and power.", + "fn fibonacci(n: u32) -> u64 {\n let (mut a, mut b) = (0u64, 1u64);\n for _ in 0..n " + "{\n let next = a + b;\n a = b;\n b = next;\n }\n a\n}\n" + "// Runs in linear time and constant space.", + "A good sourdough starter needs regular feeding, warmth, and a little patience. Mix " + "equal weights of flour and water each day, discard most of the old starter, and keep " + "the jar somewhere around twenty-four degrees. After a week it should smell pleasantly " + "sour and double in size within hours.", + "In 1969 the first humans walked on the surface of the Moon. The lunar module landed " + "in a flat region called the Sea of Tranquility, and the crew spent about two and a " + "half hours outside collecting rock samples. Millions of people watched the broadcast " + "on television.", + "Question: If a train leaves the station at 3 pm travelling at 80 kilometres per hour, " + "and a second train leaves at 4 pm at 100 kilometres per hour on the same track, when " + "does the second train catch up? Answer: the first train has an 80 kilometre head " + "start, closing at 20 kilometres per hour.", + "The city council announced that the downtown library will stay open until midnight " + "during exam week. Volunteers will serve free coffee, and extra study rooms can be " + "reserved online. Officials hope the pilot programme will become permanent next year.", + "Neural networks learn by adjusting their weights to reduce a loss function. Each " + "training step computes the gradient of the loss with respect to every weight and " + "moves the weights a small distance in the opposite direction. Repeating this over " + "millions of examples gradually produces useful behaviour.", + "She opened the old wooden chest and found a bundle of letters tied with blue string. " + "The ink had faded, but the handwriting was still clear, looping and quick. The first " + "line read simply: I hope this reaches you before the winter does.", + "import argparse\n\nparser = argparse.ArgumentParser(description='Resize images')\n" + "parser.add_argument('--width', type=int, default=640)\nparser.add_argument('--height', " + "type=int, default=480)\nargs = parser.parse_args()\nprint(f'Resizing to " + "{args.width}x{args.height}')\n", +) +```` + +- [ ] **Step 4: Modify `scripts/run_baseline.py`** + +Apply these five edits. + +(a) Add imports. Replace +```python +from dispatch.benchmark.harness import generate_with_timings, load_model +``` +with +```python +from dispatch.benchmark.agreement import ( + MIN_GATE_POSITIONS, + aggregate_gap_split, + compare_gap_split, +) +from dispatch.benchmark.gate_prompts import GATE_PROMPTS +from dispatch.benchmark.harness import generate_with_timings, load_model +``` + +(b) Add the flag. Replace +```python + parser.add_argument( + "--compare-reference", +``` +with +```python + parser.add_argument( + "--prompt-set", + choices=["default", "gate"], + default="default", + help="'gate' runs Phase 6's larger fixed prompt set and enforces its gap-split rule", + ) + parser.add_argument( + "--compare-reference", +``` + +(c) Select the prompts. Replace +```python + prompts=DEFAULT_PROMPTS, +``` +with +```python + prompts=list(GATE_PROMPTS) if args.prompt_set == "gate" else DEFAULT_PROMPTS, +``` + +(d) Compute the gap split once, from one loaded reference. Replace +```python + comparison = ( + compare_top_k_agreement(logits, load_reference(args.compare_reference)) + if args.compare_reference is not None + else {} + ) +``` +with +```python + reference = load_reference(args.compare_reference) if args.compare_reference else None + comparison = compare_top_k_agreement(logits, reference) if reference is not None else {} + gap_split = ( + aggregate_gap_split(compare_gap_split(logits, reference).values()) + if reference is not None + else None + ) +``` +and in the results dict, after `"moe_layers_patched": moe_layers_patched,` add +```python + "prompt_set": args.prompt_set, + "gap_split": None if gap_split is None else gap_split.to_dict(), +``` + +(e) Enforce the gate after the existing `mutual_top_k` check. Replace +```python + if not all(value.mutual_top_k for value in comparison.values()): + raise SystemExit( + f"{args.moe_kernel} logits disagree with {args.compare_reference} -- see {results_path}" + ) +``` +with +```python + if not all(value.mutual_top_k for value in comparison.values()): + raise SystemExit( + f"{args.moe_kernel} logits disagree with {args.compare_reference} -- see {results_path}" + ) + if args.prompt_set == "gate" and gap_split is not None: + if gap_split.positions < MIN_GATE_POSITIONS: + raise SystemExit( + f"gate compared only {gap_split.positions} positions, fewer than the " + f"{MIN_GATE_POSITIONS} an at-scale claim needs -- see {results_path}" + ) + if gap_split.large_gap_disagreements > 0: + raise SystemExit( + f"{args.moe_kernel} flips {gap_split.large_gap_disagreements} position(s) the " + f"reference was confident about (widest gap {gap_split.max_disagreement_gap:.3f}) " + f"-- a bug, not a near-tie; see {results_path}" + ) +``` + +- [ ] **Step 5: Run the tests to verify they pass** + +Run: `uv run pytest tests/unit/test_gate_prompts.py tests/unit/test_run_baseline.py tests/unit/test_agreement.py -v -m "not slow"` +Expected: all pass (the existing `test_run_baseline.py` tests still pass unmodified). + +- [ ] **Step 6: Commit** + +STATUS bullet: `- **Task 2 (gate prompt set + run_baseline gate mode)**: --prompt-set gate runs 16 fixed prompts (1,036 positions with the real tokenizer) and fails on any large-gap flip or fewer than 500 positions.` + +```bash +make check +git add src/dispatch/benchmark/gate_prompts.py tests/unit/test_gate_prompts.py scripts/run_baseline.py tests/unit/test_run_baseline.py docs/STATUS.md +git commit -m "feat: at-scale gate mode for run_baseline (16 prompts, gap-split rule)" +``` + +--- + +### Task 3: `ignore_eos` and a public `percentile` + +Dispatch's concurrency-1 reference row must generate exactly 64 tokens per request, as `vllm bench serve --ignore-eos` does, or the throughputs are not comparable. `metrics._percentile` becomes public because Task 8's serving summary reuses it. + +**Files:** +- Modify: `src/dispatch/benchmark/harness.py`, `src/dispatch/benchmark/metrics.py`, `scripts/run_baseline.py` +- Modify (insert tests): `tests/unit/test_harness.py` + +**Interfaces:** +- Consumes: Task 2's `run_baseline.py`. +- Produces: `generate_with_timings(..., ignore_eos: bool = False)`; `run_baseline(..., ignore_eos: bool = False)`; CLI `--ignore-eos`; results JSON gains `max_new_tokens` and `ignore_eos`; `dispatch.benchmark.metrics.percentile(sorted_values, fraction)`. + +- [ ] **Step 1: Write the failing tests** + +Insert into `tests/unit/test_harness.py`, immediately above `def test_load_model_omits_attn_implementation_by_default(`: + +````python +class _EosAtSecondTokenTokenizer(_FakeTokenizer): + eos_token_id = 2 # _FakeModel emits token (call_count % vocab): 1, 2, 3, ... + + +def _run_with_eos_tokenizer(*, ignore_eos: bool) -> int: + counter = iter(range(100)) + timing = generate_with_timings( + _FakeModel(), # type: ignore[arg-type] + _EosAtSecondTokenTokenizer(), # type: ignore[arg-type] + "prompt", + max_new_tokens=5, + clock_fn=lambda: float(next(counter)), + ignore_eos=ignore_eos, + ) + return timing.generated_token_count + + +def test_generate_with_timings_stops_at_eos_by_default() -> None: + assert _run_with_eos_tokenizer(ignore_eos=False) == 2 + + +def test_generate_with_timings_ignore_eos_generates_exactly_max_new_tokens() -> None: + assert _run_with_eos_tokenizer(ignore_eos=True) == 5 +```` + +- [ ] **Step 2: Run to verify failure** + +Run: `uv run pytest tests/unit/test_harness.py -v -m "not slow"` +Expected: the two new tests FAIL with `TypeError: generate_with_timings() got an unexpected keyword argument 'ignore_eos'`. + +- [ ] **Step 3: Implement** + +`src/dispatch/benchmark/harness.py`: in `generate_with_timings`' signature, add `ignore_eos: bool = False,` after `clock_fn: Callable[[], float] = time.perf_counter,`; and replace +```python + if eos_token_id is not None and next_token.item() == eos_token_id: +``` +with +```python + if not ignore_eos and eos_token_id is not None and next_token.item() == eos_token_id: +``` + +`src/dispatch/benchmark/metrics.py`: rename `_percentile` to `percentile` (its definition and its two call sites inside `summarize`). No other module imports it (`scripts/gpu/phase4_concurrency.py` has its own private copy; leave it). + +`scripts/run_baseline.py`: (a) add `ignore_eos: bool = False,` after `moe_kernel: str = "none",` in `run_baseline`'s signature; (b) replace +```python + model, tokenizer, prompt, max_new_tokens=max_new_tokens, device=device + ) +``` +with +```python + model, + tokenizer, + prompt, + max_new_tokens=max_new_tokens, + device=device, + ignore_eos=ignore_eos, + ) +``` +(c) after the `--max-new-tokens` argument add +```python + parser.add_argument( + "--ignore-eos", + action="store_true", + help="generate exactly --max-new-tokens even past an EOS (matches vllm bench serve)", + ) +``` +(d) pass `ignore_eos=args.ignore_eos,` after `moe_kernel=args.moe_kernel,` in the `run_baseline(...)` call in `main`; (e) in the results dict, after `"prompt_set": args.prompt_set,` add `"max_new_tokens": args.max_new_tokens,` and `"ignore_eos": args.ignore_eos,`. + +- [ ] **Step 3b: Run the tests to verify they pass** + +Run: `uv run pytest tests/unit/test_harness.py tests/unit/test_metrics.py tests/unit/test_run_baseline.py -v -m "not slow"` +Expected: all pass. + +- [ ] **Step 4: Commit** + +STATUS bullet: `- **Task 3 (ignore_eos, public percentile)**: the harness can generate exactly N tokens past EOS, matching vllm bench serve --ignore-eos, so dispatch's concurrency-1 reference row is comparable.` + +```bash +make check +git add src/dispatch/benchmark/harness.py src/dispatch/benchmark/metrics.py scripts/run_baseline.py tests/unit/test_harness.py docs/STATUS.md +git commit -m "feat: ignore_eos in the harness so reference rows match vllm bench serve" +``` + +--- +### Task 4: Engine contract, seeded inputs, fp32 reference, dispatch adapter + +**Files:** +- Create: `src/dispatch/benchmark/engines/__init__.py` (empty), `src/dispatch/benchmark/engines/base.py`, `src/dispatch/benchmark/engines/dispatch_kernels.py` +- Test: `tests/unit/test_engine_base.py`, `tests/unit/test_dispatch_engine.py` + +**Interfaces:** +- Consumes: `dispatch.kernels.bench.sample_topk_idx`, `dispatch.kernels.moe_forward` (`StackedExpertWeights`, `grouped_moe_routed`, `torch_grouped_matmul`, `assert_matches_reference`), `dispatch.kernels.quantization` (`QuantizedStackedExpertWeights`, `dequantize_int8`, `quantize_stacked_weights`, `grouped_moe_routed_quantized`, `torch_grouped_matmul_dequant`), `dispatch.kernels.backends` (`resolve_backend`, `resolve_quantized_backend`). +- Produces (Tasks 5-9 rely on these exact names): + - `RaceCase(x, topk_idx, topk_weight)`; `BoundLayer = Callable[[], Tensor]`; `LayerFactory = Callable[[RaceCase], BoundLayer]`. + - `MoEEngine` Protocol: `name: str`, `prepare_bf16(StackedExpertWeights) -> LayerFactory`, `prepare_int8(QuantizedStackedExpertWeights) -> LayerFactory`. `prepare_*` does one-time weight conversion; `bind` (the factory call) does per-case dtype casts; only the returned closure is ever timed. + - `make_weights(dtype, *, seed, device, hidden_size=..., intermediate_size=..., n_experts=...)`, `make_case(num_tokens, distribution, *, dtype, seed, device, hidden_size=..., n_experts=..., top_k=...)`, `reference_output(case, weights)`, `dequantized_weights(qweights)`, `fuse_gate_up(gate, up)`; constants `HIDDEN_SIZE`, `MOE_INTERMEDIATE_SIZE`, `N_ROUTED_EXPERTS`, `NUM_EXPERTS_PER_TOK`, `TOKEN_COUNTS`. + - `DispatchEngine(backend, *, block_m=16)` with `name == f"dispatch-{backend}-bm{block_m}"`; backends `torch` (CPU, for tests), `naive`, `persistent`; int8 supported for `torch` and `naive` only. + +- [ ] **Step 1: Write the failing tests** + +Create `tests/unit/test_engine_base.py`: + +````python +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F # noqa: N812 -- F is the universal PyTorch convention + +from dispatch.benchmark.engines.base import ( + dequantized_weights, + fuse_gate_up, + make_case, + make_weights, + reference_output, +) +from dispatch.kernels.moe_forward import assert_matches_reference +from dispatch.kernels.quantization import quantize_stacked_weights + +DIMS = {"hidden_size": 16, "intermediate_size": 8, "n_experts": 6} +CASE_DIMS = {"hidden_size": 16, "n_experts": 6, "top_k": 2} + + +def test_make_weights_is_seeded_and_shaped_from_the_dims() -> None: + first = make_weights(torch.float32, seed=3, device="cpu", **DIMS) + again = make_weights(torch.float32, seed=3, device="cpu", **DIMS) + other = make_weights(torch.float32, seed=4, device="cpu", **DIMS) + + assert first.gate.shape == (6, 8, 16) + assert first.up.shape == (6, 8, 16) + assert first.down.shape == (6, 16, 8) + assert torch.equal(first.gate, again.gate) + assert not torch.equal(first.gate, other.gate) + + +def test_uniform_and_zipf_cases_share_inputs_and_differ_only_in_routing() -> None: + uniform = make_case(64, "uniform", dtype=torch.float32, seed=1, device="cpu", **CASE_DIMS) + zipf = make_case(64, "zipf", dtype=torch.float32, seed=1, device="cpu", **CASE_DIMS) + + assert torch.equal(uniform.x, zipf.x) + assert torch.equal(uniform.topk_weight, zipf.topk_weight) + assert not torch.equal(uniform.topk_idx, zipf.topk_idx) + assert uniform.topk_weight.dtype == torch.float32 + assert uniform.topk_idx.dtype == torch.int64 + + +def test_zipf_routing_concentrates_load_on_low_experts() -> None: + uniform = make_case(512, "uniform", dtype=torch.float32, seed=1, device="cpu", **CASE_DIMS) + zipf = make_case(512, "zipf", dtype=torch.float32, seed=1, device="cpu", **CASE_DIMS) + + uniform_load = torch.bincount(uniform.topk_idx.reshape(-1), minlength=6) + zipf_load = torch.bincount(zipf.topk_idx.reshape(-1), minlength=6) + assert zipf_load[0] > zipf_load[-1] * 2 + assert uniform_load.max() < uniform_load.min() * 2 + + +def test_fuse_gate_up_puts_gate_first_and_is_the_layout_silu_and_mul_reads() -> None: + weights = make_weights(torch.float32, seed=0, device="cpu", **DIMS) + fused = fuse_gate_up(weights.gate, weights.up) + assert fused.shape == (6, 16, 16) + assert torch.equal(fused[:, :8], weights.gate) + assert torch.equal(fused[:, 8:], weights.up) + + x = torch.randn(3, 16) + expert = 2 + hidden = x @ fused[expert].T + fused_activation = F.silu(hidden[:, :8]) * hidden[:, 8:] + separate_activation = F.silu(x @ weights.gate[expert].T) * (x @ weights.up[expert].T) + torch.testing.assert_close(fused_activation, separate_activation) + + +def test_fuse_gate_up_also_fuses_per_channel_scales() -> None: + gate_scale, up_scale = torch.ones(6, 8), torch.full((6, 8), 2.0) + fused = fuse_gate_up(gate_scale, up_scale) + assert fused.shape == (6, 16) + assert torch.equal(fused[:, :8], gate_scale) + assert torch.equal(fused[:, 8:], up_scale) + + +def test_reference_output_is_float32_whatever_the_case_dtype() -> None: + weights = make_weights(torch.bfloat16, seed=0, device="cpu", **DIMS) + case = make_case(5, "uniform", dtype=torch.bfloat16, seed=0, device="cpu", **CASE_DIMS) + + assert reference_output(case, weights).dtype == torch.float32 + + +def test_int8_reference_uses_the_dequantized_weights_not_the_originals() -> None: + weights = make_weights(torch.float32, seed=0, device="cpu", **DIMS) + qweights = quantize_stacked_weights(weights) + case = make_case(5, "uniform", dtype=torch.float32, seed=0, device="cpu", **CASE_DIMS) + + int8_reference = reference_output(case, dequantized_weights(qweights)) + original_reference = reference_output(case, weights) + + assert not torch.equal(int8_reference, original_reference) + assert_matches_reference(int8_reference, original_reference) # close, not identical + with pytest.raises(AssertionError): + torch.testing.assert_close(int8_reference, original_reference, rtol=0, atol=0) +```` + +Create `tests/unit/test_dispatch_engine.py`: + +````python +from __future__ import annotations + +import pytest +import torch + +from dispatch.benchmark.engines.base import ( + dequantized_weights, + make_case, + make_weights, + reference_output, +) +from dispatch.benchmark.engines.dispatch_kernels import DispatchEngine +from dispatch.kernels.moe_forward import assert_matches_reference +from dispatch.kernels.quantization import quantize_stacked_weights + +DIMS = {"hidden_size": 16, "intermediate_size": 8, "n_experts": 6} +CASE_DIMS = {"hidden_size": 16, "n_experts": 6, "top_k": 2} + + +def test_name_carries_backend_and_tile_size() -> None: + assert DispatchEngine("naive", block_m=32).name == "dispatch-naive-bm32" + + +def test_bf16_layer_matches_the_fp32_reference() -> None: + weights = make_weights(torch.float32, seed=0, device="cpu", **DIMS) + case = make_case(9, "zipf", dtype=torch.float32, seed=1, device="cpu", **CASE_DIMS) + + layer = DispatchEngine("torch").prepare_bf16(weights)(case) + + assert_matches_reference(layer(), reference_output(case, weights)) + + +def test_bound_layer_is_repeatable_so_it_can_be_timed() -> None: + weights = make_weights(torch.float32, seed=0, device="cpu", **DIMS) + case = make_case(9, "uniform", dtype=torch.float32, seed=1, device="cpu", **CASE_DIMS) + layer = DispatchEngine("torch").prepare_bf16(weights)(case) + + assert torch.equal(layer(), layer()) + + +def test_int8_layer_matches_the_dequantized_reference() -> None: + weights = make_weights(torch.float32, seed=0, device="cpu", **DIMS) + qweights = quantize_stacked_weights(weights) + case = make_case(9, "zipf", dtype=torch.float32, seed=1, device="cpu", **CASE_DIMS) + + layer = DispatchEngine("torch").prepare_int8(qweights)(case) + + assert_matches_reference(layer(), reference_output(case, dequantized_weights(qweights))) + + +def test_persistent_backend_has_no_int8_kernel() -> None: + qweights = quantize_stacked_weights(make_weights(torch.float32, seed=0, device="cpu", **DIMS)) + + with pytest.raises(ValueError, match="no int8 kernel"): + DispatchEngine("persistent").prepare_int8(qweights) + + +def test_a_layer_that_disagrees_with_the_reference_is_detectable() -> None: + """The race driver's refuse-not-time gate depends on assert_matches_reference + rejecting a wrong layer; prove it does for this adapter's output shape.""" + weights = make_weights(torch.float32, seed=0, device="cpu", **DIMS) + case = make_case(9, "uniform", dtype=torch.float32, seed=1, device="cpu", **CASE_DIMS) + wrong = DispatchEngine("torch").prepare_bf16(weights)(case)() * 1.5 + + with pytest.raises(AssertionError): + assert_matches_reference(wrong, reference_output(case, weights)) +```` + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: `uv run pytest tests/unit/test_engine_base.py tests/unit/test_dispatch_engine.py -v` +Expected: collection FAIL with `ModuleNotFoundError: No module named 'dispatch.benchmark.engines'`. + +- [ ] **Step 3: Write the implementation** + +Create the empty package marker, then the two modules: + +```bash +mkdir -p src/dispatch/benchmark/engines && touch src/dispatch/benchmark/engines/__init__.py +``` + +`src/dispatch/benchmark/engines/base.py`: + +````python +"""The contract every Phase 6 race contestant implements, plus the seeded +inputs and fp32 reference they are all checked against. + +An engine is anything that computes DeepSeek's routed-expert MoE forward from +`(x, topk_idx, topk_weight)` and a set of expert weights. `prepare_*` does +the one-time weight conversion (fusing gate+up, uploading scales), and +`bind` does the per-case input conversion (dtype casts) -- both untimed -- +so the closure a benchmark times is only the engine's own routed-MoE call, +never a cast the engine's real serving path wouldn't pay. +""" + +from __future__ import annotations + +import math +from collections.abc import Callable +from dataclasses import dataclass +from typing import Protocol + +import torch + +from dispatch.kernels.bench import sample_topk_idx +from dispatch.kernels.moe_forward import ( + StackedExpertWeights, + grouped_moe_routed, + torch_grouped_matmul, +) +from dispatch.kernels.quantization import QuantizedStackedExpertWeights, dequantize_int8 + +# deepseek-ai/deepseek-moe-16b-base config.json, checked live 2026-09-15 +# (same constants as scripts/run_kernel_bench.py). +HIDDEN_SIZE = 2048 +MOE_INTERMEDIATE_SIZE = 1408 +N_ROUTED_EXPERTS = 64 +NUM_EXPERTS_PER_TOK = 6 + +# Phase 1's default sweep (1, 16, 128, 512, 2048) plus 4 and 64, so results +# extend Phase 1's tables and decode-sized batches are better covered. +TOKEN_COUNTS = (1, 4, 16, 64, 128, 512, 2048) + + +@dataclass(frozen=True) +class RaceCase: + """One benchmark input. `topk_weight` is float32 by contract; an adapter + casts it to whatever its engine wants inside `bind`.""" + + x: torch.Tensor # (num_tokens, hidden) + topk_idx: torch.Tensor # (num_tokens, top_k) int64 + topk_weight: torch.Tensor # (num_tokens, top_k) float32 + + +BoundLayer = Callable[[], torch.Tensor] +LayerFactory = Callable[[RaceCase], BoundLayer] + + +class MoEEngine(Protocol): + name: str + + def prepare_bf16(self, weights: StackedExpertWeights) -> LayerFactory: ... + + def prepare_int8(self, qweights: QuantizedStackedExpertWeights) -> LayerFactory: ... + + +def make_weights( # noqa: PLR0913 -- the real DeepSeek dims are the defaults; tests shrink them + dtype: torch.dtype, + *, + seed: int, + device: str, + hidden_size: int = HIDDEN_SIZE, + intermediate_size: int = MOE_INTERMEDIATE_SIZE, + n_experts: int = N_ROUTED_EXPERTS, +) -> StackedExpertWeights: + generator = torch.Generator(device=device).manual_seed(seed) + + def stack(n: int, k: int) -> torch.Tensor: + weights = torch.randn(n_experts, n, k, device=device, dtype=dtype, generator=generator) + return weights / math.sqrt(k) + + return StackedExpertWeights( + gate=stack(intermediate_size, hidden_size), + up=stack(intermediate_size, hidden_size), + down=stack(hidden_size, intermediate_size), + ) + + +def make_case( # noqa: PLR0913 -- the real DeepSeek dims are the defaults; tests shrink them + num_tokens: int, + distribution: str, + *, + dtype: torch.dtype, + seed: int, + device: str, + hidden_size: int = HIDDEN_SIZE, + n_experts: int = N_ROUTED_EXPERTS, + top_k: int = NUM_EXPERTS_PER_TOK, +) -> RaceCase: + """Uniform and zipf cases at the same `seed` and `num_tokens` share x and + topk_weight exactly; only the routing differs, so the two distributions + are a controlled comparison.""" + generator = torch.Generator(device=device).manual_seed(seed) + x = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype, generator=generator) + topk_idx = sample_topk_idx( + num_tokens, + n_experts, + top_k, + distribution=distribution, + generator=torch.Generator().manual_seed(seed), + ).to(device) + topk_weight = torch.rand( + num_tokens, top_k, device=device, dtype=torch.float32, generator=generator + ) + return RaceCase(x=x, topk_idx=topk_idx, topk_weight=topk_weight) + + +def dequantized_weights(qweights: QuantizedStackedExpertWeights) -> StackedExpertWeights: + """float32 weights equal to what an int8 kernel effectively multiplies by: + the int8 race's reference is the *same* quantized weights, dequantized, + so nothing should diverge beyond float precision.""" + return StackedExpertWeights( + gate=dequantize_int8(qweights.gate), + up=dequantize_int8(qweights.up), + down=dequantize_int8(qweights.down), + ) + + +def reference_output(case: RaceCase, weights: StackedExpertWeights) -> torch.Tensor: + """The fp32 eager reference every engine is checked against: the same + per-expert-loop contract DeepseekMoE.moe_infer implements, in float32.""" + weights_fp32 = StackedExpertWeights( + gate=weights.gate.float(), up=weights.up.float(), down=weights.down.float() + ) + return grouped_moe_routed( + case.x.float(), case.topk_idx, case.topk_weight.float(), weights_fp32, torch_grouped_matmul + ) + + +def fuse_gate_up(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """(E, N, K) gate and up -> (E, 2N, K) with gate's rows first. Works for + (E, N) per-channel scales too (-> (E, 2N)). This is the layout vLLM's and + SGLang's fused-MoE kernels read: silu(first half) * second half.""" + return torch.cat([gate, up], dim=1) +```` + +`src/dispatch/benchmark/engines/dispatch_kernels.py`: + +````python +"""dispatch's own kernels as a race contestant: the naive or persistent +Triton grouped-GEMM (bf16), or the int8 weight-only kernel, behind the +common engine contract. `torch` is the eager backend, which runs on CPU and +is what the unit tests use.""" + +from __future__ import annotations + +import functools + +import torch + +from dispatch.benchmark.engines.base import LayerFactory, RaceCase +from dispatch.kernels.backends import resolve_backend, resolve_quantized_backend +from dispatch.kernels.moe_forward import StackedExpertWeights, grouped_moe_routed +from dispatch.kernels.quantization import ( + QuantizedGroupedMatmul, + QuantizedStackedExpertWeights, + grouped_moe_routed_quantized, + torch_grouped_matmul_dequant, +) + + +class DispatchEngine: + def __init__(self, backend: str, *, block_m: int = 16) -> None: + self.name = f"dispatch-{backend}-bm{block_m}" + self._backend = backend + self._block_m = block_m + + def prepare_bf16(self, weights: StackedExpertWeights) -> LayerFactory: + matmul = resolve_backend(self._backend) + block_m = self._block_m + + def bind(case: RaceCase) -> functools.partial[torch.Tensor]: + return functools.partial( + grouped_moe_routed, + case.x, + case.topk_idx, + case.topk_weight.to(case.x.dtype), + weights, + matmul, + block_m=block_m, + ) + + return bind + + def prepare_int8(self, qweights: QuantizedStackedExpertWeights) -> LayerFactory: + matmul = self._quantized_matmul() + block_m = self._block_m + + def bind(case: RaceCase) -> functools.partial[torch.Tensor]: + return functools.partial( + grouped_moe_routed_quantized, + case.x, + case.topk_idx, + case.topk_weight.to(case.x.dtype), + qweights, + matmul, + block_m=block_m, + ) + + return bind + + def _quantized_matmul(self) -> QuantizedGroupedMatmul: + if self._backend == "torch": + return torch_grouped_matmul_dequant + if self._backend == "naive": + return resolve_quantized_backend() + raise ValueError( + f"no int8 kernel for backend {self._backend!r}: Phase 5a's int8 kernel is " + "built on the naive launch order only" + ) +```` + +- [ ] **Step 4: Run the tests to verify they pass** + +Run: `uv run pytest tests/unit/test_engine_base.py tests/unit/test_dispatch_engine.py -v` +Expected: 13 passed. + +- [ ] **Step 5: Commit** + +STATUS bullet: `- **Task 4 (engine contract + dispatch adapter)**: seeded inputs, fp32 reference, fused gate+up layout, and dispatch's kernels behind one MoEEngine contract; uniform and zipf cases share x and weights so only routing differs.` + +```bash +make check +git add src/dispatch/benchmark/engines tests/unit/test_engine_base.py tests/unit/test_dispatch_engine.py docs/STATUS.md +git commit -m "feat: race engine contract, seeded inputs, fp32 reference, dispatch adapter" +``` + +--- + +### Task 5: vLLM and SGLang adapters, and the engine registry + +The adapters call each engine's own `fused_experts` directly with the race's fixed `topk_ids`/`topk_weights` (no gating), so routing is identical for every contestant. APIs checked live 2026-09-19 against vllm-project/vllm `main` and 0.29.0, and sgl-project/sglang `main` and 0.5.20: + +- vLLM: `vllm.model_executor.layers.fused_moe.fused_moe.fused_experts(hidden_states, w1, w2, topk_weights, topk_ids, ..., quant_config=)`; int8 via `vllm.model_executor.layers.fused_moe.config.int8_w8a16_moe_quant_config(w1_scale, w2_scale, ...)`. +- SGLang: `sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe.fused_experts(hidden_states, w1, w2, topk_output, moe_runner_config, ..., use_int8_w8a16, per_channel_quant, w1_scale, w2_scale)` with `topk_output = sglang.srt.layers.moe.topk.StandardTopKOutput(topk_weights, topk_ids, router_logits)` and `moe_runner_config = sglang.srt.layers.moe.moe_runner.base.MoeRunnerConfig(...)`. **`inplace=False` is essential** (the default overwrites `x`). SGLang's fused-MoE reads its tensor-parallel group even on one GPU, so `init_distributed()` (mirroring SGLang's own benchmark scripts) must run first. + +These are the API shapes on `main`; the pinned release may differ. That is what Task 10's real-engine test exists to find. On this dev machine neither engine is installed, so the CPU tests inject fakes that reproduce each engine's documented layout contract in eager PyTorch. + +**Files:** +- Create: `src/dispatch/benchmark/engines/vllm_moe.py`, `src/dispatch/benchmark/engines/sglang_moe.py`, `src/dispatch/benchmark/engines/registry.py` +- Modify: `pyproject.toml` +- Test: `tests/unit/test_engine_adapters.py` + +**Interfaces:** +- Consumes: Task 4's `MoEEngine`, `RaceCase`, `BoundLayer`, `LayerFactory`, `fuse_gate_up`, `DispatchEngine`. +- Produces: `VllmEngine(*, fused_experts=None, int8_quant_config=None)` and `SglangEngine(*, fused_experts=None, topk_output_cls=None, runner_config_cls=None)` (each injectable for tests, lazily importing the real engine otherwise); `sglang_moe.init_distributed() -> None`; `registry.ENGINE_NAMES = ("dispatch-naive", "dispatch-persistent", "vllm", "sglang")`, `registry.DEFAULT_BLOCK_MS = (16, 32, 64, 128)`, `registry.build_engines(engine: str, block_ms: Sequence[int] = DEFAULT_BLOCK_MS) -> list[MoEEngine]` (dispatch returns one contestant per tile size; vllm and sglang exactly one). + +- [ ] **Step 1: Write the failing tests** + +Create `tests/unit/test_engine_adapters.py`: + +````python +"""vLLM's and SGLang's adapters, driven through fakes that reproduce each +engine's documented fused-layout contract (gate rows first in w1, per-channel +int8 scales). The fakes are eager PyTorch, so a wrong layout, a wrong scale +shape, or a mutated input in the adapter changes the output and fails here; +the real engines are checked by the `gpu`-marked test on the pod.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any, NamedTuple + +import pytest +import torch +import torch.nn.functional as F # noqa: N812 -- F is the universal PyTorch convention + +from dispatch.benchmark.engines import registry +from dispatch.benchmark.engines.base import ( + dequantized_weights, + make_case, + make_weights, + reference_output, +) +from dispatch.benchmark.engines.sglang_moe import SglangEngine +from dispatch.benchmark.engines.vllm_moe import VllmEngine +from dispatch.kernels.moe_forward import assert_matches_reference +from dispatch.kernels.quantization import quantize_stacked_weights + +DIMS = {"hidden_size": 16, "intermediate_size": 8, "n_experts": 6} +CASE_DIMS = {"hidden_size": 16, "n_experts": 6, "top_k": 2} + + +def _eager_fused_moe( + x: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w1_scale: torch.Tensor | None = None, + w2_scale: torch.Tensor | None = None, +) -> torch.Tensor: + """What vLLM's and SGLang's fused-MoE compute, in eager float32: w1 holds + gate rows then up rows, activation is silu(gate) * up, int8 weights are + scaled per output channel.""" + w1_f, w2_f = w1.float(), w2.float() + if w1_scale is not None and w2_scale is not None: + w1_f, w2_f = w1_f * w1_scale.unsqueeze(-1), w2_f * w2_scale.unsqueeze(-1) + n = w2.shape[2] + out = torch.zeros(x.shape[0], x.shape[1], dtype=torch.float32) + for token in range(x.shape[0]): + for slot in range(topk_ids.shape[1]): + expert = int(topk_ids[token, slot]) + hidden = x[token].float() @ w1_f[expert].T + activated = F.silu(hidden[:n]) * hidden[n:] + out[token] += topk_weights[token, slot].float() * (activated @ w2_f[expert].T) + return out + + +def _fake_vllm_fused_experts( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + quant_config: Any = None, +) -> torch.Tensor: + assert topk_ids.dtype == torch.int32 + assert topk_weights.dtype == torch.float32 + scales = ( + (None, None) if quant_config is None else (quant_config.w1_scale, quant_config.w2_scale) + ) + return _eager_fused_moe(hidden_states, w1, w2, topk_weights, topk_ids, *scales) + + +def _fake_vllm_int8_quant_config(*, w1_scale: torch.Tensor, w2_scale: torch.Tensor) -> Any: + return SimpleNamespace(w1_scale=w1_scale, w2_scale=w2_scale) + + +class _FakeTopKOutput(NamedTuple): + topk_weights: torch.Tensor + topk_ids: torch.Tensor + router_logits: torch.Tensor + + +class _FakeRunnerConfig: + inplace: bool + + def __init__(self, **kwargs: Any) -> None: + self.__dict__.update(kwargs) + + +def _fake_sglang_fused_experts( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_output: _FakeTopKOutput, + moe_runner_config: _FakeRunnerConfig, + use_int8_w8a16: bool = False, + per_channel_quant: bool = False, + w1_scale: torch.Tensor | None = None, + w2_scale: torch.Tensor | None = None, +) -> torch.Tensor: + assert moe_runner_config.inplace is False # the default would overwrite case.x + assert topk_output.topk_ids.dtype == torch.int32 + if use_int8_w8a16: + assert per_channel_quant + return _eager_fused_moe( + hidden_states, + w1, + w2, + topk_output.topk_weights, + topk_output.topk_ids, + w1_scale, + w2_scale, + ) + + +def _vllm() -> VllmEngine: + return VllmEngine( + fused_experts=_fake_vllm_fused_experts, int8_quant_config=_fake_vllm_int8_quant_config + ) + + +def _sglang() -> SglangEngine: + return SglangEngine( + fused_experts=_fake_sglang_fused_experts, + topk_output_cls=_FakeTopKOutput, + runner_config_cls=_FakeRunnerConfig, + ) + + +@pytest.mark.parametrize("make_engine", [_vllm, _sglang], ids=["vllm", "sglang"]) +def test_bf16_layout_conversion_reproduces_the_reference(make_engine: Any) -> None: + weights = make_weights(torch.float32, seed=0, device="cpu", **DIMS) + case = make_case(9, "zipf", dtype=torch.float32, seed=1, device="cpu", **CASE_DIMS) + + layer = make_engine().prepare_bf16(weights)(case) + + assert_matches_reference(layer(), reference_output(case, weights)) + + +@pytest.mark.parametrize("make_engine", [_vllm, _sglang], ids=["vllm", "sglang"]) +def test_int8_scale_layout_reproduces_the_dequantized_reference(make_engine: Any) -> None: + weights = make_weights(torch.float32, seed=0, device="cpu", **DIMS) + qweights = quantize_stacked_weights(weights) + case = make_case(9, "uniform", dtype=torch.float32, seed=1, device="cpu", **CASE_DIMS) + + layer = make_engine().prepare_int8(qweights)(case) + + assert_matches_reference(layer(), reference_output(case, dequantized_weights(qweights))) + + +@pytest.mark.parametrize("make_engine", [_vllm, _sglang], ids=["vllm", "sglang"]) +def test_a_swapped_gate_and_up_layout_is_caught(make_engine: Any) -> None: + """The failure this whole adapter layer exists to prevent: fusing up + before gate silently computes silu(up) * gate. The reference check must + reject it.""" + weights = make_weights(torch.float32, seed=0, device="cpu", **DIMS) + swapped = type(weights)(gate=weights.up, up=weights.gate, down=weights.down) + case = make_case(9, "uniform", dtype=torch.float32, seed=1, device="cpu", **CASE_DIMS) + + layer = make_engine().prepare_bf16(swapped)(case) + + with pytest.raises(AssertionError): + assert_matches_reference(layer(), reference_output(case, weights)) + + +@pytest.mark.parametrize("make_engine", [_vllm, _sglang], ids=["vllm", "sglang"]) +def test_bound_layer_does_not_mutate_the_case_and_is_repeatable(make_engine: Any) -> None: + weights = make_weights(torch.float32, seed=0, device="cpu", **DIMS) + case = make_case(9, "uniform", dtype=torch.float32, seed=1, device="cpu", **CASE_DIMS) + x_before = case.x.clone() + layer = make_engine().prepare_bf16(weights)(case) + + assert torch.equal(layer(), layer()) + assert torch.equal(case.x, x_before) + + +def test_registry_builds_a_tile_size_sweep_for_dispatch_and_one_engine_otherwise() -> None: + naive = registry.build_engines("dispatch-naive", (16, 64)) + assert [engine.name for engine in naive] == ["dispatch-naive-bm16", "dispatch-naive-bm64"] + assert len(registry.build_engines("vllm")) == 1 + assert len(registry.build_engines("sglang")) == 1 + + +def test_registry_rejects_unknown_engines() -> None: + with pytest.raises(ValueError, match="unknown engine"): + registry.build_engines("tensorrt") + with pytest.raises(ValueError, match="unknown engine"): + registry.build_engines("dispatch-cuda") +```` + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: `uv run pytest tests/unit/test_engine_adapters.py -v` +Expected: collection FAIL with `ModuleNotFoundError: No module named 'dispatch.benchmark.engines.vllm_moe'`. + +- [ ] **Step 3: Write the implementation** + +`src/dispatch/benchmark/engines/vllm_moe.py`: + +````python +"""vLLM's fused-MoE as a race contestant. Calls `fused_experts` directly with +the race's fixed topk_ids/topk_weights (no gating, so routing is identical to +every other contestant's). vLLM is imported lazily, on first use, so this +module -- and its tests, which inject fakes -- import on a machine without it. + +API checked live 2026-09-19 against vllm-project/vllm main and the 0.29.0 +release: `fused_experts(hidden_states, w1, w2, topk_weights, topk_ids, ..., +quant_config=)`, with weight-only int8 built by `int8_w8a16_moe_quant_config` +(per-output-channel scales: w1_scale (E, 2N), w2_scale (E, K)). +""" + +from __future__ import annotations + +import functools +from collections.abc import Callable +from typing import Any + +import torch + +from dispatch.benchmark.engines.base import BoundLayer, LayerFactory, RaceCase, fuse_gate_up +from dispatch.kernels.moe_forward import StackedExpertWeights +from dispatch.kernels.quantization import QuantizedStackedExpertWeights + + +class VllmEngine: + name = "vllm" + + def __init__( + self, + *, + fused_experts: Callable[..., torch.Tensor] | None = None, + int8_quant_config: Callable[..., Any] | None = None, + ) -> None: + self._fused_experts = fused_experts + self._int8_quant_config = int8_quant_config + + def prepare_bf16(self, weights: StackedExpertWeights) -> LayerFactory: + return self._factory(fuse_gate_up(weights.gate, weights.up), weights.down, None) + + def prepare_int8(self, qweights: QuantizedStackedExpertWeights) -> LayerFactory: + quant_config = self._resolve_int8_quant_config()( + w1_scale=fuse_gate_up(qweights.gate.scale, qweights.up.scale), + w2_scale=qweights.down.scale, + ) + w1 = fuse_gate_up(qweights.gate.data, qweights.up.data) + return self._factory(w1, qweights.down.data, quant_config) + + def _factory(self, w1: torch.Tensor, w2: torch.Tensor, quant_config: Any) -> LayerFactory: + fused_experts = self._resolve_fused_experts() + + def bind(case: RaceCase) -> BoundLayer: + return functools.partial( + fused_experts, + case.x, + w1, + w2, + case.topk_weight, + case.topk_idx.to(torch.int32), + quant_config=quant_config, + ) + + return bind + + def _resolve_fused_experts(self) -> Callable[..., torch.Tensor]: + if self._fused_experts is None: + from vllm.model_executor.layers.fused_moe.fused_moe import ( # noqa: PLC0415 + fused_experts, + ) + + self._fused_experts = fused_experts + return self._fused_experts + + def _resolve_int8_quant_config(self) -> Callable[..., Any]: + if self._int8_quant_config is None: + from vllm.model_executor.layers.fused_moe.config import ( # noqa: PLC0415 + int8_w8a16_moe_quant_config, + ) + + self._int8_quant_config = int8_w8a16_moe_quant_config + return self._int8_quant_config +```` + +`src/dispatch/benchmark/engines/sglang_moe.py`: + +````python +"""SGLang's fused-MoE as a race contestant, called through `fused_experts` +with the race's fixed topk_ids/topk_weights. SGLang is imported lazily. + +API checked live 2026-09-19 against sgl-project/sglang main and the 0.5.20 +release, whose fused-MoE code was recently reorganized: `fused_experts` lives +in `sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe` and takes a +`StandardTopKOutput` plus a `MoeRunnerConfig`. The pinned version is recorded +in every output JSON; if the pinned version's API differs, the import fails +loudly rather than falling back to another code path. + +`inplace=False` is essential: the default (`True`) writes the output into +`hidden_states`, which would corrupt `case.x` across repeated timed calls. +""" + +from __future__ import annotations + +import functools +from collections.abc import Callable +from typing import Any + +import torch + +from dispatch.benchmark.engines.base import ( + BoundLayer, + LayerFactory, + RaceCase, + fuse_gate_up, +) +from dispatch.kernels.moe_forward import StackedExpertWeights +from dispatch.kernels.quantization import QuantizedStackedExpertWeights + +_INIT_METHOD = "tcp://127.0.0.1:23456" + + +def init_distributed() -> None: + """SGLang's fused-MoE reads its tensor-parallel group even on one GPU, so + a world-size-1 group must exist first. Mirrors SGLang's own + benchmark/kernels/fused_moe_triton/ scripts. Call once, before `bind`.""" + from sglang.srt.distributed.parallel_state import ( # noqa: PLC0415 + init_distributed_environment, + initialize_model_parallel, + ) + + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group( + backend="nccl", init_method=_INIT_METHOD, world_size=1, rank=0 + ) + init_distributed_environment( + world_size=1, + rank=0, + distributed_init_method=_INIT_METHOD, + local_rank=0, + backend="nccl", + ) + initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=1) + + +class SglangEngine: + name = "sglang" + + def __init__( + self, + *, + fused_experts: Callable[..., torch.Tensor] | None = None, + topk_output_cls: Callable[..., Any] | None = None, + runner_config_cls: Callable[..., Any] | None = None, + ) -> None: + self._fused_experts = fused_experts + self._topk_output_cls = topk_output_cls + self._runner_config_cls = runner_config_cls + + def prepare_bf16(self, weights: StackedExpertWeights) -> LayerFactory: + w1 = fuse_gate_up(weights.gate, weights.up) + return self._factory(w1, weights.down, {}) + + def prepare_int8(self, qweights: QuantizedStackedExpertWeights) -> LayerFactory: + w1 = fuse_gate_up(qweights.gate.data, qweights.up.data) + quant_kwargs = { + "use_int8_w8a16": True, + "per_channel_quant": True, + "w1_scale": fuse_gate_up(qweights.gate.scale, qweights.up.scale), + "w2_scale": qweights.down.scale, + } + return self._factory(w1, qweights.down.data, quant_kwargs) + + def _factory( + self, w1: torch.Tensor, w2: torch.Tensor, quant_kwargs: dict[str, Any] + ) -> LayerFactory: + fused_experts, topk_output_cls, runner_config_cls = self._resolve() + n_experts = int(w1.shape[0]) + + def bind(case: RaceCase) -> BoundLayer: + topk_output = topk_output_cls( + topk_weights=case.topk_weight, + topk_ids=case.topk_idx.to(torch.int32), + router_logits=case.x.new_empty(0), # unused: fused_experts drops it + ) + runner_config = runner_config_cls( + num_experts=n_experts, + num_local_experts=n_experts, + top_k=int(case.topk_idx.shape[1]), + inplace=False, + ) + return functools.partial( + fused_experts, case.x, w1, w2, topk_output, runner_config, **quant_kwargs + ) + + return bind + + def _resolve( + self, + ) -> tuple[Callable[..., torch.Tensor], Callable[..., Any], Callable[..., Any]]: + fused_experts = self._fused_experts + topk_output_cls = self._topk_output_cls + runner_config_cls = self._runner_config_cls + if fused_experts is None or topk_output_cls is None or runner_config_cls is None: + from sglang.srt.layers.moe.moe_runner.base import ( # noqa: PLC0415 + MoeRunnerConfig, + ) + from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import ( # noqa: PLC0415 + fused_experts as sglang_fused_experts, + ) + from sglang.srt.layers.moe.topk import StandardTopKOutput # noqa: PLC0415 + + fused_experts = fused_experts or sglang_fused_experts + topk_output_cls = topk_output_cls or StandardTopKOutput + runner_config_cls = runner_config_cls or MoeRunnerConfig + return fused_experts, topk_output_cls, runner_config_cls +```` + +`src/dispatch/benchmark/engines/registry.py`: + +````python +"""Engine name -> contestants. vLLM and SGLang import lazily, so this module +is importable (and the driver's tests run) without either installed.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from dispatch.benchmark.engines.base import MoEEngine +from dispatch.benchmark.engines.dispatch_kernels import DispatchEngine + +ENGINE_NAMES = ("dispatch-naive", "dispatch-persistent", "vllm", "sglang") +DEFAULT_BLOCK_MS = (16, 32, 64, 128) + + +def build_engines(engine: str, block_ms: Sequence[int] = DEFAULT_BLOCK_MS) -> list[MoEEngine]: + """dispatch-* returns one contestant per tile size in `block_ms` (its own + tuning sweep); vllm and sglang return exactly one, tuned by their own + tuners through their own config folders, not by this call.""" + if engine.startswith("dispatch-"): + backend = engine.removeprefix("dispatch-") + if f"dispatch-{backend}" not in ENGINE_NAMES: + raise ValueError(f"unknown engine {engine!r}; expected one of {ENGINE_NAMES}") + return [DispatchEngine(backend, block_m=block_m) for block_m in block_ms] + if engine == "vllm": + from dispatch.benchmark.engines.vllm_moe import VllmEngine # noqa: PLC0415 + + return [VllmEngine()] + if engine == "sglang": + from dispatch.benchmark.engines.sglang_moe import SglangEngine # noqa: PLC0415 + + return [SglangEngine()] + raise ValueError(f"unknown engine {engine!r}; expected one of {ENGINE_NAMES}") +```` + +- [ ] **Step 4: Update `pyproject.toml`** + +vLLM and SGLang exist only in the pod's engines venv, so mypy must treat them as `Any`, the same boundary already drawn for `triton` and `deep_ep`. Add this block immediately above `[tool.coverage.run]`: + +```toml +# vLLM and SGLang are only installed in Phase 6's rented-pod engine venv, never +# on this dev box or CI (they pin torch==2.13.0 and need a CUDA build), so their +# adapters import them lazily and mypy treats them as Any, same boundary as +# triton and deep_ep above. +[[tool.mypy.overrides]] +module = ["vllm", "vllm.*", "sglang", "sglang.*"] +ignore_missing_imports = true +follow_imports = "skip" +``` + +The fake-engine tests take many positional arguments by construction, so extend the existing test ignore (`PLR0917` is "too many positional arguments"; `PLR0913` is already there). Replace +```toml +"tests/**/*.py" = ["PLR2004", "PLR0913"] +``` +with +```toml +"tests/**/*.py" = ["PLR2004", "PLR0913", "PLR0917"] +``` + +- [ ] **Step 5: Run the tests to verify they pass** + +Run: `uv run pytest tests/unit/test_engine_adapters.py -v` +Expected: 10 passed, including `test_a_swapped_gate_and_up_layout_is_caught` for both engines (the silent-wrong-activation failure this layer exists to prevent). + +- [ ] **Step 6: Commit** + +STATUS bullet: `- **Task 5 (vLLM/SGLang adapters + registry)**: each engine behind one adapter calling its own fused_experts with fixed routing; verified against eager fakes only -- the real engines are not exercised until the pod (Task 10).` + +```bash +make check +git add src/dispatch/benchmark/engines pyproject.toml tests/unit/test_engine_adapters.py docs/STATUS.md +git commit -m "feat: vLLM and SGLang fused-MoE adapters behind the race contract" +``` + +--- + +### Task 6: Race summarizer (tuning rule, rows, markdown) + +**Files:** +- Create: `src/dispatch/benchmark/race_summary.py` +- Test: `tests/unit/test_race_summary.py` + +**Interfaces:** +- Consumes: `TOKEN_COUNTS` from Task 4. Reads the per-engine record JSON that Task 7's `run` writes: `{"config": {"engine", "precision", "tuning_label", ...}, "results": [{"num_tokens", "distribution", "variant", "status", "mean_ms", "p50_ms", "p99_ms", "tflops"}]}`. +- Produces: `RaceRow` (frozen dataclass), `summarize_race(records) -> list[RaceRow]`, `render_markdown(rows) -> str`, `DEFAULT_BLOCK_M = 16`, `TUNING_DISTRIBUTION = "uniform"`. Encodes the pre-registered tuning rule: dispatch's "tuned" variant is chosen on uniform results only and reused for zipf. + +- [ ] **Step 1: Write the failing tests** + +Create `tests/unit/test_race_summary.py`: + +````python +from __future__ import annotations + +from typing import Any + +from dispatch.benchmark.race_summary import render_markdown, summarize_race + + +def _result( + variant: str, tokens: int, distribution: str, mean_ms: float, status: str = "ok" +) -> dict[str, Any]: + return { + "num_tokens": tokens, + "distribution": distribution, + "variant": variant, + "status": status, + "mean_ms": mean_ms, + "p50_ms": mean_ms, + "p99_ms": mean_ms * 1.2, + "tflops": 1.0, + } + + +def _record( + engine: str, results: list[dict[str, Any]], *, tuning_label: str = "sweep" +) -> dict[str, Any]: + return { + "config": {"engine": engine, "precision": "bf16", "tuning_label": tuning_label}, + "results": results, + } + + +def _dispatch_record() -> dict[str, Any]: + return _record( + "dispatch-naive", + [ + _result("dispatch-naive-bm16", 16, "uniform", 1.0), + _result("dispatch-naive-bm64", 16, "uniform", 0.6), + _result("dispatch-naive-bm16", 16, "zipf", 1.1), + _result("dispatch-naive-bm64", 16, "zipf", 0.9), + ], + ) + + +def test_dispatch_default_is_the_kernels_own_tile_size() -> None: + rows = summarize_race([_dispatch_record()]) + + default = [row for row in rows if row.tuning == "default"] + assert {row.variant for row in default} == {"dispatch-naive-bm16"} + assert {row.distribution for row in default} == {"uniform", "zipf"} + + +def test_dispatch_tuned_is_picked_on_uniform_and_reused_for_zipf() -> None: + rows = summarize_race([_dispatch_record()]) + + tuned = {row.distribution: row for row in rows if row.tuning == "tuned"} + assert tuned["uniform"].variant == "dispatch-naive-bm64" + # Not re-tuned on zipf, even though bm16's zipf number (1.1) is worse + # than bm64's (0.9) only by coincidence here: the rule is uniform-only. + assert tuned["zipf"].variant == "dispatch-naive-bm64" + assert tuned["zipf"].mean_ms == 0.9 + + +def test_the_tuning_choice_never_looks_at_zipf_results() -> None: + record = _record( + "dispatch-naive", + [ + _result("dispatch-naive-bm16", 16, "uniform", 1.0), + _result("dispatch-naive-bm64", 16, "uniform", 1.1), + _result("dispatch-naive-bm16", 16, "zipf", 5.0), + _result("dispatch-naive-bm64", 16, "zipf", 0.1), # far better on zipf + ], + ) + + tuned = {row.distribution: row for row in summarize_race([record]) if row.tuning == "tuned"} + + assert tuned["zipf"].variant == "dispatch-naive-bm16" + + +def test_engines_with_their_own_tuner_keep_the_label_they_ran_under() -> None: + rows = summarize_race( + [ + _record("vllm", [_result("vllm", 16, "uniform", 0.5)], tuning_label="tuned"), + _record("vllm", [_result("vllm", 16, "uniform", 0.8)], tuning_label="default"), + ] + ) + + assert {(row.tuning, row.mean_ms) for row in rows} == {("tuned", 0.5), ("default", 0.8)} + + +def test_refused_results_never_reach_a_table() -> None: + record = _record("vllm", [_result("vllm", 16, "uniform", 0.5, status="refused")]) + + assert summarize_race([record]) == [] + + +def test_markdown_has_one_table_per_precision_and_tuning_with_a_column_per_engine() -> None: + rows = summarize_race( + [ + _dispatch_record(), + _record("vllm", [_result("vllm", 16, "uniform", 0.5)], tuning_label="tuned"), + ] + ) + + text = render_markdown(rows) + + assert "### bf16, tuned" in text + assert "### bf16, default" in text + assert "| tokens | routing | dispatch-naive | vllm |" in text + assert "| 16 | uniform | 0.600 / 0.720 | 0.500 / 0.600 |" in text + assert "| 16 | zipf | 0.900 / 1.080 | n/a |" in text +```` + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: `uv run pytest tests/unit/test_race_summary.py -v` +Expected: collection FAIL, `ModuleNotFoundError: No module named 'dispatch.benchmark.race_summary'`. + +- [ ] **Step 3: Write the implementation** + +`src/dispatch/benchmark/race_summary.py`: + +````python +"""Turns the race driver's per-engine JSON records into comparable rows and a +markdown table. Pure, so the selection rule is unit-tested without a GPU. + +Tuning rule, fixed before any run (docs/plans/2026-09-19-phase-6-final- +benchmark-plan.md): every contestant is tuned under *uniform* routing, the +distribution the engines' own tuners use, and then timed under both +distributions. vLLM and SGLang are tuned by their own tuners (one run per +tuning label); dispatch sweeps its tile size `block_m`, and its "tuned" row +is the variant fastest on uniform routing at that token count, applied to the +zipf row too. Its "default" row is `block_m` 16, the kernel's default. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + +from dispatch.benchmark.engines.base import TOKEN_COUNTS + +DEFAULT_BLOCK_M = 16 +TUNING_DISTRIBUTION = "uniform" + + +@dataclass(frozen=True) +class RaceRow: + precision: str + num_tokens: int + distribution: str + engine: str + tuning: str # "default" | "tuned" + variant: str + mean_ms: float + p50_ms: float + p99_ms: float + tflops: float + + +def summarize_race(records: Iterable[dict[str, Any]]) -> list[RaceRow]: + """Refused results are excluded (the driver already exits non-zero on + them); only measured rows reach a table.""" + rows: list[RaceRow] = [] + for record in records: + config = record["config"] + ok = [result for result in record["results"] if result["status"] == "ok"] + if config["engine"].startswith("dispatch-"): + rows.extend(_dispatch_rows(config, ok)) + else: + rows.extend( + _row(config, result, tuning=config["tuning_label"], variant=result["variant"]) + for result in ok + ) + return rows + + +def render_markdown(rows: list[RaceRow]) -> str: + """One table per (precision, tuning): a row per (tokens, distribution), a + column per engine, each cell 'mean / p99' in milliseconds.""" + blocks: list[str] = [] + for precision, tuning in sorted({(row.precision, row.tuning) for row in rows}): + subset = [row for row in rows if (row.precision, row.tuning) == (precision, tuning)] + engines = sorted({row.engine for row in subset}) + lookup = {(row.num_tokens, row.distribution, row.engine): row for row in subset} + keys = sorted( + {(row.num_tokens, row.distribution) for row in subset}, + key=lambda key: (_token_rank(key[0]), key[1]), + ) + lines = [ + f"### {precision}, {tuning} (mean / p99 ms per routed-MoE layer call)", + "", + "| tokens | routing | " + " | ".join(engines) + " |", + "|---|---|" + "---|" * len(engines), + ] + for tokens, distribution in keys: + cells = [ + f"{found.mean_ms:.3f} / {found.p99_ms:.3f}" + if (found := lookup.get((tokens, distribution, engine))) + else "n/a" + for engine in engines + ] + lines.append(f"| {tokens} | {distribution} | " + " | ".join(cells) + " |") + blocks.append("\n".join(lines)) + return "\n\n".join(blocks) + "\n" + + +def _dispatch_rows(config: dict[str, Any], ok: list[dict[str, Any]]) -> list[RaceRow]: + rows: list[RaceRow] = [] + for num_tokens in sorted({result["num_tokens"] for result in ok}): + at_tokens = [result for result in ok if result["num_tokens"] == num_tokens] + default = [r for r in at_tokens if _block_m(r["variant"]) == DEFAULT_BLOCK_M] + rows.extend( + _row(config, result, tuning="default", variant=result["variant"]) for result in default + ) + tuning_pool = [r for r in at_tokens if r["distribution"] == TUNING_DISTRIBUTION] + if not tuning_pool: + continue + best_variant = min(tuning_pool, key=lambda r: r["mean_ms"])["variant"] + rows.extend( + _row(config, result, tuning="tuned", variant=best_variant) + for result in at_tokens + if result["variant"] == best_variant + ) + return rows + + +def _row(config: dict[str, Any], result: dict[str, Any], *, tuning: str, variant: str) -> RaceRow: + return RaceRow( + precision=config["precision"], + num_tokens=result["num_tokens"], + distribution=result["distribution"], + engine=config["engine"], + tuning=tuning, + variant=variant, + mean_ms=result["mean_ms"], + p50_ms=result["p50_ms"], + p99_ms=result["p99_ms"], + tflops=result["tflops"], + ) + + +def _block_m(variant: str) -> int: + return int(variant.rsplit("-bm", 1)[1]) + + +def _token_rank(num_tokens: int) -> int: + return TOKEN_COUNTS.index(num_tokens) if num_tokens in TOKEN_COUNTS else len(TOKEN_COUNTS) +```` + +- [ ] **Step 4: Run the tests to verify they pass** + +Run: `uv run pytest tests/unit/test_race_summary.py -v` +Expected: 6 passed, including `test_the_tuning_choice_never_looks_at_zipf_results`. + +- [ ] **Step 5: Commit** + +STATUS bullet: `- **Task 6 (race summarizer)**: the pre-registered tuning rule as code -- dispatch's tuned tile size is picked on uniform routing only and reused for zipf; refused results never reach a table.` + +```bash +make check +git add src/dispatch/benchmark/race_summary.py tests/unit/test_race_summary.py docs/STATUS.md +git commit -m "feat: race summarizer encoding the uniform-only tuning rule" +``` + +--- + +### Task 7: Race driver (`prepare` / `run` / `merge`) + +**Files:** +- Create: `scripts/run_engine_race.py` +- Test: `tests/unit/test_run_engine_race.py` + +**Interfaces:** +- Consumes: Tasks 4-6 (`make_weights`, `make_case`, `reference_output`, `dequantized_weights`, `registry`, `summarize_race`, `render_markdown`), `dispatch.kernels.bench` (`time_grouped_gemm`, `grouped_gemm_flops`, `DO_BENCH_*`, `ROUTING_DISTRIBUTIONS`, `KernelBenchmarkSummary`). +- Produces: CLI `python -m scripts.run_engine_race {prepare,run,merge}`; `prepare_inputs(out_dir, *, num_tokens, distributions, dtype, seed, device)`; `run_engines(inputs_dir, engines, *, precision, device, time_fn=None) -> list[dict]` (each result `status` is `ok` or `refused`); `describe_environment(device) -> dict`. `run` writes `/