Skip to content

Return BSHD-contiguous input grads from XPU efficient attention backward - #5433

Open
orrangetabby17 wants to merge 1 commit into
mainfrom
shuxin/fix-ut
Open

orrangetabby17 wants to merge 1 commit into
mainfrom
shuxin/fix-ut

Conversation

@orrangetabby17

@orrangetabby17 orrangetabby17 commented Sep 18, 2026

Copy link
Copy Markdown
Contributor

Motivation (Why)

_scaled_dot_product_efficient_attention_backward has a layout contract
that the XPU implementation does not honor. Its meta kernel allocates
grad_q/k/v with torch.empty_permuted((B, H, S, D), (0, 2, 1, 3))
(meta__scaled_dot_product_efficient_backward in torch/_meta_registrations.py),
i.e. BHSD sizes over BSHD-contiguous memory — strides (H*S*D, D, H*D, 1).

The XPU backward computes gradients by re-running the math forward under
autograd and returns them untouched, so they come back plain contiguous.
Eager code does not care, but torch.compile asserts the meta strides, so any
compiled use of the efficient-attention backward fails with e.g.
expected size 4==4, stride 8192==64 at dim=1 on [2, 4, 128, 64] grads.
This was first hit by
inductor/test_cuda_repro.py::CudaReproTests::test_effn_attn_uniform_zero_bias_backward.

Solution (How)

Relay grad_q/k/v out to BSHD-contiguous before returning them, reusing the
same permute/contiguous/permute normalization the XPU forward already applies
to its own output (Attention.cpp:468-469). This matches what the CUDA kernel
delivers (attention_backward.cu:1131-1132), so the single shared stride
contract is preserved.

grad_bias is deliberately left alone: the meta kernel pads only its last
dimension up to a multiple of 16 and then slices back, so a bias whose last
dimension is already 16-aligned (e.g. the [2, 4, 128, 128] in this test)
matches without a change. An unaligned bias width is a separate, untested gap.

The alternative — teaching the meta kernel to describe what XPU actually
returns, keyed off device — was rejected: it forks a stride contract CUDA,
ROCm and XPU otherwise share, and bakes a temporary property of the math-based
fallback into a device-agnostic registration.

Test Plan

Built PyTorch at 31527a43 with this branch as the third_party/torch-xpu-ops
override, then ran the upstream test that motivated this:

cd <pytorch_root>
pip install -e . -v --no-build-isolation
cd test
TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 ZE_AFFINITY_MASK=0 \
  python -m pytest -q \
  inductor/test_cuda_repro.py::CudaReproTests::test_effn_attn_uniform_zero_bias_backward

@CuiYifeng CuiYifeng left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please ensure the readability of PR description.

_scaled_dot_product_efficient_attention_backward has a layout contract
that the XPU implementation does not honor. Its meta kernel allocates
grad_q/k/v with torch.empty_permuted((B, H, S, D), (0, 2, 1, 3)), i.e.
BHSD sizes over BSHD-contiguous memory. The CUDA kernel delivers exactly
that; the XPU math-based fallback returns plain contiguous tensors, so
compiled code that asserts the meta strides fails.

The fix relays out grad_q/k/v to BSHD-contiguous before returning,
reusing the same permute/contiguous/permute normalization the XPU
forward already applies. grad_bias is left alone: the meta kernel pads
only its last dimension to a multiple of 16 and slices back.

Test Plan:

Built PyTorch and ran the upstream test that motivated this (DISABLED
for XPU as pytorch/pytorch#196553):

```bash
cd <pytorch_root>
pip install -e . -v --no-build-isolation
cd test
TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 ZE_AFFINITY_MASK=0 \
  python -m pytest -v \
  inductor/test_cuda_repro.py::CudaReproTests::test_effn_attn_uniform_zero_bias_backward
```

This change was authored with Claude Code.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@github-actions github-actions Bot added disable_e2e Disable all e2e test jobs for the PR disable_distributed Disable distributed UT test jobs for the PR labels Sep 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

disable_distributed Disable distributed UT test jobs for the PR disable_e2e Disable all e2e test jobs for the PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants