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

dt0, dt1, dt2
[torch.float32, torch.float16, torch.bfloat16]

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.0493 ms    0.680 TFLOPS
  [sycl_tla]     :   0.0390 ms    0.860 TFLOPS  speedup= 1.26x
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.00001, Mean Diff: 0.00000
  [joint_matrix] :   0.0508 ms    5.284 TFLOPS
  [sycl_tla]     :   0.0390 ms    6.879 TFLOPS  speedup= 1.30x
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.00001, Mean Diff: 0.00000
  [joint_matrix] :   0.0561 ms   19.125 TFLOPS
  [sycl_tla]     :   0.0460 ms   23.364 TFLOPS  speedup= 1.22x
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.00001, Mean Diff: 0.00000
  [joint_matrix] :   0.0568 ms   75.619 TFLOPS
  [sycl_tla]     :   0.0599 ms   71.752 TFLOPS  speedup= 0.95x
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.00001, Mean Diff: 0.00000
  [joint_matrix] :   0.3098 ms  110.915 TFLOPS
  [sycl_tla]     :   0.3280 ms  104.749 TFLOPS  speedup= 0.94x
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.00001, Mean Diff: 0.00000
  [joint_matrix] :   0.5656 ms  121.498 TFLOPS
  [sycl_tla]     :   0.6608 ms  103.999 TFLOPS  speedup= 0.86x
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.00001, Mean Diff: 0.00000
  [joint_matrix] :   0.8254 ms  124.883 TFLOPS
  [sycl_tla]     :   0.9628 ms  107.064 TFLOPS  speedup= 0.86x
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.00001, Mean Diff: 0.00000
  [joint_matrix] :   1.0782 ms  127.473 TFLOPS
  [sycl_tla]     :   1.2749 ms  107.803 TFLOPS  speedup= 0.85x
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.0518 ms    0.648 TFLOPS
  [sycl_tla]     :   0.0388 ms    0.865 TFLOPS  speedup= 1.34x
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.00006, Mean Diff: 0.00000
  [joint_matrix] :   0.0538 ms    4.994 TFLOPS
  [sycl_tla]     :   0.0311 ms    8.635 TFLOPS  speedup= 1.73x
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.01562, Mean Diff: 0.00000
  [joint_matrix] :   0.0589 ms   18.220 TFLOPS
  [sycl_tla]     :   0.0435 ms   24.674 TFLOPS  speedup= 1.35x
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.01562, Mean Diff: 0.00000
  [joint_matrix] :   0.0598 ms   71.855 TFLOPS
  [sycl_tla]     :   0.0568 ms   75.560 TFLOPS  speedup= 1.05x
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.01562, Mean Diff: 0.00000
  [joint_matrix] :   0.3002 ms  114.467 TFLOPS
  [sycl_tla]     :   0.3152 ms  109.003 TFLOPS  speedup= 0.95x
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.01562, Mean Diff: 0.00000
  [joint_matrix] :   0.5356 ms  128.308 TFLOPS
  [sycl_tla]     :   0.6108 ms  112.504 TFLOPS  speedup= 0.88x
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.03125, Mean Diff: 0.00000
  [joint_matrix] :   0.7715 ms  133.615 TFLOPS
  [sycl_tla]     :   0.8746 ms  117.852 TFLOPS  speedup= 0.88x
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.03125, Mean Diff: 0.00000
  [joint_matrix] :   1.0122 ms  135.785 TFLOPS
  [sycl_tla]     :   1.1542 ms  119.081 TFLOPS  speedup= 0.88x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt2-4096-4096-1] 
  Max Diff joint_matrix vs sycl_tla: 0.25000, Mean Diff: 0.01068
  [joint_matrix] :   0.0501 ms    0.669 TFLOPS
  [sycl_tla]     :   0.0248 ms    1.353 TFLOPS  speedup= 2.02x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt2-4096-4096-8] 
  Max Diff joint_matrix vs sycl_tla: 0.25000, Mean Diff: 0.01068
  [joint_matrix] :   0.0526 ms    5.104 TFLOPS
  [sycl_tla]     :   0.0261 ms   10.265 TFLOPS  speedup= 2.01x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt2-4096-4096-32] 
  Max Diff joint_matrix vs sycl_tla: 0.25000, Mean Diff: 0.01062
  [joint_matrix] :   0.0558 ms   19.227 TFLOPS
  [sycl_tla]     :   0.0438 ms   24.513 TFLOPS  speedup= 1.27x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt2-4096-4096-128] 
  Max Diff joint_matrix vs sycl_tla: 0.25000, Mean Diff: 0.01074
  [joint_matrix] :   0.0582 ms   73.834 TFLOPS
  [sycl_tla]     :   0.0571 ms   75.171 TFLOPS  speedup= 1.02x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt2-4096-4096-1024] 
  Max Diff joint_matrix vs sycl_tla: 0.25000, Mean Diff: 0.01062
  [joint_matrix] :   0.2943 ms  116.746 TFLOPS
  [sycl_tla]     :   0.3180 ms  108.059 TFLOPS  speedup= 0.93x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt2-4096-4096-2048] 
  Max Diff joint_matrix vs sycl_tla: 0.25000, Mean Diff: 0.01068
  [joint_matrix] :   0.5261 ms  130.628 TFLOPS
  [sycl_tla]     :   0.6336 ms  108.451 TFLOPS  speedup= 0.83x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt2-4096-4096-3072] 
  Max Diff joint_matrix vs sycl_tla: 0.25000, Mean Diff: 0.01056
  [joint_matrix] :   0.7556 ms  136.413 TFLOPS
  [sycl_tla]     :   0.9064 ms  113.730 TFLOPS  speedup= 0.83x
PASSED
test/test_matmul.py::test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf[2000-dt2-4096-4096-4096] 
  Max Diff joint_matrix vs sycl_tla: 0.25000, Mean Diff: 0.01080
  [joint_matrix] :   0.9891 ms  138.948 TFLOPS
  [sycl_tla]     :   1.1950 ms  115.008 TFLOPS  speedup= 0.83x

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.

AutoRoundBot · 2026-07-21T03:38:15Z

/azp run Unit-Test-CUDA-AutoRound

azure-pipelines · 2026-07-21T03:38:21Z

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>

Merge branch 'main' into zhenzhong/sycltla-gemm-2

261d55a

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 9 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!

AutoRoundBot commented Jul 21, 2026

Uh oh!

azure-pipelines Bot commented Jul 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 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!

AutoRoundBot commented Jul 21, 2026

Uh oh!

azure-pipelines Bot commented Jul 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Zhenzhong1 commented Jul 13, 2026 •

edited

Loading