fix(qwen4_exp): support modelopt MIXED_PRECISION checkpoints - #426
Merged
Conversation
…onfig - derive expert_quant from QuantConfig.from_hf(...).scheme_for_name(...) so modelopt MIXED_PRECISION exports (nvidia/Qwen3.8-Flash-Next-NVFP4) resolve to nvfp4 through quantized_layers instead of the "fp4" substring probe - drop the attn/dense/lm_head flags: qwen4_exp dense weights are always bf16 and nothing reads them - test parse_config against the quantization_config of each released checkpoint; list the nvidia checkpoint in docs/models.md Fixes #394
This was referenced Sep 10, 2026
gberasmus87
added a commit
to gberasmus87/FreeToken
that referenced
this pull request
Sep 10, 2026
FlashML-org#426 taught parse_config to read the checkpoint's QuantConfig, so a modelopt MIXED_PRECISION build that declares its dense attention and GDN projections FP8_PB_WO now resolves to a block-FP8 scheme, and gdn.py builds the in_proj_qkvz + in_proj_ba split for it. The reader was never taught the other half. It still fuses all four in_proj parts into one buffer, so on those checkpoints it cats fp8 qkv|z with bf16 b|a and dies in torch.cat before the model sees anything, and no weight_scale_inv ever reaches the buffers Fp8BlockLinearMethod declares. The four-way fusion cannot survive block-FP8 in any case: b|a are num_v_heads rows wide, and create_weights requires every output size to be a whole number of 128-row blocks. The split is a requirement, not an optimization. So the reader now asks which fused attention groups are block-FP8 and emits in_proj_qkvz | in_proj_ba (and q|k|v -> qkv_proj) for those, fusing each group's weight_scale_inv on the same axis as its weight. Every fp8 part is a whole number of 128-row blocks (10240|6144 for qkv|z, 12288|512|512 for q|k|v), so the per-block scales concatenate exactly alongside the rows they describe. It asks through checkpoint_quant_config -- the same QuantConfig, built from the same ModelSpec name map the engine hands the model, queried with the same attribute prefixes gdn.py and attention.py use. That is what keeps the buffers this reader emits from disagreeing with the modules the model built. It matters beyond tidiness: the block-FP8 linears have no tensor-parallel variant, so a rank that downgrades has to downgrade on both sides at once. A checkpoint that stores block-FP8 dense WITHOUT declaring it per module gets no scheme, so the model builds plain bf16 linears for it. Those weights keep the existing dequantize-at-load path, and their weight_scale_inv is dropped rather than emitted into a buffer that does not exist.
gdevenyi
added a commit
to gdevenyi/FreeToken
that referenced
this pull request
Sep 10, 2026
Brings the quantization refactor (FlashML-org#418 config/scheme/method layers, FlashML-org#427 the QuantConfig reaching the weight readers) and FlashML-org#426, which takes qwen4_exp's expert quant kind from the checkpoint QuantConfig. Resolutions, all taken from the rebased PR branches so the deploy tree and the upstream PRs stay identical where they overlap (rb/tp, rb/vision, rb/fp8 -> rb/all): - qwen4_exp attention/gdn/moe/config/model, layers/linear.py, layers/embedding.py, models/nvfp4_banks.py: the rebased versions. Notably o_proj stays LinearOProj (row-parallel); main's LinearReplicated is correct only at TP=1. - NVFP4 expert TP sharding moved to the new seam: the piece stream is sliced along the intermediate axis in nvfp4_banks, and the kernel sizes its banks from MoEConfig.local_intermediate. The old _Placer/_alloc_nvfp4_host_banks are gone with the functions they lived in. - models/quant_linear.py: deleted, as upstream did; nothing imports it. - DROPPED, superseded upstream: the FlashML-org#320 cherry-pick (_dense_is_block_fp8, _block_fp8_fusions, _load_maybe_block_fp8). FlashML-org#320 is CLOSED upstream because FlashML-org#426 does the same job through QuantConfig. This checkpoint's dense side is bf16 (a 292-entry modelopt ignore list), so none of those paths were exercised here. - KEPT: the --num-tokens/--num-pages KV reserve fix (issue FlashML-org#383, still unfixed on main), FlashML-org#231's collect_decode_freq, FlashML-org#169's prefill_warmup, and every other deploy-only commit -- they merged clean. - tests/engine/test_cache_budget.py keeps both the KV-reserve test and upstream's new slot_limit test; _resolve_auto_moe_cache_size's new method arg is optional. deploy/chatdnp stays at adc32da as the rollback point. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
This was referenced Sep 10, 2026
Draft
This was referenced Sep 11, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #394
nvidia/Qwen3.8-Flash-Next-NVFP4 is a modelopt
MIXED_PRECISIONexport: the per-module algo lives inquantized_layers, so the family's"fp4" in quant_algoprobe reported the routed experts as unquantized.parse_confignow asks the checkpoint QuantConfig for the experts' scheme (scheme_for_name("model.language_model.layers.0.mlp.experts.0.gate_proj")), which resolves NVFP4 throughquantized_layers. The attn/dense/lm_head flags are dropped: qwen4_exp dense weights are always bf16 and nothing read them.Tests:
parse_configagainst thequantization_configof each released checkpoint (RadixArk NVFP4, nvidia NVFP4, Qwen FP8, bf16).Verified end to end (offload MoE, greedy answers correct):
docs/models.md lists the nvidia checkpoint.