torch.topk on XPU returns invalid (-1) indices with NaN input at batch size >= ~512
- Dominant language
- Python
- Stars
- 113
- Forks
- 128
- Avg merge
- 5d 9h
- Merged PRs (30d)
- 112
Description
## Summary
`torch.topk` on XPU returns out-of-range indices (`-1`) when the input contains a NaN row and the batch dimension is >= ~512. CPU on the identical seed/input returns valid indices. This is not an alternate/valid NaN-ordering convention — `-1` is never a valid `topk` index, and downstream `gather()` on that index aborts the process via a device-side bounds assert rather than raising a catchable exception.
Likely related to intel/torch-xpu-ops#3777 ("[Triton] torch.topk returns wrong values on XPU", closed without a linked fix) — same op, overlapping version window — but that repro is `[8,8]` float32 with no NaN, so we haven't confirmed it's the identical underlying defect.
## Repro
```python
import torch
torch.manual_seed(0)
go_cpu = torch.randn(512, 32, dtype=torch.bfloat16)
go_cpu[0] = float("nan")
values_cpu, ids_cpu = torch.topk(go_cpu, k=4, dim=-1, sorted=False)
print("CPU row0 ids: ", ids_cpu[0].tolist())
go_xpu = go_cpu.to("xpu")
values_xpu, ids_xpu = torch.topk(go_xpu, k=4, dim=-1, sorted=False)
print("XPU row0 ids: ", ids_xpu[0].tolist())
bad = (ids_xpu < 0) | (ids_xpu >= 32)
print("XPU bad_count: ", bad.sum().item())
try:
go_cpu.gather(1, ids_xpu.cpu())
print("CPU gather with XPU's ids: succeeded (unexpected)")
except (IndexError, RuntimeError) as e:
print("CPU gather with XPU's ids: raised", type(e).__name__, "->", e)
# On XPU, the same gather aborts the process with a device_assert
# (SIGABRT) instead of raising a catchable exception.
go_xpu.gather(1, ids_xpu)
```
## Observed output
```
CPU row0 ids: [22, 21, 20, 23]
XPU row0 ids: [0, -1, -1, -1]
XPU bad_count: 3
CPU gather with XPU's ids: raised RuntimeError -> index -1 is out of bounds for dimension 1 with size 32
AssertHandler::printMessage
.../torch-xpu-ops/src/ATen/native/xpu/sycl/ScatterGatherKernels.cpp:233: operator(): global id: [1,0,0], local id: [1,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size_ && "scatter gather kernel index out of bounds"` failed.
```
Process exits with SIGABRT (134).
## Trigger conditions (from further isolation, happy to share the sweep script)
- Requires a NaN row in the input — fully finite inputs (including degenerate all-zero/all-equal ties, and +-Inf) never reproduce, at any batch size tested.
- Requires batch size >= ~512 — fails at 512/1024/2048/4096/8192, passes at 16/64/128/256.
- Reproduces with `sorted=False`; not yet swept for `sorted=True`.
## Environment
- `torch==2.13.0+xpu`
- `triton-xpu==3.7.2`
- Intel Arc Pro B70 (30.3GiB), driver `intel-opencl-icd 26.18.38308.1`
- Does **not** reproduce on `torch==2.12.0+xpu` (older stack, otherwise same environment) — bisected by dependency bump alone, no application-level code changed across the bisect.
## Downstream impact
This corrupts MoE expert routing in vLLM (`grouped_topk()`, used by DeepSeek-V2/V3-style and LiquidAI LFM2 gates) when the router's gating output happens to contain a NaN — the corrupted index is then consumed by a fused Triton kernel doing `gather`, which asserts one step later than the plain-`gather` repro above shows, but from the same root cause. We understand from internal driver-team follow-up that a fix (a loop-strength-reduction wrap-around bug in IGC's shader codegen) has landed internally but is not yet in a public/GA driver release. We have not been able to test our exact repro (bf16, NaN-triggered) against the fixed driver build to confirm it resolves this trigger.
cc @astachowiczhabana
Contributor guide
Assessment
This issue has not been assessed yet.