gemma4 CUDA: the global layers' caches in pages of the core pool, every stream leasing; sliding layers keep their rings by jschmied · Pull Request #1619 · ashhart/TensorFold

jschmied · 2026-10-29T19:55:04Z

Stacked on #1618 (only the last three commits are new). Gemma 4's global layers take their K/V from pages of the core pool, every stream leasing as it grows; sliding layers keep their rings. TENSORFOLD_GEMMA4_KV_PAGES=0 keeps the linear caches.

Paged: global layers only (every sixth: 8 of 48 on the 12B, 10 of 60 on the 31B), pages of 256 positions. One arena: per global layer a K region, then a V region; page p is rows [256p, 256(p+1)) of each.
Plan: Loaded.layout with engine_paged, attach/fits exported as lfm2_moe and Dory do; a stream's fixed bytes are its sliding rings and table row (about 480 MiB on the 12B, 1.17 GiB on the 31B).
Leasing: promised pagesFor(prompt, max_new) at the prompt pass, pages taken as rows reach them, returned at dropLane. LaneMemory and fits hold back what can't be promised; forward refuses rows past a slot's pages (PastPages).
Writes: the rows kernel writes its kr/vr scratch as before, and tf_wrows_put copies those rows into the pages (one launch per global layer per forward), so the cache bits are the linear cache's.
Reads: decode through the whole table, prompt rows through the slot's row (cuda: tree attention and head dim 512's window rows read keys from a page table (_pg kernels), SASS of the others unchanged #1618's 512 _pg kernels). The assistant and DFlash read the target's pages the same way.
Kept states: still copies; snap.each walks a paged layer through the slot's pages.
Fixes found on the GPU: the drafter's warm-up runs before attach, so build gives slot 0 a first page; the assistant's captured graphs are dropped when attach replaces the caches (graph_gen).

Checks (GB10, on main `d33dcfc`)

tf-cuda-test gemma4-pages MODEL TEXT 64 [DRAFTER]: linear vs paged at context 8192, two lanes, the pool exactly the two widest requests' pages and shuffled, so later streams lease pages earlier ones returned. Replies and last-row logits byte-equal: three prompts (200, 1,300, 700), a sampled reply, two streams in shared rounds, a state kept at 1,000 and resumed.

model	drafter	same
12B MLX-4bit	assistant	15/15
12B MLX-4bit	none	15/15
31B MLX-4bit	DFlash	15/15

gemma4-states now attaches the pool when paged. Its two "drafter context MISSING" checks fail the same way with linear caches (TENSORFOLD_GEMMA4_KV_PAGES=0).

Speed, paged vs TENSORFOLD_GEMMA4_KV_PAGES=0: tensorfold-native serve --context 16384 --parallel 4, no drafter, fresh server per arm, three rounds with the order alternating, 256 greedy tokens, two prompts per case; min / median / max:

	12B paged	12B linear	31B paged	31B linear
c=1, 2k prompt: TTFT s	1.03 / 1.10 / 1.20	1.03 / 1.11 / 1.19	2.51 / 2.72 / 2.94	2.52 / 2.72 / 2.94
c=1, 2k prompt: decode tok/s	30.5 / 31.0 / 31.2	30.8 / 30.9 / 31.6	12.34 / 12.44 / 12.53	12.33 / 12.49 / 12.66
c=1, 9k prompt: TTFT s	4.09 / 4.37 / 4.63	4.13 / 4.38 / 4.63	10.28 / 10.84 / 11.43	10.28 / 10.86 / 11.43
c=1, 9k prompt: decode tok/s	30.0 / 30.5 / 30.7	30.3 / 30.6 / 31.3	11.9 / 12.5 / 13.0	11.9 / 12.5 / 13.2
c=4, 9k prompts: tok/s total	38.0 / 39.1 / 40.2	38.2 / 39.3 / 40.5	15.09 / 15.45 / 15.83	15.08 / 15.50 / 15.89

The arms overlap in every cell; the extra put launch does not show. Not measured: vision prompts over pages.

zig build test test-kernel-contract test-golden test-server-security test-pull-security fuzz-cortheon, git diff --check: pass; lean 0 new.

ashhart · 2026-10-30T11:30:53Z

Thanks @jschmied. We ran this on a GB10 against main with no drafter, and it's clean. Gemma 4 31B 4-bit with cold prompts at 2k, 8k and 32k, and Gemma 4 12B 4-bit at 1, 2, 4 and 8 streams, both GREEN: tokens equal to main and no cell slower beyond noise.

One thing stops us landing it yet. Paging is on by default for every Gemma 4 CUDA server, and the drafted paths go through the new paged kernels too: the 12B assistant reads the global layers through the _pg decode kernel, and the 31B's DFlash verify windows go through paged decode. Those are the headline drafted numbers, and we have no Gemma drafters on our test boxes, so nobody has measured them.

Could you post, on your GB10, paged (the default) against TENSORFOLD_GEMMA4_KV_PAGES=0 on the same binary:

the 12B with its assistant at 1 and 4 streams, code and chat;
the 31B with DFlash at 1 and 4 streams, code and chat;
tok/s for each, with replies equal between the two arms and to "draft": false.

If paged is no slower within noise, this lands. If it is slower on a drafted path, keeping linear caches while a drafter is attached (or until the paged decode catches up) would let the rest land now. #1620 waits on this one.

jschmied · 2026-10-30T15:17:51Z

Thanks @ashhart. Measured on the GB10 with c460a2a, one binary for both arms, paged (default) against TENSORFOLD_GEMMA4_KV_PAGES=0. Fresh server per arm, arms alternating, three rounds each, --parallel 4 --context 8192, greedy, seed fixed, replies to 448 tokens (the 31B chat prompt ends at 303), one warm-up per server. Per-stream is the server's decode tok/s; aggregate is tokens over wall time.

cell	paged	linear	paged vs linear
12B + assistant, 1 stream, code	98.0 ±3.1	99.2 ±0.4	-1.2%
12B + assistant, 1 stream, chat	59.1 ±0.9	59.0 ±0.4	+0.3%
12B + assistant, 4 streams, code (agg)	223.1 ±5.7	225.5 ±0.9	-1.1%
12B + assistant, 4 streams, chat (agg)	192.9 ±2.7	191.8 ±2.9	+0.6%
31B + DFlash, 1 stream, code	57.6 ±1.2	58.5 ±0.9	-1.7%
31B + DFlash, 1 stream, chat	24.6 ±0.3	24.8 ±0.5	-1.0%
31B + DFlash, 4 streams, code (agg)	108.2 ±0.7	109.0 ±1.4	-0.8%
31B + DFlash, 4 streams, chat (agg)	74.3 ±0.5	75.1 ±1.1	-1.0%

Every reply is byte-identical between the two arms and to its "draft": false reply, on all six servers per model.

Six of the eight means are lower on the paged side, by 0.8 to 1.7 %, each inside the run-to-run spread. Paged wins round 2 on all four 31B cells, so I read it as noise rather than a gap. Paged is no slower within noise on the drafted paths.

jschmied · 2026-10-30T20:39:43Z

/atc exchange 7d514b9d11922e9c3a53bbd6293e53f6373fdc00710d42e937e2f21317d9c293

jschmied · 2026-10-31T01:08:51Z

I pushed d4bafe4 here: the paged decode kernels from #1618 (tree attention _pg and tf_wrows_part_512_pg) now look up a tile's page once per tile instead of once per key, with the same bits. On Gemma 4 12B's global decode the per-key lookup had cost +5.8 % at 9k × 4 rows and +7.9 % at 32k × 4 rows in the kernel alone; with the fix it is +1.9 % and +3.0 %. Only the 13 _pg functions change in SASS, and gemma4-pages is 15/15 byte-equal with and without the assistant. The 31B + DFlash single-stream code cell is still 0.7 to 2.4 % under main per round with this, so that drafted path isn't shown level yet. The numbers and checks are on #1620.

ashhart · 2026-10-31T01:33:17Z

Thanks @jschmied, and thanks for the drafted cells. Byte-identical everywhere is what we needed.

On the 31B + DFlash single-stream code cell that's still 0.7 to 2.4% under main: yes, take the cheap way out. Keep the global layers' caches linear while a drafter is attached, and paged otherwise. A default can't get slower on any cell, and paging only pays when streams share the pool. A drafter-attached server can move to pages later, in its own PR, once that cell is level.

Push that and we'll re-run the 12B + assistant and the 31B + DFlash cells, at 1 and 4 streams, code and chat, paged against TENSORFOLD_GEMMA4_KV_PAGES=0 and against main. Then #1620's Kolibri cells on top of it.

jschmied · 2026-10-31T02:33:07Z

Done in f419bed: Gemma 4's global layers lease pages only when no drafter is attached. With --drafter (the assistant or DFlash) the server keeps the linear caches, and its startup line reports them as such. TENSORFOLD_GEMMA4_KV_PAGES=0 still turns paging off everywhere.

gemma4-pages without a drafter: 15/15 byte-equal, paged against linear.
With the 12B assistant, the paged pass now opens linear caches as asked, and its replies and logits match the linear pass, 15/15.
31B + DFlash, single-stream code against main d33dcfc: two rounds each, fresh servers, order main, this, this, main. Decode tok/s was 59.3 / 57.0 and 57.7 / 57.6 on main, and 56.8 / 57.5 and 57.9 / 56.7 here. Replies match each other and "draft": false. That is inside the repeat-to-repeat spread of this cell, and it is now the linear path main runs.
zig build test test-kernel-contract test-golden test-server-security test-pull-security fuzz-cortheon and git diff --check pass; lean 0 new.

Your rerun list then covers 12B + assistant and 31B + DFlash, at 1 and 4 streams, code and chat, paged against TENSORFOLD_GEMMA4_KV_PAGES=0 and against main. With a drafter attached, both arms on this head are linear. #1620's Kolibri cells follow on top.

ashhart · 2026-10-31T03:05:18Z

Thanks @jschmied, that's the fix we wanted. The 31B + DFlash single-stream numbers now sit inside main's own spread.

We've rebased the series onto current main (94b90c9). One conflict: #1386 added Gemma 4's grammar-masked greedy pick (pickAllowed and its test) to the end of gemma4/cuda_backend.zig, where your pool field and page test go. We kept both, and the series builds the same code otherwise.

Together they take cuda_backend.zig to 666 lines, past the 600-line limit. Could you move a piece out on your branch (the pool and page-count code into a cuda_pages.zig, say) and push? You know where it splits cleanly. We've started the GB10 cells on our rebased copy in the meantime, since a move won't change them: Gemma 4 31B and 12B without a drafter, paged against main, alternating fresh servers. Then #1620's Kolibri cells on top.

jschmied · 2026-10-31T04:14:28Z

Moved in 2084804 on paged-gemma4: the pool and page-count code (pagesFor, attach, slots, leases, reach, the idle report) now lives in gemma4/cuda_pages.zig. The code itself is unchanged; cuda_backend.zig keeps pub const pagesFor and the page test, and its call sites go through pg..

Your rebased copy isn't anywhere I can reach (before this push the PR head was still f419bed, and I found no branch on ashhart/TensorFold), so I did the move on our branch. To check it applies after your rebase, I rebased f419bed onto 94b90c9 myself, resolved the conflict as you did (pickAllowed and its test kept, plus our page test), and cherry-picked the move on top. It applies without a conflict and leaves cuda_backend.zig at 598 lines there (510 on our branch).

No kernel file changes; SASS identical (0 images reached).
gemma4-pages 12B, paged against linear, with and without the assistant: replies and logits match, as before.
zig build test, the kernel-contract/golden/security/fuzz set, lean_check (0 new) and diff --check: all green.

ashhart · 2026-10-31T04:34:41Z

Thanks @jschmied. The move applies cleanly on our rebased copy: cuda_backend.zig is at 598 lines, lean is back to main's count, and the code is unchanged. Sorry the copy wasn't reachable; we push it to the PR branch when it lands. The Gemma 4 31B and 12B GB10 runs are queued against it.

ashhart · 2026-10-31T10:57:12Z

Gemma 4 31B on a GB10 without a drafter, @jschmied: your series with the split, rebased on main, against main, alternating fresh servers. Tokens equal (0 of 20 rows differ), prefill and decode level (-0.1%), sweep level at 1, 2, 4 and 8 streams (0.0 to +0.2%). GREEN. The 12B run is going now, and it lands with that. Then #1620's Kolibri cells after your merge fix.

ashhart · 2026-10-31T12:32:03Z

The 12B is back, @jschmied, and it's a hold for now. Gemma 4 12B 4-bit on a GB10, no drafter, main against this, order main, this, this, main. Tokens equal, prefill level, sweep level (-0.2 to +0.1%). But decode after long prompts is lower in both runs of this PR, in every cell:

decode tok/s after the prompt	main (runs 1, 4)	this (runs 2, 3)
16k, offset 0	49.79, 49.90	49.26, 49.36
32k, offset 0	53.63, 53.56	53.19, 53.13
64k, offset 5000	56.36, 56.20	55.07, 55.13 (-2.1%)
64k, offset 10000	42.32, 42.32	41.59, 41.73

That's about -0.5 to -2%, growing with context, and the shape is #1653's before your paged decode merge fix (7c0a80b). The 31B didn't show it, maybe because its global layers are a smaller share of its decode. If Gemma's global layers decode through the same window_rows merge, carrying that fix should level it too. Push it here and we'll re-run both.

jschmied · 2026-10-31T14:08:47Z

Thanks @ashhart. Gemma 4's global layers do decode through window_rows' paged merge (tf_wrows_merge_512_pg, and _256_pg for the 256 heads), so #1653's fix applies here as written. It's b426e52 on this branch: 7c0a80b cherry-picked, with the paged merge's float64 replay in a __noinline__ helper, so the fast path builds no page lookup.

The kernel code applied unchanged. Only the records conflicted, and I measured them again on this branch:
- nvcc 13.0: tf_wrows_merge_256_pg 48/60/60/60/60 registers and tf_wrows_merge_512_pg 56/56/56/64/64 (sm_80/86/89/120/121), no spill;
- nvcc 13.2 differs only for merge_512_pg, at 60 on sm_120/121, so that has its own 13.2 entry.
gemma4-pages on the 12B is 15/15 byte-equal, with and without the assistant.
zig build test test-kernel-contract test-golden test-server-security test-pull-security fuzz-cortheon and git diff --check pass. Lean has 0 new problems.

I haven't timed it here. Could you re-run the 12B long-context cells and the 31B?

…ry stream leasing; sliding layers keep their rings The window_pages kernels already read keys and values through a slot's page table at head dim 512, so this changes only where the global caches live and who writes them; no kernel changes. - A page is 256 positions (window_attention.page_unit: a decode chunk, whole prompt tiles) of every global layer's K and V; the arena holds each layer's K region, then its V region, page id p at rows [256 p, ...). - Loaded.layout with engine_paged, as lfm2_moe's and Dory's: open makes one slot's rings, the host plans the pool and calls attach (the arena, the slots' rings and tables). A stream is promised its prompt, reply and one verify window at its prompt pass, takes pages as its rows reach them and gives them back when it ends; fits is the host's admission. The idle server logs the pool's census. - Writes: the rows kernel writes no global cache when paged; tf_wrows_put copies its k and v rows (the same bits it would have stored) into the pages. - The decode and prompt windows of the global layers carry pg, in the target forward and in the assistant, which reads the target's pages through the same table. Kept states stay copies: a global layer's rows are walked through the slot's pages. - TENSORFOLD_GEMMA4_KV_PAGES=0 (or Options.kv_pages = false) keeps the linear caches made at open. tf-cuda-test gemma4-pages MODEL TEXT N [DRAFTER] runs the same prompts on linear and paged caches (pages leased in a shuffled order, a pool the ended streams' pages must be reused from) and wants the same replies and last-row logits, solo, sampled, two streams in a round and resumed from a kept state.

…est) before the passes

…ey (tree attention and head dim 512 _pg kernels), SASS of the others unchanged

…linear, paged only without one

…to cuda_pages.zig, the code unchanged

… a noinline helper, so its fast path builds no page lookup

ashhart · 2026-11-03T13:08:46Z

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 · 2026-11-03T18:58:45Z

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

…and Nemotron's Metal paging The switch registry grows by four rows (ashhart#1660's next group, all measured identical): TENSORFOLD_DORY_KV_PAGES (opt-in, ashhart#1284), TENSORFOLD_GEMMA4_KV_PAGES and TENSORFOLD_SPARK25_KV_PAGES (default-on, exact 0 opts out, ashhart#1619, ashhart#1653), and TF_NEMOTRON_METAL_KV_PAGES, the first metal-backend row (exact 1, ashhart#1537). A new .not_zero syntax covers Dory's on-unless-0 reader; the 0-opt-out rows reuse report_values. The four readers go through the registry's environment-first helper, so an unset variable and every set value behave byte for byte as before.

jschmied force-pushed the paged-gemma4 branch from b426e52 to 32d1f87 Compare November 2, 2026 11:30

jschmied mentioned this pull request Nov 2, 2026

Proposal: one KV admission ledger for unpaged streams: a request reserves its prompt and reply, not a whole window per --parallel stream (#781 follow-up) #1943

Open

Bizuayeu mentioned this pull request Nov 3, 2026

glm cuda: the idle KV pool's line reads even under TF_GLM_KV_SHARE #1993

Merged

8 tasks done

jschmied and others added 7 commits November 3, 2026 13:01

gemma4 states check: a paged model attaches its pool (both lanes' wid…

f36ed96

…est) before the passes

cuda: paged decode reads take a tile's page once a tile, not once a k…

19ca89b

…ey (tree attention and head dim 512 _pg kernels), SASS of the others unchanged

gemma4 cuda: a server with a drafter keeps the global layers' caches …

389368b

…linear, paged only without one

gemma4 cuda: the global pages' counts, attach, slots and leases move …

3c696c6

…to cuda_pages.zig, the code unchanged

cuda window_rows: the paged decode merge replays a non-finite head in…

3d157af

… a noinline helper, so its fast path builds no page lookup

changelog: ashhart#1619 under Unreleased

6f4afdc

ashhart force-pushed the paged-gemma4 branch from 32d1f87 to 6f4afdc Compare November 3, 2026 13:08

ashhart merged commit 14c907b into ashhart:main Nov 3, 2026
1 of 7 checks passed

devopsfortunato mentioned this pull request Nov 6, 2026

profile: register the pool-page switches for Dory, Gemma 4, Spark2.5 and Nemotron's Metal paging #2306

Open

4 tasks done

Conversation

jschmied commented Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Checks (GB10, on main d33dcfc)

Uh oh!

ashhart commented Oct 30, 2026

Uh oh!

jschmied commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jschmied commented Oct 30, 2026

Uh oh!

jschmied commented Oct 31, 2026

Uh oh!

ashhart commented Oct 31, 2026

Uh oh!

jschmied commented Oct 31, 2026

Uh oh!

ashhart commented Oct 31, 2026

Uh oh!

jschmied commented Oct 31, 2026

Uh oh!

ashhart commented Oct 31, 2026

Uh oh!

ashhart commented Oct 31, 2026

Uh oh!

ashhart commented Oct 31, 2026

Uh oh!

jschmied commented Oct 31, 2026

Uh oh!

Uh oh!

ashhart commented Nov 3, 2026

Uh oh!

ashhart commented Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jschmied commented Oct 29, 2026 •

edited

Loading

Checks (GB10, on main `d33dcfc`)

jschmied commented Oct 30, 2026 •

edited

Loading