Repository navigation
Conversation
…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.
|
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.
|
Thanks. Pushed 13a296e on top of a16a1d4 (no rebase, still on v1.1.0). Wide plan. Tests.
The served contract also checks Gates on 13a296e:
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 (
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 Prompt reuse and concurrency (
The script's own verdict line fails the same way on main and the branch ( ABBA served speed, tok/s, mean of the two starts per build (decode median a request for 1 stream, aggregate at the top
Two cells are slightly slower on the branch:
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,
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. |
…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.
|
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. |
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.zigforwardpicksbucket = 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.zigrows()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_topkandtf_moe_combinetake the forward's params and skip rows fromP[2]up. A null pointer keeps every row; the MTP head passes null, since it runs uncaptured at its exact row count.grouped.Router.routeRowsplans only the firstP[2]rows' pairs at any width:plan_rows(experts.cu):plan_kernelover the first*rows * perpairs;plan_rank_rowsandplan_scatter_rows:plan_rankandplan_scatterover the same live pairs, withplan_offsetsshared.plan_rankandplan_rank_rowsshare one device body.route,plan_kernel,plan_rank,plan_offsetsandplan_scatterare 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*_rowskernels. 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 ofmoe): 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: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:
Results:
FAIL bucket 128 at 1 live rows planned 1152 pairs (largest member 1151), want the live rows' 9.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.P[2]below the row count at rows 1, 15, 16, 17, 64 and 128: live rows meet float64, padding rows' outputs stay unwritten.routeRowsis checked againstrouteover 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_rowsandplan_scatter_rowsare registered inqwen_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-limitsover fatbins for sm_80, 86, 89, 120 and 121: every registered launch fits, every control red.Tests run on 13a296e:
zig build testfrom 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 --checkis clean.Engine output and timing,
tf-qwen35-cudaover the same 300 receipt ids, main against 13a296e, ABAB over four rounds (35B) and two (27B):step, one row alonestepbeside a second stream (2 rows, bucket 16)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:
moe-liveand the contract above.tf_moe_*and the*_rowsplan kernels are not re-derived.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
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.SUMMARYline is in the receipt.