[jschmied]

On main (14c907b, with #1619 landed): four commits. Kolibri-1's full layers take their K/V from pages of the core pool, every stream leasing as it grows; sliding layers keep their rings. TENSORFOLD_KOLIBRI_KV_PAGES=0 keeps the linear caches.

• Pages of 512 positions (tree attention's key groups). Arena, plan and leasing as in gemma4 CUDA: the global layers' caches in pages of the core pool, every stream leasing; sliding layers keep their rings #1619.
• Writes: tf_wrows_put copies the rows kernel's kn/vn into the pages.
• Reads: prompts through new tf_window_prompt_h4_pg/_h8_pg (twins of wattn_kernel<128, 8, 4|8, 8>), decode and drafted verify through 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 tree attention _pg kernels with a table in the call's stream order.
• Kept states: copies, walked through the slot's pages.
• attach drops any step graph captured before it (none is today; the guard keeps a stale graph off the freed linear caches).

Checks (GB10, nvcc 13.0, on main d33dcfc)

• SASS of window_attention.cu against gemma4 CUDA: the global layers' caches in pages of the core pool, every stream leasing; sliding layers keep their rings #1619: 35 functions over sm_80/86/89/120/121, 0 changed, 10 added, 0 unstable.
• Records of the two new kernels measured: 223 registers on sm_80/86/89, 229 on sm_120/121; kernel-contract-limits over the 5-SM fatbins shows them current.
• Kernel contract: tf_window_prompt_h4_pg over pages of 512 and 37 and _h8_pg over 64, keys shuffled with NaN in spare pages and unused slots, pass repeat, shapes, sums and nonfinite; the control (first page id a spare page's) goes red at shapes. The runner's SUMMARY caveat is in cuda: tree attention and head dim 512's window rows read keys from a page table (_pg kernels), SASS of the others unchanged #1618.
• tf-cuda-test kolibri-pages MODEL TEXT 64: linear vs paged, pool of the two widest requests shuffled, plain and drafted: replies and last-row logits byte-equal, 29/29, including shared rounds, a kept state's resume and reused pages.
• kolibri-full MODEL IDS 128: passes paged and with TENSORFOLD_KOLIBRI_KV_PAGES=0.
• zig build test test-kernel-contract test-golden test-server-security test-pull-security fuzz-cortheon, git diff --check: pass; lean 0 new.

Speed, paged vs TENSORFOLD_KOLIBRI_KV_PAGES=0 (same harness as #1619: fresh server per arm, three alternating rounds, 256 greedy tokens, two prompts per case; min / median / max):

	paged	linear
c=1, 2k prompt: TTFT s	0.99 / 1.07 / 1.12	1.01 / 1.08 / 1.14
c=1, 2k prompt: decode tok/s	51.2 / 52.2 / 53.2	52.2 / 52.4 / 53.3
c=1, 9k prompt: TTFT s	3.90 / 4.06 / 4.22	3.89 / 4.05 / 4.22
c=1, 9k prompt: decode tok/s	49.8 / 50.8 / 51.9	50.7 / 51.5 / 52.3
c=4, 9k prompts: tok/s total	44.6 / 45.6 / 46.4	45.2 / 45.9 / 46.7

TTFT is the same; decode medians are 0.3 to 1.3 % lower paged, inside the arms' spread. kolibri-full's in-process rates agree: 53.5 / 85.1 (drafted) / 124.4 (four streams) tok/s paged against 52.9 / 84.7 / 124.1 linear (that linear run on the previous main, 7c06889).

[ashhart]

Thanks @jschmied. We ran this on a GB10 (sm_121) against main, with Kolibri-1, fresh servers per arm and alternating base/new rounds. Our first run asked for 64 streams, which fit neither arm (paged needed 42.8 GiB for the other streams against main's 135.6 GiB, so the memory win is real). This run uses 12 streams at a 40k context.

Tokens are equal to main (0 of 12 prompt-test rows differ), and prefill is flat. But decode is slower everywhere, and every run of this PR came in under every run of main:

cell	main tok/s	this PR
1 stream	50.5	49.7 (-1.6%)
4 streams	141.1	139.8 (-0.9%)
12 streams	282.4	276.3 (-2.1%)
decode after a 2k prompt	52.4	51.8 (-1.0%)
decode after a 9k prompt	47.9	47.1 (-1.8%)
decode after a 32k prompt	46.0	45.7 (-0.6%)

This matches your own medians, which were 0.3 to 1.3% lower. We don't land a slower default decode, so we're parking this for now, along with #1619 underneath it. The likely cost is the full-attention layers' paged decode reads: one page-table lookup per 256-position page for each head. If the decode kernel reads a stream's pages with the table hoisted out of the inner loop, or the full layers keep a contiguous cache until the pool needs the room, decode should be level, and the memory win can land with it. We'll rerun the same cells when you push.

[jschmied]

/atc exchange b348c34b14982e4a51a78ad30b389dc4e469a386fd4655dbbcfaf0b60f918cf9

[jschmied]

Thanks @ashhart. I pushed 58a8cfb here (and the same change as d4bafe4 on #1619, since the kernels came from #1618): the tree attention _pg kernels and tf_wrows_part_512_pg now look up a tile's page once per tile, not once per key. The addresses are the same, so the bits are the same.

Where the time goes. I timed the kernels alone first, linear against _pg on the same keys with a cold L2. At Kolibri's shapes (48/4 heads, d 128, pages of 512) the tree attention _pg kernels were already level with their linear twins: within ±4 % at 2k, 9k and 32k and at 1, 4 and 12 rows, also with the pages spread over a 20,000-page arena. The per-key lookup did cost something on Gemma 4 12B's global decode (16/1 heads of 512): +5.8 % at 9k × 4 rows and +7.9 % at 32k × 4. After the fix that is +1.9 % and +3.0 %. An nsys trace of the served Kolibri decode shows the same: 255 steps took 4,816 ms paged against 4,866 ms linear, and 5,834 against 5,896 ms after a 32k prompt. The only work paging adds per token is tf_wrows_put (13 µs) and 9 µs in tf_tree_part_w1_pg, about 0.1 % of a step.

Your cells, fix against main d33dcfc. Kolibri-1, --parallel 12 --context 40960, fresh server per arm, arms alternating, three rounds, 256 greedy tokens. Decode tok/s, min / median / max:

cell	main	this PR, fixed
1 stream	52.9 / 53.7 / 53.9	53.1 / 53.2 / 53.7
4 streams	123.1 / 124.6 / 125.9	123.3 / 124.3 / 125.2
12 streams	195.6 / 197.1 / 199.6	196.2 / 199.0 / 199.0
after a 2k prompt	51.9 / 52.5 / 52.6	51.9 / 51.9 / 52.4
after a 9k prompt	49.3 / 50.0 / 50.4	49.4 / 49.5 / 50.0
after a 32k prompt	42.8 / 44.0 / 44.5	43.2 / 43.4 / 43.6

The ranges overlap in every cell. Per round, the fix comes in 0.2 to 2.0 % under main in rounds 1 and 3 and level or above in round 2. The PR's own linear caches (TENSORFOLD_KOLIBRI_KV_PAGES=0) moved -1.6 to +1.1 % against main between rounds, so this setup can't resolve a 1 % gap. Tokens are equal in every cell, arm and round.

On #1619's 31B + DFlash single-stream code cell, the fixed binary is still 0.7 to 2.4 % under main in each paired round, with replies byte-equal to main and to "draft": false. So I can't say that drafted path is level yet. If your rerun shows the same, the cheap way out is the one you suggested on #1619: keep linear caches while a drafter is attached.

Checks: SASS (nvcc 13.0, sm_80/86/89/120/121) differs only in the 13 _pg functions I changed. Their 7 kernel records are updated from kernel-contract-limits. The contract's tree attention and head dim 512 entries pass, and their 6 controls go red. kolibri-pages is 29/29 byte-equal. gemma4-pages 12B is 15/15 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 0 new.

[ashhart]

Kolibri-1 on a GB10, @jschmied: your three #1620 commits on top of our rebased #1619, against #1619 alone, so only this PR's change shows. --parallel 12 --context 40960, alternating fresh servers (base, this, this, base).

• Tokens equal (0 of 12 prompt rows differ).
• Decode after 2k, 9k and 32k prompts: level. The 9k cell's -1.3% is main's first run reading high; its last run matches this PR's.
• Sweep at 1 and 12 streams: level (-0.5%, -0.1%).
• Sweep at 4 streams: slightly lower in every cell of both runs:

4 streams, tok/s	#1619 alone (runs 1, 4)	+ #1620 (runs 2, 3)
code, T 1.0	112.0, 112.3	111.8, 111.8
chat, T 1.0	113.4, 113.7	113.2, 113.0
code, greedy	175.5, 175.7	174.9, 174.5
chat, greedy	168.1, 168.1	167.5, 167.1

That's about -0.4%, the same shape as #1653's 8-stream cost before your paged decode merge fix (7c0a80b). If Kolibri's decode goes through the same window_rows merge, carrying that fix here may well level it. Could you check, and push it here if it applies? We'll re-run the same cells.

[jschmied]

Thanks @ashhart. I checked, and #1653's merge fix doesn't apply here: Kolibri-1 decodes its full layers through tree attention, not window_rows. Tree attention's replay is already a __noinline__ helper (tree_odd), and tf_tree_merge reads no pages.

The 4-stream cost looks like registers instead. With nvcc 13.0 on sm_120/121, the head dim 128 one-warp paged twins tf_tree_shared_w1_pg and tf_tree_shared_w1s3_pg come out at 168 registers with 44 and 40 bytes of spill stores on their fast path. Their linear twins use 195 and 198 registers with no spill. sm_80/86/89 don't spill.

597d473 gives those two _pg twins __maxnreg__(232) in place of __launch_bounds__(32):

• They now use 212 and 214 registers with no spill (210 and 213 with nvcc 13.2).
• The arithmetic is unchanged, and so is every other kernel's SASS. The head dim 256 twins keep their bounds, because the budget made them worse.
• The records are updated for 13.0 and 13.2. The two kernels' .bound is now 0, since they have no __launch_bounds__.

Checks:

• kolibri-pages is 29/29 byte-equal.
• 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. Could you re-run the 4-stream cells? #1619's merge fix follows there separately.

[ashhart]

@jschmied #1619 landed today (14c907b) in its rebased form. This PR still sits on an older #1619, so its branch carries #1619's earlier Gemma 4 paging commits. Those now conflict with what landed, in gemma4/cuda_assistant.zig, cuda_backend.zig, cuda_forward.zig, cuda_native.zig, cuda_snap.zig, tests/cuda/gemma4_pages_tests.zig and tests/cuda/main.zig. Could you rebase #1620, and #1755 on top of it, onto current main, keeping only the Kolibri-1 commits and anything not already in #1619? It lands from the queue as soon as it applies; the review and device checks stand.

[jschmied]

Rebased onto main 14c907b as asked: 9d80794, four commits, only Kolibri-1's. Every #1619 commit on the old branch already has its equivalent in main, so they're gone; the four applied without conflicts.

• ecc4ff30a kolibri1 cuda: the full layers' caches in pages of the core pool
• 9165c6a15 kolibri1 cuda: attach drops step graphs captured before the arena
• f68efdb43 kernel contract: Kolibri's paged prompt rows over shuffled pages
• 9d8079414 cuda tree attention: the head dim 128 one-warp shared _pg kernels take a register budget

Checks on 9d80794 (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 (shuffled pool pages give the linear caches' replies and logits, plain and drafted, shared rounds and a kept state's resume included). #1755 follows on top once its checks finish.

[ashhart]

Thanks @jschmied for the quick rebase. Each of the four commits here, and #1755's two, has the same patch ID as the series we reviewed and ran, so it's the same code on current main. Both are in our landing queue now: #1620 first, then #1755 once it's in.

[ashhart]

A small rebase from our side, @jschmied: since your 14:15 push, main gained #2022, the verified copies on the CUDA lane hosts, which touches the same facts() line in kolibri1/backend.zig. I kept both: #2022's priority-copy fields and your slotCap(). Nothing else changed, and the four commits keep your authorship. It lands once its CI is green, then #1755.

[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