Repository navigation
CUDA input gradients for Living Weights on Flash Next: a Gated DeltaNet layer, 4-bit products, routed experts, gated attention (shared), and the hyper-connections and pocket (family) - #895
Open
ZackSample wants to merge 49 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
…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.
…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.
…t within float64's tolerance (tf-cuda-test gdn-back)
…erts with a shared one, and softmax top-k routing train_q4.cu with cuda.train_q4: a linear's input gradient on tensor cores (split K for a head), the experts' products over MLX stacks or experts_pack.cu's packed table, each pair's SwiGLU, and the top_k softmax and sigmoid shared gate back through the router. tf_train_sort_pairs takes up to 1024 experts. tf-cuda-test q4-back and moe-back check them against float64 at 64 and 256 rows, each reference against central differences; packed experts give the stacks' bits.
…d's q then its gate, RMS-normed q and k, partial RoPE) train_attn.cu with cuda.train_attn: the prepared heads again, the output gate, causal attention to 2048 keys with train.cu's key and value part, the heads' RoPE and norms undone, and the q, k and v projections' input gradients through cuda.train_q4. tf-cuda-test attn-back checks every gradient against float64 at 64 and 256 rows (24 query heads, 2 kv heads, head dim 256), the reference against central differences.
…nts, and fn_back over the core's layers flashnext_train.cu: the write-back's and read-out's backward (mix, activations, down and inject, the streams' RMS norms), and the pocket's SwiGLU neurons beside the shared expert behind its gate: the forward in two kernels (tf_fn_pocket_in, then _out, tiled for many rows; a layer with no neuron launches nothing), its input gradient, and its output vectors' gradient. flashnext.fn_back composes them with cuda.train_q4 and cuda.train_attn on one stream. tf-cuda-test fn-hc-back and fn-moe-back check against float64 at 64 and 256 rows; fn-pocket checks the two kernels against tf_fn_pocket_rows byte for byte over 444 cases; fn-back runs every input-gradient case the trainer uses.
ZackSample
force-pushed
the
lw-c-fnback
branch
from
October 16, 2026 00:44
a5d183d to
0211f4d
Compare
This was referenced Oct 16, 2026
Open
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this changes
These are the kernels Living Weights' trainer needs to run the loss's gradient back through Qwen3.8 Flash Next on CUDA.
There is no serving change: nothing on the serve path calls them, and no shipped file outside the kernel list, the
runtime's exports and
tf-cuda-testis touched. The trainer that uses them is the next PR (lw-d-fnlearn).Following #736 and the recipe's core/family split, what any family with the same layer could use sits in the core
(
zig/kernels/cuda/*.cuwith azig/src/cuda/*.zigwrapper), and only what is Flash Next's own sits infamilies/flashnext:gdn_train.cu,cuda.gdn_train). The forward fortraining is recomputed from the layer's kept projection rows (conv, q and k L2 norms, g and beta, the delta-rule
chain with a state checkpoint every 4 rows), with the serving chain's operations in its order, so
yhas the servingbits. The backward walks each chunk back after replaying its states into an L2-sized work area. Per row, with
dsthe state's gradient carried:
ds += dy_j q;dq = sum_j dy_j S_j;dd_j = ds_j . k;dv_j = beta dd_j;dbeta = sum_j dd_j (v_j - S'_j . k);dk = sum_j (delta_j ds_j - beta dd_j S'_j);ds' = ds - beta dd k^T;dg = sum ds' . S_prev;ds = g ds'. Then the gated RMS norm, the L2 norms, the gates, the causal conv and its SiLU,and the in and out projections. Gradients into a kept prefix's state and conv history are dropped, as Nemotron's
Mamba backward drops them. Shapes are parameters (key and value heads, hidden), so a Qwen3.8-27B-shaped DeltaNet uses
the same kernels.
train_q4.cu,cuda.train_q4). A linear's input gradient on tensor cores (tf_train_gemm's tiles, split K for a 248,320-row head),the experts' products over MLX stacks or over
experts_pack.cu's packed table (the engine's own copy, so no secondcopy of the experts), each pair's SwiGLU, and the top-k softmax weights and the sigmoid shared gate back through the
router.
tf_train_sort_pairsnow takes up to 1024 experts (512 and a shared one); below 256 it does what it did.train_attn.cu,cuda.train_attn), Qwen3-Next's form: each head's q then its outputgate, RMS-normed q and k, rotate-half RoPE on the first
2 halfdims. The prepared heads again, the output gate, thecausal part to 2048 keys (with
train.cu'stf_train_attn_kv), the RoPE and norms undone, and the three projectionsthrough
cuda.train_q4.flashnext_train.cu,families/flashnext/cuda_fn_back.zig). Thewrite-back
h' = h + bf16(branch inject)and the read-out (mix, its activations, down and inject, the streams' RMSnorms) undone; and Living Weights' pocket of SwiGLU neurons beside the shared expert, behind its gate: the forward in
two kernels (
tf_fn_pocket_inthen_out, tiled forms past a size, nothing launched for a layer with no neuron inuse; a port of nemotron cuda: skip empty pocket layers and spread the pocket over many blocks #787's two-kernel pocket to this family), its input gradient and its output vectors' gradient.
flashnext.fn_backcomposes these with the core's on one stream; its MoE backward runs the core's experts, then thepocket (whose shared-gate share joins the gate's gradient), then the core's router.
Every comment is one line; the block comments the derivations used to be in are in this description instead.
Receipt
lw-c-fnbacka5d183d on 7746bb9 (Flash Next on CUDA: the qwen4_exp family on the native CUDA host, every stream through the lane rounds #869), Zig 0.17.0; NVIDIA GB10 (DGX Spark, sm_121), driver580.178.04, CUDA 13.0.88, Ubuntu 24.04.5. Synthetic tensors at Flash Next's shapes; no checkpoint.
differences of a float64 forward. A gradient's error is its worst column (the largest difference in a column over that
column's largest reference value) and its relative L2 error. The bf16 products' tolerance is 3% worst column and 1%
L2; the fp32 chains' is 1e-3.
tf-cuda-test gdn-back(16 key heads, 48 value heads, hidden 2560, 4-bit group-32 projections, 64, 256 and 37 rows,the last after a prefix with a conv history and state): forward bf16 outputs within 1 ulp of float64's rounding
(
y: 25 of 227,328 values 1 ulp off), every chain gradient within 1.3e-6 worst column,dxwithin 1.5e-2 (4.7e-5with the GEMM's operands rounded as the GPU rounds them). PASS.
tf-cuda-test q4-back(the loss through a 248,320 x 2560 4-bit head, 64 and 256 rows): loss and target probabilitywithin 7.2e-8, the bf16 gradient rows within 2.6e-3 of the largest,
dmixedworst column 6.6e-3, L2 2.6e-3;reference against central differences 8e-7. 7 checks, PASS.
tf-cuda-test moe-back(512 routed experts of width 640, the shared expert, top 10, 64 and 256 rows):dxworstcolumn 2.7e-6, each slot's dL/dw 1.7e-6, the router logits' gradients 1.9e-6; the same layer over the packed table
gives the MLX stacks'
dxbit for bit; reference against central differences 1.2e-9. 11 checks, PASS.tf-cuda-test attn-back(24 query heads, 2 kv heads, head dim 256, RoPE on 64 dims, 64 and 256 rows):dxworstcolumn 1.3e-2, L2 3.1e-3;
dpq,dpk,dvwithin 1.3e-2; the prepared heads within 2^-5 of the host's; referenceagainst central differences 7.6e-9. 11 checks, PASS.
tf-cuda-test fn-hc-back(4 streams, low rank 320, a block's read-out with inject and the final mixer's, 64 and 256rows):
dbranchanddinjectwithin 1.5e-7,dhworst column 5.8e-3, L2 1.7e-3; references against centraldifferences 6e-8. 13 checks, PASS.
tf-cuda-test fn-moe-back(moe-backwith 16 pocket neurons): every gradient as above, the pocket's forward and itsoutput vectors' gradients within tolerance, the packed table again bit for bit. 16 checks, PASS.
tf-cuda-test fn-pocket 20: the two pocket kernels againsttf_fn_pocket_rowsbyte for byte in 444 cases (1 to2,048 rows across both tiling thresholds, 1 to 128 neurons, widths 2560 and 2600, fp32 and bf16 outputs, with and
without the gate and the training row scales); controls that reverse the neurons or rotate the inputs change
12,703 and 9,633 values, so the comparison sees order; 0 neurons launches nothing and keeps every bit, -0.0
included. PASS.
fn-backruns all six gradient cases one after another.over the packed table), attention 3.3 ms, a hyper-connection read-out 1.2 ms, the head 16.7 ms.
zig buildwith no-Dnvccor-Dfatbins) each new command printsSKIPand passes.zig build test-cuda-hostfrom a fresh cache: 29/29 steps, 506/506 tests pass (Flash Next on CUDA: the qwen4_exp family on the native CUDA host, every stream through the lane rounds #869's head: 503); everycommit builds with
-Dnvccand passeszig fmt --check build.zig zig;python3 tools/zig/lean_check.pygives 527problems against 527 on 7746bb9 (0 new); all 46 fatbins build for sm_86, sm_120 and sm_121 and
check_fatbinspasses.
Checklist
change.)
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (Nottouched.)
bits change. The backward runs in fp32 from bf16 operands and treats the forward's roundings as identities, as
Nemotron's trainer does.)
Not in this PR, said plainly
--slide: they are the next PR, stacked on this one and on Living Weights as a core module: the learner, pocket and fits in zig/src/living, a family described as data (Nemotron only) #878.using
cuda.train_q4's 4-bit product there instead is a speed-up I have not made.train_q4is group 32 only (Flash Next's); Nemotron's trainer keeps its own group-64 kernels.indexer's sparse attention has no backward.
Stack
flashnext.fn_backlives in that family, and the packed-expert layout it readsis the one Flash Next on CUDA: the qwen4_exp family on the native CUDA host, every stream through the lane rounds #869's engine packs.
--slidefor Flash Next, branchlw-d-fnlearn) stacks on this and on Living Weights as a core module: the learner, pocket and fits in zig/src/living, a family described as data (Nemotron only) #878.Questions
the core because Qwen3.8-27B (or any Qwen3-Next-shaped model) has the same layers; the hyper-connections and the
pocket are Flash Next's own.
form is what the float64 test feeds directly; the packed form is what the trainer uses.