[CPU][QSDPA] Fix issue with strided input (#3794)

* [cpu][qsdpa] fix issue with strided input

Author: Valentine233

test/test_ops.py

@@ -138,23 +138,27 @@ def test_quantized_scaled_dot_product_op(
         torch.manual_seed(1234)
         device = "cpu"
         if input_dtype == torch.uint8:
-            q_scale = float(1.7907238006591797)
-            k_scale = float(1.8039721250534058)
-            v_scale = float(1.839004635810852)
-            a_scale = float(0.003919653594493866)
-            o_scale = float(1.8191684484481812)
-            q_zp = int(127)
-            k_zp = int(125)
-            v_zp = int(127)
-            a_zp = int(120)
-            o_zp = int(128)
+            q_scale, k_scale, v_scale = (
+                float(1.7907238006591797),
+                float(1.8039721250534058),
+                float(1.839004635810852),
+            )
+            a_scale, o_scale = float(0.003919653594493866), float(1.8191684484481812)
+            q_zp, k_zp, v_zp, a_zp, o_zp = (
+                int(127),
+                int(125),
+                int(127),
+                int(120),
+                int(128),
+            )
             atol, rtol = 1.0, 5e-6
         else:
-            q_scale = float(5.96875)
-            k_scale = float(5.78125)
-            v_scale = float(0.98046875)
-            a_scale = float(4.84375)
-            o_scale = float(3.171875)
+            q_scale, k_scale, v_scale = (
+                float(5.96875),
+                float(5.78125),
+                float(0.98046875),
+            )
+            a_scale, o_scale = float(4.84375), float(3.171875)
             atol, rtol = 0.125, 5e-6
         q_shape = [batch_size, q_seq_len, n_head, head_dim]
         kv_shape = [batch_size, kv_seq_len, n_head, head_dim]
@@ -163,12 +167,8 @@ def test_quantized_scaled_dot_product_op(
         k = torch.randn(kv_shape, dtype=torch.float, device=device).transpose(1, 2)
         v = torch.randn(kv_shape, dtype=torch.float, device=device).transpose(1, 2)
         if input_dtype == torch.uint8:
-            q *= 100
-            k *= 100
-            v *= 100
-        q = q.to(input_dtype)
-        k = k.to(input_dtype)
-        v = v.to(input_dtype)
+            q, k, v = (t * 100 for t in (q, k, v))
+        q, k, v = (t.to(input_dtype) for t in (q, k, v))
         attn_mask = (
             torch.randn(mask_shape, dtype=mask_dtype, device=device)
             if mask_dtype is not None
@@ -181,72 +181,106 @@ def test_quantized_scaled_dot_product_op(
             attn_mask.clone() if mask_dtype is not None else None,
         )
 
+        common_kwargs = dict(
+            dropout_p=0.0,
+            is_causal=False,
+            q_scale=q_scale,
+            k_scale=k_scale,
+            v_scale=v_scale,
+            a_scale=a_scale,
+            o_scale=o_scale,
+        )
+
         if input_dtype == torch.uint8:
+            common_kwargs.update(q_zp=q_zp, k_zp=k_zp, v_zp=v_zp, a_zp=a_zp, o_zp=o_zp)
             math_ref = self._scaled_dot_product_int8_op_ref(
-                q2,
-                k2,
-                v2,
-                attn_mask=attn_mask,
-                dropout_p=0.0,
-                is_causal=False,
-                q_scale=q_scale,
-                q_zp=q_zp,
-                k_scale=k_scale,
-                k_zp=k_zp,
-                v_scale=v_scale,
-                v_zp=v_zp,
-                a_scale=a_scale,
-                a_zp=a_zp,
-                o_scale=o_scale,
-                o_zp=o_zp,
+                q2, k2, v2, attn_mask=attn_mask, **common_kwargs
             )
             actual = torch.ops.torchao.qscaled_dot_product(
-                q,
-                k,
-                v,
-                attn_mask=attn_mask_2,
-                dropout_p=0.0,
-                is_causal=False,
-                q_scale=q_scale,
-                q_zp=q_zp,
-                k_scale=k_scale,
-                k_zp=k_zp,
-                v_scale=v_scale,
-                v_zp=v_zp,
-                a_scale=a_scale,
-                a_zp=a_zp,
-                o_scale=o_scale,
-                o_zp=o_zp,
+                q, k, v, attn_mask=attn_mask_2, **common_kwargs
             )
         else:
             math_ref = self._scaled_dot_product_fp8_op_ref(
-                q2,
-                k2,
-                v2,
-                attn_mask=attn_mask,
-                dropout_p=0.0,
-                is_causal=False,
-                q_scale=q_scale,
-                k_scale=k_scale,
-                v_scale=v_scale,
-                a_scale=a_scale,
-                o_scale=o_scale,
+                q2, k2, v2, attn_mask=attn_mask, **common_kwargs
             )
             actual = torch.ops.torchao.qscaled_dot_product(
-                q,
-                k,
-                v,
-                attn_mask=attn_mask_2,
-                dropout_p=0.0,
-                is_causal=False,
-                q_scale=q_scale,
-                k_scale=k_scale,
-                v_scale=v_scale,
-                a_scale=a_scale,
-                o_scale=o_scale,
+                q, k, v, attn_mask=attn_mask_2, **common_kwargs
             )
         self.assertEqual(actual.float(), math_ref.float(), atol=atol, rtol=rtol)
 
+    @pytest.mark.skipif(
+        not torch_version_at_least("2.7.0"),
+        reason="quantized sdpa requires torch 2.7 or later",
+    )
+    @pytest.mark.skipif(not IS_LINUX, reason="only support on linux")
+    @pytest.mark.skipif(
+        "CPU" not in torch._C._dispatch_dump("torchao::qscaled_dot_product"),
+        reason="cpp kernels not built",
+    )
+    @parametrize("input_dtype", [torch.uint8, torch.float8_e4m3fn])
+    def test_quantized_scaled_dot_product_op_with_strided_inputs(
+        self,
+        input_dtype,
+    ):
+        torch.manual_seed(1234)
+        device = "cpu"
+        if input_dtype == torch.uint8:
+            q_scale, k_scale, v_scale = (
+                float(1.7907238006591797),
+                float(1.8039721250534058),
+                float(1.839004635810852),
+            )
+            a_scale, o_scale = float(0.003919653594493866), float(1.8191684484481812)
+            q_zp, k_zp, v_zp, a_zp, o_zp = (
+                int(127),
+                int(125),
+                int(127),
+                int(120),
+                int(128),
+            )
+            atol, rtol = 1.0, 5e-6
+        else:
+            q_scale, k_scale, v_scale = (
+                float(5.96875),
+                float(5.78125),
+                float(0.98046875),
+            )
+            a_scale, o_scale = float(4.84375), float(3.171875)
+            atol, rtol = 0.125, 5e-6
+        batch_size, seq_len, num_head, head_dim = 56, 100, 16, 32
+        hidden_size = num_head * head_dim
+        proj = torch.nn.Linear(hidden_size, 3 * hidden_size, bias=False)
+        input_shape = (batch_size, seq_len, num_head * head_dim)
+        qkv_shape = (batch_size, seq_len, num_head, head_dim)
+        input_tensor = torch.randn(input_shape, dtype=torch.float, device=device)
+        hidden_state = proj(input_tensor)
+        q, k, v = torch.split(hidden_state, hidden_size, dim=-1)
+        q, k, v = (t.view(*qkv_shape).transpose(1, 2) for t in (q, k, v))
+        if input_dtype == torch.uint8:
+            q, k, v = (t * 100 for t in (q, k, v))
+        q, k, v = (t.to(input_dtype) for t in (q, k, v))
+        q2, k2, v2 = (t.clone() for t in (q, k, v))
+
+        common_kwargs = dict(
+            attn_mask=None,
+            dropout_p=0.0,
+            is_causal=False,
+            q_scale=q_scale,
+            k_scale=k_scale,
+            v_scale=v_scale,
+            a_scale=a_scale,
+            o_scale=o_scale,
+        )
+        if input_dtype == torch.uint8:
+            common_kwargs.update(q_zp=q_zp, k_zp=k_zp, v_zp=v_zp, a_zp=a_zp, o_zp=o_zp)
+            math_ref = self._scaled_dot_product_int8_op_ref(q2, k2, v2, **common_kwargs)
+            actual = torch.ops.torchao.qscaled_dot_product(q, k, v, **common_kwargs)
+        else:
+            math_ref = self._scaled_dot_product_fp8_op_ref(q2, k2, v2, **common_kwargs)
+            actual = torch.ops.torchao.qscaled_dot_product(q, k, v, **common_kwargs)
+        assert actual.transpose(1, 2).is_contiguous(), "Output is not contiguous!"
+        self.assertEqual(actual.float(), math_ref.float(), atol=atol, rtol=rtol)
+
 
 instantiate_parametrized_tests(TestOps)
 

torchao/csrc/cpu/aten_kernels/quantized_sdpa.cpp

@@ -2525,11 +2525,15 @@ at::Tensor _qscaled_dot_product_cpu(
     TORCH_CHECK(q_zp == 0 && k_zp == 0 && v_zp == 0 && a_zp == 0 && o_zp == 0,
       "_qscaled_dot_product_cpu: Don't accept zero point for Float8_e4m3");
   }
+  int64_t batch_size = query.size(0);
+  int64_t num_head = query.size(1);
+  int64_t q_seq_len = query.size(2);
+  int64_t head_size = query.size(3);
 
   if (dtype == at::ScalarType::Byte) {
 #ifdef CPU_CAPABILITY_AVX512
       if (at::native::cpublas::could_pack(dtype)) {
-          at::Tensor output = at::empty_like(query, query.options()).transpose(1, 2);
+          at::Tensor output = at::empty({batch_size, q_seq_len, num_head, head_size}, query.options());
           int8_sdpa_fused_kernel(output, query, key, value,
               dropout_p, is_causal, attn_mask, scale,
               q_scale, q_zp,
@@ -2554,7 +2558,7 @@ at::Tensor _qscaled_dot_product_cpu(
 #if defined(CPUBLAS_BRGEMM_F8F8F32) && defined(CPU_CAPABILITY_AVX512)
 // CPUBLAS_BRGEMM_F8F8F32 is defined if FP8 BRGEMM is supported in PyTorch CPUBlas.
       if (at::native::cpublas::could_pack(dtype)) {
-          at::Tensor output = at::empty_like(query, query.options()).transpose(1, 2);
+          at::Tensor output = at::empty({batch_size, q_seq_len, num_head, head_size}, query.options());
           fp8_sdpa_fused_kernel(output, query, key, value,
               dropout_p, is_causal, attn_mask, scale,
               q_scale, k_scale,

torchao/ops.py

@@ -152,7 +152,7 @@ def _(
     o_scale: float = 1.0,
     o_zp: int = 0,
 ) -> Tensor:
-    return query
+    return query.transpose(1, 2).contiguous().transpose(1, 2)
 
 
 def rowwise_scaled_linear_sparse_cutlass_f8f8(