Repository navigation
gemma4, nemotron cuda: a prompt chunk nothing reads stops at its last state write - #2432
Conversation
7f9f50d to
eb2f6f1
Compare
|
Thanks, @jschmied. This is a tidy extension of #2229's tail, and the exactness evidence is thorough. Three asks. First, the CHANGELOG line: your table shows Nemotron level at 2K and Gemma 4 at 1.3%, so please say "up to 4% … level at 2K" rather than "2-4%". End it "(#2432, refs #2165). Thanks to @jschmied." and put it under |
eb2f6f1 to
26da2ea
Compare
|
Thanks, all three are done on the new head, 26da2ea, rebased onto main at c16a296. The only conflict was in CHANGELOG.md, and no code changed.
|
… state write A non-final prompt chunk's last layers feed only the head, and its logits are not read. Gemma 4 stops its last layer after the keys and values, Nemotron after the last Mamba scan. core/prompt_tail.zig decides for both; taps, a head that absorbs rows, prompt logprobs, dumps and lessons keep the whole chunk. No new kernels.
26da2ea to
681b8a6
Compare
|
Thank you @jschmied. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes. |
|
@jschmied, a follow-up from our GB10 with Gemma 4 31B (4-bit) on this change. It ran against its merge-base in the order base, new, new, base, after a warm-up run:
That's a little better than the 1.3% your receipt measured at 2K. |
|
Thanks @jschmied! This is in TensorFold 1.2.5, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.5 |
What this changes
Before: every prompt chunk ran its last layers over all its rows, although only the prompt's last chunk has its
logits read. After: on Gemma 4 and Nemotron (CUDA), a chunk whose rows nothing reads stops at its last state write.
Gemma 4 skips its last layer's q, attention, o and MLP and keeps the key and value writes; Nemotron skips everything
after the last Mamba scan, including that chunk's unused head and draw. No row is read, so unlike #2229 no tile is kept
and no kernel is added. The rule is
prompt_tail.unreadinzig/src/core/prompt_tail.zig; each family only reports itsreaders. A chunk runs whole when anything else reads it: Gemma 4's tapped last layer; Nemotron's absorbing head, taps,
prompt logprobs, dumps and lessons. The prompt's last chunk is unchanged, so a one-chunk prompt launches the same kernels.
atc overlap --jsonagainst origin/main: no entry for this branch, 0 conflicts; the result is incomplete(file and time limits). Rebased onto main c16a296; the only conflict was CHANGELOG.md.
Checks on this head
zigdiff); this head adds thereceipt, the CHANGELOG line and the ledger. The drafted runs ran on this head's code (39e33dd33, same tree outside
docs/).python3 tools/zig/lean_check.py, on the base and on this head: 5,794 files on each; 116 on the base, 116 on this head, 0 new (inherited).zig build -Dfixed-cpu test --test-timeout 10m: 589 of 589 steps, 7,418 of 7,476 tests passed, 58 skipped, 0 failed.test-kernel-contract210 of 210,test-golden322 of 335 equal, 0 unexpected,test-server-security test-pull-security29 of 29,
fuzz-cortheon20 of 20.zig fmt --check build.zig zigandgit diff --checkpass.CUDA SMs; Metal and HIP (the change is CUDA only).
Receipt
Environment: this head and base above, Zig 0.17.0, CUDA 13.0 (V13.0.88), driver 615.71.09, NVIDIA GB10 (sm_121),
Linux 7.0, 128 GB unified memory. Checkpoints: mlx-community/gemma-4-31b-it-4bit; Nemotron 3.5 Lightning MLX 4-bit,
revision d9d758fb83953437f7263256b0d96157e2a348b8. Drafts off,
--parallel 1for timing.Paths and platforms the change reaches: the CUDA serving prompt pass of both families (
gemma4/cuda_forward.ziglaunch,nemotron/cuda_prompt.zigserialChunks). Nemotron's segmented pass and the training walks are notchanged. Run on GB10 only.
Exactness:
tools/bench_concurrent.py --tokens 256 --reps 3 --temperatures 0 --levels 1,4 --alone --serialon thishead, both families: every cell alone and serial unequal 0 (30 equal per family). Base, new, new, base, against main:
replies, token hashes, seeded sampling (temperature 0.8) and two prompt cache turns equal. Nemotron: top-20 logprobs
and prompt_logprobs (8,794 rows) equal. Gemma 4 on CUDA has no logprobs; its last-row logits hash, through a local
hook that is not in this pull request, was equal on every prompt: 15 (14 distinct) and 9 (8 distinct) in two runs.
Main against this branch, same machine and checkpoint: base, new, new, base, page cache dropped before each server.
Cold TTFT, median of 3 prompts an arm,
--prompt-cache-gib 0:Nemotron at 2K: the second three-prompt run read +0.52 %, past the +0.5 % limit set for level; a 12-prompt run per
arm, limit fixed before it, read -0.35 % with equal hashes. Nemotron at 32K on this head: -2.31 %, 0.01 pp past the
preset fast bound of -2.3 %; the three runs read -2.11, -2.29 and -2.31 %.
Decode speed before and after, 128 tokens, same arms: Gemma 4 -0.34 % and -0.12 %; Nemotron +0.17 %, -0.13 % and
-0.46 %, as measured (this head's Nemotron arms: base 108.1 and 107.7, new 107.6 and 107.3 tokens/s).
Prompt speed before and after: the cold TTFT table above. The tools and raw runs are in
docs/receipts/prompt-tail-gemma-nemo/; its README maps each cell to its run.Drafted, main against this head, the table's nine prompts, 64 tokens, temperature 0, drafts accepted on every request:
replies and token hashes equal 9 of 9 on each family. Nemotron (MTP head) is level, as its head absorbs the rows.
Gemma 4's DFlash drafter taps layers 1-57 of 60, never the last, so its chunks stay unread: 32K TTFT -3.8 to -4.0 %
with drafts on (single arms).
Licence of any new weights or drafter: n/a, none.
Checklist
the new
unread, which did not exist before; it needs no device."draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.Concurrent equals solo and serial (above); prompt cache turns equal main; drafted equals main, 9 of 9 a family
(above). Drafted against
"draft": falseand resumed against fresh were not run.--drafter, noslower than without it, greedy and sampled. Drafter runs were not done; with a drafter that reads the rows the chunks
run whole by the rule above.