Repository navigation
Conversation
… sums, bit for bit A ModelOpt NVFP4 checkpoint's routed experts ran the decode kernel on prompt chunks: plans of 16 pairs an item, so a 1,536-row chunk read each expert's weights two to three times. prompt_tiles gets a FACT flag: each 16-input block is one mma of the raw codes from zero, then fma(p, its e4m3 scale, acc), blocks in ascending K, then the global scale and the epilogue, the decode form's arithmetic in its order. Two entries, tf_nvfp4_fact_swiglu and tf_nvfp4_fact_down_f32, on 64-pair plans (Nvfp4Experts.runFact); the ModelOpt expert path takes them from 64 rows a pass (TF_FLASHNEXT_FACT_ROWS, 0 for never). Without the flag the template is unchanged: the 55 existing nvfp4_experts functions' SASS equals main's on sm_80/86/89/120/121. tf-cuda-test nvfp4-fact: bytes equal to the decode form at 1 to 300 rows, FP64-bounded, the scaled-code prompt tile as a red control. Kernel contract: both entries replace the decode instances (gate 3 bit-equal), gate 4 on integer sums, the scaled-code tile as a control red at gate 3, gate 5 records from nvcc 13.0.88.
…ch.zig back to 600 lines The lean gate's new findings on this branch: multi-line comments folded to one line within 120 columns, experts.zig's runFact moved unchanged to experts_fact.zig (re-exported as Nvfp4Experts.runFact), and factRows' cache as a file-level optional. nvfp4_experts.cu changes in comments only (sm_121a SASS identical); lean_new.py: 0 new.
|
Thanks @jschmied for the NVFP4 expert prompt tiles that keep the decode form's sums, and for registering both entries with a scaled-code control that has to fail at gate 3. We are carrying both commits onto current main ourselves, kept as yours, and a separate commit of ours will move the new dispatch policy ( - Synthema, Ash's AI Assistant. |
|
A heads-up, since you asked to be told so we don't collide: we are now working on a follow-up on top of these prompt tiles, not on this PR itself. Please land #1143 as you planned; we won't push to its branch. The follow-up is the down projection at large chunks (at 4,096 rows it is the slow half, with K = 640 and two stages), kept bit-identical to the decode form, and then a separately measured opt-in FP4 x FP4 path for |
|
Data for the chunk-size question (question 1 in the description), measured on one GB10 with Bits do not depend on the chunk. All 272 token-hash comparisons are identical across 1024, 1536, 4096 and 8192. That covers greedy and seeded, c=1 and c=4, and cold, warm and resumed requests at about 10k and 38k tokens. The cuts sit on the 128-row grid and the kernels are row-invariant. So the size is a free knob as far as exactness goes.
A suggestion, if you want one: choose the size from state rather than a constant tuned on one machine. Use the largest size that fits a share of free memory at load (4096 here) when no stream is decoding, and 1024-1536 while streams decode. Since bits don't depend on the chunk, that choice can't break concurrent == solo or resumed == fresh. |
|
An addendum to the chunk data above: those numbers were measured on
|
… contract names that do not collide with main's - factRows() and the prompt-tile choice move from cuda_batch.zig to cuda_nvfp4_tile.zig (size() and run()), unchanged in what they pick, since main's cuda_batch.zig has grown and the merge reached 611 lines; it is 600 now. - contract_tests.zig: the new fact_k and facts become nvfp4_fact_k and nvfp4_facts, since main's folded-affine contract already declares both and the runner did not compile.
|
Landed on main in 93a215b..35e46c0, with your records under your authorship plus the maintainer fix-ups. Thank you. Measured with the whole group (#1139, #1140, #1143, #1147, #1164, #1169, #1171 and #1170's Flash Next slice) against main on GB10s:
- Synthema, Ash's AI Assistant. |
What this changes
On a ModelOpt NVFP4 checkpoint (
nvidia/Qwen3.8-Flash-Next-FP4-Local, local-inference-lab's and Mia-AiLab's exports)the routed experts of a prompt chunk ran the decode kernel:
nvfp4_expert_kernelon plans of 16 pairs an item, so a1,536-row chunk (about 30 pairs an expert) read each expert's weights two to three times, through a kernel without
shared-memory staging. The prompt tiles
nvfp4_experts.cualready has (tf_nvfp4_prompt_*, Nemotron's) sumdifferently (codes times scales in bf16, one fp32 sum over K) and have no fp32 down, so they could not stand in.
prompt_tilesgets aFACTflag. With it each 16-input block is one mma of the raw codes from zero and thenfma(p, its e4m3 scale, acc), blocks in ascending K, then the global scale and the epilogue: the decode form'sarithmetic in its order, so every row's bits are the decode form's at any row count. Without it the code is as
before; the 55 existing functions' SASS is identical to
main's on sm_80, 86, 89, 120 and 121.tf_nvfp4_fact_swiglu(SwiGLU gate-up to bf16, 4 warps) andtf_nvfp4_fact_down_f32(fp32 slot rows,8 warps), on plans of
promptTile(.nvfp4)= 64 pairs an item;Nvfp4Experts.runFactlaunches them.cuda_batch.zig's ModelOpt expert path takes them from 64 rows a pass (TF_FLASHNEXT_FACT_ROWS, 0 for never);below it, decode rounds and verify windows, nothing changes. Since the bits are the same either way, the threshold
only sets speed.
Overlap: Mirrdhyn's #1046 does this for the MLX affine experts (
flashnext_fexp.cu, factored g32 tiles) and theformat-independent prompt kernels; it does not touch the ModelOpt NVFP4 experts. This branch is off
mainand stackson nothing of #1046; the two touch different expert formats, and
cuda_batch.zig'smoeBlockKeepin differentbranches of its
if (lw.native_experts).Receipt
Environment: TensorFold
maind1018b31plus this branch; one DGX Spark (GB10, sm_121, CUDA 13.0, nvcc 13.0.88);checkpoint
nvidia/Qwen3.8-Flash-Next-FP4-Local@b9ce18a6(33/33 files against the publisher's hashes).Exactness:
tf-cuda-test nvfp4-fact: random NVFP4 experts (dims 256, width 160, 6 experts), gate-up and fp32 down from thefact tiles on 64-pair plans byte-equal to the decode form on 16-pair plans at 1, 7, 17, 22, 64, 65, 130 and 300
rows; each output within its float64 bound (worst 0.47 of it for gate-up, 0.01 for down). Control: the existing
scaled-code prompt tile on the same plan parts from the decode form's bytes (red).
mainon the same checkpoint and flags: every reply byte-identical (three chats, three~8k agent contexts, one 30,922-token context; 64-96 tokens each).
Prompt speed (cold,
--no-drafts --parallel 4 --context 40960,--prompt-cache-gib 0, time to first token):Per-op GPU time of a 1,536-row chunk (cudaEvents around each op class, a local measurement patch): routed experts
480 ms on
main, 330 ms here (every expert's 2.76 MB of weights read once an item; 68 GB a chunk, a 250 ms floor atthe GB10's bandwidth). With
TF_FLASHNEXT_CHUNK=4096the weights are read once per 4,096 rows: 8k prompts 4.73-5.39 s,30,922 tokens 19.93 s, replies again identical (the chunk size is not changed here).
Decode speed: unchanged (23.6 / 23.6 tok/s code / chat at one stream, 65.6 / 64.9 aggregate at four,
main23.6 /23.7 and 65.6 / 65.1); decode rows never reach 64.
Kernel contract,
tf-cuda-test kernel-contract(built-Dnvcc, ReleaseFast) on the GB10:SUMMARY 87 kernels: 87 pass gates 1-4, 0 fail, 0 known findings; 8 controls: 8 red as planted, 0 missed; 1533 limit problems(this run less the 2 entries and the control is 85 kernels and 7 controls; the 1,533 limit problems also appear on a main-based run on this box (another branch, same count), none in the two new launches, which pass live on sm_121)zig/tests/cuda/contract_nvfp4_fact.zig), each a replacement of the decode instance it stands for(
nvfp4_expert_kernel<2, 2, 4>,<1, 0, 4>): gate 3 bit-equal at Flash Next's shapes (2,560 dims, 640 wide)and the tile edges; gate 1 by the stated rule (the decode form's bytes with NaN and infinite global scales, a NaN
block scale, NaN and infinite inputs, some output non-finite); gate 4 on integer inputs whose partial sums reach
about 2^20; SwiGLU states
sums.bf16 = 30: a bf16 sum that collapses moves the output by about itself, and the epilogue's three bf16 roundings already allow 3u of it, so at most 85x is provable (measured: fp16 misses 1,192,211x, bf16 36x; the fp32 down: 272,629,752x and 1,962,934,214x).The run's 1,533 limit problems are
main's (the offline audit of every recorded launch on this nvcc).tf_nvfp4_fact_swiglutakes 255 registers and spills 116 bytes on sm_121 (the existing SwiGLU prompt tile none).Tests run:
zig build test --summary all(built-Dnvcc): 211/211 steps succeeded (the summary counts 1,462 of 1,465 tests passed, 3 skipped);tf-cuda-test nvfp4-factandkernel-contract;zig fmt --checkclean.Not run, and why: sm_80/86/89/120 only through their compiles (gate 5); Metal (no Metal change).
Gotchas
cuda_batch.zig'smoeBlockKeep, in different branches of the sameif. Flash Next CUDA: prompt processing 1.4-1.8x faster on GB10, bit for bit with 1.1.0 #1046 itself conflicts withmainright now, so I measured it on its own 1.1.0 base and didn't stack on it.runFact, and wrap thecuMemsetD8Asyncin Flash Next CUDA: no zero fills the one-rank ModelOpt pass reads back unchanged #1147'sif (native.world != 1 or native.rank != 0).tf_nvfp4_fact_swiglutakes 255 registers and spills 116 bytes on sm_121. It is still faster, and tuning comes next.Questions
These came up while measuring and are your call, not part of this change:
TF_FLASHNEXT_CHUNK=4096instead of the default 1,536 rows takes another 11-12 % off prompts on this checkpoint (numbers above), with identical replies. The cost is about 2.3 GiB more and a longer pause for other streams' decode per chunk (about 3 s instead of 1.1 s). Should the default move, or stay an option?TF_FLASHNEXT_DRAFT_HEAD=affine4is 10-13 % faster drafted at one stream on this checkpoint (code and chat), outputs equal. Should it become the default for ModelOpt checkpoints?Would an opt-inWithdrawn: I built it locally (FP4xFP4 experts under the checkpoint's input scales), and it gave no speed. At these chunk sizes the experts are bound by streaming their weights, so FP4 math saves nothing, and the extra quantize passes cost a little. It also couldn't be quality-gated, because the Flash Next backend returns no logprobs.--precision checkpointfor Flash Next be welcome?Checklist
"draft": false(8/8), concurrent equals solo (8/8), and a resumed two-turn prompt equals a fresh one (7,147 of 7,184 tokens restored), each equal tomain's output (MTP drafts on,--prompt-cache-gib 4; measured at074e0ec3, before the later lean fix, which only folds comments and moves code).SUMMARYline is in the receipt.main).nvfp4-fact's control is red on the old tile).