Flash Next on CUDA: the qwen4_exp family on the native CUDA host, every stream through the lane rounds

CHANGELOG.md

@@ -2,6 +2,15 @@
 
 Each release's page on GitHub has its notes and measurements. See [the 1.0.0 release notes](RELEASE-NOTES-1.0.0.md) for the native binary's supported models and migration details.
 
+## Unreleased
+
+- Qwen3.8 Flash Next serves on CUDA from its MLX 4-bit group-32 checkpoint with the MTP head, through the lane core: drafted windows, shared rounds of streams without drafts, kept prompt states (a drafting stream's too, with the head's cache rows and a lookahead of 1, as Nemotron's) and prompts a chunk at a time between other streams' rounds. Each prompt chunk, verify window and decode round is one batched pass whose kernels give every row the bits it gets alone, so drafted replies equal `"draft": false`, concurrent replies equal their solo runs and resumed prompts equal fresh ones; past 2,048 keys attention reads the 512 blocks the model's indexer picks, up to the 262,144-token window. Every kernel it launches is built from this repository (`flashnext_hc.cu`, `flashnext_attn.cu`, `flashnext_batch.cu`, and the core's token, RMS, routing and expert rows), so it needs no capture; `TENSORFOLD_CUDA_KERNELS` can name a captured set of the Python 0.6.1 engine's Triton kernels (scripts in `zig/tests/cuda/flashnext`), which then stands in for the glue so replies can be compared with that engine's teacher path. On a GB10 one stream decodes 84.2 tok/s drafted on code and 62.6 on chat, 35.1 without drafts, and 4 sampled streams without drafts 77.9 in all, within 2% of the captured set either way; a cold 12,001-token prompt's first token takes 10.4 s and a 45,030-token prompt 36.5 s, and another stream waits at most 1.3 s between tokens while it runs. 32 greedy tokens equal the Python engine's teacher path on four of five prompts of 19 to 12,001 tokens and part at token 15 on the 3,001-token one (the captured set gives all five); next-token perplexity over 4,096 tokens of this repository's documentation is 6.178 against the captured set's 6.168. A checkpoint of another width, shape or semantic setting (RMS epsilon, RoPE, gate and activation types, layer roles) is refused at load, by its config field. @AdrianBinDC started the port in #476, matching every layer of one decode token to host checks; see [Flash Next on CUDA](docs/recipes/flash-next-cuda.md) (#472). Thanks to @AdrianBinDC.
+- A Flash Next sequence's buffers are zeroed on the engine's compute stream. A fill with no stream returns before the memset lands and is not ordered against the non-blocking stream the passes run on, so a new sequence's first prompt chunk could read its DeltaNet state before it was zeroed; the drafted check, which zeroes the states between a prompt's serial and drafted runs, once left the serial tokens that way.
+- CUDA has a lane kernel for affine 4-bit experts over a window's or a prompt chunk's rows (`experts_rows.cu`, `cuda.experts.RowsExperts`): experts.cu's decode sums for each row, two row tiles of 16 a warp, so a warp reads an expert's weights once for 32 pairs. `tf-cuda-test experts-rows` checks it against the decode kernel's one-row launches byte for byte at 1 to 300 rows.
+- CUDA has routing rows for every family (`route_rows.cu`, `cuda.route_rows`): router logits on mma tiles, one chain over K from zero so a row's sums never depend on the row count, and softmax top-k with an optional sigmoid-gated shared expert; `tf-cuda-test route-rows` checks both against float64 (`core/route_ref.zig`). The token rows take groups of 32 and a copy count, and the RMS rows gain a form with the stored outputs' group sums (`Norms.rmsSums`); `norm-rows` covers both.
+- A Flash Next pass uploads its token ids, positions and n-gram rows from pinned staging the scratch owns until the copies land, instead of allocator memory freed when the call returned (an asynchronous copy from pageable memory is only documented as "might be synchronous"), and a failed zero fill frees the buffer it was filling.
+- `tools/zig/triton_aot_pack.py` leaves out a constexpr Triton folded from `None` (an unused pointer argument), where it stopped, and takes `--jit` once per capture, merged, so one call packs several captures; `test_aot_manifest.py` covers both.
+
 ## 1.0.11
 
 - `--cortheon`'s instruction filter keeps reading past a long word. A word longer than the filter's scan buffer ended the scan early, so an order written after one was never checked; each word now keeps at most 16 bytes, so a long word can no longer end the scan (#829). Thanks to @joeynyc.

RUNBOOK.md

@@ -123,6 +123,8 @@ The archives require glibc 2.28 or newer and an NVIDIA driver providing `libcuda
 CUDA serves Nemotron 3.5 Lightning, qualified on GB10 (`linux-aarch64`, greedy and sampled) and on Ampere cards with
 compute capability 8.6 (`linux-x86_64`, greedy: the RTX 30 series, RTX A6000, A10 and A40), and checked on an RTX PRO 6000 Blackwell (compute
 capability 12.0, `linux-x86_64`).
+Flash Next also serves on CUDA from its MLX 4-bit checkpoint, on the built-in kernels, run on a GB10: see
+[Flash Next on CUDA](docs/recipes/flash-next-cuda.md).
 
 On a GB10 arm64 host (an x86_64 host uses `linux-x86_64` in the same commands):
 

build.zig

@@ -1,4 +1,5 @@
 const std = @import("std");
+const glm_expiry_build = @import("zig/build/glm_expiry.zig");
 const builtin = @import("builtin");
 const cuda_build = @import("zig/build/cuda.zig");
 const dist_build = @import("zig/build/dist.zig");
@@ -178,6 +179,7 @@ pub fn build(b: *std.Build) void {
     const disk_run = b.addRunArtifact(disk_cpu);
     test_step.dependOn(&disk_run.step);
     b.step("test-learn-recovery", "CPU learned-state recovery and admission tests").dependOn(&disk_run.step);
+    glm_expiry_build.steps(b, test_step, lanes_cpu);
     const warning = b.addExecutable(.{ .name = "learned-warning-test", .root_module = b.createModule(.{
         .root_source_file = b.path("zig/src/learned_warning_test.zig"),
         .target = b.graph.host,

docs/recipes/README.md

@@ -8,7 +8,7 @@ the native engine, or the Python engine frozen at 0.6.6.
 | Nemotron 3.5 Lightning | [MLX](nemotron-3.5.md) |
 | Qwen3.8-27B | [MLX, quantization and CUDA](qwen3.8-27b.md) |
 | Qwen3.5-2B | [Native Metal preview, tied affine 4-bit](qwen3.5-2b-native.md) |
-| Qwen3.8 Flash Next | [MLX prefill and CUDA](qwen3.8-flash-next.md), [CUDA images](flash-next-vision.md) |
+| Qwen3.8 Flash Next | [MLX prefill and CUDA](qwen3.8-flash-next.md), [native CUDA](flash-next-cuda.md), [CUDA images](flash-next-vision.md) |
 | Ternary Bonsai 2 27B | [MLX](ternary-bonsai-2.md) |
 | GLM-5.3-Flash | [MLX on a 256 GB Mac, two-rank CUDA](glm-5.3-flash.md) |
 | Gemma 4 26B-A4B | [MLX, fused one-row decode](gemma-4.md) |

docs/recipes/adding-a-zig-family.md

@@ -38,7 +38,7 @@ Some of the host side already lives in core, but so far each piece serves one fa
 - staggered segments (`zig/src/core/segments.zig`, and `zig/src/cuda/segments.zig` on CUDA) serve Flash Next on Metal and Nemotron
   on CUDA;
 - every host returns the HTTP contract in `zig/src/core/engine_api.zig`, and `zig/src/core/lane_host.zig` serves Nemotron on Metal
-  and CUDA. Flash Next has its own host and serves one reply at a time.
+  and CUDA, and Flash Next on CUDA. Flash Next on Metal has its own host and serves one reply at a time.
 
 The GPU side is barely shared yet, and that costs time. In the GLM-5.3 port each extra row in a window cost about
 4.4 ms of GPU time, against about 2.3 ms for the row's own experts. Row kernels read every replicated weight again for
@@ -71,11 +71,11 @@ lane kernels go in, check against the engine's own plain output (step 4).
 
 | Piece | Shared now | Family copies or gaps |
 | --- | --- | --- |
-| Lane kernels | `zig/kernels/metal/ops/qmv.metal` (affine 4-, 6- and 8-bit row matmuls, MLX's arithmetic) and `zig/src/core/frags.zig` (the prompt layout on M1-M4) | Flash Next: `zig/kernels/metal/decode/fn_lane.metal` (6-bit, groups of 32, tensor ops, up to 16 rows). Kimi K3: MXFP4 experts and bf16 projections. Nemotron: generated per-shape kernels |
+| Lane kernels | `zig/kernels/metal/ops/qmv.metal` (affine 4-, 6- and 8-bit row matmuls, MLX's arithmetic) and `zig/src/core/frags.zig` (the prompt layout on M1-M4) | Flash Next: `zig/kernels/metal/decode/fn_lane.metal` (6-bit, groups of 32, tensor ops, up to 16 rows). Kimi K3: MXFP4 experts and bf16 projections. Nemotron: generated per-shape kernels. On CUDA, `zig/kernels/cuda/experts_rows.cu` runs affine 4-bit group-32 experts over a window's rows with the decode kernel's sums, and `route_rows.cu` the router logits and softmax top-k (Flash Next) |
 | GPU-side rounds | the host loop, depth rule and copy proposer in `zig/src/core/lanes`, used by Nemotron | Flash Next: its own host loop, plus `fz_accept`, the ring and event chaining (`zig/src/families/flashnext/replay.zig`). Nemotron on Metal: `zig/src/families/nemotron/gpu_round.zig` for a lone stream. Kimi K3: none |
 | Several Macs | the MCDMA fabric (`zig/src/fabric`) | Flash Next: `zig/src/families/flashnext/tp.zig` splits rows and holds the whole model on each Mac. Kimi K3: `zig/src/families/kimi_k3/parallel.zig` is a split plan with no link yet |
 | Prompt path | segments; the prompt cache, used by Flash Next, GLM, Nemotron and the Qwen3.8-27B | per-family prompt kernels and snapshots |
-| Serving | `engine_api` for every host; `lane_host` for Nemotron | Flash Next: its own host, one reply at a time |
+| Serving | `engine_api` for every host; `lane_host` for Nemotron, and for Flash Next on CUDA | Flash Next on Metal: its own host, one reply at a time |
 
 **The order.**
 1. GLM-5.3 builds the core pieces first:

docs/recipes/cuda.md

@@ -8,7 +8,7 @@ They are not the native Zig release's support matrix. The kernel paths, the prom
 (`--prefill-fp8`, `--no-prefill-fp8`, `--precision`) and the request and memory notes below are that engine's too,
 and the native binary refuses those flags.
 
-The native binary registers only Nemotron on CUDA, qualified on NVIDIA GB10 and on Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40).
+The native binary registers Nemotron on CUDA, qualified on NVIDIA GB10 and on Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40), and Flash Next from its MLX 4-bit checkpoint, run on a GB10 ([Flash Next on CUDA](flash-next-cuda.md)).
 A user's receipt for 1.0.4 on a bare-metal RTX 3090, at 225 W and 350 W, is in [#592](https://github.com/ashhart/TensorFold/issues/592).
 GLM-5.3-Flash on CUDA, EXL3 loading and DFlash2 drafting remain Python-engine features, not native upgrade paths.
 Inspect the installed binary with `tensorfold-native --version` and `tensorfold-native capabilities --json`.

docs/recipes/flash-next-cuda.md

@@ -0,0 +1,124 @@
+# Qwen3.8 Flash Next on CUDA
+
+The native CUDA build serves Flash Next (`qwen4_exp`) on one NVIDIA GPU from an MLX affine 4-bit, group-32 checkpoint
+with its MTP head, on kernels built in this repository. A set of the Python 0.6.1 engine's Triton kernels, captured
+and packed, can stand in for its glue to compare against that engine's teacher path. @AdrianBinDC started the port in
+[#476](https://github.com/ashhart/TensorFold/pull/476); it is tracked in
+[#472](https://github.com/ashhart/TensorFold/issues/472). It has been run on a GB10 (sm_121) only.
+
+## Run
+
+```bash
+tensorfold-native serve /path/to/checkpoint --temperature 0
+```
+
+- `--context N` takes up to the model's 262,144 tokens; the default is the whole window. Past 2,048 keys each attention
+  layer reads 512 blocks of 4 keys a row, chosen by the model's block indexer, as the Python engine does.
+- `--parallel N` admits N streams. Each stream's caches are its own: about 110 MiB of DeltaNet state and the n-gram
+  tail, plus 28 KiB a token of keys, values and indexer keys, and with drafts its MTP head's cache.
+- `--no-drafts` serves without the MTP head; `"draft": false` on a request does the same for that request.
+- `TF_FLASHNEXT_DRAFTS` (1 to 6, default 6) and `TF_FLASHNEXT_CONFIDENCE` (default 0.60) set the head's chain.
+- The engine uploads 73.3 GiB of the checkpoint: the n-gram tables stay mapped from the shards, and the vision tower is
+  not loaded.
+- It refuses another quantization, other shapes, and semantic settings its kernels fix (RMS epsilon, RoPE base, rotated
+  share, type and scaling, the gate and activation types, attention bias, the layer roles, the n-gram layer, kernel
+  and dilation): config.json is checked at load and the log names the field. It also refuses `--segments`, `--slide`
+  and a folder whose shared experts carry Living Weights neurons.
+- `TENSORFOLD_CUDA_KERNELS=DIR` swaps the glue for a captured set (see [A captured set](#a-captured-set)) when the set
+  holds every entry point the engine launches; a set without them (Nemotron's, say) is passed over, and the startup
+  line says which glue runs. With a set loaded, `TF_FLASHNEXT_OWNED` (`all`, or Triton names such as
+  `_chunks,_merge`) runs those entry points on our kernels and the rest on the set's, to put a difference down to one
+  kernel.
+
+## What the core runs and what the family runs
+
+The family plugs into the core's pieces and adds no scheduler:
+
+- The core's round loop and lane host drive every token. Drafted windows, shared rounds of streams without drafts, and
+  prompts a chunk at a time between other streams' rounds (`prefill_step`) all go through them.
+- The core's prompt cache keeps and resumes a stream's state at the marks it asks for. The family copies the state
+  (`cuda_snap.zig`), with the MTP head's cache rows for a drafting stream; the head reads one token past a row, so
+  states keep with a lookahead of 1, as Nemotron's do.
+- Shared row kernels: the token rows (`embed_rows.cu`, groups of 32 and a copy a hyper-connection stream), the MTP
+  head's RMS rows with their group sums (`norms.cu`), the router logits and softmax top-k (`route_rows.cu`), and the
+  routed experts over a window's or a chunk's rows on the core's plan (`cuda/grouped.zig`, `experts_rows.cu`).
+
+The family supplies its config check, weight loading and packing, the layer graph on one batched pass
+(`cuda_batch.zig`), and kernels for what is its own: the hyper-connections and the n-gram gate and conv
+(`flashnext_hc.cu`), attention with its gate and the sparse indexer's pooling, scores and block selection
+(`flashnext_attn.cu`, `flashnext_batch.cu`), the Gated DeltaNet chain, the n-gram gather and the MTP head.
+
+Every kernel keeps a row's own sums in a fixed order, so a row's bits are the same at any width: drafted windows equal
+serial decoding, rounds equal lone streams, and resumed prompts equal fresh ones. The glue kernels round to bf16 where
+the Python engine's do, with IEEE division and square root and a within-one-ulp exp; the router and the 4-bit
+hyper-connection dots run on mma tiles with K in ascending order.
+
+## Checks
+
+- `tf-cuda-test flashnext-hc`, `flashnext-attn` and `route-rows`, and `norm-rows` for the shared rows: each kernel
+  against a float64 chain that rounds where it rounds, and every row count's bits equal (tolerances in the test files).
+- `zig build flashnext-prompt -Dflashnext_model=DIR -Dflashnext_kernels=native -Dflashnext_ids=FILE
+  -Dflashnext_count=32`: 32 greedy tokens through the batched pass, against the Python engine's teacher path
+  (`zig/tests/cuda/flashnext/lr7/results.json`, `results-sparse.json`); a captured set's folder in place of `native`
+  runs its glue. `TF_FLASHNEXT_DRAFTS=6` decodes each prompt serially and drafted and compares them token for token
+  (`TF_FLASHNEXT_SEED` for keyed sampled draws).
+- `zig build flashnext-compare -Dflashnext_model=DIR -Dflashnext_kernels=SET -Dflashnext_ids=FILE`: our glue against a
+  captured set on one loaded engine: every prompt row's scores (top-1 agreement, KL, each path's perplexity on the
+  prompt's own next tokens, dense rows and rows past 2,048 apart), then each path's greedy reply and where they part.
+  With `TF_FLASHNEXT_OWNED` it compares the set against the set with those entries swapped.
+- `zig build flashnext-select -Dflashnext_fixtures=DIR`: the native block selection against Triton's lists.
+- `tf-cuda-test experts-rows`: the rows kernel against the decode kernel's one-row launches, byte for byte.
+- Served: `tools/bench_concurrent.py --alone --serial` (concurrent equal to solo, drafted equal to `"draft": false`),
+  and a prompt sent twice (the second resumed from a kept state, equal to the first).
+
+## A captured set
+
+A set is built from the Python 0.6.1 engine in a container (`IMAGE`, default tag `tensorfold-qwen38:v0.6.1`) with the
+scripts in `zig/tests/cuda/flashnext`, then packed with `tools/zig/triton_aot_pack.py`:
+
+```bash
+zig/tests/cuda/flashnext/capture.sh WORK/cap /path/to/hf/snapshot          # the model's own run, about 78 GiB of GPU
+zig/tests/cuda/flashnext/capture_rows.sh WORK/rows                          # the batched pass's forms, no checkpoint
+zig/tests/cuda/flashnext/capture_mtp.sh WORK/mtp                            # the MTP head's own kernels
+zig/tests/cuda/flashnext/capture_sparse.sh WORK/sparse                      # the sparse attention's kernels
+zig/tests/cuda/flashnext/pack.sh KERNELS WORK/cap/out/cap WORK/rows/out WORK/mtp/out WORK/sparse/out
+```
+
+`capture.sh` also writes the prompts and the Python engine's references the checks read. The set has 76 kernels;
+`capture_sparse.sh` writes the selection fixtures `flashnext-select` reads.
+
+## Measurements
+
+On a GB10 (DGX Spark, driver 580.178, CUDA 13.0), an MLX 4-bit group-32 Flash Next checkpoint with its MTP head,
+greedy unless said, `tools/bench_openai.py` (256 tokens, median of 3) and `tools/bench_concurrent.py` (256 tokens,
+aggregate tok/s), on our kernels and, for comparison, on the 76-kernel captured set:
+
+| | ours | captured |
+| --- | --- | --- |
+| One stream, drafted, code / chat | 84.2 / 62.6 | 82.5 / 62.1 |
+| One stream, `--no-drafts`, code / chat | 35.1 / 35.1 | 35.2 / 35.2 |
+| `--no-drafts`, 1 / 2 / 4 streams, sampled code | 35.0 / 53.1 / 77.9 | 35.2 / 53.0 / 77.1 |
+| drafted, 1 / 2 / 4 streams, greedy code | 83.8 / 82.7 / 82.2 | 82.3 / 81.1 / 80.7 |
+| Cold first token at 3,001 / 12,001 tokens (s) | 2.67 / 10.40 | 2.67 / 10.30 |
+| 45,030-token prompt, no cache (s) | 36.5 | 35.8 |
+
+Every concurrent reply equalled its solo run, and every drafted one its `"draft": false` run. Drafted streams verify
+their windows a pass each, so drafted throughput does not grow with streams yet. While a 12,001-token prompt runs,
+another stream's decode waits at most 1.3 s between tokens (one prompt chunk). Each stream at a 16,384-token window
+takes 0.74 GiB with drafts, 0.54 GiB without.
+
+Our kernels' sums are their own, so tokens part from the captured set's (the Python engine's teacher path) where two
+candidates are close: 32 greedy tokens equal it on four of `zig/tests/cuda/flashnext/lr7`'s five prompts and part at
+token 15 on the 3,001-token one. Over 4,096 tokens of this repository's documentation, the next-token perplexity is
+6.168 on the captured set and 6.178 on ours.
+
+## Not yet
+
+- The dense projections run a copy of the Python engine's group-32 lane matmul (`qmm.cu`); the core's `qmm_group` and
+  `lane_gemv` are built for groups of 64.
+- No CUDA graphs or GPU-side rounds; drafted streams verify their windows one pass each, and the MTP head drafts one
+  stream at a time.
+- Prompts run slower than the Python engine's: a cold 12,001-token prompt takes about 10 s against its 6.7 s.
+- Draws are made on the host from each row's scores.
+- Our glue's sums are not the captured kernels' bit for bit, so replies part from the Python engine's teacher path
+  after a while, as that engine's own served path does (10 to 36 tokens, from its routed-experts kernel's rounding).