Skip to content

gemma4, nemotron cuda: a prompt chunk nothing reads stops at its last state write - #2432

Merged
ashhart merged 1 commit into
ashhart:mainfrom
jschmied:prompt-tail-gemma-nemo
Nov 7, 2026
Merged

ashhart merged 1 commit into
ashhart:mainfrom
jschmied:prompt-tail-gemma-nemo

Conversation

@jschmied

@jschmied jschmied commented Nov 7, 2026 •

Copy link
Copy Markdown
Contributor

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.unread in zig/src/core/prompt_tail.zig; each family only reports its
readers. 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.

Checks on this head

  • Head commit: 26da2ea, on c16a296. The TTFT gates ran on 7f9f50d and earlier (same zig diff); this head adds the
    receipt, the CHANGELOG line and the ledger. The drafted runs ran on this head's code (39e33dd33, same tree outside docs/).
  • Lean checker, 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).
  • Tests run, with the command and pass, fail and skip counts: 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-contract 210 of 210, test-golden 322 of 335 equal, 0 unexpected, test-server-security test-pull-security
    29 of 29, fuzz-cortheon 20 of 20. zig fmt --check build.zig zig and git diff --check pass.
  • CI: not run yet; this head is pushed to the fork, the pull request is not opened.
  • Not run, and why: 64K prompts (servers ran at a 40,960-token window); sampled concurrent runs; other
    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 1 for timing.

  • Paths and platforms the change reaches: the CUDA serving prompt pass of both families (gemma4/cuda_forward.zig
    launch, nemotron/cuda_prompt.zig serialChunks). Nemotron's segmented pass and the training walks are not
    changed. Run on GB10 only.

  • Exactness: tools/bench_concurrent.py --tokens 256 --reps 3 --temperatures 0 --levels 1,4 --alone --serial on this
    head, 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:

    2K 16K 32K
    Gemma 4 31B, earlier base -1.32 % -3.09 % -3.83 %
    Gemma 4 31B, this head -3.88 %
    Nemotron 3.5 Lightning, earlier base +0.00 %, +0.52 % -1.90 %, -2.20 % -2.11 %, -2.29 %
    Nemotron 3.5 Lightning, this head -2.31 %

    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 title and description match the head commit, including the receipt and the stack.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. n/a: the test covers
    the new unread, which did not exist before; it needs no device.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths in the diff or its commits. No AI attribution lines.
  • 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.
    Concurrent equals solo and serial (above); prompt cache turns equal main; drafted equals main, 9 of 9 a family
    (above). Drafted against "draft": false and resumed against fresh were not run.
  • No precision traded for speed. If the bits change, the description says which sums change and why.
  • Decode and prompt processing are no slower than main at one stream and the widest admitted; with --drafter, no
    slower 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.

@ashhart

ashhart commented Nov 7, 2026

Copy link
Copy Markdown
Owner

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 ## Unreleased, since 1.2.4 has shipped. Second, please attach the comparison tool and the raw runs behind the TTFT table. Third, one drafted run on each family, main against branch, with replies equal. By our reading, a drafter's head makes the chunk a reader, so this should only confirm that.

@jschmied
jschmied force-pushed the prompt-tail-gemma-nemo branch from eb2f6f1 to 26da2ea Compare November 7, 2026 14:23
@jschmied

jschmied commented Nov 7, 2026

Copy link
Copy Markdown
Contributor Author

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.

  • CHANGELOG: the line is now at the top of ## Unreleased: "... long cold prompts start up to 4 % sooner (Gemma 4 1.3 % at 2K, Nemotron level at 2K) with identical replies ... (gemma4, nemotron cuda: a prompt chunk nothing reads stops at its last state write #2432, refs qwen3_5 CUDA: the prompt's last layer computes q/o/MLP only for the kept row's tile, logits bit-equal #2165). Thanks to @jschmied."
  • Tool and raw runs: they are in docs/receipts/prompt-tail-gemma-nemo/. That has the arm runners and the comparison gates (tools/), every arm's JSON, the server logs and each gate's output for every cell of the table (runs/), plus a README that maps each cell to its run. Local paths are redacted, and Nemotron's prompt_logprobs arrays are replaced by their sha256, since they run to about 3.7 MB per arm.
  • Drafted runs: main against branch, drafts on, on the table's nine prompts, 64 tokens each at temperature 0. Every request had drafts accepted, and the replies and token hashes are equal 9 of 9 for each family.
    • Nemotron (MTP head): TTFT is level, from -1.0 to +0.4 %. That fits your reading: the head absorbs the rows, so every chunk runs whole.
    • Gemma 4 (z-lab DFlash): your reading doesn't hold here. The drafter taps layers 1, 12, 23, 35, 46 and 57 of 60 and never the last one, so the chunks stay unread and the tail is still skipped with drafts on. Branch TTFT at 32K was -3.8 to -4.0 % against main, with equal replies. These are single arms, not ABBA.
    • Drafted counts per prompt differ a little between the arms while the tokens stay the same: Gemma 4 by up to 19 drafted and 1 accepted, Nemotron on one prompt by 1 accepted. I haven't measured the cause, and there is no main-against-main control.
  • Checks on this head: zig build test (7,418 of 7,476 passed, 58 skipped), test-kernel-contract, test-golden (0 unexpected), test-server-security, test-pull-security, fuzz-cortheon, zig fmt --check and lean (116 on base and head, 0 new) all pass.

… 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.
@ashhart
ashhart force-pushed the prompt-tail-gemma-nemo branch from 26da2ea to 681b8a6 Compare November 7, 2026 20:13
@ashhart
ashhart merged commit a4236ad into ashhart:main Nov 7, 2026
3 of 7 checks passed
@ashhart

ashhart commented Nov 7, 2026

Copy link
Copy Markdown
Owner

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.

@ashhart

ashhart commented Nov 8, 2026

Copy link
Copy Markdown
Owner

@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:

  • Prompt: +1.5% at 2K (992.8 to 1007.6 tok/s) and +2.3% at 8K (918.3 to 939.1).
  • Decode: level, 26.2 tok/s at 2K and 11.7 at 8K.
  • Streams 1 to 8: level, 85.7 to 85.6 aggregate at 8.
  • Tokens: 0 of 20 prompt rows differ across runs. Chat qualified with thinking on and off.

That's a little better than the 1.3% your receipt measured at 2K.

@ashhart

ashhart commented Nov 8, 2026

Copy link
Copy Markdown
Owner

Thanks @jschmied! This is in TensorFold 1.2.5, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.5

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