intel / intel/auto-round

[Bug]: woq_gemm's m > 1 path submits its oneDNN matmul to a cached stream bound to the first queue ever seen, so it is never recorded into a torch.xpu graph and every replay returns zeros

Open
#2,206 5 comments 0 reactions 1 assignee Claimed by @luoyu-intel View on GitHub
Dominant language
Python
Stars
1.6k
Forks
175
Avg merge
1d 18h
Merged PRs (30d)
99

Description

### Environment

auto-round-lib 0.14.2 from PyPI, as shipped inside `vllm/vllm-openai-xpu:nightly` (image id c48edf76bb9f). torch 2.13.0+xpu. Intel Arc Pro B70 (BMG-G31), host kernel 7.1.0-070100, NEO 26.27.39122.11. `icpx` is not on PATH in that image.

That last detail decides which code path runs. `fallback_compute_type_if_needed` (`auto_round_kernel/utils.py:54-60`) downgrades int8 to fp16 only when `is_b70()` and not `is_oneapi_ge_2026()`, and `is_oneapi_ge_2026` (`utils.py:25-37`) shells out to `icpx --version` and returns False when icpx is absent. So on this image `post_init` resolves `cdt/wdt/sdt = fp16 int4 fp16`, and the reproducer below prints that so nobody has to guess. On a B70 with oneAPI 2026 on PATH you would get `cdt = "int8"` and exercise a different branch. See the scope note at the end.

### Bug

Captured into a `torch.xpu` graph, `ark.woqgemm` behaves two different ways depending on m.

At m = 1 it is recorded correctly. Replay recomputes from the current input and matches eager bit for bit on every replay.

At m > 1 the matmul is not recorded at all. Every replay leaves the output at exactly zero, in every element, unchanged no matter what the input is. Tested at m = 2, 3, 4, 8 and 16.

Two observations locate it. The output tensor already holds the correct answer the instant the capture region closes at m > 1, so the work ran during capture instead of being recorded. At m = 1 the same tensor is still zero at that point, which is what recording without executing looks like. On replay the two invert.

### Mechanism

The oneDNN stream is cached per device and bound permanently to the first queue that device was ever seen with.

`DnnlContext::check_dnnl_device` (`auto_round_kernel/wrapper/include/utils.hpp:124-145`) keys `dev_engine_map` and `dev_stream_map` on a hash of the **device UUID** (`utils.hpp:134-135`), not on the queue. It constructs the stream only on first sight of that device, at `utils.hpp:141`:

```cpp
dev_stream_map[key] = dnnl::sycl_interop::make_stream(dev_engine_map[key], *q);
```

Every later call for the same device returns that stream regardless of which queue was passed (`get_stream`, `utils.hpp:119-122`).

The Python wrapper does the right thing: `__init__.py:105` passes the live queue, `torch.xpu.current_stream().sycl_queue`. Inside the m > 1 fallback (`xpu_wrapper.hpp:697-706`), `unpackq` at `:705` is submitted to that queue and is recorded into the graph. But `DnnlWrapper::gemm` at `:706` reaches the stream through `GETCTX()` (`dnnl_wrapper.hpp:17-19`, used at `:30`) and submits with it at `dnnl_wrapper.hpp:67`, `matmul_prim.execute(stream, matmul_args)`. That stream is the stale cached one. During capture it is not the capture queue, so the matmul goes to a normally-executing queue and never enters the graph.

We tested that directly rather than leaving it as a reading of the source. The cache binds on first sight, so if the first call for a device happens *inside* the capture region, the cached stream should be the capture queue and the matmul should be recorded. Two fresh processes, identical except for whether an eager warmup runs before capture, both at m = 4:

```
warm (warmup first, then capture) replay 0.000000, 0/16384 nonzero, frozen, wrong
cold (first call inside capture) replay matches eager exactly (diff 0.000000),
responds to input, 16384/16384 nonzero
```

Moving which queue the device is first seen with is enough to fix it, which is what the cache being keyed on the device rather than the queue predicts.

One prediction did not hold, noted so nobody has to rediscover it: we expected eager to break in the cold case, since eager would then submit through the capture queue. It does not. A capture queue is still an ordinary queue once capture ends, and `torch.xpu.synchronize()` is device-wide, so eager stays correct. That does not affect the result above.

That accounts for the exact-zero result rather than merely a wrong one. The output buffer is allocated as `C = torch.zeros(m, n, dtype=A.dtype, device=A.device)` at `__init__.py:341`. The graph records the zero-fill and the unpack, and nothing that writes C. Replay reproduces exactly that.

It also accounts for the boundary. `woq_gemv` (`xpu_wrapper.hpp:516`) declines anything but m = 1 with `if (m > 1) return -2;` at `:518`, and on the m = 1 path everything is submitted through `q` directly and oneDNN is never touched.

### Reproducer

No vLLM, no checkpoint, synthetic weights. There is no index buffer to get wrong: a sym int4 `QuantLinearGPTQ` registers exactly `qweight`, `qzeros`, `scales` and `bias` (`qlinear.py:139-166`), there is no `g_idx` and no desc_act. Random int32 in `qweight` is a valid arbitrary 4-bit weight, since `unpack_to_8bit_signed` masks to nibbles and `qlinear.py:205` maps them to [-8, 7]. With `sym=True` the zeros are discarded at `qlinear.py:197-198` and never reach the kernel. The blob-size guard at `xpu_wrapper.hpp:690-696` does not fire.

```python
import shutil, torch
from auto_round_kernel.qlinear import QuantLinearGPTQ

DEV, K, N, G = 'xpu', 4096, 4096, 128
torch.manual_seed(0)

lin = QuantLinearGPTQ(bits=4, group_size=G, sym=True, in_features=K, out_features=N,
bias=True, weight_dtype=torch.float16)
for _, buf in lin.named_buffers():
if buf.dtype == torch.int32:
buf.copy_(torch.randint(-2**31, 2**31 - 1, buf.shape, dtype=torch.int64).to(torch.int32))
else:
buf.copy_((torch.rand(buf.shape, dtype=torch.float32) * 0.01).to(buf.dtype))
lin = lin.to(DEV).eval()
lin.post_init()
print(f'cdt={lin.cdt} wdt={lin.wdt} sdt={lin.sdt} icpx={shutil.which("icpx") is not None}')

def case(m, replays=6):
x = torch.randn(m, K, device=DEV, dtype=torch.float16) * 0.1
s = torch.xpu.Stream()
with torch.xpu.stream(s):
for _ in range(3):
lin(x)
torch.xpu.synchronize()

a, b = lin(x).float().clone(), lin(x).float().clone() # eager determinism control
torch.xpu.synchronize()
print(f'm={m} eager spread {(a - b).abs().max().item():.6f}')

g = torch.xpu.XPUGraph()
with torch.xpu.graph(g):
y = lin(x)
torch.xpu.synchronize()
print(f'm={m} |y| right after capture {y.float().abs().max().item():.6f}')

for _ in range(replays):
x.copy_(torch.randn(m, K, device=DEV, dtype=torch.float16) * 0.1)
g.replay()
torch.xpu.synchronize()
got = y.float().clone()
ref = lin(x).float()
torch.xpu.synchronize()
print(f' m={m} |replay| {got.abs().max().item():.6f} '
f'nonzero {int((got != 0).sum())}/{got.numel()} '
f'rel {((got - ref).abs().max() / ref.abs().max()).item():.4f}')

with torch.no_grad():
for m in (1, 2, 3, 4, 8, 16):
case(m)
```

Output on the B70, first replay of each m:

```
cdt=fp16 wdt=int4 sdt=fp16 icpx=False

m=1 eager spread 0.000000 |y| right after capture 0.000000
m=1 |replay| 0.687012 nonzero 4096/4096 rel 0.0000
m=2 eager spread 0.000000 |y| right after capture 0.850586
m=2 |replay| 0.000000 nonzero 0/8192 rel 1.0000
m=3 eager spread 0.000000 |y| right after capture 0.807617
m=3 |replay| 0.000000 nonzero 0/12288 rel 1.0000
m=4 eager spread 0.000000 |y| right after capture 0.785156
m=4 |replay| 0.000000 nonzero 0/16384 rel 1.0000
m=8 eager spread 0.000000 |y| right after capture 0.769043
m=8 |replay| 0.000000 nonzero 0/32768 rel 1.0000
m=16 eager spread 0.000000 |y| right after capture 0.703125
m=16 |replay| 0.000000 nonzero 0/65536 rel 1.0000
```

Eager is bit-deterministic at every m, so the comparison holds. All 6 replays behave the same at each m: m = 1 bit-exact, everything above it exactly zero and frozen.

### Why this reaches users

vLLM advertises ARK integration in this repo's own README ecosystem table (`INCXPUARKLinearMethod`). With an AutoRound int4 checkpoint it selects this kernel for every quantized linear, so any decode step carrying more than one token gets zeros out of the entire model and the output is garbage.

To be accurate about how often that fires: cudagraph capture is not on by default on XPU. vLLM gates it behind `VLLM_XPU_ENABLE_XPU_GRAPH=1` (`platforms/xpu.py:300`, which otherwise forces `cudagraph_mode` to NONE). So this affects users who have turned capture on, which is the configuration anyone chasing decode throughput ends up in, rather than every XPU user out of the box. Two ordinary ways to hit it: two concurrent requests, or any speculative decoding verify step. Single-request non-speculative decode is m = 1 throughout, which is why it looks fine until the server has more than one user. Filed on the vLLM side at https://github.com/vllm-project/vllm/issues/53211.

### Possible fixes

Key the stream cache on the queue rather than on the device UUID, or construct a stream per call from the queue that was passed in. Either makes the submission land where the caller asked.

Short of that, having `woq_gemm` detect an active capture and return an error, the way `woq_gemv` already declines m > 1 at `xpu_wrapper.hpp:518`, would at least turn silent wrong output into something a caller can handle.

### Scope, and what we did not test

We only measured the fp16 branch. Reading the source, the int8 branch looks like it should be fine: `xpu_wrapper.hpp:707-722` calls `DnnlWrapper::woq_s8` at `:721`, and that function (`dnnl_wrapper.hpp:229`) splits on `#if ARK_XPU`. The XPU side, `dnnl_wrapper.hpp:237-238`, uses `sycl_dyn_quant_s8` and `sycl_igemm_s8s8`, which submit to `q` and never touch `GETCTX()`. The `dyn_quant_s8` and `igemm_s8s8` pair that does use the cached stream is in the `#else`, so an XPU build should not reach it. We have not run the int8 path and are not claiming anything about it either way. We cannot: oneAPI 2026 is not in the published image and we have no build against it.

If that reading is right, this is narrower than it first appears. `fallback_compute_type_if_needed` selects fp16 only on a B70 without oneAPI 2026, so that specific combination may be the only one that reaches the affected path.

Everything cited above is from the shipped 0.14.2 sdist, which is what we tested. Current master has moved some of this: the scratch call is different, and the oneDNN gemm is gated behind `ARK_DNNL` against an `ARK_SYCL_TLA` alternative, so a master build with TLA may not reach `DnnlWrapper::gemm` at all. Worth re-checking there before assuming it is unfixed or already fixed.

The reproducer needs Battlemage hardware and torch 2.13+xpu. We have not tried it on A770 or B580; note that `is_b70()` is False there, so reaching the fp16 branch on those cards would require pinning `compute_type` through the low-level API rather than going through `QuantLinearGPTQ`.

### Related

There is no capture harness in the test suite: `test/test_weightonly.py` parametrizes m over 1 to 1024 but every case is eager, so nothing today would catch this. The reproducer above is standalone and needs no checkpoint, but it does need a capture harness that does not currently exist here, so it is not a drop-in CI addition.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.