diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 1bce039cb..95d0cfb37 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -91,6 +91,66 @@ def detect_compressed_tensors_nvfp4(hf_config: Any) -> bool: return saw_nvfp4 +def _nvfp4_global_reciprocal(hf_config: Any) -> bool: + """Whether the routed-experts' per-tensor global scale needs to be reciprocated + at load time. + + llm-compressor stores the QUANT-side scale (the multiplier baked into the local + fp8 scales before the FP4 cast); the dequant kernel divides by it, so the loader + must reciprocate 1/x. Modelopt's NVFP4 stores the DEQUANT-side divisor directly + and the kernel multiplies (no transform). + + The ``format: nvfp4-pack-quantized`` string on the config_group is set by BOTH + exporters, so it alone can't disambiguate. The ground-truth signal is the on-disk + tensor naming: llm-compressor uses ``weight_packed`` + ``weight_global_scale`` + + ``input_global_scale``; modelopt uses ``weight`` + ``weight_scale_2`` + ``input_scale``. + We probe the safetensors index for one routed-expert tensor and pick the convention + from its sibling suffixes. + + Returns True (reciprocate) for llm-compressor naming, False (use as-is) for + modelopt naming. Returns False when nothing is found (caller can't decide, so the + bank loader keeps the on-disk value verbatim -- the modelopt default).""" + # Cheap pre-check: ``quant_method`` should be compressed-tensors for either case. + quant = getattr(hf_config, "quantization_config", None) + if not quant: + return False + get = quant.get if isinstance(quant, dict) else (lambda k, d=None: getattr(quant, k, d)) + if str(get("quant_method") or "").lower() != "compressed-tensors": + return False + + # Probe the safetensors index for an expert tensor and inspect its sibling suffixes. + # Cheap read (one json parse) and bounds-checked (only the first hit is needed). + try: + from pathlib import Path as _Path + from freetoken.utils.hf import download_hf_weight as _download + folder = _Path(_download(hf_config._name_or_path + if hasattr(hf_config, "_name_or_path") else None)) + idx_path = folder / "model.safetensors.index.json" + if not idx_path.exists(): + return False # single-file or no index -- fall back to safe default + import json as _json + weight_map = _json.loads(idx_path.read_text())["weight_map"] + except Exception: + return False # any read error -- don't reciprocate (modelopt default is safer) + + # Find one expert key in any shard; the suffix names reveal the layout. + has_packed = False # llm-compressor writes weight_packed + has_scale2 = False # modelopt writes weight_scale_2 + for name in weight_map: + if ".mlp.experts." in name and ".gate_proj." in name: + has_packed = has_packed or name.endswith(".weight_packed") + has_scale2 = has_scale2 or name.endswith(".weight_scale_2") + if has_packed or has_scale2: + break + # Disambiguate: llm-compressor uses weight_packed (and weight_global_scale); modelopt + # uses weight + weight_scale_2 directly. A safetensors checkpoint that has weight_packed + # for routed experts is llm-compressor; one that has weight_scale_2 is modelopt. + if has_packed: + return True + return False + + + @dataclass(frozen=True) class RotaryConfig: head_dim: int @@ -270,6 +330,11 @@ class ModelConfig: expert_quant: str = "none" # NVFP4 routed-expert GEMM backend (--nvfp4-backend); injected from EngineConfig. nvfp4_backend: str = "triton" + # Reciprocate the per-tensor NVFP4 expert global scale at load (llm-compressor NVFP4 + # stores the QUANT-side scale; the dequant kernel multiplies by the DEQUANT-side value, + # so we read 1/x from disk). False for modelopt NVFP4 (dequant-side value is on disk). + # See ``_nvfp4_global_reciprocal`` for the heuristic. + nvfp4_global_reciprocal: bool = False # Block size (out, in) for block-wise weight quantization (fp8_block: (128, 128)). weight_block_size: tuple[int, int] | None = None # Weight quantization of the *dense* attention / GatedDeltaNet projections (separate diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 6b933ff1d..173aa0693 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -130,9 +130,13 @@ def load_nvfp4_expert_source_banks( raise ValueError(f"{spec.desc}: unknown NVFP4 expert projection {proj!r}") kind = _canon_kind(spec, match.group("kind")) if kind == "weight_scale_2": + # modelopt's ``weight_scale_2`` and llm-compressor's ``weight_global_scale`` + # (aliased via spec.kind_map) both carry the per-row global. global_shards[shard].append((name, match, bank_layer)) elif kind in {"weight", "weight_scale"}: weight_shards[shard].append((name, match, bank_layer)) + elif kind in {"input_scale", "input_global_scale"}: + pass # activation scale (W4A8 / W8A8); consumed with its .weight else: raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}") @@ -254,9 +258,13 @@ def load_nvfp4_expert_source_banks_parallel( continue kind = _canon_kind(spec, match.group("kind")) if kind == "weight_scale_2": + # modelopt names the per-row global ``weight_scale_2``; llm-compressor names it + # ``weight_global_scale`` (after ``_canon_kind`` maps via spec.kind_map, or alias). global_names_by_shard[shard].append(name) elif kind in {"weight", "weight_scale"}: weight_info[name] = (match, bank_layer) + elif kind in {"input_scale", "input_global_scale"}: + pass # activation scale (W4A8 / W8A8); consumed with its .weight, not a bank tensor else: raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}") diff --git a/python/freetoken/models/qwen3_5_moe/config.py b/python/freetoken/models/qwen3_5_moe/config.py index 2ac4b607f..f40c49fc4 100644 --- a/python/freetoken/models/qwen3_5_moe/config.py +++ b/python/freetoken/models/qwen3_5_moe/config.py @@ -7,7 +7,9 @@ LinearGatedDeltaGroupConfig, ModelConfig, RotaryConfig, + _nvfp4_global_reciprocal, detect_compressed_tensors_nvfp4, + detect_expert_quant, ) @@ -39,6 +41,15 @@ def _fp8_block_quant(hf_config: Any) -> tuple[str, tuple[int, int] | None]: return "none", None +def _has_moe_experts(hf_config: Any) -> bool: + """True if the model declares any routed experts (``text_config.num_experts > 0``). + Used to disambiguate dense compressed-tensors NVFP4 (Qwen3.6-27B) from MoE mixed-precision + compressed-tensors (Ornith: FP8 attn/shared_expert + NVFP4 experts) so per-group quant + detection isn't overridden by a blanket "everything is NVFP4" assumption.""" + text = getattr(hf_config, "text_config", hf_config) + return (getattr(text, "num_experts", 0) or 0) > 0 + + def _expert_quant(hf_config: Any) -> str: """Quantization format of the *routed* experts (the only weights served from the offload cache). The nvidia/modelopt checkpoints are either plain NVFP4 (``quant_algo`` @@ -60,6 +71,12 @@ def _expert_quant(hf_config: Any) -> str: return "nvfp4" if "fp8" in expert_algo: return "fp8" + if str(get("quant_method") or "").lower() == "compressed-tensors": + # llm-compressor: the routed experts' group carries ``format: nvfp4-pack-quantized``. + # ``detect_expert_quant`` (top-level) already returns "nvfp4" for this -- mirror it + # here so the offload expert bank activates. + if detect_expert_quant(hf_config) == "nvfp4": + return "nvfp4" return "none" @@ -108,17 +125,31 @@ def _attn_quant(hf_config: Any) -> str: """Per-tensor FP8 on the *dense* attention/GDN projections. The modelopt ``MIXED_PRECISION`` checkpoints tag ``self_attn.{q,k,v,o}_proj`` and ``linear_attn.{in_proj_qkv,in_proj_z,out_proj}`` with ``quant_algo`` ``FP8`` (fp8-e4m3 - weight + a scalar ``weight_scale``; W8A16). Returns ``"fp8_pertensor"`` when present, - else ``"none"`` (NVFP4 dense weights -- shared_expert/lm_head -- stay dequant-at-load).""" + weight + a scalar ``weight_scale``; W8A16). llm-compressor mixed-precision checkpoints + (Ornith) use a per-tensor-FP8 ``config_groups`` entry targeting the same projections. + Returns ``"fp8_pertensor"`` when present, else ``"none"``.""" get = _quant_accessor(hf_config) if get is None: return "none" layers = get("quantized_layers") or {} - if not isinstance(layers, dict): - return "none" - for name, spec in layers.items(): - algo = str((spec or {}).get("quant_algo", "")).lower() - if algo == "fp8" and (".self_attn." in name or ".linear_attn." in name): + if isinstance(layers, dict): + for name, spec in layers.items(): + algo = str((spec or {}).get("quant_algo", "")).lower() + if algo == "fp8" and (".self_attn." in name or ".linear_attn." in name): + return "fp8_pertensor" + # llm-compressor path: a config_group with format "float-quantized" and per-tensor + # strategy (``group_size is None and strategy == "tensor"``) targeting ``.self_attn.`` + # or ``.linear_attn.`` is per-tensor FP8 (W8A16). Targets are regex strings + # (``re:.*\.self_attn\.``), so the substring check uses the unescaped leaf path. + for g in (get("config_groups") or {}).values(): + if not g: + continue + if not any((".self_attn" in t or "linear_attn" in t) for t in (g.get("targets") or [])): + continue + w = g.get("weights") or {} + if str(w.get("type", "")).lower() != "float" or int(w.get("num_bits", 0) or 0) != 8: + continue + if w.get("group_size") is None and w.get("strategy") == "tensor": return "fp8_pertensor" return "none" @@ -170,19 +201,36 @@ def parse_config(hf_config: Any) -> ModelConfig: # Dense attention/GDN quant is independent of the routed experts (block-fp8 already # quantizes both, so only probe for per-tensor FP8 when experts aren't block-fp8). attn_quant = "none" if expert_quant == "fp8_block" else _attn_quant(hf_config) - # NVFP4 checkpoints store the dense MLP projections (shared_expert; dense non-MoE MLP) as - # packed FP4 exactly like the routed experts -- independent of whether attention is FP8 - # (mixed) or bf16 (pure NVFP4). Keep them native FP4 (W4A16) whenever the experts are - # NVFP4. The lm_head is detected separately (only the mixed checkpoint quantizes it). - # MoE-NVFP4 keeps the shared_expert dense MLP native FP4 (expert_quant=="nvfp4"); a dense - # (non-MoE) modelopt checkpoint instead tags the bare .mlp.{gate,up,down}_proj as NVFP4. - dense_quant = "nvfp4" if expert_quant == "nvfp4" else _dense_mlp_quant(hf_config) + # Default: shared expert / dense MLP match the routed experts' quant (NVFP4 when experts + # are NVFP4). The modelopt dense (non-MoE) checkpoint sets a per-projection override via + # ``quantized_layers`` that ``_dense_mlp_quant`` picks up. llm-compressor mixed-precision + # MoE (Ornith: NVFP4 experts + FP8 attn + FP8 shared_expert) is the one true exception: + # its config_groups put the shared_expert in the FP8 group, so we explicitly read that + # and override dense_quant to "none" -- the FP8 path is selected via attn_quant. + if expert_quant == "nvfp4": + mlp_q = _dense_mlp_quant(hf_config) + dense_quant = mlp_q if mlp_q != "none" else "nvfp4" + if _compressed_tensors_nvfp4(hf_config) and _has_moe_experts(hf_config): + quant = getattr(hf_config, "quantization_config", None) + get = quant.get if isinstance(quant, dict) else (lambda k, d=None: getattr(quant, k, d)) + for g in (get("config_groups") or {}).values(): + if not g: + continue + if any("shared_expert" in t for t in (g.get("targets") or [])): + if str(g.get("format") or "").lower() == "float-quantized": + dense_quant = "none" + break + else: + dense_quant = _dense_mlp_quant(hf_config) lm_head_quant = _lm_head_quant(hf_config) - # compressed-tensors NVFP4 (dense Qwen3.6-27B): the attention (q/k/v/o, GDN out_proj) AND - # the dense MLP are W4A16 NVFP4; GDN in_proj_*, lm_head, norms stay bf16. Wire the shared - # W4A16 kernels (attn_quant=="nvfp4" routes the attention/GDN linears through them too). - if _compressed_tensors_nvfp4(hf_config): + # compressed-tensors NVFP4 *dense* (Qwen3.6-27B-NVFP4): every Linear is W4A16 NVFP4 + # (q/k/v/o, GDN out_proj, dense MLP). GDN in_proj_*, lm_head, norms stay bf16. Wire the + # shared W4A16 kernels (attn_quant=="nvfp4" routes the attention/GDN linears through + # them too). MoE compressed-tensors checkpoints (Ornith: FP8 attn + FP8 shared_expert + # + NVFP4 experts) skip the override -- the per-group detectors above already chose + # the right per-tensor quant for each block. + if _compressed_tensors_nvfp4(hf_config) and not _has_moe_experts(hf_config): attn_quant = "nvfp4" dense_quant = "nvfp4" lm_head_quant = "none" @@ -257,6 +305,7 @@ def parse_config(hf_config: Any) -> ModelConfig: attn_quant=attn_quant, dense_quant=dense_quant, lm_head_quant=lm_head_quant, + nvfp4_global_reciprocal=_nvfp4_global_reciprocal(hf_config), ) diff --git a/python/freetoken/models/qwen3_5_moe/moe.py b/python/freetoken/models/qwen3_5_moe/moe.py index b5ab1cf3d..ea0b4ddcb 100644 --- a/python/freetoken/models/qwen3_5_moe/moe.py +++ b/python/freetoken/models/qwen3_5_moe/moe.py @@ -3,43 +3,41 @@ from typing import TYPE_CHECKING import torch -from freetoken.layers import ( - BaseOP, - LinearColParallelMerged, - LinearReplicated, - LinearRowParallel, - make_moe_layer, - silu_and_mul, -) - -from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged, Fp8BlockLinear +from freetoken.layers import BaseOP, LinearReplicated, make_moe_layer, silu_and_mul if TYPE_CHECKING: from freetoken.models.config import ModelConfig class _SharedExpert(BaseOP): - """Always-present shared SwiGLU expert of width ``shared_expert_intermediate_size``.""" + """Always-present shared SwiGLU expert of width ``shared_expert_intermediate_size``. + The quant type follows ``dense_quant`` (NVFP4 W4A16 for modelopt/shared_expert NVFP4 + and dense Qwen3.6-27B; per-tensor FP8 W8A8 for llm-compressor mixed-precision Ornith), + then ``attn_quant`` (per-tensor FP8 when shared_expert is in the FP8 group), then + ``expert_quant`` (block FP8), else bf16. The dispatch lives in models/quant_linear.py + so every dense col-merged linear (attn, GDN, shared_expert) picks the same kernel.""" def __init__(self, config: ModelConfig, hidden_size: int, intermediate_size: int): - if getattr(config, "expert_quant", "none") == "fp8_block": - self.gate_up_proj = Fp8BlockColMerged( - hidden_size, [intermediate_size, intermediate_size], has_bias=False - ) - self.down_proj = Fp8BlockLinear(intermediate_size, hidden_size, has_bias=False) - elif getattr(config, "dense_quant", "none") == "nvfp4": - # NVFP4 checkpoint: keep the shared expert's NVFP4 weights native (W4A16). - from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged, Nvfp4DenseLinear - - self.gate_up_proj = Nvfp4DenseColMerged( - hidden_size, [intermediate_size, intermediate_size], has_bias=False - ) - self.down_proj = Nvfp4DenseLinear(intermediate_size, hidden_size, has_bias=False) - else: - self.gate_up_proj = LinearColParallelMerged( - hidden_size, [intermediate_size, intermediate_size], has_bias=False - ) - self.down_proj = LinearRowParallel(intermediate_size, hidden_size, has_bias=False) + from freetoken.models.quant_linear import make_col_merged_quant, make_replicated_quant + + # ``dense_quant`` wins for the NVFP4 path (shared_expert NVFP4 in modelopt mixed + # precision and in dense Qwen3.6-27B); ``attn_quant`` for FP8 paths. + # The factory's pre-NVFP4 check (``expert_quant == "fp8_block"``) still fires for + # block-fp8 models; per-tensor FP8 takes the attn_quant path; everything else bf16. + dense_quant = getattr(config, "dense_quant", "none") + attn_quant = getattr(config, "attn_quant", "none") + expert_quant = getattr(config, "expert_quant", "none") + # For NVFP4 dense, the factory's check is on ``attn_quant``; promote ``dense_quant`` + # to ``attn_quant`` for the call so a single dispatch covers both. Block-fp8 and + # per-tensor-fp8 are unchanged. + fused_attn = "nvfp4" if dense_quant == "nvfp4" else attn_quant + self.gate_up_proj = make_col_merged_quant( + expert_quant, fused_attn, hidden_size, + [intermediate_size, intermediate_size], has_bias=False, + ) + self.down_proj = make_replicated_quant( + expert_quant, fused_attn, intermediate_size, hidden_size, has_bias=False, + ) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down_proj.forward(silu_and_mul(self.gate_up_proj.forward(x))) diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index d07f18cd7..8a898fbfb 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -25,7 +25,7 @@ from freetoken.utils import cached_load_hf_config, download_hf_weight from tqdm import tqdm -from .config import _compressed_tensors_nvfp4, parse_config +from .config import _compressed_tensors_nvfp4, _has_moe_experts, parse_config # Expert weights are stored pre-fused per layer: experts.gate_up_proj / experts.down_proj. _PACKED_EXPERT_PATTERN = re.compile( @@ -39,7 +39,8 @@ _NVFP4_EXPERT_RE = re.compile(r"\.mlp\.experts\.\d+\.") _NVFP4_EXPERT_KEY_RE = re.compile( r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." - r"(?Pgate_proj|up_proj|down_proj)\.(?Pweight|weight_scale|weight_scale_2)$" + r"(?Pgate_proj|up_proj|down_proj)\.(?Pweight|weight_scale|weight_packed|" + r"weight_scale_2|weight_global_scale|input_scale|input_global_scale)$" ) _NVFP4_SOURCE_SPEC = Nvfp4ExpertSourceSpec( key_pattern=_NVFP4_EXPERT_KEY_RE, @@ -47,10 +48,13 @@ layer_to_bank=lambda layer, config: layer, # every layer is MoE desc="Qwen3.5 NVFP4 experts", ) -# Suffixes of the per-tensor modelopt quant scales; consumed alongside their ``.weight``, -# never yielded on their own. -_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") - +# Suffixes of the per-tensor quant scales; consumed alongside their ``.weight`` (modelopt's +# ``weight_scale``/``weight_scale_2``/``input_scale``; llm-compressor's ``weight_packed``/ +# ``weight_global_scale``/``input_global_scale``), never yielded on their own. +_SCALE_SUFFIXES = ( + ".weight_scale", ".weight_scale_2", ".weight_global_scale", + ".input_scale", ".input_global_scale", +) # Gemma-style (1+weight) RMSNorm weights. Excludes GDN gated norm (linear_attn.norm), # which is a standard weight*x norm. _GEMMA_NORM_SUFFIXES = ( @@ -122,9 +126,12 @@ def _load_maybe_quantized(f, raw_name: str, keyset: set[str]) -> torch.Tensor: if not raw_name.endswith(".weight"): return tensor base = raw_name[: -len(".weight")] - if base + ".weight_scale_2" in keyset: # NVFP4 (two-level block scale) + # NVFP4 (two-level block scale): modelopt names the per-row global ``weight_scale_2``; + # llm-compressor names it ``weight_global_scale``. Either presence -> dequantize as NVFP4. + s2_key = next((k for k in (".weight_scale_2", ".weight_global_scale") if base + k in keyset), None) + if s2_key is not None: return _dequant_nvfp4_weight( - tensor, f.get_tensor(base + ".weight_scale"), f.get_tensor(base + ".weight_scale_2") + tensor, f.get_tensor(base + ".weight_scale"), f.get_tensor(base + s2_key) ) if base + ".weight_scale" in keyset: # FP8 (per-tensor scale) return _dequant_fp8_weight(tensor, f.get_tensor(base + ".weight_scale")) @@ -179,9 +186,11 @@ def iter_weights( ) -> Iterator[tuple[str, torch.Tensor]]: hf_config = cached_load_hf_config(model_path) config = parse_config(hf_config) - if _compressed_tensors_nvfp4(hf_config): + if _compressed_tensors_nvfp4(hf_config) and not _has_moe_experts(hf_config): # Dense compressed-tensors NVFP4 (e.g. Qwen3.6-27B): attn (q/k/v/o, GDN out_proj) + - # dense MLP are W4A16 NVFP4; GDN in_proj_*, lm_head, norms bf16. + # dense MLP are W4A16 NVFP4; GDN in_proj_*, lm_head, norms bf16. MoE compressed-tensors + # checkpoints (Ornith: FP8 attn + FP8 shared_expert + NVFP4 experts) fall through to + # the modelopt branch below -- it has the FP8 attn path and the offload expert bank. yield from _iter_weights_compressed_tensors( model_path, device, include_non_moe=include_non_moe, include_moe_experts=include_moe_experts, @@ -251,8 +260,12 @@ def iter_weights( # NVFP4 dense projections kept native (W4A16) where the model expects them # (shared_expert); everything else dequantizes to bf16 below as before. - if (dense_nvfp4 or lmhead_nvfp4) and name.endswith(".weight") \ - and raw_name[: -len(".weight")] + ".weight_scale_2" in keyset: + # Detect by either modelopt's ``weight_scale_2`` or llm-compressor's + # ``weight_global_scale`` -- both name the per-row global scale. + if (dense_nvfp4 or lmhead_nvfp4) and name.endswith(".weight") and ( + raw_name[: -len(".weight")] + ".weight_scale_2" in keyset + or raw_name[: -len(".weight")] + ".weight_global_scale" in keyset + ): emit = _dense_nvfp4_emit( f, name[: -len(".weight")], raw_name[: -len(".weight")], shared_nvfp4=dense_nvfp4, lmhead_nvfp4=lmhead_nvfp4, @@ -304,6 +317,11 @@ def iter_weights( ".linear_attn.in_proj_qkvz": ( ".linear_attn.in_proj_qkv", ".linear_attn.in_proj_z", ), + # shared_expert gate|up merge -> shared_expert.gate_up_proj (per-tensor FP8 W8A16). + # Used by llm-compressor Ornith (FP8 attn/shared_expert + NVFP4 experts). + ".mlp.shared_expert.gate_up_proj": ( + ".mlp.shared_expert.gate_proj", ".mlp.shared_expert.up_proj", + ), } # bf16 (unquantized) GDN b|a projections fused -> in_proj_ba (matches the fp8 split). _PT_BF16_FUSE: dict[str, tuple[str, ...]] = { @@ -368,10 +386,13 @@ def _pt_fp8_fuse(base: str, weight: torch.Tensor, scalar: torch.Tensor, def _nvfp4_parts(f, raw_base: str): """Load a native NVFP4 weight as ``(packed uint8 [O, IN//2], block scale fp8 [O, IN//16], - per-output-row global fp16 [O])`` -- the dense W4A16 kernels' expected buffers.""" + per-output-row global fp16 [O])`` -- the dense W4A16 kernels' expected buffers. + The per-row global is ``weight_scale_2`` in modelopt and ``weight_global_scale`` in + llm-compressor; either is accepted.""" w = f.get_tensor(raw_base + ".weight") # uint8 packed FP4 (2 codes/byte) s = f.get_tensor(raw_base + ".weight_scale") # fp8-e4m3 per-16 block scale - g2 = f.get_tensor(raw_base + ".weight_scale_2") # per-tensor global scalar + g2_key = ".weight_scale_2" if raw_base + ".weight_scale_2" in f.keys() else ".weight_global_scale" + g2 = f.get_tensor(raw_base + g2_key) # per-tensor global scalar g = g2.reshape(1).to(torch.float16).expand(w.shape[0]).contiguous() return w, s, g @@ -544,6 +565,13 @@ def _iter_weights_attn_fp8( _CT_NVFP4_FUSE: dict[str, tuple[str, ...]] = { ".self_attn.qkv_proj": (".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj"), ".mlp.gate_up_proj": (".mlp.gate_proj", ".mlp.up_proj"), + # shared_expert gate|up merge -> shared_expert.gate_up_proj (W4A16 NVFP4). + # llm-compressor single-group checkpoints (kj_v2) put the shared expert in the same + # NVFP4 group as the routed experts; modelopt mixed_precision stores shared_expert as + # packed FP4 too. Fused cat on the output dim keeps each part's global scale exact. + ".mlp.shared_expert.gate_up_proj": ( + ".mlp.shared_expert.gate_proj", ".mlp.shared_expert.up_proj", + ), } _CT_BF16_FUSE: dict[str, tuple[str, ...]] = { ".linear_attn.in_proj": ( @@ -1060,6 +1088,31 @@ def _load(sink) -> None: return ExpertBanks("bf16", banks, streamed=layer_sink is not None) +def _spec_for(config) -> Nvfp4ExpertSourceSpec: + """Per-checkpoint NVFP4 source spec: same key pattern + role map, with + ``global_reciprocal`` set from the parsed config (``ModelConfig.nvfp4_global_reciprocal``, + populated by ``_nvfp4_global_reciprocal`` in models/config.py). The on-disk global scale + is the QUANT-side scale for llm-compressor NVFP4 (the dequant kernel multiplies by the + DEQUANT-side value, so the loader must reciprocate 1/x from disk) and the DEQUANT-side + divisor for modelopt (kernel multiplies, no transform). + + ``kind_map`` aliases the llm-compressor per-expert naming onto the modelopt canonical + kinds the bank builder dispatches on: ``weight_packed`` -> ``weight`` (uint8 packed FP4 + bytes), ``weight_global_scale`` -> ``weight_scale_2`` (per-row global), and the + activation scales map to themselves (the bank builder skips them as activation-only). + """ + from dataclasses import replace + return replace( + _NVFP4_SOURCE_SPEC, + global_reciprocal=bool(getattr(config, "nvfp4_global_reciprocal", False)), + kind_map={ + "weight_packed": "weight", + "weight_global_scale": "weight_scale_2", + "input_global_scale": "input_scale", # alias; bank builder still skips input scales + }, + ) + + def load_nvfp4_expert_sources( model_path: str, config, *, layer_sink=None ) -> dict[str, torch.Tensor]: @@ -1068,7 +1121,7 @@ def load_nvfp4_expert_sources( return load_nvfp4_expert_source_banks( model_path, config, - _NVFP4_SOURCE_SPEC, + _spec_for(config), drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, @@ -1084,7 +1137,7 @@ def load_nvfp4_expert_sources_parallel( return load_nvfp4_expert_source_banks_parallel( model_path, config, - _NVFP4_SOURCE_SPEC, + _spec_for(config), drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), workers=workers,