feat(qwen3_5_moe): read every checkpoint layout through the QuantConfig

[jason-fxz]

What

qwen3_5_moe picked one of four weight readers by string flags (attn_quant, dense_quant, lm_head_quant) that its config derived from the HF quantization_config on its own. Since #418 every layer builds its buffers from config.quant.scheme_for(prefix), so the flags and the readers disagreed with the model on every compressed-tensors export (#238, #263, #344, #381, #410, #437, and the PRs #208, #275, #296, #390, #413, #415).

One reader now asks the same scheme_for for every Linear module and fills exactly the buffers the model declared:

• qwen3_5_moe/weight.py: one _DenseReader replaces the four readers and their fusion tables. Tensors are matched to the roles of the module's scheme, packed projections are concatenated per role, dtypes and scale shapes are checked against the scheme, per-tensor scales are broadcast per row (the shape handling from fix(qwen3_5_moe): accept a per-channel fp8 weight_scale in the dense loader #415). A module the config calls bf16 but whose tensors are quantized is rejected; a module the family serves bf16 (the routers) is dequantized.
• qwen3_5_moe/config.py: the string detectors are gone; expert_quant comes from the QuantConfig, as in qwen4_exp.
• Each dialect declares STORAGE, the checkpoint tensor behind every role of every kind it exports (weight_packed / weight_global_scale for llm-compressor, weight / weight_scale_2 for ModelOpt), with the reciprocal convention on the entry. The NVFP4 expert reader spec is derived from it.
• ct_set: a compressed-tensors targets / ignore name covers that module alone, not its children, as in llm-compressor and vLLM. Exports list every unquantized module including containers (layers.N.linear_attn, mlp.experts.N), and the old ancestor match turned the quantized projections under them into bf16. Target class names other than Linear fail closed.
• Nvfp4LinearMethod declares input_scale when the scheme carries it, so W4A4 checkpoints load in full; no kernel reads it yet. The gemma4, glm4_moe and muse_glimmer readers load it too.
• loader.safetensors_weight_map reads the index or, without one, the shard headers; it replaces the hard index read in nvfp4_banks.py that failed on single-file checkpoints.
• glm5_next keeps its own copy of _fp8_block_quant instead of importing it from qwen3_5_moe.

Compatibility

An FTW written by an older build from an NVFP4 dense checkpoint (for example nvidia/Qwen3.6-27B-NVFP4, nvidia/Qwen3.6-35B-A3B-NVFP4, RadixArk/Qwen3.8-27B-NVFP4, nvidia/Gemma-4-31B-IT-NVFP4, nvidia/GLM-4.7-NVFP4) lacks the input_scale tensors the model now declares and fails with a KeyError. scripts/ftw_hotfix.py --ftw <dir> --source <hf dir> adds them.

Tested

H100 80 GB, driver 580.95, CUDA 13.0, torch 2.11, on top of main @ fb7f732. tests/e2e/test_aime.py (AIME25 problem 1, pass@3) with FREETOKEN_TEST_MODEL=<dir>, plus FREETOKEN_TEST_MOE_CACHE_AUTO=1 for MoE.

checkpoint	layout	result
sakamakismile/Qwen3.6-27B-NVFP4	compressed-tensors NVFP4, GDN quantized	pass (main: KeyError in_proj_qkvz)
unsloth/Qwen3.6-35B-A3B-NVFP4-Fast	compressed-tensors channel-fp8 + NVFP4, fp8 lm_head	pass
RedHatAI/Qwen3.6-35B-A3B-NVFP4	compressed-tensors NVFP4	pass
primitive-ai/Ornith-1.5-35B-A3B-mixed-NVFP4-FP8	compressed-tensors static fp8 + NVFP4	pass
kyaky/Qwen3.6-35B-A3B-Uncensored-NVFP4	compressed-tensors block-fp8 + NVFP4, single file	pass
unsloth/Qwen3.8-27B-NVFP4	compressed-tensors channel-fp8 + NVFP4, dense	pass
nvidia/Qwen3.6-35B-A3B-NVFP4	ModelOpt FP8 + NVFP4	pass
nvidia/Qwen3.6-27B-NVFP4	ModelOpt FP8 + NVFP4, dense	pass
RadixArk/Qwen3.8-27B-NVFP4	ModelOpt FP8 + NVFP4, dense	pass
Qwen/Qwen3.6-35B-A3B-FP8	block-fp8	pass
Qwen/Qwen3.8-27B-FP8	block-fp8, dense	pass
Qwen/Qwen3.6-27B	bf16, dense	pass
Qwen/Qwen3.5-2B	bf16, dense, tied embeddings	pass
Qwen/Qwen3.6-35B-A3B	bf16	pass
nvidia/Gemma-4-31B-IT-NVFP4	ModelOpt NVFP4 dense MLP (gemma4 reader)	pass
nvidia/Qwen3.6-27B-NVFP4 FTW from an older build		KeyError ... input_scale as expected; pass after ftw_hotfix.py --source
RedHatAI/Muse-Glimmer-30B-NVFP4	compressed-tensors NVFP4 (muse_glimmer reader)	loads; no answer within the 16k-token cap, same as main
nvidia/GLM-4.7-NVFP4	ModelOpt NVFP4 shared experts (glm4_moe reader)	loads (dense pass); AIME not run, 215 GB of experts do not fit this host's RAM

pytest tests/models tests/moe/test_offload.py tests/engine/test_cache_budget.py: 255 passed.

Closes #238, closes #344, closes #381. Implements the migration proposed in #437, left open for the int4 discussion. Addresses the compressed-tensors NVFP4 side of #263 (sakamakismile/Huihui-Qwen3.6-35B-A3B-abliterated-NVFP4; the joshebbs export with the nested ModelOpt config is not covered) and of #410 (the sahilchachra, sakamakismile and Ttimms KAT-Coder exports have regular configs; doth4580's format-only config is not covered). Not #252: its ornith-ai/Ornith-1.5-35B-A3B-FP8 quantizes the routed experts per-channel fp8 too, and there is no fp8 per-tensor MoE method yet. Supersedes #208, #275, #296, #390, #415 (its per-channel scale handling is included) and the reader and detection parts of #413 (its KAT-Coder handling is not included).

[salekseev]

Ran #438 on sm_89 with MoE offload, since the table in the PR is H100. It works out of the box, 953659d over the gfb7f732de nightly wheels.

RTX 4080 SUPER, 16 GiB, sm_89, TP=1, --moe-backend offload, 3600 expert slots, 62 GB host RAM.

checkpoint	dialect	result
unsloth/Qwen3.6-35B-A3B-NVFP4-Fast	CT, fp8 per-channel dense + NVFP4	healthy in 40 s; decode 123.46 tok/s; prefill @8k 3648 tok/s, 8/8; coherent; 2/2 parallel tool calls
nvidia/Qwen3.6-35B-A3B-NVFP4	ModelOpt MIXED_PRECISION	healthy in 50 s, 2.07 GiB free, generates correctly -- no regression
primitive-ai/Ornith-1.5-35B-A3B-mixed-NVFP4-FP8	CT, fp8 per-tensor dense + NVFP4	healthy in 50 s, 1.50 GiB free, generates correctly

The 8192-token prefill is the canary I care about on a 16 GiB card, because that is where the GDN chunked-prefill workspace has OOM'd before (#401). 8/8 at 3648 tok/s is the best I have recorded for this checkpoint.

This is now my serving default rather than a one-off test, so it has continuous use on sm_89. At the model's full 262,144-token context it measures 122.4 tok/s decode and 3398 tok/s prefill at 8k, 8/8, which is at or slightly above what I was getting before this branch.

I confirmed your point 1 directly rather than taking it on trust. checkpoint_quant_config on this branch:

Avesed/Qwen3.6-35B-A3B-INT4-W4A16      compressed-tensors -> NotImplementedError: weight scheme {...}
cyankiwi/Qwen3.6-35B-A3B-AWQ-4bit      awq                -> NotImplementedError: 'awq' is not supported
cyankiwi/Qwen-AgentWorld-35B-A3B-...   compressed-tensors -> NotImplementedError: weight scheme {...}
nvidia / unsloth / primitive-ai        modelopt / CT      -> ModelOptConfig / CompressedTensorsConfig

You are right that the legacy detectors were never protecting those three. I had proposed keeping them as a fallback and that was wrong: they only ever loaded for me because of local int4 patches, not because of anything on main. My "never downgrade a verdict" rule was guarding something that does not exist upstream. I have retired those checkpoints rather than keeping a fallback alive, and I am treating #396 as where that work belongs, with the scope you laid out. The repack step in particular I had not thought about, and it is the part that would have bitten me.

Two findings from prototyping the same migration before you posted, in case either is useful.

The ignore ancestor-semantics fix is the one that matters most, and it is worth a regression test if there isn't one already. It is not a corner case: unsloth lists ...layers.N.linear_attn beside the genuinely-bf16 in_proj_b/in_proj_a/norm, and the Ornith mixed export lists all 10,240 ...experts.M parents. Under ancestor semantics both read as bf16, and on a 16 GiB card that surfaces as KeyError: '...linear_attn.in_proj.weight' -- the GDN builds a fused in_proj because scheme_for('...in_proj_qkvz') returned None, while the reader emits the split names. It took me a while to connect those two.

Separately, the layers/quantization matchers built from closures cannot be pickled: AttributeError: Can't pickle local object 'name_set.<locals>.hit'. It does not bite today because the worker rebuilds the config rather than receiving it, but if a QuantConfig is ever placed on a config object that crosses into the scheduler worker it will, and the failure mode is silent-unquantized rather than a pickling error. Making those matchers classes with __call__ would pre-empt it cheaply.

One more datapoint on this as a base to build on: I carry a per-tensor fp8 KV change of my own, and rebasing it onto this branch needed hand-work on only two hunks, both in _adjust_config and both from the moe_backend -> moe_strategy rename shifting context. Nothing interacted with the scheme reader.

#415 is closed, since this carries it. Thanks for crediting the _per_row_scale logic. Happy to keep reporting sm_89 numbers as this evolves, and to run needle-recall or BFCL on it if that would be useful for the PR table.