Repository navigation
nemotron: Super's MTP drafts on the draft vocabulary's FP8 rows while a stream's text is covered, same replies (stacked on #1006) - #1011
Merged
Conversation
…d of its own, tf_super_topk's bits tf_super_topk takes a row's picks one argmax at a time on one warp. tf_super_topk_few, which calls of up to launch.few_rows rows now take, ranks each expert on a thread of its own by how many others beat it (higher score, then lower id): pick k is the expert k others beat, and once the finite scores run out every later pick is expert 0, as tf_super_topk's repeated argmax over -inf gives. Scores, sums and weights are the same operations in the same order. Prompt chunks keep tf_super_topk. glue_ref.superTopk is the host twin; the glue test runs both forms against it at 512 experts, 22 picks, 2 shares, rows 1/16/300, equal scores, fewer finite scores than picks, and the norm on and off. The glue tests' Rig moves to glue_rig.zig, which keeps glue_tests.zig under 600 lines.
…mns a warp, four steps in flight, same bits ck_experts.cu's grouped kernel is a template over a warp's n8 tiles and the steps of codes, scales and weights it keeps in registers ahead of the mma. Prompt chunks keep 8 tiles and one step; calls of up to launch.few_rows token rows take the _few pair, 4 tiles and 4 steps. Each column still takes its 64-input steps in order through the same mma_fp4, and the epilogue is unchanged. Ck.down takes the call's token rows to choose. The modelopt test checks up and down at 1 and 16 token rows byte-equal to the same pairs inside a 40-row call.
…-equal to the kernels they replace tf_super_topk_few replaces tf_super_topk at decode-sized calls: gate 1 holds it to tf_super_topk's bytes and to glue_ref.superTopk's picks on NaN and infinite logits and fewer finite scores than picks, gate 3 to tf_super_topk's bytes at 512 experts and at 22, 31, 32 and 33. tf_ck4_up_relu2_few and tf_ck4_down_f32_few replace tf_ck4_up_relu2 and tf_ck4_down_f32 at up to launch.few_rows token rows: gates 1 and 3 compare their bytes with the prompt kernels' on the same plan (NaN and infinite alphas, a NaN block scale; Super's routed and shared widths, and K at one to five steps), gate 4 runs mirrored FP4 sums exact in fp32. Ck.run and Ck.warp_columns are public for the test to launch either form. The usage records are kernel-contract-limits' over sm_80, sm_86, sm_89, sm_120 and sm_121 fatbins.
tf_nemo_scan loaded every step's inputs into shared memory at once, so a call took at most 32 kept and window rows (windows of 16). It now loads and steps them 32 at a time; each step's arithmetic and order are unchanged, so a call of up to 32 steps runs as before and longer calls give serial's bits. The glue test runs windows of 32 rows over 31 kept ones (51 and 63 steps), alone and as a stream of a shared window.
…bytes tf_nemo_scan registers with the kernel contract under a stated rule: gates 1 and 3 hold its y, the window's dt and the state after the kept rows to glue_ref.scan's bytes (NaN wherever the twin gives NaN), at Super's 128 heads with 1 to 32 window rows after 0 to 31 kept ones, so calls either side of a 32-step stage and up to 63 steps. Gate 1 plants NaN in an x, +Inf in a gate and in a dt input, and -Inf in a kept state. Each run restores the state and dt planes the kernel writes, so the five repeats read the same inputs. The usage record is kernel-contract-limits' over sm_80, sm_86, sm_89, sm_120 and sm_121 fatbins.
…fts a round) state.max_rows is 32 and state.windowRows(c) the width a checkpoint takes: 32 for a latent MoE (ModelOpt's FP8 lane tiles of 16 and 32 rows, bf16 and NVFP4 lanes, and per-pair experts give a row the same bits at any width), 16 for every other one, whose MLX affine-4 decode products (Nano's ModelOpt MTP head among them) tile 16 rows. Window buffers, the Mamba replay planes, captured windows and head levels, the cost table, the depth rule, copied chains, grammar windows, the lane facts and the pinned layouts follow it; the latent shared-expert plans cover max_rows, the non-latent ones stay at 16 rows. --mtp-drafts takes up to max_chain (31), and a lone stream's cap stops at its checkpoint's window.
…e NVFP4 experts and the latent top-k Ck chose the _few expert kernels, and the latent router the thread-an-expert top-k, for calls of up to launch.few_rows (16) token rows, the widest window when they came; Super's windows of 17 to 32 rows fell to the prompt kernels. Expert calls of up to Ck.few_tokens (32) token rows and top-k calls of up to state.max_rows rows now take the decode forms. Each form gives its prompt kernel's bits, so only time changes. The modelopt test checks 32 token rows byte-equal to the same pairs inside a 40-row call, the glue test runs both top-k forms at 32 rows against the host twin, and the contract holds the three decode forms to their prompt kernels' bytes at 32 rows.
… and grow with what copies keep A copied chain is verified whole, without the depth rule's pricing, and with 32-row windows it could take up to 31 tokens of the context's continuation, so every wrong copy cost a 32-row window. A copy now proposes at least 15 tokens (a 16-row window's), and after each copy round at most twice the tokens that round kept, up to the window's 31. The MTP head's drafts keep the depth rule.
…egime Costs.measured pooled every row's added ms so it never rises with rows, as shared experts make it. Super's windows past launch.few_rows (16) rows switch to wider kernels (the router's logits, the lanes' 32-row tiles), so row 17 adds more ms than row 16 did, and pooling spread that step back over rows 3 to 16: the depth rule then priced mid-size windows above their timed cost and drafted shallower. Costs.measured now pools within regimes of a given width; the row that starts a regime keeps its own ms. A 16-width table is unchanged. The new draft_depth test runs a 32-width series with a step at 17 rows.
A window's ms follows its tokens, through the experts its rows route to. measure prefilled windowRows tokens of its text before the windows it timed, so a 32-row engine timed its 1-16 row windows on later text than a 16-row engine, where the same kernels take longer. The windows now follow the text's first 16 tokens at any width: rows 1-16 are timed on the same tokens as without the wider windows, and the depth rule prices them the same.
tf_nemo_scan read the state into registers before its first stage's inputs, so the 32 state registers were live through those loads and a window of 16 rows or fewer, one stage, ran slower than the unstaged scan it replaces. The first stage now loads its inputs, then the state, then steps, as the unstaged scan did; later stages follow in a loop. Each step's arithmetic and order are unchanged. The contract's usage record follows the new register counts.
…s FP8 output rows The draft head was the target's own, so every chain level read the whole FP8 output head. It now packs the checkpoint's FP8 lm_head rows for the embedded draft vocabulary into a lane weight of their own and uploads their vocabulary ids; a level proposes over that subset and the target still verifies every proposal with its full head, so replies are unchanged.
…ntext tokens are draft ids, else on the whole head The draft vocabulary's head proposes only its 49,152 ids, so on text outside them a chain stops at the first uncovered token. Each round a stream now takes its head from its context's tail: the subset while the last 16 tokens (the prompt's, until 16 are emitted) are all draft ids, else the target's whole head, which stays resident, so nothing is loaded twice. The choice is the stream's own (Rounds.vocab, mtp.Vocab) and goes with every launch, chain level and captured graph; batched chains project and sample each stream with its own head, and shared rounds are priced at the whole head's costs. cuda_costs.measure times the windows once and each head's levels between the same widths, and caches each head's table under its own key, so each head's depth rule prices its rounds by its own level.
Owner
|
Thank you, @plotarmordev, for the full/subset FP8 draft heads and their kernel controls. We merged the work for 1.1.0 in 4aa3f91. Our fixes publish a head only after construction succeeds, drain before failed cleanup and validate gathered rows before reading or allocating. Integrated fixtures follow real restore, binding, chains and draws through mixed and shrinking batches. The full head remains available when the context lacks subset coverage. GB10 Super VL loaded and passed 17/17, concurrent 2/4 and chat 12/12 after repair; full/subset FP8-head device coverage remains unrecorded. Posted by Ash's AI Assistant Synthena |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Branch
nemotron/draft-subset-head. Stacked on #1006 (nemotron/verify-32-rowsat fa2af58, itself on #999 and#997); review only the last 2 commits.
What this changes
Each MTP chain level ends in an output projection. Super's draft head was the target's own FP8 head, so every level
read all 131,072 rows. TensorFold embeds a 49,152-id draft vocabulary for Super; a level that reads only those rows
costs less, but it can propose only those ids.
cuda_modelopt.fp8DraftHeadpackslm_head's FP8 rows for the draft ids into a lane weight of their own (thehead's scales) and uploads their vocabulary ids.
cuda_weights.gatherRowsreads a tensor once and gathers rows byid;
loadMtpand the NVFP4draftHeadnow use it too.the prompt's tail until 16 tokens are emitted (
draft_ids.Coverage). Otherwise it drafts on the target's wholehead, which stays resident. The choice is the stream's own:
Rounds.vocab/mtp.Vocabis passed to everylaunch, chain level and captured graph, not stored on the engine-wide head. Batched chains project and sampleeach stream with its own head. A shared round's lane costs use the whole head's level.
cuda_costs.measuretimes the windows once and each head's levels between the same widths. It caches each head'stable under its own key (the key's draft width) and returns both (
Measured). Each head's depth rule pricesrounds by its own level.
Both commits belong in this pull request. With only the first one, every stream drafts on the subset, and on text
outside the draft ids (Spanish, Japanese, and so on) a chain stops at the first token the subset cannot propose.
Receipt
Raw receipts:
receipts/nemotron/draft-subset-head(build.txtmaps every file to its binary).Environment: this branch at
f58dc04b4a98ed2e5b89c17bb5966ef58a1df4c7against its basefa2af585fc52b888bbba09e379367b8e0792115b(nemotron: Super verifies windows of up to 32 rows on CUDA, same replies (stacked on #999, #997) #1006's tip). Both were built from clean checkouts withzig build fatbins -Dsm=121 -Dnvcc=…,zig build nativeandzig build install(
-Doptimize=fast -Dstrip=true -Dtarget=aarch64-linux-gnu.2.28 -Dcpu=baseline -Dversion=1.0.12, zig 0.17.0).tensorfold-nativesha256: base097904fb69e20ae74986e50ab438db6abec16a68aaa0ce74d51ef4656513179b(the binarynemotron: Super verifies windows of up to 32 rows on CUDA, same replies (stacked on #999, #997) #1006's receipts ran), branch
3eefef75610ce5b42262f9c4ad5a347c8d36f04ba62022811edf3afe782a2596; branchtf-cuda-teste6b6ae3535a660af76f552879af6e841812b2fa4ad07e62b1c0b1a8426842448. A DGX Spark (GB10) and a secondDGX Spark, CUDA 13.0 (nvcc 13.0.r13.0), driver 580.173.02.
nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-NVFP4@82d93699acd3be9e4617c53360fe5b3f5dbeb1b3, served astensorfold-native serve <checkpoint> --name super --no-update-check --prompt-cache-gib 0on loopback withTF_HEAT_HIGH=85 TF_HEAT_LOW=75, everything else default (MTP drafts on). On each Spark, six holds, each a freshserver, in the order base, branch, branch, base, branch, base:
core/on the first Spark andlang/on the second.Each build's window-cost cache files (both heads' keys) were deleted before its first hold on each Spark, so its
tables are measured fresh; the later holds read them. Every receipt JSON carries
commit,binary_sha256andbuild.Exactness, every hold: drafted ==
"draft": falsefor code, chat, count, tool_json and repetition (core/) and forthe English, Spanish, French, German, Italian and Japanese essay prompts and the multilingual translation prompt
(
lang/). Each timed prompt's warmup and repetitions give the same hash.bench_concurrent --alone --serial: 30of 30 alone and 4 of 4 serial. The English, German, Japanese and multilingual prompts run four at once, twice a
hold, equal their solo replies: 24 of 24 streams on each build. All 12 prompts' token hashes are equal on both
builds (code
51e4c630298a14a1…, chat161034fd0a090a34…, count1b3d78d3b1bfd175…, Japanese07f8403ff9649a39…; all incore/summary.mdandlang/summary.md).Seeded sampling (temperature 1), one more hold of each build on the second Spark (
sampled/):bench_concurrent --temperatures 1.0 --alone --serial20 of 20 alone and 16 of 16 serial. Each of the sevenlanguage prompts drafted ==
"draft": false. The English, German, Japanese and multilingual prompts run four atonce equal their solo replies. All seven hashes are equal on both builds.
GPU tests (no kernel changes):
kernel-contractSUMMARY 16 kernels: 16 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 0 limit problems;modeloptPASS (its FP8 lane repack equals the host packer, and FP8 lane rows 1-16 are byte-equal).Failing before: the new host test (
draft coverage reads only the context's last tokens) fails on the base, whichhas no
draft_ids.Coverage(host/before.diff,host/before.test-cuda-host.log).Draft level cost from each fresh table (ms), on both Sparks:
Window ms are shared by both heads. First Spark, base / branch: 46.49 / 46.37 at 1 row, 86.07 / 86.07 at 8,
117.74 / 118.71 at 16, 122.16 / 122.23 at 17 and 164.91 / 164.49 at 32.
Decode speed on the first Spark, medians of 3 a hold, tok/s (every run in
core/summary.md):fibonacci-raw, 256)gpu-chat-no-think, 256)Count drafts deeper on the cheaper level: 13.33 → 13.79 tokens a round (acceptance 0.963 → 0.976). Repetition
takes 8.53 → 8.26 tokens a round (acceptance 0.527 → 0.502) and is still faster.
bench_concurrentaggregate tok/s,median of the three holds, base → branch: code x1 46.6 → 50.5, code x4 135.8 → 141.2, chat x1 35.9 → 38.6, chat x4
97.7 → 100.8.
Languages on the second Spark, medians of 3 a hold, tok/s (every run in
lang/summary.md). No language is belowthe base, and every prompt takes the same tokens a round on both builds:
Acceptance is unchanged on French, German and Italian and moves on English (0.633 → 0.623), Spanish
(0.660 → 0.668), Japanese (0.786 → 0.789) and the multilingual prompt (0.726 → 0.712): the branch drafts the rounds
whose last 16 tokens are draft ids on the subset, whose depth rule prices a level lower.
Prompt speed:
bench_cache.py --require-cold --tokens 16 --repeats 2 --streams 1on an 8,186-token chat prompt.Cold
ttft_sof the rows with no heat wait logged in the request: base 4.471, 4.468, 4.455, 4.477, 4.459; branch4.472, 4.459, 4.461, 4.452. The rows that waited net 4.45-4.47 s after the logged wait, except 2.7 s where the gate's
wait overlapped GPU work (one row on each of two branch holds) and one branch row that nets 4.549 (12.549 - 8.0).
This branch does not change the prompt path.
Memory and start: model 71.03 → 71.22 GiB (the subset's 49,152 FP8 rows of 4,096), 0.40 GiB a stream on both.
A first start that measures the tables takes 62 s against 66 s on the base (60 s against 64 s on the second
Spark). Later starts: 20 and 24 s against 22 and 22 s (18 and 18 s against 16 and 16 s); the branch captures each
head's level graphs.
Tests run:
zig build test1381 of 1388 passed, 7 skipped, 0 failed (the base: 1380 of 1387; the new test is theone above);
zig build test-kernel-contract18 of 18;zig build test-golden309 of 319 equal, 0 unexpecteddifferences;
zig build fuzz-cortheon7 of 7;zig fmt --check zig build.zigclean;tools/zig/lean_check.py517 problems, the same set as the base, 0 new.
Not run, and why: Metal (no Mac here; the change is CUDA-only). TP=2 across two GB10s: not run. Checkpoints without
Super's FP8
lm_head(Nano and the NVFP4 draft head) keep their single draft head with no coverage set, so theydraft as before; not run on a GPU (no weights here). Resumed against fresh prompts: not exercised separately.
Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (Drafted vs"draft": falseand concurrent vs alone, including English and non-English prompts at once, measured above; resume not exercised separately.)SUMMARYline is in the receipt. (No kernel is new or changed; the CUDASUMMARYline is above.)ttft_sabove.)