diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 0e3ab6a51..5dccc7cec 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -320,8 +320,272 @@ def _load(sink) -> int: } + +# -------------------------------------------------------------------------------------- +# compressed-tensors NVFP4 experts (llm-compressor layout): weight_packed (uint8) + +# weight_scale (fp8-e4m3 block) + scalar weight_global_scale. The dequant global is the +# reciprocal of the stored quant-side global (vLLM inverts it identically), broadcast +# per output row. Same native source banks as the modelopt path; single-file (no index) +# checkpoints are supported. [mixed-ct block-fp8 dense + nvfp4 experts] +# -------------------------------------------------------------------------------------- + + +def _checkpoint_weight_map(folder: str) -> dict[str, str]: + """name -> shard basename, from model.safetensors.index.json when present, else from + the single-file safetensors header(s).""" + index = os.path.join(folder, "model.safetensors.index.json") + if os.path.exists(index): + with open(index, encoding="utf-8") as f: + return json.load(f)["weight_map"] + import glob + import struct + + weight_map: dict[str, str] = {} + for shard in sorted(os.path.basename(p) for p in glob.glob(os.path.join(folder, "*.safetensors"))): + with open(os.path.join(folder, shard), "rb") as fh: + n = struct.unpack(" fill rows + if role == "gate": + gate_up_scale[bank_layer_id][expert, :I] = tensor + gate_up_global[bank_layer_id][expert, :I].fill_(global_val) + elif role == "up": + gate_up_scale[bank_layer_id][expert, I:] = tensor + gate_up_global[bank_layer_id][expert, I:].fill_(global_val) + else: + down_scale[bank_layer_id][expert] = tensor + down_global[bank_layer_id][expert].fill_(global_val) + + +def _ct_sort_shards(weight_map, spec): + """Bucket matched expert tensors: weight_global_scale -> globals pass; weight_packed / + weight_scale -> bulk pass. Returns (weight_shards, global_names_by_shard).""" + weight_shards: dict[str, list[tuple[str, re.Match[str], int]]] = collections.defaultdict(list) + global_shards: dict[str, list[str]] = collections.defaultdict(list) + for name, shard in weight_map.items(): + match = spec.key_pattern.match(name) + if match is None: + continue + # Bounds checking against num_layers happens in the parallel builder (which has the + # config); the serial build trusts the contiguous 0..L-1 Qwen3.5 MoE layer ids. + bank_layer = spec.layer_to_bank(int(match.group("layer")), None) + if bank_layer is None: + continue + kind = match.group("kind") + if kind == "weight_global_scale": + global_shards[shard].append(name) + elif kind in ("weight_packed", "weight_scale"): + weight_shards[shard].append((name, match, bank_layer)) + else: + raise ValueError(f"{spec.desc}: unknown CT expert tensor kind {kind!r}") + return weight_shards, global_shards + + +def _ct_load_globals(folder, global_shards, spec, drop_page_cache): + """Per-(layer, expert, proj) dequant global = 1/weight_global_scale, as python floats.""" + globals_map: dict[tuple, float] = {} + for shard in sorted(global_shards): + path = os.path.join(folder, shard) + drop_page_cache(path) + with safetensors.safe_open(path, framework="pt", device="cpu") as f: + for name in global_shards[shard]: + m = spec.key_pattern.match(name) + wg = f.get_tensor(name).reshape(-1).to(torch.float32) + assert wg.numel() == 1, ( + f"{spec.desc}: expected scalar weight_global_scale, got {tuple(wg.shape)}" + ) + globals_map[(int(m.group("layer")), int(m.group("expert")), m.group("proj"))] = ( + (1.0 / wg[0]).to(torch.float16).item() + ) + drop_page_cache(path) + return globals_map + + +def load_nvfp4_ct_expert_source_banks( + model_path: str, + config, + spec: Nvfp4ExpertSourceSpec, + *, + drop_page_cache: DropPageCache, + primary: bool, + layer_sink=None, +) -> dict[str, list[torch.Tensor]]: + """Serial per-shard build of the native NVFP4 source banks from a compressed-tensors + (llm-compressor) NVFP4 expert checkpoint.""" + folder = download_hf_weight(model_path) + weight_map = _checkpoint_weight_map(folder) + + E = config.num_experts + H = config.hidden_size + I = config.moe_intermediate_size + num_layers = _num_moe_layers(config) + + for shard in sorted(set(weight_map.values())): + drop_page_cache(os.path.join(folder, shard)) + + weight_shards, global_shards = _ct_sort_shards(weight_map, spec) + globals_map = _ct_load_globals(folder, global_shards, spec, drop_page_cache) + + _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) + banks = _ct_banks_tuple(_hb) + + from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline + + def _load(sink) -> int: + tracker = LayerCompletionTracker(E * 6, _hb, sink) + placed = 0 + for shard in tqdm(sorted(weight_shards), desc=f"Loading {spec.desc}", disable=not primary): + path = os.path.join(folder, shard) + with safetensors.safe_open(path, framework="pt", device="cpu") as f: + for name, match, bank_layer_id in weight_shards[shard]: + layer = int(match.group("layer")) + expert = int(match.group("expert")) + proj = match.group("proj") + role = spec.proj_to_role[proj] + kind = match.group("kind") + g = globals_map[(layer, expert, proj)] + _ct_place(f.get_tensor(name), kind, role, bank_layer_id, expert, I, banks, g) + tracker.note(bank_layer_id) + placed += 1 + drop_page_cache(path) + return placed + + if layer_sink is not None: + placed = _load(layer_sink) + else: + with PinPipeline() as pins: + placed = _load(pins) + + expected = num_layers * E * 6 + assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" + gate_up_packed, gate_up_scale, gate_up_global, down_packed, down_scale, down_global = banks + return { + "gate_up_packed": gate_up_packed, + "gate_up_scale": gate_up_scale, + "gate_up_global": gate_up_global, + "down_packed": down_packed, + "down_scale": down_scale, + "down_global": down_global, + } + + +def load_nvfp4_ct_expert_source_banks_parallel( + model_path: str, + config, + spec: Nvfp4ExpertSourceSpec, + *, + drop_page_cache: DropPageCache, + primary: bool, + workers: int = 8, + chunk: int = 8 << 20, + layer_sink=None, +) -> dict[str, list[torch.Tensor]]: + """Parallel counterpart of load_nvfp4_ct_expert_source_banks (byte-for-byte same + placement): bulk weight_packed/weight_scale via the chunked O_DIRECT reader; tiny + globals stay serial.""" + from freetoken.models.weight import iter_expert_tensors_parallel + + folder = download_hf_weight(model_path) + weight_map = _checkpoint_weight_map(folder) + + E = config.num_experts + H = config.hidden_size + I = config.moe_intermediate_size + num_layers = _num_moe_layers(config) + + weight_info: dict[str, tuple[re.Match[str], int]] = {} + global_shards: dict[str, list[str]] = collections.defaultdict(list) + for name, shard in weight_map.items(): + match = spec.key_pattern.match(name) + if match is None: + continue + bank_layer = spec.layer_to_bank(int(match.group("layer")), config) + if bank_layer is None: + continue + _bank_layer(spec, int(match.group("layer")), config) # bounds check + kind = match.group("kind") + if kind == "weight_global_scale": + global_shards[shard].append(name) + elif kind in ("weight_packed", "weight_scale"): + weight_info[name] = (match, bank_layer) + else: + raise ValueError(f"{spec.desc}: unknown CT expert tensor kind {kind!r}") + + globals_map = _ct_load_globals(folder, global_shards, spec, drop_page_cache) + + _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) + banks = _ct_banks_tuple(_hb) + + from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline + + def _load(sink) -> int: + tracker = LayerCompletionTracker(E * 6, _hb, sink) + placed = 0 + for name, tensor in iter_expert_tensors_parallel( + folder, lambda n: n in weight_info, workers=workers, chunk=chunk + ): + match, bank_layer_id = weight_info[name] + layer = int(match.group("layer")) + expert = int(match.group("expert")) + proj = match.group("proj") + role = spec.proj_to_role[proj] + kind = match.group("kind") + g = globals_map[(layer, expert, proj)] + _ct_place(tensor, kind, role, bank_layer_id, expert, I, banks, g) + tracker.note(bank_layer_id) + placed += 1 + return placed + + if layer_sink is not None: + placed = _load(layer_sink) + else: + with PinPipeline() as pins: + placed = _load(pins) + + expected = num_layers * E * 6 + assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" + gate_up_packed, gate_up_scale, gate_up_global, down_packed, down_scale, down_global = banks + return { + "gate_up_packed": gate_up_packed, + "gate_up_scale": gate_up_scale, + "gate_up_global": gate_up_global, + "down_packed": down_packed, + "down_scale": down_scale, + "down_global": down_global, + } + + __all__ = [ "Nvfp4ExpertSourceSpec", "load_nvfp4_expert_source_banks", "load_nvfp4_expert_source_banks_parallel", + "load_nvfp4_ct_expert_source_banks", + "load_nvfp4_ct_expert_source_banks_parallel", ] diff --git a/python/freetoken/models/qwen3_5_moe/config.py b/python/freetoken/models/qwen3_5_moe/config.py index 2ac4b607f..6da671303 100644 --- a/python/freetoken/models/qwen3_5_moe/config.py +++ b/python/freetoken/models/qwen3_5_moe/config.py @@ -68,6 +68,30 @@ def _expert_quant(hf_config: Any) -> str: _compressed_tensors_nvfp4 = detect_compressed_tensors_nvfp4 +def _compressed_tensors_mixed_fp8block(hf_config: Any) -> bool: + """compressed-tensors mixed-precision checkpoint: routed experts are NVFP4 (4-bit, + group_size 16, strategy tensor_group) while the dense projections (self_attn + q/k/v/o, linear_attn in_proj_*/out_proj, mlp.shared_expert gate/up/down) are 128x128 + block FP8 (8-bit, strategy block, W8A8). Distinct from pure compressed-tensors NVFP4 + (detect_compressed_tensors_nvfp4), where the dense side is packed FP4 as well.""" + get = _quant_accessor(hf_config) + if get is None: + return False + if str(get("quant_method") or "").lower() != "compressed-tensors": + return False + saw_fp8_block = saw_nvfp4 = False + for gspec in (get("config_groups") or {}).values(): + w = (gspec or {}).get("weights") or {} + bits = int(w.get("num_bits", 0) or 0) + strategy = str(w.get("strategy", "")).lower() + group_size = int(w.get("group_size", 0) or 0) + if bits == 8 and strategy == "block": + saw_fp8_block = True + if bits == 4 and group_size == 16 and strategy == "tensor_group": + saw_nvfp4 = True + return saw_fp8_block and saw_nvfp4 + + def _lm_head_quant(hf_config: Any) -> str: """Whether the checkpoint stores ``lm_head`` as NVFP4. modelopt MIXED_PRECISION lists it in the per-layer ``quantized_layers`` map (``W4A16_NVFP4``); pure-NVFP4 checkpoints have no @@ -164,9 +188,14 @@ def parse_config(hf_config: Any) -> ModelConfig: else {k: v for k, v in rope_params.items() if not isinstance(v, (list, dict))} ) + mixed_ct = _compressed_tensors_mixed_fp8block(hf_config) expert_quant, weight_block_size = _fp8_block_quant(hf_config) if expert_quant == "none": expert_quant = _expert_quant(hf_config) # nvfp4 / mixed-precision modelopt + if mixed_ct: + # Routed experts are packed NVFP4 (offload banks); dense side is 128x128 block FP8. + expert_quant = "nvfp4" + weight_block_size = (128, 128) # 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) @@ -182,10 +211,16 @@ def parse_config(hf_config: Any) -> ModelConfig: # 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): + if _compressed_tensors_nvfp4(hf_config) and not mixed_ct: attn_quant = "nvfp4" dense_quant = "nvfp4" lm_head_quant = "none" + if mixed_ct: + # Dense attn/GDN projections and the shared expert stay native 128x128 block FP8 + # (W8A8); lm_head / embeddings / norms are bf16. + attn_quant = "fp8_block" + dense_quant = "fp8_block" + lm_head_quant = "none" # Dense variants (e.g. Qwen3.6-27B) report num_experts==0: route the decoder MLP through # the dense Qwen3_5DenseMLP instead of the MoE block. diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index 2e7320051..7e759e3b2 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -74,7 +74,7 @@ def __init__( # qkv|z carry a weight scale (block-fp8 weight_scale_inv, or per-tensor FP8 # weight_scale); b|a stay bf16. Both quant modes therefore split the four-way # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM). - self._block_fp8 = expert_quant == "fp8_block" + self._block_fp8 = expert_quant == "fp8_block" or attn_quant == "fp8_block" self._pertensor_fp8 = attn_quant == "fp8_pertensor" self._fp8 = self._block_fp8 or self._pertensor_fp8 diff --git a/python/freetoken/models/qwen3_5_moe/moe.py b/python/freetoken/models/qwen3_5_moe/moe.py index b5ab1cf3d..db5bac411 100644 --- a/python/freetoken/models/qwen3_5_moe/moe.py +++ b/python/freetoken/models/qwen3_5_moe/moe.py @@ -22,7 +22,9 @@ class _SharedExpert(BaseOP): """Always-present shared SwiGLU expert of width ``shared_expert_intermediate_size``.""" def __init__(self, config: ModelConfig, hidden_size: int, intermediate_size: int): - if getattr(config, "expert_quant", "none") == "fp8_block": + if getattr(config, "expert_quant", "none") == "fp8_block" or getattr( + config, "dense_quant", "none" + ) == "fp8_block": self.gate_up_proj = Fp8BlockColMerged( hidden_size, [intermediate_size, intermediate_size], has_bias=False ) diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index d07f18cd7..0d2339da6 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -20,12 +20,18 @@ ) from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, + load_nvfp4_ct_expert_source_banks, + load_nvfp4_ct_expert_source_banks_parallel, load_nvfp4_expert_source_banks, ) 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_mixed_fp8block, + _compressed_tensors_nvfp4, + parse_config, +) # Expert weights are stored pre-fused per layer: experts.gate_up_proj / experts.down_proj. _PACKED_EXPERT_PATTERN = re.compile( @@ -38,7 +44,7 @@ # head's ``mtp.layers.N.mlp.experts.*`` tensors (served text-only, dropped). _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"^model\.(?:language_model\.)?layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." r"(?Pgate_proj|up_proj|down_proj)\.(?Pweight|weight_scale|weight_scale_2)$" ) _NVFP4_SOURCE_SPEC = Nvfp4ExpertSourceSpec( @@ -47,6 +53,20 @@ layer_to_bank=lambda layer, config: layer, # every layer is MoE desc="Qwen3.5 NVFP4 experts", ) +# compressed-tensors NVFP4 experts (llm-compressor): weight_packed (uint8) + weight_scale +# (fp8-e4m3 block) + scalar weight_global_scale (quant-side; the dequant global is its +# reciprocal, broadcast per output row). [mixed-ct block-fp8 dense + nvfp4 experts] +_NVFP4_CT_EXPERT_KEY_RE = re.compile( + r"^model\.(?:language_model\.)?layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." + r"(?Pgate_proj|up_proj|down_proj)\." + r"(?Pweight_packed|weight_scale|weight_global_scale)$" +) +_NVFP4_CT_SOURCE_SPEC = Nvfp4ExpertSourceSpec( + key_pattern=_NVFP4_CT_EXPERT_KEY_RE, + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, # every layer is MoE + desc="Qwen3.5 compressed-tensors 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") @@ -179,6 +199,16 @@ def iter_weights( ) -> Iterator[tuple[str, torch.Tensor]]: hf_config = cached_load_hf_config(model_path) config = parse_config(hf_config) + if _compressed_tensors_mixed_fp8block(hf_config): + # Mixed compressed-tensors: dense (attn/GDN/shared_expert) weights are 128x128 + # block FP8; routed experts are packed NVFP4 and excluded from this pass (they + # are built by the NVFP4 offload-bank provider). [mixed-ct block-fp8 dense + nvfp4 experts] + yield from _iter_weights_fp8( + model_path, device, + include_non_moe=include_non_moe, include_moe_experts=False, + ct_scale_suffix=True, + ) + return if _compressed_tensors_nvfp4(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. @@ -718,7 +748,7 @@ def _is_expert(raw_name: str) -> bool: # Routed-expert checkpoint key (per-expert, un-fused). ``mtp.layers...`` is excluded by the # ``model.language_model.`` anchor, so the parallel reader only sees the real experts. _FP8_EXPERT_RE = re.compile( - r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." + r"^model\.(?:language_model\.)?layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." r"(?Pgate|up|down)_proj\.(?Pweight|weight_scale_inv)$" ) @@ -751,7 +781,8 @@ def _fp8_fuse(base: str, suf: str, tensor: torch.Tensor, buf: dict) -> tuple[str def _iter_weights_fp8( - model_path: str, device: torch.device, *, include_non_moe: bool, include_moe_experts: bool = False + model_path: str, device: torch.device, *, include_non_moe: bool, include_moe_experts: bool = False, + ct_scale_suffix: bool = False, ) -> Iterator[tuple[str, torch.Tensor]]: """Yield the block-fp8 weights, renamed + fused to the model buffers. @@ -774,6 +805,14 @@ def _iter_weights_fp8( name = _rename(raw_name) if name is None or ".mlp.experts." in name: continue # routed experts handled below / by the offload cache + if ct_scale_suffix: + # W8A8 activations are quantized dynamically: no stored input scales. + if name.endswith((".input_scale", ".input_global_scale")): + continue + # compressed-tensors names block scales .weight_scale (bf16); the + # block-fp8 kernels/fusion expect .weight_scale_inv. + if name.endswith(".weight_scale"): + name = name[: -len(".weight_scale")] + ".weight_scale_inv" tensor = f.get_tensor(raw_name) base, suf = _split_kind(name) fused = _fp8_fuse(base, suf, tensor, fuse_buf) @@ -1065,6 +1104,16 @@ def load_nvfp4_expert_sources( ) -> dict[str, torch.Tensor]: """Build the CPU NVFP4 expert source banks for the offload cache (gate/up fused on the output-row axis, down separate; weight_scale_2 carried as the per-row global scale).""" + if _compressed_tensors_mixed_fp8block(cached_load_hf_config(model_path)): + # compressed-tensors NVFP4 layout (weight_packed/weight_scale/weight_global_scale). + return load_nvfp4_ct_expert_source_banks( + model_path, + config, + _NVFP4_CT_SOURCE_SPEC, + drop_page_cache=drop_page_cache, + primary=get_tp_info().is_primary(), + layer_sink=layer_sink, + ) return load_nvfp4_expert_source_banks( model_path, config, @@ -1081,6 +1130,17 @@ def load_nvfp4_expert_sources_parallel( """parallel: same NVFP4 source banks via the common chunked multi-threaded reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel + if _compressed_tensors_mixed_fp8block(cached_load_hf_config(model_path)): + return load_nvfp4_ct_expert_source_banks_parallel( + model_path, + config, + _NVFP4_CT_SOURCE_SPEC, + drop_page_cache=drop_page_cache, + primary=get_tp_info().is_primary(), + workers=workers, + chunk=chunk, + layer_sink=layer_sink, + ) return load_nvfp4_expert_source_banks_parallel( model_path, config, diff --git a/tests/models/test_qwen3_5_moe_ct_mixed_detect.py b/tests/models/test_qwen3_5_moe_ct_mixed_detect.py new file mode 100644 index 000000000..23086baea --- /dev/null +++ b/tests/models/test_qwen3_5_moe_ct_mixed_detect.py @@ -0,0 +1,83 @@ +"""Detection of compressed-tensors mixed-precision checkpoints (NVFP4 experts + +128x128 block-FP8 dense side), e.g. kyaky/Qwen3.6-35B-A3B-Uncensored-NVFP4. + +These are distinct from pure compressed-tensors NVFP4 checkpoints (where the +dense side is packed FP4 as well): they use llm-compressor config_groups with an +8-bit block strategy for attn/shared-expert and a 4-bit group_size=16 +tensor_group strategy for the routed experts.""" + +from __future__ import annotations + +from types import SimpleNamespace + +from freetoken.models.qwen3_5_moe.config import _compressed_tensors_mixed_fp8block + + +def _hf(qc): + return SimpleNamespace(quantization_config=qc) + + +_CT_MIXED = { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "config_groups": { + "group_0": { + "weights": {"num_bits": 8, "strategy": "block", "block_structure": [128, 128]}, + "targets": ["re:.*shared_expert.*", "re:.*self_attn.*"], + }, + "group_1": { + "weights": {"num_bits": 4, "strategy": "tensor_group", "group_size": 16}, + "targets": ["re:.*mlp.experts.*"], + }, + }, +} + +_CT_PURE_NVFP4 = { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "config_groups": { + "group_0": { + "weights": {"num_bits": 4, "strategy": "tensor_group", "group_size": 16}, + "targets": ["re:.*"], + }, + }, +} + +_CT_FP8_ONLY = { + "quant_method": "compressed-tensors", + "config_groups": { + "group_0": { + "weights": {"num_bits": 8, "strategy": "group", "group_size": 128}, + "targets": ["re:.*"], + }, + }, +} + +_MODELOPT_MIXED = { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "model.layers.0.mlp.experts": {"quant_algo": "NVFP4", "group_size": 16}, + "model.layers.0.mlp.shared_expert.gate_proj": {"quant_algo": "FP8"}, + }, +} + + +def test_ct_mixed_fp8block_detected(): + assert _compressed_tensors_mixed_fp8block(_hf(_CT_MIXED)) is True + + +def test_pure_ct_nvfp4_not_flagged(): + assert _compressed_tensors_mixed_fp8block(_hf(_CT_PURE_NVFP4)) is False + + +def test_ct_fp8_only_not_flagged(): + assert _compressed_tensors_mixed_fp8block(_hf(_CT_FP8_ONLY)) is False + + +def test_modelopt_mixed_not_matched(): + # modelopt MIXED_PRECISION (e.g. Apodex) has no config_groups -> handled elsewhere + assert _compressed_tensors_mixed_fp8block(_hf(_MODELOPT_MIXED)) is False + + +def test_no_quant_config(): + assert _compressed_tensors_mixed_fp8block(SimpleNamespace()) is False