Zig CUDA: Qwen 3.8 Flash Next (qwen4_exp)

This work is already underway. Please do not start a second Zig CUDA port of this family.

Flash Next is already served on Metal. This issue is the same family on the Zig CUDA engine. The Python engine stays frozen. The checkpoint model_type is qwen4_exp. The weights are Vontra/Qwen3.8-Flash-Next-MLX-4bit-MTP at revision 2b170fa6309d5d1ee380b35636075fac7945f286.

The family is not registered in the server yet. Registration waits until one prompt matches the Python engine with drafts off.

What has matched so far is one decode token against a host replay of the same row. That is not the full contract: drafted output equal to "draft": false, a resumed prompt equal to a fresh one, and concurrent equal to solo.

Matched on that token:

The embedding row.
Layer 0: both hyper-connections, linear attention, and the experts.
Layer 3 full attention, exact.
Layer 7 full attention: one projected value is 1 step off, and the prepared heads, the one-key output, the gate, and the output row are exact.

Still to do:

Full attention on the other 10 layers.
Layer 2's n-gram table.
The final mix and the word scores.
One real prompt compared with the Python engine, drafts off.

The branch is zig-flashnext-flashnext-cuda. It is not pushed, and there is no pull request yet.

I read #286 and the open pull requests. None of them is this port.

Update: still on zig-flashnext-flashnext-cuda. Not pushed.

Full attention on the same one decode token now matches through layer 39. Next is layer 43, then 47. After those: layer 2's n-gram table, the final mix and the word scores, then one real prompt against the Python engine with drafts off.

A few of those layers are 1 bf16 step off on one projected value, one prepared key, or one gate. The rest of each check is exact. This is still a host replay of one token, not the full contract.

@AdrianBinDC

Hi @AdrianBinDC, thanks for taking this on. How is it going?

We need this. We run Seesa on DGX Sparks (GB10) and want to move our Flash Next serving to the native engine. Our target is two Sparks per site, so after single-GPU registration, two ranks on GB10, --ple-on-ssd (or its native equivalent) and concurrent requests in shared rounds are what we'd be waiting for.

We'd like to help with the work, not just test it. We won't start a second port, but if any of it would be useful to hand off, we can take it on against your branch. For example: the remaining layers, layer 2's n-gram table, the two-rank path on GB10, or the SSD n-gram reads. We're already working on the Zig CUDA port of Qwen3.8-27B, so we know the CUDA family layout. Tell us which pieces you'd rather share and how you'd like to coordinate.

We can also bring compute: a pair of GB10s, plus an RTX PRO 6000 (96 GB) and an RTX 5090. Whenever you have a branch, we can run the contract checks on real hardware: drafted equal to "draft": false, resumed equal to fresh, concurrent equal to solo, and the public prompt bench before and after against Python 0.6.6 on the same checkpoint (2b170fa), with pass, fail and skip counts in the CONTRIBUTING format.

@AdrianBinDC

@AdrianBinDC we've run Flash Next on a single GB10 in vLLM since August, and the n-gram table was the hardest part there. If it helps, we can take the GB10 side of layer 2's table against your branch once it's pushed. We won't start anything separate.

What we measured in vLLM, on the checkpoint we serve there (its n-gram table is 47.7 GiB):

A pinned host copy doesn't fit. Body 69.4 GiB + pinned table 47.7 GiB + the loader's page cache went past the 121 GiB pool and hung the box.
Reading it in place from the mapped checkpoint works. That's [Qwen4Exp] Checkpoint-mapped PLE storage for unified-memory GPUs (DGX Spark) vllm-project/vllm#58439, which we have run in production since 2026-09-23. It needs no copy of the table, and swap went from about 50 GiB to about 6.
How much of it stays resident depends on the KV budget. With vLLM's default KV, the table was barely resident: about 30 major faults per decode step, and about 60 ms/step against 53.7 ms with a smaller KV cache. Readahead for a step's cold pages is [Qwen4Exp] Checkpoint-mapped PLE: read a decode step's cold pages with readahead vllm-project/vllm#58835.

We can also run your branch on our GB10 for receipts: Python vs Zig equality, memory, and faults per step.

@AdrianBinDC

@AdrianBinDC we've been running Flash Next on two DGX Sparks (TP=2, the Python engine) and built some Zig CUDA pieces for our own deployment before we saw this issue. We won't open a competing port or PR for this family. Everything is on a public branch in case any of it saves you time, and we'd be glad to rebase whatever is useful onto yours once it's pushed: https://github.com/BobClawblaw/TensorFold/tree/flashnext-zig-port

What's there, all checked against the Python engine on GB10:

The Python extension kernels built from source by the Zig build. zig/tests/cuda/flashnext/gen_copies.py copies the device code of the CUDA extensions Flash Next launches (experts, experts_prefill, gdn_prefill, qmm, qmm_prefill, qwen4_exp gdn and gdn_io) into zig/kernels/cuda/flashnext, keeping the anonymous namespace and instantiating each specialization the Python build instantiated, in that build's order. Built with the PyTorch image's nvcc, all 103 kernels' SASS equals the Python build's. Renaming the namespace or sorting the instantiations changed the SASS of two kernels; that is the one non-obvious part.
A packed Triton set for Flash Next (173 variants of 29 functions, aot.json + cubins from Nemotron's packer). _pool passes ROPE=None, which the packer can't pack; tools/flashnext_aot_pack.py packs it as an empty constant, which aot.zig matches by name.
Two ranks on CUDA: zig/src/cuda/tp_link.zig (a TCP control link: rank 0 listens and rank 1 connects; the NCCL unique id crosses it), and native --tp/--rank/--master/--master-port with a follow loop on rank 1.
A replay harness (zig/tests/cuda/fn_*.zig) that runs the Python engine's traced programs from Zig on two ranks. It reproduces a 2,695-token prefill and 48 greedy tokens exactly, with MTP drafts matching Python round for round (25 rounds, 38 drafts, 22 accepted). It's a test oracle, not an engine.

We also hit what @jschmied describes: with the 47.7 GiB n-gram table, NCCL failed with "unhandled system error" when the page cache filled the 121 GiB pool. Connecting NCCL before paging the table in, and dropping caches before a run, fixed it for us.

@AdrianBinDC

@AdrianBinDC an update on our branch, in case any of it is useful for your port. As before, we won't open a competing port or PR for this family; this is for you to take, or ignore.

Since our last comment the branch serves Flash Next natively on two DGX Sparks (TP=2, GB10) for our own deployment. It's all on https://github.com/BobClawblaw/TensorFold/tree/flashnext-zig-port, plus two open PRs on the fork: shared rounds and an int4 KV cache.

What it does, each item checked against the Python engine on the same checkpoint and GB10 pair:

Exactness. A 2,695-token prefill plus 48 tokens with MTP drafts equals the Python engine token for token. Drafted output equals undrafted. A turn resumed from a kept prompt state equals a fresh prefill. Each concurrent stream equals its solo run, including 8 prompts admitted in one pass.

Concurrency. Up to 16 streams in shared rounds, with drafts batched across streams. Aggregate tok/s at 16 streams, greedy, 256 tokens a reply:

	prose	code	structured	list
Python engine	302	509	383	545
native, per-stream launches	263	417	316	443
native, batched launches (PR Cannot get more than 8% GPU utilization on M4 with 48 gb (even using --resident-budget 30 GiB for example) #2)	296	509	383	519

At 8 streams native is at or slightly above Python. Single-stream long-context decode is about 60 tok/s against Python's 33-38 at 26k-180k tokens.

Sampling. Each rank keeps its top-64 candidates and log-sum-exp, and the draw runs on the host over both ranks' candidates. It falls back to the full row when top_p/min_p reach past 64. 20 of 20 sampled replies equal Python's.
Structured output through an xgrammar shim; masks for every draft row. I'll open a separate issue for the engine-independent part.
Prompt states kept per stream and resumed one token before the prompt's end, the same rule as the Python engine.
Images, video and HTTPS image URLs (--vision, --vision-urls).
int4 attention caches (PR issue running server #1). On this model they save memory but not time: attention reads at most 2,051 selected keys per row either way.

One lesson that will apply to your two-rank path: agree cache growth on both ranks before rank 0 sends a round's frame. Ours sent the frame first and grew the cache inside the forward. When rank 0's budget refused a growth, rank 1 had already run the forward with its collectives and rank 0 had not. The next round's collectives then paired with the wrong ones: an illegal address at int8, an indefinite hang at int4. We fixed it by reserving growth for both ranks before the frame, and ending only the refused stream.

The limits, so you know what you'd be taking:

Two ranks only.
Only our converted INT4-Mixed AutoRound checkpoint, not the MLX one this issue targets.
Kernels are replayed from the Python engine's captured Triton set, qualified on sm_121 only.
Several files are well over the 600-line limit.

So the branch is a reference and a parts bin rather than something to merge as it stands. @philip-pentatonic, since two Sparks per site with shared rounds is your target, it may be a useful stopgap to run, with those limits.

@AdrianBinDC

@AdrianBinDC a short status update on our branch, plus three lessons from this week that should apply to your port's two-rank path.

The port now runs on TensorFold v1.0.4: https://github.com/BobClawblaw/TensorFold/tree/native-v1.0.4 (git diff v1.0.4 native-v1.0.4). It's checked at int8 and int4 cache on two GB10 ranks against the Python engine:

the reply is equal token for token;
drafted equals undrafted;
concurrent equals solo;
resumed equals fresh;
the full quality suite scores the same as the Python engine.

Single-stream decode on our converted checkpoint is 77 / 126 / 99 / 126 tok/s (prose / code / structured / list), up from 67 / 106 / 85 / 107 earlier today. 16 streams run 298 / 507 / 384 / 517. As before, it's a parts bin for your port, not a competing PR.

1. A collective inside a captured graph is slower than the same collective eager. Time per partial-sum all-gather between the two ranks:

Rows	Eager	Inside one graph	Graph split at the gather
1	15.9 us	25.1 us	15.4 us
6	37.6 us	44.3 us	37.3 us
16	33.1 us	43.8 us	34.9 us

With 96 gathers per forward plus the draft steps, that adds up to about 1 ms per round, more than we expected.
NCCL protocol and channel settings didn't help. NCCL_GRAPH_MIXING_SUPPORT=0 closes most of the gap, but it isn't safe when graph and eager collectives mix.
Ending each captured graph at a gather, and running the gather between graph segments, recovered it. Payload, sum and order are unchanged, so outputs are bit-identical.

2. One captured step can serve every sequence. We had keyed graphs by sequence, so each request captured 36-41 graphs of its own.

At capture, we read each kernel node's arguments (cuFuncGetParamInfo / cuKernelGetParamInfo) and note the 8-byte words that point into the capturing sequence's memory.
Launching the plan for another sequence rewrites only those (cuGraphExecKernelNodeSetParams). We keep at most 96 plans, least recently used first out.
Requests now do 0-9 captures. This only paid off once the collectives were out of the graphs, because NCCL nodes can't be relocated.
Together with 1, single-stream decode went up 15-21%, with outputs still bit-identical.

3. Graphs cached per sequence pile up. Pooled spares and kept prompt states held 161 graphs over 9 sequences after a few minutes of mixed traffic, against 18-73 on a fresh server. Captures got slower too (3 ms to 6-7 ms each), and one stream lost 7-10% until a restart. Free a sequence's graphs when it's released, or share them as in 2.

Also, a bug pattern to check if you build engine structs with gpa.create and then assign fields one by one: field defaults never apply. One of ours was a timing flag that came up set from leftover allocator memory, which put CUDA events around every verify step.

I'll open a separate issue for 1 and 2, since they apply to every CUDA family that captures graphs.

@AdrianBinDC

@AdrianBinDC, a reference for when this port reaches sm_110 (Jetson AGX Thor), since #203 was closed in favour of this issue. We are not starting a competing port; this is only what we have on Thor, in case any of it helps:

Fork (frozen Python line, not an upstream release): https://github.com/mcclanahanaman/TensorFold/tree/thor-release (tag thor-release-2026-10-10, Apache-2.0, upstream notices kept). docs/THOR-FORK.md lists every Thor change and its kill switch.
What is in it: NVFP4 routed-expert prompt kernels (W4A16), dense/DeltaNet/sparse-attention prompt kernels, decode graphs and a device-fed MTP draft chain, n-gram tables read from NVMe through a small row cache, one KV pool with admission control (503 + Retry-After) and /v1/capacity, and vision host-memory budgeting.
How it was checked: every step against the previous build with byte-identical replies, then a gated release run (decode both orders, prefill, vision, soaks, an 8-session stress test) and a text-quality suite.
Measured on one Thor (gated, sparkDash protocol, aggregate tok/s at 1/4/8 requests): prose 62.8 / 133.3 / 196.6, code 90.3 / 191.8 / 271.5; prefill 3,113 / 3,310 / 2,755 tok/s at 8k / 32k / 128k.
Recipe and full method: https://github.com/mcclanahanaman/Qwen3.8-Flash-Next-Jetson-AGX-Thor-TensorFold

Happy to share kernel notes or traces for the native port if they help. Credit to @ashhart for TensorFold and to Mia AI Lab for the recipes and benchmark protocol this builds on.

@AdrianBinDC

@AdrianBinDC, I'd like to help with this port, not start another one.

My setup: one RTX PRO 6000 Blackwell (sm_120, 96 GB, held at 350 W) with 89 GiB of host RAM. I serve Flash Next in production on Pennyroyal/SGLang 3.0.1 from jpezzulli/OrcaRouter-Qwen3.8-Flash-Next-Uncensored-ModelOpt-NVFP4 (revision f24d2b68). It has the RadixArk layout: NVFP4 routed experts; bf16 attention, DeltaNet, shared expert and head; and an F8_E4M3 n-gram table (47.7 GiB). My workload is one long-context agent at 400K–512K tokens (YaRN ×2 over the native 262,144) plus up to 15 short subagents running at the same time.

Work that may carry over. I keep an NVFP4 KV cache for the QSA layers in Lesilva/qwen38-flashnext-nvfp4-kv. It packs E2M1 values with one E4M3 scale per 16. QSA's top-k and prefix gathers dequantize only the selected keys, so the attention kernels themselves are unchanged. Against FP8 KV it uses about 41% less KV memory and costs 1–2% of decode speed; paired NLL from 8K to 128K showed no measurable difference, and 490K three-needle retrieval passes. That matches what @BobClawblaw saw with int4 here: attention reads at most about 2K selected keys per row, so a smaller cache is a capacity win, not a speed win. The same repo also has a split-K fix for the indexer's prefill scoring. Past about 200K tokens, its logits budget left about 8 CTAs per launch on 188 SMs. Splitting the key tiles made one launch 8.5× faster at 480K and cut time to first token for a 480K prompt from 74.8 s to 61.6 s, with bit-identical logits. If the Zig engine goes past 262,144, it will probably hit the same wall.

What I could take on, against your branch once it's pushed:

Receipts on sm_120. Python 0.6.6 against Zig on one PRO 6000, in the CONTRIBUTING format: drafted equal to "draft": false, resumed equal to fresh, concurrent equal to solo, plus decode and prompt speed. The card serves production, so these would run in scheduled windows rather than on demand.
Loading ModelOpt NVFP4 Flash Next checkpoints on the Zig CUDA engine, after the MLX 4-bit path is registered. I'd reuse Nemotron's modelopt and nvfp4 pieces where they fit, and qualify against both a public export (RadixArk/Qwen3.8-Flash-Next-NVFP4 at 7b719225) and the checkpoint above.
YaRN rope scaling past 262,144, and a quantized attention cache, as separate PRs, if you think either belongs in the family. The Metal runner fixes attention capacity at 262,144, so I'd like to know whether a longer window is in scope before I start.

Two questions:

Which of these would you rather keep, and which would you be glad to hand off?
Is a shared KV pool on the roadmap for the CUDA families: one long stream plus many short ones, without each stream growing its own contiguous cache? If Zig CUDA ports: shared layers below lanes.Backend, and one place to coordinate them #548 is the place for that, I'll take it there.

@AdrianBinDC

@AdrianBinDC, I built on your #476 branch instead of starting another port, and wanted to share where it got to in case any of it is useful. Your commits are kept with your authorship and everything is on top of them, rebased onto main (now 1.0.9): https://github.com/ZackSample/TensorFold/tree/flashnext-cuda-living

Two fixes that may matter for #476 directly

Token ids and positions were uploaded with a blocking copy that wasn't ordered with the work stream, so a step could read the previous step's ids. That gave wrong tokens intermittently. Fix: ZackSample@e9c79de1
Prompts over 512 tokens: attention now runs over every prompt chunk, not just the last.

What runs now on one GB10 (Unbound-LR7, MLX affine 4-bit g32)

Token for token with the Python engine's teacher path: short, mid, long, 3k and 12k-token prompts, 32/32 tokens each.
A batched multi-row prompt pass: prefill ~1,260 tok/s at 12k tokens (Python is still about 2x faster there).
MTP drafts with sampling: drafted equals "draft": false greedy and seeded. Single stream ~31 tok/s plain, ~45-48 prose and ~73-80 code drafted.
The sparse attention indexer, so the full 262k window works (checked on a 45k-token request), plus an optional int8 KV cache.
Concurrent streams through the lane core with prompt reuse: concurrent equals solo, resumed equals fresh; 31 / 51 / 73 tok/s at 1 / 2 / 4 streams.
The family reports the memory it uploads (73.3 GiB) rather than the checkpoint's file size, so the server's load check no longer refuses it.

Known gaps

Replies match the teacher path but drift from Python's served path after 10-36 tokens, from the routed-experts kernel's rounding form (inherited).
It still replays Triton kernels captured from the Python engine, as Zig CUDA: one drafts-off Flash Next prompt matches #476 does.

Living Weights on Flash Next. The same branch has the pocket learner running on Flash Next. It has backward kernels for DeltaNet, the hyper-connections, the MoE and sparse attention, and the learner made family-neutral. It learns facts about the user and keeps them across restarts. While testing I also found a forgetting case in the learner's quiet-refit that I think Nemotron shares: a later lesson can silence an earlier fact's neurons while the refit's own check still passes. That one I'll write up separately with a repro.

Happy for you to take whatever helps, or I can split it into small PRs stacked on #476 in whatever order suits you, or keep it to myself until #476 lands. Your call.

@ZackSample

Hello @ZackSample and @AdrianBinDC,

Thank you for the Flash Next CUDA port on Adrian's #476 branch, and for keeping his authorship.

A family follows docs/recipes/adding-a-zig-family.md. Each step is built once, in the core, and the family plugs in. The family supplies its config, weight map and layer graph, and kernels only for what is its own (adding-a-zig-family.md:66). Shared lane verification, caches and per-stream state stay in the core. Model operators and weights stay in the family (CONTRIBUTING.md:30). The family uses that scheduler (CONTRIBUTING.md:31).

A pull request against main, with the receipts you listed, is welcome. Timing stays open.

Thank you both.

@AdrianBinDC

Thank you, @AdrianBinDC, and everyone who carried the Flash Next CUDA work forward, preserving the original authorship and tests.

1.1.0 ships the native CUDA family through #952, with #869's tests and credit retained. Prompt reuse, drafting and concurrent requests use the shared lane core. This closes the CUDA family port. Living Weights and longer-window or reduced-precision cache proposals remain separate work. Please reopen if the released family still fails your pinned checkpoint or request.

NVIDIA GB10 affine-4 MTP passed 17/17, concurrent equality and chat 12/12; image warm-turn/JPEG gaps and the NVFP4 pair remain limited.

Posted by Ash's AI Assistant Synthena

Navigation Menu

Repository navigation

Zig CUDA: Qwen 3.8 Flash Next (qwen4_exp) #472

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Zig CUDA: Qwen 3.8 Flash Next (qwen4_exp) #472

Description

Activity

AdrianBinDC commented on Oct 7, 2026

philip-pentatonic commented on Oct 8, 2026

jschmied commented on Oct 8, 2026

BobClawblaw commented on Oct 8, 2026

BobClawblaw commented on Oct 9, 2026

BobClawblaw commented on Oct 10, 2026

mcclanahanaman commented on Oct 10, 2026

Lesilva commented on Oct 12, 2026

ZackSample commented on Oct 14, 2026

ashhart commented on Oct 14, 2026

ashhart commented on Oct 18, 2026

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions