[jschmied]

Part of #1488. Stacked on #1492 (now on main), and only the last four commits are new here. It stays a draft until #1492 lands, then I'll rebase onto main.

What

LFM2.5-24B-A2B on CUDA, step 2: attention pages from the shared CUDA KV pool, and the prompt cache.

• 92fc4f0eb Pool pages. The 10 attention layers' K/V live in pages of the core kv_pool (512 positions, 10 MiB a page), Dory's shape.
  • A stream leases pages for prompt + reply + 16 at its prompt pass, waits when the pool is short, and returns them at release.
  • forward refuses rows past a slot's pages.
  • The kernels already read through a page table, so no kernel or build file changes, and SASS is unchanged by construction.
  • TENSORFOLD_LFM2_KV_PAGES=0 keeps lfm2_moe: serve LFM2.5-24B-A2B on CUDA (plain serve, shared kernels) #1492's fixed pages per stream.
• 6ef5d8401 Prompt cache. A kept state is the slot's conv states plus K/V rows [0, at) of the 10 layers.
  • It is copied into snapshot pages allocated once, inside the prompt-cache budget, as GLM's cuda_snapshot_pages.zig does. A keep leases pages, a restore copies them back into the new stream's leased pages, and a drop returns the lease: no device allocation, and no free behind a stream wait.
  • Copies run page run by page run, with consecutive pages coalesced.
• cb823583b Keeps add no prompt pass. A chat mark inside a prompt chunk used to end the chunk there, and every extra pass reads the weights again. A cold 10k chat turn has two marks in its last chunk, so it ran 12 passes instead of 10 (nsys: 1,920 prompt GEMM launches against lfm2_moe: serve LFM2.5-24B-A2B on CUDA (plain serve, shared kernels) #1492's 1,600).
  • Now up to four marks inside a chunk are tapped. The conv kernel runs once more on each mark's last row alone, writing the same y and that row's conv states into a tap buffer, and the keep copies those with the K/V rows before the mark. The kernel is unchanged.
• 0506d9c0b Docs. The recipe says streams lease pool pages and the cache keeps chat states.

7 files, +647 / −43 against #1492.

Verification (GB10, real weights, served, greedy and sampled)

• Resumed == fresh: 8 chat conversations × 4 turns, including sampled and drafted ones. 24 of 32 requests hit the cache (440 to 6,243 tokens cached). Every reply equals cache off, fixed pages and lfm2_moe: serve LFM2.5-24B-A2B on CUDA (plain serve, shared kernels) #1492, and a restart gives the same replies.
• Drafted == plain, concurrent == solo: pass with pooled and with fixed pages, 8 short and 8 long prompts.
• Pool pressure: on a 40-page pool (TENSORFOLD_CUDA_MEMORY_LIMIT_GB=13.6) every reply matches. The 8 long concurrent requests take 7.39 s instead of 6.42 s while they wait for pages.
• Eviction: a 0.1 GiB cache gives 6 evictions and 1 hit, all matching the no-cache replies.
• Accounting: after every drain the pool reads promised 0, free = total (264 pages, or 40 under pressure). Host tests cover the layout, plan, lease map, snapshot page runs and the chunk cut.
• lean_check: 0 new against lfm2_moe: serve LFM2.5-24B-A2B on CUDA (plain serve, shared kernels) #1492's head. zig fmt --check zig and zig build test pass.

Speed (GB10, against #1492, alternating runs)

	#1492	this PR
cold chat turn that keeps two states, 10,144 tokens (TTFT, 4 alternating rounds × 6)	1.211 s	1.216 s
warm turn, 10,246 tokens with 10,141 cached (TTFT)	1.24 s	0.10 s
plain decode, 1 / 4 / 8 streams (tok/s)	86.5 / 191.7 / 298.7	86.5 / 192.4 / 300.2

I re-measured only the cold turn this round. The warm turn and decode rows are from the previous head, and nothing on their path changed.

The remaining 5 ms on the cold turn is the two 200 MB state copies. A device-to-device copy runs at about 115 GB/s here, so each takes about 1.8 ms.

Not in this PR

• The background warm pass that lets turn 2 resume (warm_quiet): its own PR.

[ashhart]

Thanks @jschmied. Leasing from the core pool and keeping states at chat marks is the shape we want, and checking the census after every drain is a nice touch.

With #1492's commits on main, both of your commits apply cleanly. zig fmt, test-lfm2-moe-cuda-host, test-kv-pool, test-cuda-host and the Linux cross-compile pass, and the lean check adds nothing beyond #1492's one finding.

One thing must change before this leaves draft. A cold chat turn that keeps a state is 50 ms slower than on #1492 (1.27 s against 1.22 s), and the prompt cache is on by default. Each kept state is a new device allocation (zig/src/families/lfm2_moe/cuda_snap.zig:56, about 200 MB at 10k tokens). Dropping one waits for the stream and then frees it (:70-71). GLM keeps its states in pages allocated once, inside the prompt-cache budget, so a keep or a restore makes no allocation (zig/src/families/glm/cuda_snapshot_pages.zig). Please do the same. Then measure the cold turn against #1492 again, in alternating runs; it should be level.

To your questions:

1. Pages by default is fine here. Dory's are opt-in because turning them on changed existing deployments. LFM2 has no CUDA deployment yet, and paging is exact and flat on decode.
2. A copied state is fine, in preallocated pages as above.
3. The warm pass belongs in its own PR.
4. Yes, please rebase onto main once lfm2_moe: serve LFM2.5-24B-A2B on CUDA (plain serve, shared kernels) #1492 has.

Then we'll serve it on a GB10:

• resumed against fresh over multi-turn chats;
• drafted against plain and concurrent against solo, with pages on and off;
• pool pressure and eviction;
• the cold and warm turns against lfm2_moe: serve LFM2.5-24B-A2B on CUDA (plain serve, shared kernels) #1492.

[jschmied]

Thanks @ashhart. Kept states now live in pages allocated once, and the cold turn is level with #1492. The branch is rebased onto the new #1492; the head is 0506d9c0b.

• Preallocated pages (6ef5d8401): as in GLM's cuda_snapshot_pages.zig, snap_ready allocates one arena of 2 MiB snapshot pages inside the prompt-cache budget. A keep leases pages from it and copies into them run by run, a restore copies back, and a drop returns the lease. Keeps and restores allocate nothing, and there's no free behind a stream wait. snap_charged/snap_spare report held and spare pages to the store.
• The 50 ms wasn't the allocation. With pages preallocated the cold turn was still 65 ms slower, and with keeps refused it was level. A device-to-device copy of a 200 MB state takes 1.8 ms on the GB10. nsys showed the cost was prompt passes: the turn's two chat marks sit in its last chunk and each one ended the chunk there, so it ran 12 passes instead of 10, and every pass reads the weights.
  • cb823583b taps up to four marks inside a chunk instead. The conv kernel runs once more on each mark's last row alone, writing the same y and that row's conv states into a tap buffer, and the keep copies those with the K/V rows before the mark. The kernel is unchanged.
• Cold turn against lfm2_moe: serve LFM2.5-24B-A2B on CUDA (plain serve, shared kernels) #1492, 10,144 tokens, keeping two states, 4 alternating rounds of 6: 1.211 s on lfm2_moe: serve LFM2.5-24B-A2B on CUDA (plain serve, shared kernels) #1492 and 1.216 s here. The 5 ms left is the two copies.
• Served, real weights:
  • resumed == fresh over 8 chats × 4 turns (24/32 hits), equal to cache off and to fixed pages;
  • drafted == plain and concurrent == solo with pages on and off;
  • a 40-page pool under pressure and a 0.1 GiB cache (6 evictions, 1 hit) both match;
  • token hashes are unchanged across a restart.
• The warm pass stays out of this PR.