Repository navigation
nemotron CUDA: the router top-k is moe_glue's sigmoid top-k (#548, step 2) - #1354
Conversation
|
Review of #1353 and #1354 (#548 steps 1 and 2) against main e9b69ed. Moving Nemotron's norms and router onto the shared glue is a good direction, and nothing here is lossy. But both steps change Nemotron's output bytes, and the float64 receipts get slightly worse rather than better:
A refactor that changes bytes needs its error to be no worse than before, and we can't take a prefill slowdown. Two ways forward: keep Nemotron's current summation order inside the shared kernels (a variant that reproduces today's bytes), or show per-projection-class worst-column float64 error, prompt and decode separately, that is no worse. Also, dropping the |
|
Thanks for moving the slice sum into the top-k and for restoring the guard with a test. We checked the head against our 26 Oct review: the extra launch is gone ( Three asks:
Once decode is level we'll A/B Lightning on a GB10 and run Bolt's token checks. |
458dd2d to
9d142df
Compare
|
Thanks. All three are done. #1354 is rebased onto 1. Rebase.
2. Decode. You were right about the loads. Against The last step-2 commit (
At Bolt's shapes (top 9, 8 slices), 1-16 rows: 5.0-5.7 µs at 240 experts (main 6.3), 5.3-6.7 µs at 496 (main 7.1). The change is bit-equal. Picks and weights are byte-equal to the loop form over 6,000 calls: 1-300 rows, 33-512 experts, top 1-16, 1-8 slices, with ties, NaN, ±inf, -0 and all-underflow rows, with and without bias. On Lightning, the kl dump sha ( Records, re-derived:
SASS against Served A/B, Lightning MLX 4-bit.
Plain decode is faster at 1 and 4 streams, and cold TTFT is level or faster. On 3. Description. It is updated: no combine against a row of ones, Bolt: its router and latent sum move here, and the receipts so far were Lightning's. I ran its token checks on
Other checks on
|
|
Thanks @jschmied. The three asks are answered:
We'll record the score kernels' nvcc 13.3 limits on our side. Two things before we run it:
Then, on a GB10:
The |
9d142df to
e6e217b
Compare
|
#1354 is restacked onto #1353's Top-k. Replayed on
Tokens (#1354 merged into
Served A/B against
Checks:
|
step 2) Nemotron's tf_nemo_topk/_wide become moe_glue.topk's warp form through a new tf_moe_topk_sigmoid_score_warp (ranks sigmoid(logit) + bias, any number of shared slots); tf_moe_topk_sigmoid_warp is unchanged. The K-sliced logits are summed by moe_glue's combine against a row of ones (byte-equal), and tf_latent_sum becomes moe_glue.combine (byte-equal). Nemotron's bytes change once: expf and w/sum replace tf_exp and w/(sum+1e-20). Other fatbins' SASS is unchanged.
…shhart#548 step 2) tf_nemo_topk normalized over the picks' sum + 1e-20; the shared score top-k dropped it, so a row whose picked sigmoids all underflow to 0 in fp32 got 0 / 0 = NaN weights. tf_moe_topk_sigmoid_score_warp adds it back, and the float64 references (glue_ref.topk, the contract's score check) divide by the same sum + 1e-20. A new glue case routes three rows of logits -100: the ids match float64's, the routed weights are 0 and the shared slots 1.
…hart#548 step 2) Step 2 summed the router's K-slice partials with a combine launch (fp32 weights 1.0 from Buffers.ones) before the top-k, one launch per MoE layer more than main's tf_nemo_topk. tf_moe_topk_sigmoid_score_warp now takes `slices` and adds the partials [slices, rows, experts] in slice order before the sigmoid, the same chain as the combine's fma(1, p, acc): the bytes stay, the launch and Buffers.ones go. Route.slices (default 1) carries it; only the score form takes it. The glue test checks the partials unsummed and routes the all-underflow rows as two slices of -50. Records: score_warp 64/64/64/112/112 registers on 13.0, 64/64/64/128/128 on 13.2.
… step 2) lean_check flagged the new multi-line and over-long comments; the text is shorter, the code unchanged.
…ums (ashhart#548 step 2) Plain decode was 0.9-1.3% slower than main at 1 and 4 streams. The score top-k added the router's K slices in a runtime loop a lane's expert at a time, so each lane waited on up to 16 x 6 dependent loads; main's tf_nemo_topk issues every slice's load first. score_warp now does the same (part[P][8], the slices summed in the same order afterwards), and a second instance, tf_moe_topk_sigmoid_score_warp_128, holds 4 experts a lane for up to 128 experts (Lightning) so its loads fit in 62 registers or fewer; past 128 the 16-a-lane form runs (210-219 registers, no spills). The host refuses more than 8 slices. Bits are unchanged: picks and weights are byte-equal to the loop form over 13,000 calls (1-300 rows, 33-512 experts, top 1-16, 1-8 slices, ties, NaN, +-inf, -0, all-underflow rows, with and without bias). At Lightning's shape (128 experts, top 6 + 2 shared, 6 slices) a call takes 2.0-2.3 us at 1-16 rows (main's tf_nemo_topk 3.4, the loop form 5.5-5.8); at 2,048 and 8,192 rows 6.0 and 50 us (main 12.5 and 86).
…qual main's (ashhart#548 step 2) Replayed on main's own router outputs, the score top-k picked the same experts but 1-ulp different weights on up to 14 % of the prompt's slots: its sigmoid took expf, tf_nemo_topk's takes glue_math's tf_sigmoid. score_warp now calls tf_glue::tf_sigmoid, and a NaN score plus bias ranks as -inf, as tf_nemo_topk ranks it. The glue test checks the weights against main's host reference bit for bit again (the float64 twin and Rig.near go), the underflow case too.
Main moved glue_math.cuh to zig/kernels/cuda/common/ after this branch added its include; every other unit includes it from there. Co-authored-by: Jürgen Schmied <600316+jschmied@users.noreply.github.com>
e6e217b to
0066773
Compare
Stacked on #1353 (step 1); its six commits come first here.
What moves
Step 2 of #548. Nemotron's router top-k moves onto
moe_glue.tf_nemo_topkandtf_nemo_topk_wide(nemotron_route.cu) are replaced bymoe_glue.topk's warp form with two new entry points:tf_moe_topk_sigmoid_score_warp_128(up to 128 experts, 4 a lane: Lightning) andtf_moe_topk_sigmoid_score_warp(up to 512, 16 a lane: Bolt, Super). Both rank by sigmoid(logit) + bias, as NemotronH does. The existingtf_moe_topk_sigmoid_warp(Kolibri, Flash-Next) ranks by logit + bias and keeps its SASS.slicesup to 8, the host refuses more). Every slice's load and the bias's are issued before the sums, astf_nemo_topkdid, so no launch sums the slices first.e..e+shared-1. The routed weights arew / (sum + 1e-20)× scale: picks that all underflow to 0 weigh 0, not NaN.moe_glue.Routegainsscore,shared_slotsandslices.tf_latent_sum(the latent experts' sum, Bolt and Super) becomesmoe_glue.combinewith bf16 output. This is byte-equal: the fma chain is the same and fma is commutative in its multiplicands.glue_ref.topkis a float64 twin of the rule. The glue test checks the ids exactly and the weights to within 2^-20 relative.routeUnderflowroutes three rows whose logits are two slices of -50 each: the ids match float64's, the routed weights are 0 and the shared slots are 1. Without the+ 1e-20it goes red.Bytes: this step changes Nemotron's bytes. The old kernel used
tf_exp; the shared kernels useexpf. The picks are identical except at near-ties, but the weights move by a few ulp. Other families are untouched. On sm_80/86/89/120/121 (nvcc 13.0), againstmain:moe_glue.fatbin's existing functions keep their SASS and the two score kernels are added;norms.fatbingains step 1's five functions with the rest unchanged;nemotron_routeandnemotron_opsonly lose functions;contract_controlsgains step 1's control; every other fatbin is byte-identical.Kernel contract
_128, and Bolt (496 experts, top 9), which runs the 16-a-lane kernel, with the tiled edges of 8 to 33 experts on_128.tf_moe_topk_sigmoid_score_warp219/219/219/210/210 regs,_12862/62/62/48/48, static 0, bound 256, no spills. nvcc 13.2 compiles_128the same and the 16-a-lane kernel at 219/219/219/203/203, which is in its 13.2 set. Thetf_nemo_topkand_widerecords are removed, the 13.3 one fortf_nemo_topk_wideincluded. nvcc 13.3 is not installed here, so its compile of the score kernels is unchecked.kernel-contracton the touched kinds: both score routes, the 4 combine kinds, step 1's 4 moe norm kinds pass gates 1-4, and both controls are red as planted.kernel-contract-limitsat 13.0 and 13.2: every record of this PR is current; the failing and stale lines are other families' (qwen35, qwen4, GLM, ...), none in an image this PR touches.Nemotron's gates (Lightning MLX 4-bit, GB10)
Distance to the float64 CPU reference, over the same 3 passages × 320 rows as step 1 (kl dump sha
19475174fbfe, unchanged by the load hoist):Paired over the 960 rows, rel-L2 moves by +0.00063 ± 0.00096 (2 se). Argmax misses go from 41 to 48: 25 rows miss only here and 18 only on
main, McNemar p = 0.36. Both are within noise. The teacher-forced loss over seven public texts (3,584 tokens) is +0.0024 ± 0.0047 againstmain.Bolt
Bolt's router and latent sum move here too (496- to 240-expert layers on the 16-a-lane kernel, its latent sum on
moe_glue.combine). The receipts above are Lightning's; Bolt was token-checked separately onNVIDIA-Nemotron-3.5-Bolt-74B-A9B-24GB(IQ2_XXS experts, NVFP4 Mamba) on a GB10, at458dd2d8c; the group-norm commit since is bit-equal. Bolt has no MTP head, so the drafted lanes, prompt cuts, resume, paged and segments checks do not apply; every check that runs without drafts does:mainmain(Soft rain on red leaves, …)The glue test also runs the router, top-k and latent sum at Bolt's shapes against the host references.
Speed
The score top-k alone, at Lightning's shape (128 experts, top 6 + 2 shared, normalized ×2.5, 6 slices), a CUDA graph of back-to-back launches on a GB10, µs a call:
tf_nemo_topk_128)At Bolt's shapes (top 9, 8 slices) the 16-a-lane kernel takes 5.0-5.7 µs at 1-16 rows for 240 experts (main 6.3) and 5.3-6.7 µs for 496 (main 7.1). The loop form and the new one are byte-equal over 6,000 calls (1-300 rows, 33-512 experts, top 1-16, 1-8 slices, ties, NaN, ±inf, -0, all-underflow rows, with and without bias).
Served A/B on a GB10, Lightning MLX 4-bit,
main41abcf402against9d142df90, 12 rounds. Each arm and round gets a fresh server, in alternating order. Cold TTFT is the median of 3 a round with drafts on; plain decode is a fresh--no-draftsserver, 256 tokens, 3 reps. The table gives the median of the paired per-round ratios:Drafted decode moves with acceptance, since the bytes differ from
main's, so it is not a speed comparison here.On the previous head (
458dd2d8c, 14 rounds) cold TTFT was 0.1-0.3% slower; an nsys trace put that on step 1's Mamba group norm. Step 1 now runs it four groups a block (#1353's8a77e3ec0, bit-equal), and TTFT is level or faster at 2k, 8k and 32k.Receipt
9d142df90onorigin/main41abcf402: nemotron CUDA: RMSNorm rows run on the shared norms and combine (#548, step 1) #1353's six commits, then step 2's six. GB10 (sm_121), CUDA 13.0, Zig 0.17.0.zig build test -Dnvcc=/usr/local/cuda/bin/nvcc: exit 0. lean_check: 0 new.