[Kernel] Reject unsupported layouts in CUTLASS 3.x scaled_mm

.buildkite/test_areas/kernels.yaml

@@ -659,7 +659,7 @@ steps:
         2)
           pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2'
           pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
-          pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
+          pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8 or cutlass_subset'
           pytest -v -s tests/kernels/quantization/test_block_fp8.py::test_w8a8_block_fp8_cutlass_matmul
           pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
           pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py

csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu

@@ -1,3 +1,5 @@
+#include <cstdint>
+
 #include <cuda.h>
 #include <cudaTypedefs.h>
 
@@ -194,6 +196,26 @@ bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
   return false;
 }
 
+namespace {
+
+void check_c3x_scaled_mm_operands(const torch::stable::Tensor& c,
+                                  const torch::stable::Tensor& a,
+                                  const torch::stable::Tensor& b) {
+  auto check = [](const torch::stable::Tensor& t, const char* name, int dim) {
+    STD_TORCH_CHECK(t.size(dim) <= 1 || t.stride(dim) == t.size(1 - dim),
+                    "CUTLASS 3.x scaled_mm requires packed operands; ", name,
+                    " is not packed");
+    STD_TORCH_CHECK(reinterpret_cast<std::uintptr_t>(t.data_ptr()) % 16 == 0,
+                    "CUTLASS 3.x scaled_mm requires 16-byte-aligned pointers; ",
+                    name, " is misaligned");
+  };
+  check(a, "A", 0);
+  check(b, "B", 1);
+  check(c, "output", 0);
+}
+
+}  // namespace
+
 void cutlass_scaled_mm(torch::stable::Tensor& c, torch::stable::Tensor const& a,
                        torch::stable::Tensor const& b,
                        torch::stable::Tensor const& a_scales,
@@ -221,13 +243,15 @@ void cutlass_scaled_mm(torch::stable::Tensor& c, torch::stable::Tensor const& a,
 
 #if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
   if (version_num >= 120) {
+    check_c3x_scaled_mm_operands(c, a, b);
     cutlass_scaled_mm_sm120(c, a, b, a_scales, b_scales, bias);
     return;
   }
 #endif
 
 #if defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM100
   if (version_num >= 100 && version_num < 120) {
+    check_c3x_scaled_mm_operands(c, a, b);
     cutlass_scaled_mm_sm100(c, a, b, a_scales, b_scales, bias);
     return;
   }
@@ -237,6 +261,7 @@ void cutlass_scaled_mm(torch::stable::Tensor& c, torch::stable::Tensor const& a,
 #if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
   if (version_num >= 90 && version_num < 100) {
     // Hopper
+    check_c3x_scaled_mm_operands(c, a, b);
     cutlass_scaled_mm_sm90(c, a, b, a_scales, b_scales, bias);
     return;
   }
@@ -425,6 +450,7 @@ void cutlass_scaled_mm_azp(torch::stable::Tensor& c,
 
 #if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
   if (version_num >= 90) {
+    check_c3x_scaled_mm_operands(c, a, b);
     cutlass_scaled_mm_azp_sm90(c, a, b, a_scales, b_scales, azp_adj, azp, bias);
     return;
   }

tests/kernels/quantization/test_cutlass_scaled_mm.py

@@ -586,6 +586,13 @@ def test_cutlass_subset():
     scale_a = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
     scale_b = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
 
+    if capability_int >= 90:
+        with pytest.raises(
+            RuntimeError, match="requires packed operands; A is not packed"
+        ):
+            ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype=torch.bfloat16)
+        return
+
     out = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype=torch.bfloat16)
     baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype=torch.bfloat16)