[gilby]

What this changes

Fixes #1566. On a prompt chunk of more than one row, the MTP layer now writes every row's latent, keys and index entries as before, then runs attention, V expansion and o_proj for the last row only (at its own position, causal limit and token list), and the residual, post-norm, FFN and their gathers on that one row. Nothing else of that pass was ever read back. Decode and verify windows and one-row prompts take the existing path.

Receipt

• Environment: three DGX Spark (GB10, 128 GB), driver 580, CUDA 13.0, ConnectX-7 RoCE through a switch, NCCL over IB; nvidia/GLM-5.3-Flash-FP4-Local @ 76348ddb, --context 500000 --parallel 4 --prompt-cache-gib 0 --kv-cache fp8, TF_GLM_INDEX_CUDA=1, TP=3; zig build native -Dsm=121. The exactness and speed numbers below were measured on main a881b6f with this change, before the rebase onto 089a18d (this base); the Tests line says what ran on the rebased head. Main's mtpMode(..., propose) (skipped unallocated proposals) conflicted with the change: propose == false keeps main's cache-only commit call, and propose == true, which the prompt paths use, takes this change's last-row path (last_row is now propose and prefill and n > 1).

• Exactness, tools/bench_concurrent.py --levels 1,4 --tokens 256 --reps 3 --temperatures 1.0,0 --alone --serial: 0 unequal and 0 failed in all 16 cells, and every completion hashes the same as main's.

• Decode speed (bench_concurrent, greedy code): 54.4 tok/s at one stream and 117.7–118.3 at four, with and without.

• Prompt speed, cold first token, tools/bench_cache.py --require-cold --repeats 1 --streams 1 --tokens 16, main and this branch in one session:

  prompt tokens	main	this branch
  32,768	22.01 s	18.99 s (−13.7%)
  41,943	26.74 s	24.62 s (−7.9%)
  166,700	104.68 s	99.39 s (−5.1%)

  With TF_GLM_PROF=1 the MTP section of a 4,096-row prompt chunk reads about 32 ms with this change against about 190 ms before.

• Tests, on the rebased head b9b7644 (main 089a18d), on an Apple M5 Max (host tests only, no GPU work): zig build test: exit 0, 611/611 steps (180 of 184 tests in that invocation, 4 skipped); zig build test-cuda-host: 287/287 steps, 96/96 tests; zig fmt --check clean; python3 tools/zig/lean_check.py: 424 problems, the same as main. The earlier run on a GB10 (Linux aarch64, at 7c06889: test-cuda-host 270/270 steps, 3,843 of 3,852 tests passed, 9 skipped) was not repeated on the rebased head.

• Not run: the byte-equality and speed bench on the rebased head (the GB10s are serving a pin now), so those numbers are from before the rebase; TP=2 speed (the path is the same at any rank count; every rank takes it together because they share the prefill flag and the row count). Metal is not affected (CUDA family path).

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false and concurrent equals solo (the bench above). Resumed against fresh was not run separately for this change; every row's MTP cache writes are kept, so a resume reads the same caches.
• No precision traded for speed. The bits do not change: the last row's attention reads the same caches and token list as in the whole-chunk call, and no other row's result was read.
• No GPU kernel added or changed (existing kernels on one row).
• Prompt processing is no slower than the last release.
• Platforms: CUDA on GB10 at TP=3. Metal not affected.
• New tests fail before the change, pass after it: none added; the change is a pure saving with no new failure to catch, and the bench above checks the bytes.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[ashhart]

Hi @gilby, a status update on #1577.

We read the diff against the GLM CUDA code around it and found no way it changes reply tokens: the cache bytes
are the same, and the last-row path makes the same calls in the same order. What's still owed is a run on two
GB10s at once (TP=2), which your receipts didn't cover. It runs with nvidia/GLM-5.3-Flash-FP4-Local, which isn't
on our test machines yet. If you used a specific revision of it, please tell us which.

[ashhart]

@gilby, this conflicts with main. Could you rebase onto current main and confirm the checkpoint revision (76348ddb)? Then we'll run it at TP=2 on two GB10s.

[gilby]

Thanks @ashhart. Rebased onto current main (089a18df) as b9b7644e; the conflict was in cuda_forward.zig with your mtpMode(..., propose): propose == false keeps your cache-only commit call and propose == true, which the prompt paths use, takes this change's last-row path. The checkpoint is nvidia/GLM-5.3-Flash-FP4-Local at revision 76348ddb, the same one the receipt names. On the rebased head: zig build test exit 0 (611/611 steps), zig build test-cuda-host 287/287 steps and 96/96 tests on an Apple M5 Max (host tests only), zig fmt clean, lean 424 like main. The body says what was not repeated: the GB10 bench and host run are from before the rebase, since our three Sparks serve a pin now. The TP=2 run is yours as you said.

[ashhart]

Thanks @gilby: the 17:34 asks are answered and nothing more is owed from you. Our TP=2 run on two GB10s (main against this head, tokens and speed) is queued at the front. It lands when that's green.

[ashhart]

Our TP=2 run on two GB10s (GLM-5.3-Flash NVFP4, main against this head, order base-new-new-base, --context 65536 --parallel 2): GREEN. 52 of 52 token checks equal. Prompts +2.6% at 2K, +2.9% at 8K, +2.7% at 32K (1,571.7 -> 1,612 tok/s at 2K). Decode level (chat +0.0%/-0.1%, code -0.1%/-0.1%). It's in the landing queue.

[ashhart]

Thank you @gilby. 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]

Thanks @gilby! This is in TensorFold 1.2.2, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.2