Add shape validation for mat1Scale and mat2Scale before kernel launch

Author: danisereb

csrc/mxfp8_gemm_cutlass.cu

@@ -100,6 +100,7 @@ void mxfp8_bmm_impl(TensorView mat1, TensorView mat2, TensorView mat1Scale, Tens
   CHECK_INPUT_AND_TYPE(mat2Scale, SF_DTYPE);
 
   int64_t m, n, k, b;
+  // Scale validation for swizzled (1D) and non-swizzled (2D) layouts.
   if (mat1.ndim() == 2) {
     TVM_FFI_ICHECK_EQ(mat2.ndim(), 2) << "mat2 must be a matrix";
     // mat2 is passed as b.T, but TensorView reads underlying storage as [N, K]
@@ -130,6 +131,89 @@ void mxfp8_bmm_impl(TensorView mat1, TensorView mat2, TensorView mat1Scale, Tens
     TVM_FFI_LOG_AND_THROW(NotImplementedError) << "mat1 must be a matrix or a batch of matrices";
   }
 
+  constexpr int64_t sfVecSize = 32;  // MXFP8 block size
+  auto scale_len = [&](int64_t dim) { return (dim + sfVecSize - 1) / sfVecSize; };
+  auto swizzled_len = [&](int64_t rows, int64_t cols) {
+    auto pad_up = [](int64_t value, int64_t multiple) {
+      return (value + multiple - 1) / multiple * multiple;
+    };
+    int64_t padded_rows = pad_up(rows, 128);
+    int64_t padded_cols = pad_up(cols, 4);
+    return padded_rows * padded_cols;
+  };
+
+  if (mat1.ndim() == 2) {
+    const int64_t k_scales = scale_len(k);
+    if (mat1Scale.ndim() == 1) {
+      int64_t expected = swizzled_len(m, k_scales);
+      TVM_FFI_ICHECK_EQ(mat1Scale.size(0), expected)
+          << "mxfp8_bmm_impl: mat1Scale size mismatch, expected " << expected << ", got "
+          << mat1Scale.size(0);
+    } else {
+      TVM_FFI_ICHECK_EQ(mat1Scale.ndim(), 2)
+          << "mxfp8_bmm_impl: mat1Scale must be 1D (swizzled) or 2D (non-swizzled), got "
+          << mat1Scale.ndim();
+      TVM_FFI_ICHECK_EQ(mat1Scale.size(0), m)
+          << "mxfp8_bmm_impl: mat1Scale size mismatch, expected " << m << ", got "
+          << mat1Scale.size(0);
+      TVM_FFI_ICHECK_EQ(mat1Scale.size(1), k_scales)
+          << "mxfp8_bmm_impl: mat1Scale size mismatch, expected " << k_scales << ", got "
+          << mat1Scale.size(1);
+    }
+
+    if (mat2Scale.ndim() == 1) {
+      int64_t expected = swizzled_len(n, k_scales);
+      TVM_FFI_ICHECK_EQ(mat2Scale.size(0), expected)
+          << "mxfp8_bmm_impl: mat2Scale size mismatch, expected " << expected << ", got "
+          << mat2Scale.size(0);
+    } else {
+      TVM_FFI_ICHECK_EQ(mat2Scale.ndim(), 2)
+          << "mxfp8_bmm_impl: mat2Scale must be 1D (swizzled) or 2D (non-swizzled), got "
+          << mat2Scale.ndim();
+      TVM_FFI_ICHECK_EQ(mat2Scale.size(0), n)
+          << "mxfp8_bmm_impl: mat2Scale size mismatch, expected " << n << ", got "
+          << mat2Scale.size(0);
+      TVM_FFI_ICHECK_EQ(mat2Scale.size(1), k_scales)
+          << "mxfp8_bmm_impl: mat2Scale size mismatch, expected " << k_scales << ", got "
+          << mat2Scale.size(1);
+    }
+  } else {
+    const int64_t k_scales = scale_len(k);
+    if (mat1Scale.ndim() == 1) {
+      int64_t expected = swizzled_len(b * m, k_scales);
+      TVM_FFI_ICHECK_EQ(mat1Scale.size(0), expected)
+          << "mxfp8_bmm_impl: mat1Scale size mismatch, expected " << expected << ", got "
+          << mat1Scale.size(0);
+    } else {
+      TVM_FFI_ICHECK_EQ(mat1Scale.ndim(), 2)
+          << "mxfp8_bmm_impl: mat1Scale must be 1D (swizzled) or 2D (non-swizzled), got "
+          << mat1Scale.ndim();
+      TVM_FFI_ICHECK_EQ(mat1Scale.size(0), b)
+          << "mxfp8_bmm_impl: mat1Scale batch size mismatch, expected " << b << ", got "
+          << mat1Scale.size(0);
+      TVM_FFI_ICHECK_EQ(mat1Scale.size(1), scale_len(m))
+          << "mxfp8_bmm_impl: mat1Scale size mismatch, expected " << scale_len(m) << ", got "
+          << mat1Scale.size(1);
+    }
+
+    if (mat2Scale.ndim() == 1) {
+      int64_t expected = swizzled_len(b * n, k_scales);
+      TVM_FFI_ICHECK_EQ(mat2Scale.size(0), expected)
+          << "mxfp8_bmm_impl: mat2Scale size mismatch, expected " << expected << ", got "
+          << mat2Scale.size(0);
+    } else {
+      TVM_FFI_ICHECK_EQ(mat2Scale.ndim(), 2)
+          << "mxfp8_bmm_impl: mat2Scale must be 1D (swizzled) or 2D (non-swizzled), got "
+          << mat2Scale.ndim();
+      TVM_FFI_ICHECK_EQ(mat2Scale.size(0), b)
+          << "mxfp8_bmm_impl: mat2Scale batch size mismatch, expected " << b << ", got "
+          << mat2Scale.size(0);
+      TVM_FFI_ICHECK_EQ(mat2Scale.size(1), scale_len(n))
+          << "mxfp8_bmm_impl: mat2Scale size mismatch, expected " << scale_len(n) << ", got "
+          << mat2Scale.size(1);
+    }
+  }
+
   // No heuristic for now, we rely on the autotuner to select the best tactic.
   if (tactic == -1) {
     tactic = 0;