XPU oneDNN: bind matmul execution to live SYCL queue (graph-capture safe)

[Copilot]

Description

Issue #2206 reports that XPU oneDNN stream caching was keyed too coarsely (device UUID), so later calls could execute on a stale queue instead of the active torch.xpu graph-capture queue; m > 1 WOQ GEMM then escaped capture and replay produced frozen/zero outputs.
This PR makes oneDNN execution queue-correct while preserving engine reuse.

• Root cause / scope

  • DnnlContext cached dnnl::stream per device; stream stayed bound to first queue seen.
  • woq_gemm (m > 1, fp16/fp32 oneDNN path) executed via that stale stream.

• Queue-correct execution model

  • Reworked DnnlContext:
    • cache engine by (device, sycl::context) (reused for performance),
    • construct stream from the current call queue (live queue binding).
  • Scratch buffer keying aligned to context-aware identity to avoid cross-context reuse.

• Call-chain cleanup

  • Removed implicit GETCTX() usage from active oneDNN GEMM/quant paths.
  • gemm/dyn_quant_s8/igemm_s8s8 now explicitly derive engine+stream from passed q.

• Regression coverage

  • Added XPU test: eager warmup on one queue, graph capture on another queue, m > 1, replay with changed inputs; replay output must track eager reference (not frozen).

auto& eng = *DnnlContext::Instance()->get_eng(q);
auto stream = DnnlContext::Instance()->get_stream(q);  // bound to live queue
matmul_prim.execute(stream, matmul_args);

Type of Change

Bug fix

Related Issues

Tracked in repository issue tracker (auto-linking handled by system).

Checklist Before Submitting

• My code has been tested locally.
• Documentation has been updated as needed.
• New or updated tests are included where applicable.
• The CUDA CI has passed. You can trigger it by commenting /azp run Unit-Test-CUDA-AutoRound.

[Copilot]

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Pull request overview

Fixes XPU oneDNN graph-capture correctness by ensuring oneDNN execution is bound to the live SYCL queue (avoiding stale cached streams) while keeping engine reuse, and adds a regression test to prevent frozen outputs on replay.

Changes:

• Reworked DnnlContext to cache engines per (device, sycl::context) and create streams from the current queue.
• Aligned scratch-buffer keying to be context-aware to avoid cross-context reuse.
• Added an XPU regression test covering warmup-on-one-queue + capture-on-another + replay-with-new-inputs.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
auto_round_extension/ark/test/test_weightonly.py	Adds a regression test ensuring WOQ GEMM graph capture uses the live queue and replay updates outputs.
auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp	Refactors oneDNN engine/stream management to be queue-correct and context-aware, plus scratch-key changes.
auto_round_extension/ark/auto_round_kernel/wrapper/include/dnnl_wrapper.hpp	Removes GETCTX() macro and explicitly derives engine + live stream from the passed queue.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

[luoyu-intel]

@copilot
Request changes: the (device UUID, SYCL context) key correctly prevents scratch-buffer reuse across different SYCL contexts, and the live-queue oneDNN stream fix addresses the original stale-stream graph-capture bug. However, this is still insufficient for vLLM's execution model.

A vLLM worker can use multiple SYCL queues concurrently while those queues belong to the same SYCL context. With the current design, those queues still share the same (context, buf_loc) scratch allocation. If queue A has an in-flight oneDNN operation using the buffer, queue B can reuse it, or enter the growth path and release it here:

release(it->second, q);
auto newptr = allocate(size, q);

This can cause cross-stream data corruption or use-after-free.

Please make scratch allocation lifetime-safe before merging. The preferred options are:

1. Track the last-use sycl::event for each allocation and defer reuse/free until that event completes; or
2. Allocate scratch independently per queue/in-flight operation, with a pool-level mutex protecting the pointer/size maps.

If the queue-aware approach is selected, the scratch key should include a stable queue identity in addition to (device, context, buf_loc). If event-based reuse is selected, the API should return a scratch lease/allocation and provide a way for each caller to associate the submitted oneDNN event with that allocation.

Please also add a regression test with two queues from the same SYCL context and overlapping asynchronous oneDNN operations, including the buffer-growth path. The current test validates cross-context isolation and live-queue graph capture, but does not cover same-context multi-queue scratch reuse or in-flight release.

[Copilot]

@copilot
Request changes: the (device UUID, SYCL context) key correctly prevents scratch-buffer reuse across different SYCL contexts, and the live-queue oneDNN stream fix addresses th...

Addressed in fa4b7ec: I added a pool-level mutex around scratch pointer/size maps and strengthened the same-context two-queue regression to exercise overlapping async submissions plus the growth path (pre-sized small on one queue, then larger request), while keeping the live-queue graph-capture checks.

[luoyu-intel]

Thanks for addressing the review feedback.

The latest changes correctly address the main correctness issues:

• oneDNN engines are reused per (device, SYCL context);
• oneDNN streams are created from the live call queue;
• scratch-buffer identity is now queue-aware, preventing reuse across different queues in the same SYCL context;
• scratch map accesses are protected by a mutex;
• the growth path waits before releasing an existing allocation;
• regression coverage now includes same-context multi-queue scratch isolation and graph-capture behavior.

This is important for vLLM, where a single worker may use multiple SYCL queues concurrently.

A few non-blocking follow-ups remain:

1. The scratch-pool mutex is held while waiting on q->wait() during buffer growth. This guarantees correctness, but can serialize unrelated scratch allocations and introduce host-side latency. A future improvement could detach the old allocation under the lock, release the lock, and then wait/free it, or use event-based deferred reclamation.

2. Queue and context identities are retained by process-lifetime singletons. This is acceptable for the usual vLLM worker lifetime, but repeated creation of queues or contexts could grow the cached metadata and retain SYCL resources longer than necessary. An explicit teardown or cleanup policy would improve long-running or dynamic-runtime scenarios.

3. The queue-identity implementation should continue to be validated against the supported oneAPI/SYCL runtimes, especially the equality semantics used to distinguish queues. The current regression test is useful and should be retained.

4. It would be valuable to run an end-to-end vLLM XPU test covering graph capture/replay together with concurrent same-context streams, rather than only testing the lower-level debug probes.

These are follow-up improvements rather than blockers for the original stale oneDNN stream issue. With the current queue-aware scratch management and regression coverage, the PR now addresses the reported graph-capture correctness problem.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.