[upstream_ut] test_max_reads_limits_fusion (OverFusionTest) precision failure on PVC due to sdpa flash attention sycltla
- Dominant language
- Python
- Stars
- 113
- Forks
- 128
- Avg merge
- 5d 13h
- Merged PRs (30d)
- 107
Description
## Summary
`OverFusionTest.test_max_reads_limits_fusion` in `test/inductor/test_mix_order_reduction.py` fails on Intel PVC (Data Center GPU Max 1100) due to a precision issue in the SYCL Flash Attention backward path with wheels building upon LTS driver.
**The root cause is that SYCLTLA on PVC with LTS driver does not currently guarantee precision/numerical determinism.** This means the gradient accumulation in the XPU Flash Attention backward kernel can produce results that vary beyond the test tolerance (`tol=5e-2`).
When oneDNN adds training (backward) support for this attention pattern on PVC, the dispatch will go through oneDNN instead of the current SYCL TLA path, which is expected to resolve the precision issue.
## Root Cause Confirmed via Ablation
The SDPA backward path is confirmed as the root cause by ablation: **replacing `F.scaled_dot_product_attention` with `sdpa_kernel(SDPBackend.MATH)` makes the test pass**, even in the CI Docker environment (Ubuntu 24.04 Noble) with CI build wheels (MKL 2024.2). All other kernels in the backward pass (RMSNorm backward, MLP backward, mix_order_reduction triton kernels) remain unchanged in the ablation.
| Test | Backend | Result |
|---|---|---|
| `OverFusionTest.test_max_reads_limits_fusion` | OVERRIDEABLE (XPU Flash / SYCL TLA) | **FAIL** |
| `OverFusionSdpaAblationTest.test_max_reads_limits_fusion_sdpa_math` | MATH (no SYCL TLA) | **PASS** |
This isolates the precision failure exclusively to `_scaled_dot_product_fused_attention_overrideable_xpu` backward (SYCL TLA on PVC build upon LTS driver).
The ablation test code (not in upstream pytorch, for local reproduction only):
```python
class OverFusionSdpaAblationTest(TestBase):
@inductor_config.patch(
{
"triton.mix_order_reduction": True,
"triton.mix_order_reduction_max_reads": 10,
"force_disable_caches": True,
}
)
def test_max_reads_limits_fusion_sdpa_math(self):
if not HAS_GPU:
self.skipTest("requires GPU")
from torch.nn.attention import sdpa_kernel, SDPBackend
num_heads = 8
num_kv_heads = 4
dim = 512
head_dim = dim // num_heads
class Attention(nn.Module):
def __init__(self):
super().__init__()
self.c_q = nn.Linear(dim, dim, bias=False)
self.c_k = nn.Linear(dim, num_kv_heads * head_dim, bias=False)
self.c_v = nn.Linear(dim, num_kv_heads * head_dim, bias=False)
self.proj = nn.Linear(dim, dim, bias=False)
def forward(self, x):
B, T, D = x.shape
q = self.c_q(x).reshape(B, T, num_heads, head_dim)
k = self.c_k(x).reshape(B, T, num_kv_heads, head_dim)
v = self.c_v(x).reshape(B, T, num_kv_heads, head_dim)
q = F.rms_norm(q, (q.size(-1),))
k = F.rms_norm(k, (k.size(-1),))
q = q.transpose(1, 2)
k = k.transpose(1, 2).repeat_interleave(num_heads // num_kv_heads, dim=1)
v = v.transpose(1, 2).repeat_interleave(num_heads // num_kv_heads, dim=1)
# Force MATH backend to bypass XPU Flash Attention (OVERRIDEABLE / SYCL TLA)
with sdpa_kernel(SDPBackend.MATH):
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
return self.proj(y.transpose(1, 2).reshape(B, T, D))
class Block(nn.Module):
def __init__(self):
super().__init__()
self.attn_norm = nn.RMSNorm(dim)
self.mlp_norm = nn.RMSNorm(dim)
self.attn = Attention()
self.fc1 = nn.Linear(dim, dim * 4, bias=False)
self.fc2 = nn.Linear(dim * 4, dim, bias=False)
def forward(self, x):
x = x + self.attn(self.attn_norm(x))
h = self.mlp_norm(x)
x = x + self.fc2(F.leaky_relu(self.fc1(h), negative_slope=0.5).square())
return x
model = nn.Sequential(*[Block() for _ in range(3)]).to(GPU_TYPE).bfloat16()
x = torch.randn(8, 2048, dim, device=GPU_TYPE, dtype=torch.bfloat16, requires_grad=True)
dy = torch.randn_like(x)
out_ref = model(x)
out_ref.backward(dy)
grad_ref = x.grad.clone()
x.grad = None
compiled = torch.compile(model, dynamic=False, fullgraph=True)
out_act = compiled(x)
out_act.backward(dy)
grad_act = x.grad.clone()
self.assertTrue(same(grad_ref, grad_act, tol=5e-2))
```
## Reproducibility: CD rolling driver vs LTS2 driver (CI)
The failure is **build-environment-dependent**:
| Build environment | Result |
|---|---|
| Ubuntu 22.04 LTS (local pt_nightly wheel, `manylinux_2_28`, MKL 2026.0) | **PASS** |
| Ubuntu 24.04 LTS2 (CI wheel `linux-noble`, MKL 2024.2, Python 3.10) | **FAIL** |
The CI wheel is built on Ubuntu 24.04 Noble (`linux-noble-xpu-n-py3.10`, tag `linux_x86_64`, not manylinux). The locally-verified passing build uses `manylinux_2_28_x86_64` with a newer MKL runtime (2026.0 vs 2024.2). The oneDNN version and git hash are identical (`v3.12.0 / 80afa710`) between both builds.
The difference in MKL runtime version likely affects the SYCL TLA kernel code path on PVC, resulting in different gradient accumulation precision between the two builds.
## Background
The test exercises a **training** workload (forward + backward + grad check) using:
- GQA attention (num_heads=8, num_kv_heads=4) with QK-norm
- `F.scaled_dot_product_attention(q, k, v, is_causal=True)` in bfloat16
- `torch.compile` with `mix_order_reduction=True`, `mix_order_reduction_max_reads=10`
- Assertion: compiled grad matches eager grad within `tol=5e-2`
On XPU (PVC), `F.scaled_dot_product_attention` dispatches via the `OVERRIDEABLE` backend (SDPBackend=4) through the following chain:
```
F.scaled_dot_product_attention(q, k, v, is_causal=True) # Python entry
-> _fused_sdp_choice() = 4 (OVERRIDEABLE) # backend selection
-> aten._scaled_dot_product_flash_attention.default # AOT post_grad graph
-> Inductor extern kernel (pass-through, no triton)
-> aten._scaled_dot_product_fused_attention_overrideable # XPU dispatch
-> _scaled_dot_product_fused_attention_overrideable_xpu # mkldnn/xpu/Attention.cpp
```
The precision gap appears in the backward pass where TLA-based gradient accumulation is used. Since **SYCL TLA on PVC does not currently guarantee reproducible precision**, the grad values computed by the compiled path can exceed the `tol=5e-2` threshold when compared to the eager reference.
## oneDNN training support (expected fix path)
Currently the XPU Flash Attention backward on PVC goes through a SYCL TLA kernel path that does not guarantee precision. Once **oneDNN adds training (backward) support** for this attention pattern on PVC, the dispatch will route through oneDNN instead of the SYCL TLA path. This is the expected resolution path — at that point the precision should meet the `tol=5e-2` requirement and this test can be re-enabled on XPU.
## Build environments compared
| Item | LTS (PASS) | LTS2 / CI (FAIL) |
|---|---|---|
| PyTorch version | `2.14.0.dev20260620+xpu` | `2.14.0a0+git02fc2d3` |
| PyTorch commit | `7f2583984d` | `02fc2d3511` |
| Wheel platform tag | `manylinux_2_28_x86_64` | `linux_x86_64` (noble) |
| Build OS | Ubuntu 22.04 LTS | Ubuntu 24.04 LTS2 (Noble) |
| MKL version | 2026.0 (Build 20260401) | 2024.2 (Build 20240605) |
| oneDNN version | v3.12.0 / `80afa710` | v3.12.0 / `80afa710` |
| SYCL compiler | 20260000 | 20260000 |
| GCC | 13.3 | 13.3 |
| Python | 3.12 | 3.10 |
## Upstream tracking
This test was disabled in upstream PyTorch CI via:
- Disable issue: https://github.com/pytorch/pytorch/issues/181699
A follow-up PyTorch PR adds an explicit `@skipIfXpu` annotation referencing this issue (intel/torch-xpu-ops#4094), and will close the disable issue once the XPU-side precision is resolved:
- PR: https://github.com/pytorch/pytorch/pull/187822
## Expected resolution
Once oneDNN supports the training (backward) path for this Flash Attention pattern on PVC, the dispatch will bypass the SYCL TLA kernel and the precision issue will be resolved. At that point:
1. Remove the `@skipIfXpu` annotation from the test in upstream PyTorch
2. Close this issue and pytorch/pytorch#181699
Contributor guide
Assessment
This issue has not been assessed yet.