[FlexAttention] explicilty create grad_q w/ strides (#152641)

drisspg · pytorchmergebot · commit a6ea63a841ca · 2025-05-02T22:57:26.000Z
Fixes: #147463 There is a mismatch between inductor's lowering for empty_like and it does not match the behavior of eager. The strides do not match preserve format #144699 Pull Request resolved: #152641 Approved by: https://github.com/xmfan
diff --git a/torch/_higher_order_ops/flex_attention.py b/torch/_higher_order_ops/flex_attention.py
@@ -781,11 +781,12 @@ def sdpa_dense_backward(
 ]:
     from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex
 
-    Bq, _, _, qk_head_dim = query.shape
+    Bq, Hq, seq_len_q, qk_head_dim = query.shape
     Bkv, Hkv, seq_len_kv, v_head_dim = value.shape
 
     # Get outputs before calling repeat interleave and permute to input stride orders
-    actual_grad_query = torch.empty_like(query)
+    actual_grad_query = query.new_empty((Bq, Hq, seq_len_q, qk_head_dim))
+    actual_grad_query = _permute_strides(actual_grad_query, query.stride())
 
     actual_grad_key = key.new_empty((Bq, Hkv, seq_len_kv, qk_head_dim))
     actual_grad_key = _permute_strides(actual_grad_key, key.stride())
diff --git a/torch/_inductor/kernel/flex_attention.py b/torch/_inductor/kernel/flex_attention.py
@@ -38,7 +38,6 @@
     _full,
     check_and_broadcast_indices,
     empty,
-    empty_like,
     empty_strided,
     expand,
     index_output_size_and_inner_fn,
@@ -2539,7 +2538,14 @@ def flex_attention_backward(*args, **kwargs):
     grad_lse_exp2, delta = maybe_realize([grad_lse_exp2, delta])
 
     # # see NOTE:[TritonTemplates with multiple outputs]
-    grad_query = empty_like(query)
+    query_size = [Bq, Hq, seq_len_q, qk_head_dim]
+    grad_query_strides = infer_dense_strides(query_size, query.get_stride())
+    grad_query = empty_strided(
+        query_size,
+        stride=[sympy.sympify(s) for s in grad_query_strides],
+        dtype=query.get_dtype(),
+        device=query.get_device(),
+    )
 
     # Construct output layout with stride order matching value
     value_size = [Bq, Hkv, seq_len_kv, v_head_dim]
