[jschmied]

Part of #1657: the native-CUDA adapter over prompt_radix, off by default, with Kolibri-1 as its first user. Stacked on #1659 (merged) and on #1620.

Stacking. Based on #1620 at 9d8079414 (rebased onto main 14c907b). The first four commits are #1620's; the last two are this PR:

• native cuda: TF_KV_SHARE=n ... is the adapter, the same commit as in spark25 cuda: TF_KV_SHARE keeps prompts as the pool's pages in the prompt_radix store (on #1653) #1756 (Spark2.5, now directly on main since spark25 CUDA: the full layers' caches in pages of the core pool, every stream leasing; sliding layers keep their rings #1653 landed). Whichever PR lands first carries the adapter.
• kolibri1 cuda: ... is the family wiring.

What it does

• TF_KV_SHARE=n turns it on for a family on the core pool that declares share_pool and share_attach. That family gets one prompt_radix.Store over its pool (native/cuda_prompt_pages.zig).
• The tree holds the prompt pages themselves. The family's snap_* decls supply the state past a page edge, and n caps those states.
• The host makes no prompt_cache store, and the cache budget becomes pool pages.
• The family calls core/kv_prompts.zig:
  • begin: the resume point, adopted pages and the planned grid's marks;
  • keep at each mark;
  • report.
• A kept state's ring rows go in pool pages the store holds (Pool.hold). hold asks the store to reclaim first, as promise does.
• Unset or 0: as before.

Gates (GB10, Kolibri-1)

gate	result
tf-cuda-test kolibri-pages, final binary	TF_KV_SHARE=4: PASS twice (and once more on the head before the rebase onto today's main), resume at 1024 plain and drafted, reply and logits same as linear. Unset: PASS
chat suite, 64 turns (content + reasoning)	nocache / copy / share: 64/64 equal. Cached tokens: copy 230,719, share 218,112
long chat, 48 turns	48/48 equal. Cached tokens: copy 228,065, share 221,184
capacity, round-robin, 4 GiB, 72 turns	72/72 equal. Copy store: 24 states, 140,500 cached tokens. Share: 76 states, 159,744 cached tokens, failed 0
evictions, round-robin, 1 GiB	72/72 equal. Copy: 7 states, 3,479 cached tokens. Share: 40 states, 81,920 cached tokens, evicted 56, failed 0
v.sh (lean 0 new, zig build test, CI set, diff-check)	green on the final head

Share misses are the first turns: the system prompt is shorter than one 512-token page.

Default path (TF_KV_SHARE unset, --prompt-cache-gib 0). Base (the stack without these two commits) vs this head, fresh server per arm, ABBA over 2 rounds (4 at 4 streams), aggregate tok/s, greedy token_sha equal in every cell:

streams	delta of medians (chat/code, T 1.0/greedy)
1	+0.05 / +0.15 / +0.21 / −0.32 %
4	−0.14 / −0.00 / −0.05 / +0.03 % (4 rounds: rounds 1–2 had head −0.3..−0.9 %, rounds 3–4 +0.1..+0.4 %)
8	−0.08 / −0.01 / +0.17 / +0.05 %

The default path runs no new code in the step: the store is absent and the branches are on prompts == null.

Left. Gemma 4 wiring, after #1619.

[ashhart]

Heads-up @jschmied: #1756 is merged (6d3fa01), so main has the native CUDA TF_KV_SHARE adapter, and this PR's adapter commit is already there. What's left here is Kolibri's own wiring on top of #1620. When #1620 is ready we'll take them together.

[ashhart]

Thanks @jschmied. The share receipts are strong: 76 kept states against 24 at 4 GiB, and 81,920 cached tokens against 3,479 at 1 GiB, with every turn equal. Since #1756 put the adapter on main, this PR is now Kolibri's wiring on top of #1620. When you restack: 1. Drop the adapter commit. Main has it (40c03bf13). 2. #1620 turns Kolibri's pages on by default, so it needs its own sweep against main. When we parked it on 30 Oct, decode was 0.6-2.1% slower. If the once-a-tile page read fixes that, a decode sweep at 1 to 64 streams against main, level or better, settles it. The default-path table here is against #1620's stack, so it doesn't cover that. 3. With TF_KV_SHARE set, three of #1657's gates: concurrent equal to solo at 1, 2, 4 and 8; the lane pages at idle the same before and after shedding; and the cache bytes for each kept state before and after.

Then we'll run both on our GB10s and take them together.

[jschmied]

Thanks @ashhart. Both are restacked on main 1dfc78c55: #1620 is 05f1761df, this PR is d8b9d83e3.

1. Restack. The adapter commit is dropped (main's 40c03bf13; the range-diff is context only), and so are #1618's four commits, which main has. #1620 also carries #1619's cuda_pages.zig move for Gemma 4 (3f97e7489), so cuda_backend.zig stays at 599 lines. lean 0 new, zig build test and the CI set pass on both heads.

2. #1620 against main, decode sweep (GB10, Kolibri-1, 4 ABBA rounds, 8 fresh servers per arm, % delta of the medians, PR vs main). 64 streams don't fit on this box in either arm: at --parallel 64 the PR wants 42.8 GiB with 30.2 free and main also fails at 8192. So this is --parallel 24 --context 8192, at 1 to 24 streams:

streams	chat T 1.0	chat greedy	code T 1.0	code greedy
1	−0.28	−0.09	+0.19	−0.18
2	−0.12	−0.20	−0.30	−0.33
4	−0.20	−0.13	−0.40	−0.18
8	−0.29	+0.31	−0.15	−0.06
16	−0.40	−0.10	−0.19	−0.11
24	−0.35	−0.18	−0.09	−0.26

• Level: every cell is within −0.40 % to +0.31 %, and no cell is below −1 % in two or more rounds. At 24 streams, code greedy is 527.5 tok/s on main and 526.2 on the PR.
• Greedy token_sha256 is equal between the arms in all 12 greedy cells.
• SM clocks 2489–2496 MHz on every server. The first main server ran about 12 % slow in every cell at normal clocks (cause unknown); it only moves round 1, and the medians are robust to it.

At 4 streams that is −0.20/−0.13/−0.40/−0.18, against about −0.4 % in your 31 Oct table (on #1619) and 0.6–2.1 % when it was parked.

3. #1657's gates with TF_KV_SHARE (this PR's binary, numbers from the server's prompt cache: and kv pool idle lines). TF_KV_SHARE=n caps the store at n states, so at 4, with 24 conversations taking turns, every turn evicts and nothing hits. I ran 4 and 256; the earlier "share" receipts were 256.

• Concurrent == solo at 1, 2, 4, 8: 0 unequal at both 4 and 256, code and chat, T 1.0 and greedy. The sweep prompts are 55 tokens, so the store keeps nothing there. With prompts that share a 2,560-token prefix, every concurrent request resumed from the store at 2,560 and matched its solo run.
• Lane pages at idle, before and after shedding: total − free − store is 0 and promised is 0 on all 72 idle lines of every run, before and after the first eviction (turn 5 at share 4), with no UNEVEN line.
• Cache bytes per kept state (72 turns, every reply equal to the no-cache run):

run	kept states	cached tokens	MiB per state
copy, 4 GiB	24	140,500	166.4
share 256, 4 GiB	76	159,744	82.6
copy, 1 GiB	7	—	134.9
share 256, 1 GiB	40	81,920	89.0
share 4	4	0	172.5

One thing I haven't explained yet: at share 256 and 4 GiB the store holds 324 pages (6,280 MiB), more than the 4 GiB copy budget. The pool is 673 pages with share on, against 264 without, and its idle census stays even. So the share store seems to be bounded by the pool and the state cap, not by the copy budget. Do you want the share store held to the same byte budget?

The receipts and tools are local; I can attach any of them.

[ashhart]

@jschmied, an apology and a plan. While your #1620 and #1755 were waiting on us, one of our helpers built the same Kolibri-1 store wiring as our own PR #1891. Before that, the same happened for Gemma 4 with #1883, which landed tonight beside your #1619. That's our coordination slip, not yours: your PRs came first.

What we found comparing them: #1755's Kolibri commit and #1891 are two implementations of the same wiring (the same share_pool/share_attach declarations, the blobPages × pageBytes charge, kv_prompts begin/keep, and ring rows held in pool pages). The differences are that #1620 pages Kolibri's full layers for every stream by default (512-position pages, TENSORFOLD_KOLIBRI_KV_PAGES=0 to turn off), while ours pages only under TF_KV_SHARE.

So we're closing #1891 in favour of yours. #1620 now add/add-conflicts with #1883 on gemma4/cuda_pages.zig. If it helps, we'll do the rebase of #1619, #1620 and #1755 onto current main ourselves, with your authorship kept, and post the branches for you to check. Or restack them yourself and we'll gate them straight away. Since #1620 and #1619 page by default, we'll run each on a GB10 against main before landing. Which do you prefer?

[jschmied]

Thanks, and no harm done. Please go ahead and rebase #1619, #1620 and #1755 onto current main yourselves. Post the branches, and I'll check each against its PR's head: the same changes, nothing from #1883 or #1891 lost or doubled up on gemma4/cuda_pages.zig. I'll also run each one on our GB10 against main and post the numbers. Paging by default in #1619 and #1620 is fine to keep, with your own GB10 runs before landing as you planned.

[ashhart]

@jschmied, the three rebases are up, on main e3ea6b717, with your authorship kept on every commit:

PR	branch	tip	build	fmt	lean
#1619	lead/rebase-pr1619	32d1f87e3	pass	pass	no new problems
#1620	lead/rebase-pr1620	adad065ac	pass	pass	no new problems
#1755	lead/rebase-pr1755	f5f3f930b	pass	pass	no new problems

Build means check-native-cuda-host compile-cuda-tests for Linux. Nothing has run on a GPU, and the two unit tests we adapted (below) were only built, not run.

#1620 and #1755 replayed cleanly. Patches are identical except one context line in registry.zig's 13.2 set, where we kept both sides. #1755's adapter commit stays dropped (it's main's 40c03bf13).

#1619 is where the judgment was. #1883 already contains most of its Gemma paging, built from the same lines, so #1619 shrinks to what main lacks:

• Dropped as already on main: 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 four commits, and 1e9a7a2be (window pages loading, the first page before attach, the assistant's graphs). We kept Gemma 4 on CUDA: prompts in the shared KV store with TF_KV_SHARE (part of #1657) #1883's relocated hook rather than the PR's graph_gen, which would have been a second copy of the same mechanism.
• Kept from your PR: Options.kv_pages and pagesOn (TENSORFOLD_GEMMA4_KV_PAGES), shufflePages, the idle-census report, the pagesFor and paged-snapshot unit tests (adapted to main's each(..., from, ...)), gemma4_pages_tests.zig, and b426e525c's nvcc 13.2 tf_wrows_merge_512_pg entry. The kernel itself is on main already.
• gemma4/cuda_pages.zig: Gemma 4 on CUDA: prompts in the shared KV store with TF_KV_SHARE (part of #1657) #1883's file is kept byte for byte. Your leasing helpers duplicated main's cuda_shared.zig, so shufflePages, report and the pagesFor test now live there, and cuda.zig exports shared. report's check became prompts != null or free == total, as in your Spark 2.5 version, so it doesn't log UNEVEN under TF_KV_SHARE.
• The switch: paged = TF_KV_SHARE > 0 or pagesOn(o). With TF_KV_SHARE unset you get your default paging; with it set, main's behaviour is unchanged.

Please check these in particular:

1. With TF_KV_SHARE unset, Gemma now pages by default through main's paged path, not your PR's own code. Your GB10 results were on your code, and main's was only device-checked with TF_KV_SHARE set. One behaviour differs: main refuses to reopen a stream that already has a lane (StreamAlreadyOpen), where yours reused the lane.
2. TF_KV_SHARE pages even with TENSORFOLD_GEMMA4_KV_PAGES=0, so run the gemma4-pages check with TF_KV_SHARE unset.
3. In two intermediate commits (283b4bc8e, 66dd75a1a), default paging meets main's SharedDFlashNotServed, so gemma4-states with DFlash would fail if you bisect there. It's fine at the tip, where a drafter keeps linear caches.
4. The kept 13.2 merge_512_pg record was measured on your PR's base, and we haven't rechecked it on current main.

If they match your PRs, push them to your branches (or say so, and we'll update the PRs ourselves). We'll run each on our GB10 against main before landing.

[jschmied]

Thanks for doing the rebases. I've checked them and pushed them to the three PRs: #1619 at 32d1f87e3, #1620 at adad065ac, #1755 at f5f3f930b.

Against our heads: git range-diff shows #1620's and #1755's commits identical to ours, apart from the registry.zig context line you mentioned. #1619 is as you describe it.

On our GB10 (tf-cuda-test, ReleaseFast, nvcc 13.0 unless noted):

check	build	result
gemma4-pages 12B, TF_KV_SHARE unset (default paging through main's path)	#1619 tip	PASS, 15/15 replies and logits byte-equal to linear
gemma4-pages 12B + assistant, TF_KV_SHARE unset	#1619 tip	PASS, 15/15; the drafter keeps linear caches
gemma4-states 31B 4-bit + z-lab DFlash	#1619 tip	PASS: drafter context kept at 1000 and 1536, resumed == fresh
gemma4-pages 12B, TF_KV_SHARE unset	#1755 tip	PASS, 15/15
kolibri-pages	#1755 tip	PASS, 29/29
kolibri-pages, TF_KV_SHARE=4	#1755 tip	PASS, 33/33, the share store's cases included
tf_wrows_merge_512_pg registers, nvcc 13.2	#1619 sources	56 / 56 / 56 / 60 / 60 on sm_80/86/89/120/121, matching the kept 13.2 record (nvcc 13.0 gives 56 / 56 / 56 / 64 / 64, its own record)

On your point 3: the tip is fine. I didn't bisect the two intermediate commits.

One gap in the test, not a regression. gemma4-pages 12B + assistant with TF_KV_SHARE=1 fails its first check: "open made the caches asked for (a drafter: linear)". Under TF_KV_SHARE the model pages with a drafter too, as on main, but the test expects a drafter to mean linear (gemma4_pages_tests.zig:92). It should either skip under TF_KV_SHARE or expect paging there. I can push that one-line fix if you'd like it.

Gates on each tip: lean 0 new, zig build test and the CI set pass, git diff --check clean.

[ashhart]

@jschmied, our GB10 runs for the Kolibri-1 part of your stack are in, both on Kolibri-1 with base-new-new-base:

• Store off (defaults, kolibri1 CUDA: the full layers' caches in pages of the core pool, every stream leasing; sliding layers keep their rings #1620's head against main): GREEN. Flat everywhere, with identical tokens: 2K prompt 2,322.1 → 2,312.8 tok/s, decode 55.1 → 55.3. So nothing changes unless the store is asked for.
• Store on (TF_KV_SHARE=8, this head against main): REVIEW for one cell. Decode after a cold 2,048-token prompt is 55.7 → 53.8 tok/s (−3.5%), with both new runs (53.9, 53.6) below both base runs (55.8, 55.5). Tokens are identical, and every other cell is within noise.

Since the store is opt-in, that doesn't hold the stack back. Could you say whether the 2K decode cost is expected, for example from the pool lookup on each step, or something to shave? And add a line to the docs so users who turn it on know. Gemma 4's run is next.

[jschmied]

Thanks @ashhart. I measured it on a GB10 with Kolibri-1, at this head. It's a prompt-side cost, not a per-step one: decode doesn't change with the store on.

• Decode: after cold ~2K chat prompts (4 prompts, order off, on, on, off), decode is the same within noise: 54.3–54.9 tok/s with the store off against 54.4–54.7 on, and 51.2–51.8 against 51.5–51.7.
• Time to first token: about 10% longer with the store on (0.78–0.80 against 0.87–0.89 s, every prompt after the first in all four runs; the first prompt is the same). That's the cost your 2K cell picks up when decode is timed with the prompt included.
• Where it comes from: how the prompt is cut, not extra work.
  • Without the store, every request after the first resumes at its 45-token shared start and runs the rest in one pass of about 2,040 rows.
  • The store keeps whole 512-token pages, so that start is never reused, and a 2,092-token prompt runs as 2,048 + 44 rows. The extra short pass costs about 0.6 ms per layer in the expert prompt kernel.
  • The 2,048-row pass from position 0 also costs more per call than the store-off pass from position 45 (2.36 against 2.05 ms). I put that down to content, because with matched input the kernel doesn't differ: exact-length random-id prompts with no shared start (2,047, 2,048 and 2,049 tokens) take the same time to first token with the store on and off (0.82–0.87 against 0.83–0.89 s), and the expert kernel at 2,048 rows takes 2,586 against 2,579 µs.
• Not the pool: --parallel 2 shrinks the store-on pool from 24.9 to 18.5 GiB and the store-off pool from 10.1 to 2.5 GiB, and neither moves the per-call times.
• The pool lookup per step: nsys shows no change in the decode kernels, tree attention included.

So it's expected given the page size. It could be shaved by letting the store keep a short tail below a page, so a shared start of a few dozen tokens still hits. That would be a follow-up. For now I've added the line to docs/known-issues.md (2274a43).

[jschmied]

Restacked on #1620's rebased head 9d8079414 (main 14c907b), as asked: 6b5bb306d. The two commits applied without conflicts:

• d350e4496 kolibri1 cuda: TF_KV_SHARE keeps prompts as the pool's pages, a kept state's rings in store-held pages
• 6b5bb306d docs: known issues: Kolibri-1's TF_KV_SHARE costs a cold chat prompt its sub-page shared start

Checks on 6b5bb30 (GB10): lean 424 on the base, 424 on the head, 0 new; zig build test, the CI set (test-kernel-contract test-golden test-server-security test-pull-security fuzz-cortheon) and git diff --check pass; tf-cuda-test kolibri-pages passes plain and with TF_KV_SHARE=4.

[ashhart]

Thanks @jschmied. #1620 has landed, so we rebased just this PR's two commits onto main and pushed them here, with the changelog line on top. It's queued to land after the current group.

[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.3, out now: https://github.com/ashhart/TensorFold/releases/tag/v1.2.3