Add NVFP4/FP8 (ModelOpt) weight reconstruction foundation for Qwen3.6

src/mobius/_configs/_quantization.py

@@ -53,6 +53,27 @@ def from_transformers(cls, hf_config) -> QuantizationConfig | None:
         if not isinstance(qc, dict):
             return None
         method = qc.get("quant_method", "none")
+        # NVIDIA ModelOpt NVFP4/FP8 checkpoints (e.g. quantized Qwen3.6) encode
+        # weights as packed E2M1 (fp4) / float8 with block + global scales — a
+        # layout the INT4 ``QuantizedLinear``/``MatMulNBits`` path below would
+        # silently mis-dequantize. Checked before the ``quant_method == "none"``
+        # early-return because ModelOpt may name its scheme only via
+        # ``quant_algo``/``quant_cfg`` (no ``quant_method``). Fail loudly: the
+        # reconstruction math is available in :mod:`mobius.integrations.modelopt`,
+        # but the full weight-load integration + native routed-expert NVFP4 QMoE
+        # emission (CUDA/Blackwell, ``onnxruntime_USE_FP4_QMOE=ON``) are not wired.
+        from mobius.integrations.modelopt import is_modelopt_quant_config
+
+        if is_modelopt_quant_config(qc):
+            raise NotImplementedError(
+                "NVIDIA ModelOpt NVFP4/FP8 checkpoints are not yet fully "
+                "supported for export. The weight-reconstruction core lives in "
+                "mobius.integrations.modelopt (dequantize_nvfp4 / dequantize_fp8); "
+                "remaining work is checkpoint weight-loading and native "
+                "routed-expert NVFP4 QMoE emission (CUDA/Blackwell, "
+                "onnxruntime_USE_FP4_QMOE=ON). Export the unquantized (bf16) "
+                "checkpoint instead, or quantize the bf16 export via Olive."
+            )
         if method == "none":
             return None
         # FP8 per-tensor quantization (float8_e4m3fn + scalar scale)

src/mobius/_configs_test.py

@@ -837,6 +837,24 @@ def test_from_transformers_fp8_returns_none(self):
         )()
         assert QuantizationConfig.from_transformers(hf) is None
 
+    def test_from_transformers_modelopt_nvfp4_raises(self):
+        """ModelOpt NVFP4/FP8 checkpoints fail loudly rather than mis-quantizing.
+
+        The INT4 path would silently mis-dequantize packed E2M1/float8 weights,
+        so ``from_transformers`` raises until the ModelOpt loader is wired.
+        """
+        import pytest
+
+        for qc in (
+            {"quant_method": "modelopt"},
+            {"quant_method": "modelopt", "quant_algo": "NVFP4"},
+            {"quant_algo": "W4A16_NVFP4"},
+            {"quant_algo": "FP8"},
+        ):
+            hf = type("HFConfig", (), {"quantization_config": qc})()
+            with pytest.raises(NotImplementedError, match="ModelOpt"):
+                QuantizationConfig.from_transformers(hf)
+
     def test_from_transformers_to_dict_object(self):
         """HF QuantizationConfig objects have a to_dict() method."""
         inner = type(

src/mobius/integrations/modelopt/__init__.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""NVIDIA TensorRT Model Optimizer (ModelOpt) checkpoint support.
+
+ModelOpt exports mixed-precision NVFP4 + FP8 checkpoints (e.g. Qwen3.6). This
+package provides the weight-reconstruction math used to consume those
+checkpoints:
+
+- NVFP4 (``W4A16_NVFP4``): block-16 E2M1 4-bit weights with FP8-E4M3 block
+  scales and a per-tensor FP32 global scale (``weight_scale_2``).
+- FP8 (``E4M3``): per-tensor scaled float8 weights.
+
+The dequantization functions reconstruct BF16 weights so the standard mobius
+build path (plain ``Linear`` / ``bf16`` graph) can consume ModelOpt checkpoints
+without a native FP8/NVFP4 kernel. Native routed-expert NVFP4 QMoE emission
+(the CUDA-only, Blackwell ``QMoE`` ``quant_type="nvfp4"`` op) is intentionally
+out of scope here — see the module docstring in :mod:`._dequant`.
+"""
+
+from __future__ import annotations
+
+from mobius.integrations.modelopt._dequant import (
+    FP4_E2M1_LUT,
+    dequantize_fp8,
+    dequantize_nvfp4,
+    is_modelopt_quant_config,
+    unpack_nvfp4_codes,
+)
+
+__all__ = [
+    "FP4_E2M1_LUT",
+    "dequantize_fp8",
+    "dequantize_nvfp4",
+    "is_modelopt_quant_config",
+    "unpack_nvfp4_codes",
+]

src/mobius/integrations/modelopt/_dequant.py

@@ -0,0 +1,131 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""NVFP4 / FP8 weight reconstruction for NVIDIA ModelOpt checkpoints.
+
+NVIDIA TensorRT Model Optimizer (ModelOpt) exports mixed-precision checkpoints
+where different modules use different numeric formats. For Qwen3.6 the layout is:
+
+- **Routed MoE experts, shared expert, lm_head**: ``W4A16_NVFP4`` — block-16
+  E2M1 (fp4) 4-bit weights, FP8-E4M3 per-block scales, and a per-tensor FP32
+  global scale stored as ``weight_scale_2``.
+- **Attention / linear-attention projections**: ``FP8`` (E4M3) — float8 weights
+  with a per-tensor ``weight_scale``.
+
+ONNX Runtime has no FP8 attention GEMM, so — following ORT GenAI's ModelOpt
+loader — the dense FP8 projections and the NVFP4 shared-expert / lm_head are
+**dequantized back to BF16**, which reconstructs them exactly and lets the
+standard (BF16) build path consume the checkpoint. Only the *routed* MoE experts
+are consumed natively by the CUDA QMoE ``quant_type="nvfp4"`` op; that native
+emission is Blackwell/``onnxruntime_USE_FP4_QMOE=ON``-only and is intentionally
+NOT implemented here (it cannot be built or verified without that runtime).
+
+The functions below are the numeric core of the loader and are format-faithful
+to ModelOpt's on-disk encoding (verified in ``_dequant_test.py``).
+
+E2M1 (fp4) magnitude table, indexed by the low 3 bits of each 4-bit code
+(bit 3 is the sign): ``{0, 0.5, 1, 1.5, 2, 3, 4, 6}``.
+"""
+
+from __future__ import annotations
+
+import ml_dtypes
+import numpy as np
+
+# E2M1 (fp4) decoded magnitudes, indexed by ``code & 0x7`` (bit 3 is the sign).
+FP4_E2M1_LUT = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=np.float32)
+
+# NVFP4 pins the weight block size to 16 K-elements per FP8-E4M3 block scale.
+NVFP4_BLOCK_SIZE = 16
+
+
+def unpack_nvfp4_codes(packed_nk2: np.ndarray) -> np.ndarray:
+    """Unpack a ModelOpt NVFP4 weight tensor to per-element E2M1 codes.
+
+    ``packed_nk2`` is uint8 ``[N, K/2]`` where each byte holds two adjacent
+    K-axis E2M1 codes for the same output row ``N`` (low nibble = even ``K``,
+    high nibble = odd ``K``) — the layout ModelOpt writes. Returns uint8 codes
+    ``[N, K]`` in ``0..15``.
+    """
+    packed = np.ascontiguousarray(packed_nk2).astype(np.uint8)
+    low = packed & 0x0F
+    high = packed >> 4
+    n = packed.shape[0]
+    # Interleave (low, high) along a new trailing axis, then flatten to [N, K].
+    codes = np.stack((low, high), axis=-1).reshape(n, -1)
+    return np.ascontiguousarray(codes)
+
+
+def dequantize_nvfp4(
+    weight_u8: np.ndarray,
+    block_scale_e4m3: np.ndarray,
+    global_scale: float | np.floating,
+) -> np.ndarray:
+    """Reconstruct a BF16 weight from ModelOpt NVFP4 tensors.
+
+    Args:
+        weight_u8: uint8 ``[N, K/2]`` packed E2M1 codes (low nibble = even
+            ``K``, high nibble = odd ``K``).
+        block_scale_e4m3: FP8-E4M3 ``[N, K/16]`` per-block scales
+            (``ml_dtypes.float8_e4m3fn`` or a uint8 raw-code view).
+        global_scale: per-tensor FP32 scalar.
+
+    Returns:
+        ``ml_dtypes.bfloat16`` array ``[N, K]`` where
+        ``w = e2m1(code) * e4m3(block_scale[n, k // 16]) * global_scale``.
+    """
+    codes = unpack_nvfp4_codes(weight_u8).astype(np.int64)  # [N, K]
+    mag = FP4_E2M1_LUT[codes & 0x7]  # [N, K]
+    val = np.where((codes & 0x8) > 0, -mag, mag).astype(np.float32)  # [N, K]
+
+    block_scale = _to_float32(block_scale_e4m3)  # [N, K/16]
+    k = codes.shape[1]
+    n_blocks = block_scale.shape[1]
+    if k % n_blocks != 0:
+        raise ValueError(f"NVFP4 K={k} is not divisible by the block count {n_blocks}.")
+    block_scale = np.repeat(block_scale, k // n_blocks, axis=1)  # [N, K]
+
+    dequant = val * block_scale * np.float32(global_scale)
+    return dequant.astype(ml_dtypes.bfloat16)
+
+
+def dequantize_fp8(
+    weight_f8: np.ndarray,
+    weight_scale: float | np.floating,
+) -> np.ndarray:
+    """Reconstruct a BF16 weight from an FP8 (E4M3) weight + per-tensor scale.
+
+    ``w = e4m3(weight) * weight_scale``.
+    """
+    weight = _to_float32(weight_f8)
+    return (weight * np.float32(weight_scale)).astype(ml_dtypes.bfloat16)
+
+
+def is_modelopt_quant_config(quantization_config: dict | None) -> bool:
+    """Return ``True`` if a HF ``quantization_config`` is a ModelOpt export.
+
+    ModelOpt writes ``quant_method="modelopt"`` (newer HF serialisations) or a
+    ``quant_algo`` / ``quant_cfg`` naming an ``NVFP4`` / ``FP8`` scheme. This
+    keeps detection tolerant of both spellings.
+    """
+    if not isinstance(quantization_config, dict):
+        return False
+    method = str(quantization_config.get("quant_method", "")).lower()
+    if method == "modelopt":
+        return True
+    algo = str(
+        quantization_config.get("quant_algo") or quantization_config.get("quant_cfg") or ""
+    ).upper()
+    return "NVFP4" in algo or "FP8" in algo
+
+
+def _to_float32(arr: np.ndarray) -> np.ndarray:
+    """View/convert an FP8-E4M3 (or raw uint8 code) tensor as float32.
+
+    Raw ``uint8`` inputs are reinterpreted as ``float8_e4m3fn`` byte codes
+    before widening, matching how ModelOpt stores block scales.
+    """
+    arr = np.asarray(arr)
+    if arr.dtype == np.uint8:
+        arr = arr.view(ml_dtypes.float8_e4m3fn)
+    return arr.astype(np.float32)

src/mobius/integrations/modelopt/_dequant_test.py

@@ -0,0 +1,101 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Tests for ModelOpt NVFP4 / FP8 weight reconstruction."""
+
+from __future__ import annotations
+
+import ml_dtypes
+import numpy as np
+
+from mobius.integrations.modelopt import (
+    FP4_E2M1_LUT,
+    dequantize_fp8,
+    dequantize_nvfp4,
+    is_modelopt_quant_config,
+    unpack_nvfp4_codes,
+)
+
+
+def _e2m1_code(sign: int, mag_index: int) -> int:
+    """Build a 4-bit E2M1 code from a sign bit and a magnitude index (0..7)."""
+    return (sign << 3) | (mag_index & 0x7)
+
+
+def test_fp4_lut_values():
+    assert FP4_E2M1_LUT.tolist() == [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
+
+
+def test_unpack_nvfp4_codes_splits_nibbles():
+    # byte = low | (high << 4); low nibble is the even-K code.
+    packed = np.array([[0x21, 0xC6]], dtype=np.uint8)  # (1, 0x21=33, 0xC6=198)
+    codes = unpack_nvfp4_codes(packed)
+    # 0x21 -> low=1, high=2 ; 0xC6 -> low=6, high=12
+    assert codes.tolist() == [[1, 2, 6, 12]]
+    assert codes.dtype == np.uint8
+
+
+def test_dequantize_nvfp4_known_values():
+    # K=4 codes: k0=+1.0, k1=-0.5, k2=+4.0, k3=-2.0
+    k0 = _e2m1_code(0, 2)  # mag 1.0
+    k1 = _e2m1_code(1, 1)  # -0.5
+    k2 = _e2m1_code(0, 6)  # 4.0
+    k3 = _e2m1_code(1, 4)  # -2.0
+    byte0 = k0 | (k1 << 4)
+    byte1 = k2 | (k3 << 4)
+    weight_u8 = np.array([[byte0, byte1]], dtype=np.uint8)
+
+    block_scale = np.array([[2.0]], dtype=ml_dtypes.float8_e4m3fn)  # one block
+    global_scale = 0.5
+
+    out = dequantize_nvfp4(weight_u8, block_scale, global_scale)
+    assert out.dtype == ml_dtypes.bfloat16
+    # val * 2.0 (block) * 0.5 (global) == val
+    expected = np.array([[1.0, -0.5, 4.0, -2.0]], dtype=np.float32)
+    np.testing.assert_array_equal(out.astype(np.float32), expected)
+
+
+def test_dequantize_nvfp4_block_scale_repeat():
+    # Two 16-element blocks with distinct block scales exercise repeat_interleave.
+    n, k = 1, 32
+    # All codes = +1.0 (index 2) -> byte 0x22 packs two such codes.
+    weight_u8 = np.full((n, k // 2), 0x22, dtype=np.uint8)
+    block_scale = np.array([[1.0, 4.0]], dtype=ml_dtypes.float8_e4m3fn)
+    out = dequantize_nvfp4(weight_u8, block_scale, 1.0).astype(np.float32)
+    assert out.shape == (1, 32)
+    # First 16 elements scaled by 1.0, next 16 by 4.0.
+    np.testing.assert_array_equal(out[0, :16], np.ones(16, dtype=np.float32))
+    np.testing.assert_array_equal(out[0, 16:], np.full(16, 4.0, dtype=np.float32))
+
+
+def test_dequantize_nvfp4_raw_uint8_block_scale():
+    # Block scales may arrive as raw uint8 e4m3 code bytes; result must match a
+    # typed float8 view.
+    k0 = _e2m1_code(0, 2)
+    weight_u8 = np.array([[k0 | (k0 << 4)]], dtype=np.uint8)  # K=2, both +1.0
+    typed = np.array([[3.0]], dtype=ml_dtypes.float8_e4m3fn)
+    raw = typed.view(np.uint8)
+    a = dequantize_nvfp4(weight_u8, typed, 1.0).astype(np.float32)
+    b = dequantize_nvfp4(weight_u8, raw, 1.0).astype(np.float32)
+    np.testing.assert_array_equal(a, b)
+    np.testing.assert_array_equal(a, np.array([[3.0, 3.0]], dtype=np.float32))
+
+
+def test_dequantize_fp8_per_tensor_scale():
+    weight = np.array([[1.0, -2.0, 0.5, 3.0]], dtype=ml_dtypes.float8_e4m3fn)
+    out = dequantize_fp8(weight, 0.25)
+    assert out.dtype == ml_dtypes.bfloat16
+    expected = np.array([[0.25, -0.5, 0.125, 0.75]], dtype=np.float32)
+    np.testing.assert_array_equal(out.astype(np.float32), expected)
+
+
+def test_is_modelopt_quant_config():
+    assert is_modelopt_quant_config({"quant_method": "modelopt"})
+    assert is_modelopt_quant_config({"quant_algo": "NVFP4"})
+    assert is_modelopt_quant_config({"quant_cfg": "W4A16_NVFP4"})
+    assert is_modelopt_quant_config({"quant_algo": "FP8"})
+    # Non-ModelOpt schemes and empty configs are rejected.
+    assert not is_modelopt_quant_config({"quant_method": "gptq"})
+    assert not is_modelopt_quant_config({"quant_method": "awq"})
+    assert not is_modelopt_quant_config({})
+    assert not is_modelopt_quant_config(None)