Skip to content

qwen3_5 CUDA: a captured bucket's padding rows no longer go through the routed experts (35B, every multi-row forward) - #1047

Closed
jschmied wants to merge 2 commits into
ashhart:mainfrom
jschmied:q35-moe-live-rows
Closed

jschmied wants to merge 2 commits into
ashhart:mainfrom
jschmied:q35-moe-live-rows

Conversation

@jschmied

@jschmied jschmied commented Oct 18, 2026 •

Copy link
Copy Markdown
Contributor

What this changes

On CUDA, qwen3_5_moe (Qwen3.6-35B-A3B) runs every forward of more than one row from a captured graph whose row count is a bucket, and the padding rows were routed through the experts.

cuda_engine.zig forward picks bucket = if (lone) 1 else if (total <= 16) 16 else most_verify (32) on v1.1.0 (zig/src/families/qwen3_5/cuda_engine.zig:432); 1.1.1 captures up to 128 rows. The kernels are supposed to read the true row count from the params' word 2 (P[2]). The norm and the embedding do. cuda_routed.zig rows() did not: it passed the bucket's row count to the router, the top-k, the plan and combine. The padding rows hold stale activations from earlier forwards, so they were routed to real experts, and the plan added those experts' weight reads to every MoE layer. A step beside one other stream (2 live rows) cost about as much as a 16-row window.

The padding came in with c3861d3 ("qwen3_5: DFlash2 drafts for Qwen3.8-27B on CUDA"). It is unchanged on v1.1.0 (76aafcf).

The fix (two commits):

  • tf_moe_router, tf_moe_topk and tf_moe_combine take the forward's params and skip rows from P[2] up. A null pointer keeps every row; the MTP head passes null, since it runs uncaptured at its exact row count.
  • grouped.Router.routeRows plans only the first P[2] rows' pairs at any width:
    • up to 1,024 pairs, plan_rows (experts.cu): plan_kernel over the first *rows * per pairs;
    • past 1,024 pairs (more than 113 rows at nine pairs a row, so 1.1.1's 128-row buckets), plan_rank_rows and plan_scatter_rows: plan_rank and plan_scatter over the same live pairs, with plan_offsets shared. plan_rank and plan_rank_rows share one device body.
  • Launches stay sized for the bucket, so the captured graphs are still valid and read the count at each replay. Padding rows get no picks, no plan items, no expert work and no combine.
  • route, plan_kernel, plan_rank, plan_offsets and plan_scatter are untouched. The SASS of every function in the experts image that main has is byte-identical to main on sm_80, 86, 89, 120 and 121 (45 of 45 functions); the only new functions are the three *_rows kernels. The other families that use the plan (GLM, Nemotron, Flash Next, Drama 3) see no change.

No sums change order. Every live row's bits are unchanged.

Receipt

  • Environment: Linux, NVIDIA GB10 (DGX Spark, sm_121), CUDA 13.0 nvcc (13.0.88), Zig 0.17.0, -Dsm=121 -Doptimize=ReleaseFast. Main is v1.1.0 (76aafcf); this branch is 13a296e. Qwen3.6-35B-A3B MLX 4-bit with its MTP layer; Qwen3.8-27B MLX 4-bit with DFlash2.

  • Regression test, tf-cuda-test moe-live (also run at the end of moe): the routed chain (router, top-k, plan, affine-4 SwiGLU experts, combine) is captured once a bucket and replayed at changing live counts in this order:

    • bucket 1 at 1; bucket 16 at 1, 2, 3, 7, 15, 16, 3, 14, 2; bucket 32 at 17, 18, 31, 32, 3, 30, 17;
    • bucket 64 at 33, 1, 64, 40, 63, 9;
    • bucket 128 (wide plan) at 1, 9, 33, 64, 127, 128, 113, 114, 2, 120, 64 (113 is the last live count under 1,024 pairs, 114 the first over).

    Before each replay the padding is poisoned: input rows past the live count are NaN, +huge and -huge in turn, stale picks name a real expert, stale weights and logits are NaN, and the experts' pair outputs hold a sentinel. Each replay must:

    • plan exactly the live rows' pairs (sum of the items' counts = live x 9, largest member below that);
    • leave every padding pair's expert output untouched (no expert ran on padding);
    • give live rows equal byte for byte to one-row uncaptured runs;
    • leave the padding rows of the output unwritten.

    Results:

    • This branch: 5 of 5 PASS.
    • The previous head a16a1d4 (small plan only): buckets 1 to 64 PASS, then FAIL bucket 128 at 1 live rows planned 1152 pairs (largest member 1151), want the live rows' 9.
    • Main (the same test on main's launches, without the params): FAIL bucket 16 at 1 live rows planned 104 pairs (largest member 143), want the live rows' 9.
  • Kernel contract (zig build kernel-contract -Dnvcc=..., 13a296e): SUMMARY 55 kernels: 55 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 0 limit problems.

    • The served-kernel checks run router, top-k and combine with P[2] below the row count at rows 1, 15, 16, 17, 64 and 128: live rows meet float64, padding rows' outputs stay unwritten.
    • routeRows is checked against route over the live pairs at every live count of buckets 1, 16, 32, 64, 114 and 128 (byte-equal, five repeats, guards). Buckets 114 and 128 take the wide plan, so this also checks the wide plan against the one-block plan wherever the live pairs fit in 1,024.
    • plan_rows, plan_rank_rows and plan_scatter_rows are registered in qwen_served_records.zig. tf_moe_router's sm_80 record goes from 40 to 32 registers; every other existing record is unchanged.
    • tf-cuda-test kernel-contract-limits over fatbins for sm_80, 86, 89, 120 and 121: every registered launch fits, every control red.
  • Tests run on 13a296e:

    • zig build test from an empty local and global cache: 152/152 steps, 2,413 of 2,422 tests passed, 9 skipped, 0 failed.
    • zig build test-kernel-contract test-qwen-served-contract-host: 12/12 steps, 46/46 passed. zig build test-exl3-prompt-host: 7/7 passed.
    • tf-cuda-test: moe 9, moe-live 5, moe-glue 1, experts-rows 7, route-rows 1, experts3 8, tree 1, glue 1, sample 1, smoke 7 and graph 3 PASS lines, 0 FAIL, 0 SKIP.
    • zig fmt --check is clean.
  • Engine output and timing, tf-qwen35-cuda over the same 300 receipt ids, main against 13a296e, ABAB over four rounds (35B) and two (27B):

    • every step, window, tree, beside and prompt output file is byte-identical to main (35B: 43 of 43 files; 27B: 12 of 12); the 35B MTP self-check passes 3/3 on both.
    ms a forward, 35B main this branch
    step, one row alone 13.50 13.55
    step beside a second stream (2 rows, bucket 16) 28.70 16.10
    windows of 2 20.30 15.98
    windows of 3 21.43 17.33
    windows of 8 24.86 22.06
    windows of 15 30.44 28.10
    windows of 16 (bucket full) 27.16 27.20
    16 rows beside one (17 rows, bucket 32) 40.78 33.03

    The 27B (dense, no routed experts) is unchanged within 0.1% in every cell.

  • Served results, prompt reuse and concurrency, and the served ABBA timings are in the comments below.

  • Not run, and why:

    • 1.1.1's 128-row buckets on a served model: this branch is on v1.1.0, whose buckets stop at 32, so served runs never take the wide plan. The wide plan is covered by moe-live and the contract above.
    • Other GPUs. The RTX 4090/5090 and other sm_80/86/89/120 cards were not run (only their limit records were checked).
    • Other CUDA versions. nvcc 13.2 and 13.3 were not run: only 13.0 is installed here. Their record sets for tf_moe_* and the *_rows plan kernels are not re-derived.
    • Windows, Metal and HIP. No Metal or HIP source changes. HIP's routed MoE sizes its buffers by bucket but routes and plans the true rows, so it is not affected.

Not in this PR

The other CUDA families with MoE (GLM, Nemotron, Flash Next, Drama 3) do not have this flaw. None of them pads a forward to a row bucket: they have no CUDA graphs, or graphs keyed on the exact row count. Their router, top-k and plan already see the live rows.

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
  • No precision traded for speed. No sums change and every kept row's bits are unchanged.
  • A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt.
  • Prompt processing is no slower than the last release.
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

…perts

A forward of more than one row replays a graph captured for 16 or 32 rows, and
the kernels read the true row count from the params' word 2. The norm and
embedding did; the routed experts did not: router, top-k, the plan and combine
took the bucket's row count, so each padding row (stale activations from an
earlier forward) was routed to real experts, and the plan added those experts'
weight reads to every MoE layer.

tf_moe_router, tf_moe_topk and tf_moe_combine take the params and skip rows
from word 2 up (null keeps all rows, as the MTP head calls them), and
plan_rows, plan_kernel over the first live rows' pairs, plans only those. The
launches stay sized for the bucket, so captured graphs are unchanged; every
kept row's bits are unchanged.

tf-cuda-test moe-live (also run by moe): the routed chain captured at buckets
1, 16 and 32 and replayed at live counts 1..32, with non-zero padding rows,
plans only the live rows' pairs and equals one-row runs byte for byte. The
served contract checks the three kernels at live counts below the rows
against float64 with the padding rows unwritten, and plan_rows against
plan_kernel over the live pairs; plan_rows and tf_moe_router's sm_80 record
are registered.
@ashhart

ashhart commented Oct 18, 2026

Copy link
Copy Markdown
Owner

Thanks @jschmied. The runtime row guards and stable grouped plan preserve the small path's live arithmetic, and the replay test is useful.

One integration fix is needed: 1.1.1, shipping shortly, captures up to 128 rows, so nine slots can exceed the small-plan limit, and routeRows then ignores the live count. Please make the wide plan live-aware too and extend captured tests beyond the old buckets, with poisoned padding and changing replay counts.

We'll carry it onto 1.1.2 ourselves, keeping the prompt and MTP arithmetic cutoffs. Please attach the pending served, cache/concurrency and ABBA results from this head; the prototype timings remain separate. Targeting 1.1.2 after those gates.

Posted by Ash's AI Assistant Synthena

routeRows took the live count only on the one-block plan (up to 1,024 pairs).
A bucket of more than 113 rows (nine pairs a row) took route, so its padding
rows' stale picks were ranked, offset and scattered, and their experts ran.

plan_rank_rows and plan_scatter_rows are plan_rank and plan_scatter over the
first *rows * per pairs. routeRows takes them past 1,024 pairs; the launches
stay sized for the bucket and plan_offsets is shared, so a captured graph
reads the count at each replay. plan_rank keeps its code through a shared
device body; route, and every other family's plan, is unchanged.

tf-cuda-test moe-live captures buckets 1, 16, 32, 64 and 128 and replays each
at changing live counts (128 at 1, 9, 33, 64, 127, 128, 113, 114, 2, 120, 64).
Each replay poisons the padding: rows of NaN and +-huge, stale picks naming a
real expert, NaN weights and logits, sentinel expert outputs. Only the live
rows' pairs are planned, no expert writes a padding pair, the padding rows of
the output stay unwritten, and the live rows equal one-row uncaptured runs
byte for byte. The served contract checks routeRows against route over the
live pairs at every live count of buckets 1, 16, 32, 64, 114 and 128; the two
kernels are registered.
@jschmied

Copy link
Copy Markdown
Contributor Author

Thanks. Pushed 13a296e on top of a16a1d4 (no rebase, still on v1.1.0).

Wide plan. routeRows now reads the live count at every width. Past 1,024 pairs (more than 113 rows at nine pairs a row) it runs plan_rank_rows and plan_scatter_rows: plan_rank and plan_scatter over the first *rows * per pairs, with plan_offsets shared and the launches still sized for the bucket. plan_rank and plan_rank_rows share one device body. route and the existing plan kernels are unchanged: across sm_80, 86, 89, 120 and 121, all 45 experts-image functions that main has have byte-identical SASS to main. The two new kernels are registered, and kernel-contract-limits passes on all five SMs.

Tests. moe-live now captures buckets 1, 16, 32, 64 and 128 and replays each graph at changing live counts. The 128 graph runs at 1, 9, 33, 64, 127, 128, 113, 114, 2, 120, 64; 113 is the last count under 1,024 pairs and 114 the first over. Before each replay the padding is poisoned: input rows are NaN, +huge and -huge in turn, stale picks name a real expert, weights and logits are NaN, and expert outputs hold a sentinel. Each replay must plan only live x 9 pairs, leave every padding pair's expert output untouched, leave padding output rows unwritten, and give live rows equal byte for byte to one-row uncaptured runs.

  • 13a296e: 5/5 PASS
  • a16a1d4: buckets 1 to 64 pass, then FAIL bucket 128 at 1 live rows planned 1152 pairs (largest member 1151), want the live rows' 9
  • main: FAIL bucket 16 at 1 live rows planned 104 pairs (largest member 143), want the live rows' 9

The served contract also checks routeRows against route over the live pairs at every live count of buckets 1, 16, 32, 64, 114 and 128. That runs the wide plan against the one-block plan wherever the live pairs fit in 1,024.

Gates on 13a296e:

  • zig build test from empty caches: 152/152 steps, 2,413/2,422 passed, 9 skipped
  • test-kernel-contract + test-qwen-served-contract-host: 46/46
  • kernel-contract: SUMMARY 55 kernels: 55 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 0 limit problems
  • tf-cuda-test moe, moe-live, moe-glue, experts-rows, route-rows, experts3, tree, glue, sample, smoke, graph: 0 FAIL
  • zig fmt --check: clean

Results from 13a296e. Main is v1.1.0 (76aafcf). GB10, CUDA 13.0. Served runs are ABBA: main, branch, branch, main, a fresh server each. v1.1.0 buckets stop at 32 rows, so served runs never reach the wide plan. The wide plan is covered by the tests above.

Served token_sha against main (bench_openai 2 prompts x T0/T1 x 3 seeds; bench_concurrent with --alone --serial):

model, mode openai concurrent concurrent == alone == serial
35B MTP, parallel 1 / 2 / 4 / 8 12/12, 12/12, 12/12, - 32/32, 96/96, 160/160, 288/288 0 unequal, 0 failed
35B plain, parallel 1 / 2 / 4 / 8 12/12, 12/12, 12/12, - 32/32, 96/96, 160/160, 288/288 0 unequal, 0 failed
27B DFlash2, parallel 4 12/12 160/160 0 unequal, 0 failed
27B plain, parallel 4 12/12 160/160 0 unequal, 0 failed

Each cell holds for both branch starts against the first main start. The second main start also equals the first. Drafted equals plain: 35B MTP vs plain token_sha are equal at every --parallel (openai 12/12, concurrent 32/96/160/288), and 27B DFlash2 vs plain openai is 12/12.

Prompt reuse and concurrency (tools/zig/reuse_served.py, --context 65536 --parallel 2): record on a server with --prompt-cache-gib 8, replay the same 43 requests on a server with --prompt-cache-gib 0, compare full token digests. This covers 9-turn agent sessions up to 34k tokens, alternating conversations, an edited transcript, thinking off, a cancel and resend, and two concurrent requests.

resumed == fresh resumed from cache drafted == plain branch == main
35B MTP, branch 43/43 40 23/23 43/43
35B MTP, main 43/43 40 23/23
27B DFlash2, branch 43/43 40 23/23 43/43
27B DFlash2, main 43/43 40 23/23

The script's own verdict line fails the same way on main and the branch (3/43 ... as cached as planned). It expects resume points at exact token counts, and this server keeps states on 64-token boundaries (for example cached 8576 of a planned 8652). Every reply still equals fresh. bench_cache (2,571-token prompt, below the 4,096-token minimum the cache keeps) gives first == repeated == concurrent on both builds: 35B MTP and plain at --parallel 1, 27B DFlash2.

ABBA served speed, tok/s, mean of the two starts per build (decode median a request for 1 stream, aggregate at the top --parallel):

35B main branch
MTP p1, 1 stream, code / chat T0 94.5 / 71.3 127.6 / 101.5
MTP p2, aggregate code / chat T0 160.3 / 119.3 196.2 / 167.0
MTP p4, aggregate code / chat T0 264.4 / 198.0 274.2 / 231.9
MTP p8, aggregate code / chat T0 360.3 / 310.2 357.1 / 307.9
MTP p8, aggregate code / chat T1 244.8 / 224.5 248.2 / 230.2
plain p1, 1 stream, code / chat T0 75.9 / 75.5 75.4 / 75.0
plain p2, aggregate code / chat T0 74.3 / 69.5 134.2 / 131.9
plain p4, aggregate code / chat T0 145.6 / 142.6 235.2 / 227.9
plain p8, aggregate code / chat T0 272.2 / 263.5 383.4 / 363.6
cold TTFT, 2,571 tokens, MTP / plain (ms) 1336.8 / 1257.9 1334.9 / 1258.3

Two cells are slightly slower on the branch:

  • MTP p8 T0 aggregate: -0.9% code, -0.7% chat.
  • plain p1 one stream: -0.7%.

The engine-level step is 13.50 vs 13.55 ms (below), so the lone step costs about 0.4%.

27B (dense, no routed experts): plain 13.3 tok/s one stream on both builds; p4 aggregate T0 47.0 / 46.3 main vs 47.0 / 46.4 branch. DFlash2 p4 aggregate T0 89.0 / 75.2 main vs 89.2 / 75.5 branch. Cold TTFT 3414.7 vs 3420.5 ms.

Engine level, tf-qwen35-cuda over 300 ids, ABAB 4 rounds (35B), ms a forward, median:

35B main branch
step alone 13.50 13.55
step beside a second stream (2 rows) 28.70 16.10
windows of 2 / 3 / 8 20.30 / 21.43 / 24.86 15.98 / 17.33 / 22.06
windows of 15 / 16 30.44 / 27.16 28.10 / 27.20
17 rows (16 beside one, bucket 32) 40.78 33.03

Every step, window, tree, beside and prompt output file is byte-identical to main: 35B 43/43, 27B 12/12. The 35B MTP self-check passes 3/3 on both builds. The 27B is within 0.1% in every cell. The description now carries only numbers from 13a296e. The earlier prototype table is gone from it.

@ashhart ashhart mentioned this pull request Oct 19, 2026
ashhart added a commit that referenced this pull request Oct 21, 2026
…dits

The 1.1.2 section had no entries for the contributed pull requests carried into this release, so neither their
changes nor their authors reached the release notes. Add an entry for each, written from its commits: Flash Next's
CUDA prompt tiles and structured output (#1046, #1062), --thinking-reserve, the long-loop guard and the --vision
template diagnostic (#1064, #1075, #1076), GLM's Metal eh_proj tile and CUDA waiting ranks (#1066, #1053), M1's
simdgroup cap (#1055), Nemotron's CUDA image and shared-window fixes (#1057, #1058), Qwen's live-row expert routing
(#1047), the shared CUDA operators and contract fixture (#1056, #1068, #1070, #1071, #1073), the ROCm receipt,
drafter and draft-count work (#1043, #1044, #1045), the ModelOpt recipe numbers (#1042), the Windows host build
(#1034) and the session store's allocation-free rollback (#637). Thank #643's and #1060's reporters as the other
entries thank theirs.
@ashhart

ashhart commented Oct 21, 2026

Copy link
Copy Markdown
Owner

The routed-expert padding fix is included in 1.1.2 (https://github.com/ashhart/TensorFold/releases/tag/v1.1.2), with your original authorship retained. Captured bucket padding rows no longer enter the router, top-k or combine work. Closing this as carried, and thank you.

It landed as b2cb75d and 2d5894b.

- Synthema, Ash's AI Assistant.

@ashhart ashhart closed this Oct 21, 2026
@jschmied
jschmied deleted the q35-moe-live-rows branch October 23, 2026 21:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants