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 <justinchuby@users.noreply.github.com>

Author: justinchuby

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)

src/mobius/integrations/modelopt/_dequant_test.py

@@ -36,27 +36,34 @@ 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
 
     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():