diff --git a/src/mobius/_configs/_quantization.py b/src/mobius/_configs/_quantization.py index 69432ebf5..f4d8a307e 100644 --- a/src/mobius/_configs/_quantization.py +++ b/src/mobius/_configs/_quantization.py @@ -74,11 +74,43 @@ def from_transformers(cls, hf_config) -> QuantizationConfig | None: "onnxruntime_USE_FP4_QMOE=ON). Export the unquantized (bf16) " "checkpoint instead, or quantize the bf16 export via Olive." ) + # Block-scaled fp8 (E4M3 weight + 2D UE8M0 block scale) and packed-fp4 + # routed experts (I8-packed E2M1 nibbles + UE8M0 micro-scale) are a + # mixed-precision layout this INT4/per-tensor path cannot load — the + # packed [out, in/2] fp4 expert vs its logical [out, in] initializer + # produces a confusing "Weight shape mismatch". Detect it by property + # (not model name, not the ``quant_method`` string) and fail closed with + # a typed, actionable blocker naming the real layout + the runtime ABI + # gap. Checked before the ``quant_method == "none"`` early-return because + # a checkpoint can advertise fp4 experts via top-level ``expert_dtype`` + # while leaving ``quant_method`` unset. + from mobius.integrations._block_quant import BlockQuantScheme + + scheme = BlockQuantScheme.from_quantization_config( + qc, expert_dtype=getattr(hf_config, "expert_dtype", None) + ) + if scheme is not None: + from mobius.integrations._block_quant import BlockQuantExportError + + raise BlockQuantExportError( + "Block-scaled FP8 / packed-FP4 checkpoint is not loadable by the " + "INT4/per-tensor quantization path. Detected " + f"quant_method={scheme.quant_method!r}, " + f"weight_block_size={list(scheme.weight_block_size) or None}, " + f"expert_dtype={scheme.expert_dtype!r}: block-FP8 projections " + "(E4M3 weight + 2D UE8M0 block scale) and/or FP4-packed routed " + "experts (I8-packed E2M1 nibbles + UE8M0 micro-scale). Parse and " + "validate these by property with mobius.integrations._block_quant " + "(BlockQuantScheme / classify_tensor / QuantizedTensorDescriptor); " + "the routed-expert emission gate (plan_routed_expert_bank) reports " + "the exact onnx-genai nxrt ABI gap. Native export is blocked until " + "the runtime gains a block-FP8 / planar-FP4 BlockFormat." + ) if method == "none": return None - # FP8 per-tensor quantization (float8_e4m3fn + scalar scale) - # is handled by dtype casting in _assign_weight(), not by - # QuantizedLinear block quantization. + # Per-tensor fp8 (float8_e4m3fn + a scalar scale) is handled by dtype + # casting in _assign_weight(), so it returns None here. (Block-scaled + # fp8 was already routed to the typed blocker above.) if method == "fp8": return None return cls( diff --git a/src/mobius/integrations/_block_quant.py b/src/mobius/integrations/_block_quant.py new file mode 100644 index 000000000..db9bbd813 --- /dev/null +++ b/src/mobius/integrations/_block_quant.py @@ -0,0 +1,824 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Block-scaled FP8 / packed-FP4 quantized weight loading contract. + +Some recent checkpoints (e.g. DeepSeek-V4) publish *mixed-precision* weights +whose real numeric layout is not captured by the coarse +:class:`~mobius._configs.QuantizationConfig` (which only distinguishes int4 +GPTQ/AWQ-style block quant, GGUF, and per-tensor fp8). Three distinct weight +families coexist in one checkpoint and must be told apart **by tensor +properties, never by model name**: + +1. **Ordinary** ``bf16`` / ``f16`` / ``f32`` tensors (router gate, norms, + biases, sinks). No paired scale. +2. **Block-FP8 projections** — an ``F8_E4M3`` weight of *logical* shape paired + with an ``F8_E8M0`` (UE8M0, exponent-only) 2D block-scale of shape + ``[ceil(out / bs0), ceil(in / bs1)]`` (``weight_block_size`` from the HF + ``quantization_config``, e.g. ``[128, 128]``). +3. **FP4-packed routed experts** — an ``I8`` weight that stores **two E2M1 + (fp4) nibbles per byte**, so its *packed* shape is the logical shape with the + last dim halved, paired with an ``F8_E8M0`` micro-scale of shape + ``[out, logical_in / 32]`` (one UE8M0 exponent per output row per 32 logical + input elements). Numerically this is **MXFP4** (E2M1 + block-32 + E8M0), not + NVFP4 (block-16 + E4M3 block-scale + FP32 global scale). + +This module is the *clean, breaking* descriptor + load contract for those +families. It intentionally does **not** dequantize to float, does not copy +weights, and preserves raw bytes. It classifies and validates tensors by +property, loads their raw bytes lazily (bounded to one tensor at a time), and +exposes a byte-exact expert-major bank-stacking primitive. + +Crucially, the routed-expert *emission gate* +(:func:`plan_routed_expert_bank`) proves whether the onnx-genai ``nxrt`` +runtime can represent these banks and, when it cannot, **fails closed with a +typed :class:`BlockQuantExportError` naming the exact ABI gap** rather than +emitting an unrunnable node. As of this writing the ``nxrt`` block-quant ABI +(``crates/onnx-runtime-ep-cpu/src/kernels/block_quantized_{matmul,moe}.rs``) +accepts only the interleaved llama.cpp ``block_mxfp4`` layout and the ``iq*`` +GGUF formats — it has **no block-FP8 format** and **no planar-FP4 bank +layout** — so both quantized families here are typed-rejected. +""" + +from __future__ import annotations + +import dataclasses +import enum +import json +import pathlib +import struct + +__all__ = [ + "QuantKind", + "BlockQuantScheme", + "QuantizedTensorDescriptor", + "BlockQuantError", + "BlockQuantValidationError", + "BlockQuantExportError", + "PackedExpertBank", + "SAFETENSORS_DTYPE_BYTES", + "classify_tensor", + "validate_descriptor", + "pair_weight_scales", + "build_descriptors", + "read_safetensors_header", + "raw_tensor_span", + "read_raw_tensor_bytes", + "LazyRawTensor", + "stack_expert_bank", + "runtime_representation_gap", + "plan_routed_expert_bank", + "NXRT_BLOCK_FORMATS", +] + + +# --------------------------------------------------------------------------- +# Typed errors +# --------------------------------------------------------------------------- + + +class BlockQuantError(Exception): + """Base class for block-quant load/emit failures.""" + + +class BlockQuantValidationError(BlockQuantError, ValueError): + """A tensor's declared metadata is internally inconsistent. + + Raised on a logical-vs-packed shape contradiction, a missing/duplicate or + mismatched scale, or a scale grid that does not match the weight's block + geometry. This is a hard reject — never a silently-reinterpreted tensor. + """ + + +class BlockQuantExportError(BlockQuantError, NotImplementedError): + """A validated bank cannot be represented by the target runtime ABI. + + Raised by the emission gate instead of emitting an unrunnable node. The + message names the exact ABI gap (which format/layout the runtime lacks) so + the blocker is actionable rather than a confusing downstream shape error. + """ + + +# --------------------------------------------------------------------------- +# safetensors dtype table + tiny header reader (byte-exact, no data read) +# --------------------------------------------------------------------------- + +#: Bytes per stored element for each safetensors dtype string. ``I8`` counts a +#: single byte even though an FP4-packed tensor stores *two* logical E2M1 codes +#: per byte — that 2x is captured by the packed-vs-logical shape, not here. +SAFETENSORS_DTYPE_BYTES: dict[str, int] = { + "F64": 8, + "F32": 4, + "F16": 2, + "BF16": 2, + "I64": 8, + "I32": 4, + "I16": 2, + "I8": 1, + "U8": 1, + "BOOL": 1, + "F8_E4M3": 1, + "F8_E5M2": 1, + "F8_E8M0": 1, +} + +_FLOAT_DTYPES = frozenset({"F64", "F32", "F16", "BF16"}) + + +def _num_elements(shape: tuple[int, ...]) -> int: + n = 1 + for d in shape: + n *= int(d) + return n + + +def read_safetensors_header(path: str | pathlib.Path) -> dict: + """Read a safetensors file's JSON header only (no tensor data). + + Returns the raw header dict mapping each tensor name to + ``{"dtype", "shape", "data_offsets"}`` plus the ``__metadata__`` block. + Only the 8-byte length prefix and the header JSON are read from disk. + """ + with open(path, "rb") as f: + (header_len,) = struct.unpack(" tuple[str, tuple[int, ...], int, int]: + """Return ``(dtype, shape, abs_start, abs_end)`` for one tensor's raw bytes. + + ``abs_start``/``abs_end`` are absolute byte offsets into *path* (the + safetensors ``data_offsets`` are relative to the end of the header, so the + ``8 + header_len`` base is added). Only the header is read. + """ + with open(path, "rb") as f: + (header_len,) = struct.unpack(" bytes: + """Read exactly one tensor's raw on-disk bytes, byte-for-byte. + + No dtype interpretation, no cast, no copy-to-float. The returned ``bytes`` + are identical to the tensor's storage in the shard, so a subsequent write + preserves the quantized payload exactly. + """ + _dtype, _shape, start, end = raw_tensor_span(path, key) + with open(path, "rb") as f: + f.seek(start) + return f.read(end - start) + + +@dataclasses.dataclass(frozen=True) +class LazyRawTensor: + """A bounded, byte-preserving handle to one tensor in one shard. + + Holds only the shard path + key + header metadata; the payload is read on + demand via :meth:`read`, so peak resident memory is one tensor, not the + whole checkpoint. ``num_bytes`` is known from the header without reading + data. + """ + + path: str + key: str + dtype: str + shape: tuple[int, ...] + num_bytes: int + + @classmethod + def open(cls, path: str | pathlib.Path, key: str) -> LazyRawTensor: + dtype, shape, start, end = raw_tensor_span(path, key) + return cls(path=str(path), key=key, dtype=dtype, shape=shape, num_bytes=end - start) + + def read(self) -> bytes: + """Read and return this tensor's raw bytes (byte-exact).""" + data = read_raw_tensor_bytes(self.path, self.key) + if len(data) != self.num_bytes: + raise BlockQuantValidationError( + f"Tensor {self.key!r} in {self.path}: header declared " + f"{self.num_bytes} bytes but read {len(data)}" + ) + return data + + +# --------------------------------------------------------------------------- +# Quantization scheme parsed from the HF quantization_config (by properties) +# --------------------------------------------------------------------------- + + +class QuantKind(enum.Enum): + """Property-classified quantization family of a single tensor.""" + + ORDINARY = "ordinary" # bf16/f16/f32, no scale + BLOCK_FP8 = "block_fp8" # F8_E4M3 weight + 2D UE8M0 block scale + FP4_PACKED = "fp4_packed" # I8-packed E2M1 nibbles + 1D UE8M0 micro-scale + UNSUPPORTED = "unsupported" # recognized-but-unhandled / malformed + + +#: Micro-scale block length (logical input elements per UE8M0 exponent) that +#: marks an MXFP4-style fp4 tensor. NVFP4 would instead use 16 + an E4M3 scale. +MXFP4_MICROSCALE_BLOCK = 32 + + +@dataclasses.dataclass(frozen=True) +class BlockQuantScheme: + """The checkpoint-wide quantization scheme, parsed by properties. + + Derived from the HF ``quantization_config`` plus the top-level + ``expert_dtype`` — never from the model name. ``weight_block_size`` is the + block-FP8 projection block geometry; an empty tuple means per-tensor fp8 + (which this contract does not own — that stays ``QuantizationConfig`` None). + """ + + quant_method: str + weight_fmt: str | None = None # e.g. "e4m3" + scale_fmt: str | None = None # e.g. "ue8m0" + weight_block_size: tuple[int, ...] = () + activation_scheme: str | None = None + expert_dtype: str | None = None # e.g. "fp4" + + @property + def is_block_scaled_fp8(self) -> bool: + """True for block-scaled fp8 (a non-empty ``weight_block_size``).""" + return self.quant_method == "fp8" and len(self.weight_block_size) > 0 + + @property + def has_packed_fp4_experts(self) -> bool: + """True when routed experts are packed fp4 (``expert_dtype`` fp4-like).""" + return (self.expert_dtype or "").lower() in {"fp4", "nvfp4", "mxfp4"} + + @property + def is_owned(self) -> bool: + """True when this scheme is one this block-quant contract handles.""" + return self.is_block_scaled_fp8 or self.has_packed_fp4_experts + + @classmethod + def from_quantization_config( + cls, qc: dict | None, *, expert_dtype: str | None = None + ) -> BlockQuantScheme | None: + """Parse a ``quantization_config`` dict; ``None`` if not owned here. + + Returns ``None`` for absent configs, per-tensor fp8 (no + ``weight_block_size``), and non-fp8 methods without fp4 experts. + """ + if not isinstance(qc, dict): + if expert_dtype and str(expert_dtype).lower() in {"fp4", "nvfp4", "mxfp4"}: + return cls(quant_method="none", expert_dtype=str(expert_dtype).lower()) + return None + block = qc.get("weight_block_size") + block_t: tuple[int, ...] = ( + tuple(int(x) for x in block) if isinstance(block, (list, tuple)) else () + ) + scheme = cls( + quant_method=str(qc.get("quant_method", "none")), + weight_fmt=(str(qc["fmt"]).lower() if qc.get("fmt") is not None else None), + scale_fmt=( + str(qc["scale_fmt"]).lower() if qc.get("scale_fmt") is not None else None + ), + weight_block_size=block_t, + activation_scheme=( + str(qc["activation_scheme"]) + if qc.get("activation_scheme") is not None + else None + ), + expert_dtype=(str(expert_dtype).lower() if expert_dtype else None), + ) + return scheme if scheme.is_owned else None + + @classmethod + def from_hf_config(cls, hf_config) -> BlockQuantScheme | None: + """Parse from a HF config object (reads ``quantization_config``).""" + qc = getattr(hf_config, "quantization_config", None) + if qc is not None and hasattr(qc, "to_dict"): + qc = qc.to_dict() + return cls.from_quantization_config( + qc, expert_dtype=getattr(hf_config, "expert_dtype", None) + ) + + +# --------------------------------------------------------------------------- +# The breaking per-tensor descriptor +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class QuantizedTensorDescriptor: + """A complete, self-describing load contract for one quantized tensor. + + Carries the *logical* shape (what the math sees), the *packed* on-disk + shape (what the bytes are), the weight/scale dtypes and names, the block + geometry in logical-element units, the exact byte counts, and the tensor's + structural role. This is deliberately verbose and breaking: consumers must + read explicit fields, never re-derive layout from a name. + """ + + name: str + kind: QuantKind + weight_dtype: str + logical_shape: tuple[int, ...] + packed_shape: tuple[int, ...] + weight_num_bytes: int + is_routed_expert: bool + is_shared_expert: bool + block_shape: tuple[int, ...] | None = None + scale_name: str | None = None + scale_dtype: str | None = None + scale_shape: tuple[int, ...] | None = None + scale_layout: str | None = None + scale_num_bytes: int | None = None + microscale_kind: str | None = None # "mxfp4" | "nvfp4" | None + unsupported_reason: str | None = None + + @property + def pack_factor(self) -> int: + """Logical elements per stored element (2 for nibble-packed fp4).""" + stored = _num_elements(self.packed_shape) + return _num_elements(self.logical_shape) // stored if stored else 1 + + +def _expert_role(name: str) -> tuple[bool, bool]: + """Return ``(is_routed_expert, is_shared_expert)`` from the module path. + + This is a *structural* graph property (the standard HF MoE naming: + ``...experts....`` routed, ``...shared_experts...`` shared), not a + model-name allowlist — the numeric *kind* is classified separately from + dtype/scale properties. + """ + is_shared = "shared_expert" in name + is_routed = (".experts." in name) and not is_shared + return is_routed, is_shared + + +def classify_tensor( + name: str, + weight_dtype: str, + weight_shape: tuple[int, ...], + *, + scale_dtype: str | None = None, + scale_shape: tuple[int, ...] | None = None, + scale_name: str | None = None, + scheme: BlockQuantScheme | None = None, +) -> QuantizedTensorDescriptor: + """Classify one tensor into a :class:`QuantizedTensorDescriptor` by property. + + Uses only the weight dtype/shape and the paired scale dtype/shape (plus the + scheme's block size for block-FP8). Never inspects the model name for the + numeric family. Unrecognized combinations return ``kind=UNSUPPORTED`` with a + reason instead of guessing. + """ + weight_shape = tuple(int(d) for d in weight_shape) + scale_shape = tuple(int(d) for d in scale_shape) if scale_shape is not None else None + is_routed, is_shared = _expert_role(name) + w_elem_bytes = SAFETENSORS_DTYPE_BYTES.get(weight_dtype) + scale_elem_bytes = SAFETENSORS_DTYPE_BYTES.get(scale_dtype) if scale_dtype else None + weight_num_bytes = _num_elements(weight_shape) * (w_elem_bytes or 0) + scale_num_bytes = ( + _num_elements(scale_shape) * scale_elem_bytes + if scale_shape is not None and scale_elem_bytes is not None + else None + ) + + def _desc( + kind: QuantKind, + *, + logical: tuple[int, ...], + block: tuple[int, ...] | None = None, + layout: str | None = None, + microscale: str | None = None, + reason: str | None = None, + ) -> QuantizedTensorDescriptor: + return QuantizedTensorDescriptor( + name=name, + kind=kind, + weight_dtype=weight_dtype, + logical_shape=logical, + packed_shape=weight_shape, + weight_num_bytes=weight_num_bytes, + is_routed_expert=is_routed, + is_shared_expert=is_shared, + block_shape=block, + scale_name=scale_name, + scale_dtype=scale_dtype, + scale_shape=scale_shape, + scale_layout=layout, + scale_num_bytes=scale_num_bytes, + microscale_kind=microscale, + unsupported_reason=reason, + ) + + # Ordinary float tensors carry no scale. + if weight_dtype in _FLOAT_DTYPES and scale_dtype is None: + return _desc(QuantKind.ORDINARY, logical=weight_shape) + + # FP4-packed: I8 nibbles + UE8M0 micro-scale. + if weight_dtype == "I8" and scale_dtype == "F8_E8M0": + if len(weight_shape) < 1: + return _desc( + QuantKind.UNSUPPORTED, logical=weight_shape, reason="scalar I8 weight" + ) + logical = (*weight_shape[:-1], weight_shape[-1] * 2) # last dim halved on disk + # MXFP4 pins the micro-scale block to 32 logical input elements per row; + # this is a format property, not something inferred from a (possibly + # wrong) scale grid — validate_descriptor checks the scale against it. + block_len = MXFP4_MICROSCALE_BLOCK + return _desc( + QuantKind.FP4_PACKED, + logical=logical, + block=(1, block_len), + layout=f"microscale_1x{block_len}_ue8m0", + microscale="mxfp4", + ) + + # Block-FP8: E4M3 weight + 2D UE8M0 block scale. + if weight_dtype == "F8_E4M3" and scale_dtype == "F8_E8M0": + bs = scheme.weight_block_size if scheme and scheme.weight_block_size else None + if ( + bs is None + and scale_shape is not None + and len(weight_shape) == len(scale_shape) == 2 + ): + # Infer square block geometry from the ceil-divided scale grid. + bs0 = -(-weight_shape[0] // scale_shape[0]) if scale_shape[0] else 0 + bs1 = -(-weight_shape[1] // scale_shape[1]) if scale_shape[1] else 0 + bs = (bs0, bs1) + block = tuple(int(x) for x in bs) if bs else None + layout = ( + f"block{block[0]}x{block[1]}_ue8m0" + if block and len(block) == 2 + else "block2d_ue8m0" + ) + return _desc(QuantKind.BLOCK_FP8, logical=weight_shape, block=block, layout=layout) + + # Recognized-but-unhandled or malformed combinations fail closed. + if weight_dtype.startswith("F8") and scale_dtype is None: + return _desc( + QuantKind.UNSUPPORTED, + logical=weight_shape, + reason=f"{weight_dtype} weight without a paired block scale", + ) + if weight_dtype == "I8" and scale_dtype is None: + return _desc( + QuantKind.UNSUPPORTED, + logical=weight_shape, + reason="I8 weight without a paired UE8M0 micro-scale (cannot be fp4)", + ) + return _desc( + QuantKind.UNSUPPORTED, + logical=weight_shape, + reason=f"unrecognized weight/scale dtype pair ({weight_dtype}, {scale_dtype})", + ) + + +def validate_descriptor(desc: QuantizedTensorDescriptor) -> None: + """Validate a descriptor's internal shape/scale consistency; raise on error. + + Checks (per kind): logical-vs-packed shape relation, scale presence + dtype, + scale grid vs block geometry, and byte counts. Raises + :class:`BlockQuantValidationError` on any contradiction. + """ + if desc.kind is QuantKind.UNSUPPORTED: + raise BlockQuantValidationError( + f"{desc.name}: unsupported tensor ({desc.unsupported_reason})" + ) + + if desc.kind is QuantKind.ORDINARY: + if desc.scale_name is not None or desc.scale_shape is not None: + raise BlockQuantValidationError( + f"{desc.name}: ordinary float tensor must not carry a scale" + ) + if desc.logical_shape != desc.packed_shape: + raise BlockQuantValidationError( + f"{desc.name}: ordinary tensor logical {desc.logical_shape} != " + f"packed {desc.packed_shape}" + ) + return + + # Both quantized kinds require a paired scale. + if desc.scale_shape is None or desc.scale_dtype is None: + raise BlockQuantValidationError( + f"{desc.name}: {desc.kind.value} tensor has no paired scale" + ) + if desc.scale_dtype != "F8_E8M0": + raise BlockQuantValidationError( + f"{desc.name}: expected UE8M0 (F8_E8M0) scale, got {desc.scale_dtype}" + ) + + if desc.kind is QuantKind.FP4_PACKED: + if len(desc.packed_shape) != 2 or len(desc.logical_shape) != 2: + raise BlockQuantValidationError(f"{desc.name}: fp4 experts must be 2D") + out_l, in_l = desc.logical_shape + out_p, in_p = desc.packed_shape + if out_l != out_p: + raise BlockQuantValidationError( + f"{desc.name}: fp4 output dim mismatch logical {out_l} vs packed {out_p}" + ) + if in_l != in_p * 2: + raise BlockQuantValidationError( + f"{desc.name}: fp4 packed input {in_p} must be logical {in_l} / 2 " + f"(two E2M1 nibbles per int8 byte)" + ) + block_len = desc.block_shape[1] if desc.block_shape else MXFP4_MICROSCALE_BLOCK + if block_len <= 0 or in_l % block_len != 0: + raise BlockQuantValidationError( + f"{desc.name}: logical input {in_l} not divisible by micro-scale block {block_len}" + ) + expected_scale = (out_l, in_l // block_len) + if desc.scale_shape != expected_scale: + raise BlockQuantValidationError( + f"{desc.name}: fp4 scale shape {desc.scale_shape} != expected " + f"{expected_scale} (out, logical_in / {block_len})" + ) + return + + # BLOCK_FP8 + if len(desc.packed_shape) != 2 or len(desc.scale_shape) != 2: + raise BlockQuantValidationError( + f"{desc.name}: block-fp8 weight and scale must both be 2D" + ) + if desc.logical_shape != desc.packed_shape: + raise BlockQuantValidationError( + f"{desc.name}: block-fp8 logical {desc.logical_shape} must equal packed " + f"{desc.packed_shape} (E4M3 is not sub-byte packed)" + ) + if not desc.block_shape or len(desc.block_shape) != 2: + raise BlockQuantValidationError(f"{desc.name}: block-fp8 needs a 2D block_shape") + out_d, in_d = desc.logical_shape + bs0, bs1 = desc.block_shape + if bs0 <= 0 or bs1 <= 0: + raise BlockQuantValidationError(f"{desc.name}: invalid block_shape {desc.block_shape}") + expected_scale = (-(-out_d // bs0), -(-in_d // bs1)) + if desc.scale_shape != expected_scale: + raise BlockQuantValidationError( + f"{desc.name}: block-fp8 scale grid {desc.scale_shape} != expected " + f"{expected_scale} (ceil(out/{bs0}), ceil(in/{bs1}))" + ) + + +# --------------------------------------------------------------------------- +# Index-level pairing + classification +# --------------------------------------------------------------------------- + +_WEIGHT_SUFFIX = ".weight" +_SCALE_SUFFIX = ".scale" + + +def pair_weight_scales( + header_index: dict[str, tuple[str, tuple[int, ...]]], +) -> dict[str, str | None]: + """Pair each ``.weight`` with its ``.scale`` sibling. + + ``header_index`` maps ``name -> (dtype, shape)``. Returns ``{weight_name: + scale_name | None}`` for every ``.weight`` key. Raises + :class:`BlockQuantValidationError` on an *orphan* scale (a ``.scale`` whose + ``.weight`` sibling is absent) — a duplicate/misnamed scale that would + otherwise be silently dropped. + """ + weights = {k for k in header_index if k.endswith(_WEIGHT_SUFFIX)} + scales = {k for k in header_index if k.endswith(_SCALE_SUFFIX)} + pairing: dict[str, str | None] = {} + for w in sorted(weights): + sibling = w[: -len(_WEIGHT_SUFFIX)] + _SCALE_SUFFIX + pairing[w] = sibling if sibling in scales else None + orphans = sorted( + s for s in scales if (s[: -len(_SCALE_SUFFIX)] + _WEIGHT_SUFFIX) not in weights + ) + if orphans: + raise BlockQuantValidationError( + f"{len(orphans)} orphan scale tensor(s) with no matching .weight: {orphans[:5]}" + ) + return pairing + + +def build_descriptors( + header_index: dict[str, tuple[str, tuple[int, ...]]], + scheme: BlockQuantScheme | None = None, + *, + validate: bool = True, +) -> dict[str, QuantizedTensorDescriptor]: + """Classify every ``.weight`` tensor in a header index into a descriptor. + + ``header_index`` maps ``name -> (dtype, shape)`` (the header-only triple a + caller gets from a safetensors reader). Scale tensors are consumed as pairs + and not returned standalone. When *validate* is True each descriptor is run + through :func:`validate_descriptor`, so a malformed group raises atomically. + """ + pairing = pair_weight_scales(header_index) + descriptors: dict[str, QuantizedTensorDescriptor] = {} + for weight_name, scale_name in pairing.items(): + w_dtype, w_shape = header_index[weight_name] + if scale_name is not None: + s_dtype, s_shape = header_index[scale_name] + else: + s_dtype, s_shape = None, None + desc = classify_tensor( + weight_name, + w_dtype, + w_shape, + scale_dtype=s_dtype, + scale_shape=s_shape, + scale_name=scale_name, + scheme=scheme, + ) + if validate and desc.kind is not QuantKind.UNSUPPORTED: + validate_descriptor(desc) + descriptors[weight_name] = desc + return descriptors + + +# --------------------------------------------------------------------------- +# Byte-exact expert-major bank stacking (reusable lowering primitive) +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class PackedExpertBank: + """A byte-exact expert-major concatenation of per-expert packed weights. + + ``data`` is ``num_experts`` equal-length payloads laid out expert-major + (expert 0's bytes, then expert 1's, ...). No re-quantization or re-ordering + within an expert occurs, so the original per-expert bytes are recoverable + by slicing. This is the reusable primitive a future planar->interleaved + transcode or a native bank emitter consumes. + """ + + num_experts: int + per_expert_num_bytes: int + per_expert_packed_shape: tuple[int, ...] + weight_dtype: str + data: bytes + + def expert_bytes(self, i: int) -> bytes: + """Return expert ``i``'s original bytes (byte-exact slice).""" + if not 0 <= i < self.num_experts: + raise IndexError(f"expert {i} out of range [0, {self.num_experts})") + off = i * self.per_expert_num_bytes + return self.data[off : off + self.per_expert_num_bytes] + + +def stack_expert_bank( + per_expert_bytes: list[bytes], + *, + per_expert_packed_shape: tuple[int, ...], + weight_dtype: str, +) -> PackedExpertBank: + """Concatenate per-expert packed bytes expert-major, byte-for-byte. + + Every expert must contribute exactly the same number of bytes (a ragged + bank is a hard error, never zero-padded). The result preserves each + expert's payload verbatim. + """ + if not per_expert_bytes: + raise BlockQuantValidationError("cannot stack an empty expert bank") + n0 = len(per_expert_bytes[0]) + for i, b in enumerate(per_expert_bytes): + if len(b) != n0: + raise BlockQuantValidationError( + f"ragged expert bank: expert 0 has {n0} bytes but expert {i} has {len(b)}" + ) + return PackedExpertBank( + num_experts=len(per_expert_bytes), + per_expert_num_bytes=n0, + per_expert_packed_shape=tuple(int(d) for d in per_expert_packed_shape), + weight_dtype=weight_dtype, + data=b"".join(per_expert_bytes), + ) + + +# --------------------------------------------------------------------------- +# Runtime (nxrt) emission gate — prove representability or typed-reject +# --------------------------------------------------------------------------- + +#: Block formats the onnx-genai ``nxrt`` CPU kernel's ``BlockFormat::parse`` +#: accepts (``crates/onnx-runtime-ep-cpu/src/kernels/block_quantized_matmul.rs``). +#: MXFP4 here means the *interleaved* llama.cpp ``block_mxfp4`` layout +#: (``QK=32``, 17 bytes/block: 1 E8M0 byte + 16 nibble bytes) packed into a +#: single tensor — NOT a planar (separate nibble + separate scale) layout. +NXRT_BLOCK_FORMATS: frozenset[str] = frozenset( + { + "mxfp4", + "iq4_nl", + "iq4_xs", + "iq3_s", + "iq3_xxs", + "iq2_s", + "iq2_xs", + "iq2_xxs", + "iq1_s", + "iq1_m", + } +) + + +def runtime_representation_gap( + desc: QuantizedTensorDescriptor, *, runtime: str = "nxrt" +) -> str | None: + """Return a precise ABI-gap string if *desc* is not runtime-representable. + + Returns ``None`` when the tensor could be emitted for *runtime* today. + Purely a property check — it never emits a node. Only ``nxrt`` is modelled. + """ + if runtime != "nxrt": + return f"unknown runtime {runtime!r}; only 'nxrt' representability is modelled" + + if desc.kind is QuantKind.ORDINARY: + return None + + if desc.kind is QuantKind.BLOCK_FP8: + return ( + "nxrt has no block-FP8 BlockFormat: its BlockFormat::parse accepts only " + f"{sorted(NXRT_BLOCK_FORMATS)} and there is no E4M3-weight x 2D-UE8M0-block-scale " + "dequant path in block_quantized_{matmul,moe}.rs. Emitting a block-fp8 " + f"projection ({desc.name}, block {desc.block_shape}) would be unrunnable." + ) + + if desc.kind is QuantKind.FP4_PACKED: + return ( + "nxrt MXFP4 requires the interleaved single-tensor llama.cpp block_mxfp4 " + "layout (QK=32, 17 bytes/block = 1 E8M0 byte + 16 nibble bytes); this " + f"checkpoint stores fp4 experts *planar* ({desc.name}: I8 packed " + f"{desc.packed_shape} nibbles + a separate F8_E8M0 {desc.scale_shape} block-32 " + "micro-scale). No planar-fp4 bank ABI exists, and a planar->interleaved " + "transcode is unproven (E2M1 nibble order + E8M0 exponent bias vs llama.cpp " + "must be verified). Emitting with the split tensors would be unrunnable." + ) + + return f"{desc.name}: unsupported tensor ({desc.unsupported_reason})" + + +def plan_routed_expert_bank( + expert_descriptors: list[QuantizedTensorDescriptor], + *, + runtime: str = "nxrt", + per_expert_bytes: list[bytes] | None = None, +) -> PackedExpertBank: + """Prove a routed-expert bank is runtime-representable, else typed-reject. + + Validates that every routed expert shares one packed shape / dtype / scale + layout, then checks :func:`runtime_representation_gap`. Because *runtime* + cannot represent either quantized family today, this raises + :class:`BlockQuantExportError` naming the exact gap rather than emitting an + unrunnable node. It never falls back to a dense per-expert graph and never + dequantizes. + + When (and only when) a future runtime *can* represent the bank, the caller + supplies the byte-exact per-expert payloads via *per_expert_bytes* and a + :class:`PackedExpertBank` is returned — the signature is stable for that + path and for Deckard's #593 integration. + """ + if not expert_descriptors: + raise BlockQuantValidationError("cannot plan an empty routed-expert bank") + + non_routed = [d.name for d in expert_descriptors if not d.is_routed_expert] + if non_routed: + raise BlockQuantValidationError( + f"plan_routed_expert_bank got non-routed tensor(s): {non_routed[:5]}" + ) + + first = expert_descriptors[0] + for d in expert_descriptors: + validate_descriptor(d) + if (d.kind, d.packed_shape, d.weight_dtype, d.scale_layout) != ( + first.kind, + first.packed_shape, + first.weight_dtype, + first.scale_layout, + ): + raise BlockQuantValidationError( + "mixed expert bank: all routed experts must share kind/packed_shape/" + f"dtype/scale layout; {d.name} differs from {first.name}" + ) + + gap = runtime_representation_gap(first, runtime=runtime) + if gap is not None: + raise BlockQuantExportError( + f"Routed-expert bank ({len(expert_descriptors)} x {first.kind.value}) is not " + f"representable by the {runtime!r} runtime, so no BlockQuantizedMoE node is " + f"emitted (fail closed, no dense fallback, no dequantization). ABI gap: {gap} " + "Resolving this needs a runtime ABI extension (a planar-fp4 / block-fp8 " + "BlockFormat) or a proven byte-exact planar->interleaved transcode primitive." + ) + + # Representable path (not reachable for today's nxrt ABI). Build the bank + # only from caller-supplied byte-exact payloads — never from placeholders. + if per_expert_bytes is None: + raise BlockQuantValidationError( + "runtime can represent the bank but no per_expert_bytes were supplied to pack it" + ) + if len(per_expert_bytes) != len(expert_descriptors): + raise BlockQuantValidationError( + f"per_expert_bytes count {len(per_expert_bytes)} != " + f"{len(expert_descriptors)} expert descriptors" + ) + return stack_expert_bank( + per_expert_bytes, + per_expert_packed_shape=first.packed_shape, + weight_dtype=first.weight_dtype, + ) diff --git a/src/mobius/integrations/_block_quant_test.py b/src/mobius/integrations/_block_quant_test.py new file mode 100644 index 000000000..5a51416b0 --- /dev/null +++ b/src/mobius/integrations/_block_quant_test.py @@ -0,0 +1,575 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the block-scaled FP8 / packed-FP4 quantized weight contract. + +Covers property-based classification, logical-vs-packed shape validation, scale +pairing, byte-preservation, bounded lazy loading, the byte-exact expert-major +bank stacking primitive, and the runtime emission gate's typed reject. Real +DeepSeek-V4 ``quantization_config`` + a measured slice of the checkpoint index +metadata are used as fixtures (no full-weight download); tiny synthetic packed +safetensors exercise byte-level behaviour. +""" + +from __future__ import annotations + +import json +import pathlib +import struct + +import pytest +import torch +from safetensors.torch import save_file + +from mobius.integrations._block_quant import ( + BlockQuantExportError, + BlockQuantScheme, + BlockQuantValidationError, + LazyRawTensor, + QuantizedTensorDescriptor, + QuantKind, + build_descriptors, + classify_tensor, + pair_weight_scales, + plan_routed_expert_bank, + read_raw_tensor_bytes, + read_safetensors_header, + runtime_representation_gap, + stack_expert_bank, + validate_descriptor, +) + +# --------------------------------------------------------------------------- +# Real DeepSeek-V4 metadata fixtures (measured from the checkpoint headers). +# These are the exact quantization_config + a slice of index dtype/shape triples; +# no weight bytes are needed to classify/validate. +# --------------------------------------------------------------------------- + +REAL_QUANT_CONFIG = { + "quant_method": "fp8", + "fmt": "e4m3", + "scale_fmt": "ue8m0", + "weight_block_size": [128, 128], + "activation_scheme": "dynamic", +} +REAL_EXPERT_DTYPE = "fp4" + +# name -> (dtype, shape) exactly as read from the real safetensors headers. +REAL_HEADER_SLICE: dict[str, tuple[str, tuple[int, ...]]] = { + "layers.0.ffn.experts.0.w1.weight": ("I8", (2048, 2048)), + "layers.0.ffn.experts.0.w1.scale": ("F8_E8M0", (2048, 128)), + "layers.0.ffn.experts.0.w2.weight": ("I8", (4096, 1024)), + "layers.0.ffn.experts.0.w2.scale": ("F8_E8M0", (4096, 64)), + "layers.0.ffn.experts.0.w3.weight": ("I8", (2048, 2048)), + "layers.0.ffn.experts.0.w3.scale": ("F8_E8M0", (2048, 128)), + "layers.0.ffn.shared_experts.w1.weight": ("F8_E4M3", (2048, 4096)), + "layers.0.ffn.shared_experts.w1.scale": ("F8_E8M0", (16, 32)), + "layers.0.ffn.shared_experts.w2.weight": ("F8_E4M3", (4096, 2048)), + "layers.0.ffn.shared_experts.w2.scale": ("F8_E8M0", (32, 16)), + "layers.0.ffn.gate.weight": ("BF16", (256, 4096)), + "layers.0.attn_norm.weight": ("BF16", (4096,)), + "layers.0.attn.wq_a.weight": ("F8_E4M3", (1024, 4096)), + "layers.0.attn.wq_a.scale": ("F8_E8M0", (8, 32)), +} + +REAL_CHECKPOINT = pathlib.Path( + "/datadisks/disk5/justinchu/onnx-genai-models/deepseek-v4-flash/checkpoint" +) + + +def _real_scheme() -> BlockQuantScheme: + scheme = BlockQuantScheme.from_quantization_config( + REAL_QUANT_CONFIG, expert_dtype=REAL_EXPERT_DTYPE + ) + assert scheme is not None + return scheme + + +# --------------------------------------------------------------------------- +# Scheme parsing +# --------------------------------------------------------------------------- + + +class TestBlockQuantScheme: + def test_real_config_parses_owned(self): + scheme = _real_scheme() + assert scheme.quant_method == "fp8" + assert scheme.weight_block_size == (128, 128) + assert scheme.scale_fmt == "ue8m0" + assert scheme.expert_dtype == "fp4" + assert scheme.is_block_scaled_fp8 + assert scheme.has_packed_fp4_experts + assert scheme.is_owned + + def test_per_tensor_fp8_not_owned(self): + # No weight_block_size => ordinary per-tensor fp8 => not this contract. + assert BlockQuantScheme.from_quantization_config({"quant_method": "fp8"}) is None + + def test_absent_config_not_owned(self): + assert BlockQuantScheme.from_quantization_config(None) is None + + def test_fp4_experts_without_method(self): + scheme = BlockQuantScheme.from_quantization_config( + {"quant_method": "none"}, expert_dtype="fp4" + ) + assert scheme is not None and scheme.has_packed_fp4_experts + + def test_from_hf_config_reads_expert_dtype(self): + hf = type("HF", (), {})() + hf.quantization_config = REAL_QUANT_CONFIG + hf.expert_dtype = "fp4" + scheme = BlockQuantScheme.from_hf_config(hf) + assert scheme is not None and scheme.is_owned + + +# --------------------------------------------------------------------------- +# Property-based classification against real metadata +# --------------------------------------------------------------------------- + + +class TestClassifyRealMetadata: + def test_fp4_packed_expert(self): + scheme = _real_scheme() + d = classify_tensor( + "layers.0.ffn.experts.0.w1.weight", + "I8", + (2048, 2048), + scale_dtype="F8_E8M0", + scale_shape=(2048, 128), + scale_name="layers.0.ffn.experts.0.w1.scale", + scheme=scheme, + ) + assert d.kind is QuantKind.FP4_PACKED + assert d.packed_shape == (2048, 2048) + assert d.logical_shape == (2048, 4096) # last dim doubled (2 nibbles/byte) + assert d.block_shape == (1, 32) + assert d.microscale_kind == "mxfp4" + assert d.is_routed_expert and not d.is_shared_expert + assert d.pack_factor == 2 + validate_descriptor(d) + + def test_block_fp8_shared_expert(self): + scheme = _real_scheme() + d = classify_tensor( + "layers.0.ffn.shared_experts.w1.weight", + "F8_E4M3", + (2048, 4096), + scale_dtype="F8_E8M0", + scale_shape=(16, 32), + scale_name="layers.0.ffn.shared_experts.w1.scale", + scheme=scheme, + ) + assert d.kind is QuantKind.BLOCK_FP8 + assert d.logical_shape == d.packed_shape == (2048, 4096) + assert d.block_shape == (128, 128) + assert d.is_shared_expert and not d.is_routed_expert + assert d.pack_factor == 1 + validate_descriptor(d) + + def test_block_fp8_attention_projection(self): + scheme = _real_scheme() + d = classify_tensor( + "layers.0.attn.wq_a.weight", + "F8_E4M3", + (1024, 4096), + scale_dtype="F8_E8M0", + scale_shape=(8, 32), + scheme=scheme, + ) + assert d.kind is QuantKind.BLOCK_FP8 + assert not d.is_routed_expert and not d.is_shared_expert + validate_descriptor(d) + + def test_ordinary_router_and_norm(self): + scheme = _real_scheme() + gate = classify_tensor("layers.0.ffn.gate.weight", "BF16", (256, 4096), scheme=scheme) + norm = classify_tensor("layers.0.attn_norm.weight", "BF16", (4096,), scheme=scheme) + assert gate.kind is QuantKind.ORDINARY and gate.scale_name is None + assert norm.kind is QuantKind.ORDINARY + validate_descriptor(gate) + validate_descriptor(norm) + + def test_build_descriptors_over_real_slice(self): + scheme = _real_scheme() + descs = build_descriptors(REAL_HEADER_SLICE, scheme) + # Only .weight keys become descriptors; scales are consumed as pairs. + assert all(name.endswith(".weight") for name in descs) + kinds = {name: d.kind for name, d in descs.items()} + assert kinds["layers.0.ffn.experts.0.w2.weight"] is QuantKind.FP4_PACKED + assert kinds["layers.0.ffn.shared_experts.w2.weight"] is QuantKind.BLOCK_FP8 + assert kinds["layers.0.ffn.gate.weight"] is QuantKind.ORDINARY + # Routed vs shared classification is structural. + assert descs["layers.0.ffn.experts.0.w1.weight"].is_routed_expert + assert descs["layers.0.ffn.shared_experts.w1.weight"].is_shared_expert + + +# --------------------------------------------------------------------------- +# Unsupported / malformed inputs fail closed +# --------------------------------------------------------------------------- + + +class TestUnsupportedAndValidation: + def test_i8_without_scale_is_unsupported(self): + d = classify_tensor("w.weight", "I8", (4, 8)) + assert d.kind is QuantKind.UNSUPPORTED + assert "micro-scale" in d.unsupported_reason + with pytest.raises(BlockQuantValidationError): + validate_descriptor(d) + + def test_fp8_weight_without_scale_is_unsupported(self): + d = classify_tensor("w.weight", "F8_E4M3", (4, 8)) + assert d.kind is QuantKind.UNSUPPORTED + + def test_unknown_dtype_pair_is_unsupported(self): + d = classify_tensor( + "w.weight", "I16", (4, 8), scale_dtype="F8_E8M0", scale_shape=(4, 1) + ) + assert d.kind is QuantKind.UNSUPPORTED + + def test_wrong_fp4_scale_shape_raises(self): + # logical_in = 16 -> expected scale last dim 16/32 is invalid; use a + # deliberately wrong scale grid. + d = classify_tensor( + "e.experts.0.w1.weight", + "I8", + (8, 32), # logical (8, 64) + scale_dtype="F8_E8M0", + scale_shape=(8, 3), # wrong: expected (8, 2) + ) + assert d.kind is QuantKind.FP4_PACKED + with pytest.raises(BlockQuantValidationError, match="scale shape"): + validate_descriptor(d) + + def test_wrong_block_fp8_scale_grid_raises(self): + scheme = BlockQuantScheme.from_quantization_config( + {"quant_method": "fp8", "weight_block_size": [128, 128]} + ) + d = classify_tensor( + "attn.wq_a.weight", + "F8_E4M3", + (1024, 4096), + scale_dtype="F8_E8M0", + scale_shape=(9, 32), # wrong: expected (8, 32) + scheme=scheme, + ) + with pytest.raises(BlockQuantValidationError, match="scale grid"): + validate_descriptor(d) + + def test_ordinary_must_not_carry_scale(self): + d = QuantizedTensorDescriptor( + name="x.weight", + kind=QuantKind.ORDINARY, + weight_dtype="BF16", + logical_shape=(4, 4), + packed_shape=(4, 4), + weight_num_bytes=32, + is_routed_expert=False, + is_shared_expert=False, + scale_name="x.scale", + scale_shape=(1, 1), + ) + with pytest.raises(BlockQuantValidationError): + validate_descriptor(d) + + +# --------------------------------------------------------------------------- +# Scale pairing: missing / duplicate / orphan +# --------------------------------------------------------------------------- + + +class TestScalePairing: + def test_pairs_weight_with_scale(self): + index = { + "a.weight": ("F8_E4M3", (4, 4)), + "a.scale": ("F8_E8M0", (1, 1)), + "b.weight": ("BF16", (2, 2)), + } + pairing = pair_weight_scales(index) + assert pairing == {"a.weight": "a.scale", "b.weight": None} + + def test_orphan_scale_raises(self): + index = { + "a.weight": ("F8_E4M3", (4, 4)), + "a.scale": ("F8_E8M0", (1, 1)), + "ghost.scale": ("F8_E8M0", (1, 1)), # no matching .weight + } + with pytest.raises(BlockQuantValidationError, match="orphan"): + pair_weight_scales(index) + + def test_missing_scale_for_quantized_weight_raises_on_validate(self): + # A block-fp8 weight whose scale is absent classifies UNSUPPORTED and + # cannot be validated as a real tensor. + index = {"a.weight": ("F8_E4M3", (4, 4))} + descs = build_descriptors(index, _real_scheme(), validate=False) + assert descs["a.weight"].kind is QuantKind.UNSUPPORTED + with pytest.raises(BlockQuantValidationError): + validate_descriptor(descs["a.weight"]) + + +# --------------------------------------------------------------------------- +# Byte-preserving raw reader + bounded lazy loader (synthetic safetensors) +# --------------------------------------------------------------------------- + + +def _write_synthetic_experts(directory: pathlib.Path, n_experts: int = 3) -> pathlib.Path: + """Write a tiny safetensors file with fp4-packed experts + real dtypes.""" + state: dict[str, torch.Tensor] = {} + torch.manual_seed(0) + for e in range(n_experts): + # logical (8, 64) -> packed int8 (8, 32); scale e8m0 (8, 2). + state[f"layers.0.ffn.experts.{e}.w1.weight"] = torch.randint( + -128, 127, (8, 32), dtype=torch.int8 + ) + state[f"layers.0.ffn.experts.{e}.w1.scale"] = torch.randint( + 0, 254, (8, 2), dtype=torch.uint8 + ).view(torch.float8_e8m0fnu) + # One block-fp8 shared projection with real E4M3 storage. + state["layers.0.ffn.shared_experts.w1.weight"] = torch.randint( + 0, 200, (4, 4), dtype=torch.uint8 + ).view(torch.float8_e4m3fn) + state["layers.0.ffn.shared_experts.w1.scale"] = torch.randint( + 0, 254, (2, 2), dtype=torch.uint8 + ).view(torch.float8_e8m0fnu) + path = directory / "model.safetensors" + save_file(state, str(path)) + return path + + +class TestRawReaderAndLazy: + def test_raw_bytes_are_byte_exact(self, tmp_path): + path = _write_synthetic_experts(tmp_path) + key = "layers.0.ffn.experts.0.w1.weight" + raw = read_raw_tensor_bytes(path, key) + # Compare against the tensor's own byte view (uint8) — must be identical. + from safetensors import safe_open + + with safe_open(path, framework="pt") as h: + t = h.get_tensor(key) + assert raw == t.view(torch.uint8).contiguous().numpy().tobytes() + + def test_lazy_tensor_knows_size_from_header(self, tmp_path): + path = _write_synthetic_experts(tmp_path) + lazy = LazyRawTensor.open(path, "layers.0.ffn.experts.0.w1.weight") + assert lazy.dtype == "I8" + assert lazy.shape == (8, 32) + assert lazy.num_bytes == 8 * 32 # 1 byte/int8 + # Reading is deferred and repeatable (byte-exact each time). + first = lazy.read() + assert len(first) == lazy.num_bytes + assert lazy.read() == first + + def test_e8m0_scale_bytes_preserved(self, tmp_path): + path = _write_synthetic_experts(tmp_path) + raw = read_raw_tensor_bytes(path, "layers.0.ffn.experts.0.w1.scale") + assert len(raw) == 8 * 2 # e8m0 is 1 byte/elem + + def test_header_only_reader_does_not_scan_data(self, tmp_path): + path = _write_synthetic_experts(tmp_path) + header = read_safetensors_header(path) + assert "layers.0.ffn.experts.0.w1.weight" in header + assert header["layers.0.ffn.experts.0.w1.weight"]["dtype"] == "I8" + + +# --------------------------------------------------------------------------- +# Byte-exact expert-major bank stacking +# --------------------------------------------------------------------------- + + +class TestExpertBankStacking: + def test_stack_is_byte_exact_and_recoverable(self, tmp_path): + path = _write_synthetic_experts(tmp_path, n_experts=4) + per_expert = [ + read_raw_tensor_bytes(path, f"layers.0.ffn.experts.{e}.w1.weight") + for e in range(4) + ] + bank = stack_expert_bank( + per_expert, per_expert_packed_shape=(8, 32), weight_dtype="I8" + ) + assert bank.num_experts == 4 + assert bank.per_expert_num_bytes == 8 * 32 + assert bank.data == b"".join(per_expert) + for e in range(4): + assert bank.expert_bytes(e) == per_expert[e] + + def test_ragged_bank_raises(self): + with pytest.raises(BlockQuantValidationError, match="ragged"): + stack_expert_bank( + [b"\x00\x01", b"\x02"], per_expert_packed_shape=(1, 2), weight_dtype="I8" + ) + + def test_empty_bank_raises(self): + with pytest.raises(BlockQuantValidationError): + stack_expert_bank([], per_expert_packed_shape=(1, 2), weight_dtype="I8") + + +# --------------------------------------------------------------------------- +# Runtime emission gate: typed reject with exact ABI gap +# --------------------------------------------------------------------------- + + +def _routed_fp4_descs(n: int = 3) -> list[QuantizedTensorDescriptor]: + scheme = _real_scheme() + return [ + classify_tensor( + f"layers.0.ffn.experts.{e}.w1.weight", + "I8", + (2048, 2048), + scale_dtype="F8_E8M0", + scale_shape=(2048, 128), + scheme=scheme, + ) + for e in range(n) + ] + + +class TestEmissionGate: + def test_gap_none_for_ordinary(self): + d = classify_tensor("gate.weight", "BF16", (8, 8)) + assert runtime_representation_gap(d) is None + + def test_gap_block_fp8_names_missing_format(self): + scheme = _real_scheme() + d = classify_tensor( + "attn.wq_a.weight", + "F8_E4M3", + (1024, 4096), + scale_dtype="F8_E8M0", + scale_shape=(8, 32), + scheme=scheme, + ) + gap = runtime_representation_gap(d) + assert gap is not None and "block-FP8" in gap + + def test_gap_fp4_names_planar_vs_interleaved(self): + d = _routed_fp4_descs(1)[0] + gap = runtime_representation_gap(d) + assert gap is not None + assert "planar" in gap and "block_mxfp4" in gap + + def test_plan_routed_bank_typed_rejects_fp4(self): + with pytest.raises(BlockQuantExportError) as ei: + plan_routed_expert_bank(_routed_fp4_descs(3)) + msg = str(ei.value) + assert "not representable" in msg + assert "no dense fallback" in msg + + def test_plan_rejects_mixed_bank(self): + scheme = _real_scheme() + experts = _routed_fp4_descs(2) + odd = classify_tensor( + "layers.0.ffn.experts.2.w2.weight", + "I8", + (4096, 1024), # different packed shape + scale_dtype="F8_E8M0", + scale_shape=(4096, 64), + scheme=scheme, + ) + with pytest.raises(BlockQuantValidationError, match="mixed expert bank"): + plan_routed_expert_bank([*experts, odd]) + + def test_plan_rejects_non_routed(self): + scheme = _real_scheme() + shared = classify_tensor( + "layers.0.ffn.shared_experts.w1.weight", + "F8_E4M3", + (2048, 4096), + scale_dtype="F8_E8M0", + scale_shape=(16, 32), + scheme=scheme, + ) + with pytest.raises(BlockQuantValidationError, match="non-routed"): + plan_routed_expert_bank([shared]) + + def test_plan_empty_raises(self): + with pytest.raises(BlockQuantValidationError): + plan_routed_expert_bank([]) + + +# --------------------------------------------------------------------------- +# Integration with QuantizationConfig.from_transformers (typed blocker) +# --------------------------------------------------------------------------- + + +class TestFromTransformersBlocker: + def _hf(self, **kw): + o = type("HF", (), {})() + for k, v in kw.items(): + setattr(o, k, v) + return o + + def test_block_scaled_fp8_raises_typed(self): + from mobius._configs import QuantizationConfig + + hf = self._hf(quantization_config=REAL_QUANT_CONFIG, expert_dtype="fp4") + with pytest.raises(BlockQuantExportError, match="Block-scaled FP8"): + QuantizationConfig.from_transformers(hf) + + def test_per_tensor_fp8_still_returns_none(self): + from mobius._configs import QuantizationConfig + + hf = self._hf(quantization_config={"quant_method": "fp8", "bits": 8}) + assert QuantizationConfig.from_transformers(hf) is None + + def test_gptq_unaffected(self): + from mobius._configs import QuantizationConfig + + hf = self._hf(quantization_config={"quant_method": "gptq", "bits": 4}) + assert QuantizationConfig.from_transformers(hf).quant_method == "gptq" + + +# --------------------------------------------------------------------------- +# Opt-in: exercise the real checkpoint headers when the disk is present. +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + not REAL_CHECKPOINT.is_dir(), reason="real DeepSeek-V4 checkpoint not mounted" +) +class TestRealCheckpointHeaders: + def _index(self) -> dict: + return json.loads((REAL_CHECKPOINT / "model.safetensors.index.json").read_text()) + + def test_real_layer0_classifies_and_gate_rejects(self): + scheme = _real_scheme() + wm = self._index()["weight_map"] + # Build a header index for layer 0 by reading only shard headers. + keys = [ + k + for k in wm + if k.startswith("layers.0.") + and (".experts.0." in k or "shared_experts" in k or k.endswith("gate.weight")) + ] + header_index: dict[str, tuple[str, tuple[int, ...]]] = {} + hdr_cache: dict[str, dict] = {} + for k in keys: + shard = wm[k] + if shard not in hdr_cache: + hdr_cache[shard] = read_safetensors_header(REAL_CHECKPOINT / shard) + e = hdr_cache[shard][k] + header_index[k] = (e["dtype"], tuple(e["shape"])) + descs = build_descriptors(header_index, scheme) + routed = [ + d for d in descs.values() if d.is_routed_expert and d.name.endswith("w1.weight") + ] + assert routed and all(d.kind is QuantKind.FP4_PACKED for d in routed) + with pytest.raises(BlockQuantExportError): + plan_routed_expert_bank(routed) + + def test_real_expert_bytes_preserved(self): + wm = self._index()["weight_map"] + key = "layers.0.ffn.experts.0.w1.weight" + shard = REAL_CHECKPOINT / wm[key] + raw = read_raw_tensor_bytes(shard, key) + _dtype, shape, start, end = _span(shard, key) + assert len(raw) == end - start + assert shape == (2048, 2048) + + +def _span(path, key): + with open(path, "rb") as f: + (n,) = struct.unpack("