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.

[luoyu-intel]

@copilot Updated the regression test to use deterministic inputs. The replay input is derived from the eager input with a fixed offset (replay_input = eager_input + 0.75), avoiding unseeded randomness while ensuring the two replay cases are observably different.

[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.