fix(models): better support for mixed-precision compressed-tensors NVFP4

[Sam-Izdat]

Summary

Adds support for compressed-tensors checkpoints that mix per-tensor
FP8 attention/shared-expert projections with NVFP4 routed
experts - a layout produced by llm-compressor and found in the wild.

Without this, such checkpoints either crash on load (FP8 attn gets
torch.cat'd with bf16 weights) or produce garbled output (the dequant
kernel multiplies by the per-row global scale, but llm-compressor
stores the QUANT-side scale rather than the DEQUANT-side divisor that
modelopt stores).

The fix is in five files, +236/-63 lines. The key idea is to use the
on-disk tensor naming as the ground truth for which convention is
in use (weight_packed -> llm-compressor, weight_scale_2 -> modelopt)
rather than relying on the config-side format: nvfp4-pack-quantized
string (which both exporters set).

Test models

HF repo	Layout	Result
nvidia/Qwen3.6-35B-A3B-NVFP4	FP8 attn + NVFP4 shared-expert + NVFP4 experts (modelopt mixed_precision)	regression-clean, coherent
primitive-ai/Ornith-1.5-35B-A3B-agentic-NVFP4-FP8	FP8 attn + FP8 shared-expert + NVFP4 experts (compressed-tensors mixed-precision)	coherent (was: garbled output)

Tested with --moe-backend offload --expert-load parallel on RTX 3060
12 GB (nvfp4_backend='triton'). Ornith decoding at ~30-40+ tok/s with minimal tuning.

Fixed

1. Crash on mixed-precision attn+shared-expert — the
   compressed-tensors iter doesn't know about FP8. Routing
   compressed-tensors MoE to the modelopt iter (which handles FP8
   attn) when the shared-expert is in the FP8 group avoids the
   cat-of-fp8-and-bf16 crash.

2. Missing global-reciprocal for llm-compressor NVFP4 experts —
   the dequant kernel multiplies by the per-row global; modelopt
   stores the DEQUANT-side divisor directly, but llm-compressor
   stores the QUANT-side scale that must be reciprocated. A
   per-checkpoint flag on the NVFP4 source spec picks the right
   convention.

3. Unknown NVFP4 kind input_global_scale (parallel build) —
   the bank dispatch only knew weight_scale_2 (modelopt). Alias
   weight_global_scale -> weight_scale_2 via spec.kind_map;
   skip input_scale / input_global_scale (activation scales,
   not bank tensors).

4. Per-expert naming regex too narrow for parallel build — the
   expert key pattern only matched weight | weight_scale | weight_scale_2. Extended to cover the full llm-compressor
   per-expert naming set so the parallel bank builder's
   weight_info populates correctly.

5. Iter dispatch routed mixed-precision MoE to the
   compressed-tensors iter, which holds more GPU state per-shard
   than the default iter (caused OOM on 12 GB). Reverted to the
   default iter for that case.

6. The on-disk-naming heuristic for reciprocal was too coarse.
   format: nvfp4-pack-quantized is set by BOTH llm-compressor and
   modelopt re-exports. Now probes the safetensors index for
   weight_packed (llm-compressor signature) vs weight_scale_2
   (modelopt signature) and picks the convention from the data, not
   the config claim.

Files changed (5, +236/-63)

• python/freetoken/models/config.py — _nvfp4_global_reciprocal
  heuristic probing the safetensors index;
  ModelConfig.nvfp4_global_reciprocal field
• python/freetoken/models/nvfp4_banks.py — bank kind dispatch
  handles weight_global_scale (via spec.kind_map) and skips
  activation scales; serial and parallel paths updated in lockstep
• python/freetoken/models/qwen3_5_moe/config.py — _has_moe_experts
  helper; _attn_quant and _expert_quant extended for
  compressed-tensors; dense_quant probe for FP8 shared-expert;
  parse_config wiring
• python/freetoken/models/qwen3_5_moe/moe.py — _SharedExpert
  refactored to use the quant_linear factory, with
  attn_quant='fp8_pertensor' dispatch for the mixed-precision
  shared-expert path
• python/freetoken/models/qwen3_5_moe/weight.py — per-checkpoint
  _spec_for with kind_map; _load_maybe_quantized and
  _nvfp4_parts accept llm-compressor naming; _PT_FP8_FUSE and
  _CT_NVFP4_FUSE extended for shared-expert gate|up fuse; expanded
  _NVFP4_EXPERT_KEY_RE regex

Backwards compatibility

• NVIDIA Qwen3.6-35B-A3B-NVFP4 (modelopt) is regression-clean.
  No NVFP4 reciprocal applied (the data shows weight_scale_2 only,
  which signals the modelopt convention).
• The default iter (the one most production single-group NVFP4 MoE
  exports use) is unchanged in behavior.

Notes

The heuristic that disambiguates llm-compressor from modelopt
probes the safetensors index for a routed-expert tensor's sibling
suffixes. This is a one-time read at model-load time; no per-tensor
overhead. If both weight_packed and weight_scale_2 are absent
(e.g. some other compressed-tensors flavor), the heuristic falls
back to the safe default of "no reciprocal" same as
global_reciprocal=False upstream.

[salekseev]

Thank you for this change. We used it as the base for per-channel compressed-tensors FP8 support, and we found one gap.

The new branch in _attn_quant tests strategy == "tensor". A checkpoint that uses strategy: "channel" does not match the test, so it still gets "none".

We checked unsloth/Qwen3.6-35B-A3B-NVFP4-Fast against your head commit:

tree	attn_quant
main af71ba4	none
#390 head	none

The value does not change. This value is a routing key, so the checkpoint never reaches _iter_weights_attn_fp8.

The correction is one token:

if w.get("group_size") is None and w.get("strategy") in ("tensor", "channel"):

This is safe. Fp8PerTensorLinear.weight_scale is per output row. A true per-tensor weight holds the same scalar in every row. Both strategies use the same layer, and no kernel change follows.

One more change is necessary with it. _per_row_scale in weight.py does scale.reshape(1). A per-channel scale is [rows, 1], so the reshape fails:

RuntimeError: shape '[1]' is invalid for input of size 248320

We have both changes on a branch with tests. Tell us if you want them in this pull request, and we will send them. We will not open a competing pull request.

The full analysis and the measurements are in #252.

---

The change is on a branch, if you want to cherry-pick it.

• Commit: f4e3270
• Branch: upstream/ct-fp8-channel-strategy
• Base: this pull request's head (6b78a6c)

The commit holds the one-token change and a new test file, tests/models/test_qwen3_5_moe_config.py. The test covers the geometries that must route to the per-tensor fp8 layer, and the geometries that must not: a group or block scale, a group_size with a per-channel name, 4-bit weights, int weights, and a group that targets only the experts.

Use it, change it, or ignore it, as you prefer.

---

One more observation, for your consideration. It is not a defect in this pull request.

The new branch does not consult the checkpoint's ignore list. compressed-tensors gives ignore precedence over a group's targets, and llm-compressor often pairs a broad target with an ignore list that carves modules back out. A module in ignore is bf16 on disk. If the branch claims such a module as fp8, the loader builds an fp8 layer for a bf16 weight, and the load stops on the dtype check.

models/qwen4_exp/config.py already does this for its own family. It reads get("ignore") and derives each flag from a probe module name. The comment there gives the reason: derive every flag from that list instead of assuming the split.

The matcher must differ, though. qwen4_exp uses fnmatch, which suits ModelOpt's plain module names. compressed-tensors writes regex entries such as re:.*lm_head, and fnmatch("lm_head", "re:.*lm_head") returns False. So a compressed-tensors check needs its own match, not a reuse of that helper.

Neither checkpoint we tested is affected. Neither lists these modules in ignore, and we have not seen a checkpoint that hits this. We mention it only because the same gap applies to every group the new branch inspects.

[Romeriz]

Hi -- thanks for this. Before I open a complementary PR, let me flag the overlap and a
couple of gaps so we don't ship conflicting changes in the same functions.

What overlaps (same base main, so the hunks collide):

• _CT_NVFP4_FUSE: I also add the identical .mlp.shared_expert.gate_up_proj gate|up entry.
• _expert_quant: I also recognize compressed-tensors MoE -> nvfp4.
• Per-expert llm-compressor bank naming: I map weight_packed -> weight,
  weight_global_scale -> weight_scale_2 with global_reciprocal too (I use a separate
  per-expert CT spec + a stacked spec selected by probing the weight map, rather than one
  widened regex + a config reciprocal flag).
• iter_weights MoE-CT routing is where we differ architecturally: you gate the CT reader
  to not _has_moe_experts (MoE CT goes to the fp8-aware modelopt branch); I keep MoE CT
  in the CT reader and teach it to skip experts, dequant an fp8/bf16 GDN, dequant NVFP4
  embed_tokens, and keep lm_head native when lm_head_quant == "nvfp4".

Why I don't think either patch alone covers the other's models:

• The exports I fixed (AEON-7/Qwen3.6-35B-A3B-heretic-NVFP4, Ttimms/KAT-Coder-V2.5-Dev-REAP-50-NVFP4A16,
  doth4580/Kwaipilot-KAT-Coder-V2.5-Dev-NVFP4-MIXED) keep the whole dense side NVFP4-
  .weight_packed (attn + shared expert), so routing them to the modelopt reader (which
  reads .weight-named NVFP4) doesn't parse them; and AEON is a single 23 GB file with no
  index, where an index-only reciprocal probe returns the modelopt default (no
  reciprocal) -> wrong dequant.
• Your Ornith-style export (FP8-attn config group + NVFP4 experts) needs the fp8 dense
  routing I don't implement.

Suggestion: rather than merge overlapping diffs, reconcile the MoE-CT dense routing
per variant (NVFP4-.weight_packed dense -> the CT reader; FP8-config_groups dense ->
the fp8 reader), and make the reciprocal probe fall back to reading the single-file
safetensors header. I'm happy to rebase my branch on yours once it's in a state you're
happy with, or fold my delta in. Want me to post the precise diff-by-hunk comparison, or
shall I just open my PR referencing this one so the maintainers can compare?

(tagging for visibility; happy to take this to the developer Slack if easier)

[Sam-Izdat]

@salekseev @Romeriz - Thank you both for feedback and corrections. I agree that we should consolidate, and either way is fine by me. I can fold in your changes if you like, or you can take what you need from here for one of your PRs, and I can close this one -- just let me know. Just a heads up: I'm sitting on an ancient workstation with a 3060 for compute, so that's the extent of the hardware testing I can do, if you need to target more demanding models.

[salekseev]

@Sam-Izdat — please don't close this on my account. Two of my four commits turned out not to touch anything you or @Romeriz touch, so I've just opened them as standalone PRs against main:

• fix(qwen3_5_moe): accept a per-channel fp8 weight_scale in the dense loader #415 — _per_row_scale accepting a [rows, 1] per-channel scale
• feat(qwen3_5_moe): keep a natively-fp8 lm_head native instead of failing to load #416 — an Fp8LMHead, so a natively-FP8 head loads instead of dying on Unexpected keys in state_dict: ['lm_head.weight_scale']

I checked rather than assumed: this PR's diff has no _per_row_scale, no Fp8LMHead and no _lm_head_quant in it, and neither appears in Romeriz's overlap list. They also don't care which way the MoE-CT routing question lands, which is why I pulled them out — they can be reviewed while the interesting argument continues.

The strategy: "channel" token from my earlier comment is a change to your branch, not a PR of mine. Take it or leave it; I'm not going to open something that competes with this.

I'm deliberately not weighing in on the routing overlap between you two. I've only exercised this against two checkpoints and I'd be guessing about the exports Romeriz is targeting. That call probably wants @jason-fxz, since as far as I can tell none of the three of us can merge anything here.

One correctness thing, because it's easy to lose in a structural discussion

Romeriz's point about the index probe is right, and I think it's sharper than it came across. _nvfp4_global_reciprocal says the right thing in its own docstring — the ground truth is the on-disk tensor naming — but it reads that naming out of model.safetensors.index.json:

if not idx_path.exists():
    return False  # single-file or no index -- fall back to safe default
except Exception:
    return False  # any read error -- don't reciprocate (modelopt default is safer)

For a single-file llm-compressor export, False picks the ModelOpt convention, so the reciprocal is skipped and the weights come out wrong with no error anywhere — just degraded output. The comment's "safer" holds for ModelOpt but inverts for the exporter this PR exists to support, and a silent wrong-numerics default is a rough one to debug from a bug report. Romeriz's AEON case (single 23 GB file, no index) hits it directly.

Worth considering: you already have the ground truth in your hand at read time. In my commit for the packed shared expert I pick the parts reader because the tensor I'm looking at literally ends in .weight_packed, straight out of f.keys() on the open file — no index, no config flag, no heuristic, and it can't be wrong for a single-file export because there's nothing to probe. Reading the safetensors header directly, as Romeriz suggested, would get you the same guarantee while keeping the config flag.

Offer

I've got an RTX 4080 SUPER (16 GiB, sm_89) with both unsloth/Qwen3.6-35B-A3B-NVFP4-Fast and nvidia/Qwen3.6-35B-A3B-NVFP4 on disk, plus needle-recall, BFCL tool-calling and evalscope harnesses already wired up for this engine. If a 3060 is the constraint, point me at a branch and I'll run the 35B-class checkpoints against it and post numbers. Same offer to you, @Romeriz — I can't obtain your NVFP4 exports quickly, but if the dense-side routing question comes down to "does this actually load and generate sensibly", I can be the box that answers it.

[salekseev]

Ran this on hardware. Took @Romeriz's suggestion and read a header instead of an index — more on that below, because there's a reproducible bug in it.

#390 loads and serves a real mixed-precision export

I pulled primitive-ai/Ornith-1.5-35B-A3B-mixed-NVFP4-FP8 (22.55 GiB) specifically because it exercises this PR and nothing of mine: its fp8 group is strategy: "tensor" with real module regexes, so your detection matches it as shipped, and its lm_head is in ignore, so my Fp8LMHead never fires. Detector output:

attn_quant='fp8_pertensor'  expert_quant='nvfp4'  lm_head_quant='none'  dense_mlp='none'

It loads through Loading mixed-fp8 weights, comes up on the fi attention backend, and generates coherently. RTX 4080 SUPER, 16 GiB, sm_89, TP=1, 3600 expert slots, fp8 KV at 262,144 tokens.

	this checkpoint	unsloth/…-NVFP4-Fast	nvidia/…-NVFP4
prefill @8k	3688 tok/s (8/8)	3311	3248–3618
decode @1200	106.95	120.05	123.4
longctx @32k	101.63	—	117.1
needle recall	9/9 to 176k	9/9 to 248k	9/9 to 248k
BFCL parallel	92.50	95.00 (mean of 2)	94.50 (mean of 2)
free VRAM	1.46 GiB	2.13 GiB	—

Prefill is the fastest of the three. BFCL 92.50 against unsloth's 95.00 is 5 flipped entries out of 200, inside the 8-flip noise floor I've measured on this category, and it's a single run where the others have two or three — so I'd call that unresolved rather than a regression. It's also a different base model (Ornith-1.5, not Qwen3.6), scored with its own tokenizer and chat template.

Two notes so nobody over-reads the table. The decode gap is mostly the bf16 lm_head — 0.947 GiB read per decoded token here versus unsloth's native-fp8 0.474 GiB. The only other difference is the shared expert (fp8 here, NVFP4 there, worth roughly +52 MiB/token), so the head is about 90% of the ~525 MiB traffic delta. That's two checkpoints differing in more than one variable, not a controlled experiment. And my build also carries an unrelated fp8-KV change, so don't read these as clean-tree numbers.

The reciprocal probe is wrong for single-file exports, and here's a repro

Romeriz was right, and it's worse than "falls back to a default" — the fallback silently picks the opposite convention with no error anywhere. Your own docstring already names the ground truth ("the on-disk tensor naming"); the problem is only that _nvfp4_global_reciprocal reads that naming out of model.safetensors.index.json rather than out of the tensors.

Using the checkpoint above, which is genuinely llm-compressor (weight_packed present, weight_scale_2 absent), I symlinked the six shards into a temp dir with the config and no index:

expert gate_proj suffixes on disk:
  ['input_global_scale', 'weight', 'weight_global_scale', 'weight_packed', 'weight_scale']
  -> llm-compressor naming (weight_packed): True
  -> ModelOpt naming (weight_scale_2):      False

WITH index:    _nvfp4_global_reciprocal = True     <- correct
WITHOUT index: _nvfp4_global_reciprocal = False    <- wrong
  (dir holds 6 shards, so the ground truth is present -- it is just never read)

Same tensors both times. Only the index file differs. False means "don't reciprocate", i.e. ModelOpt's convention applied to an llm-compressor checkpoint — degraded output, nothing raised, and nothing in a bug report to point at. That's Romeriz's AEON case (single 23 GB file, no index) reproduced on a published checkpoint.

Worth noting the except Exception: return False path has the same shape, so a permissions error or an unreadable index lands in the same place. If it's useful I'm happy to send a patch that reads the safetensors header directly and keeps the config flag as an override — but it's your PR, so say the word rather than me pushing at it.

Where I've landed

@Sam-Izdat — #415 and #416 are open against main and don't touch anything in here, so nothing you do with this PR is blocked on them. Still not opening anything that competes with this one, and still not going to referee the routing question between you and Romeriz; that wants @jason-fxz.

Offer stands on the hardware. I've got this box plus three 35B-class checkpoints on disk and the harnesses wired up, so if either of you wants a branch run through load + needle + BFCL, point me at it.