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.
…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: |
|
Thanks for running it on a second GB10, @Mirrdhyn. Those receipts are useful, and the profile matches what we saw: prompts are where this engine trails Python, with Please take the g32 prompt path. Since ash has picked #952 as the Flash Next CUDA base, with affine single-rank qualified first, I'd aim it at #952's affine path rather than this branch. |
|
Thanks @Mirrdhyn for the second-GB10 checks and profile, and @ZackSample for coordinating the base. The g32 prompt work fits the plan on #952's affine path, with #869's fixtures and credit retained. Please keep it a focused shared-kernel increment, with prompt rows matching their single-row computation across chunk sizes, tails, resets and neighboring streams. Register the changed launches and retain independent numerical controls. The profile motivates that work but does not establish a speedup yet. Images, cache reuse, draft/plain and concurrent/solo qualification still belong to the final combined family. Posted by Ash's AI Assistant Synthena |
…d a drain that fails before it succeeds #952 already carries #869's staging fix. Its regression now also fails the first copy beside the second copy and the record, treats a failed copy as possibly queued, tells an event wait from the stream drain it requires, and adds a failed drain then a retry. The guard stays set until the drain succeeds. The production staging code is unchanged. Co-authored-by: ZackSample <134018356+ZackSample@users.noreply.github.com>
…ared Qwen-VL tower, and our follow-ups
|
Thank you, @ZackSample, for retaining staging until copies complete and the failing-copy regression. Your fix reached 1.1.0 through #952 in 53abe45. Follow-up 5a82c35, with your co-author credit, also fails the first copy, treats a failed copy as possibly queued and covers a drain that fails before its retry succeeds. The guard stays armed until completion is proven. We are closing this PR as carried through #952. Host failure/custody controls and GB10 serving passed; physical first-copy/failed-drain fault injection is not separately recorded. Posted by Ash's AI Assistant Synthena |
|
@ZackSample Thank you and @ashhart for carrying the torch forward on this! |
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.