diff --git a/.agents/benchmark-record.md b/.agents/benchmark-record.md index 9cf2a1522..4caa8b338 100644 --- a/.agents/benchmark-record.md +++ b/.agents/benchmark-record.md @@ -19752,3 +19752,68 @@ the residual lives inside the reference's own spread. Evidence: `dgx:~/work/dspark-w6/interleaved5.log`, `oracle5_rep{1..5}.json`, `diag.log`. + +## KERNEL-MOE-ROUTER-WARP — first DEVICE run of a kernel that had never been compiled; kernel-level 1.363x ESTABLISHED, step-level NOT SEPARABLE; and the canonical 35B grid is STALE (2026-08-12, `row/A-35B-RESIDUAL`, base `origin/main` `bbc482a2`, source `6c3be5c3`, GB10 sm_121a, #378) + +Ran under punch-list item 5 of [roadmap-v1-completion.md](specs/roadmap-v1-completion.md) §3 +(`ROAD-V1-A`, "close the 35B c1/c2 residual"). Two findings, one of them a record correction +that matters more than the measurement. + +**1. The warp-shuffle MoE router had NOT landed.** The premise this task was dispatched on — +that the router landed and took 35B c1/c4 to ~0.98x — is false in both halves. `git +merge-base --is-ancestor` says none of `9f7afa7a`, `cce81c7e`, `be6a1f57`, `4f4ee725` is an +ancestor of `bbc482a2`, and `git log -S` on `MoeRouterTopKWarpKernel`, +`MoeRouterWarpEnabled`, `VT_MOE_ROUTER_WARP`, `MoeRouterWarpTreeSum` and `moe_router_warp` +returns **0 commits on main** for every one of them. The work sits on a local-only branch +`row/KERNEL-MOE-ROUTER-WARP`, never pushed to origin, no PR, 9 commits behind main, and +`cce81c7e` states outright that the `.cu` was never compiled or run. The tree never claimed +otherwise: `KERNEL-MOE-ROUTER-WARP` appears in no matrix and no doc on main. So the 35B +residual was, in this part, a **landing** problem and not an optimisation problem. + +The `~0.98x / 0.87x / 0.92x` figures in the dispatch appear to be the **`SPEC-DSPARK`** row +(`0.870x-0.981x` -> `0.986x-0.995x`, #442) — speculative decoding, a different row — not the +router and not the 35B grid. + +**2. The canonical 35B grid on record is STALE.** `docs/BENCHMARKS.md` carries "CANONICAL +2026-08-11 @`348c265d`: 0.918x-0.972x (c1 0.9708, c2 0.9293, c4 0.9719, c8 0.9183, c16 +0.9264, c32 0.9377)". `348c265d` is an ancestor of main, but **136 `src/`+`include/` commits +have landed since**, including `972a152c` *"perf(35B): shared-expert down_proj emits bf16 — +one CastF32 per layer-step gone (+2.05% c8)"*, which merged **1 h 36 min after** the grid was +captured and reports +2.05% c8 / +0.79% c4 measured, bit-identical. Holding the vLLM arm +fixed that implies c8 ~0.937 and c4 ~0.980 today — an **estimate, not a measurement**. So the +honest answer to "where is the 35B gap today" is: **nobody knows, and a regrid is owed**; the +last measured value is stale in our favour by at least one landed lever. No axis is claimed +closed and none is claimed open beyond what `348c265d` measured. + +**3. The router, measured for the first time.** Full method, gate table and caveats in +[moe-router-topk-single-warp.md](specs/moe-router-topk-single-warp.md) §12. Summary: + +- It **compiles** — not previously known. Correctness holds everywhere: gates 1/2/3/6 green, + 35B `test_qwen36_paged_engine` **315/315 on both arms of the same binary**, plus six more + green legs during the A/B. Gate-1 mutation RED re-proven (76 assertions, 2 cases). +- Which kernel ran is proven **by name** in `nsys --cuda-graph-trace=node`: + `MoeRouterTopKWarpKernel<__nv_bfloat16,(int)8>` (VPT=8 = E=256) in the ON arm with no block + kernel present, and the converse under `VT_MOE_ROUTER_WARP=0`. 1280 calls in every leg. +- **Kernel-level, 3 reps/arm order-alternated, one flock:** warp 10.422 ms (10.307/10.529) + vs block 14.201 ms (14.141/14.275) — **-3.779 ms, -26.6%, 1.363x, bands non-overlapping.** +- **Step-level: NOT SEPARABLE, not claimed.** Total GPU-busy read 1601.890 vs 1619.395 ms + (ratio 0.9892) but the *unchanged* kernels — identical code, identical 155,310 launches — + differ 13.726 ms between arms with a 24.140 ms spread across the six legs, **6.4x the + 3.779 ms effect**. The window is prefill+decode aggregated, which dilutes 40 router + calls/step across ~1600 ms. Sized against #378's 16.1863 ms/step decode GPU-busy the saving + is **0.083-0.118 ms/step = ~0.5-0.75% of decode GPU-busy**, i.e. **below** #378's ~1.5% + estimate: the warp kernel closes ~79% of the gap to `topkGating` (12.98 -> 8.14 us/call + mean vs 6.85), not all of it. +- Therefore **no default-flip credit** under spec gate 7, and #378 cannot close `ROAD-V1-A` + alone — as its own §10 already said. + +**Denominator note.** No vLLM arm was run here, so no ratio-vs-vLLM is asserted. The pinned +oracle source is at the pin (`5559679229bc`, asserted by commit), but the shared +`~/venvs/vllm-oracle` symlink currently resolves to `vllm-oracle-v0.25.0-stage`, **not** the +pin; it was left untouched for other sessions. Any denominator taken through that symlink is +the 0.25.0 rollback — the caveat #378 already carries. + +**Next traceable hypotheses:** a decode-only profiler window to turn the 0.5-0.75% estimate +into a measurement; then the gate-GEMM epilogue fold and batching the 40 per-step launches. +No ceiling is declared. **Owed:** gate 5 `compute-sanitizer memcheck`; the 35B canonical +regrid on current main; a vLLM denominator re-measured against the pin. diff --git a/.agents/roadmap_v1.md b/.agents/roadmap_v1.md index b7cfe8ea2..4eb69289e 100644 --- a/.agents/roadmap_v1.md +++ b/.agents/roadmap_v1.md @@ -110,6 +110,7 @@ issue is not yet placed. Keyed record: update in place, never append. | [#264](https://github.com/mudler/vllm.cpp/issues/264) | `SAMPLE-LOGPROB-TOKEN-IDS` | `logprob_token_ids`: generative scoring over an explicit token set is unported | feature | | [#365](https://github.com/mudler/vllm.cpp/issues/365) | `PERF-27B-DENSE-MARLIN-GATEUP` | 27B gap decomposed vs vLLM's own decode profile; dense W4A16 MLP bypassed the fused gate_up seam (spec `specs/perf-27b-dense-marlin-gateup.md`) | bug | | [#362](https://github.com/mudler/vllm.cpp/issues/362) | — | First VALID 27B per-kernel decode attribution: whole-run trace windowed by the profiler's own cudaProfilerStart/Stop | bug | +| [#378](https://github.com/mudler/vllm.cpp/issues/378) | `KERNEL-MOE-ROUTING` | MoE router top-k as one warp per token, byte-exact via a stride-32 lane map; a lever inside the routing family, so no new stable row (spec `specs/moe-router-topk-single-warp.md`) | perf | | [#223](https://github.com/mudler/vllm.cpp/issues/223) | `SAMPLE-PROMPT-LOGPROBS` | `prompt_logprobs` is a silent no-op: the runner never computes logits at prompt positions | bug | | [#294](https://github.com/mudler/vllm.cpp/issues/294) | `SERVE-ASYNC-LLM` | `test_async_llm`: reusing an aborted request id races the core abort — 26% failure rate under contention, on main | bug | | [#222](https://github.com/mudler/vllm.cpp/issues/222) | `SERVE-GATE-ONLINE` | The online-serving gate aborts on current main: harness looks for `examples/server`, the build produces `examples/vllm-server` — paths repaired by `2b262622`/`8fce04d3`; residue is the drift guard, the `27n` key and a model gate that recorded a skip as passed | bug | diff --git a/.agents/specs/moe-router-topk-single-warp.md b/.agents/specs/moe-router-topk-single-warp.md new file mode 100644 index 000000000..1a98bf194 --- /dev/null +++ b/.agents/specs/moe-router-topk-single-warp.md @@ -0,0 +1,701 @@ +# MoE router top-k: one warp per token, zero barriers, byte-exact + +**Issue:** [#378 — MoE router top-k is 1.9x slower than vLLM's `topkGating` +(12.98 vs 6.85 us/call) — and the "byte-exactness forbids it" conclusion was too +strong](https://github.com/mudler/vllm.cpp/issues/378) + +**Row:** `KERNEL-MOE-ROUTING` — the existing kernel-matrix row that owns router +top-k, align, permute/unpermute and combine, whose current claim is +`CLAIM-MOE-DECODE-PARALLEL-1` (landed by `6a8c5cf9`). This work is a lever +*inside* that family, not a new kernel family, so it registers no new matrix +row. Work branch and role claim: `KERNEL-MOE-ROUTER-WARP`. + +**Parity pin:** vLLM `555967922` (0.26.0.dev0), per +[upstream-sync.md](../upstream-sync.md) + +**Lifecycle:** `READY`. This spec authorizes an implementation and a focused +correctness gate. It authorizes **no** performance claim, no default-flip +credit, and no benchmark record; the operator owns the A/B, the `ncu` +attribution and the token gate. + +## 1. Scope + +Add a second CUDA realization of the **ungrouped** MoE router top-k +(`vt::MoeRouterTopK`, `MoeScoringFunc` softmax, `num_expert_group == 0`) that +runs **one warp per token with the whole logit row in registers**: no shared +memory, no `__syncthreads()`, and one global read of the row instead of two. +It must be **bit-identical** to the incumbent `MoeRouterTopKKernel` +on every input, including NaN, Inf, exact ties and `k > E`. + +In scope: + +- `MoeRouterTopKWarpKernel` in `src/vt/cuda/cuda_moe.cu` for + `VPT ∈ {1,2,4,8}`, i.e. `E ∈ {32,64,128,256}`. +- A portable, host-compilable header carrying the lane map, the two register + reduction trees and the selector, so the numerics-critical part is testable + without a GPU. +- Dispatch from `RouterDispatch`, default ON, with a same-binary rollback. +- Tests: the reduction-order equivalence proof (host, runnable anywhere) and + extra CUDA parallel-vs-serial cases at the real gate shape. + +Explicitly out of scope: + +- The grouped-topk (`noaux_tc`) router. It is a separate kernel and is not + touched, so DeepSeek-class routing stays byte-identical by construction. +- The `Serial` template path. It stays exactly as it is: it is the byte-exact + oracle the parity test compares against, and a change there would invalidate + the oracle rather than test the candidate. +- The CPU reference (`cpu_ops.cpp MoeRouterTopKKernel`). +- The two structurally larger prizes named in #378 and **not attempted here**: + folding the router into the preceding gate-GEMM epilogue, and batching the + 40 per-step launches. Each needs its own row. +- Any arithmetic change. See §4. + +## 2. The measured defect + +From the 35B-A3B both-arms decode trace in #378 +(`nvidia/Qwen3.6-35B-A3B-NVFP4`@`491c2f1e`, GB10, batch-1): + +| kernel | ms/step | calls | us/call | +|---|---:|---:|---:| +| ours `vt::cuda::MoeRouterTopKKernel` | 0.5193 | 40 | 12.98 | +| vLLM `vllm::moe::topkGating<8,256,4,16,32,int,__nv_bfloat16,…>` | 0.2740 | 40 | 6.85 | + +1.9x, ~0.245 ms/step, **~1.5% of GPU-busy**. Both numbers are **PROVISIONAL**: +the vLLM arm ran against the 0.25.0 rollback, not the recorded pin (#375). They +are the motivation for this row, not evidence for it. + +The cost is structural, and re-verified against the current tree at +`6dbedf9f`. `src/vt/cuda/cuda_moe.cu:61-201`, launched at `:420-432` with +**one block of 256 threads per token** and `E * sizeof(float)` dynamic shared +memory. At `E=256`, one token costs: + +| phase | anchor | `__syncthreads()` | +|---|---|---:| +| block-tree max over 256 lanes | `:73-79` | 1 + 8 + 1 | +| second global read of the logit row + `expf` | `:83-87` | — | +| block-tree sum over 256 lanes | `:89-95` | 1 + 8 + 1 | +| normalize | `:103` | 1 | +| k=8 rounds x (warp argmax + leader scan) | `:139-195` | 8 x 2 | +| **total** | | **37** | + +plus 3 KiB of dynamic shared memory (`sp[E]`), 1 KiB `red[256]`, 1 KiB +`redi[256]`, and **two** full global reads of the row (`:71` and `:83`). + +`topkGating` has **zero** barriers and **zero** shared memory: the row lives in +`float row_chunk[VPT]` registers +(`csrc/libtorch_stable/moe/topk_softmax_kernels.cu:352` @ `555967922`), the +reductions are `VLLM_SHFL_XOR_SYNC_WIDTH` butterflies (`:417-421`, `:433-437`), +and the top-k masks the winner in the owning lane's own register (`:566-577`). + +## 3. Why the previous conclusion was too strong + +`6a8c5cf9` parallelized the argmax and recorded: + +> vLLM's 6.7us register-fused topkGating reorders the softmax reduction so it is +> off-limits under byte-exactness — only the serial argmax was removable + +and `cuda_moe.cu:161-163` states the softmax trees are deliberately untouched. +That conclusion is correct **for vLLM's own lane partition** and false in +general, because it conflates "register-resident" with "reassociated". + +vLLM's partition really does reassociate. At the traced instantiation +`topkGating<8,256,4,16,32,…>`: `ELTS_PER_LDG = 16/sizeof(bf16) = 8`, +`THREADS_PER_ROW = 256/8 = 32`, `LDG_PER_THREAD = 1`, and +`first_elt_read_by_thread = thread_group_idx * 8` (`:344-346`), so **lane `L` +owns the contiguous experts `[8L, 8L+8)`**. Summing eight contiguous experts in +a lane and then butterflying is a different association than our block tree, so +it lands on a different last ulp. That is a property of *that map*, not of +register residency. + +**A different lane map is bit-exact.** In the block tree + +```c +for (int s = kBlock / 2; s > 0; s /= 2) { + if (threadIdx.x < s) red[threadIdx.x] += red[threadIdx.x + s]; + __syncthreads(); +} +``` + +levels `s = 128, 64, 32` all have `s ≡ 0 (mod 32)`, so `t` and `t + s` always +have the **same lane id** `L = t mod 32`. Those three levels therefore never +cross a lane: they combine exactly the entries `red[L + 32q]`, `q ∈ [0,8)`, and +they combine them by the standard halving recursion on `q`. Levels +`s = 16 … 1` operate entirely inside warp 0 and are what +`__shfl_down_sync(0xffffffffu, v, s)` reproduces. + +So: **give lane `L` slot `q` the expert `L + 32q` and reduce the slots with the +same halving tree, then shuffle.** Identical float operations, identical +association, identical operands — bit-identical by structural congruence, not by +an appeal to associativity. Zero barriers and zero shared memory fall out. + +Written in the leaf order of the tree, the offsets are +`32 * {0,4,2,6,1,5,3,7}`, the bit-reversal of `{0..7}` (#378's phrasing). The +stride form above is the same statement and is the one implemented, because it +is checkable by inspection against the incumbent loop. + +The derivation is worked per `VPT` in §5. It is **not** assumed for any `VPT` +that is not shown there, and a `VPT` that cannot be shown exact is not +dispatched. + +The top-k argmax needs none of this. It is a reduction over the total order +"higher value, then lower expert index", which is associative and commutative, +so any grouping yields the same `(value, index)` — the argument already recorded +at `cuda_moe.cu:156-163` and unchanged here. + +## 4. Arithmetic held fixed — this is a SHAPE port, not a MATH port + +The candidate keeps every one of the incumbent's arithmetic decisions. vLLM +does four of them differently, and porting any of them changes tokens: + +| decision | ours (`cuda_moe.cu`) | vLLM @ `555967922` | verdict | +|---|---|---|---| +| normalize | `pj = sum > 0.f ? sp[j] / sum : 0.f` — a **divide** (`:99`) | `reciprocal_row_sum = 1.f/row_sum; row_chunk *= reciprocal` (`:445-451`) | KEEP OURS | +| `sum <= 0` guard | present (`:99`) | absent | KEEP OURS | +| renormalize | `weights[j] /= denom` (`:198`) | `scale = routed_scaling_factor; scale /= denom; output *= scale` (`:582-590`) | KEEP OURS | +| `denom <= 0` | `if (!(denom > 0.f)) denom = 1.f` (`:197`) | `denom = selected_sum > 0.f ? selected_sum : 1.f` (`:584`) | equivalent; keep our spelling | +| mask value | `-INFINITY` (`:189`) | `-10000.f`, and only `if (k_idx+1 < k)` (`:570-576`) | KEEP OURS | +| max seed | `m = -INFINITY` then fold (`:71`) — erases NaN | `thread_max = row_chunk[0]` then fold (`:410-414`) | KEEP OURS | +| NaN/Inf clamp | `if (!isfinite(pj)) pj = 0.f`, **after** normalize (`:100-101`) | `if (isnan||isinf) 0.f`, after normalize (`:465-471`) | already agrees | +| `best < 0` | `indices = -1`, `weights = -INFINITY` (`:189-191`) | `indices = NUM_EXPERTS` / `-1` for pad rows (`:559-561`) | KEEP OURS | + +The renormalize row is independently load-bearing beyond this row: the pin at +`555967922` folds `routed_scaling_factor / denom` into one multiply where the +prior `e24d1b24` pin divided. Ours divides. **Advancing this op past the +0.25.0 arithmetic is a token-affecting event regardless of this row** — relevant +to #375 and to roadmap C10, and out of scope here. + +## 5. The lane map, derived per dispatched VPT + +Notation: `l[·]` are the logits of one token; `E_j = expf(l[j] - mx)`; +`R_j = fmaxf(-INFINITY, l[j])`; `⊔` is `fmaxf`; the block kernel's per-thread +seeds are `e[t]` (sum) and `m[t]` (max), `t ∈ [0,256)`. + +For every dispatched `E ≤ 256 = blockDim`, the strided seed loops at `:71` and +`:83` give each thread **at most one** expert, so + +``` +m[t] = R_t and e[t] = (+0.0f) + E_t for t < E +m[t] = -INFINITY, e[t] = +0.0f for t >= E (loop body never runs) +``` + +**The sum seed is `(+0.0f) + E_t`, not `E_t`** (corrected 2026-08-11; the +original text wrote `e[t] = E_t`). The incumbent is +`float acc = 0.0f; ... acc += ex` (`cuda_moe.cu:84-89`), so every occupied +thread's seed is an addition against `+0.0f`, and the candidate's per-lane tree +starts from the loaded value instead. `(+0.0f) + x` differs from `x` for exactly +two classes of `x`: `-0.0f`, which `expf` never produces (it never sets the sign +bit), and a **signalling** NaN, which the addition quiets and a bare copy does +not — and `expf` returns only quiet NaNs. So the two spellings coincide on every +value these seeds can hold. That is a **third instance of the same contained +class** as the two zero-leaf arguments below and in §8.1: an operation that is +formally distinguishable but whose distinguishing inputs are unreachable here. +It is called out rather than glossed because "unreachable" is the load-bearing +word, and a future change to what feeds the seed would have to re-establish it. + +The three cross-warp levels give, for `t = L < 32`: + +``` +C(L) = ((e[L] + e[L+128]) + (e[L+64] + e[L+192])) + + ((e[L+32] + e[L+160]) + (e[L+96] + e[L+224])) +``` + +**VPT = 8 (E = 256).** No padding; `e[L+32q] = E_{L+32q}` for all `q ∈ [0,8)`. +Halving tree on `v[q] = E_{L+32q}`: + +``` +s=4: v0 = E_L + E_L+128 v1 = E_L+32 + E_L+160 + v2 = E_L+64 + E_L+192 v3 = E_L+96 + E_L+224 +s=2: v0 = (E_L + E_L+128) + (E_L+64 + E_L+192) + v1 = (E_L+32 + E_L+160) + (E_L+96 + E_L+224) +s=1: v0 = v0 + v1 +``` + +which is `C(L)` term for term, parenthesis for parenthesis. **EXACT.** + +**VPT = 4 (E = 128).** `e[t] = +0.0f` for `t ∈ [128,256)`, so + +``` +C(L) = ((E_L + 0) + (E_L+64 + 0)) + ((E_L+32 + 0) + (E_L+96 + 0)) +``` + +Dropping the zero leaves is bit-exact here because `x + (+0.0f) == x` for every +`x` these seeds can hold: `expf` returns `+0.0f` or a positive value and never +`-0.0f` (the one float for which `x + 0.0f` flips the sign bit), and `±INFINITY` +and NaN are absorbed unchanged. A NaN sum cannot leak either: it is clamped at +`:100`. So `C(L) = (E_L + E_{L+64}) + (E_{L+32} + E_{L+96})`, which is the +halving tree on `v[q] = E_{L+32q}`, `q ∈ [0,4)`. **EXACT.** + +**VPT = 2 (E = 64).** `C(L) = ((E_L + 0) + (0 + 0)) + ((E_{L+32} + 0) + (0 + 0)) += E_L + E_{L+32}`, the halving tree on two slots. **EXACT.** + +**VPT = 1 (E = 32).** `C(L) = E_L`, a zero-level tree. **EXACT.** + +**Max.** Identical structure with `⊔` for `+`. The per-element seed +`R_j = fmaxf(-INFINITY, l[j])` is applied **before** the tree because that is +what the incumbent does (`float m = -INFINITY; ... m = fmaxf(m, Load(...))`, +`cuda_moe.cu:70-71`), and reproducing the incumbent verbatim is the whole method +of this row. **It stays.** Dropping the `-INFINITY` pad leaves for `VPT < 8` is +bit-exact because `x ⊔ -INFINITY == x` for every non-NaN `x` including `-0.0f` +(verified exhaustively over all 2^32 non-NaN `x` during review). + +**Corrected 2026-08-11 — what the seed does, and does not, buy.** The original +text said the seed "is what erases NaN" and that reproducing it "is what makes +an all-NaN row behave identically". Review falsified the second half by +mutation: deleting the seed from `MoeRouterWarpTreeMax` fails **only** on the +`mx` intermediate; the weights and the indices stay byte-identical on every +case, `all-nan` included. `fmaxf` already returns the non-NaN operand, so the +tree erases a NaN whether or not the leaves were seeded; the seed changes the +result only when a lane holds *nothing but* NaN, and then only by making `mx` +`-INFINITY` instead of NaN — which `expf(l - mx)` washes out, because +`expf(NaN - anything)` is NaN either way, every prob is then clamped to `0.0f` +at `:100`, and the tie-break hands back `0,1,…,k-1` regardless. + +So the seed's **only** guard in the gate is the `mx` intermediate comparison in +`tests/vt/test_moe_router_warp_map.cpp`, itself under the `SameIntermediate` +exemption of §8.1. That is a thin guard, and it is deliberately named as thin: +the reason the seed must not be removed is fidelity to the incumbent, not an +output the test can see. + +**Warp levels.** `s = 16,8,4,2,1` combine `t` and `t+s` with `t < s < 32`, all +inside warp 0, and are reproduced by `v = op(v, __shfl_down_sync(0xffffffffu, v, +s))`. Lane 0's dependency cone stays within lanes `< 16` after the first step, +so the lanes whose `__shfl_down_sync` returns their own value never enter it — +exactly as the block tree leaves `red[t >= s]` stale and reads only `red[0]`. + +**VPT > 8 is not dispatched.** At `E > 256` the seed loop at `:83` accumulates +several experts per thread in ascending order, so `e[t]` is itself an +association we would have to reproduce. Not derived, not dispatched. + +## 6. Design + +New header `src/vt/cuda/moe_router_warp.h`, host-compilable (the +`#if defined(__CUDACC__) __host__ __device__ #endif` idiom already used by +`src/vt/cuda/gdn_decode_fused.h:57`), carrying: + +- `MoeRouterWarpExpert(lane, slot) -> lane + 32 * slot` — **the lane map**. +- `MoeRouterWarpTreeSum` / `MoeRouterWarpTreeMax` — the halving trees + of §5, `MoeRouterWarpTreeMax` applying the `fmaxf(-INFINITY, ·)` seed first. +- `MoeRouterWarpValuesPerThread(e) -> VPT` for `e ∈ {32,64,128,256}`, else `0`. +- `MoeRouterWarpFlagIsOn(const char*)` — the selector, default ON, `"0"` off, + mirroring `Fa2PrefillEnabled()` (`cuda_paged_attn.cu:2504-2507`). + +Putting the map and the trees in a header is deliberate: it is the part that can +be wrong in a way no compiler catches, and it is the part a host test can +execute without a GPU. + +`MoeRouterTopKWarpKernel` in `cuda_moe.cu`: + +- `dim3(32, kRouterWarpsPerCta)` with `kRouterWarpsPerCta = 4`, one token per + warp, `grid = ceil(T / 4)` — the same geometry as + `topkGating<…,WARPS_PER_CTA=4,…>` (`ROWS_PER_WARP = 1`, `ROWS_PER_CTA = 4`, + `topk_softmax_kernels.cu:311-317`). The out-of-range exit is warp-uniform, so + a launched warp always has all 32 lanes active and `0xffffffffu` masks are + valid. +- One coalesced pass loads `p[q] = Load(lrow, MoeRouterWarpExpert(lane, q))`. + Slot `q` reads the 32 consecutive experts `[32q, 32q+32)`, so every load is + fully coalesced. This trades vLLM's single 16-byte vectorized `LDG` for `VPT` + coalesced scalar loads; the bytes moved are unchanged, and it is still **one** + read of the row against the incumbent's two. +- Max, `expf`, sum, normalize and the k rounds all run out of `p[]`. +- Winner masking is `if (best >= 0 && lane == (best & 31))` followed by an + unrolled `for (q) if (q == slot) p[q] = -INFINITY;`. The unrolled compare is + required: a runtime index into a per-thread array forces it to local memory and + would spill the whole row, which is the entire point of the kernel. +- Lane 0 writes `weights`/`indices` and accumulates `denom` in `k` order, and + applies the `denom <= 0 -> 1` guard and the final divide — the same statements + as `:189-199`. + +Dispatch in `LaunchRouter`: when `!serial`, the selector is on, and +`MoeRouterWarpValuesPerThread(e) != 0`, launch the warp kernel; otherwise fall +through to the **unchanged** block kernel. `RouterDispatch`'s grouped branch is +untouched, so `num_expert_group > 0` never reaches this code. + +The selector is read fresh per launch (a `getenv` on a host path that runs once +per MoE layer per step, ~40 per 16 ms step) so an in-process test can flip it, +matching `Fa2PrefillEnabled()`. Under CUDA-graph capture the selector is read at +capture time and the graph bakes the chosen kernel, which is how every other +lever in this file behaves. + +## 7. Risks + +| risk | why it is contained | +|---|---| +| The lane map is wrong for some `VPT` | Derived per `VPT` in §5 and executed by the host equivalence test in §8; an underived `VPT` is not dispatched | +| A register array spills to local memory, erasing the win | Every index into `p[]` is a compile-time constant under `#pragma unroll`; the masking step is an unrolled compare, not a dynamic index | +| `__shfl_*_sync` on a partly-exited warp | The only early exit is `row >= T`, which is warp-uniform | +| The lever changes the grouped router | It is a different kernel behind a different branch of `RouterDispatch`; not edited | +| The lever changes the oracle | The `Serial` path is not edited; the parity test still compares against it | +| A NaN/Inf row behaves differently | The clamp, the `sum > 0` guard and the `-INFINITY` max seed are all reproduced verbatim; covered by a dedicated test case | +| It is fast but not byte-exact, and someone widens a tolerance | The parity test is `memcmp` on f32 weights and `==` on i32 indices. Widening it is forbidden; see §10 | +| It is byte-exact but not faster | Then it is a neutral change and takes no credit. §9 | + +## 8. Tests + +**RED-first, two vehicles, because the box that implements this has no `nvcc` +and no GPU.** That limitation is stated in the evidence, not worked around. + +**(a) Host equivalence test** — `tests/vt/test_moe_router_warp_map.cpp`, runs +anywhere, registered in `tests/CMakeLists.txt` with +`${CMAKE_SOURCE_DIR}/src` on the include path (the pattern of +`test_gdn_decode_fused` at `tests/CMakeLists.txt:966-967`). It contains: + +- a literal transcription of the incumbent block kernel (`cuda_moe.cu:61-201`) + over a 256-entry array, including the pad seeds and the `t < s` update rule; +- a warp model that calls the **shipped** `MoeRouterWarpExpert` and + `MoeRouterWarpTree{Sum,Max}` from the header, with `__shfl_down_sync` emulated + by its defined semantics (lane `l` reads lane `l+off`, or itself when + `l+off >= 32`); +- for `E ∈ {32,64,128,256}`, `k ∈ {1,8}` and `k > E`, `renormalize ∈ {0,1}`, and + **both f32 and bf16-rounded logit arms**: `mx`, `sum`, the full weights vector + by `memcmp`, and the indices vector must be identical; +- adversarial rows: exact-tie storms, all-NaN, all-Inf, mixed NaN/Inf, and + `-0.0f`; +- a **permanent discrimination assertion**: the same warp model instantiated + with the contiguous map `lane * VPT + slot` (vLLM's own partition) must + produce a *different* weight bit pattern on at least one seeded random row. + A test that cannot see that difference cannot see a reduction-order defect, + and this pins the discriminating power in the gate instead of in a one-off + transcript; +- the selector truth table. + +The bf16 arm is not optional. A bf16 store has hidden a defect in this tree +before; f32 weights are compared bitwise here, but the *logits* arm still +matters because bf16 rounding manufactures the exact ties the tie-break must +resolve. + +**(b) CUDA parallel-vs-serial test** — extending the existing case at +`tests/vt/test_ops_moe_grouped.cpp:503`. That case's comment at `:556` says +`// 35B routing shape: E=128, top-8`, which is **wrong**: the gate model is +`E=256` (`num_experts=256`, `num_experts_per_tok=8`; 40 MoE layers is the +observed 40 calls/step). E=256 was covered only by a hand-built near-tie +pattern. Fix the comment and add: + +- random-logit `E=256` top-8, and an exact-tie storm at `E=256`; +- `E ∈ {32,64}` — every dispatched `VPT` gets device coverage; +- `k=1` and `k > E` (the `-1` sentinel with `weights = -INFINITY`) — **the + `k > E` half is retracted for the device case by §8.2; the op rejects it**; +- an all-NaN row and an all-Inf row. + +Land these against **current** code first; they only add coverage and must be +GREEN before the kernel exists. Then get the device RED by mutation: build the +warp kernel with the contiguous map and confirm the weights `memcmp` fails. + +**Ordering.** (a) and the extra (b) cases go in green against current code. +Then the mutation red. Then the correct map, green. + +### 8.1 Amendment from implementation: two IEEE-unspecified intermediates + +Implementing (a) surfaced two facts the spec did not anticipate. Both are +recorded here because they change what the test may assert, and a reviewer must +be able to check the reasoning rather than take the relaxation on trust. + +The **weights and indices comparison is unaffected** — it stays a strict, total +`memcmp` / `==`, and it holds at `-O0`, `-O1`, `-O2` and `-O3`. What moved is +only the two internal intermediates `mx` and `sum`, which the test pins purely +because they localize a defect faster than the output does. + +1. **The NaN payload of `sum` is not an algorithm property.** IEEE-754 leaves + the payload of `NaN + NaN` unspecified and a compiler may commute the + operands of a float add freely. Measured here: `red[0] = red[0] + red[1]` and + `n[0] = v[0] + v[1]`, over the identical operands `0x7fc00000` and + `0xffc00000` and in the identical source-level operand order, produce + `0x7fc00000` and `0xffc00000` — one loop got vectorized, the other did not. + Unobservable in the output: `sum` is consumed only through the predicate + `sum > 0.0f` (`cuda_moe.cu:99`), which is false for **every** NaN payload, so + every probability takes the `: 0.0f` arm identically. +2. **The sign of a zero `mx` is a HOST CODEGEN artifact** (reworded 2026-08-11; + the original text framed it as an unspecified IEEE tie that "the algorithm + may resolve either way", which is wrong). The algorithm has no such freedom. + Review checked exhaustively that `fmaxf(x, -INFINITY) == x` **bitwise** for + all 2^32 non-NaN `x`, so the seed reproduces each leaf exactly; past the seed + both models apply `fmaxf` to *identical pairs in identical order*. Whatever a + conforming `fmaxf` does with the `(+0.0f, -0.0f)` tie — and C does leave that + unspecified — it does the same thing on both sides, so the two transcriptions + cannot disagree by any property of the algorithm. + + What actually disagrees is the **host compiler**: gcc vectorises one + transcription and not the other, and the vector max instruction does not + implement `fmaxf`'s zero handling. Hence the `-O`-dependence — the same + source passes at `-O1` and fails at `-O0` and `-O2`. It is a fact about this + test's two host loops, not about the kernels, and it says nothing about what + the device will compute. Unobservable in the output either way: `mx` is + consumed only as `expf(logit - mx)` (`cuda_moe.cu:84`); `x - (+0.0f)` and + `x - (-0.0f)` differ for exactly one input, `x == -0.0f`, and `expf` maps + both zeros to the same `1.0f`, so `sp[]` is bit-identical. + +Neither is a tolerance and neither may be widened further. Both proofs are +asserted executably by the test case *"the intermediate exemptions cannot reach +the output"*; if either stops holding, the exemption is invalid and the +intermediate comparisons go back to a strict `memcmp` — they do **not** get +loosened again to accommodate a failure. + +This also sharpens what the device gate in (b) must and must not compare: the +existing parallel-vs-serial case compares only weights and indices, which is +exactly right, and it must stay that way. + +### 8.2 Amendment: `k > E` is not an input this op has + +§8(b) above asked the device case to add "`k=1` and `k > E` (the `-1` sentinel +with `weights = -INFINITY`)". The `k > E` half of that was wrong, and it made +the device case RED: `vt::MoeRouterTopK` validates + +```c +VT_CHECK(args.top_k >= 1 && args.top_k <= e, + "moe_router_topk: top_k must be in [1, num_experts]"); +``` + +(`src/vt/ops.cpp`) **before** dispatch, so the arm threw and aborted the case +partway — with every assertion that had already run still reporting *passed*, +exactly the `grep 'assertions:'` trap §9 gate 2 warns about. The reachable domain of +this op is `1 <= top_k <= num_experts`, on rank-2 contiguous same-device +tensors with `weights`/`indices` shaped `[T, top_k]`. + +**The sentinel is dead through the public op, and that is a fact about the op, +not a gap in the test.** Within `1 <= k <= E` it cannot fire: after +`sum > 0 ? sp[j]/sum : 0` and the `isfinite` clamp every prob is finite and +`>= 0` (`cuda_moe.cu`), the argmax seeds `best_v = -INFINITY`, so each round +finds *some* unmasked expert; only masking all `E` of them leaves `best == -1`, +which needs a `(k+1)`-th round. `best < 0` is therefore a defensive guard. + +Driving it from the device case was also unsafe. That case calls the oracle +`vt::cuda::MoeRouterTopKSerialCuda` **directly**, bypassing the validator, and +the `Serial` branch's `sp[best] = -INFINITY` carries no `best >= 0` guard — a +`k > E` round writes shared memory at `sp[-1]`, which gate 5's +`compute-sanitizer memcheck` would rightly flag. The oracle is not editable +under §1, so the fix is to stop feeding it an input the op forbids. + +Resolution: the device case drops the `k > E` arm and instead pins the contract +with one `CHECK_THROWS_AS`, so the reason the sentinel is unreachable is +asserted rather than assumed. The **host** case keeps its `k > E` arm: it models +the two kernels directly, it is the only vehicle that can compare their +`best >= 0` guards, and its block transcription carries the guard so it never +indexes `-1`. `k=1`, `E ∈ {32,64,128,256}` and the all-NaN / all-Inf / mixed +rows are all inside the reachable domain and stay on device unchanged — the +validator constrains shapes, dtypes and `top_k`, never logit *values*. + +### 8.3 Amendment from review: three gate-integrity repairs (2026-08-11) + +An adversarial review of `be6a1f57` found no correctness defect in the kernel or +the lane map, and three defects in what the gate could *see*. All three are +fixed; none touches kernel arithmetic. + +1. **The device case did not pin `VT_MOE_ROUTER_WARP`, so it could not say which + kernel it had tested.** `MoeRouterWarpEnabled()` (`cuda_moe.cu`) is a fresh + `getenv` per launch, and the parallel-vs-serial sweep neither set, cleared nor + asserted the variable. With `0` in the ambient environment — which §9 gates + 6/7 *instruct* the operator to export, in the shell they then run `ctest` in — + every `run()` exercised the block kernel, i.e. block-vs-serial, green since + `6a8c5cf9`, reporting the identical `14 cases / 1137 assertions / SUCCESS`. + Green could not distinguish "the warp kernel is byte-exact" from "the warp + kernel never ran". The sweep is now hoisted into a lambda and run **twice** + under `vt_test::ScopedMoeRouterWarp` (`tests/vt/moe_router_warp_env.h`), + pinned `"1"` and `"0"`, each arm asserting the pinned state before it runs, so + the candidate **and** the rollback are both gated and the ambient environment + is irrelevant. **This doubles the case's assertion count by design** — gate 2 + reads a changed count as a red flag, and this is the change that explains it. + The pin is a shared header rather than a per-file copy precisely so the + portable companion, which runs with no GPU, executes the *same* pin the device + gate uses. + +2. **Both host models shared one argmax comparator**, so a comparator divergence + between `cuda_moe.cu:169`/`:186` (block) and `:302` (warp) was invisible by + construction: the models would have agreed because they ran the same code. + Measured: dropping the `cur_i < 0 ||` sentinel clause from the shared helper + left the suite at `7 cases / 3806 assertions / SUCCESS`. Each model now + carries its own transcription (`BlockArgmaxTakes` / `WarpArgmaxTakes`), and a + dedicated case pins them equal over every `(value, index)` pair that matters + (`±INFINITY`, signed zeros, NaN, and the `-1` sentinel); the same mutation + applied to the warp side alone is now RED. + +3. **The discrimination case asserted `sum`-or-`weights`, not `weights`.** Its + stated purpose is to prove the *weights* `memcmp` has discriminating power — + the comparison the device gate makes — but the `||` was satisfiable by a + contiguous map that moved the denominator without ever reaching the output. + It now asserts on `weights` alone. The property holds: of 16 seeded rows, + `2 / 2 / 6` differ in the weight bits at `VPT = 2 / 4 / 8`. + +## 9. Gates + +Correctness first. Nothing below authorizes a performance claim on its own. + +1. `test_moe_router_warp_map` green, with the mutation red captured on the same + binary. +2. `test_ops_moe_grouped` green on a real GPU, with its assertion count recorded + — a *changed* count is a red flag, not a pass, and `grep 'assertions:'` alone + has hidden thrown cases here before, so `Status` and the test-case counts are + recorded too. +3. `test_ops_moe`, `test_ops_moe_router_grouped`, `test_ops_moe_grouped_bf16` + unchanged. +4. Clean CUDA `-Werror` rebuild — not incremental; a header changes here. +5. `compute-sanitizer memcheck` clean on the router fixture. +6. The 35B-A3B token gate: greedy output byte-identical with the lever ON and + with `VT_MOE_ROUTER_WARP=0`, **same binary**. Byte-exactness is the claim, so + any divergence rejects the candidate outright rather than opening a + tolerance discussion. +7. Only then: same-binary A/B on an idle GB10 with the lever the only + difference, plus a same-tool `nsys` slice showing the router kernel's + us/call, and an `ncu` attribution of where the remaining gap sits. Report the + value and the ratio; a kernel-level win with no step-level movement is + diagnostic evidence, not a default-flip credit. + +The vLLM denominator must be **re-measured against the rebuilt pin** before any +number from #378 is recorded as established. + +## 10. Stop conditions + +- If the map cannot be shown byte-exact for some `E`, **narrow the dispatch** to + the `VPT` that can and record which and why. Never re-ratify a golden to + accommodate a reduction-order change, and never widen the `memcmp` — that is + the developer's decision, and it is a `NEEDS_DECISION`, not an implementation + choice. +- If a test outside the ones this spec adds goes red, stop and report. +- If the kernel is byte-exact but measures neutral or slower, record it as a + falsified hypothesis in `## Outcome` with the regime it was measured in, keep + the block kernel as the default, and do **not** declare a ceiling: the next + named hypotheses are the gate-GEMM epilogue fold and batching the 40 launches + (§1), and #378 sizes the whole prize at ~1.5% of GPU-busy, which does not + reach parity alone. +- Anything requiring a public API change, an arithmetic change from §4, a pin + advance, or a keyed-record edit beyond this row returns `NEEDS_DECISION`. + +## 11. Evidence obligations + +Record: the immutable SHA; the exact commands; the build type (a Release/NDEBUG +green over an assert-firing bug is a latent failure); the host, driver and +contention state; the assertion and test-case counts for every suite run; the +mutation transcript; and, separately, everything that could **not** be verified +on the implementing host. An implementer's report is an input; the operator +reruns the gate. + +**Which golden gate is evidence for this row, and which is not.** +`test_qwen36_paged_engine` (35B-A3B, `num_experts=256`, `num_experts_per_tok=8`) +is **the** relevant model-level gate: E=256 is a dispatched width, the router is +ungrouped, so the warp kernel is on the path a greedy golden actually walks. +`test_qwen27_paged_engine` is **NEUTRAL** — 27B is dense and has no ungrouped +MoE router at all, so it exercises none of this code and cannot be cited as +evidence for this row in either direction. A green there is not a pass for +`KERNEL-MOE-ROUTER-WARP`, and a red there is somebody else's bug. + +**Counts to expect after the §8.3 repairs.** `test_moe_router_warp_map` moves +from `7 cases / 3806 assertions` to `9 cases / 4597 assertions` (4598 when +`VT_MOE_ROUTER_WARP` is set in the ambient environment — the restore check +asserts one extra time when there is a prior value to restore). +`test_ops_moe_grouped`'s router case runs its sweep twice, so its assertion count +roughly doubles; the previously recorded `14 cases / 1137 assertions` is +superseded and a device re-run owes the new numbers. + +## 12. First DEVICE verification (2026-08-12, GB10 `dgx.casa`) + +Everything below is the first time this kernel has been **compiled or executed at +all**. `cce81c7e` shipped it unbuilt and unrun ("no nvcc, no GPU"), and the branch +never landed, so until now the row's only evidence was host arithmetic. + +**Build.** Source `6c3be5c3` (this branch on main `bbc482a2`), clean tree. +`RelWithDebInfo`, pinned `/usr/local/cuda-13.0/bin/nvcc`, +`VLLM_CPP_CUDA_ARCHITECTURES=121a`, `VLLM_CPP_CUTLASS_DIR=$HOME/cutlass-4.5.0`, +`VLLM_CPP_TRITON=ON`, `VLLM_CPP_BENCH_PROFILE_CONTROL=OFF`. Configure log verified +to print `FlashAttention-2 prefill/decode: ENABLED for arch(es) [121a]` and the +`sm_121a` vendored Triton AOT lines. **The `.cu` compiles**; that was not known. + +**Oracle identity, asserted by commit.** `~/work/vllm-src-5559679` HEAD = +`5559679229bc961848b121ccdeaa8fa5d79bec98` = the pin. The venv carrying it is +`~/venvs/vllm-oracle-next` (`0.23.1rc1.dev1511+g555967922` — a setuptools_scm +nearest-tag artifact, not a mismatch). **The shared `~/venvs/vllm-oracle` symlink +currently points at `vllm-oracle-v0.25.0-stage`, i.e. NOT the pin.** Left +untouched, since other sessions depend on it; recorded because any denominator +taken through that symlink is the 0.25.0 rollback — exactly the caveat #378 carries. + +**Gates 1-3, 6 — all GREEN, one `flock $HOME/gpu.lock` held for the group.** + +| gate | suite | cases | assertions | Status | +|---|---|---|---|---| +| 1 | `test_moe_router_warp_map` | 9 | 4597 | SUCCESS | +| 2 | `test_ops_moe_grouped` (self-pins both arms) | 14 | 1907 | SUCCESS | +| 2 | same, ambient `VT_MOE_ROUTER_WARP=0` | 14 | 1907 | SUCCESS | +| 3 | `test_ops_moe` | 9 | 33451 | SUCCESS | +| 3 | `test_ops_moe_router_grouped` | 14 | 3880 | SUCCESS | +| 3 | `test_ops_moe_grouped_bf16` | 7 | 19 | SUCCESS | +| 6 | `test_qwen36_paged_engine`, warp ON | 2 | 315 | SUCCESS | +| 6 | `test_qwen36_paged_engine`, `=0` rollback | 2 | 315 | SUCCESS | + +Gate 1's predicted `9 / 4597` and gate 6's `315 / 315` land exactly on the counts +§11 predicted. Gate 2's 1907 is the post-§8.3 doubled count replacing the +superseded `14 / 1137`. Gate 1's mutation RED was re-proven on the same source: +replacing the halving tree with a left-linear fold fails **76 assertions across 2 +cases** at `-O2` (§8.3 recorded 32 against the smaller pre-repair suite). The 35B +gate is confirmed to have actually RUN, not skipped — it loaded +`nvidia/Qwen3.6-35B-A3B-NVFP4@491c2f1e` and emitted its real continuation. Gate 4 +(clean `-Werror` CUDA build) is satisfied by construction: a from-scratch tree. +**Gate 5 (`compute-sanitizer memcheck`) was NOT run and is owed.** + +**Which kernel actually ran — proven by name, not inferred.** `nsys +--cuda-graph-trace=node`, same binary, both arms: + +- warp arm: `MoeRouterTopKWarpKernel<__nv_bfloat16, (int)8>` — VPT=8, the E=256 + dispatch the gate model uses — 1280 calls, and **no** block kernel present. +- `=0` arm: `MoeRouterTopKKernel<__nv_bfloat16, (bool)0>`, 1280 calls, and **no** + warp kernel present. + +That closes the F1 hazard at the model level: the 315/315 is not a green that +could have come from either kernel. + +**Kernel-level A/B — ESTABLISHED.** 3 reps/arm, order-alternated `on off off on +on off`, one lock, 1280 router calls in every leg: + +| arm | mean | min | max | spread | +|---|---|---|---|---| +| warp | **10.422 ms** | 10.307 | 10.529 | 0.222 | +| block | **14.201 ms** | 14.141 | 14.275 | 0.134 | + +**-3.779 ms, -26.6%, ratio 0.7339 (1.363x on the kernel)**, bands +non-overlapping. All six legs also passed the token gate, so correctness holds +across six further independent runs. + +**Step-level movement — NOT SEPARABLE, and the window is why.** Total GPU-busy +read warp 1601.890 ms vs block 1619.395 ms (-17.505 ms, ratio 0.9892). **That +number is not attributable and is not claimed.** The kernels this change does not +touch — identical code, identical 155,310 launches every leg — differ by 13.726 ms +between arms, and their full spread across the six legs is 24.140 ms, i.e. **6.4x +the 3.779 ms effect**. An earlier single pair showed it more crudely still: an +unchanged `gemvx::kernel` moved 15.7 ms at zero call-count change. This +instrument's systematic error swamps the effect, so the apparent 1.011x is an +artifact of which legs landed where. + +The deeper reason is the window: this is the whole paged-engine test, prefill and +decode aggregated, so the router's 40 calls/step are diluted across ~1600 ms. +Sized against the decode step instead — 40 calls/step, and #378's measured +16.1863 ms/step decode GPU-busy — the saving is **0.083-0.118 ms/step, i.e. +~0.5-0.75% of decode GPU-busy** (medians and means respectively). That is +**smaller than #378's ~1.5% estimate**, which assumed the whole 6.13 us/call gap +to `topkGating` was recoverable; the warp kernel closes roughly 79% of it +(12.98 -> 8.14 us/call mean against vLLM's 6.85), not all of it. + +**Consequences, stated plainly.** + +1. Per gate 7 this is a kernel-level win **without** demonstrated step-level + movement, and therefore **not a default-flip credit**. The default is left as + `cce81c7e` authored it (ON) because that is the reviewed author's decision and + the standing parity-enabler policy, *not* because this measurement earned it. + Whether default-ON is justified for a byte-identical, strictly-cheaper kernel + (no shared memory, no barriers, one global read instead of two) whose + step-level effect is real but under 1% is a **`NEEDS_DECISION` for the + operator**, not an implementer's call. +2. #378 cannot close `ROAD-V1-A` alone and never could; at ~0.5-0.75% of decode + GPU-busy it does not reach a 3-8% grid gap. §10 already said so. + +**Next traceable hypotheses** (no ceiling is declared): a decode-only profiler +window to convert the 0.5-0.75% estimate into a measurement; then the two levers +§1 names and does not attempt — folding the router into the preceding gate-GEMM +epilogue, and batching the 40 per-step launches, where the remaining ~1.3 us/call +against `topkGating` and most of the launch/tail cost live. + +**Owed:** gate 5 `compute-sanitizer memcheck`; the decode-only window; and a vLLM +denominator re-measured against the pin rather than the 0.25.0 rollback. + +## Outcome + +Pending. Filled in when the row reaches `DONE`, with what was measured, what was +rejected and why, and why the default is set the way it is. §12 supplies the +correctness half and the kernel-level measurement; the default-flip credit and +the decode-only attribution remain open. diff --git a/docs/BENCHMARKS.md b/docs/BENCHMARKS.md index 7236fbcb6..241331200 100644 --- a/docs/BENCHMARKS.md +++ b/docs/BENCHMARKS.md @@ -17,7 +17,7 @@ | **Record/checker repair 2026-08-07–08** | Gates fixed. Public: `VT_GEMMA4_EXPERT_VRAM_MB` (positive-MiB LRU cap; unset/0 unlimited), `VT_SERVER_MAX_{PROMPT_CHARS,NEW_TOKENS}` (200000/4096; 0 disables); nine Gemma4/ROCm tuners internal. No runtime/perf change. | | **vLLM** | Qwen3.6-27B NVFP4 `unsloth` @`890bdef7`, GB10 | ahead 4.5% at c1, **tie** at c2 to c32 | identical | | **vLLM** | Qwen3.6-27B NVFP4 `nvidia` @`0893e160` (ModelOpt `modelopt_mixed`), GB10 | **0.937x to 0.956x, BEHIND** and FLAT over c1-c32 (canonical 6-point 2026-08-11). The 2026-08-10 c1 0.838 did NOT reproduce (#349); gap NOT attributed | near-tie | -| **vLLM** | Qwen3.6-35B-A3B NVFP4 `nvidia` @`491c2f1e`, GB10 | **CANONICAL 2026-08-11 @`348c265d`: 0.918x-0.972x** over c1-c32 (first c16/c32); best c4 0.9719. Supersedes the ad-hoc grid | near-tie | +| **vLLM** | Qwen3.6-35B-A3B NVFP4 `nvidia` @`491c2f1e`, GB10 | **CANONICAL 2026-08-11 @`348c265d`: 0.918x-0.972x** over c1-c32; best c4 0.9719. STALE: +136 src commits, incl. +2.05% c8; regrid owed | near-tie | | **vLLM** | DeepSeek-V2-Lite (MLA), GB10 | 0.86x to 0.95x throughput, TTFT wins at c4/c8 | identical | | **vLLM** | Laguna-S-2.1 NVFP4 (118B/8B MoE), GB10 | **parity+, 1.03x** (44.46 vs 43.10 tok/s, byte-exact, default config; bf16 weights now device-resident) | near-tie | | **llama.cpp** | Qwen3.5-2B GGUF, CPU aarch64 | 20-core Arm/i8mm: prefill **1.18x ahead**, decode tie, memory parity. RPi5/A76: vllm.cpp is **0.461x prefill / 0.653x decode+E2E**, but uses **24.2% less RSS** | byte-identical on both Arm lanes | diff --git a/docs/ENVIRONMENT.md b/docs/ENVIRONMENT.md index cd8287f8a..9e932c217 100644 --- a/docs/ENVIRONMENT.md +++ b/docs/ENVIRONMENT.md @@ -104,6 +104,7 @@ portable/reference path. In normal operation leave them unset. | `VT_GDN_POSTCONV_TOKEN_TILE` | `0` (CUDA GDN prefill experiment) | A non-`0` value selects the byte-identical 16-token, per-head fused post-conv schedule. It remains opt-in pending the combined K4+post-conv Pareto measurement in issue #206 | | `VT_MODELOPT_W4A4` | `0` (Qwen3.6 dense ModelOpt NVFP4) | ModelOpt NVFP4 checkpoints ship a per-tensor `input_scale` next to every projection. Consuming it sets `Nvfp4Weight::alpha`, which flips `IsTrueW4A4()` and routes the weight to the fp4-ACTIVATION GEMM; on `nvidia/Qwen3.6-27B-NVFP4` that produced incoherent text, so the default leaves `alpha` at 0 and takes the W4A16 weight-only dispatcher (verified coherent). Set `1` to consume `input_scale` and take the W4A4 path | | `VT_LMHEAD_FP4` | **on** (Qwen3.6 dense NVFP4 `lm_head`) | Keeps a ModelOpt/compressed-tensors NVFP4 output head PACKED (`Qwen3_5DenseWeights::lm_head_fp4`) so the logits GEMM reads `K*N/2 + K*N/16` bytes per step instead of the `2*K*N` of a dequantized bf16 operand (~0.715 GB vs ~2.543 GB on `nvidia/Qwen3.6-27B-NVFP4`), and the operand keeps its on-disk `[N,K]` orientation instead of forcing the row-major NN GEMM that has no `nvjet_sm121` kernel. Mirrors vLLM, which resolves a quantized `lm_head` through `ModelOptNvFp4W4A16LinearMethod` (`modelopt.py:2491-2496,2508-2536`) and never materializes bf16 (`logits_processor.py:98-133`). Measured peak host RSS 21.06 -> 19.36 GiB (**-1.70 GiB**: the bf16 head is 2,543,206,400 B = 2.368 GiB, the packed head 715,264,000 B = 0.666 GiB) — measured before #150 changed the RSS accounting, so the figure is owed a re-measurement. On a backend with NO fp4 GEMM the head additionally keeps ONE dequantized `2*K*N` bf16 operand, built at prepare time; no other NVFP4 projection keeps one, so a quantized tower is never expanded (issue #203). The head is W4A16 under BOTH spellings: the on-disk activation divisor (`input_scale` / `input_global_scale`) is dropped for the head unless `VT_MODELOPT_W4A4=1`, because vLLM's `ModelOptNvFp4W4A16LinearMethod` deletes it (`modelopt.py:1365`). `=0` is the same-binary rollback to dequantize-at-load. BF16, FP8, GGUF and tied heads are unaffected either way (row `PERF-27B-LMHEAD-FP4`, issue #213) | +| `VT_MOE_ROUTER_WARP` | **on** (CUDA, ungrouped MoE router, `E` in {32,64,128,256}) | Runs the router top-k as ONE WARP per token with the whole logit row in registers — no shared memory, no `__syncthreads()`, and one global read of the row instead of two — where the incumbent 256-thread block kernel spends 37 barriers and 3 KiB of dynamic shared memory per token at `E=256`. BYTE-IDENTICAL, not approximately: the block reduction's `s = 128, 64, 32` levels are all multiples of the warp width, so they never cross a lane and are exactly a per-lane halving tree over `expert = lane + 32*slot`, which makes the float additions, their operands and their association identical (derived per width in `.agents/specs/moe-router-topk-single-warp.md` §5, executed with no GPU by `tests/vt/test_moe_router_warp_map.cpp`). This is where the earlier "vLLM's `topkGating` reorders the softmax so it is off-limits" conclusion was too strong — that holds for vLLM's own CONTIGUOUS lane map, not for this one. `=0` is the same-binary rollback to the block kernel. Grouped (`noaux_tc`) routing, the `Serial` byte-exact oracle, and every other `E` are untouched either way (issue #378) | | `VT_FA2_PREFILL` | on (CUDA) | The portable prefill attention instead of the vendored FA2 | | `VT_FA2_DECODE` | on (CUDA) | The portable decode attention instead of the vendored FA2 | | `VT_FA2_DECODE_4B` | on (CUDA, Qwen3.5-4B) | The portable paged decode attention instead of the ratio-4 vendored FA2 path; the 27B and 35B selectors are unchanged | diff --git a/docs/STATUS.md b/docs/STATUS.md index 5d8404fdc..fd6a139fd 100644 --- a/docs/STATUS.md +++ b/docs/STATUS.md @@ -89,7 +89,7 @@ token-for-token correctness against the pinned oracle. | Capability | State | Notes | |---|---|---| | Qwen3.6-27B (NVFP4) text generation | Correctness-complete; speed is CHECKPOINT-dependent | Token-exact GB10 on both. `unsloth` @`890bdef7` beats vLLM every c (1.007-1.045x), 115/124; `nvidia` @`0893e160` **flat 0.937-0.956 c1-c32** (#349; 0.838 void) | -| Qwen3.6-35B-A3B (NVFP4, GDN MoE) | Correctness-complete; **canonical 0.918-0.972x c1-c32** @`348c265d` (first c16/c32) | Token-exact SYNC+ASYNC; `VT_ASYNC_DEVICE_MIRROR` ON fixes async batch-1 token-0 degeneration | +| Qwen3.6-35B-A3B (NVFP4, GDN MoE) | Correctness-complete; **canonical 0.918-0.972x c1-c32** @`348c265d`, now STALE (+136 src commits, incl. a +2.05% c8 lever); regrid owed | Token-exact SYNC+ASYNC; `VT_ASYNC_DEVICE_MIRROR` ON fixes async batch-1 token-0 degeneration. Router warp kernel (#378) device-gated 315/315 both arms, kernel 1.363x, step-level not separable | | Qwen3 / Qwen2 dense (BF16) | Correctness-complete, speed-pending. Async-serving P0 FIXED (#323: the decode graph replayed stale HOST ids, now declines while the mirror is live; async 7/7 incl. Llama/Mistral/InternLM2) | Near-tie-robust token-exact vs vLLM (Qwen3-0.6B, Qwen3-4B); c1 effective parity, c8 decode residual. **Async device-mirror (`ROW-SERVE-ASYNC-DENSE-MIRROR`, `f9c969ae`): the #31 fix ported to the classic dense family, dgx-VERIFIED.** The shared dense `EmbedInto` (qwen3.cpp) raced the async combine's device input-ids write against a stale host upload → token-0 degeneration on the depth-2 AsyncLLM serving path (quant-independent). `EmbedInto` now consumes the device override published by `ForwardQwen3ForCausalLM`'s `DeviceTokenIdsScope` (27B-dense template); gate `test_qwen3_dense_async_serving` RED on `VT_ASYNC_DEVICE_MIRROR=0`, GREEN default, byte-identical mirror-off. dgx GB10: async gate RED→GREEN 0.6B+4B, SACRED 0.6B+4B 184/184 unchanged (byte-neutral sync path), memcheck 0 errors; Yi30/Qwen3-8B-MXFP4 default-config e2e coherent + 3/4 token-exact (p2 = oracle-ratified near-tie, gap 0.0000), closing the QUANT-CT-MXFP4 async-default residual. RESIDUAL: sibling InternLM2/Mistral/Llama scope one-liner; W4 bench RAN; FA2 GQA-swap default-ON, c2-c8 <1.0x. `FLASH-PTXAS` #82: codegen at PARITY (no ptxas lever); gap=engine context. **D1 (2026-07-31, `CLAIM-D1-BF16-MERGED-QKV`): the bf16 merged-QKV path (`Qwen3QkvMergeEnabled`/`VT_QWEN3_QKV_MERGE`) is now default-ON** — one `vt::MatmulBT` over the merged `[qdim+2kdim,H]` owner + a contiguous `vt::QkvSplit` (OLMo-2 exemplar), replacing three per-shard GEMMs. Bit-exact GEMM math (A/B unit `test_ops_qkv_merge` byte-identical, RED-first); the wider-N cuBLASLt K-reduction flips the 0.6B genuine bf16 near-tie so the SACRED 0.6B golden was regenerated (all tokens within the near-tie band, max 0.125 nats), while Qwen3-4B is byte-neutral (0 diffs, stays STRICT). Re-gated 0.6B 16/16 + 4B 16/16; consistency/launch-count fold (measured NEUTRAL on 4B decode), no new throughput owed | | Qwen3.5-4B BF16 direct-load on discrete CUDA | Correct; throughput/host PSS ahead, acceptance `PENDING`; latency/VRAM open | Atomic pretoken exact. Ratios: tput 1.0283x; TTFT/TPOT/E2E 1.0853/1.0165/1.0288x slower; VRAM +118.7 MiB. GDN local stack retained ([data](bench-evidence/qwen35-4b-sm120-main-20260807.md)) | | Qwen3-Coder-30B-A3B MoE (BF16) | Correctness-complete, speed-pending | Near-tie-robust token-exact 6/6; 11 of 16 binding grid cells at or above vLLM. **D1 (2026-07-31): inherits the default-ON bf16 merged-QKV via the shared dense `AttnBlock` — byte-neutral (0 token diffs, golden UNCHANGED); re-gated 6/6** | diff --git a/src/vt/cuda/cuda_moe.cu b/src/vt/cuda/cuda_moe.cu index 1e6537418..94dce54a5 100644 --- a/src/vt/cuda/cuda_moe.cu +++ b/src/vt/cuda/cuda_moe.cu @@ -9,9 +9,11 @@ #include #include +#include #include #include +#include "vt/cuda/moe_router_warp.h" #include "vt/ops.h" namespace vt::cuda { @@ -200,6 +202,133 @@ __global__ void MoeRouterTopKKernel(float* weights, int32_t* indices, const Tin* } } +// ─── Single-warp router top-k (KERNEL-MOE-ROUTER-WARP, issue #378) ───────── +// One WARP per token with the whole logit row in registers: no shared memory, +// no __syncthreads(), and ONE global read of the row instead of two. At E=256 +// the block kernel above spends 37 barriers and 3 KiB of dynamic shared memory +// on a few hundred comparisons; this spends none. +// +// BYTE-IDENTICAL to MoeRouterTopKKernel, and the reason is +// structural, not an appeal to associativity. In that kernel's block reduction +// the levels s = 128, 64, 32 are all multiples of the warp width, so `t` and +// `t + s` always share a lane id: those three levels never cross a lane, and +// they combine exactly red[L + 32q] by the standard halving recursion on q. +// Levels s = 16..1 live inside warp 0 and are what __shfl_down_sync +// reproduces. So `MoeRouterWarpExpert(lane, q) == lane + 32q` plus +// `MoeRouterWarpTree{Sum,Max}`'s halving tree performs the IDENTICAL float +// operations on the IDENTICAL operands in the IDENTICAL association. The full +// derivation, one VPT at a time, is in +// .agents/specs/moe-router-topk-single-warp.md §5; the reduction-order claim is +// EXECUTED (with no GPU needed) by tests/vt/test_moe_router_warp_map.cpp. +// +// This is where 6a8c5cf9's "vLLM's topkGating reorders the softmax reduction so +// it is off-limits" (see :156-163) is too strong. It is true of vLLM's OWN lane +// map -- at topkGating<8,256,4,16,32,...> lane L owns the CONTIGUOUS experts +// [8L, 8L+8) (topk_softmax_kernels.cu:344-346), which genuinely reassociates -- +// and false of this one. +// +// SHAPE port only. Every arithmetic decision below is the block kernel's, NOT +// vLLM's: the divide (not a reciprocal multiply), the sum>0 guard, the isfinite +// clamp after normalize, the -INFINITY mask, the -INFINITY max seed that erases +// NaN, the denom<=0 -> 1 guard, the trailing divide, and the best<0 sentinel. +// vLLM does five of those differently and porting any of them changes tokens; +// the spec §4 tabulates each against its upstream file:line. +template +__launch_bounds__(kMoeRouterWarpsPerCta* kMoeRouterWarpWidth) __global__ + void MoeRouterTopKWarpKernel(float* weights, int32_t* indices, const Tin* logits, int64_t t, + int k, bool renormalize) { + constexpr int kE = kMoeRouterWarpWidth * VPT; + const int lane = static_cast(threadIdx.x); + // One token per warp; the exit is warp-UNIFORM, so a warp that survives it + // has all 32 lanes active and the 0xffffffffu shuffle masks are valid. + const int64_t row = static_cast(blockIdx.x) * kMoeRouterWarpsPerCta + + static_cast(threadIdx.y); + if (row >= t) return; + const Tin* lrow = logits + row * kE; + + // The ONLY read of the logit row. Slot q covers the 32 consecutive experts + // [32q, 32q+32), so every load is fully coalesced. + float p[VPT]; +#pragma unroll + for (int q = 0; q < VPT; ++q) p[q] = Load(lrow, MoeRouterWarpExpert(lane, q)); + + // Max: per-lane halving tree (with the -INFINITY seed that erases NaN), + // then the five in-warp levels. Congruent to :70-79. + float m = MoeRouterWarpTreeMax(p); +#pragma unroll + for (int off = kMoeRouterWarpWidth / 2; off > 0; off >>= 1) { + m = fmaxf(m, __shfl_down_sync(0xffffffffu, m, off)); + } + const float mx = __shfl_sync(0xffffffffu, m, 0); // the block read red[0] here + + // exp(logit - max) stays in registers -- no second global read, no sp[]. + // Congruent to :82-94. +#pragma unroll + for (int q = 0; q < VPT; ++q) p[q] = expf(p[q] - mx); + float s = MoeRouterWarpTreeSum(p); +#pragma unroll + for (int off = kMoeRouterWarpWidth / 2; off > 0; off >>= 1) { + s += __shfl_down_sync(0xffffffffu, s, off); + } + const float sum = __shfl_sync(0xffffffffu, s, 0); + + // Normalize: the DIVIDE, the sum>0 guard, the clamp after (:97-103). +#pragma unroll + for (int q = 0; q < VPT; ++q) { + float pj = sum > 0.0f ? p[q] / sum : 0.0f; + if (!isfinite(pj)) pj = 0.0f; + p[q] = pj; + } + + // Greedy top-k. The argmax is a reduction over the total order "higher value, + // then lower expert index", which IS associative and commutative, so this + // grouping matches the block's per-thread/warp/leader grouping exactly -- + // the argument already recorded at :156-163 and unchanged here. + float denom = 0.0f; // meaningful on lane 0 only, accumulated in k order + for (int j = 0; j < k; ++j) { + float lv = -INFINITY; + int li = -1; +#pragma unroll + for (int q = 0; q < VPT; ++q) { // ascending expert index within the lane + if (p[q] > lv) { // strict `>` -> lowest index at the max + lv = p[q]; + li = MoeRouterWarpExpert(lane, q); + } + } +#pragma unroll + for (int off = kMoeRouterWarpWidth / 2; off > 0; off >>= 1) { + const float ov = __shfl_down_sync(0xffffffffu, lv, off); + const int oi = __shfl_down_sync(0xffffffffu, li, off); + if (ov > lv || (ov == lv && oi >= 0 && (li < 0 || oi < li))) { + lv = ov; + li = oi; + } + } + const float best_v = __shfl_sync(0xffffffffu, lv, 0); + const int best = __shfl_sync(0xffffffffu, li, 0); + // Exclude the winner from later rounds, in the owning lane's own register. + // The slot compare is UNROLLED on purpose: a runtime index into a per-thread + // array forces it to local memory, which would spill the whole row and cost + // exactly the register residency this kernel exists for. + if (best >= 0 && lane == (best & (kMoeRouterWarpWidth - 1))) { + const int slot = best / kMoeRouterWarpWidth; +#pragma unroll + for (int q = 0; q < VPT; ++q) { + if (q == slot) p[q] = -INFINITY; + } + } + if (lane == 0) { + weights[row * k + j] = best_v; // -INFINITY when best < 0 + indices[row * k + j] = static_cast(best); // -1 when best < 0 + denom += best_v; + } + } + if (lane == 0 && renormalize) { + if (!(denom > 0.0f)) denom = 1.0f; // denom<=0 -> 1 guard, as :197 + for (int j = 0; j < k; ++j) weights[row * k + j] /= denom; + } +} + // ─── Grouped-topk (`noaux_tc`) router (W3) ───────────────────────────────── // Mirrors the CPU reference (cpu_ops.cpp MoeRouterGroupedTopKKernel), itself a // 1:1 port of grouped_topk_router.py:106-161 @ e24d1b24. This is a SEPARATE @@ -417,9 +546,60 @@ void LaunchGroupedRouter(cudaStream_t s, Tensor& weights, Tensor& indices, Check(cudaGetLastError(), "moe_router_grouped_topk launch"); } +// VT_MOE_ROUTER_WARP (default ON, "0" restores the block kernel for a +// same-binary A/B). Read fresh per launch -- a getenv on a host path that runs +// once per MoE layer per step -- so an in-process test can flip it, matching +// Fa2PrefillEnabled() (cuda_paged_attn.cu:2504-2507). Under CUDA-graph capture +// it is read at capture time and the graph bakes the chosen kernel, which is +// how every other lever in this backend behaves. +bool MoeRouterWarpEnabled() { + return MoeRouterWarpFlagIsOn(std::getenv("VT_MOE_ROUTER_WARP")); +} + +// Returns false when this (E) is not one of the widths whose byte-exactness is +// derived, so the caller falls through to the UNCHANGED block kernel. +template +bool LaunchRouterWarp(cudaStream_t s, Tensor& weights, Tensor& indices, const Tensor& logits, + int64_t t, int64_t e, int k, bool renorm) { + const int vpt = MoeRouterWarpValuesPerThread(e); + if (vpt == 0) return false; // decide BEFORE touching the tensors + const dim3 block(kMoeRouterWarpWidth, kMoeRouterWarpsPerCta); + const unsigned grid = + static_cast((t + kMoeRouterWarpsPerCta - 1) / kMoeRouterWarpsPerCta); + float* w = weights.Ptr(); + int32_t* idx = indices.Ptr(); + const Tin* lg = logits.Ptr(); + switch (vpt) { + case 1: + MoeRouterTopKWarpKernel<<>>(w, idx, lg, t, k, renorm); + break; + case 2: + MoeRouterTopKWarpKernel<<>>(w, idx, lg, t, k, renorm); + break; + case 4: + MoeRouterTopKWarpKernel<<>>(w, idx, lg, t, k, renorm); + break; + case 8: + MoeRouterTopKWarpKernel<<>>(w, idx, lg, t, k, renorm); + break; + default: + return false; + } + Check(cudaGetLastError(), "moe_router_topk_warp launch"); + return true; +} + template void LaunchRouter(cudaStream_t s, Tensor& weights, Tensor& indices, const Tensor& logits, int64_t t, int64_t e, int k, bool renorm, bool serial) { + // The single-warp kernel is byte-identical (see MoeRouterTopKWarpKernel) but + // NEVER replaces `serial`: that path is the byte-exact ORACLE the parity test + // compares against, so changing it would invalidate the oracle instead of + // testing the candidate. + if (!serial && MoeRouterWarpEnabled() && + LaunchRouterWarp(s, weights, indices, logits, t, e, k, renorm)) { + return; + } const size_t shmem = static_cast(e) * sizeof(float); if (serial) { MoeRouterTopKKernel<<(t), kBlock, shmem, s>>>( diff --git a/src/vt/cuda/moe_router_warp.h b/src/vt/cuda/moe_router_warp.h new file mode 100644 index 000000000..b26844e01 --- /dev/null +++ b/src/vt/cuda/moe_router_warp.h @@ -0,0 +1,171 @@ +// vllm.cpp original (vt runtime, inventory deviation §9.1). Portable contract +// for the single-warp MoE router top-k (spec: +// .agents/specs/moe-router-topk-single-warp.md, issue #378). The CUDA kernel +// itself stays in cuda_moe.cu; what lives here is the part that can be WRONG in +// a way no compiler catches -- the lane->expert map and the two register +// reduction trees -- so a host test can execute it without a GPU. +// +// SHAPE port of vLLM's register-resident topkGating +// (csrc/libtorch_stable/moe/topk_softmax_kernels.cu:279-592 @ 555967922): one +// warp per token, no shared memory, no barriers. NOT a math port -- every +// arithmetic decision stays the incumbent block kernel's; see the spec §4 for +// the point-by-point table of where vLLM differs and why we do not follow it. +// +// WHY THE MAP IS WHAT IT IS. The incumbent block reduction is +// +// for (int s = kBlock / 2; s > 0; s /= 2) { +// if (threadIdx.x < s) red[threadIdx.x] += red[threadIdx.x + s]; +// __syncthreads(); +// } +// +// over kBlock == 256 lanes (cuda_moe.cu:90-93). Levels s = 128, 64, 32 are all +// multiples of the 32-wide warp, so `t` and `t + s` always have the SAME lane id +// `L = t % 32`: those three levels never cross a lane. They combine exactly the +// entries red[L + 32q] for q in [0, 8), and they combine them by the standard +// halving recursion on q. Levels s = 16..1 run entirely inside warp 0 and are +// what __shfl_down_sync(0xffffffffu, v, s) reproduces. +// +// So if lane L holds slot q == expert (L + 32q) and reduces the slots with the +// SAME halving tree, the result is bit-identical to the block tree by +// structural congruence: identical operands, identical operations, identical +// association. Not by an appeal to associativity -- float addition is not +// associative, which is exactly why the previous row (6a8c5cf9) believed a +// register-resident router was off-limits. It was off-limits for vLLM's OWN +// partition (lane L owns the CONTIGUOUS experts [VPT*L, VPT*L + VPT), +// topk_softmax_kernels.cu:344-346), which does reassociate. A different map +// does not. +// +// Written in the tree's leaf order the offsets are 32 * {0,4,2,6,1,5,3,7}, the +// bit-reversal of {0..7}. The stride form implemented here is the same +// statement and is the one that can be checked by eye against the incumbent +// loop. +#pragma once + +#include +#include + +#if defined(__CUDACC__) +#define VT_MOE_ROUTER_HD __host__ __device__ +// nvcc only. The host build runs -Wall -Wextra -Werror, and -Wall enables +// -Wunknown-pragmas, so a bare `#pragma unroll` here would fail the host TU +// that the portable test compiles. +#define VT_MOE_ROUTER_UNROLL _Pragma("unroll") +#else +#define VT_MOE_ROUTER_HD +#define VT_MOE_ROUTER_UNROLL +#endif + +namespace vt::cuda { + +// The warp width the map is derived against. The whole derivation above depends +// on the incumbent block's kBlock (256) being a multiple of this, and on the +// three cross-warp levels therefore being lane-local. +inline constexpr int kMoeRouterWarpWidth = 32; + +// Warps per CTA, one token per warp. Mirrors the traced upstream instantiation +// topkGating<8, 256, /*WARPS_PER_CTA=*/4, 16, 32, ...> (ROWS_PER_WARP == 1, +// ROWS_PER_CTA == 4, topk_softmax_kernels.cu:311-317). +inline constexpr int kMoeRouterWarpsPerCta = 4; + +// THE LANE MAP. Lane `lane` holds, in register slot `slot`, the expert +// `lane + 32 * slot`. Equivalently: slot q of the warp covers the 32 +// CONSECUTIVE experts [32q, 32q + 32), so every load is fully coalesced. +// +// This is deliberately NOT vLLM's contiguous `lane * VPT + slot`. That map +// reassociates the softmax sum and is the mutation the host test discriminates +// against; see MoeRouterWarpContiguousExpert below. +VT_MOE_ROUTER_HD inline constexpr int MoeRouterWarpExpert(int lane, int slot) { + return lane + kMoeRouterWarpWidth * slot; +} + +// vLLM's own partition, present ONLY so the test can instantiate it and prove +// the oracle can tell the two apart. Never dispatched. +VT_MOE_ROUTER_HD inline constexpr int MoeRouterWarpContiguousExpert(int lane, int slot, + int values_per_thread) { + return lane * values_per_thread + slot; +} + +// Values (experts) per lane for a row of `e` experts, or 0 when the row is not +// eligible for the warp kernel. +// +// Only E in {32, 64, 128, 256} is admitted, because those are the widths whose +// bit-exactness is DERIVED in the spec (§5), one VPT at a time. E > 256 is +// excluded on purpose: above the incumbent's 256-wide block the seed loops at +// cuda_moe.cu:71,83 accumulate SEVERAL experts per thread in ascending order, +// so e[t] is itself an association that would have to be reproduced. Anything +// not derived falls through to the unchanged block kernel. +VT_MOE_ROUTER_HD inline constexpr int MoeRouterWarpValuesPerThread(int64_t e) { + return e == 32 ? 1 : (e == 64 ? 2 : (e == 128 ? 4 : (e == 256 ? 8 : 0))); +} + +// Per-lane sum over the slots, in the incumbent block tree's association. +// +// Congruent to the s = 128, 64, 32 levels of cuda_moe.cu:90-93 restricted to +// one lane. For VPT < 8 the levels the block spent on its zero-padded seeds +// (threads t >= e never enter the loop body at :83, so red[t] keeps its +0.0f) +// are dropped, which is bit-exact because x + (+0.0f) == x for every value +// these seeds can hold: expf yields +0.0f or a positive value and never -0.0f, +// the one float whose sign bit that addition would flip. +// +// The OCCUPIED seeds rest on the same argument, which the spec originally +// glossed. The incumbent's per-thread seed is not the loaded value but +// `float acc = 0.0f; ... acc += ex` (cuda_moe.cu:84-89), i.e. (+0.0f) + E_t, +// while `r[q] = v[q]` below is a bare copy. Those differ only for -0.0f (which +// expf never produces) and for a SIGNALLING NaN, which the addition would quiet +// and the copy would not (expf returns only quiet NaNs). Same contained class as +// the pad leaves above; spec §5. +template +VT_MOE_ROUTER_HD inline float MoeRouterWarpTreeSum(const float (&v)[VPT]) { + float r[VPT]; + VT_MOE_ROUTER_UNROLL + for (int q = 0; q < VPT; ++q) r[q] = v[q]; + VT_MOE_ROUTER_UNROLL + for (int s = VPT / 2; s > 0; s >>= 1) { + VT_MOE_ROUTER_UNROLL + for (int q = 0; q < s; ++q) r[q] += r[q + s]; + } + return r[0]; +} + +// Per-lane max over the slots, in the incumbent block tree's association. +// +// The fmaxf(-INFINITY, .) seed is applied per element BEFORE the tree because +// that is what the incumbent does (`float m = -INFINITY; ... m = fmaxf(m, +// Load(...))` at cuda_moe.cu:70-71). Verbatim reproduction is the reason, and +// the only reason. An earlier version of this comment claimed the seed "is what +// makes an all-NaN row behave identically"; review FALSIFIED that by mutation. +// Deleting the seed changes ONLY the `mx` intermediate -- the weights and the +// indices stay byte-identical on every case, all-NaN included. fmaxf already +// returns the non-NaN operand, so the tree erases NaN with or without the seed; +// the seed bites only when a lane holds nothing BUT NaN, and then only by making +// mx -INFINITY instead of NaN, which expf(l - mx) washes out (NaN either way -> +// every prob clamped to 0.0f at :100). Do not "simplify" it away: its only guard +// in the gate is the mx intermediate check in +// tests/vt/test_moe_router_warp_map.cpp (spec §5, §8.1). +// +// Dropping the block's -INFINITY pad leaves for VPT < 8 is bit-exact because +// x = fmaxf(x, -INFINITY) for every non-NaN x, including -0.0f (checked +// exhaustively over all 2^32 non-NaN x during review), and after the seed no +// NaN remains. +template +VT_MOE_ROUTER_HD inline float MoeRouterWarpTreeMax(const float (&v)[VPT]) { + float r[VPT]; + VT_MOE_ROUTER_UNROLL + for (int q = 0; q < VPT; ++q) r[q] = fmaxf(-INFINITY, v[q]); + VT_MOE_ROUTER_UNROLL + for (int s = VPT / 2; s > 0; s >>= 1) { + VT_MOE_ROUTER_UNROLL + for (int q = 0; q < s; ++q) r[q] = fmaxf(r[q], r[q + s]); + } + return r[0]; +} + +// VT_MOE_ROUTER_WARP selector. DEFAULT ON; "0" restores the incumbent block +// kernel for a same-binary A/B. Same spelling as the sibling default-on levers +// in this backend (cuda_paged_attn.cu:2504-2507 Fa2PrefillEnabled), and read +// fresh per launch so an in-process test can flip it. +inline bool MoeRouterWarpFlagIsOn(const char* value) { + return value == nullptr || value[0] != '0'; +} + +} // namespace vt::cuda diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index a072b5ccd..6a716dce1 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -1018,6 +1018,16 @@ target_include_directories(test_gdn_packed_decode_triton PRIVATE ${CMAKE_SOURCE_ # in cuda_gdn.cu and is covered by same-process byte-exact CUDA comparisons. vllm_cpp_add_test(test_gdn_decode_fused vt/test_gdn_decode_fused.cpp) target_include_directories(test_gdn_decode_fused PRIVATE ${CMAKE_SOURCE_DIR}/src) +# CPU-tier REDUCTION-ORDER contract for the single-warp MoE router top-k +# (VT_MOE_ROUTER_WARP, default ON; spec moe-router-topk-single-warp.md, #378). +# The claim is that the lane map `expert = lane + 32*slot` plus a per-lane +# halving tree performs the bit-identical float additions as the incumbent +# 256-thread block reduction in src/vt/cuda/cuda_moe.cu -- pure IEEE-754, so it +# is provable with no GPU and no nvcc. Reaches the CUDA-free map/tree/selector +# header under src/; the executing kernel is covered on-device by the +# parallel-vs-serial byte-exactness case in test_ops_moe_grouped. +vllm_cpp_add_test(test_moe_router_warp_map vt/test_moe_router_warp_map.cpp) +target_include_directories(test_moe_router_warp_map PRIVATE ${CMAKE_SOURCE_DIR}/src) # CPU-tier contract for the two NUMERICS-NEUTRAL FP4-quant vectorized-load # fast-path flags (VT_FP4_QUANT_FAST + VT_SILU_FP4_FAST, both default OFF / # '1'-opt-in). Reaches the internal header under src/ (not public API); the @@ -1106,6 +1116,13 @@ add_test(NAME test_ops_fp8_cutlass_plan_cache_on set_tests_properties(test_ops_fp8_cutlass_plan_cache_on PROPERTIES ENVIRONMENT "VT_FP8_PLAN_CACHE=1") vllm_cpp_add_test(test_ops_moe_grouped vt/test_ops_moe_grouped.cpp) +# Reaches the CUDA-free VT_MOE_ROUTER_WARP selector header under src/, through +# the shared test-only pin tests/vt/moe_router_warp_env.h: the parallel-vs-serial +# router case PINS that toggle both ways and ASSERTS the pinned state, so it can +# no longer degenerate to block-vs-serial when the ambient environment says 0 +# (spec moe-router-topk-single-warp.md §9, #378). Same include pattern as +# test_moe_router_warp_map above, which executes that pin without a GPU. +target_include_directories(test_ops_moe_grouped PRIVATE ${CMAKE_SOURCE_DIR}/src) vllm_cpp_add_test(test_ops_moe_grouped_bf16 vt/test_ops_moe_grouped_bf16.cpp) vllm_cpp_add_test(test_ops_moe_grouped_bf16_gate_up_silu vt/test_ops_moe_grouped_bf16_gate_up_silu.cpp) diff --git a/tests/vt/moe_router_warp_env.h b/tests/vt/moe_router_warp_env.h new file mode 100644 index 000000000..8595d43fa --- /dev/null +++ b/tests/vt/moe_router_warp_env.h @@ -0,0 +1,70 @@ +// vllm.cpp original. TEST-ONLY scoped pin of the VT_MOE_ROUTER_WARP selector +// (spec: .agents/specs/moe-router-topk-single-warp.md, issue #378). +// +// WHY THIS EXISTS. vt::cuda::MoeRouterWarpEnabled() (src/vt/cuda/cuda_moe.cu) +// does a FRESH getenv on every launch, so which router kernel +// `vt::MoeRouterTopK` runs is decided by the AMBIENT environment of whoever +// started ctest. A device parallel-vs-serial sweep that neither sets nor +// asserts that variable therefore proves nothing about the warp kernel: exported +// `VT_MOE_ROUTER_WARP=0` — which spec §9 gates 6/7 tell the operator to export +// for the A/B — silently degenerates the sweep to block-vs-serial, which has +// been green since 6a8c5cf9, and it reports the IDENTICAL case and assertion +// counts. Green then cannot distinguish "the warp kernel is byte-exact" from +// "the warp kernel never ran". +// +// So the toggle is PINNED around the sweep, both ways, and the pinned state is +// ASSERTED rather than assumed. Same scoped setenv/restore idiom as +// tests/vt/test_ops_nvfp4_fp4.cpp:1268-1287 and tests/vt/test_ops_gdn.cpp:1350-1370. +// +// It is a shared header, not a per-file copy, on purpose: the device gate that +// needs the pin cannot be compiled without nvcc and a GPU, while the portable +// companion tests/vt/test_moe_router_warp_map.cpp runs anywhere — so the host +// test executes the very same pin the device gate uses, instead of a +// character-identical copy of it that nothing checks. +#pragma once + +#include +#include + +#include "vt/cuda/moe_router_warp.h" + +namespace vt_test { + +// RAII: pin VT_MOE_ROUTER_WARP to `value` (nullptr = unset it) and restore +// whatever was there before — including its absence — on scope exit. +class ScopedMoeRouterWarp { + public: + static constexpr const char* kName = "VT_MOE_ROUTER_WARP"; + + explicit ScopedMoeRouterWarp(const char* value) { + const char* current = std::getenv(kName); + had_ = current != nullptr; + if (had_) saved_ = current; + Apply(value); + } + ~ScopedMoeRouterWarp() { Apply(had_ ? saved_.c_str() : nullptr); } + + ScopedMoeRouterWarp(const ScopedMoeRouterWarp&) = delete; + ScopedMoeRouterWarp& operator=(const ScopedMoeRouterWarp&) = delete; + + // Exactly what the launcher's per-call getenv will resolve to RIGHT NOW: + // MoeRouterWarpEnabled() is MoeRouterWarpFlagIsOn(getenv(kName)) and nothing + // else (cuda_moe.cu). true -> vt::MoeRouterTopK dispatches the WARP kernel + // for E in {32,64,128,256}; false -> it falls through to the block kernel. + static bool EffectiveFlag() { + return vt::cuda::MoeRouterWarpFlagIsOn(std::getenv(kName)); + } + + private: + static void Apply(const char* value) { + if (value != nullptr) + setenv(kName, value, 1); + else + unsetenv(kName); + } + + std::string saved_; + bool had_ = false; +}; + +} // namespace vt_test diff --git a/tests/vt/test_moe_router_warp_map.cpp b/tests/vt/test_moe_router_warp_map.cpp new file mode 100644 index 000000000..5738974af --- /dev/null +++ b/tests/vt/test_moe_router_warp_map.cpp @@ -0,0 +1,698 @@ +// vllm.cpp original. Portable contract for the single-warp MoE router top-k +// (spec: .agents/specs/moe-router-topk-single-warp.md, issue #378). +// +// WHAT THIS TEST IS. The claim the row rests on is a REDUCTION-ORDER claim: +// that a 32-lane, register-resident router with the lane map +// `expert = lane + 32*slot` performs the bit-identical sequence of float +// additions and fmaxf calls as the incumbent 256-thread block kernel +// (src/vt/cuda/cuda_moe.cu:61-201). That claim is pure IEEE-754 single +// precision arithmetic, so it is checkable on a host with no GPU and no nvcc, +// which is exactly what this file does. Both models below are plain float, and +// every operation is an add or an fmaxf -- no multiply, hence no FMA +// contraction, hence the host evaluates the association as written. +// +// WHAT THIS TEST IS NOT. It does not prove the CUDA kernel compiles, launches, +// or that the device agrees with the host. The device-side proof is the +// parallel-vs-serial byte-exactness case in test_ops_moe_grouped.cpp, which +// needs a GPU. This file proves the ORDER is right; that one proves the KERNEL +// is right. Neither substitutes for the other. +// +// The warp model calls the SHIPPED map and the SHIPPED reduction trees from +// src/vt/cuda/moe_router_warp.h, so a defect introduced there fails here. The +// block model is a literal transcription of the incumbent kernel; if the +// incumbent ever changes, this test goes red, which is the point. +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "moe_router_warp_env.h" +#include "vt/cuda/moe_router_warp.h" +#include "vt/dtype.h" + +using vt::cuda::kMoeRouterWarpWidth; +using vt::cuda::MoeRouterWarpContiguousExpert; +using vt::cuda::MoeRouterWarpExpert; +using vt::cuda::MoeRouterWarpFlagIsOn; +using vt::cuda::MoeRouterWarpTreeMax; +using vt::cuda::MoeRouterWarpTreeSum; +using vt::cuda::MoeRouterWarpValuesPerThread; + +namespace { + +constexpr int kBlock = 256; // cuda_moe.cu:20 +constexpr int kWarp = 32; + +struct RouterOut { + float mx = 0.0f; + float sum = 0.0f; + std::vector weights; + std::vector indices; +}; + +// THE TWO COMPARATORS ARE TRANSCRIBED SEPARATELY, ONE PER KERNEL. +// +// Higher value wins; on an exact tie the lower expert index wins; -1 means +// "this lane holds nothing". The block kernel and the warp kernel spell that +// out independently, and an EARLIER revision of this file had both models call +// one shared helper — which makes a comparator divergence between the two +// KERNELS invisible by construction: the models would agree because they run +// the same code, not because the kernels agree. They are character-identical +// today; that is a fact to PIN, not to build the test on. +// +// Block kernel, cuda_moe.cu:169 (warp-shuffle argmax) and :186 (thread 0's scan +// over the eight warp leaders) — identical to each other there. +bool BlockArgmaxTakes(float other_v, int other_i, float cur_v, int cur_i) { + return other_v > cur_v || + (other_v == cur_v && other_i >= 0 && (cur_i < 0 || other_i < cur_i)); +} + +// Warp kernel, cuda_moe.cu:302. Transcribed separately ON PURPOSE; the test case +// "the two kernels' argmax comparators agree ..." below is what holds them equal. +bool WarpArgmaxTakes(float other_v, int other_i, float cur_v, int cur_i) { + return other_v > cur_v || + (other_v == cur_v && other_i >= 0 && (cur_i < 0 || other_i < cur_i)); +} + +// --------------------------------------------------------------------------- +// Model A: the incumbent block kernel, MoeRouterTopKKernel, one +// block of 256 threads per token. Transcribed statement for statement from +// src/vt/cuda/cuda_moe.cu:61-201. `e <= kBlock` is the regime the warp kernel +// is dispatched in, and the strided seed loops then give each thread at most +// one expert. +RouterOut BlockModel(const std::vector& logits, int e, int k, bool renormalize) { + RouterOut out; + out.weights.assign(static_cast(k), 0.0f); + out.indices.assign(static_cast(k), 0); + std::vector sp(static_cast(e), 0.0f); + std::vector red(kBlock, 0.0f); + + // cuda_moe.cu:70-79 -- max, seeded from -INFINITY (which ERASES NaN). + for (int t = 0; t < kBlock; ++t) { + float m = -INFINITY; + for (int j = t; j < e; j += kBlock) m = fmaxf(m, logits[static_cast(j)]); + red[static_cast(t)] = m; + } + for (int s = kBlock / 2; s > 0; s /= 2) { + for (int t = 0; t < s; ++t) { + red[static_cast(t)] = + fmaxf(red[static_cast(t)], red[static_cast(t + s)]); + } + } + out.mx = red[0]; + + // cuda_moe.cu:82-94 -- exp(logit - max) into shared, block-summed. This is + // the SECOND full read of the logit row. + for (int t = 0; t < kBlock; ++t) { + float acc = 0.0f; + for (int j = t; j < e; j += kBlock) { + const float ex = expf(logits[static_cast(j)] - out.mx); + sp[static_cast(j)] = ex; + acc += ex; + } + red[static_cast(t)] = acc; + } + for (int s = kBlock / 2; s > 0; s /= 2) { + for (int t = 0; t < s; ++t) red[static_cast(t)] += red[static_cast(t + s)]; + } + out.sum = red[0]; + + // cuda_moe.cu:97-103 -- a DIVIDE (not a reciprocal multiply), the sum>0 + // guard, then the isfinite clamp AFTER normalize. + for (int j = 0; j < e; ++j) { + float pj = out.sum > 0.0f ? sp[static_cast(j)] / out.sum : 0.0f; + if (!std::isfinite(pj)) pj = 0.0f; + sp[static_cast(j)] = pj; + } + + // cuda_moe.cu:137-199 -- parallel greedy top-k: per-thread local argmax, a + // warp-shuffle argmax, then thread 0's serial scan over the 8 warp leaders. + constexpr int kWarps = kBlock / kWarp; + std::vector lv(kBlock); + std::vector li(kBlock); + float leader_v[kWarps]; + int leader_i[kWarps]; + float denom = 0.0f; + for (int j = 0; j < k; ++j) { + for (int t = 0; t < kBlock; ++t) { + float v = -INFINITY; + int i = -1; + for (int idx = t; idx < e; idx += kBlock) { + if (sp[static_cast(idx)] > v) { + v = sp[static_cast(idx)]; + i = idx; + } + } + lv[static_cast(t)] = v; + li[static_cast(t)] = i; + } + for (int w = 0; w < kWarps; ++w) { + const int base = w * kWarp; + for (int off = kWarp / 2; off > 0; off >>= 1) { + float nv[kWarp]; + int ni[kWarp]; + for (int lane = 0; lane < kWarp; ++lane) { + const int src = lane + off < kWarp ? lane + off : lane; + float cv = lv[static_cast(base + lane)]; + int ci = li[static_cast(base + lane)]; + const float ov = lv[static_cast(base + src)]; + const int oi = li[static_cast(base + src)]; + if (BlockArgmaxTakes(ov, oi, cv, ci)) { + cv = ov; + ci = oi; + } + nv[lane] = cv; + ni[lane] = ci; + } + for (int lane = 0; lane < kWarp; ++lane) { + lv[static_cast(base + lane)] = nv[lane]; + li[static_cast(base + lane)] = ni[lane]; + } + } + leader_v[w] = lv[static_cast(base)]; + leader_i[w] = li[static_cast(base)]; + } + float best_v = leader_v[0]; + int best = leader_i[0]; + for (int w = 1; w < kWarps; ++w) { + if (BlockArgmaxTakes(leader_v[w], leader_i[w], best_v, best)) { + best_v = leader_v[w]; + best = leader_i[w]; + } + } + if (best >= 0) sp[static_cast(best)] = -INFINITY; + out.weights[static_cast(j)] = best_v; + out.indices[static_cast(j)] = static_cast(best); + denom += best_v; + } + if (renormalize) { + if (!(denom > 0.0f)) denom = 1.0f; + for (int j = 0; j < k; ++j) out.weights[static_cast(j)] /= denom; + } + return out; +} + +// __shfl_down_sync(0xffffffffu, v, off) semantics: lane `l` receives lane +// `l+off`, or its OWN value when `l+off >= 32`. Only lane 0 is ever read, and +// after the first level lane 0's dependency cone stays inside lanes < 16, so +// the self-returning lanes never enter it -- exactly as the block tree leaves +// red[t >= s] stale and reads only red[0]. +template +void ShflDownReduce(float (&v)[kWarp], Op op) { + for (int off = kWarp / 2; off > 0; off >>= 1) { + float n[kWarp]; + for (int lane = 0; lane < kWarp; ++lane) { + n[lane] = op(v[lane], v[lane + off < kWarp ? lane + off : lane]); + } + for (int lane = 0; lane < kWarp; ++lane) v[lane] = n[lane]; + } +} + +// --------------------------------------------------------------------------- +// Model B: the candidate warp kernel, one warp per token, the whole row in +// `p[lane][slot]` registers. Uses the SHIPPED reduction trees; `map` supplies +// the lane->expert map so the test can also instantiate the WRONG one. +template +RouterOut WarpModel(const std::vector& logits, int k, bool renormalize, Map map) { + RouterOut out; + out.weights.assign(static_cast(k), 0.0f); + out.indices.assign(static_cast(k), 0); + const int e = kWarp * VPT; + + float p[kWarp][VPT]; + for (int lane = 0; lane < kWarp; ++lane) { + for (int q = 0; q < VPT; ++q) { + p[lane][q] = logits[static_cast(map(lane, q))]; + } + } + + float m[kWarp]; + for (int lane = 0; lane < kWarp; ++lane) m[lane] = MoeRouterWarpTreeMax(p[lane]); + ShflDownReduce(m, [](float a, float b) { return fmaxf(a, b); }); + out.mx = m[0]; + + float s[kWarp]; + for (int lane = 0; lane < kWarp; ++lane) { + for (int q = 0; q < VPT; ++q) p[lane][q] = expf(p[lane][q] - out.mx); + s[lane] = MoeRouterWarpTreeSum(p[lane]); + } + ShflDownReduce(s, [](float a, float b) { return a + b; }); + out.sum = s[0]; + + for (int lane = 0; lane < kWarp; ++lane) { + for (int q = 0; q < VPT; ++q) { + float pj = out.sum > 0.0f ? p[lane][q] / out.sum : 0.0f; + if (!std::isfinite(pj)) pj = 0.0f; + p[lane][q] = pj; + } + } + + float denom = 0.0f; + for (int j = 0; j < k; ++j) { + float lv[kWarp]; + int li[kWarp]; + for (int lane = 0; lane < kWarp; ++lane) { + float v = -INFINITY; + int i = -1; + for (int q = 0; q < VPT; ++q) { // ascending expert index within the lane + if (p[lane][q] > v) { + v = p[lane][q]; + i = map(lane, q); + } + } + lv[lane] = v; + li[lane] = i; + } + for (int off = kWarp / 2; off > 0; off >>= 1) { + float nv[kWarp]; + int ni[kWarp]; + for (int lane = 0; lane < kWarp; ++lane) { + const int src = lane + off < kWarp ? lane + off : lane; + float cv = lv[lane]; + int ci = li[lane]; + if (WarpArgmaxTakes(lv[src], li[src], cv, ci)) { + cv = lv[src]; + ci = li[src]; + } + nv[lane] = cv; + ni[lane] = ci; + } + for (int lane = 0; lane < kWarp; ++lane) { + lv[lane] = nv[lane]; + li[lane] = ni[lane]; + } + } + const float best_v = lv[0]; + const int best = li[0]; + // Mask the winner in the owning lane's own register. The kernel uses the + // closed form for the shipped map (lane == best & 31, slot == best >> 5); + // the model searches so it stays correct for the mutation map too. + if (best >= 0) { + for (int lane = 0; lane < kWarp; ++lane) { + for (int q = 0; q < VPT; ++q) { + if (map(lane, q) == best) p[lane][q] = -INFINITY; + } + } + } + out.weights[static_cast(j)] = best_v; + out.indices[static_cast(j)] = static_cast(best); + denom += best_v; + } + if (renormalize) { + if (!(denom > 0.0f)) denom = 1.0f; + for (int j = 0; j < k; ++j) out.weights[static_cast(j)] /= denom; + } + (void)e; + return out; +} + +auto StridedMap = [](int lane, int slot) { return MoeRouterWarpExpert(lane, slot); }; + +// Bitwise, not approximate. The whole row is a byte-exactness claim. +bool BitEqual(const std::vector& a, const std::vector& b) { + if (a.size() != b.size()) return false; + return std::memcmp(a.data(), b.data(), a.size() * sizeof(float)) == 0; +} + +bool BitEqual(float a, float b) { return std::memcmp(&a, &b, sizeof(float)) == 0; } + +// For an INTERNAL intermediate that can legitimately be NaN. The PAYLOAD of a +// NaN produced by a float addition is not a property of the algorithm: +// IEEE-754 leaves the payload of NaN+NaN unspecified, and a compiler may +// commute the operands of an add freely. Measured on this host (gcc, -O1): +// `red[0] = red[0] + red[1]` and `n[0] = v[0] + v[1]` over the identical +// operands 0x7fc00000 and 0xffc00000 -- same source-level operand order -- +// produce 0x7fc00000 and 0xffc00000 respectively, because one loop got +// vectorized and the other did not. +// +// The SECOND exemption is the sign of a zero `mx`, and it is a HOST CODEGEN +// artifact, not a property of the algorithm. The algorithm has no freedom here: +// fmaxf(x, -INFINITY) == x BITWISE for every non-NaN x (all 2^32 of them, +// checked exhaustively during review), and past the seed both trees apply fmaxf +// to identical pairs in identical order — so whatever a conforming fmaxf does +// with the (+0.0f, -0.0f) tie, it does the same thing in both models. What +// actually differs is gcc VECTORIZING one transcription and not the other, and +// the vector max instruction does not implement fmaxf's zero handling; measured, +// this test's `signed-zeros` row passes at -O1 and fails at -O0 and -O2 on +// identical source. +// +// Neither exemption is a tolerance: one is a bit IEEE-754 leaves unspecified, +// the other is the host compiler's, both are proven below to be unobservable in +// the OUTPUT, and the weights and indices stay a strict, total comparison that +// holds at -O0, -O1 and -O2. +bool SameIntermediate(float a, float b) { + if (std::isnan(a) || std::isnan(b)) return std::isnan(a) && std::isnan(b); + if (a == 0.0f && b == 0.0f) return true; // +0.0f and -0.0f: an fmaxf tie + return BitEqual(a, b); +} + +// Round through bf16 exactly as the kernel's Load(const __nv_bfloat16*) does. +std::vector RoundToBf16(const std::vector& v) { + std::vector out(v.size()); + for (size_t i = 0; i < v.size(); ++i) out[i] = vt::BF16ToF32(vt::F32ToBF16(v[i])); + return out; +} + +// One (E, k, renormalize, logits) case, run through both models and compared +// bit for bit on every observable. +void CheckCase(const char* label, int e, int k, bool renormalize, + const std::vector& logits) { + const RouterOut block = BlockModel(logits, e, k, renormalize); + RouterOut warp; + switch (MoeRouterWarpValuesPerThread(e)) { + case 1: warp = WarpModel<1>(logits, k, renormalize, StridedMap); break; + case 2: warp = WarpModel<2>(logits, k, renormalize, StridedMap); break; + case 4: warp = WarpModel<4>(logits, k, renormalize, StridedMap); break; + case 8: warp = WarpModel<8>(logits, k, renormalize, StridedMap); break; + default: FAIL("E is not a dispatched width: " << e); return; + } + INFO("case: " << std::string(label)); + CAPTURE(e); + CAPTURE(k); + CAPTURE(renormalize); + // `mx` is never NaN by construction: the fmaxf(-INFINITY, .) seed erases + // every NaN before the tree ever runs, so a NaN here would itself be a + // failure. Its zero SIGN is exempt (SameIntermediate) and provably cannot + // reach the output -- see the dedicated test case below. + CHECK(SameIntermediate(block.mx, warp.mx)); + CHECK_FALSE(std::isnan(block.mx)); + // `sum` may legitimately be NaN (an all-NaN or Inf-minus-Inf row), and then + // only its NaN-ness is an algorithm property; see SameIntermediate. This + // weakens NOTHING observable: `sum > 0.0f` is false for EVERY NaN payload, so + // a NaN sum sends every probability through the `: 0.0f` arm of + // cuda_moe.cu:99 identically in both models -- which is exactly why the + // WEIGHTS below stay a strict, unconditional memcmp and still pass on those + // rows. The row's claim is about weights and indices; `sum` is checked only + // because pinning the intermediate localizes a defect faster. + CHECK(SameIntermediate(block.sum, warp.sum)); + // The actual contract. Strict, unconditional, never widened. Weights can + // never be NaN (probs are isfinite-clamped at cuda_moe.cu:100 and the + // exhausted-expert sentinel is exactly -INFINITY), so this comparison is + // total. + CHECK(BitEqual(block.weights, warp.weights)); + CHECK(block.indices == warp.indices); +} + +std::vector RandomRow(int e, uint32_t seed) { + std::vector v(static_cast(e)); + std::mt19937 rng(seed); + std::uniform_real_distribution d(-4.0f, 4.0f); + for (auto& x : v) x = d(rng); + return v; +} + +} // namespace + +TEST_CASE("router warp lane map is lane + 32*slot, and covers the row exactly once") { + // The map must be a BIJECTION onto [0, E) for every dispatched VPT -- + // a duplicated expert would be silently selectable twice, and a missing one + // could never be selected at all. + for (int vpt : {1, 2, 4, 8}) { + const int e = kMoeRouterWarpWidth * vpt; + std::vector hits(static_cast(e), 0); + for (int lane = 0; lane < kMoeRouterWarpWidth; ++lane) { + for (int slot = 0; slot < vpt; ++slot) { + const int expert = MoeRouterWarpExpert(lane, slot); + REQUIRE(expert >= 0); + REQUIRE(expert < e); + ++hits[static_cast(expert)]; + } + } + CAPTURE(vpt); + for (int j = 0; j < e; ++j) CHECK(hits[static_cast(j)] == 1); + } + // Slot q covers the 32 CONSECUTIVE experts [32q, 32q+32) -- the coalescing + // property, and the reason this is not vLLM's contiguous partition. + static_assert(MoeRouterWarpExpert(0, 0) == 0); + static_assert(MoeRouterWarpExpert(31, 0) == 31); + static_assert(MoeRouterWarpExpert(0, 1) == 32); + static_assert(MoeRouterWarpExpert(31, 7) == 255); + static_assert(MoeRouterWarpContiguousExpert(1, 0, 8) == 8); // vLLM's map, for contrast +} + +// The two intermediate exemptions in SameIntermediate are claims about the +// OUTPUT, so they are asserted rather than merely asserted-in-a-comment. If +// either of these ever stops holding, the exemption is invalid and the +// intermediate comparisons must go back to a strict memcmp -- they must never +// be widened further to accommodate a failure. +TEST_CASE("the intermediate exemptions cannot reach the output") { + // (1) SIGN OF A ZERO `mx`. mx is consumed only as `expf(logit - mx)` + // (cuda_moe.cu:84). Subtracting +0.0f and -0.0f differ for exactly one input, + // -0.0f, and expf maps both zeros to the same 1.0f -- so sp[] is bit-identical + // whichever zero the fmaxf tie returned. + for (float x : {-3.5f, -1.0f, -0.0f, 0.0f, 1.0f, 7.25f, INFINITY, -INFINITY}) { + const float minus_pos = x - 0.0f; + const float minus_neg = x - (-0.0f); + CAPTURE(x); + if (!BitEqual(minus_pos, minus_neg)) { + // The one permitted disagreement, and only for x == -0.0f. + CHECK(BitEqual(x, -0.0f)); + CHECK(BitEqual(minus_pos, -0.0f)); + CHECK(BitEqual(minus_neg, 0.0f)); + } + CHECK(BitEqual(expf(minus_pos), expf(minus_neg))); + } + CHECK(BitEqual(expf(-0.0f), 1.0f)); + CHECK(BitEqual(expf(0.0f), 1.0f)); + + // (2) NaN PAYLOAD of `sum`. sum is consumed only through the predicate + // `sum > 0.0f` (cuda_moe.cu:99), which is false for EVERY NaN, so every + // probability takes the `: 0.0f` arm regardless of payload. + for (uint32_t payload : {0x7fc00000u, 0xffc00000u, 0x7fc0deadu}) { + float nan_value; + std::memcpy(&nan_value, &payload, sizeof(float)); + REQUIRE(std::isnan(nan_value)); + CAPTURE(payload); + CHECK_FALSE(nan_value > 0.0f); + } +} + +TEST_CASE("only the derived widths dispatch to the warp kernel") { + CHECK(MoeRouterWarpValuesPerThread(32) == 1); + CHECK(MoeRouterWarpValuesPerThread(64) == 2); + CHECK(MoeRouterWarpValuesPerThread(128) == 4); + CHECK(MoeRouterWarpValuesPerThread(256) == 8); + // Everything else falls through to the unchanged block kernel. 512 and 1024 + // are excluded DELIBERATELY: above the incumbent's 256-wide block the seed + // loop accumulates several experts per thread, an association the spec does + // not derive. + for (int64_t e : {int64_t{0}, int64_t{1}, int64_t{8}, int64_t{16}, int64_t{31}, int64_t{33}, + int64_t{48}, int64_t{96}, int64_t{160}, int64_t{192}, int64_t{255}, + int64_t{257}, int64_t{384}, int64_t{512}, int64_t{1024}}) { + CAPTURE(e); + CHECK(MoeRouterWarpValuesPerThread(e) == 0); + } +} + +TEST_CASE("VT_MOE_ROUTER_WARP is on by default and '0' rolls back") { + CHECK(MoeRouterWarpFlagIsOn(nullptr)); + CHECK(MoeRouterWarpFlagIsOn("")); + CHECK(MoeRouterWarpFlagIsOn("1")); + CHECK(MoeRouterWarpFlagIsOn("on")); + CHECK_FALSE(MoeRouterWarpFlagIsOn("0")); +} + +// The comparators are the ONE thing the two models are not allowed to differ +// on, because the row's claim is that the two KERNELS pick the same expert. They +// are transcribed separately (see BlockArgmaxTakes / WarpArgmaxTakes) so that a +// divergence between cuda_moe.cu:169/:186 and :302 shows up HERE instead of +// being absorbed by a shared helper. This case is what pins them equal, over the +// values the argmax can actually carry: -INFINITY (the seed and the mask value), +// signed zeros, a NaN (which the isfinite clamp forbids downstream but which a +// defect could reintroduce), and the -1 "this lane holds nothing" sentinel. +TEST_CASE("the two kernels' argmax comparators agree on every (value, index) pair") { + const float values[] = {-INFINITY, -1.0f, -0.0f, 0.0f, 1.0f, INFINITY, std::nanf("")}; + const int idx[] = {-1, 0, 1, 7}; + for (float ov : values) { + for (int oi : idx) { + for (float cv : values) { + for (int ci : idx) { + CAPTURE(ov); + CAPTURE(oi); + CAPTURE(cv); + CAPTURE(ci); + CHECK(BlockArgmaxTakes(ov, oi, cv, ci) == WarpArgmaxTakes(ov, oi, cv, ci)); + } + } + } + } +} + +// THE SELECTOR PIN, executed. vt::cuda::MoeRouterWarpEnabled() (cuda_moe.cu) is +// a FRESH getenv per launch, so the AMBIENT environment decides which kernel +// vt::MoeRouterTopK runs. The on-device parallel-vs-serial gate in +// tests/vt/test_ops_moe_grouped.cpp used to neither set nor assert it, so an +// exported `VT_MOE_ROUTER_WARP=0` — exactly what spec §9 gates 6/7 tell the +// operator to export for the A/B — degenerated that sweep to block-vs-serial and +// reported the identical counts. It now runs its sweep TWICE under +// vt_test::ScopedMoeRouterWarp, and asserts the pinned state each time. +// +// That pin is shared code, not a copy, so this host case IS its contract test: +// the device file cannot be compiled without nvcc, this one runs anywhere. If +// the pin ever stops overriding the ambient value, this case goes red on a box +// whose environment disagrees with the arm — which is the whole failure mode. +TEST_CASE("pinning VT_MOE_ROUTER_WARP overrides the ambient value, then restores it") { + using vt_test::ScopedMoeRouterWarp; + const char* ambient = std::getenv(ScopedMoeRouterWarp::kName); + const bool ambient_had = ambient != nullptr; + const std::string ambient_saved = ambient_had ? std::string(ambient) : std::string(); + + // Whatever the shell said, the ON arm runs the WARP kernel ... + { + ScopedMoeRouterWarp pin("1"); + CHECK(ScopedMoeRouterWarp::EffectiveFlag()); + } + // ... and the OFF arm runs the incumbent block kernel. + { + ScopedMoeRouterWarp pin("0"); + CHECK_FALSE(ScopedMoeRouterWarp::EffectiveFlag()); + } + // Nesting resolves inner-first and unwinds to the outer pin, so an arm cannot + // leak into the arm that follows it. + { + ScopedMoeRouterWarp outer("0"); + { + ScopedMoeRouterWarp inner("1"); + CHECK(ScopedMoeRouterWarp::EffectiveFlag()); + } + CHECK_FALSE(ScopedMoeRouterWarp::EffectiveFlag()); + } + // Unset is restored as ABSENT, not as the empty string (which reads as ON). + { + ScopedMoeRouterWarp pin(nullptr); + CHECK(std::getenv(ScopedMoeRouterWarp::kName) == nullptr); + CHECK(ScopedMoeRouterWarp::EffectiveFlag()); // default ON + } + + const char* after = std::getenv(ScopedMoeRouterWarp::kName); + CHECK((after != nullptr) == ambient_had); + if (ambient_had && after != nullptr) CHECK(ambient_saved == std::string(after)); +} + +TEST_CASE("warp reduction trees reproduce the block tree bit for bit (random rows)") { + for (int e : {32, 64, 128, 256}) { + for (int k : {1, 2, 8}) { + for (bool renormalize : {false, true}) { + for (uint32_t seed = 0; seed < 8; ++seed) { + const std::vector f32 = RandomRow(e, 1234u + seed + 97u * static_cast(e)); + CheckCase("random-f32", e, k, renormalize, f32); + // bf16 LOGIT arm. The weights are f32 either way, but bf16 rounding + // manufactures the exact ties the tie-break has to resolve, and this + // project has a recorded case where a bf16 store hid a defect that + // only an f32 arm caught -- so both arms run, always. + CheckCase("random-bf16", e, k, renormalize, RoundToBf16(f32)); + } + } + } + } +} + +TEST_CASE("warp reduction trees reproduce the block tree on adversarial rows") { + for (int e : {32, 64, 128, 256}) { + const size_t n = static_cast(e); + for (bool renormalize : {false, true}) { + // Exact-tie storm: blocks of identical logits, so many experts tie at the + // max and the lowest-index tie-break decides every slot. + { + std::vector v(n); + for (size_t j = 0; j < n; ++j) v[j] = static_cast((j / 4) % 5); + CheckCase("tie-storm", e, 8, renormalize, v); + } + // Every expert identical: the whole row is one tie. + { + std::vector v(n, 0.5f); + CheckCase("all-equal", e, 8, renormalize, v); + } + // k > E: the -1 sentinel path (indices -1, weights -INFINITY). This is a + // KERNEL-level input only — vt::MoeRouterTopK rejects k > E outright + // ("top_k must be in [1, num_experts]", src/vt/ops.cpp), and within + // 1 <= k <= E the sentinel cannot fire because every prob is finite and + // >= 0, so only exhausting all E experts leaves best == -1. The two + // kernels' `best >= 0` guards still have to AGREE, and this is the only + // place that can compare them: the on-device parallel-vs-serial case in + // test_ops_moe_grouped.cpp cannot, since the op throws first (and its + // serial oracle would write sp[-1]). Keep this arm; do not add a device one. + { + CheckCase("k-gt-E", e, e + 3, renormalize, RandomRow(e, 4242u)); + } + // k == 1. + { + CheckCase("k-eq-1", e, 1, renormalize, RandomRow(e, 777u)); + } + // All NaN. The -INFINITY max seed erases every NaN, so mx is -INFINITY, + // the sum is NaN, the clamp zeroes every prob, and the tie-break must + // then hand back 0,1,2,...,k-1. + { + std::vector v(n, std::nanf("")); + CheckCase("all-nan", e, 8, renormalize, v); + } + // All +Inf, all -Inf, and a mix with NaN. + { + CheckCase("all-pos-inf", e, 8, renormalize, std::vector(n, INFINITY)); + CheckCase("all-neg-inf", e, 8, renormalize, std::vector(n, -INFINITY)); + std::vector v(n, 1.0f); + v[0] = std::nanf(""); + v[1] = INFINITY; + v[n - 1] = -INFINITY; + CheckCase("nan-inf-mix", e, 8, renormalize, v); + } + // Signed zeros: -0.0f is the one float for which x + 0.0f flips the sign + // bit, which is what makes dropping the block's zero pad leaves legal. + { + std::vector v(n, -0.0f); + for (size_t j = 0; j < n; j += 2) v[j] = 0.0f; + CheckCase("signed-zeros", e, 8, renormalize, v); + } + // Huge magnitudes: the max subtraction is what keeps expf finite. + { + std::vector v(n); + for (size_t j = 0; j < n; ++j) v[j] = 80.0f + static_cast(j); + CheckCase("huge-magnitude", e, 8, renormalize, v); + } + } + } +} + +// THE DISCRIMINATION TEST. A parallel-vs-serial comparison that cannot tell the +// two lane maps apart is not testing the reduction order, and would report +// GREEN for a kernel that quietly reassociates the softmax sum -- which is the +// exact defect this row exists to avoid. So pin the oracle's discriminating +// power in the gate rather than in a one-off mutation transcript. +// +// The contiguous map (lane L owns experts [VPT*L, VPT*L + VPT)) is vLLM's OWN +// partition, topk_softmax_kernels.cu:344-346. It is a correct softmax; it is +// simply a DIFFERENT association, so it lands on a different last ulp. +TEST_CASE("the contiguous (vLLM) lane map is detectably different, so the oracle bites") { + for (int vpt : {2, 4, 8}) { // VPT == 1 has one slot per lane: no association to differ + const int e = kMoeRouterWarpWidth * vpt; + auto contiguous = [vpt](int lane, int slot) { + return MoeRouterWarpContiguousExpert(lane, slot, vpt); + }; + int differing_rows = 0; + for (uint32_t seed = 0; seed < 16; ++seed) { + const std::vector v = RandomRow(e, 555u + seed); + const RouterOut block = BlockModel(v, e, 8, true); + RouterOut wrong; + switch (vpt) { + case 2: wrong = WarpModel<2>(v, 8, true, contiguous); break; + case 4: wrong = WarpModel<4>(v, 8, true, contiguous); break; + default: wrong = WarpModel<8>(v, 8, true, contiguous); break; + } + // WEIGHTS ALONE. This case exists to prove the WEIGHTS memcmp — the + // comparison the row's claim actually rests on, and the one the on-device + // parallel-vs-serial gate makes — has discriminating power. An earlier + // revision wrote `!BitEqual(sum) || !BitEqual(weights)`, which pins only + // "sum-or-weights": a contiguous map that moved the denominator but never + // reached the output would have satisfied it, and the weights memcmp would + // still have been unproven. + if (!BitEqual(block.weights, wrong.weights)) ++differing_rows; + } + CAPTURE(vpt); + CHECK(differing_rows > 0); + } +} diff --git a/tests/vt/test_ops_moe_grouped.cpp b/tests/vt/test_ops_moe_grouped.cpp index a45aef6d8..6f4fd90c8 100644 --- a/tests/vt/test_ops_moe_grouped.cpp +++ b/tests/vt/test_ops_moe_grouped.cpp @@ -10,12 +10,14 @@ #include #include #include +#include #include #include #include #include #include +#include "moe_router_warp_env.h" #include "vllm/model_executor/model_loader/mxfp4_dequant.h" #include "vllm/model_executor/model_loader/nvfp4_dequant.h" #include "vt/cuda/moe_decode_ref.h" @@ -495,6 +497,15 @@ TEST_CASE("moe_grouped_gemm_nvfp4 validates shapes loudly (CPU dispatch)") { // softmax + comparison-only argmax with the same lowest-index tie-break); this // test pins it on adversarial inputs — exact ties, near-ties, uneven loads, and // M in {1,8,16}, for f32 and bf16 logits (bf16 rounding manufactures ties). +// +// It is ALSO the on-device gate for the single-warp router (VT_MOE_ROUTER_WARP, +// issue #378), which is what `vt::MoeRouterTopK` dispatches by default for +// E in {32,64,128,256}: that kernel's whole claim is that it is byte-identical +// to this same serial reference, so it must survive every case here. The +// portable companion tests/vt/test_moe_router_warp_map.cpp proves the REDUCTION +// ORDER with no GPU; this one proves the KERNEL. Neither substitutes for the +// other, and neither comparison is ever loosened: weights are compared with +// memcmp and indices with ==. // Guarded on VLLM_CPP_CUDA: the parallel-vs-serial cross-check calls the // CUDA-only reference vt::cuda::MoeRouterTopKSerialCuda (cuda_moe.cu), so a // runtime HasCuda() skip is not enough — the symbol is undefined at link time @@ -507,6 +518,43 @@ TEST_CASE("CUDA moe_router_topk parallel == serial byte-for-byte (adversarial)") } Backend& gpu = vt::GetBackend(DeviceType::kCUDA); + // ── The kernels' `-1` sentinel is UNREACHABLE through this op. ──────────── + // An earlier revision of this case ran `k = E + 3` to exercise the `best < 0` + // path (indices -1, weights -INFINITY, denom<=0 -> 1). That is not an input + // vt::MoeRouterTopK has: it validates + // VT_CHECK(args.top_k >= 1 && args.top_k <= e, + // "moe_router_topk: top_k must be in [1, num_experts]"); + // (src/vt/ops.cpp) BEFORE dispatch, so the arm threw, this whole case aborted + // partway, and the byte-exactness sweep below never finished — while every + // assertion that HAD run still reported passed. Within `1 <= k <= E` the + // sentinel cannot fire either: after `sum > 0 ? sp/sum : 0` and the isfinite + // clamp every prob is finite and >= 0 (cuda_moe.cu), the argmax seeds + // best_v = -INFINITY, and only masking ALL E experts leaves nothing greater — + // which needs a (k+1)-th round, i.e. k > E. The sentinel is a defensive guard, + // not a reachable behaviour of this op. + // + // It was also UNSAFE to drive from here: the serial oracle this case compares + // against is called directly (vt::cuda::MoeRouterTopKSerialCuda), bypassing + // the validator, and its `sp[best] = -INFINITY` carries no `best >= 0` guard + // — so a k > E round writes shared memory at sp[-1]. That is the oracle, so + // the fix is to stop feeding it an input the op forbids, never to edit it. + // + // The two kernels' sentinel guards are still compared, at the level where the + // sentinel exists: tests/vt/test_moe_router_warp_map.cpp models the block and + // warp kernels directly and keeps its `k > E` arm. Here we pin the contract + // that makes the sentinel dead. + { + QueueGuard gq(gpu); + const int64_t T = 2, E = 32; + const int k = static_cast(E) + 3; + DeviceTensor dlog(gpu, gq.q, DType::kF32, {T, E}); + DeviceTensor w_bad(gpu, gq.q, DType::kF32, {T, static_cast(k)}); + DeviceTensor i_bad(gpu, gq.q, DType::kI32, {T, static_cast(k)}); + const vt::MoeRouterTopKArgs args{k, false}; + CHECK_THROWS_AS(vt::MoeRouterTopK(gq.q, w_bad.tensor(), i_bad.tensor(), dlog.tensor(), args), + std::runtime_error); + } + auto run = [&](int64_t T, int64_t E, int k, bool renorm, DType dt, const std::vector& logits) { QueueGuard gq(gpu); @@ -550,41 +598,111 @@ TEST_CASE("CUDA moe_router_topk parallel == serial byte-for-byte (adversarial)") CHECK(idiff == 0); }; + // ── The sweep runs TWICE, with VT_MOE_ROUTER_WARP PINNED both ways. ─────── + // vt::cuda::MoeRouterWarpEnabled() (cuda_moe.cu) is a FRESH getenv per launch, + // so WHICH kernel vt::MoeRouterTopK runs below is decided by the environment + // ctest was started in. An earlier revision of this case neither set, cleared + // nor asserted that variable, which made it unable to say what it had tested: + // with `VT_MOE_ROUTER_WARP=0` exported — exactly what spec §9 gates 6/7 tell + // the operator to export for the same-binary A/B, in the same shell — every + // `run()` below silently exercised the BLOCK kernel, i.e. block-vs-serial, + // green since 6a8c5cf9, reporting the identical case and assertion counts. A + // green run could not distinguish "the warp kernel is byte-exact" from "the + // warp kernel never ran". + // + // So: pin it, assert the pinned state, and run the sweep under BOTH arms — + // "1" pins the candidate warp kernel, "0" pins the same-binary rollback, and + // each must be byte-identical to the serial oracle. The pin restores the + // ambient value (or its absence) on scope exit. + // + // This DOUBLES this case's assertion count on purpose. Spec §9 gate 2 treats a + // changed count as a red flag; the change is this, and it is the point. + // + // The body below keeps its original indentation so this change reads as what + // it is — a pure wrapper — and a reviewer can see that not one line of the + // sweep itself moved. + auto sweep = [&]() { for (DType dt : {DType::kF32, DType::kBF16}) { for (bool renorm : {true, false}) { for (int64_t T : {int64_t{1}, int64_t{8}, int64_t{16}}) { - // 35B routing shape: E=128, top-8. Random distinct logits. - { - const int64_t E = 128; - std::vector lg(static_cast(T * E)); - std::mt19937 rng(1234u + static_cast(T)); - std::uniform_real_distribution d(-4.0f, 4.0f); - for (auto& v : lg) v = d(rng); - run(T, E, 8, renorm, dt, lg); - } - // Exact-tie storm: blocks of identical logits so many experts tie at - // the max; the tie-break (lowest index) must agree across paths. - { - const int64_t E = 128; - std::vector lg(static_cast(T * E)); - for (int64_t t = 0; t < T; ++t) - for (int64_t e = 0; e < E; ++e) - lg[static_cast(t * E + e)] = static_cast((e / 4) % 5); // 5 tie groups - run(T, E, 8, renorm, dt, lg); - } - // Uneven load + near-ties around the top-k boundary (E=256, top-8). - { - const int64_t E = 256; - std::vector lg(static_cast(T * E)); - for (int64_t t = 0; t < T; ++t) - for (int64_t e = 0; e < E; ++e) - lg[static_cast(t * E + e)] = - (e < 12 ? 3.0f : 0.0f) + 1e-4f * static_cast((e * 7 + t) % 3); - run(T, E, 8, renorm, dt, lg); + // E in {32,64,128,256}: EVERY width the single-warp router dispatches + // (VPT = E/32 in {1,2,4,8}, VT_MOE_ROUTER_WARP, issue #378), so the + // byte-exactness of each derived lane map is pinned on device and not + // only by the portable reduction-order test. + // + // The 35B-A3B gate model is E=256 top-8, NOT the E=128 this case used + // to claim: num_experts=256, num_experts_per_tok=8, and its 40 MoE + // layers are the 40 router calls/step the decode trace shows. E=256 had + // only a hand-built near-tie pattern here and never random logits. + for (int64_t E : {int64_t{32}, int64_t{64}, int64_t{128}, int64_t{256}}) { + // Random distinct logits. + { + std::vector lg(static_cast(T * E)); + std::mt19937 rng(1234u + static_cast(T) + 7919u * static_cast(E)); + std::uniform_real_distribution d(-4.0f, 4.0f); + for (auto& v : lg) v = d(rng); + run(T, E, 8, renorm, dt, lg); + run(T, E, 1, renorm, dt, lg); // k=1 + // NO k > E arm here — it is not an input this op has. See the + // contract block at the top of this case. + } + // Exact-tie storm: blocks of identical logits so many experts tie at + // the max; the tie-break (lowest index) must agree across paths. + { + std::vector lg(static_cast(T * E)); + for (int64_t t = 0; t < T; ++t) + for (int64_t e = 0; e < E; ++e) + lg[static_cast(t * E + e)] = + static_cast((e / 4) % 5); // 5 tie groups + run(T, E, 8, renorm, dt, lg); + } + // Uneven load + near-ties around the top-k boundary. + { + std::vector lg(static_cast(T * E)); + for (int64_t t = 0; t < T; ++t) + for (int64_t e = 0; e < E; ++e) + lg[static_cast(t * E + e)] = + (e < 12 ? 3.0f : 0.0f) + 1e-4f * static_cast((e * 7 + t) % 3); + run(T, E, 8, renorm, dt, lg); + } + // Degenerate rows. The -INFINITY max seed ERASES NaN, so an all-NaN + // row normalizes to all zeros and the tie-break must hand back + // 0,1,...,k-1; Inf rows exercise the same clamp from the other side. + // These are the rows CUDA-graph padding actually produces. + { + run(T, E, 8, renorm, dt, + std::vector(static_cast(T * E), std::nanf(""))); + run(T, E, 8, renorm, dt, + std::vector(static_cast(T * E), INFINITY)); + run(T, E, 8, renorm, dt, + std::vector(static_cast(T * E), -INFINITY)); + std::vector lg(static_cast(T * E), 1.0f); + for (int64_t t = 0; t < T; ++t) { + lg[static_cast(t * E)] = std::nanf(""); + lg[static_cast(t * E + 1)] = INFINITY; + lg[static_cast(t * E + E - 1)] = -INFINITY; + } + run(T, E, 8, renorm, dt, lg); + } } } } } + }; // sweep + + { + // ON: vt::MoeRouterTopK dispatches MoeRouterTopKWarpKernel for every E here. + vt_test::ScopedMoeRouterWarp pin("1"); + REQUIRE(vt_test::ScopedMoeRouterWarp::EffectiveFlag()); + sweep(); + } + { + // OFF: the same binary falls through to the unchanged block kernel. This arm + // is what the pre-existing evidence actually covered, and it stays covered. + vt_test::ScopedMoeRouterWarp pin("0"); + REQUIRE_FALSE(vt_test::ScopedMoeRouterWarp::EffectiveFlag()); + sweep(); + } } #endif // VLLM_CPP_CUDA