qwen3_5_moe: support unsloth/Qwen3.8-27B-NVFP4 (mixed-precision NVFP4+FP8)

[chrisqianz]

Problem

Mixed-precision compressed-tensors exports (e.g. unsloth/Qwen3.8-27B-NVFP4) cannot load: the dense pass assumes either pure NVFP4 (weight_packed) or pure bf16, but unsloth's layout stores attention linears as weight-only FP8 (per-row scale), most dense-MLP layers as NVFP4, a few MLP layers as FP8, and lm_head as FP8. Raw e4m3 tensors die in the bf16 fusion (torch.cat fp8/bf16 promotion), and bf16 dequant of lm_head/MLP doubles memory.

Fix (per-module native storage — no new kernels)

Reuses the existing Fp8PerTensor* kernels; no shared kernel / modelopt-pass changes:

• models/config.py: dense_mlp_storage per-layer override on ModelConfig
• models/loader.py: ShardReader.has() for scale-sibling detection
• qwen3_5_moe/config.py: _compressed_linear_storage sniffs the safetensors index weight_map (order-independent per-layer classification: weight_packed → nvfp4, scaled → fp8, plain → bf16) + FP8 lm_head detection
• qwen3_5_moe/moe.py + model.py: overridden layers build shared-expert / dense-MLP / lm_head linears as native W8A16
• qwen3_5_moe/weight.py: dense pass keeps scaled FP8 linears native (fused q/k/v → qkv_proj, in_proj_qkv/z → in_proj_qkvz, gate/up → gate_up_proj, singletons o_proj / GDN out_proj / down_proj / lm_head with per-row fp32 scales); unscaled fp8 dequantizes to bf16 (per-tensor / per-row / block scale). Official pure-NVFP4 checkpoints are unaffected (pass is a no-op there).
• docs/models.md: add unsloth/Qwen3.8-27B-NVFP4 to the known-good list

Verification (RTX 5090 D 32GB, driver 595.84, torch 2.11 cu130)

• Before: torch.cat fp8/bf16 promotion crash on the mixed layout
• After: native resident 21.809 GiB (vs ~54 GB if all dequantized to bf16)
• config tests 8/8, weight tests 14/14
• full-pass key/shape/dtype vs model state dict: 0 mismatch / 0 stray / 0 missing
• ft serve e2e: 21.8 GB load, KV 16K, CUDA graph capture, /health ok, real chat generation

Usage note: on a 32 GB card use ft serve --model unsloth/Qwen3.8-27B-NVFP4 --num-tokens 32768 --max-prefill-length 1024 (the large-vocab logits buffer OOMs at the default 8192-token prefill).

[yuyi2439]

maybe fix #238

[dcutugno]

Second data point for this PR, on a different variant of the same family.

Checkpoint: a third-party Qwen3.8-27B multimodal NVFP4 export — architectures: ["Qwen3_5ForConditionalGeneration"], model_type: qwen3_5, compressed-tensors 0.18.0, quantization_config.format == "mixed-precision". Same layout family you describe, but the vision-tower variant rather than the text-only one.

Env: freetoken 0.1.2+g932e87e8a, Python 3.12, Windows 11, sm_120.

Both failure modes you describe reproduce:

1. Stock: the fp8/bf16 promotion crash in the dense pass (tracked as [BUG] qwen3_5 dense: mixed-precision NVFP4 crashes in ct_bf16_fuse (Float8 × BFloat16 promotion) #238).
2. Working around it by dequantizing the FP8 group to bf16 at load — the approach in fix: [BUG] qwen3_5 dense: mixed-precision NVFP4 crashes in ct_bf16_fuse (Float8 × BFloat16 promotion) #275 — clears that crash and converts cleanly, but then dies in model construction:

File "freetoken/kernel/triton/nvfp4_linear.py", line 867, in load_state_dict
  s = state_dict.pop(_concat_prefix(prefix, "weight_scale"))
KeyError: 'model.layers.0.linear_attn.out_proj.weight_scale'

because detect_compressed_tensors_nvfp4() is true on any NVFP4 group and qwen3_5_moe/config.py then forces attn_quant = "nvfp4". That corroborates the per-module classification this PR introduces — no single global scalar describes the checkpoint.

Exact per-layer split, since the description says "most dense-MLP layers as NVFP4, a few MLP layers as FP8":

modules	group	quant
MLP gate/up/down_proj, layers 0-55	group_1 (168 targets)	nvfp4-pack-quantized
MLP gate/up/down_proj, layers 56-63	group_0	float-quantized FP8
linear_attn.{in_proj_qkv,in_proj_z,out_proj}, 48 layers	group_0	FP8
self_attn.{q,k,v,o}_proj, 16 layers	group_0	FP8

232 FP8 targets / 168 NVFP4. FP8 weight_scale is per-row BF16 [out, 1]; NVFP4 carries fp8 weight_scale plus fp32 weight_global_scale / input_global_scale. transform_config and sparsity_config are both empty.

One variant difference worth handling: here lm_head is not FP8. It's untied (tie_word_embeddings: false) and sits in the 304-entry ignore list, so it stays bf16. If the FP8 lm_head detection assumes presence, this is the checkpoint where it's absent — the "plain -> bf16" branch should cover it, but it's untested against this layout.

On footprint: dequantizing the 232 FP8 tensors takes them from 8.71 GiB to 17.42 GiB, giving a 30.06 GiB FTW from a 23 GiB source. Consistent with your 21.8 vs ~54 GB figures.

I can't run this one — 16 GB card, and 21.8 GiB resident won't fit — so I can't verify generation. Happy to test conversion and load-path behaviour on this variant if useful.

[dcutugno]

Follow-up to my comment above — I applied this PR's sources onto the released build (dropping the branch .py files in place of the Cython-compiled modules) and tested it on the multimodal variant. Mostly good news, plus one gap.

Working:

• This PR's own suites: 22/22 (8 config + 14 weight).
• ft checkpoint: 21.35 GiB FTW, 1027 weight tensors, 32.4 s — against 30.06 GiB when the FP8 group is dequantized to bf16 instead. Matches your 21.8 vs ~54 GB result.
• ft serve pointed directly at the HF source dir: loads and generates coherent text. So the per-module classification and the native FP8 path both work on this variant.

The gap: the ft checkpoint -> ft serve <ftw-dir> round-trip fails.

KeyError: 'model.layers.0.linear_attn.in_proj.weight'

_compressed_linear_storage classifies layers by sniffing model.safetensors.index.json. A converted FTW directory carries freetoken_weight.json instead and has no safetensors index, so the classification silently falls back:

config read from	attn_quant	dense_mlp_storage
HF source dir (index present)	fp8_pertensor	populated
converted FTW dir (no index)	nvfp4	None

With attn_quant back to nvfp4, gdn.py takes the else branch and builds the fused bf16 in_proj, while the FTW was written with in_proj_qkvz + per-row fp32 scales. Hence the missing key.

test_index_unavailable_falls_back_to_native covers the fallback itself, but in the FTW case the fallback is silently wrong rather than safe: the storage decision was already made at conversion time and then lost. Persisting the resolved dense_mlp_storage / attention storage into the FTW metadata (or re-deriving it from the FTW's own tensor names, where in_proj_qkvz vs in_proj is already unambiguous) would close it.

A concrete fix, tested. The FTW already carries what is needed: freetoken_weight.json records a dtype per tensor, and those dtypes describe the storage unambiguously.

model.layers.0.mlp.gate_up_proj.weight           uint8           -> nvfp4
model.layers.0.mlp.gate_up_proj.weight_scale     float8_e4m3fn
model.layers.56.mlp.gate_up_proj.weight          float8_e4m3fn   -> fp8
model.layers.56.mlp.gate_up_proj.weight_scale    float32
model.layers.0.linear_attn.in_proj_qkvz.weight   float8_e4m3fn   -> attention fp8
lm_head.weight                                   bfloat16        -> lm_head bf16

So _compressed_linear_storage can fall back to freetoken_weight.json when model.safetensors.index.json is absent, applying the same per-layer rule against dtypes rather than key suffixes: uint8 -> nvfp4, float8* -> fp8, anything else -> bf16; the presence of linear_attn.in_proj_qkvz at all implies fp8 attention. It returns the same 4-tuple, so nothing downstream changes.

I ran exactly that as a monkeypatch over this branch, and a converted FTW then loads and generates correctly with no other changes.

The alternative -- persisting the resolved (attn, dense fallback, per-layer overrides, lm_head) tuple into the FTW metadata at conversion -- needs fewer heuristics, but it requires a metadata field and leaves already-converted FTWs broken, whereas the dtype fallback repairs those too.

Not a blocker for ft serve straight off the HF checkpoint, which is how the PR description verifies it — but anyone converting first hits it.

Environment note in case it is useful: this ran on a 16 GB card, which is under the ~21.4 GiB resident. It works on Windows only because WDDM spills into shared system memory, at roughly 1.3 tok/s, with --max-running-requests 1 --cache-type naive --num-tokens 4096 --max-prefill-length 512 and FREETOKEN_MAMBA_SSM_DTYPE=bfloat16 (the fp32 GDN recurrent-state pool wants 3.52 GiB on its own at the defaults). Not a configuration I would suggest supporting — just context for the numbers above.