From 5fb09c26e72a4864c24e81ca46fbc0c7737a3222 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 1 Aug 2026 05:01:47 +0000 Subject: [PATCH 1/2] Add NVFP4/FP8 (ModelOpt) weight reconstruction foundation for Qwen3.6 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit NVIDIA TensorRT Model Optimizer (ModelOpt) exports mixed-precision NVFP4 + FP8 checkpoints (e.g. quantized Qwen3.6): routed MoE experts / shared expert / lm_head as W4A16_NVFP4 (block-16 E2M1 weights, FP8-E4M3 block scales, per-tensor FP32 global scale weight_scale_2), and attention/linear-attn projections as FP8 (E4M3, per-tensor scale). Add the numeric core of the ModelOpt loader under mobius/integrations/modelopt: - dequantize_nvfp4(): E2M1 codes x FP8-E4M3 block scale x FP32 global -> BF16 - dequantize_fp8(): FP8 (E4M3) x per-tensor scale -> BF16 - unpack_nvfp4_codes(), is_modelopt_quant_config(), FP4_E2M1_LUT Following ORT GenAI's loader, dense FP8 and NVFP4 shared-expert/lm_head weights are reconstructed to BF16 exactly (ORT has no FP8 attention GEMM), so the standard BF16 build path can consume the checkpoint. Fully unit-tested against hand-computed references (numeric-only; no runtime dependency). Guard QuantizationConfig.from_transformers: ModelOpt schemes now raise a clear NotImplementedError instead of being silently routed through the INT4 MatMulNBits path (which would mis-dequantize packed E2M1/float8 weights). The check runs before the quant_method=="none" early-return because ModelOpt may name its scheme only via quant_algo/quant_cfg. Out of scope (documented, follow-up): full checkpoint weight-load integration and native routed-expert NVFP4 QMoE emission (CUDA/Blackwell, onnxruntime_USE_FP4_QMOE=ON) — neither buildable nor verifiable without that runtime and a real ModelOpt checkpoint. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/_configs/_quantization.py | 21 +++ src/mobius/_configs_test.py | 18 +++ src/mobius/integrations/modelopt/__init__.py | 37 +++++ src/mobius/integrations/modelopt/_dequant.py | 131 ++++++++++++++++++ .../integrations/modelopt/_dequant_test.py | 101 ++++++++++++++ 5 files changed, 308 insertions(+) create mode 100644 src/mobius/integrations/modelopt/__init__.py create mode 100644 src/mobius/integrations/modelopt/_dequant.py create mode 100644 src/mobius/integrations/modelopt/_dequant_test.py diff --git a/src/mobius/_configs/_quantization.py b/src/mobius/_configs/_quantization.py index 40556e0b..69432ebf 100644 --- a/src/mobius/_configs/_quantization.py +++ b/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) diff --git a/src/mobius/_configs_test.py b/src/mobius/_configs_test.py index da97a82e..ba1af0ab 100644 --- a/src/mobius/_configs_test.py +++ b/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( diff --git a/src/mobius/integrations/modelopt/__init__.py b/src/mobius/integrations/modelopt/__init__.py new file mode 100644 index 00000000..a7a2c099 --- /dev/null +++ b/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", +] diff --git a/src/mobius/integrations/modelopt/_dequant.py b/src/mobius/integrations/modelopt/_dequant.py new file mode 100644 index 00000000..ce1255e0 --- /dev/null +++ b/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) diff --git a/src/mobius/integrations/modelopt/_dequant_test.py b/src/mobius/integrations/modelopt/_dequant_test.py new file mode 100644 index 00000000..0a26b79f --- /dev/null +++ b/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) From d08c344a8bab240d1e5d2347dc23fc1a6333c1ce Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sun, 2 Aug 2026 15:41:27 +0000 Subject: [PATCH 2/2] fix(modelopt): harden NVFP4 dequant shape validation and tests Address PR review feedback on the NVFP4 dequant path: - unpack_nvfp4_codes now rejects non-2D input with a clear ValueError so a loader shape/dtype mistake surfaces instead of being silently reinterpreted via the uint8 cast. - dequantize_nvfp4 enforces the fixed NVFP4 16-element block size; a derived size other than 16 means mismatched weight/scale shapes and now raises. - Rewrite the known-value and raw-uint8 block-scale tests to use full K=16 blocks (spec-realistic), asserting trailing zero-code reconstruction. - Fix stale "repeat_interleave" comment to "np.repeat". - Add regression tests for the two new validation errors. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/integrations/modelopt/_dequant.py | 21 ++++++++++- .../integrations/modelopt/_dequant_test.py | 37 +++++++++++++++---- 2 files changed, 48 insertions(+), 10 deletions(-) diff --git a/src/mobius/integrations/modelopt/_dequant.py b/src/mobius/integrations/modelopt/_dequant.py index ce1255e0..55c958d0 100644 --- a/src/mobius/integrations/modelopt/_dequant.py +++ b/src/mobius/integrations/modelopt/_dequant.py @@ -46,7 +46,16 @@ def unpack_nvfp4_codes(packed_nk2: np.ndarray) -> np.ndarray: 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``. + + Raises: + ValueError: if ``packed_nk2`` is not a 2D ``[N, K/2]`` array. Loader + code should surface an upstream shape/dtype mistake (e.g. passing + already-unpacked codes) rather than silently reinterpreting it. """ + if packed_nk2.ndim != 2: + raise ValueError( + f"NVFP4 packed codes must be 2D [N, K/2], got shape {packed_nk2.shape}." + ) packed = np.ascontiguousarray(packed_nk2).astype(np.uint8) low = packed & 0x0F high = packed >> 4 @@ -81,9 +90,17 @@ def dequantize_nvfp4( 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: + if n_blocks == 0 or 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] + block_size = k // n_blocks + if block_size != NVFP4_BLOCK_SIZE: + # NVFP4 pins the block size to 16; a different derived size means the + # weight/scale shapes are mismatched (silently-wrong reconstruction). + raise ValueError( + f"NVFP4 block size must be {NVFP4_BLOCK_SIZE}, got {block_size} " + f"(K={k}, block scales={n_blocks})." + ) + block_scale = np.repeat(block_scale, block_size, axis=1) # [N, K] dequant = val * block_scale * np.float32(global_scale) return dequant.astype(ml_dtypes.bfloat16) diff --git a/src/mobius/integrations/modelopt/_dequant_test.py b/src/mobius/integrations/modelopt/_dequant_test.py index 0a26b79f..8a4ac7f7 100644 --- a/src/mobius/integrations/modelopt/_dequant_test.py +++ b/src/mobius/integrations/modelopt/_dequant_test.py @@ -36,14 +36,17 @@ def test_unpack_nvfp4_codes_splits_nibbles(): def test_dequantize_nvfp4_known_values(): - # K=4 codes: k0=+1.0, k1=-0.5, k2=+4.0, k3=-2.0 + # One full NVFP4 block (K=16). First 4 codes are known non-zero values; + # the trailing 12 codes are zero (index 0 -> magnitude 0.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 + zero = _e2m1_code(0, 0) # 0.0 + # Pack 16 codes into 8 bytes (low nibble = even K, high nibble = odd K). byte0 = k0 | (k1 << 4) byte1 = k2 | (k3 << 4) - weight_u8 = np.array([[byte0, byte1]], dtype=np.uint8) + weight_u8 = np.array([[byte0, byte1] + [zero | (zero << 4)] * 6], dtype=np.uint8) block_scale = np.array([[2.0]], dtype=ml_dtypes.float8_e4m3fn) # one block global_scale = 0.5 @@ -51,12 +54,16 @@ def test_dequantize_nvfp4_known_values(): 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) + expected = np.zeros((1, 16), dtype=np.float32) + expected[0, :4] = [1.0, -0.5, 4.0, -2.0] np.testing.assert_array_equal(out.astype(np.float32), expected) + # Trailing 12 elements reconstruct to exactly zero. + np.testing.assert_array_equal(out.astype(np.float32)[0, 4:], np.zeros(12)) def test_dequantize_nvfp4_block_scale_repeat(): - # Two 16-element blocks with distinct block scales exercise repeat_interleave. + # Two 16-element blocks with distinct block scales exercise the np.repeat + # per-block scale broadcast. 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) @@ -70,15 +77,29 @@ def test_dequantize_nvfp4_block_scale_repeat(): 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 float8 view. One full NVFP4 block (K=16), all codes +1.0. + k0 = _e2m1_code(0, 2) # +1.0 + weight_u8 = np.full((1, 8), k0 | (k0 << 4), dtype=np.uint8) # K=16, all +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)) + np.testing.assert_array_equal(a, np.full((1, 16), 3.0, dtype=np.float32)) + + +def test_unpack_nvfp4_codes_rejects_non_2d(): + # Loader code should surface a shape mistake, not silently reinterpret it. + with np.testing.assert_raises(ValueError): + unpack_nvfp4_codes(np.zeros(8, dtype=np.uint8)) # 1D + + +def test_dequantize_nvfp4_rejects_non16_block(): + # A derived block size other than 16 indicates mismatched weight/scale shapes. + weight_u8 = np.full((1, 2), 0x22, dtype=np.uint8) # K=4 + block_scale = np.array([[1.0]], dtype=ml_dtypes.float8_e4m3fn) # 1 block -> size 4 + with np.testing.assert_raises(ValueError): + dequantize_nvfp4(weight_u8, block_scale, 1.0) def test_dequantize_fp8_per_tensor_scale():