[jschmied]

Branch fn-paged-kv f23b5ab5 on main d1018b31. Part of #781's follow-ups (other families on the core pool).

What this changes

The seven flashnext_attn.cu kernels that touch a stream's caches now take a Pg argument: an int page table, the
elements in a page, and the positions in a page. They are fn_attn_prep(_media), fn_qsa_pool(_media),
fn_qsa_scores and fn_attn_chunks(_ahead). pg_at() finds each row:

• with no table, it is the old expression;
• with a table, it is the start of the page that holds the position, plus the row inside that page.

A page of n positions holds, for each attention layer:

• the keys, values and indexer keys of its n positions;
• the pooled rows of their n / 4 blocks.

A cache pointer is its layer's part of page 0 (layout in cuda_pages.zig). Only addresses change; no sum changes
order. So paged caches give the contiguous bytes at any page size that holds whole blocks.

• TF_FLASHNEXT_KV_PAGE_TEST=n reads each stream's caches through a reversed table of n-position pages. It is for
  checks only.
• fn_attn_chunks looks up a key's row once per tile (key_at in shared memory, beside key_ids), and the pool kernel
  once per block of four; the contiguous path pays one uniform branch.
• Prompt-cache saves and restores copy page runs in the contiguous layout, so kept states move between the two forms.
• A captured Triton set refuses a page table.

Receipt

• Environment: TensorFold main d1018b31 plus this branch; one DGX Spark (GB10, sm_121, CUDA 13.0, nvcc 13.0.88);
  nvidia/Qwen3.8-Flash-Next-FP4-Local@b9ce18a6.

• Kernel contract: tf-cuda-test flashnext-pages, also run inside kernel-contract.

  • Six cases: cached lengths up to 9,000; pages of 16, 48, 128, 256 and 512 positions; two streams shuffled through
    one arena.
  • Paged output equals contiguous over three repeats for: queries, indexer queries, block scores, dense and sparse
    attention rows, the stores, and the pooled rows.
  • Attention rows stay within a float64 bound, and the guards are intact.
  • Controls are red: a swapped table over placed pages, and a planted row.
  • Gate 5: new nvcc 13.0 records for the five kernels whose register counts moved. All seven pass live. The 13.2
    override set still needs regenerating with nvcc 13.2, which is not on this box.

• Served, greedy, --parallel 4 --context 40960 --no-drafts, ABAB against main (c=1 tok/s code/chat, TTFT):

  arm	replies	c=1	~8k	31k
  main, run 1 / 2	-	23.46/23.40, 23.38/23.48	6.04-7.00 s	25.67 / 26.01 s
  this, contiguous, run 1 / 2	equal (3 chats, three ~8k, one 31k)	23.51/23.49, 23.55/23.44	5.97-6.91 s	25.64 / 25.74 s
  this, page test 16	equal	23.56/23.38	6.05-6.90 s	26.08 s

  The page test's first version (one page-table lookup per 16-byte load) cost +3-4 % at 31k; one lookup per key per tile
  removes that.

• Gates with MTP drafts, page test 16, compared with contiguous:

  • drafted == plain: 8/8;
  • concurrent == solo: 8/8;
  • resumed == fresh: yes, with 7,147 tokens restored into pages;
  • every output equals contiguous.

• Tests: full zig build test -Dnvcc=/usr/local/cuda/bin/nvcc exit 0, zig fmt --check clean, lean check 0 new
  findings; kernel-contract: 85/85 kernels pass gates 1-4, 7 controls red (exit 1 only from main's 877 limit
  problems).

Gotchas

• Off by default: only the check switch TF_FLASHNEXT_KV_PAGE_TEST uses pages here. The pool itself is the stacked follow-up, Flash Next CUDA: the streams beside the first take KV pages from the core pool #1170.
• Gate-5 records: the nvcc 13.2 override set for the five kernels whose registers moved isn't regenerated, because there is no nvcc 13.2 here. The 13.0 records are new.
• Captured Triton sets refuse a page table.
• Overlap: [Draft] EXL3 Flash Next on CUDA: TP=2, BF16 paging, cache reuse, graphs #1113 pages Flash Next KV to host RAM for its EXL3 path. That is a different mechanism, but the same family, so expect textual overlap. No conflict with Flash Next CUDA: NVFP4 experts on prompt tiles with the decode form's sums, bit for bit #1143, Flash Next CUDA: no zero fills the one-rank ModelOpt pass reads back unchanged #1147, Flash Next CUDA: a shared round's greedy streams draft their MTP heads together #1155 or Flash Next CUDA: the hyper-connection normalize over 256-column slices #1164.

Checklist

• Every token still goes through the lane rounds.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. All
  three were measured on the page test.
• No precision traded for speed. The bits do not change.
• New and changed kernels are registered with the kernel contract, with red controls.
• Prompt processing is no slower than the last release (contiguous within main's spread).
• Platforms: CUDA on GB10 (sm_121), one rank; other SMs through recorded compiles; no Metal change.
• Comments are one line. No measurements or history in the source (lean check 0 new).
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[ashhart]

Thanks @jschmied for the page table, for keeping every sum in its old order so paged bytes equal contiguous ones, and for the swapped-table and planted-row controls. While this was open we built a Flash Next page pool of our own on the KV pool line for the next release, along the same lines as your Qwen pages; since yours came first, we'll take this PR and #1170 as the base, kept as yours, and fold what ours adds in a separate commit on top. That commit also threads the page table through the attention paths main has gained since your base (fn_attn_chunks_one and fn_attn_fused in zig/kernels/cuda/flashnext_attn.cu, four-head QSA through qsa_tile::scores, and the grouped prompt loop in cuda_batch.zig), adds paged cases for the image prep and pooling kernels, runs the page suite through the kernel contract's five gates, and makes testPages (zig/src/families/flashnext/cuda_pages.zig:105) allocate before its zero fill and drain before any free. Then it goes to the GB10 kernel contract and served-equality runs, pages against contiguous, and lands with that line. You need not rebase; if you are in the middle of a change here or on #1170, please tell us so we do not collide.

- Synthema, Ash's AI Assistant.

[ashhart]

Landed on main in 93a215b..35e46c0, with your records under your authorship plus the maintainer fix-ups. Thank you. Measured with the whole group (#1139, #1140, #1143, #1147, #1164, #1169, #1171 and #1170's Flash Next slice) against main on GB10s:

• Qwen3.8 Flash Next FP4-Local: tokens identical (0 of 20 rows differ). Prefill +23.7% to +28.9% from 2k to 64k (16k: 1356.6 → 1748.4 tok/s). Decode +1.0% to +2.2% at 1, 2, 4 and 8 streams.
• Flash Next MLX 4-bit, alternating base/new/new/base: tokens identical, prefill +0.4% to +0.8%, decode within noise.

- Synthema, Ash's AI Assistant.