cuda flashnext: --learn-conversations keeps Flash Next's states on disk (#840) by jschmied · Pull Request #1359 · ashhart/TensorFold

jschmied · 2026-10-26T05:55:01Z

Split at review: turning on CUDA --spill-gib / --checkpoint-slots moved to #1413; this PR is now only Flash Next's disk image.

Part of #840, the "Spill of evicted states" row (#68, @gilby; claimed in #840, comment 6399090685). Stacked on #1358, whose kept_disk hook and grid fix it reuses; the last two commits are new here.

Flash Next on CUDA refused --learn-conversations. Its kept state (cuda_snap.zig) is a device buffer. native/cuda_flashnext_disk.zig lays it out as the rows store's image:

32 bytes of words: the n-gram history, the position, and whether the MTP head's rows follow.
The 36 DeltaNet conv tails and recurrent states, as the state's own parts.
For each of the 12 attention layers: the key rows, the value rows, the indexer-key rows and the pooled indexer rows (four rows to one), all as row parts in 256-token blocks.
The PLE conv tail, as an own part.
With the MTP head: the head's key, value, indexer and pooled rows.

The copies go through the engine's stream, DtoH out and HtoD in, and a state read back is a new device buffer. The flag is refused with --tp, because rank states stay on their ranks. A state kept from a stream that did not draft lacks the head's rows, so it stays in RAM. The identity covers the checkpoint, the kernels, the device, the driver, whether the MTP head drafts, and the hash of a probe state at row 4136.

Kept state and the three parts that restoring a context depends on (SGLang #46422):

QSA indexer cache: the ikc and pooled rows, written as row parts.
PLE conv state: tail, written as an own part.
n-gram history: history, written in the words.
MTP head: its cache rows are part of the image whenever the server drafts.

Red controls: a gate-only build (not part of this PR) zeroes one part after each restore (TF_GATE_DROP=history|index|ple). Each one must break resumed == fresh:

dropped on restore	text == fresh	logits hash == fresh
n-gram history	6/6	0/6
indexer cache (`ikc` + `pooled`)	5/6	0/6
PLE tail	5/6	0/6
nothing (fn-bf16-contig)	6/6	6/6

The text alone misses some of these breaks, so the logits hash is required.

Gates (GB10, nvidia/Qwen3.8-Flash-Next-FP4-Local, MTP head drafting)

The method is the same as in #1358. Each comparison checks the tokens and a hash of the last prompt row's logits from a gate-only log line. Flash Next keeps bf16 caches only, so it has no fp8 row.

Gate	Flash Next
read from disk after restart == fresh (contiguous)	2/2 + 4/4
read from disk after restart == fresh (`TF_FLASHNEXT_KV_PAGES=256`, `_ALL=1`)	2/2 + 4/4
written contiguous, read by a paged-all server (same identity, so it reads) == paged fresh	2/2 + 4/4
resumed at c=1/2/4 == solo	10/10
fresh at c=1/2/4 == solo	20/20
drafted (MTP head) == `--no-drafts`	10/10
`TF_IMPRINT_VERIFY` read-back	byte-equal
corrupt `.blk` / `.st`	refused (RowsHash), recomputed, == fresh
`--learn-gib 0.8`	786 MB, then 642 MB on disk; == fresh
`--learn-daily-gib 0.5`	DailyBudgetSpent, RAM only; == fresh

Turn 2 of the 9.6k-token conversation: 0.49 s read back (390 MiB in 0.15 s) vs 5.32 s fresh.

Turn-2 TTFT of a 29,455-token chat after a restart (2 disk + 2 cold, alternating): 1.37 s and 1.96 s read back from disk (29,426 tokens, 974 MiB in 0.38 to 0.50 s), against 16.6 s and 18.4 s cold. All four runs gave the same text and logits hash.

jschmied · 2026-10-26T15:27:08Z

Follow-up on the exported shape: cuda_snap.move used to build its pieces and sizes inline, while the disk image (native/cuda_flashnext_disk.zig) restated them from shape, and only their total was tested. Now cuda_snap.order(at, head) is move's own table: for each piece, its layers, bytes per layer, row and pooled block, in move's order. move copies by that table and refuses a sequence whose buffers or page rows differ from it (SnapShape). A new test checks the image's parts against the table part by part, at 0, 1, 3, 4, 255, 256, 1000, 4136 and 4137 tokens, with and without the MTP head:

sizes, rows, pooled blocks and running offsets;
the attention K/V rows, the indexer key rows, the pooled indexer rows (complete blocks only), the DeltaNet conv tails and recurrent states, the PLE tail, and the MTP head's four pieces.

All of them already matched; no byte of the image changes. The Flash Next probe gets the same single cleanup path as the qwen3_5 one. Both changes are in one new commit, "cuda flashnext: a kept state's pieces in one table, and the disk probe releases its stream once".

Commits in this PR. Besides the review commit above, this PR adds two commits over #1358 (qwen-kv-disk-qwen). "cuda flashnext: --learn-conversations keeps Flash Next's states on disk" adds native/cuda_flashnext_disk.zig (the image, the device copies, the probe and the identity facts), exports the shape in families/flashnext/cuda_snap.zig, attaches through flashnext/cuda_launch.zig, and touches native/cuda_disk.zig (keptDisk) and native/cuda_refusal.zig (--tp refused with the flag). "cuda: serve --spill-gib and --checkpoint-slots for Nemotron, the Qwen hybrids and Flash Next" changes native/cuda_disk.zig (opening through rows.attach with the spill mode and slot count), native/cuda.zig and native/cuda_refusal.zig (a family without a disk image refuses --spill-gib; --checkpoint-slots without a disk store is refused).

Gates: zig build test passes (the new parts test included), and lean shows 0 new problems. Flash Next's disk gate did not run again: the checkpoint needs about 116 GB free and the box had 85 GB. No byte of the image or of move's copies changes; the image's size and order are checked on the host, and move still asserts its total.

ashhart · 2026-10-26T15:46:47Z

Review is on #1355, covering the whole stack #1355–#1359.

jschmied · 2026-10-26T18:40:30Z

Split: this PR is now only the Flash Next CUDA disk image (two commits on #1358). Turning on CUDA spill moved to #1413.

Flash Next declares kept_disk in its family namespace (flashnext/cuda_lanes.zig), as qwen3_5 does. native/cuda_disk.keptDisk no longer names a family: return if (@hasDecl(F, "kept_disk")) F.kept_disk else null;.
The image moved into the family as families/flashnext/cuda_disk.zig, beside cuda_snap.zig. The flashnext module gets the server's engine_api from build/cuda.zig's family(), the same one qwen35 and the server link, so the rows store's types are one type set.
cuda_lanes.zig was at 600 lines; it drops its unused cuda_config import to make room.
The order table and the probe's single cleanup path (the second commit) are unchanged.

zig build test passes on this branch and lean shows 0 new problems. The Flash Next disk gate still needs about 116 GB free beside the box's other work; I did not run it on this split.

Against #1411 (GLM-5.3-Flash): this PR adds nothing beyond #1358's four one-line conflicts, which are described there. kept_disk takes an optional attach, so GLM's per-rank image can be declared the same way.

ashhart · 2026-10-27T14:23:29Z

Thanks for the split, and for moving kept_disk into the family. keptDisk no longer names anyone, and the order
table tested part by part is a good guard. Three small things before we run it:

zig/build/cuda.zig reaches 602 lines on main (601 on top of cuda qwen3_5: --learn-conversations keeps the 27B's and 35B's states on disk (#840) #1358 with cuda qwen3_5: the prompt cache keeps KV pages on the device (TF_QWEN_KV_SHARE, #1156) #1357's one-line import), so please give a
line back. Moving flashnextModule and engines into their own build file would also make room for the other open
PRs that touch this file.
The description still contains the "Second commit: --spill-gib and --checkpoint-slots" section, which now belongs
to cuda: serve --spill-gib and --checkpoint-slots for Nemotron, the Qwen hybrids and Flash Next (#840, split from #1359) #1413.
The --tp refusal in cuda_refusal.zig still names Flash Next. A flag on its kept_disk would keep the refusal
in the family.

Your device receipts are from the pre-split head. We'll run the write, restart and read cell (contiguous, paged, and
the cross read) on this head.

jschmied · 2026-10-27T16:29:06Z

Thanks. Restacked on #1358's new head (402e64fbd, on main 59044397a) and pushed as 2f8d48d0d. The force-push is from the restack.

A line back in cuda.zig. flashnextModule and engines now live in zig/build/cuda_engines.zig, together with glmCuda and lfm2Module, which engines needs. zig/build/cuda.zig is 548 lines (was 601), which leaves room for the other open PRs that touch it.
The --tp refusal. cuda_refusal.zig no longer names Flash Next. Flash Next's kept_disk (Kept(F)) declares one_rank = true. Any family whose kept states declare it is refused with --tp plus --learn-conversations. Host tests cover a ranked family with the flag (refused under --tp, fine without it) and one without (fine under --tp). The Flash Next case still refuses.

Checks: lean 0 new; zig build test passes on this head. ATC shows cuda.zig conflicts with BlivionIaG's EXL3 Flash Next stack (#1441–#1446). Those come from the flashnextModule(..., api) argument this PR already had, not from the move.

ashhart · 2026-10-27T17:35:17Z

Thanks for the split; the description now says spill moved to #1413. Two things before it can land:

Unfinished copies can escape on error. Kept.copy queues a copy for each range, but a later try can return before the final synchronization (zig/src/families/flashnext/cuda_disk.zig:44, :62, :66), and the rows store may then reuse that staging buffer (core/rows_read.zig:54, core/rows_write.zig:193). With a host driver fake that accepts the first transfer and rejects the second, both directions return with copies still pending ("expected 0, found 2"). Draining the stream on the error path, as the family's other cleanup does, fixes it, with a regression that fails before and passes after (the persistence suite goes from 17/18 to 18/18). We have that as a small separate commit and are happy to carry it, or you can take the same shape.
zig/build/cuda.zig is at 602 lines on this head (one new lean finding). Taking cuda qwen3_5: --learn-conversations keeps the 27B's and 35B's states on disk (#840) #1358's latest commit and removing one blank line between the CUDA aliases brings it to 600, with no build change.

After that, the remaining gate is a current-head device run on a GB10 (the earlier receipts predate the split).

jschmied · 2026-10-27T17:48:51Z

Thanks, and good catch on the pending copies.

Unfinished copies on error. Please carry your commit; the drain on the error path plus the fake-driver regression is the right shape, and I'd rather not write a second copy of it.
zig/build/cuda.zig. On this PR's head 2f8d48d0d it is 548 lines: flashnextModule, engines, glmCuda and lfm2Module moved to the new zig/build/cuda_engines.zig (60 lines). Merged with main 6eaa4cd5c it stays at 548. If your lean still reports 602 on that head, could you say which file and commit it reads? I can't reproduce it here.

jschmied · 2026-10-27T21:13:12Z

Restacked on #1358's rebased head (f593a527e, on main 5897e0bca) as e018f176e, with no conflicts. Lean 0 new, zig build test -Dnvcc passes, and zig/build/cuda.zig is 548 lines. Your drain-on-error commit isn't on the branch yet; it should apply on top unchanged.

ashhart · 2026-10-27T21:27:57Z

Thanks for the quick restack. e018f176e sits on #1358's rebased head on current main, zig/build/cuda.zig is 548 lines, and lean finds nothing new, so our 602 didn't come from this head. The one_rank flag and its tests read well.

We'll add the drain-on-error commit on top, with its fake-driver regression. The head still returns from Kept.copy with copies queued (zig/src/families/flashnext/cuda_disk.zig:53, :64 and :68), and on current main the gather's buffer is an arena that the writer frees on error (zig/src/core/rows_write.zig:243), so we want it in before this lands.

Then we'll run the write, restart and read cell on a GB10 (contiguous, paged and the cross read), with drafted against plain and concurrent against solo.

Flash Next on CUDA refused --learn-conversations. Its kept state (cuda_snap.zig) is one device buffer, and flashnext/cuda_disk.zig lays it out as the rows store's image: - 32 bytes of words: the n-gram history, the position, and whether the MTP head's rows follow. - The DeltaNet conv tails and recurrent states, as the state's own parts. - Each attention layer's key, value, indexer-key and pooled indexer rows, as row parts in 256-token blocks. - The PLE conv tail. - With the head: its key, value, indexer and pooled rows. So the QSA indexer cache, the PLE state and the n-gram history all go with a state. The copies run on the engine's stream; a state read back is a new device buffer, restored as any other, and one whose words name another position or head is refused. The family declares kept_disk, as qwen3_5 does, and the server's engine_api reaches the flashnext module for it (build/cuda.zig: family() hands it out; cuda_lanes.zig drops its unused cuda_config import to stay within 600 lines). Its identity covers the checkpoint, the kernels, the device, the driver, whether the head drafts, and the hash of a probe state at row 4136. A state from a stream that did not draft lacks the head's rows this server's image names, so it stays in RAM. --tp is refused with the flag, because rank states stay on their ranks. Co-authored-by: gilby <880013+gilby@users.noreply.github.com>

…e releases its stream once cuda_snap.move built its pieces and their sizes in place, and the disk image (cuda_disk.zig) restated them from `shape`; only their total was tested. `order(at, head)` is now move's table: each piece's layers, bytes a layer, row and pooled block, in move's order. move copies by it and refuses a sequence whose buffers or page rows differ from it (SnapShape); a test checks the image's parts against it part by part (sizes, rows, blocks, offsets) at 0 to 4137 tokens, with and without the MTP head's rows. The probe, as the qwen3_5 and Nemotron ones: the pass and its wait in one block whose defer releases the stream, a kept state dropped by a defer of its own.

… Flash Next's one-rank refusal is its states' flag flashnextModule, engines, glmCuda and lfm2Module move from zig/build/cuda.zig to zig/build/cuda_engines.zig, which gives cuda.zig 53 lines back (548). cuda_refusal.zig no longer names Flash Next for --tp with --learn-conversations: a kept_disk that declares one_rank (Flash Next's, whose states are one tensor rank's) is refused under --tp, any family's.

…he copies it queued Kept(F).copy (flashnext/cuda_disk.zig) queues a device copy for each range of a kept state's image, cuMemcpyDtoHAsync out and cuMemcpyHtoDAsync in, and waits for the stream once, after the last. When the driver refused a later copy, or a later range ran past the state (BadKeptRange), copy returned with the earlier copies still queued on `host`, the rows store's staging. The writer frees that staging with its arena on the error (rows_write.zig, stageJob), so a copy out still in flight could write freed memory. A copy in was drained only because the reader drops the blank state it was filling. copy now drains the stream on every error, as cuda_snap.save does: an errdefer cuda.completion.drain waits for the stream, else the context, else exits. The test is named "Flash Next completion:", so test-flashnext-completion-host runs it, as does test-cuda-host. It drives the image's gather and scatter with a fake driver whose copies queue the first transfer and refuse the next, and then with a second range past the state's bytes. Each of the four calls returns its error with one copy queued and none left pending. Without the drain, each leaves one pending (expected 0, found 1). Co-authored-by: jschmied <600316+jschmied@users.noreply.github.com>

…sk (#1359)

ashhart · 2026-10-28T03:38:55Z

Merged, thank you @jschmied, and @gilby for your part. It ships in 1.1.8, with the changelog entry in 08f1539.

To land it as a merge after #1358, we rebased your branch onto main (your three commits as they were) and added one commit of ours on top, crediting you: when a kept state's disk copy fails part way, it now drains the copies it already queued.

Our device run on a GB10 with Qwen3.8-Flash-Next (MTP), with contiguous and with paged KV: for a 16,384-token and a 14,000-token conversation, a server with --learn-conversations wrote the state, and a fresh server on the same folder read it back after the restart (15,939 of 15,946 prompt tokens cached at 16k, first reply in 0.36-0.44 s against about 10 s fresh). Every reply was byte-equal to main's fresh reply. The full test suite, the Linux CUDA cross-compiles and the release path check pass on the merged tree.

ashhart · 2026-10-28T07:55:04Z

This is in TensorFold 1.1.8, released today. Thank you, @jschmied and @gilby!

jschmied force-pushed the qwen-kv-disk-spill branch from 56fba04 to c903158 Compare October 26, 2026 08:40

ashhart mentioned this pull request Oct 26, 2026

GLM-5.3-Flash on CUDA cannot keep conversation states on disk: an 18.7k-token conversation evicted by one 60k filler recomputes its next turn in 15.6 s (0.49 s resumed), and 15.4 s after a restart #1312

Open

jschmied force-pushed the qwen-kv-disk-spill branch from c903158 to a6d5ee5 Compare October 26, 2026 15:26

wojo mentioned this pull request Oct 26, 2026

GLM-5.3-Flash on CUDA keeps conversation states on disk, every rank's part (--learn-conversations) #1411

Open

jschmied force-pushed the qwen-kv-disk-spill branch from a6d5ee5 to 7a83b05 Compare October 26, 2026 18:39

jschmied mentioned this pull request Oct 26, 2026

cuda: serve --spill-gib and --checkpoint-slots for Nemotron, the Qwen hybrids and Flash Next (#840, split from #1359) #1413

Merged

jschmied changed the title ~~cuda flashnext: --learn-conversations keeps Flash Next's states on disk, and CUDA serves --spill-gib (#840)~~ cuda flashnext: --learn-conversations keeps Flash Next's states on disk (#840) Oct 26, 2026

BlivionIaG mentioned this pull request Oct 26, 2026

EXL3 Qwen3.8-Flash-Next on CUDA sm_120: expert offload, TP=2, graphs and MTP (pilot; the chain that replaces draft #1113) #1419

Open

jschmied force-pushed the qwen-kv-disk-spill branch from 7a83b05 to 2f8d48d Compare October 27, 2026 16:27

jschmied force-pushed the qwen-kv-disk-spill branch from 2f8d48d to e018f17 Compare October 27, 2026 21:12

jschmied and others added 4 commits October 28, 2026 02:33

ashhart force-pushed the qwen-kv-disk-spill branch from e018f17 to fe6e62c Compare October 28, 2026 03:38

ashhart merged commit 968cf27 into ashhart:main Oct 28, 2026
8 of 9 checks passed

ashhart added a commit that referenced this pull request Oct 28, 2026

Changelog: --learn-conversations keeps Flash Next's CUDA states on di…

08f1539

…sk (#1359)

ashhart mentioned this pull request Oct 28, 2026

Python engine features to carry into the native engine #840

Open

Navigation Menu

Repository navigation

cuda flashnext: --learn-conversations keeps Flash Next's states on disk (#840) - #1359

cuda flashnext: --learn-conversations keeps Flash Next's states on disk (#840)#1359
ashhart merged 4 commits into
ashhart:mainfrom
jschmied:qwen-kv-disk-spill

jschmied commented Oct 26, 2026 •

edited

Loading

Uh oh!

jschmied commented Oct 26, 2026

Uh oh!

ashhart commented Oct 26, 2026

Uh oh!

jschmied commented Oct 26, 2026

Uh oh!

ashhart commented Oct 27, 2026

Uh oh!

jschmied commented Oct 27, 2026

Uh oh!

ashhart commented Oct 27, 2026

Uh oh!

jschmied commented Oct 27, 2026

Uh oh!

jschmied commented Oct 27, 2026

Uh oh!

ashhart commented Oct 27, 2026

Uh oh!

Uh oh!

ashhart commented Oct 28, 2026

Uh oh!

ashhart commented Oct 28, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

jschmied commented Oct 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Gates (GB10, nvidia/Qwen3.8-Flash-Next-FP4-Local, MTP head drafting)

Uh oh!

jschmied commented Oct 26, 2026

Uh oh!

ashhart commented Oct 26, 2026

Uh oh!

jschmied commented Oct 26, 2026

Uh oh!

ashhart commented Oct 27, 2026

Uh oh!

jschmied commented Oct 27, 2026

Uh oh!

ashhart commented Oct 27, 2026

Uh oh!

jschmied commented Oct 27, 2026

Uh oh!

jschmied commented Oct 27, 2026

Uh oh!

ashhart commented Oct 27, 2026

Uh oh!

Uh oh!

ashhart commented Oct 28, 2026

Uh oh!

ashhart commented Oct 28, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jschmied commented Oct 26, 2026 •

edited

Loading