Add MXFP8 scale checks in _check_mm_mxfp8_problem_size

Author: danisereb

flashinfer/gemm/gemm_base.py

@@ -2513,6 +2513,45 @@ def _check_mm_mxfp8_problem_size(
     if b_descale.dtype != torch.uint8:
         raise ValueError(f"b_descale must be a uint8 tensor, got {b_descale.dtype=}")
 
+    sf_vec_size = 32
+    if a_descale.ndim == 1:
+        expected_len = _mxfp8_swizzled_scale_len(a.shape[0], a.shape[1])
+        if a_descale.shape[0] != expected_len:
+            raise ValueError(
+                "a_descale shape mismatch for swizzled layout. "
+                f"Expected {(expected_len,)}, got {a_descale.shape}."
+            )
+    elif a_descale.ndim == 2:
+        expected_shape = (a.shape[0], a.shape[1] // sf_vec_size)
+        if a_descale.shape != expected_shape:
+            raise ValueError(
+                "a_descale shape mismatch for non-swizzled layout. "
+                f"Expected {expected_shape}, got {a_descale.shape}."
+            )
+    else:
+        raise ValueError(
+            f"a_descale must be 1D (swizzled) or 2D (non-swizzled), got {a_descale.shape}."
+        )
+
+    if b_descale.ndim == 1:
+        expected_len = _mxfp8_swizzled_scale_len(b.shape[1], b.shape[0])
+        if b_descale.shape[0] != expected_len:
+            raise ValueError(
+                "b_descale shape mismatch for swizzled layout. "
+                f"Expected {(expected_len,)}, got {b_descale.shape}."
+            )
+    elif b_descale.ndim == 2:
+        expected_shape = (b.shape[0] // sf_vec_size, b.shape[1])
+        if b_descale.shape != expected_shape:
+            raise ValueError(
+                "b_descale shape mismatch for non-swizzled layout. "
+                f"Expected {expected_shape}, got {b_descale.shape}."
+            )
+    else:
+        raise ValueError(
+            f"b_descale must be 1D (swizzled) or 2D (non-swizzled), got {b_descale.shape}."
+        )
+
     if out is not None:
         expected_shape = (a.shape[0], b.shape[1])
         if out.shape != expected_shape:

tests/gemm/test_mm_mxfp8.py

@@ -214,7 +214,7 @@ def test_mm_mxfp8_invalid_ndim():
     a_descale = a_scale.view(1, -1, 1)
     b_descale = b_scale.view(1, -1, 1)
     with pytest.raises(
-        AssertionError, match="a_descale must be 1D \(swizzled\) or 2D \(non-swizzled\)"
+        ValueError, match="a_descale must be 1D \(swizzled\) or 2D \(non-swizzled\)"
     ):
         mm_mxfp8(
             a_mx,