Repository navigation
Qwen3.6-35B-A3B on CUDA: serve NVIDIA's NVFP4 checkpoint as downloaded - #1662
Conversation
|
Thank you for this. It's a careful, well-tested piece of work, and serving NVIDIA's checkpoint as downloaded is something we want. One housekeeping note: the six commits were authored with an email address that isn't your GitHub noreply one, and we only land noreply identities. Your earlier commits here use 106273+laudney@users.noreply.github.com, so we've re-authored our prepared copy of the branch with that address. The code is byte for byte the same, and it will land under your name. If you push again, please use that address so it stays that way. We're now running the regression half on our GB10: Qwen3.6-35B-A3B MLX 4-bit with MTP against main, at 1 and 4 streams and cold prompts at 2k, 8k and 32k, with replies token-equal. Your receipts cover the new NVFP4 path. Two small optional things, if you touch it again:
|
|
Thank you, and thanks for re-authoring the prepared copy. Everything I push from here on uses 106273+laudney@users.noreply.github.com, and a pre-push check on my side now refuses anything else for this fork. Both points are taken, without rewriting the branch:
The two follow-ups under #1656 are open as stacked PRs: #1674 (DFlash2 drafter, one commit of its own) and #1675 (1M-token window with YaRN, nine commits of its own). Each description says which commits to review. |
The Qwen3.6 MoE loader will pack gate and up as two parts of one NVFP4 expert table with the shared expert last, a packer path no test reached. The device test now checks it byte for byte with guards, and the decode-sums contract entries serve Qwen's 512 by 2048 shapes while the K32 tile keeps Flash Next's.
The MTP drafter of a ModelOpt checkpoint will quantize its bf16 layer on the host and needs the same staging and packing the MLX loader runs after its tensor reads. packLinear and packTable now hold that tail; linear and expertTable call them, unchanged.
nvidia/Qwen3.6-35B-A3B-NVFP4 was refused because the routed experts read only MLX tables. Admission now checks every expert projection's effective ModelOpt recipe and keeps the router bf16; the loader packs each projection of all 257 experts on the GPU into Nvfp4Experts' tables with the shared expert last; decode rows run the decode kernel and prompt tiles its decode-sum form, so a pair has the same bits in any window. Rules ported from ashhart#345 by gilby.
nvidia/Qwen3.6-35B-A3B-NVFP4 ships its MTP layer in bf16 inside its own shards, while the CUDA drafter reads only mtp-4bit.safetensors. Without that file, the loader now quantizes the mtp.* tensors at load with the shared affine-4 quantizer, as Nemotron's ModelOpt drafter does, and packs the same tables; a drafter-only exception the user chose, since every drafted token is verified.
The recipe, the README's support table and the changelog now name NVIDIA's NVFP4 checkpoint on the native CUDA engine: how to serve it with its own MTP drafter, what runs in which precision, the drafter-only exception, credit to ashhart#345, and the GB10 receipts for decode, concurrency, exactness and cold prefill.
capabilities --json listed only mlx-q4g64 for qwen3_5_moe on CUDA, though the family now reads NVIDIA's MIXED_PRECISION and single-format NVFP4 exports. The CUDA MoE family names modelopt-mixed and modelopt-nvfp4 beside it, with a test; HIP's registry keeps its own list.
The bf16 MTP loader stages each matrix in one f32 scratch buffer, sized for q with its gate and for an expert, while fc's two halves are hidden by hidden: a model whose attention width is under half its hidden size would overrun it. The buffer now holds the largest matrix the loader stages (fc's halves, q, k and v, o, an expert); for Qwen3.6 the size is unchanged, as q is still the largest.
|
Thanks @laudney. This shipped in TensorFold 1.2.0 today, with your credit in the release notes and the changelog: https://github.com/ashhart/TensorFold/releases/tag/v1.2.0 |
|
Thank you @laudney. This is in TensorFold 1.2.0, released today: https://github.com/ashhart/TensorFold/releases/tag/v1.2.0 It landed with your commits under your own name, so you're listed as a TensorFold contributor, and it's credited to you in the release notes and the changelog. |
Part of #1656 (1 of 3: NVFP4, then DFlash2, then a 1M window with YaRN).
Why
Serving the weights a lab or NVIDIA actually publishes is what makes TensorFold benchmarkable and adoptable:
nvidia/Qwen3.6-35B-A3B-NVFP4is the checkpoint SGLang, vLLM and TensorRT-LLM serve on Blackwell. Running the same bytes is the only like-for-like comparison: no conversion step, no different quantizer, nothing to explain away in a speed or quality table.UnsupportedQwenQuantization: the routed experts read MLX tables only), so a GB10 user has to fetch a different model before they can try TensorFold at all.What
The dense Qwen3.8-27B already serves ModelOpt checkpoints (
cuda_modelopt.zig) and the grouped NVFP4 expert kernel already exists; this PR wires the MoE to them:modelopt_tests.zig,contract_nvfp4_fact.zig): two-part NVFP4 expert packing (gate and up from separate tensors into one pair table) and the shared-last stride, at Qwen's shapes.cuda_weights.zig): the affine-4 packer splits its parts from its tensor reads, so the ModelOpt MTP drafter can pack quantized parts; MLX bytes unchanged.cuda_geometry.zig,cuda_modelopt.zig,cuda_modelopt_experts.zig,cuda_routed.zig,cuda_engine.zig): admission reads every expert projection's effective recipe (W4A16 NVFP4 routed and shared experts, FP8 attention and DeltaNet projections, NVFP4 head, bf16 router and gates) and refuses a checkpoint that quantizes one expert, the shared expert or the router differently, before loading. Experts pack into the grouped NVFP4 tables with the shared expert last, as in the MLX layout; prompt tiles use the kernel's decode-sum form, so an expert pair has the same bits in a prompt chunk, a verify window and a one-row step. Host tests incuda_modelopt_moe_test.zig.cuda_mtp_bf16.zig): the checkpoint keeps its MTP layer in bf16; withoutmtp-4bit.safetensorsbeside it, the drafter quantizes that layer to affine-4 at load, as Nemotron's ModelOpt drafter does. This changes only which rows a round proposes, never a reply; the startup log says so.docs/recipes/qwen3.6-moe.md, a README support row and a changelog entry.capabilities --jsonnamesmodelopt-mixedandmodelopt-nvfp4forqwen3_5_moeon CUDA besidemlx-q4g64.The rules (bf16 activations for the W4A16 layers, the shared expert last, the MTP layer drafting through the target's head) come from @gilby's Python port, #345.
Gate (one DGX Spark: GB10, sm_121, CUDA 13.0; checkpoint revision
1355db6a)tf-cuda-test modeloptpasses (the two-part repack and Qwen shapes included)."draft": falseruns (6 of 6: prose and code, greedy and two seeded temperatures); concurrent replies equal their solo runs (40 of 40 at one and four streams, greedy and sampled); a second turn that resumes from 8,192 cached tokens equals its run on a server with the prompt cache off, drafted and plain.TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP(revision f84b054) gives the same replies on this branch as onmainf350cd3 (6 of 6 token hashes).zig build test(5,604 passed, 27 skipped) andzig build test-cuda-host(89/89) pass;zig fmt --checkandgit diff --checkare clean;lean_check.pyreports 426 problems, the same asmain.Numbers
The run on record is the recipe's (
docs/recipes/qwen3.6-moe.md), the same numbers as the changelog: 256 tokens, greedy unless noted, thinking off.--no-drafts(tools/bench_openai.py)tools/bench_concurrent.py)Cold prefill (a fresh server per size,
--prompt-cache-gib 0): 0.38 s to first token at 2,048 tokens, 1.42 s at 8,192, 6.59 s at 32,768 and 15.91 s at 65,536. A re-run of the gate on this branch overmainf350cd3 passed the same exactness checks, with decode within 3% of these.Overlap
atc overlapandgit merge-treeagainst every open pull request: this PR adds no conflicts. #1335 sharescuda_native_moe.zigbut already conflicts withmainin the same files without this change.The DFlash2 drafter (#1674) and the 1M YaRN window (#1675) follow as stacked PRs under the same issue.
Follow-up commit
fix(qwen3_5): size the MTP loader's scratch for fc's halvestakes the review's first point: the scratch now holds the largest matrix the loader stages (hidden × max(hidden, 2·aw, kvw, intermediate)); the size is unchanged for Qwen3.6.