Decoupling oneDNN by Zhenzhong1 · Pull Request #2049 · intel/auto-round

Zhenzhong1 · 2026-07-13T07:00:34Z

Description

Decoupling oneDNN

Changes:

Adds ARK_DNNL CMake gating and conditional compilation to allow XPU builds without oneDNN when SYCL-TLA is enabled.
Introduces DeviceMemoryPool and refactors multiple call sites to use it instead of DnnlContext scratch allocations.
Extends SYCL-TLA dense GEMM and Python wrappers/tests to support FP32 (plus adds an XPU INT8-weight GEMM wrapper via SYCL).

Progress

Moffied CmakeLists
CPU build: ARK_DNNL=ON by default
XPU build: ARK_XPU & ARK_DNNL=OFF & ARK_SYCL_TLA=ON by default.
Support FP32 sycl_tla_dense_gemm
Decoupling oneDNN
replace joint_maxtrix

Tests

pytest -vvs ./test/test_matmul.py

pytest -vvs ./test/test_weightonly.py

Perf

test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt0-4096-4096-1] 
  Max Diff joint_matrix vs sycl_tla: 0.00000, Mean Diff: 0.00000
  [joint_matrix] :   0.0517 ms    0.649 TFLOPS
  [sycl_tla]     :   0.0380 ms    0.883 TFLOPS  speedup= 1.36x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt0-4096-4096-8] 
  Max Diff joint_matrix vs sycl_tla: 0.00006, Mean Diff: 0.00000
  [joint_matrix] :   0.0525 ms    5.114 TFLOPS
  [sycl_tla]     :   0.0381 ms    7.046 TFLOPS  speedup= 1.38x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt0-4096-4096-32] 
  Max Diff joint_matrix vs sycl_tla: 0.01562, Mean Diff: 0.00000
  [joint_matrix] :   0.0601 ms   17.870 TFLOPS
  [sycl_tla]     :   0.0462 ms   23.249 TFLOPS  speedup= 1.30x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt0-4096-4096-128] 
  Max Diff joint_matrix vs sycl_tla: 0.01562, Mean Diff: 0.00000
  [joint_matrix] :   0.0596 ms   72.122 TFLOPS
  [sycl_tla]     :   0.0590 ms   72.789 TFLOPS  speedup= 1.01x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt0-4096-4096-1024] 
  Max Diff joint_matrix vs sycl_tla: 0.01562, Mean Diff: 0.00000
  [joint_matrix] :   0.3000 ms  114.550 TFLOPS
  [sycl_tla]     :   0.3150 ms  109.088 TFLOPS  speedup= 0.95x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt0-4096-4096-2048] 
  Max Diff joint_matrix vs sycl_tla: 0.01562, Mean Diff: 0.00000
  [joint_matrix] :   0.5345 ms  128.563 TFLOPS
  [sycl_tla]     :   0.6101 ms  112.643 TFLOPS  speedup= 0.88x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt0-4096-4096-3072] 
  Max Diff joint_matrix vs sycl_tla: 0.03125, Mean Diff: 0.00000
  [joint_matrix] :   0.7707 ms  133.756 TFLOPS
  [sycl_tla]     :   0.8763 ms  117.631 TFLOPS  speedup= 0.88x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt0-4096-4096-4096] 
  Max Diff joint_matrix vs sycl_tla: 0.03125, Mean Diff: 0.00000
  [joint_matrix] :   1.0118 ms  135.836 TFLOPS
  [sycl_tla]     :   1.1525 ms  119.251 TFLOPS  speedup= 0.88x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt1-4096-4096-1] 
  Max Diff joint_matrix vs sycl_tla: 0.00000, Mean Diff: 0.00000
  [joint_matrix] :   0.0497 ms    0.676 TFLOPS
  [sycl_tla]     :   0.0247 ms    1.359 TFLOPS  speedup= 2.01x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt1-4096-4096-8] 
  Max Diff joint_matrix vs sycl_tla: 0.00001, Mean Diff: 0.00000
  [joint_matrix] :   0.0519 ms    5.175 TFLOPS
  [sycl_tla]     :   0.0259 ms   10.354 TFLOPS  speedup= 2.00x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt1-4096-4096-32] 
  Max Diff joint_matrix vs sycl_tla: 0.00001, Mean Diff: 0.00000
  [joint_matrix] :   0.0543 ms   19.768 TFLOPS
  [sycl_tla]     :   0.0433 ms   24.820 TFLOPS  speedup= 1.26x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt1-4096-4096-128] 
  Max Diff joint_matrix vs sycl_tla: 0.00001, Mean Diff: 0.00000
  [joint_matrix] :   0.0571 ms   75.225 TFLOPS
  [sycl_tla]     :   0.0574 ms   74.770 TFLOPS  speedup= 0.99x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt1-4096-4096-1024] 
  Max Diff joint_matrix vs sycl_tla: 0.00001, Mean Diff: 0.00000
  [joint_matrix] :   0.3092 ms  111.126 TFLOPS
  [sycl_tla]     :   0.3275 ms  104.924 TFLOPS  speedup= 0.94x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt1-4096-4096-2048] 
  Max Diff joint_matrix vs sycl_tla: 0.00001, Mean Diff: 0.00000
  [joint_matrix] :   0.5637 ms  121.910 TFLOPS
  [sycl_tla]     :   0.6587 ms  104.328 TFLOPS  speedup= 0.86x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt1-4096-4096-3072] 
  Max Diff joint_matrix vs sycl_tla: 0.00001, Mean Diff: 0.00000
  [joint_matrix] :   0.8248 ms  124.967 TFLOPS
  [sycl_tla]     :   0.9621 ms  107.143 TFLOPS  speedup= 0.86x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt1-4096-4096-4096] 
  Max Diff joint_matrix vs sycl_tla: 0.00001, Mean Diff: 0.00000
  [joint_matrix] :   1.0765 ms  127.672 TFLOPS
  [sycl_tla]     :   1.2840 ms  107.036 TFLOPS  speedup= 0.84x
PASSED

Usage

ARK_DNNL=1 python setup.py bdist_wheel

Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com>

AutoRoundBot · 2026-07-13T09:28:38Z

/azp run Unit-Test-CUDA-AutoRound

azure-pipelines · 2026-07-13T09:28:46Z

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

Copilot

Pull request overview

This PR decouples oneDNN from the ARK kernel build (especially for XPU), introduces a shared scratch-memory pool, and extends SYCL-TLA dense GEMM to accept FP32 inputs (along with corresponding Python/test updates).

Changes:

Adds ARK_DNNL CMake gating and conditional compilation to allow XPU builds without oneDNN when SYCL-TLA is enabled.
Introduces DeviceMemoryPool and refactors multiple call sites to use it instead of DnnlContext scratch allocations.
Extends SYCL-TLA dense GEMM and Python wrappers/tests to support FP32 (plus adds an XPU INT8-weight GEMM wrapper via SYCL).

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
auto_round_extension/ark/auto_round_kernel/CMakeLists.txt	Adds `ARK_DNNL` option + build-time constraints for XPU backend selection and conditional oneDNN FetchContent/linking.
auto_round_extension/ark/setup.py	Adds env flag parsing and passes `ARK_DNNL`/`ARK_SYCL_TLA` toggles into the XPU CMake build.
auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp	Adds `DeviceMemoryPool` and refactors `DnnlContext` scratch allocation to use it.
auto_round_extension/ark/auto_round_kernel/wrapper/include/xpu_wrapper.hpp	Makes oneDNN/SYCL-TLA includes conditional and routes scratch allocations + S8 path through new wrappers/pool.
auto_round_extension/ark/auto_round_kernel/wrapper/include/cpu_wrapper.hpp	Switches scratch allocation to `DeviceMemoryPool`.
auto_round_extension/ark/auto_round_kernel/sdpa.cpp	Switches varlen workspace scratch allocation to `DeviceMemoryPool`.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_dense_gemm.hpp	Adds FP32 support to SYCL-TLA dense GEMM and updates error messaging.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp	New SYCL-based S8 wrapper used by XPU paths.
auto_round_extension/ark/auto_round_kernel/ark.cpp	Adds conditional XPU backend dispatch (oneDNN vs SYCL-TLA) and uses `SyclS8Wrapper` for XPU S8.
auto_round_extension/ark/auto_round_kernel/init.py	Updates `matmul_sycl_tla` Python API dtype validation to include FP32.
auto_round_extension/ark/test/test_matmul.py	Expands SYCL-TLA test parameterization to include FP32 and adjusts backend-compare tolerances.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

AutoRoundBot · 2026-07-14T07:51:21Z

/azp run Unit-Test-CUDA-AutoRound

azure-pipelines · 2026-07-14T07:51:28Z

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

Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com>

AutoRoundBot · 2026-07-16T09:43:50Z

/azp run Unit-Test-CUDA-AutoRound

azure-pipelines · 2026-07-16T09:43:57Z

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

Zhenzhong1 added 3 commits July 13, 2026 11:45

support fp32 in sycltla

b827f56

Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com>

update CMakeLists & setup.py

14bbb03

Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com>

decompling oneDNNL

03f0307

Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com>

Zhenzhong1 marked this pull request as ready for review July 13, 2026 08:34

Zhenzhong1 requested a review from Copilot July 14, 2026 03:02

Copilot started reviewing on behalf of Zhenzhong1 July 14, 2026 03:02 View session

Copilot AI reviewed Jul 14, 2026

View reviewed changes

Comment thread auto_round_extension/ark/setup.py Outdated

Comment thread auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp

Zhenzhong1 and others added 2 commits July 14, 2026 15:02

modify the enable_dnnl variable

c6076c3

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

add license

0f86148

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Zhenzhong1 added 2 commits July 16, 2026 16:16

sycltla igemm_s8s8 inference pass

d13c291

Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com>

add igemm_sycltla&joint_matrix

bc7476e

Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com>

Zhenzhong1 force-pushed the zhenzhong/sycltla-gemm-2 branch from e1c638b to bc7476e Compare July 16, 2026 08:30

perf ipmprove

f18a732

Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Decoupling oneDNN#2049

Decoupling oneDNN#2049
Zhenzhong1 wants to merge 8 commits into
mainfrom
zhenzhong/sycltla-gemm-2

Zhenzhong1 commented Jul 13, 2026 •

edited

Loading

Uh oh!

AutoRoundBot commented Jul 13, 2026

Uh oh!

azure-pipelines Bot commented Jul 13, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

AutoRoundBot commented Jul 14, 2026

Uh oh!

azure-pipelines Bot commented Jul 14, 2026

Uh oh!

AutoRoundBot commented Jul 16, 2026

Uh oh!

azure-pipelines Bot commented Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

Zhenzhong1 commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Changes:

Progress

Tests

Perf

Usage

Uh oh!

AutoRoundBot commented Jul 13, 2026

Uh oh!

azure-pipelines Bot commented Jul 13, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

AutoRoundBot commented Jul 14, 2026

Uh oh!

azure-pipelines Bot commented Jul 14, 2026

Uh oh!

AutoRoundBot commented Jul 16, 2026

Uh oh!

azure-pipelines Bot commented Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Zhenzhong1 commented Jul 13, 2026 •

edited

Loading