Skip to content

Flash Next CUDA: the streams beside the first take KV pages from the core pool - #1170

Closed
jschmied wants to merge 11 commits into
ashhart:mainfrom
jschmied:fn-paged-kv-pool
Closed

Flash Next CUDA: the streams beside the first take KV pages from the core pool#1170
jschmied wants to merge 11 commits into
ashhart:mainfrom
jschmied:fn-paged-kv-pool

Conversation

@jschmied

@jschmied jschmied commented Oct 21, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Branch fn-paged-kv-pool 336cb3f7. It is #1156's ce5b78d7, then #1169 (85aeb4bf, the same change as f23b5ab5), then this change. It uses #1156's
served() and host plan.

What this changes

TF_FLASHNEXT_KV_PAGES=n is off by default; n must be a multiple of 4. One stream keeps its own caches. Every other
stream takes a pool lane (cuda_kv_lanes.zig), which holds only its page table while idle:

  • Its attention caches are n-position pages of one arena, enough for its prompt, its reply and a drafted window. With
    the MTP head a page holds the head's cache too, as the layer after the model's.
  • Its DeltaNet states, convolution and n-gram tails and the head's buffers (window streams, scores, saved rows, spare
    states; Draft.initIn) are one run of consecutive pool pages, claimed when the stream opens (Pool.claimRun) and
    given back when it closes. Only addresses move.

The host plans the pool from Loaded.layout (fixed: the table and image rows; pages a token; the run) and admits through
fits. A request waits (LaneMemory) when no lane, pages or run is free. A lane's pages and run go back once the
engine's stream drains. Prompt-cache copies of the head's cache follow its pages. The layout is made after the head is
loaded, so the run covers every buffer the head maps (initIn refuses a short run).

TF_FLASHNEXT_KV_PAGES_ALL=1 puts every stream on pages, for checks. Two ranks, --slide and a captured set refuse
the variable.

Receipt

  • At n=256, --parallel 4 --context 40960, MTP: "3 more taking 594 pages of 256 rows (4.36 GiB) as they need them,
    0 MiB fixed each". Before the run change each lane held 369 MiB from load (states, head cache and buffers).
  • With MTP drafts:
    • concurrent == alone 4/4 and serial 1/1, outputs equal to contiguous;
    • drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh (7,147 tokens), all equal to contiguous.
  • Pages only (--parallel 5, 4 lanes, 792 pages, 8 requests, some waiting): the same bytes, and a kept state with its
    head cache restored into a lane.
  • Without drafts, with the states in runs: see Gotchas (within noise of contiguous, bytes equal).
  • Tests: full zig build test -Dnvcc=/usr/local/cuda/bin/nvcc exit 0 (the admission test sets the new field), fmt
    clean, lean 0 new findings; kernel-contract 85/85, 7 controls red.

Gotchas

Checklist

  • Every token still goes through the lane rounds.
  • Drafted == plain, resumed == fresh, concurrent == solo, as measured above.
  • 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.

jschmied added 10 commits October 21, 2026 09:21
@jschmied
cuda qwen3_5: attention caches readable through a page table
3f5b87d 
Seg gains pt (an int page table), ps (elements a page) and ptok (positions a page); kv_row() gives
every key/value row address: with pt 0 the old expression (cap rows a head), else the page holding
the position, [layer][key, value][head][ptok][256] inside it. rope's store, tf_q35_attn_partial and
keep's tf_q35k_compact read and write through it. Only addresses change, never the order a sum
takes, so paged results are byte-equal to contiguous at any page size.

Cache.paged maps a cache's rows to pages; TF_QWEN_KV_PAGE_TEST=n gives each cache's own buffer a
reversed table of n-position pages (checks only). Prompt-cache saves and restores copy page runs into
the same host layout, so kept states move between paged and contiguous caches. --kv-cache fp8 keeps
contiguous rows and refuses a page table.

tf-cuda-test qwen-pages (also in kernel-contract): six shapes, cached lengths to 9000, pages of 16,
48, 128, 256 and 512 positions, two streams shuffled through one arena: rotated queries, attention
rows, rope's stores and keep's moves equal contiguous, three repeats, a float64 bound, guards and
inputs intact; red: a swapped table over placed pages, a planted row.
GB10, Qwen3.8-27B MLX-4bit: contiguous == main (3/3 chat, 3/3 8k prefill), 13.8 tok/s; page test 16
and 256 == contiguous, 13.8-14.0 tok/s; with DFlash2 at --parallel 8, drafted == plain 8/8,
concurrent == solo 8/8, resumed == fresh.
@jschmied
cuda qwen3_5: the streams beside the first take KV pages from the cor…
d5189b3 
…e pool (TF_QWEN_KV_PAGES)

TF_QWEN_KV_PAGES=n (off by default): the first stream keeps its reserved window; every other
stream takes n-position pages of one device arena from core kv_pool, as many as its prompt, reply
and one more forward need (cuda_pages.zig: layout, Pages, lanes of Cache.initPaged). Admission
goes through fits; a request with no lane or pages free waits (LaneMemory). A closed stream's
pages go back once the engine's stream drains. qwen3_5_moe serves the same pool.
TF_QWEN_KV_PAGES_ALL=1 puts every stream on pages (checks). --kv-cache fp8 and --slide refuse it.

The native host's plan: cuda_kv_pool.served() adds a wrapping drafter's per-stream buffers to each
stream's fixed part, and a family with no free stream pays its first window before pages are
planned. Nemotron (a free stream, no wrapper) plans as before.

GB10, Qwen3.8-27B MLX-4bit, n=256, --parallel 4 at 32768: 3 more streams on 387 pages (6.05 GiB)
instead of three 2.5 GiB windows; solo bytes == contiguous (3/3 chat, 3/3 8k prefill), 13.9-14.0
tok/s, c=4 48.6 (code) / 48.2 (chat) against 48.4 / 46.9 contiguous. With DFlash2 at --parallel 8:
drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh, outputs equal the page test's.
Pages only (8 requests on 7 lanes, one waiting): the same bytes, a kept state restored into a lane.
@jschmied
cuda qwen3_5: with a KV pool, warm the shared shapes on lanes opened …
6037621 
…for one row

Reservation in pool mode reserves the first stream's window and warms the 2..N-segment shapes on
pool lanes opened for one row and closed again, so the first concurrent round runs no shape cold.
GB10, Qwen3.8-27B MLX-4bit, --parallel 8, a fresh server's first round of 8 requests: TTFT max
0.97/0.97 s warmed, 0.98/0.98 unwarmed, 0.99 contiguous (qwen captures no graphs, so a cold shape
cost nothing visible); the bytes of every arm equal.
@jschmied
cuda qwen3_5: --kv-cache fp8 through the page table too
5c09685 
qwen35_kv8.cu's store, tf_q35_attn_partial8 and tf_q35k_compact8 take a row's e4m3 codes and fp32
scale through row8(): cap rows a head as before, or with a page table the row's page, laid out as a
layer's block at ptok rows (Seg.ps in bytes for fp8). Prompt-cache copies run a page at a time for
either format (the piece's layout at cap = ptok inside each page), the pool's token bytes follow the
format, and TF_QWEN_KV_PAGES no longer refuses --kv-cache fp8 (only --slide, whose learner reads
whole windows).

tf-cuda-test qwen-pages runs the six shapes and both controls for bf16 and fp8 (the fp8 reference
dequantizes e4m3 x scale in float64): paged == contiguous in every shape, controls red.
GB10, Qwen3.8-27B MLX-4bit, --kv-cache fp8: contiguous == the previous kernels (3/3 chat, 3/3 8k
prefill), 14.0/13.7 tok/s against 14.05/13.6, c=4 48.9/47.5 against 49.2/47.7; page test 16 and
pool 256 == contiguous (pool: 3 streams on 387 pages, 3.07 GiB); with DFlash2 at --parallel 8,
contiguous and pages-only both drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh, and
equal to each other.
@jschmied
cuda qwen3_5: one-line comments on row8 and the pool's reservation
9a99d1f
@jschmied
cuda qwen3_5: a pool lane's recurrent state is a run of pool pages wh…
1ac2075 
…ile its stream runs

A pool lane held its GDN recurrences and convolution tails, the one-row rollback copy, the verify
log and the drafter's tap ring from load on, whether a stream ran on it or not (558 MiB a lane on
the 27B). Now they are one run of consecutive pages in the KV arena, taken when a stream opens the
lane and given back when it closes: the lane keeps only its page table. Only the buffers' addresses
move; no kernel changes.

core kv_pool: Layout.run_bytes (0 by default, so Nemotron's plan is unchanged) adds runPages() to a
request's need, to each stream's minimum and to the window bound; Pool.findRun / claimRun /
releaseRun take k consecutive free pages from the top (the free stack hands out KV pages from the
bottom) out of the stream's promise, held as leased so the store never reclaims them. A pool too
fragmented for a run makes the request wait (LaneMemory), as one without pages does.

GB10, Qwen3.8-27B MLX-4bit, TF_QWEN_KV_PAGES=256, --parallel 4 at 32768: "3 more taking 492 pages
of 256 rows (7.69 GiB), 0 MiB fixed each" (was 387 pages and 558 MiB fixed each); solo bytes ==
contiguous 3/3 chat, 3/3 8k prefill, also pages only; 13.9/14.1 tok/s, c=4 48.8/48.3. With DFlash2
at --parallel 8, pages only, bf16 and fp8: drafted == plain 8/8, concurrent == solo 8/8, resumed ==
fresh (into a lane), outputs equal the previous runs'; 8 requests on 7 lanes, the waiting one 1.7 s.
@jschmied
cuda qwen3_5: a DFlash2 stream's ring is part of its pool lane's run;…
51f8354 
… the pool planned after the drafter

The pool's layout was planned when the family opened, before a drafter set the target's taps, so a
lane's run left out the tap ring that mapRun then mapped after the state: with --drafter, the tap
ring ran past the run (into the run above, or past the arena for the top one). The DFlash2 gates of
1ac2075 passed regardless; that was luck, not a guarantee.

Now the native host asks a family that places a drafter's buffers itself (pagedLayout) for its
layout once the drafter is attached: qwen3_5 puts the target's tap ring and the DFlash2 ring in each
lane's run and keeps the rest of the drafter's bytes a stream (its frame scratch's share, an MTP
head's buffers) fixed. mapRun refuses a run shorter than its parts (QwenRunTooShort). The drafter
borrows a lane's ring for a stream on a pool lane (a context with a borrowed buffer) and keeps its
own rings only for the first stream.

GB10, Qwen3.8-27B MLX-4bit, TF_QWEN_KV_PAGES=256, DFlash2, --parallel 4 at 32768: "3 more taking 522
pages (8.16 GiB), 23 MiB fixed each" (was 169 MiB fixed each); bytes == contiguous (3/3 chat, 3/3
8k prefill, pages only); DFlash2 decode code/chat 59.9/44.4 tok/s pool, 60.3/44.5 pages only
(--parallel 5), 60.5/44.6 contiguous; c=4 161.5/123.1, 164.2/123.4 against 164.2/123.6. Pages only
at --parallel 8, bf16 and fp8: drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh, equal
to the earlier runs; no run refused.
@jschmied
cuda qwen3_5: the attention partials test the page table once a block…
ce5b78d 
…, not once a key

On Qwen3.6-35B-A3B (decode at 84 tok/s, attention a larger share of a step than on the 27B) the
contiguous path paid kv_row's page-table test inside the key loop: an ABAB against main measured
-2.7..-3.6 % at one stream and -1.5..-2.1 % at four. tf_q35_attn_partial and tf_q35_attn_partial8
now branch once: without a table they take main's K and V (and scale) bases and step a row at a
time, with one they read each row through kv_row / row8. The addresses are the same, so the bytes.

GB10, Qwen3.6-35B-A3B MLX-4bit (MTP), --parallel 4 at 32768, ABAB against main, code/chat tok/s:
main 84.2/75.0 and 83.3/75.1, this 82.7/74.5 and 84.0/75.2 (c=4 163.7/140.4, 163.2/143.8 against
161.0/142.1, 163.5/143.7); fp8 84.0/75.1 main, 84.4/75.3 this. Bytes == main (3/3 chat, 3/3 8k
prefill) for the 35B in bf16 and fp8 and for the 27B; pages only == main; tf-cuda-test qwen-pages
passes in both formats with its controls red.
@jschmied
cuda flashnext: attention caches readable through a page table
85aeb4b 
flashnext_attn.cu's cache kernels take a Pg (an int page table, the elements a page holds, the positions
a page holds): fn_attn_prep(_media)'s key, value and indexer-key stores, fn_qsa_pool(_media)'s reads and
pooled writes, fn_qsa_scores' pooled reads and fn_attn_chunks(_ahead)'s key and value tiles find a row
with pg_at(): the old expression when the table is null, else the page holding the position plus its
row there. fn_attn_chunks looks a key's row up once a tile (key_at, beside key_ids), the pool once a
block of four. A page of n positions holds, for each attention layer, its keys, values and indexer
keys of n positions and the pooled rows of their n / 4 blocks (cuda_pages.zig); a cache pointer is its
layer's part of page 0. Only addresses change, never a sum's order, so paged caches give contiguous
bytes at any page size that holds whole blocks. Without our glue (a captured Triton set) a page table
is refused.

TF_FLASHNEXT_KV_PAGE_TEST=n gives each stream's caches one buffer read through a reversed table of
n-position pages (checks only). Prompt-cache saves and restores copy page runs into the same layout,
so kept states move between paged and contiguous caches.

tf-cuda-test flashnext-pages (also in kernel-contract): six cases, cached lengths to 9000, pages of 16,
48, 128, 256 and 512 positions, two streams shuffled through one arena: queries, indexer queries,
block scores, dense and sparse attention rows, the stores and pooled rows equal contiguous over three
repeats; attention rows within a float64 bound; guards intact. Red: a swapped table over placed pages,
a planted row. Records: the sm_80..121 registers and static bytes of the kernels that moved (nvcc 13.0).

GB10, nvidia/Qwen3.8-Flash-Next-FP4-Local, --parallel 4 --context 40960, --no-drafts, ABAB against
main: c=1 23.51/23.49 and 23.55/23.44 tok/s against 23.46/23.40 and 23.38/23.48; 8k 5.97-6.91 s against
6.03-7.00 s, 31k 25.64/25.74 s against 25.67/26.01 s; replies equal (3 chats, three ~8k, one 31k).
Page test 16: the same bytes, 23.56/23.38 tok/s, 31k 26.08 s. With MTP drafts: drafted == plain 8/8,
concurrent == solo 8/8, resumed == fresh (7,147 cached tokens restored into pages), equal to contiguous.
@jschmied
cuda flashnext: the streams beside the first take KV pages from the c…
336cb3f 
…ore pool (TF_FLASHNEXT_KV_PAGES)

TF_FLASHNEXT_KV_PAGES=n (off by default, n a multiple of 4): one stream keeps its own caches as
before; every other stream takes a pool lane (cuda_kv_lanes.zig). A lane holds only its page table
while idle. Its stream's attention caches are n-position pages of one arena, as many as its prompt,
reply and a drafted window need; with the MTP head a page holds the head's cache too, as the layer
after the model's. Its DeltaNet states, convolution and n-gram tails and the head's buffers (window
streams, scores, saved rows, spare states; Draft.initIn) are one run of consecutive pool pages,
claimed when the stream opens (Pool.claimRun) and given back when it closes, so only addresses move.
The native host plans the pool from Loaded.layout (fixed: the table and image rows; a token's pages;
the run) and admits through fits; a request with no lane, pages or free run waits (LaneMemory). A
lane's pages and run go back once the engine's stream drains. Prompt-cache copies of the head's
cache follow its pages. TF_FLASHNEXT_KV_PAGES_ALL=1 puts every stream on pages (checks). Two ranks,
--slide and a captured Triton set refuse it.

GB10, nvidia/Qwen3.8-Flash-Next-FP4-Local, n=256, --parallel 4 --context 40960, MTP: "3 more taking
594 pages of 256 rows (4.36 GiB) as they need them, 0 MiB fixed each" (was 369 MiB fixed each, the
head's cache and buffers and the states held from load). With MTP: concurrent == alone 4/4, serial
1/1, outputs equal to contiguous; drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh
(7,147 tokens), equal to contiguous. Pages only (--parallel 5: 4 lanes, 792 pages, 8 requests, some
waiting): the same bytes, a kept state (and its head cache) restored into a lane. Without drafts:
solo bytes == contiguous, c=4 65.0/64.5 against 65.4/65.0 (measured before states moved into runs).
This was referenced Oct 21, 2026
Closed
Closed
@jschmied
cuda qwen3_5: the reuse tests' caches say they have no page table
e0e69c5 
Two cuda_reuse tests build a Cache as undefined and set only the fields they meant to read. Since
the page table (3f5b87d), restoreCache also reads Cache.paged to decide how to copy the rows, so
those tests read undefined memory: on CI's x86_64 Debug run the bytes formed a page table and
piece()'s offsets overflowed ("a failed restore's host copy stays through a drop ..."). On aarch64
the same bytes read as null, so the test passed here (with the stack filled with 0xaa it still
reads null: the optional's null is a niche value of its bool). Both tests now set paged = null.
@jschmied

jschmied commented Oct 21, 2026

Copy link
Copy Markdown
Contributor Author

The zig job failed on a qwen3_5 test this PR carries from #1156's branch: two cuda_reuse tests built a Cache as undefined and never set paged, so restoreCache read undefined memory. In CI's x86_64 Debug build the bytes looked like a page table and piece() overflowed. e0e69c52 sets paged = null in both tests (the same commit as 0ceb61f2 on #1156's branch). zig build test-cuda-host passes locally; the bytes never looked like a page table on aarch64, so CI is the real check.

@jschmied jschmied mentioned this pull request Oct 21, 2026
Closed
@ashhart

ashhart commented Oct 22, 2026

Copy link
Copy Markdown
Owner

Thanks @jschmied for the pool lanes, for moving each lane's states into a run so an idle lane holds only its table, and for tracking down the reuse tests' CI failure yourself in e0e69c52. The eight Qwen commits this branch carries are already on our KV pool line for the next release, kept as yours (see #1156), so we'll carry the three Flash Next commits on top of them and #1169, kept as yours, with a separate commit of ours that folds in the Flash Next pool we had built in parallel there. That commit routes the pools' close, failed-open and constructor error paths through cuda.completion.drain before anything goes back to the pool or is freed (for example zig/src/families/flashnext/cuda_kv_lanes.zig:236), runs the pool's fills on the engine's stream, and keeps zig/src/families/flashnext/cuda_lanes.zig and zig/src/native/cuda.zig under 600 lines after the rebase. Then it goes to the GB10 kernel contract and served-equality runs on Flash Next, pages against contiguous, and lands with that line. You need not rebase; if you are in the middle of a change, please tell us so we do not collide.

- Synthema, Ash's AI Assistant.

ashhart added a commit that referenced this pull request Oct 25, 2026
@ashhart
Maintainer fix-up for #1170: port paged paths and prove pool ownershi…
7deb9e9 
…p before reuse

Retain current grouped, fused and one-buffer routes with dedicated paged launches. Check page dimensions and admission, keep constructor fills on the engine stream, and drain uploads, failed opens and closes before releasing their owners.

Add initialized host fault fixtures and numerical, media, snapshot and accumulator contract witnesses. Keep unknown paged compiler records pending. Preserve the three Qwen byte pins only where they exactly match the earlier reviewed kernel blobs.
This was referenced Oct 25, 2026
Closed
Closed
Closed
Closed
Closed
Closed
@ashhart

ashhart commented Oct 25, 2026

Copy link
Copy Markdown
Owner

The Flash Next slice of this PR landed with #1169 in 93a215b..35e46c0. Tokens are identical on Flash Next FP4 and MLX, and prefill is up 24-29% on FP4 with the rest of that group. The Qwen slice, with the page-table fix that recovered the 27B prefill, didn't compile under real nvcc on our GB10 (qwen35_kv8.cu: widen8 undefined after the rebase). We're fixing that, and I'll close this PR when the Qwen slice lands.

- Synthema, Ash's AI Assistant.

ashhart added a commit that referenced this pull request Oct 25, 2026
@ashhart
Maintainer fix-up for #1170: port paged paths and prove pool ownershi…
793d191 
…p before reuse

Retain current grouped, fused and one-buffer routes with dedicated paged launches. Check page dimensions and admission, keep constructor fills on the engine stream, and drain uploads, failed opens and closes before releasing their owners.

Add initialized host fault fixtures and numerical, media, snapshot and accumulator contract witnesses. Keep unknown paged compiler records pending. Preserve the three Qwen byte pins only where they exactly match the earlier reviewed kernel blobs.
@ashhart

ashhart commented Oct 25, 2026

Copy link
Copy Markdown
Owner

Both slices have now landed. The Flash Next slice went in with #1169 (93a215b..35e46c0). The Qwen slice, the streams beside the first taking KV pages from the core pool (TF_QWEN_KV_PAGES), went in now in bcfaea6..8086c20, with your records under your authorship. Thank you, @jschmied.

Four maintainer children on the Qwen side:

  • 27B prefill: restored main's contiguous tensor-core attention for flat batches. Only paged or mixed batches take the page-run path, which fixed an earlier 27B prefill loss.
  • Compile fix: restored the scalar E4M3 widen8 that qwen35_kv8.cu needs under real nvcc.
  • File size: split the grown cuda_native.zig adapter.
  • Contract: listed the two new legacy attention launches as awaiting resource receipts.

On a GB10 with Qwen3.8-27B MLX 4-bit, alternating base/new:

  • Tokens: identical (0 of 20 rows differ), including concurrent against solo.
  • Prefill: flat over three alternating runs each at 2k and 8k (1280.9 → 1280.4, 1260.0 → 1258.5), and within 0.3% to 64k.
  • Sweep: flat at 1, 2, 4 and 8 streams.

- Synthema, Ash's AI Assistant.

@ashhart ashhart closed this Oct 25, 2026
@BobClawblaw BobClawblaw mentioned this pull request Oct 26, 2026
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@jschmied @ashhart