[Kernel] Honor leading strides in SM90 tensorwise CUTLASS scaled_mm

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh

@@ -89,6 +89,13 @@ void cutlass_gemm_caller(torch::stable::Tensor& out,
       cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, L));
   StrideD d_stride =
       cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, L));
+
+  // Preserve the runtime leading strides of valid padded tensor views. L is
+  // fixed to 1 by get_problem_shape(), so there is no batch stride to forward.
+  cute::get<0>(a_stride) = a.stride(0);
+  cute::get<0>(b_stride) = b.stride(1);
+  cute::get<0>(c_stride) = out.stride(0);
+  cute::get<0>(d_stride) = out.stride(0);
   StrideAux aux_stride = d_stride;
 
   auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
@@ -107,4 +114,4 @@ void cutlass_gemm_caller(torch::stable::Tensor& out,
                                   epilogue_args);
 }
 
-}  // namespace vllm::c3x
+}  // namespace vllm::c3x

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh

@@ -262,6 +262,17 @@ void cutlass_gemm_caller_sm90_fp8(torch::stable::Tensor& out,
       StrideC{},
       swap_ab ? cute::make_shape(n, m, 1) : cute::make_shape(m, n, 1));
 
+  // Preserve the runtime leading strides of valid padded tensor views. The
+  // swapped epilogue treats the row-major output as its column-major
+  // transpose, whose leading stride is the second stride component.
+  cute::get<0>(a_stride) = a.stride(0);
+  cute::get<0>(b_stride) = b.stride(1);
+  if constexpr (swap_ab) {
+    cute::get<1>(c_stride) = out.stride(0);
+  } else {
+    cute::get<0>(c_stride) = out.stride(0);
+  }
+
   auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
   auto b_ptr = static_cast<ElementAB*>(b.data_ptr());
   auto c_ptr = static_cast<ElementD*>(out.data_ptr());

csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu

@@ -207,7 +207,7 @@ void cutlass_scaled_mm(torch::stable::Tensor& c, torch::stable::Tensor const& a,
   // Check for strides and alignment
   STD_TORCH_CHECK(a.stride(1) == 1 && c.stride(1) == 1);  // Row-major
   STD_TORCH_CHECK(b.stride(0) == 1);                      // Column-major
-  STD_TORCH_CHECK(c.stride(0) % 16 == 0 &&
+  STD_TORCH_CHECK(a.stride(0) % 16 == 0 && c.stride(0) % 16 == 0 &&
                   b.stride(1) % 16 == 0);  // 16 Byte Alignment
 
   if (bias) {
@@ -396,7 +396,7 @@ void cutlass_scaled_mm_azp(torch::stable::Tensor& c,
   // Check for strides and alignment
   STD_TORCH_CHECK(a.stride(1) == 1 && c.stride(1) == 1);  // Row-major
   STD_TORCH_CHECK(b.stride(0) == 1);                      // Column-major
-  STD_TORCH_CHECK(c.stride(0) % 16 == 0 &&
+  STD_TORCH_CHECK(a.stride(0) % 16 == 0 && c.stride(0) % 16 == 0 &&
                   b.stride(1) % 16 == 0);  // 16 Byte Alignment
   STD_TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
 

tests/kernels/quantization/test_cutlass_scaled_mm.py

@@ -561,23 +561,53 @@ def test_cutlass_int8_azp(
         )
 
 
-# Test working with a subset of A and B
-def test_cutlass_subset():
-    big_m, big_n, big_k = 1024, 1024, 1024
-    m, n, k = 512, 512, 512
+@pytest.mark.parametrize("padded_tensor", ["a", "b", "out"])
+@pytest.mark.parametrize("input_dtype", ["int8", "fp8"])
+@pytest.mark.parametrize("m", [32, 512])
+def test_cutlass_strided_subsets(padded_tensor: str, input_dtype: str, m: int):
+    if input_dtype == "fp8" and not current_platform.has_device_capability(89):
+        pytest.skip("FP8 is not supported on this GPU type.")
 
-    whole_a = to_int8(torch.randn((big_m, big_k), device="cuda") * 5)
-    whole_b = to_int8(torch.randn((big_n, big_k), device="cuda").t() * 5)
-    a = whole_a[0:m, 0:k]
-    b = whole_b[0:k, 0:n]
+    big_m, big_n, big_k = 1024, 1024, 1024
+    n, k = 512, 512
+
+    quantize = to_fp8 if input_dtype == "fp8" else to_int8
+    whole_a = quantize(torch.randn((big_m, big_k), device="cuda") * 5)
+    whole_b = quantize(torch.randn((big_n, big_k), device="cuda").t() * 5)
+    a = whole_a[0:m, 0:k].contiguous() if padded_tensor != "a" else whole_a[0:m, 0:k]
+    b = (
+        whole_b[0:k, 0:n].t().contiguous().t()
+        if padded_tensor != "b"
+        else whole_b[0:k, 0:n]
+    )
 
     scale_a = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
     scale_b = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
 
-    out = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype=torch.bfloat16)
+    if padded_tensor == "out":
+        whole_out = torch.empty((m, big_n), device="cuda", dtype=torch.bfloat16)
+        out = whole_out[0:m, 0:n]
+        torch.ops._C.cutlass_scaled_mm(out, a, b, scale_a, scale_b, None)
+    else:
+        out = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16)
     baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype=torch.bfloat16)
 
-    torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
+    if input_dtype == "fp8":
+        torch.testing.assert_close(out, baseline, rtol=5e-1, atol=1.5e-1)
+    else:
+        torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
+
+
+def test_cutlass_rejects_misaligned_a_leading_stride():
+    m = n = k = 512
+    whole_a = to_int8(torch.randn((m, k + 1), device="cuda") * 5)
+    a = whole_a[:, :k]
+    b = to_int8(torch.randn((n, k), device="cuda").t() * 5)
+    scale_a = torch.ones((1, 1), device="cuda", dtype=torch.float32)
+    scale_b = torch.ones((1, 1), device="cuda", dtype=torch.float32)
+
+    with pytest.raises(RuntimeError):
+        ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16)
 
 
 # Test to make sure cuda graphs work