Repository navigation
Flash Next on CUDA: the qwen4_exp family on the native CUDA host, every stream through the lane rounds - #869
ZackSample wants to merge 45 commits into
Conversation
A resident prompt state died (its slot's rows overwritten, or its pool blocks given up) left the Store serving the dead entry: every later lookup found it, its restore failed, and the pass prefilled the whole prompt from the start — the repeated full-refill cascade seen on live servers with long conversations. The backend now tells the prompt cache the moment a state goes stale (Slots.expire), so the Store drops that entry immediately and the next turn keeps a fresh state instead of paying for the same failure again.
The prefill path read a snap after pool admission could free it, and overwrite walked snaps while expiry dropped them. A resident state that dies now expires its cache entry at the host boundary; each lookup borrows its entry and snap until resumed or release ends that ownership; reclamation runs at host boundaries and skips active borrowers.
# Conflicts: # zig/src/core/prompt_cache.zig
|
Thanks, @ZackSample, for moving this onto the shared lane host and documenting the captured-kernel dependency clearly. Needs changes at e4e1926 before it can be our default Flash Next CUDA implementation.
Two source issues also need fixing. Please also retain #835's saved-prefix borrowing fix for deferred admission, and confirm the lifetime of temporary asynchronous upload sources at |
…roup-32 token, so the CUDA replay starts from the row the Python engine reads.
… so the linear-attention input is that mixed row.
…f the mixed row on one token, so the DeltaNet input is that projection.
…y match the host chain on one token, all 6144 gated values and all 2560 projected values, so the mixer branch is that row.
…, the write-back exact and its sums within 1 ulp, one normed value 1 step off, and the mixed row within 2 steps, so the expert input is that row.
…he shared expert exact, router logits within 29184 ulps, two weights 1 step off, the activations exact, and every down value the same bf16 with fp32 within 16384 ulps.
…the projection, the prepared heads, the one-key output, the gate and the output row all exact.
…one projected value 1 step off, and the prepared heads, the one-key output, the gate and the output row all exact.
… one projected value 1 step off, and the prepared heads, the one-key output, the gate and the output row all exact.
… the projection, the prepared heads, the one-key output, the gate and the output row all exact.
… the projection, the prepared heads, the one-key output, the gate and the output row all exact.
… the projection, the prepared heads, the one-key output, the gate and the output row all exact.
… two projected values 1 step off, one gate value 1 step off with group sums within 1024 ulps, and the prepared heads, the one-key output and the output row all exact.
… one gate value 1 step off with group sums within 8192 ulps, and the projection, the prepared heads, the one-key output and the output row all exact.
… one prepared key value 1 step off, and the projection, the other prepared heads, the one-key output, the gate and the output row all exact.
… one projected value 1 step off, and the prepared heads, the one-key output, the gate and the output row all exact.
… one projected value 1 step off, and the prepared heads, the one-key output, the gate and the output row all exact.
… two projected values 1 step off, and the prepared heads, the one-key output, the gate and the output row all exact.
…all 2560 values and the group sums exact.
…token, all 2560 mixed values exact, and three of 248320 scores 1 step off.
…ing every expert onto the GPU once.
…m; prompts past 512 keys; LR7 capture. The ids and position uploads were synchronous copies, which do not wait for the non-blocking compute stream, so a later row overwrote a queued row's token id or position: wrong tokens on most prompts. Both now go on the stream. attnStep launches one 512-key chunk per started chunk of the row's keys (as attention.attention does) instead of always one, so prompts run to the captured NCH * 512 keys; the lanes engine defaults to that window and refuses a larger one. The prompt check prints each run's top five scores and its timing. zig/tests/cuda/flashnext: capture.py and capture.sh (Python engine 0.6.1, drafts off, greedy, 1536 cache slots) and the Unbound-LR7 prompts with their references; flashnext-prompt matches the teacher paths on all three.
…ules and device weight bytes from what it declares; Flash Next registers The registry gains flashnext.native. What was Nemotron's by name is now read from the family's declarations: its prompt-cache copies (snap_bytes, snap_save, snap_restore, snap_drop, and snap_ready's spare buffers), the store's rules (cache_rules; without them a state keeps at any row with the MTP head's one token on, ashhart#670), the quiet wait before a warm pass and the lone driver's hand-over budget. Nemotron declares the copies, the wait and the budget, so its paths are unchanged; a family without copies gets no prompt cache, and the startup line says so. A family that keeps part of its checkpoint off the device says what it uploads (deviceWeightBytes) and admission counts that, not every shard: Flash Next leaves its n-gram tables mapped and loads no vision tower, so the MLX 4-bit checkpoint uploads 73.3 of its 105.4 GiB, where counting every shard made the default reserve on a GB10 refuse it by 0.1 GiB. --slide is refused before anything loads for a family that declares no learner. The Flash Next family names what the host reads: one segment, its prompt rows, each stream's bytes, its refusals in words, and the native host's lone-driver signature.
…the decode kernel's sums: experts_rows.cu experts_rows.cu is experts.cu's decode expert kernel with two 16-row tiles a warp, so a warp reads each group of an expert's weights once for up to 32 pair rows. Every row keeps the decode kernel's arithmetic: the same group sums, the same two m16n8k16 steps from zero a group, the same fma epilogue order and the same SwiGLU, built with experts.cu's flags. A row's bits therefore do not depend on how many rows share its window or chunk, which is what the lane core's exactness contract needs of a lane kernel. It is instantiated at groups of 32 (gate and up stacked with SwiGLU, and the down projection to fp32), the MLX 4-bit group-32 experts Flash Next reads. cuda.experts.RowsExperts resolves the two instances and launches them on the shared plan (cuda.grouped.Router at its 32-pair tile). `tf-cuda-test experts-rows` packs random group-32 experts with experts_pack.cu, routes 1, 2, 5, 16, 33, 100 and 300 rows (three routed experts and a shared one a row) and checks every pair row of the rows kernel against experts.cu's one-row launches of the same row, byte for byte, for the SwiGLU rows and the fp32 down rows.
…and --jit is repeatable
A Triton kernel called with None for a pointer it does not read gets that argument folded into the kernel as a
constexpr None; Flash Next's _pool kernel has two (ROPE and DELTA). The packer stopped on them ("constexpr None has
no Zig form"); it now leaves them out of the variant's constexprs, since there is nothing for a launch to match.
--jit is repeatable like --manifest, and the captures' jit.json files are merged, so one call packs several captures
whose kernels come from different JIT functions. A single --jit, or none, works as before.
test_aot_manifest.py's fake kernel now has a None constexpr (the packer before this change stops on it) and packs once
more with a second capture's jit.json beside the first, which must give the same set.
…the sparse indexer and the 262,144-token window cuda_batch.zig runs a chunk of rows through each layer at once, on kernels that keep a row's bits whatever the row count: the group-32 lane tile (qmm.cu), the core's grouped expert plan (cuda.grouped.Router) and rows kernel (cuda.experts.RowsExperts, experts_rows.cu), the DeltaNet chain over the chunk (flashnext_batch.cu: fn_gdn_front, fn_gdn_rows without block barriers, fn_gdn_norm, each value by the one-row chain's expression), attention over each row's causal prefix, and the captured Triton forms at more than one row (capture_rows.sh). A chunk therefore leaves the logits, KV cache, DeltaNet conv windows and states and the n-gram tail exactly as feeding its rows one at a time. The n-gram table stays mapped for the engine's life and a chunk's rows are gathered on host threads while the GPU runs layer 0. Engine.round runs several sequences' next tokens as the rows of one pass: the stateless kernels over every row, each row's n-gram conv, DeltaNet step, rope, cache writes and attention on its own sequence at its own position. Past 2,048 keys each attention layer reads 512 blocks of 4 keys a row, as the model does: the pooled indexer keys (_pool), the block scores (_scores), a native selection (fn_qsa_select: the lists Triton's _select and _select_tiles make, checked by flashnext-select on capture_sparse.sh's fixtures with many ties) and the sparse _chunks/_merge (QSA 1, NCH 5). The family serves the model's 262,144-token window. The family's own pieces only: config.json's quantization and shapes are checked at load against the ones the kernels and buffers are built for (cuda_config.zig: a 6-bit checkpoint or another width is refused by name); the shards are read with core.safetensors.File; a folder whose shared experts carry Living Weights neurons is refused in words. The row-by-row path Adrian's prompt check was built on moves out of the engine into cuda_rows.zig, where flashnext-prompt (TF_FLASHNEXT_PROMPT=rows) still uses it as the batched pass's reference. Captures: capture_rows.sh and capture_sparse.sh record the extra Triton forms, and pack.sh packs any set of captures with tools/zig/triton_aot_pack.py. flashnext-prompt (32 greedy tokens) matches the Python engine's teacher path on short, mid and long, and on long3k (3,001 tokens) and long12k (12,001 tokens) past the indexer's budget (results-sparse.json).
…he DeltaNet rollback The head (cuda_mtp.zig, as the Python engine's mtp.py and decode.draft) reads the next token's embedding and the main model's four streams after the last MoE write-back, runs its full-attention layer with its own hyper-connections, experts and KV cache, then its mixer and the lm_head rows of the 79,591-id draft vocabulary. A chain keeps its first draft and stops before a later one whose temperature-1 probability falls under the confidence (6 drafts and 0.60 by default; TF_FLASHNEXT_DRAFTS and TF_FLASHNEXT_CONFIDENCE). The head runs over a prompt a chunk at a time beside the main pass (Ops.prefillChunk), so it keeps only a chunk's streams, not the whole prompt's. A verify window (the pending token and the drafts) is one batched pass whose rows keep the bits they have alone, each row drawn at its own keyed position with the draw serial decoding makes. The window runs the DeltaNet chain out of place: kept whole, its end states swap in; cut, the kept rows replay from the window's start states and the projection rows the pass kept, the conv windows and the n-gram tail shift by the kept rows alone. capture_mtp.sh records the head's own Triton forms. flashnext-prompt with TF_FLASHNEXT_DRAFTS decodes each prompt serially and drafted and compares them token for token (TF_FLASHNEXT_SEED for keyed sampled draws). The server does not use the head yet.
… windows, shared rounds, kept prompt states and prompt steps The lane backend (cuda_lanes.zig) plugs into the core's round loop, prompt cache and lane host; the family adds no scheduler. Each stream owns its caches (the 36 DeltaNet conv windows and states, the 12 attention layers' KV, indexer keys and pooled blocks, the n-gram tail and history), keyed by the core's stream. - With the MTP head, a drafting stream's verify window is one batched pass; several streams' windows come in one verify call, and each window keeps its own n-gram rows so another stream's pass can reuse the scratch before it is kept. - Streams without drafts (--no-drafts, or "draft": false) take their next token together in one decode round (Engine.round) of up to 32 rows: each stream's reply is the one it gets alone. - prefill_step runs a prompt a chunk at a time (the engine's 1,536 rows, cut at the cache's marks), so the lane host runs the other streams' rounds between a long prompt's chunks instead of after the whole prompt. - The prompt cache keeps a stream's state at the marks the store asks for (cuda_snap.zig: the DeltaNet states, the attention rows and complete pooled blocks, the n-gram tail, and for a drafting stream the MTP head's cache rows too), and a later prompt restores it and runs the rest; the batched pass gives the same bits however a prompt is cut. The head reads one token past a row, so states keep with the core's lookahead of 1, as Nemotron's do. A drafting stream resumes only from a state that holds the head's rows; one without them serves a stream without drafts. - A stream's bytes are its caches and, with the head, its head cache and a chunk's streams; the pass's scratch is reserved once at load.
…cannot read them before the fill lands A fill with no stream (DeviceBuffer.fill8(..., null), cuMemsetD8) runs on the legacy stream and, for device memory, returns before the memset is done; the engine's passes run on a non-blocking stream, which the legacy stream does not order against. A new sequence's DeltaNet conv windows and states, attention caches and n-gram tail, a drafting stream's head cache, and the pass scratch's zero row were filled that way, so the first chunk of a prompt could read a buffer before it was zeroed, for example the previous sequence's DeltaNet state. Every fill now goes on the compute stream (and the scratch's inverse frequencies upload there too), ahead of the passes that read them, as the token-id upload does. The drafted check showed it: it zeroes the DeltaNet states between a prompt's serial and drafted runs, and in one run of six prompts the prose prompt's drafted decode left the serial tokens at token 9 (113 drafts where its other runs make 111). A variant that skips the reset, so the drafted run starts from a state that is not zero, gives a mismatch of the same kind; with the fills ordered, the check's repeats are equal, under CPU load too. The same pattern was in the row path the prompt check uses, which is fixed alike.
…the recipe index, RUNBOOK and CHANGELOG docs/recipes/flash-next-cuda.md says what the CUDA engine serves and how: the checkpoint it takes and the ones it refuses, the captured kernel set and how to capture and pack one, the checks (the prompt and drafted checks, the block selection, the rows kernel, served exactness) and measurements on a GB10, with what it does not do yet. The Flash Next, CUDA and zig-family pages point to it, and CHANGELOG has the Unreleased entries.
… epsilon, RoPE, gates, activations, layer roles cuda_config.zig checked the shapes and the quantization but let other settings through while the kernels evaluate fixed constants. semantics() now refuses, by field name, a checkpoint whose text config changes what the CUDA path hard-codes: rms_norm_eps (taken as fp32 1e-6), RoPE (rope_theta 10,000,000 in rope_parameters or at the top level, partial_rotary_factor 0.25, a type other than default, a scaling factor or rope_scaling), output_gate_type (sigmoid), hidden_act (silu), attention_bias, attn_output_gate, norm_topk_prob, mamba_ssm_dtype (float32), tied embeddings, dedicated MTP embeddings, indexer_budget and indexer_kv_heads, the n-gram layer (ple_layer_ids [2], ple_embed_dim, ple_conv_kernel_size 4, ngram_size 3, which is the conv's dilation), the layer roles (three linear-attention layers, then one full) and the MTP layer's (one full-attention layer at the same RoPE base). An absent field means the Python engine's default, as its config reader takes it. The n-gram hash's own fields (heads_per_ngram, the vocabulary base and divisor, split parts, seed) are read and honoured by the n-gram reader; mrope_section moves image rows only. Tests change each setting in turn and expect its name back.
zeros() counted a buffer only after its fill, so when the fill failed the errdefer freed the buffers before it and leaked the one just allocated. The loop is now made(): a buffer is counted as soon as it is allocated, before its fill, so a failure frees every buffer made, the last one too. A test with a fake buffer fails the fill at the first, second and fourth of four and checks none is left live.
…s until the copies land Pass.stage() copied a chunk's token ids and attention positions, and chunk() the rows' n-gram embeddings and their group sums, from allocator memory freed when the function returned, with cuMemcpyHtoDAsync on the compute stream. For pageable memory CUDA documents that call only as "might be synchronous": the driver stages the bytes before it returns on the drivers we run, but nothing promises it, and freed memory may be reused before a copy that was not staged lands. The scratch now owns one pinned buffer for all four (ids, positions, n-gram rows and sums, at 16-byte-aligned offsets) and an event recorded after each pass's last upload; the next pass waits on the event before it writes the buffer (the gather thread writes the n-gram rows into it directly). The row-by-row check path's two uploads from stack values now wait for their copy before the frame goes.
…te_rows.cu) tf_route_logits gives OUT[m, e] = x[m] . w[e] in fp32 for bf16 rows and router rows on mma tiles: one chain over K in 16-wide steps from zero, the inputs staged with cp.async and read back with ldmatrix. A row's sums never depend on the row count or the tile, so the few-row form (up to 16 rows, three 256-wide stages in flight) and the prompt form (two 128-wide stages) give the same bits. tf_route_topk takes a row's top k of its first E logits (the lower id among equal ones) with weights bf16(e^(l_k - l_0) / their sum, summed in pick order), and optionally a shared expert as slot k at bf16(sigmoid(bf16(logit E))), the softmax routing Qwen-style MoE families use. core/route_ref.zig is the float64 reference, and tf-cuda-test route-rows checks both kernels against it at 1, 3, 16, 17 and 40 rows (logits within the fp32 dot's bound, picks equal to the reference rule's on the kernel's own logits, a row of ties too, weights and the shared gate within half a bf16 ulp plus 2^-18), and every row count's bits equal.
…ith their group sums embed_rows.cu's affine 4-bit rows take the scale group (32 or 64) and a copy count (EmbedRows.affine4Groups; affine4 keeps groups of 64 and one copy, so Nemotron's launches are unchanged). norms.cu gains tf_rms_rows_sums (Norms.rmsSums): bf16 rows read x_ld apart, y = bf16(x * rsqrt(mean(x^2) + eps) * w) with fp32 weights and IEEE division and square root, then each group's sum of the stored y, for a group-quantized matmul's input. tf-cuda-test norm-rows checks both against float64 at every row count (Flash Next's 2,560-wide embedding in 4 copies, its MTP head's 2,560- and 10,240-wide norms) with the bits of every row count equal.
…attn.cu, against float64 Our kernels for every entry point the batched pass took from a captured Triton set, with the Python engine's formulas and bf16 rounding points, IEEE division and square root and the core's within-one-ulp exp (glue_math.cuh), built without contraction: - flashnext_hc.cu: the hyper-connection write-back (stream sums alone, the MoE slots, K slices), the down projection with its normed rows (the tiled 4-bit dots on mma tiles, groups in order, per K slice), the slices' sum with the bf16 SiLU and inject gates, the up projection with the stream mix, the n-gram gate and dilated conv, and the MTP head's stream adds; - flashnext_attn.cu: the q, k and indexer heads' norm and RoPE with the cache writes, the output gate, the indexer's pooled blocks and block scores, and attention over 512-key chunks (q . k and p . v on mma tiles, an online softmax with fp32 sums, dense or through a sparse row's list) with the chunks merged in order. Every kernel keeps a row's own sums in a fixed order, so a row's bits are the same at any row count; attention has a form that reads the next tile ahead for a decode round's few rows, with the same sums. cuda_glue.zig launches them (Glue), with the core's token rows, RMS rows and route rows for the embedding, the MTP head's norms and the router. tf-cuda-test flashnext-hc and flashnext-attn check each kernel at the model's shapes against a float64 chain that rounds where the kernel rounds: each bf16 output within one bf16 ulp (an intermediate rounding can tip either way), fp32 sums within their sum-of-|terms| bound, attention within 2^-8 of sum p |v| / l plus half an ulp (p enters p . v as bf16), copies and the stream adds exact, and the bits of 1, 3 and 17 rows (attention: 1, 3 and 20, both forms) equal.
…y entry point is used for comparison The engine no longer needs TENSORFOLD_CUDA_KERNELS: every launch Tri made from a captured cubin now runs our kernels (cuda_glue.zig) unless a set is loaded. A set given by TENSORFOLD_CUDA_KERNELS, or found beside the binary, replaces them only when it holds all 15 entry points the pass launches (17 with the MTP head); one without them, Nemotron's say, is passed over with a log line, and the startup line says "own glue kernels" (native/cuda.zig asks the family's Loaded.own_glue; aot.Set.has names a set's functions). flashnext-prompt takes native in place of a set's folder. flashnext-compare loads the engine once and runs a prompt on our glue and on a captured set: every prompt row's scores (top-1 agreement, the largest score difference, KL, each path's perplexity on the prompt's own next tokens), then each path's greedy reply and the token where they part.
…ue beside a captured set With a captured set loaded, TF_FLASHNEXT_OWNED (all, or a list of Triton names such as _chunks,_merge) makes those entry points run our kernels and the rest the set's, so a difference can be put down to one kernel. Glue reads it once at load into a mask; without a set every entry is ours, as before. flashnext-compare takes it on its own side, so it compares the captured set against the captured set with the named entries swapped.
…for comparison The recipe, the RUNBOOK line, the CUDA and Flash Next pages and the CHANGELOG say the engine needs no capture: every kernel is built from the repository, TENSORFOLD_CUDA_KERNELS names a captured set only to compare against the Python engine's teacher path, and TF_FLASHNEXT_OWNED swaps single entry points. The recipe lists the new tf-cuda-test cases and flashnext-compare, the semantic settings the config check refuses, measurements on our kernels beside the captured set's, and where our tokens part from the teacher path. The family recipe's lane-kernel row names route_rows.cu.
e4e1926 to
7746bb9
Compare
|
Thanks, @ashhart. 7746bb9 is pushed: the branch is rebased onto 027ddea ( 1. Startup without a capture (
The captured path stays for comparison: a set given by Each new kernel has a Owned against captured, measured on the same build:
2. 3. 4. #835's saved-prefix borrowing for deferred admission. It is in the base: the branch sits on 027ddea, so the lane host's borrow 5. The temporary upload sources at Receipts on 7746bb9. On our kernels: drafted equals serial on 6 prompts greedy and 5 seeded; every We have not inspected the captured cubins either; with this change the default path no longer depends on them. |
|
Thank you, @ZackSample, for replacing the capture dependency with owned implementations for all 17 native serving entries. At 7746bb9, the epsilon, gate, theta and zero-fill cleanup findings are closed. Two paths still need attention. Please send a follow-up with exact layer-label validation and a drain before reuse or free after partial-upload failure. Posted by Ash's AI Assistant Synthena |
…rnels run The check compared each label only with linear_attention, so a full-attention layer labelled sliding_attention (or any other word) was accepted and served as full attention. Every label now has to be linear_attention or full_attention, as the layer's place in the stack says. The new test relabels full and linear layers in turn and expects layer_types back; on the old check the full layer relabelled sliding_attention came back accepted.
…he staging is reused or freed A pass queues its uploads from the scratch's pinned staging, then records an event behind them, and the next pass waits on that event before it writes the staging again. If a later upload or the record failed, an earlier upload of the same pass stayed queued while the event still marked the previous pass, so the next pass, or the scratch's free, could change or release bytes a copy had yet to read. The staging now notes an upload queued after the last record and, until a record covers it, waits by draining the stream instead of on the event. Scratch.free waits the same way before it frees the staging. The new test runs a clean pass and then one whose second upload, or whose record, fails, against a stand-in queue, and expects nothing queued at the next wait; waiting on the event alone, as before, leaves one copy queued.
…w-c-fnback (ashhart#869 and the backward kernels) Flash Next's trainer needs both open PRs: the CUDA family (ashhart#869) and the learner in zig/src/living (ashhart#878). The two touched zig/build/cuda.zig and native/cuda.zig: the build keeps both the flashnext and living modules, and the native host keeps @hasDecl(F, "slide") because Flash Next does not declare one until the next commits.
|
Thanks, @ashhart. 4b46157 is pushed with both, a commit each: 1. Exact layer labels (e7163c5, 2. A drain after a partial upload (4b46157, Receipts on 4b46157 (GB10, sm_121, driver 580.178.04, CUDA 13.0.88; Unbound-LR7 2ae3a2ff):
#895 and #896 sit on this branch, so they are rebased onto 4b46157 and pushed (lw-c-fnback 0211f4d, Heads-up: our two DGX Sparks have been running 24/7 for a week on this, so they're getting a break. Replies on these PRs may be slower for a few days. |
|
Thank you, @ZackSample, we will carry the Flash Next CUDA tests and staging fixtures with your credit into the selected Flash Next CUDA base; the landing order and remaining checks are in our reply on #952. Posted by Ash's AI Assistant Synthena |
|
Ran #869 (4b46157) on a second GB10 (sm_121, CUDA 13.2,
An nsys profile of a 12,689-token prompt shows the GPU busy 99% of the time. 35% of it goes to If nobody is on it, I'd like to take the prompt path for g32: |
5 participants
What this changes
Adds a CUDA engine for Qwen3.8 Flash Next (
qwen4_exp) totensorfold-native: an MLX affine 4-bit, group-32checkpoint with its MTP head, served through the core's lane rounds, lane host and prompt cache, with drafted windows,
shared rounds, kept prompt states and prompt chunks between other streams' rounds. Every kernel it launches is built
from this repository, so it boots with no capture.
TENSORFOLD_CUDA_KERNELScan name a captured set of the Python0.6.1 engine's Triton kernels, which then stands in for our glue so its tokens can be compared with that engine's
teacher path;
TF_FLASHNEXT_OWNEDswaps single entry points back to ours beside a set. It builds on @AdrianBinDC's#476 and answers #472. His commits are here with his authorship and messages, rebased onto
mainwith #835's cacheguard (027ddea): five needed their
zig/build/cuda.zigorzig/tests/cuda/copies.pyhunks fitted to whatmainhasnow, and his formatting commit is left out, since
mainhas the same line from his 7bbe220.Following
docs/recipes/adding-a-zig-family.mdand CONTRIBUTING.md:30-31, the family supplies its config check, weightmap, layer graph and kernels for what is its own; shared pieces are the core's:
rows for a drafting stream, lookahead 1 as Nemotron);
prefill_step;embed_rows.cu, now in groups of 32 with a copy a stream), RMS rows with their groupsums (
norms.cu, the MTP head's norms), a newroute_rows.cufor router logits and softmax top-k with a sharedgate, and the routed experts over a window's rows on the core's plan (
cuda/grouped.zig) and a new core lane kernel,experts_rows.cu(cuda.experts.RowsExperts), which gives every row the decode kernel's sums;flashnext_hc.cu), attention with itsoutput gate and the sparse indexer's pooling and scores (
flashnext_attn.cu), the indexer's block selection and theGated DeltaNet chain (
flashnext_batch.cu), and the MTP head;core.safetensors.File; a captured set is packed withtools/zig/triton_aot_pack.pyandlaunched through
cuda.aot.Set.Every kernel keeps a row's own sums in a fixed order, so a row's bits are the same at any row count: drafted windows
equal serial decoding, rounds equal lone streams, resumed prompts equal fresh ones. Our glue rounds to bf16 where the
Python engine's kernels do, with IEEE division and square root and the core's within-one-ulp exp (
glue_math.cuh);the router and the hyper-connections' 4-bit dots run on mma tiles with K in ascending order.
Commits:
host checks one by one, then a drafts-off prompt token for token.
native/cuda.zigreads a family's prompt-cache copies, cache rules and device weight bytes from what it declares(Nemotron's paths unchanged); Flash Next registers.
experts_rows.cuandtf-cuda-test experts-rows.triton_aot_pack.pyleaves out constexprs Triton folded fromNoneand takes a--jitper capture.roles, the n-gram layer), by field name.
route_rows.cu(router logits on mma tiles, softmax top-k with a shared gate) withcore/route_ref.zigandtf-cuda-test route-rows.norm-rowscovers both.flashnext_hc.cuandflashnext_attn.cu, withtf-cuda-test flashnext-hcandflashnext-attn.flashnext-compare.TF_FLASHNEXT_OWNED, our glue for the named entry points beside a captured set.docs/recipes/flash-next-cuda.md, the Flash Next, CUDA and recipe-index pages, RUNBOOK, CHANGELOG.Receipt
Environment: TensorFold
flashnext-cuda7746bb9 on 027ddea (mainc7137bb, v1.0.11 + 19, with glm: a kept state that goes stale drops its cache entry at once #835's cache guard),Zig 0.17.0; NVIDIA GB10 (DGX Spark, sm_121), driver 580.178.04, CUDA 13.0.88, Ubuntu 24.04.5; checkpoint
leimroth-lab/Unbound-LR7revision 2ae3a2ff (Qwen3.8 Flash Next, MLX affine 4-bit group 32, MTP head). "Captured"is the same build with
TENSORFOLD_CUDA_KERNELSnaming the 76-kernel set packed from 4 captures of the Python 0.6.1engine;
mainhas no Flash Next on CUDA to compare against.Our kernels, each against float64 at the model's shapes (
tf-cuda-test):route-rows(logits within the fp32dot's bound, picks equal to the reference rule's, weights and the shared gate within half a bf16 ulp plus 2^-18),
norm-rows(the group-32 token rows with copies, RMS rows with group sums),flashnext-hcandflashnext-attn(each bf16 output within one bf16 ulp of a float64 chain that rounds where the kernel rounds, fp32 sums within their
sum-of-|terms| bound, attention within 2^-8 of sum p|v|/l plus half an ulp, copies and stream adds exact). Every case
also checks that 1, 3 and 17 rows (attention: 1, 3 and 20, both of its forms) give each row the same bits.
Exactness on our kernels (the default):
TF_FLASHNEXT_DRAFTS=6: drafted equals serial token for token on 6 prompts greedy and 5 seeded.tools/bench_concurrent.py --alone --serial, 1/2/4 streams, code and chat, greedy and T=1, drafts on and off:every cell equal to its solo run (112 of 112) and drafted cells equal to
"draft": false(30 of 30), 0 failed.drafts and without, alone and with 4 at once; every
token_shaequal to the fresh run.--context 65536: the sametoken_shadrafted,"draft": falseand both resumed (45,025 cached).flashnext-selectequals Triton's block lists.Against the Python engine's teacher path,
flashnext-prompt(32 greedy tokens): captured 32/32 on all five prompts(19, 52, 600, 3,001 and 12,001 tokens). Ours: 32/32 on the 19-, 52-, 600- and 12,001-token prompts; on the
3,001-token prompt ours parts at token 15, where the captured path's best two scores are 0.375 apart.
Ours against the captured set, every prompt row teacher-forced on one loaded engine (
flashnext-compare):On natural documentation the two agree to 0.15% in perplexity, with a median per-token NLL difference of 1e-6 nats.
The Python-reference prompts are hypersensitive: with everything else captured, swapping in any single one of our
kernels that is not bit-identical moves the 3,001-token prompt's perplexity as far as the whole set does (from
134.3 to between 119 and 183; past 2,048 tokens, from 262 to between 209 and 654, where all of ours gives 483), and
on the 600-token prompt from 15.38 to between 12.9 and 16.2. Four of our entry points give the captured cubins'
bits exactly, so swapping them alone changes nothing (KL 0):
_embed,_router,_pooland_scores.Decode and prompt speed, the same build and session (
tools/bench_openai.py256 tokens median of 3;bench_concurrent.pyaggregate tok/s;bench_cache.py --require-cold):--no-drafts, code / chat--no-drafts1 / 2 / 4 streams, code T=1--no-drafts1 / 2 / 4 streams, code greedy"draft": false(s)flashnext-prompton the five prompts: prompt 92.5 / 150.1 / 662.6 / 2,483 / 9,574 ms ours against 90.8 / 154.5 /690.3 / 2,504 / 9,533 captured; decode 29.9 / 31.6 / 30.5 / 33.2 / 31.9 ms a token against 30.3 / 31.9 / 32.6 /
32.0 / 30.7 (the replies differ on the 3,001-token prompt). With the prompt cache on, the 45,030-token prompt comes
back resumed at 45,025 tokens in 0.11 s without drafts and 0.17 s drafted.
Memory: the model takes 75.8 GiB with the head (74.2 without); a stream at a 16,384-token window 0.74 GiB with drafts
and 0.54 without; at 65,536, 2.15 GiB.
Tests:
zig build test-cuda-hostfrom a fresh cache: 29/29 steps, 503/503 tests pass, 0 skipped, and each commit after Zig CUDA: one drafts-off Flash Next prompt matches #476's builds (zig build) and passes them;zig build nativebuilds from 8391cac on (Zig CUDA: one drafts-off Flash Next prompt matches #476's tip 0d04c43 and the commit after it stop atnative/cuda.zig'sF.max_segments, which the registry commit supplies);tf-cuda-testsmoke 7, symbols 1, experts-rows 7, norms 1, norm-rows 1, sample 1, glue 1, pocket3, quant-ref 4, modelopt 4, exlr8 7, route-rows 1, flashnext-hc 1, flashnext-attn 1: all pass, 0 fail, 0 skip;
lean_check.py0 new problems against 027ddea;zig fmt --checkclean; all 42 fatbins build for sm_86, sm_120 and sm_121. A captured set without FlashNext's entry points is passed over at load and the server answers on its own glue.
Not run: Metal (no change there; the Mac Flash Next host is untouched), sm_86 and sm_120 cards (none here), CUDA
graphs (none in this engine).
Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.the captured kernels:
_embed,_router,_pooland_scoresgive the same bits; the others change at the ulplevel, from IEEE division and square root and a within-one-ulp exp where Triton uses its approximate forms (the
norms' 1/sqrt, sigmoid, SiLU and softmax exps, the output gate), from fixed-order block and warp sums where Triton
reduces in its own tree (the write-back's and n-gram gate's squared sums, the group sums the 4-bit dots take, the
attention rows' p sums), from attention tiles of 32 keys where Triton's are 64 (the online softmax rescales at other
points), from explicit fused multiply-adds in the dots' epilogues and the merge, and from sincosf with full range
reduction for RoPE. Which tokens that moves is in the receipt.
cold 12,001-token prompt is 1% slower, 10.40 s against 10.30 s.)
Not in this PR, said plainly
token 15. The captured set, given by
TENSORFOLD_CUDA_KERNELS, still gives that path's tokens (32/32 on all five).qmm.cu); the core'sqmm_groupandlane_gemvare built for groups of 64. I have not checked whether a group-32qmm_groupinstance gives the same bits.throughput does not grow with streams, and the MTP head drafts one stream at a time (Nemotron batches them, Nemotron CUDA: a shared round's MTP heads draft every stream in one batch #768).
DeviceBuffer.fill8(..., null)returns before the memset lands on device memory and is not ordered with anon-blocking stream. This PR stops Flash Next using it; a doc line on
fill8may save the next family the same hunt.