fp32 `scaled_dot_product_attention` hangs the GPU when attention scores exceed 2^30 elements
- Dominant language
- Python
- Stars
- 113
- Forks
- 128
- Avg merge
- 5d 9h
- Merged PRs (30d)
- 112
Description
## Summary
On Intel XPU, `F.scaled_dot_product_attention` in **float32** hangs the compute engine as soon as `batch * num_q_heads * seq_len^2` exceeds **2^30**. The driver force-resets the engine and the device is unusable for the rest of the process:
```
RuntimeError: Enqueue process failed.
level_zero backend failed with error: 20 (UR_RESULT_ERROR_DEVICE_LOST)
```
```
xe 0000:17:00.0: [drm] GT0: Engine reset: engine_class=ccs, logical_mask: 0x1, guc_id=12
xe 0000:17:00.0: [drm] exec queue reset detected
```
fp16/bf16 are fine at the same shapes — including cases *larger in bytes* than the failing fp32 one — so this looks like 32-bit index overflow in the fp32 (materializing) SDPA path rather than OOM.
| | |
|---|---|
| torch | `2.15.0.dev20260829+xpu` |
| device | Intel(R) Arc(TM) Pro B60 Graphics (Battlemage G21, `8086:e211`), 25.6 GB |
| driver | `xe`, DRM 1.1.0 |
| reference | NVIDIA A100 80GB PCIe, torch `2.13.0a0+8145d630e8.nv26.06` — all cases pass |
## Reproduction
[`sdpa_fp32_hang_repro.py`](#repro-script) — no arguments, device detected at runtime, runs unchanged on CUDA and XPU.
```bash
python sdpa_fp32_hang_repro.py
```
Each case runs in a fresh subprocess: a hang kills the device for the entire process, so the remaining cases would all fail spuriously otherwise.
### Intel XPU
```
torch 2.15.0.dev20260829+xpu | xpu | Intel(R) Arc(TM) Pro B60 Graphics
dtype batch seq_len variant elements GiB result
---------------------------------------------------------------
float32 2 3072 - 603,979,776 2.25 OK out=(2, 32, 3072, 128) sum=17185.1348
float32 2 3584 - 822,083,584 3.06 OK out=(2, 32, 3584, 128) sum=-23022.0820
float32 2 4088 - 1,069,551,616 3.98 OK out=(2, 32, 4088, 128) sum=12605.1348 <- just below 2**30
float32 2 4104 - 1,077,940,224 4.02 RuntimeError: level_zero backend failed with error: 20 (UR_RESULT_ERROR_DEVICE_LOST) <- just above 2**30
float32 2 4185 - 1,120,910,400 4.18 RuntimeError: level_zero backend failed with error: 20 (UR_RESULT_ERROR_DEVICE_LOST) <- shape taken from a real model
float32 1 4185 - 560,455,200 2.09 OK out=(1, 32, 4185, 128) sum=-4510.2246 <- half the batch -> below 2**30
float32 2 4185 no-mask 1,120,910,400 4.18 RuntimeError: level_zero backend failed with error: 20 (UR_RESULT_ERROR_DEVICE_LOST) <- mask is not the trigger
float32 2 4185 no-gqa 1,120,910,400 4.18 RuntimeError: level_zero backend failed with error: 20 (UR_RESULT_ERROR_DEVICE_LOST) <- GQA is not the trigger
float16 2 4185 - 1,120,910,400 2.09 OK out=(2, 32, 4185, 128) sum=-3002.4832
float16 2 6144 - 2,415,919,104 4.50 OK out=(2, 32, 6144, 128) sum=-31061.2637 <- 2.4e9 elements / 4.5 GiB, still fine
bfloat16 2 4185 - 1,120,910,400 2.09 OK out=(2, 32, 4185, 128) sum=-2970.9536
```
### CUDA (reference — every case passes)
```
torch 2.13.0a0+8145d630e8.nv26.06 | cuda | NVIDIA A100 80GB PCIe
dtype batch seq_len variant elements GiB result
---------------------------------------------------------------
float32 2 3072 - 603,979,776 2.25 OK out=(2, 32, 3072, 128) sum=-13910.6260
float32 2 3584 - 822,083,584 3.06 OK out=(2, 32, 3584, 128) sum=15218.8047
float32 2 4088 - 1,069,551,616 3.98 OK out=(2, 32, 4088, 128) sum=-12881.0137 <- just below 2**30
float32 2 4104 - 1,077,940,224 4.02 OK out=(2, 32, 4104, 128) sum=-1405.4749 <- just above 2**30
float32 2 4185 - 1,120,910,400 4.18 OK out=(2, 32, 4185, 128) sum=-18660.0000 <- shape taken from a real model
float32 1 4185 - 560,455,200 2.09 OK out=(1, 32, 4185, 128) sum=-15130.3271 <- half the batch -> below 2**30
float32 2 4185 no-mask 1,120,910,400 4.18 OK out=(2, 32, 4185, 128) sum=-13046.5605 <- mask is not the trigger
float32 2 4185 no-gqa 1,120,910,400 4.18 OK out=(2, 32, 4185, 128) sum=-18665.0000 <- GQA is not the trigger
float16 2 4185 - 1,120,910,400 2.09 OK out=(2, 32, 4185, 128) sum=-18661.3516
float16 2 6144 - 2,415,919,104 4.50 OK out=(2, 32, 6144, 128) sum=8029.9731 <- 2.4e9 elements / 4.5 GiB, still fine
bfloat16 2 4185 - 1,120,910,400 2.09 OK out=(2, 32, 4185, 128) sum=-18654.8945
```
The same 11 cases, including the four that kill the XPU, all complete on an A100. (The `sum=` values are not comparable between the two tables — the inputs are generated on-device, so the two platforms draw different random numbers. Only OK-vs-failure matters here.)
### What the table shows
- **The boundary is exactly 2^30 elements**: `seq_len=4088` (0.4 % below) passes, `seq_len=4104` (0.4 % above) hangs.
- **Elements, not bytes**: fp16 at 4.50 GiB works while fp32 at 4.02 GiB hangs. Halving the batch at `seq_len=4185` — same per-sample shape and dtype — drops the count below 2^30 and it passes.
- **Not the mask, not GQA**: the `no-mask` and `no-gqa` rows hang identically.
- **Not a bad card**: reproduced on 5 different physical cards.
- **Not OOM**: the cards have 25.6 GB; a real OOM raises `UR_RESULT_ERROR_OUT_OF_RESOURCES` cleanly without resetting the engine.
## Where this came from
`transformers` `tests/models/emu3/test_modeling_emu3.py::Emu3IntegrationTest::test_model_generation_batched` — `BAAI/Emu3-Chat-hf` ships `dtype: float32`, and with a batch of 2 and a 4185-token prompt the first attention layer lands at 1.12e9 elements and hangs.
Note this test was baselined on Intel XPU CI in April 2025 ([transformers#37662](https://github.com/huggingface/transformers/pull/37662)) at the same fp32 4185-token shape, so this may be a regression or specific to Battlemage rather than a long-standing XPU limitation.
## Two things that make this easy to misdiagnose
- **The reported error is misleading.** Once the device is lost, rendering a traceback calls `repr()` on every tensor, filling the output with dozens of `UR_RESULT_ERROR_OUT_OF_RESOURCES` (40) messages that look exactly like OOM. The real error is the first one, `DEVICE_LOST` (20).
- **The failure survives the process.** A later run can fail much earlier — e.g. during `from_pretrained` at `tensor.to(device)` — looking like an unrelated bug. Anything measured after a hang is untrustworthy until the device recovers.
## Repro script
```python
"""fp32 SDPA hangs the GPU once attention scores exceed 2**30 elements.
python sdpa_fp32_hang_repro.py
Device is detected at runtime (CUDA or XPU). No arguments.
Each case runs in a fresh subprocess on purpose: a hang raises
UR_RESULT_ERROR_DEVICE_LOST and kills the device for the whole process, so
every later case in that process would fail too and the table would be
meaningless.
"""
import os
import subprocess
import sys
import torch
import torch.nn.functional as F
Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128
# Must come before torch.xpu.is_available(): calling is_available() first makes
# the process segfault at exit on this build, long after the compute is done.
if hasattr(torch, "xpu"):
try:
torch.xpu.init()
except Exception:
pass
if torch.cuda.is_available():
DEV, NAME, SYNC = torch.device("cuda"), torch.cuda.get_device_name(0), torch.cuda.synchronize
elif torch.xpu.is_available():
DEV, NAME, SYNC = torch.device("xpu"), torch.xpu.get_device_name(0), torch.xpu.synchronize
else:
sys.exit("no CUDA or XPU device available")
# (dtype, seq_len, batch, mask, gqa, note)
CASES = [
("float32", 3072, 2, True, True, ""),
("float32", 3584, 2, True, True, ""),
("float32", 4088, 2, True, True, "just below 2**30"),
("float32", 4104, 2, True, True, "just above 2**30"),
("float32", 4185, 2, True, True, "shape taken from a real model"),
("float32", 4185, 1, True, True, "half the batch -> below 2**30"),
("float32", 4185, 2, False, True, "mask is not the trigger"),
("float32", 4185, 2, True, False, "GQA is not the trigger"),
("float16", 4185, 2, True, True, ""),
("float16", 6144, 2, True, True, "2.4e9 elements / 4.5 GiB, still fine"),
("bfloat16", 4185, 2, True, True, ""),
]
def run_case(dtype_name, n, b, use_mask, use_gqa):
dtype = getattr(torch, dtype_name)
torch.manual_seed(0)
query = torch.randn(b, Q_HEADS, n, HEAD_DIM, dtype=dtype, device=DEV)
key = torch.randn(b, KV_HEADS, n, HEAD_DIM, dtype=dtype, device=DEV)
value = torch.randn(b, KV_HEADS, n, HEAD_DIM, dtype=dtype, device=DEV)
kwargs = {}
if use_gqa:
kwargs["enable_gqa"] = True
else: # expand KV by hand so enable_gqa is never involved
key = key.repeat_interleave(Q_HEADS // KV_HEADS, dim=1)
value = value.repeat_interleave(Q_HEADS // KV_HEADS, dim=1)
mask = torch.ones(b, 1, n, n, dtype=torch.bool, device=DEV).tril() if use_mask else None
out = F.scaled_dot_product_attention(query, key, value, attn_mask=mask, **kwargs)
SYNC()
print(f"OK out={tuple(out.shape)} sum={out.float().sum().item():.4f}")
def main():
print(f"torch {torch.__version__} | {DEV.type} | {NAME}\n")
header = f"{'dtype':9} {'batch':>5} {'seq_len':>7} {'variant':10} {'elements':>13} {'GiB':>6} result"
print(header)
print("-" * len(header))
for i, (dtype_name, n, b, use_mask, use_gqa, note) in enumerate(CASES):
elements = b * Q_HEADS * n * n
gib = elements * (4 if dtype_name == "float32" else 2) / 2**30
variant = "-" if use_mask and use_gqa else ("no-mask" if not use_mask else "no-gqa")
proc = subprocess.run(
[sys.executable, __file__],
capture_output=True,
text=True,
env={**os.environ, "SDPA_REPRO_CASE": str(i)},
)
ok = [ln for ln in proc.stdout.splitlines() if ln.startswith("OK")]
if ok:
# Trust the printed result: a crash *after* it is a teardown bug, not this one.
result = ok[-1] + ("" if proc.returncode == 0 else f" (exit {proc.returncode} at teardown)")
else:
lines = [ln for ln in proc.stderr.strip().splitlines() if ln.strip()]
result = lines[-1] if lines else f"exit {proc.returncode}"
print(f"{dtype_name:9} {b:>5} {n:>7} {variant:10} {elements:>13,} {gib:>6.2f} {result}"
+ (f" <- {note}" if note else ""), flush=True) # fmt: skip
if __name__ == "__main__":
case = os.environ.get("SDPA_REPRO_CASE")
if case is None:
main()
else:
dtype_name, n, b, use_mask, use_gqa, _ = CASES[int(case)]
run_case(dtype_name, n, b, use_mask, use_gqa)
```
Contributor guide
Research direction
Start with the embedded sdpa_fp32_hang_repro.py and the F.scaled_dot_product_attention entry point; run it on Intel XPU, focusing on the fp32 cases just below and above 2**30. Trace the XPU fp32/materializing SDPA path to identify where the failure occurs. Done means the affected cases complete without DEVICE_LOST while the CUDA reference and fp16/bf16 behavior remain valid.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100