[jschmied]

What this changes

A prompt chunk keeps one row's logits. Every layer's keys, values and recurrences are still written for every row. Past the last layer's keys and values, though, its queries' attention, gate, o projection and MLP feed nothing but that row. So in the last layer, k and v still run over the whole chunk, and q, attention, attn_gate, o and the MLP run over the last 128 rows only.

• 128 rows, not one. At 128 rows, q, o and the MLP stay on the prompt GEMMs (rows > 16), and each row keeps its own per-row bits. One row would drop to the decode kernels.

• cuda_prompt.runFrom runs the prompt attention from one chunk row onward: each row over its own keys (len + from), with the same key tiles. run is now runFrom(.., 0, rows), paged and whole.

• Dense and routed MLPs both take the tail. The glue launches (attn_gate, mlp_act, routing) get no device row count for the tail, because their rows come from their own pointers.

• The decision lives in core/prompt_tail.zig, so any family can share it: a chunk past one tile, in one part, whose last layer fits a row range and has no other readers, takes the tail from total - 128. Its host tests cover the short chunk, a cut chunk, readers and a layer that doesn't fit.

• Qwen 3.5's facts, qwen3_5/cuda_prompt_tail.zig: hidden reads, learner hooks, pockets, int8, images and a tapped last layer are readers; odd head groups and a last layer that is not attention don't fit. MTP taps the last layer, so a 35B served with its MTP head keeps the full layer. DFlash2's taps (5..61) do not refuse it. A host test plants each fact and checks that each one refuses the tail.

• Issue: Fixes qwen3_5 CUDA: the prompt's last layer computes q/o/MLP only for the kept row's tile, logits bit-equal #2165.

• Stacked on: main 981fe201d, one commit.

• atc overlap --json: no conflicts against open pull requests (qwen35: dense BF16 loading and native prompt logprobs #2020 and CUDA: #1764's BF16 prompt tile for Qwen ModelOpt, Nemotron and Dory, plus two exact Qwen prompt kernels (#1762) #2008 share files and merge cleanly).

Checks on this head

• Head commit: 77c3f9de0.
• Lean checker on the base and on this head: 425 and 425, 0 new.
• Tests run: zig build test RC 0, 0 errors; the no-GPU CI set (test-kernel-contract test-golden test-server-security test-pull-security) RC 0; git diff --check clean. tf-cuda-test qwen-pages on a GB10 (sm_121): every paged == contiguous case passes, bf16 and fp8, including the runFrom rows, and its controls are red as planted (swapped table, planted row, uncounted tail keys).
• CI: green on the GB10 for the set above.
• Not run, and why: other CUDA chips (only a GB10 here); the 35B with its MTP head gets no change by design (the guard keeps the full last layer).

Receipt

• Environment: TensorFold main 981fe201d + this commit; Zig 0.17.0, nvcc 13.0; GB10 (sm_121), clock capped at 2,300 MHz, page cache dropped before every server start; TensorFold/Qwen3.8-27B-MLX-4bit with its DFlash2 drafter, and Qwen3.6-35B-A3B MLX 4-bit without a drafter.
• Paths reached: qwen3_5 CUDA prompt chunks (dense 27B and routed 35B MLPs), paged and whole KV, bf16 and fp8.
• Exactness: the last row's logits are bit-equal to main. On a GB10, a local logits hook hashed every prefill's last-row logits (never committed). The hashes were the same on main, a main copy and this head, across all 18 servers per model: 27B and 35B, at 2K, 8K, 32K and the warm turn. The greedy texts were also the same. The kernel contract's qwen35 pages check now compares runFrom rows against the whole chunk's, paged and whole, bf16 and fp8. Its planted control (the start's keys left uncounted) must move those rows.

TTFT on a GB10, clock capped at 2,300 MHz, page cache dropped per start. Six ABBA rounds of fresh servers, each with A (main 981fe20), A' (a copy of main's binary) and B (this head):

model	cell	A' vs A	B vs A	B faster
Qwen3.8 27B 4-bit, DFlash2	2K	+0.24 %	−1.23 %	6/6
	8K	−0.04 %	−1.31 %	6/6
	32K	+0.08 %	−1.62 %	6/6
	warm turn	+0.10 %	−1.48 %	5/6
Qwen3.6 35B-A3B 4-bit, no drafter	2K	−0.28 %	−1.32 %	6/6
	8K	−0.02 %	−2.09 %	6/6
	32K	−0.07 %	−2.90 %	6/6
	warm turn	−0.02 %	−2.17 %	6/6

Decode cells (2K, 16K, four streams) stay within ±0.2 % on both models.

• Licence of new weights: n/a.

Co-credit to jungjipdo: their kept-row prompt tail in the cudafast Qwen3.8 27B challenge contributed the attention from the kept row's own query (runFrom). The 128-row q/o/MLP tile is from our entry in the same challenge.

[ashhart]

Thanks @jschmied. Keeping k and v over the whole chunk while only the kept row's tile runs the rest is a clean cut, and the planted uncounted-keys control plus 18 servers of bit-equal last-row logits per model make it easy to trust. We'll add the changelog line, crediting @jungjipdo for runFrom. Since this changes the default path for every Qwen 3.5 and 3.6 CUDA prompt, it gets one run on our GB10 before it lands; it's queued.

[ashhart]

Our GB10 run, main against this head, Qwen3.8-27B: verdict GREEN. 32K prompts are 1.6% to 2.0% faster (about 1,284 to 1,306 tok/s), decode after them is level, and the 1- and 4-stream chat and code sweeps are within 0.1%. No token differences across runs in the 20 prompt rows I checked. A small gain with nothing lost, and exact. Thanks @jschmied: it's queued for after the 1.2.3 cut.

[ashhart]

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]

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

[ashhart]

Thanks @jungjipdo! This is in TensorFold 1.2.4, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.4