Repository navigation
Zig CUDA: one drafts-off Flash Next prompt matches - #476
AdrianBinDC wants to merge 23 commits into
Conversation
…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.
…inear-attention files.
…, 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.
|
Need to make some changes. Not ready for PR. Will work on it later today. |
|
Thank you, @AdrianBinDC. Eight greedy token ids matching the Python engine on a GB10, with the forward split into files under 600 lines, is real progress for Flash Next on CUDA. Since it's still a draft and you said you wanted to change things, here is where it stands for landing. The family registers with drafts off, and every model on the native engine serves through the lane rounds with drafts on (wide windows, rewind, the MTP head, prompt reuse and concurrent streams), so registration would wait for those. The group-32 instances added to experts.cu and experts_pack.cu are additive, and we'll check Nemotron's kernels stay unchanged when it lands. We'll catch it up with main ourselves when you're ready; no need to rebase. When we land it, your commits get your GitHub noreply address instead of the personal one, still under your name. |
|
@ashhart Apologies. I had to go on a last minute work trip and just got back. I'll continue working on it this weekend. |
|
Thank you, @AdrianBinDC, for the first CUDA forward for Flash Next and the prompt comparison that gave it a starting point. We are closing this branch because every engine we ship runs on the lane engine with drafts on, and this one decodes with drafts off only. Your work is not lost: #869 carries your rebased commits under your name and adds the lane host and the MTP path. Before it lands, it needs a default that boots without captured kernels, and image input in our own kernels. We would be glad to have your kernel work and exactness checks there. Posted by Ash's AI Assistant Synthena |
2 participants
Draft. I want to shake this down and clean it up before marking it ready later today (US Eastern). The branch is 85 commits behind
zig-flashnext, so catching up is part of that cleanup.Refs #472.
What this changes
Zig CUDA for Qwen 3.8 Flash Next (
qwen4_exp), drafts off. The prompt forward is split into files under about 600 lines, and all 513 experts (512 routed plus the shared expert) are packed onto the GPU once at load. The family is registered so themlx-q4g32checkpoint can open. There is no MTP head in this lane.Receipt
Environment
b8f08c4tensorfold-qwen38:v0.6.1Vontra/Qwen3.8-Flash-Next-MLX-4bit-MTP, revision2b170fa6309d5d1ee380b35636075fac7945f286Exactness
One drafts-off prompt, greedy, seed 0: "Say hello in one word." (58 tokens). The Python engine and the Zig check both returned
1596, 1144, 310, 5707, 310, 1156, 25, 328.Decode speed
Not run.
Prompt processing speed
Not run.
Tests
zig build testexited 0. It printed no pass, fail, or skip counts.Not run
tools/bench_concurrent.pytools/bench_openai.pytools/prefill_cold.pyChecklist
lanerounds"draft": false