Repository navigation
Conversation
…eight-only The 27B's NVFP4 MLPs and head, per-tensor FP8 mixers and bf16 a/b projections load as qlinear weights (the lane matmul for decode rows, prompt16 for prompt chunks past 16 rows), packed on the GPU by modelopt_pack. Activations stay bf16: input_scale is read but unused (--precision full). HF's RMSNorm weights get MLX's 1 + w in bf16 (bit-equal to the MLX checkpoint's 161 norms), the bf16 embedding is copied by embed_rows, and a bucket's padding rows read token 0. core/modelopt reads ModelOpt's per-tensor FP8 recipe.
qmmf keeps each row in its own accumulators at every row tile, so a window's rows need no 16-row calls: four streams of DFlash2 windows read the weights once (code 113.7 -> 171.6 tok/s on GB10, replies byte-identical to the 16-row calls).
…FP8 tile qmmf's FP4, FP8 and BF16 tiles writing bf16 (16, 32 and 64 rows, one CTA or a cluster of 8) and prompt16's FP8 tile run gates 1-4 at the 27B's reductions (5120, 6144, 17408) and rows 1 to 300; the 18 bf16-out tiles' nvcc 13.0 records join the registry. The token rows take embed_rows' registered copy kernel (one copy).
ck_dense4.cu multiplies tf_ck_quant4's rows by tf_ck_pack4's tiles of one expert on the block-scaled FP4 mma, bf16(alpha * sum) out: a warp kernel (16 rows by 32 columns) and a staged 128 x 128 prompt kernel with the same K-ordered mma chain, so a row's bits are the same at every row count. qlinear takes them for an nvfp4 weight with an input scale when a Linear carries ck4 (as FP8 weights take ck's lane8). Both pass gates 1-4 at the 27B's shapes (the prompt kernel bit-equal to the warp kernel); their nvcc 13.0 records join the registry.
The checkpoint's own math: rows quantized under each layer's static input scale, to e4m3 for the FP8 mixers (ck's lane8) and to NVFP4 for the MLPs and the head (ck_dense4), the NVFP4 weights packed for the FP4 mma. --precision full (and no flag) keeps bf16 rows against the stored weights. Refused for MLX checkpoints and before compute capability 12.
A ModelOpt config with quant_algo NVFP4 throughout (Inferact's export) quantizes every linear but the names in its ignore list, which stay bf16. A MIXED_PRECISION layer recorded as W4A16_NVFP4 (r0b0tlab's) takes bf16 rows in either precision: weight-only is its own recipe.
Their kinds come from the tensors the index names: weight_packed NVFP4 (weight_global_scale divides, so the lane takes its fp32 reciprocal), an e4m3 weight with a bf16 scale a row FP8 on the FP8G lane (the row's scale every 64 inputs), else bf16. --precision checkpoint is refused for them: their rows are quantized per token at run time, which is not built.
A row tile's units ran its column tiles in turn, so every weight tile came from DRAM once a row tile. qmm_frag::tile_of now takes about 4 MiB of NVFP4 rows a band and each band's column tiles in turn, so a weight tile is read once a band: 4096 x 17408 x 5120 135 -> 242 TFLOPS and 4096 x 5120 x 17408 128 -> 220 on GB10 (FP4 mma peak 490). Each output's mma chain is unchanged: still bit-equal to tf_ck4_dense (contract, 17 shapes).
ck's lane8 slices K by the weight's shape, so prompt rows paid the split (synthetic: sk 2 to 4 for the 27B's mixers ran 86-116 TFLOPS, its sk 1 shape 145). Ck.lane8Prompt runs the 128-row tile with K in one slice; qlinear takes it for FP8 prompt rows past 16 when a Linear sets ck_prompt_whole, which only the ModelOpt 27B does, so Nemotron's prompt bytes are unchanged. The 27B's 8k prompt in --precision checkpoint drops from 4.2-4.5 s to 3.5-3.7 s (main + ashhart#1078-ashhart#1080); drafted == plain, concurrent == solo and resumed == fresh still hold. A contract entry checks the tile at the 27B's prompt shapes against float64.
…ck_lane8t) A compressed-tensors export's FP8 math: tf_ck_quant8t takes a row's amax, its scale amax / 448 (IEEE division) and e4m3(x / scale) to nearest even, saturating, in quant8's fragment order; tf_ck8t_16/64/128 multiply those rows by e4m3 weights (the FP8 pack) with K whole and write bf16((sum * row scale) * column scale). Every output is one K-ordered mma chain, so the 64- and 128-row tiles are bit-equal to the 16-row one. Ck.quant4Global quantizes NVFP4 rows under a global scale as given (a compressed-tensors input_global_scale has no exact reciprocal in fp32: 61 of 337 of unsloth's do not round-trip). Contract: the three tiles pass gates 1-4 at the unsloth 27B's FP8 shapes; gate-5 records from fatbins for sm_80/86/89/120/121. tf-cuda-test checkpoint checks the quantizer against its definition and the tiles against each other and float64.
unsloth's export in its own math: FP8 layers (a scale a column) take rows quantized a token at a time on ck_lane8t; NVFP4 layers take NVFP4 rows under input_global_scale on ck_dense4, alpha 1 / (input_global_scale * weight_global_scale) in fp32; layers without input activations (the head) stay weight-only. On a GB10 (unsloth's export with an NVFP4 head): DFlash2 code 51.6 -> 60.7 tok/s, four streams 141.0 -> 169.6, an 8k prompt 12.3-14.3 -> 11.1-12.8 s against --precision full; drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh.
|
I pushed two commits on top of |
|
Thanks @jschmied for the ModelOpt and compressed-tensors readers, the checkpoint-math kernels and the detailed per-export numbers. We are carrying all eleven commits onto current main ourselves, kept as yours, and a separate commit of ours will give - Synthema, Ash's AI Assistant. |
|
Please keep whole-K, with the 16/17-row seam check and the float64 comparison against split-K that you describe. What I have measured: on the 27B ModelOpt export in |
|
The isolated number, as promised. Whole-K FP8 prompts on vs off, same build otherwise (banded prompt tiles on in both arms). 27B ModelOpt export (
Whole-K saves 0.31-0.37 s at 8k (about 9 %) and 1.33-1.35 s at 32k (about 8 %), against a same-arm spread of at most 0.07 s at 8k and 0.20 s at 32k. One caveat: this base is #1080's head |
|
Thanks, Jürgen. This is careful work, and the A/B on a GB10 came back clean. With the branch on main, the Qwen3.8-27B MLX 4-bit cells are token-identical to main (0 of 8 rows differ). Prefill at 8k and 32k and the 1-8 stream sweeps are all within 0.3%. We've settled a project rule this week. TensorFold is exact serial decode made faster, so we're not taking lossy activation formats, even as an opt-in. Here's how that applies to this PR:
We're preparing the split on our side and will keep your authorship on every record that stays. I'll link the landed commits here. If you'd rather split it yourself, just say so and we'll take your branch instead. - Synthema, Ash's AI Assistant. |
|
Landed on main (d1f58cf..3d219d2), as described earlier: your weight-only admission, loading, packing and What we measured on a GB10 with the landed tree:
- Synthema, Ash's AI Assistant. |
What this changes
Implements the proposal in #1144, steps 1-6, one or two commits a step (step 1 is
148ca51f, step 3b6f82573+dcb4d64f, step 694a507e3). I'll split them into separate PRs if you prefer.nvidia/Qwen3.8-27B-NVFP4and the other public NVFP4 exports of the 27B now load on the native CUDA engine. On main they stop atUnsupportedQwenQuantization.cuda_geometry.zig,cuda_modelopt.zig,core/modelopt.zig): a ModelOpt config gives each projection a kind:W4A16_NVFP4); per-tensor FP8 (core/modeloptnow reads theFP8recipe); or bf16.quantized_layers(nvidia, RadixArk, r0b0tlab), or from a single-format NVFP4 config'signorelist (Inferact).weight_packedNVFP4, FP8 with a scale a row, or bf16.cuda_weights.zig): each projection is staged as stored and packed on the GPU bymodelopt_pack(fp4Lane,fp8Lane,bf16Lane).weight_global_scaledivides, so the lane takes its fp32 reciprocal.1 + w, rounded to bf16 as MLX's conversion does. All 161 are bit-equal toTensorFold/Qwen3.8-27B-MLX-4bit's.conv1d[C, 1, 4]holds the same bytes as MLX's[C, 4, 1].embed_rows' registered copy kernel. A bucket's padding rows read token 0.--precision full(the default here): bf16 rows against the stored weights.qlinear.Weight(Proj.ql). Decode rows and verify windows take the lane matmul in one call. Rows stay in their own accumulators at every row tile, so a window needs no 16-row calls. Prompt chunks past 16 rows takeprompt16.input_scaleis read.--precision checkpoint: rows are quantized under each layer's static input scale, as the checkpoint's runtimes do:ck's lane8;tf_ck_quant4), on the newck_dense4.cu(below).W4A16_NVFP4layers stay weight-only, which is their recipe.qlinear.Lineargetsck4/ck4_rows: annvfp4weight with an input scale takes the checkpoint math when a Linear carries them, as FP8 weights already do withck.Projgains a third variant, and the MLX paths' code and launches are as on main.Receipt
d1018b31(main) plus this branch. One DGX Spark (GB10, sm_121, CUDA 13.0, nvcc 13.0.88). Drafterz-lab/Qwen3.8-27B-DFlash2. The MLX baseline isTensorFold/Qwen3.8-27B-MLX-4biton main's build.nvidia/Qwen3.8-27B-NVFP4@dbb8f445(19/19 files against HF's sha256 and git-blob hashes), plusRadixArk/Qwen3.8-27B-NVFP4(the August revision, NVFP4 head), r0b0tlab's, Inferact's and unsloth's exports from local archives, each sha256-checked against its archive manifest. RadixArk's, r0b0tlab's and unsloth's revisions were not recorded at download time, so I can't pin them.nvidia/Qwen3.8-27B-NVFP4in both precisions (the other exports below in one mode each):"draft": false: 8/8 (four prompts, greedy and sampled at seed 7).nvidia/Qwen3.8-27B-NVFP4,bench_openai.py/bench_concurrent.py(256 tokens, greedy), tok/s:fullcheckpointfulltakes 11.8-13.6 s andcheckpoint10.8-12.6 s, against MLX's 11.7-13.7 s. With qwen CUDA: TF_DRAFT_NODES above 15 serves; a stop or cancel lands between prompt chunks #1078-qwen CUDA: long-prompt prefill 5.4x faster (4096-row chunks, tensor-core attention, the delta-rule chain) #1080 merged (2bd805a4), 8k takes 6.1-6.3 s infull, 3.5-3.7 s incheckpoint(with the banded FP4 tile and FP8 prompts with K whole, below) and 6.4-6.9 s on MLX, and 32k (32,706 tokens) takes 26.8, 16.4 and 28.3 s.full11.7 / 69.9 / 192.7,checkpoint11.8 / 62.8 / 173.8 (gates incheckpoint).W4A16_NVFP4MLP and head):full11.7 / 70.1 / 188.6,checkpoint11.7 / 64.7 / 185.6 (gates incheckpoint).ignorelist):full9.6 / 42.3 / 120.1,checkpoint9.7 / 54.9 / 157.2 (gates incheckpoint).full10.9 / 51.6 / 141.0 (gates infull),checkpoint11.1 / 60.7 / 169.6 (gates incheckpoint, step 7).TensorFold/Qwen3.8-27B-MLX-4biton this branch gives main's bytes (3/3 chat replies, 3/3 8k continuations of 64 tokens).fullfullcheckpointdocs/kernel-contract.md),zig build kernel-contract -Dnvccon the GB10:SUMMARY 95 kernels: 95 pass gates 1-4, 0 fail, 0 known findings; 7 controls: 7 red as planted, 0 missed; 1533 limit problems. Main's run gives the same 1533 limit problems (records still pending on main); this branch adds none. New entries:qmmf's FP4, FP8 and BF16 tiles writing bf16, one CTA and a cluster of 8, at the 27B's reductions (5120, 6144, 17408) and rows 1-64;prompt16's FP8 tile at 17-300 rows;tf_ck4_denseandtf_ck4_dense_prompt, the prompt kernel bit-equal to the warp kernel at every shape;tf_ck8_128with K whole at the 27B's FP8 prompt shapes.tf-cuda-test kernel-contract-limitsover fatbins built for sm_80, 86, 89, 120 and 121, every record current..cufile changes (git diff main -- zig/kernelsaddsck_dense4.cuonly), so the SASS of every existing kernel is main's.zig build testpasses; the new host tests have been mutation-checked (red when broken).zig fmt --checkis clean, andlean_check.pyadds no findings.--tp 2and--visionon these checkpoints (refused, as before).check_cuda_sources.pyneeds clang++, which this box lacks.Overlap with #1078-#1080
They merge with this branch without textual conflicts. One semantic join: #1079's draft-head span split (
cuda_draft_frame.zighead) needs.ql => return full(d, rows). On the merged tree every gate above holds.Step 3:
ck_dense4A dense projection for ModelOpt NVFP4 layers under
--precision checkpoint: rows quantized bytf_ck_quant4under the layer's static input scale, times the weights packed bytf_ck_pack4as one expert, on the block-scaled FP4 mma (sm_120a / sm_121a). It is the NVFP4 experts' math (ck_experts.cu) without a routing plan. Nothing but the 27B (step 4) calls it.zig/kernels/cuda/ck_dense4.cu, a new image (arch_specific, likeck_experts). Each output isbf16(alpha * sum), wherealpha = input_scale * weight_scale_2(ck.alpha) and the sum is one mma chain over K's 64-input steps in order, from zero.tf_ck4_dense: a warp per 16 rows by 32 columns, four steps read ahead (the_fewexperts' structure).tf_ck4_dense_prompt: a block per 128 x 128 unit, four staged steps, rows in bands of about 4 MiB (qmm_frag::tile_of), so a band's weight tiles stay in L2. Same chain, so it is bit-equal totf_ck4_denseat every shape.zig/src/cuda/ck_dense.zig:Dense4.load/matmul(the staged kernel from 17 rows) andweight, the lane view of a ck-packed NVFP4 weight carrying its input scale.qlinear.Lineargetsck4andck4_rows. Annvfp4weight with an input scale takestf_ck_quant4thenDense4when a Linear carries them, as FP8 weights takeck's lane8.decodeF32/promptF32refuse ck-packed weights..cufile changes, so their SASS is main's.Ckis unchanged, and so are Nemotron's calls.Receipt
d1018b31plus this branch. One DGX Spark (GB10, sm_121, CUDA 13.0, nvcc 13.0.88).zig build kernel-contract -Dnvcc:SUMMARY 94 kernels: 94 pass gates 1-4, 0 fail, 0 known findings; 7 controls: 7 red as planted, 0 missed; 1533 limit problems(main: 85 kernels, 1533; this branch adds no limit problem). Two new entries,zig/tests/cuda/contract_ck_dense.zig:tf_ck4_denseat the 27B's NVFP4 shapes (rows 1-64; K 5120 and 17408).tf_ck4_dense_promptat 65-300 rows, bit-equal totf_ck4_denseat every served shape and edge (17 shapes).tf-cuda-test kernel-contract-limitsover fatbins for sm_80, 86, 89, 120 and 121 (116 registers for the warp kernel and 120 for the prompt kernel on sm_120/121; sm_80-89 build the trap), every record current.The bf16 m16n8k16 peak here is 123 TFLOPS, and FP4 m16n8k64's 490.
nvidia/Qwen3.8-27B-NVFP4): MLP GEMMs (tf_ck_quant4included) take 1.6-1.9 s incheckpoint. That was 2.6-2.8 s before the bands, andprompt16takes 3.6-3.8 s with bf16 rows infull. The whole 8k prompt drops from 5.1-5.3 s to 4.2-4.5 s, and 32k from 22.7 s to 18.5 s.Step 6: FP8 checkpoint-math prompts with K whole
ck's FP8 lane8 slices K by the weight's shape, so that a row's bits don't depend on the row count. At prompt sizes the split is pure cost. A synthetic bench on the 27B's mixer shapes at 4096 rows: sk 2 runs 113-117 TFLOPS, sk 4 86, and the sk 1 shape 145.Ck.lane8Prompt: the existing 128-row tile (tf_ck8_128) launched with K in one slice. Same kernel, so no new SASS.qlinear.Linear.ck_prompt_whole(default off): FP8 prompt rows past 16 takelane8Promptinstead oflane8. Only the ModelOpt 27B sets it, so Nemotron's prompt bytes are unchanged.prompt16rows already do in the 27B'sfullmode. Chunks are cut at the same places in a fresh and in a resumed prompt, so resumed == fresh holds.contract_ck_lane.zig's newwholeentry checks the tile at the 27B's prompt shapes (17-300 rows, K 5120 and 6144) against float64. It replaces no kernel.Receipt
d1018b31plus this branch. One DGX Spark (GB10, sm_121, nvcc 13.0.88).nvidia/Qwen3.8-27B-NVFP4,--precision checkpoint.SUMMARY 95 kernels: 95 pass gates 1-4, 0 fail, 0 known findings; 7 controls: 7 red as planted, 0 missed; 1533 limit problems(main: 85 kernels, 1533). The new entry is within 0.980 of the float64 bound at 12 shapes, its non-finite classes match float64's, and its sums are exact at K 6144.zig build test,zig build test-kernel-contract,zig fmt --check,lean_check.py(444 findings, as on main).Step 7:
checkpointfor compressed-tensors exports (added after the first push)Two commits on top of
94a507e3:2d4f3161and8c78b842.ck_lane8t.cu: FP8 rows quantized a token at a time (scaleamax/448, as compressed-tensors' dynamic recipe), times FP8 weights with a scale a column, in one ordered pass over K. The 16-, 64- and 128-row tiles give the same bytes.Ck.quant4Global: takes an NVFP4 input global scale as stored. compressed-tensors stores the divisor, and a double reciprocal isn't exact: 61 of unsloth's 337 global scales don't round-trip.--precision checkpointis now admitted for compressed-tensors exports. Layers stored without activation scales (unsloth's head) stay weight-only.Receipt
SUMMARY 98 kernels: 98 pass gates 1-4, 0 fail, 0 known findings; 7 controls: 7 red as planted, 0 missed; 1533 limit problems(main's 1,533). Gate-5 records for sm_80/86/89/120/121.tf-cuda-test checkpointgives bytes equal to its definition. The tiles are within the float64 bound.checkpoint: drafted == plain 8/8, concurrent == solo 8/8, resumed == fresh.fullcheckpointzig build testpasses,zig fmtis clean, andlean_checkadds no findings.Gotchas
full(exact). The Python engine defaulted tocheckpoint. See the issue.1 + wis rounded to bf16 at load. That is bit-equal to the MLX checkpoint, not to HF's fp32(1 + w)..qlcase above. The best prompt numbers need those PRs.--kv-cache q4port touches the same family files (cuda_engine.zig,cuda_layers.zig). I expect textual overlap only.Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo, in both modes.full).--precision checkpointis opt-in and changes the sums on purpose: rows are quantized under each layer's static input scale, as the checkpoint's own runtimes do.SUMMARYline is in the receipt.