Skip to content

[ARK] Optimize SYCL-TLA woqgemm_s8 on B70 - #2133

Open
Zhenzhong1 wants to merge 14 commits into
mainfrom
zhenzhong/woqgemm_s8_update
Open

[ARK] Optimize SYCL-TLA woqgemm_s8 on B70#2133
Zhenzhong1 wants to merge 14 commits into
mainfrom
zhenzhong/woqgemm_s8_update

Conversation

@Zhenzhong1

@Zhenzhong1 Zhenzhong1 commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Summary

Optimize ARK XPU woqgemm_s8 with a SYCL-TLA S8 GEMM path for BMG/B70, and enable device-specific SYCL target selection during ARK extension build.

This PR adds a fused K-block S8 GEMM path for weight-only quantization, improves W2/W4-to-S8 rescale routing, and updates build target detection for B60/B70/PVC.

related PR: #2049

Changes

  • Add SYCL-TLA S8 WOQ GEMM K-block path.

  • Add host-side validation for S8 K-block launch.

    • blocksize must divide k.
    • K-block blocksize must be a multiple of 64, matching the fixed SYCL-TLA MMA K tile.
  • Refine XPU S8 compute routing.

    • W2/W4 may rescale to S8 when ARK_RESCALE is enabled.
    • W8A8 keeps the original scale and blocksize, and does not use W2/W4 rescale metadata.
    • ARK_AUTO_S8=-1 is the default and maps W2/W4 rescale to full-K S8.
    • ARK_AUTO_S8=0 disables W2/W4 rescale.
    • Positive ARK_AUTO_S8 values select a larger effective rescale block when valid.
  • Improve SYCL target setup.

    • Detect B70/G31, B60/G21, and PVC in setup.py.
    • Pass DPCPP_SYCL_TARGET into the XPU CMake build.
    • Map bmg to bmg-g21,bmg-g31 for AOT compilation.
  • Remove the old B70 oneAPI < 2026 runtime fallback that forced XPU int8 compute to fp16. B70 now uses the SYCL-TLA int8 path, and ARK/vLLM builds have moved to oneAPI 2026.

ARK_AUTO_S8 Behavior

ARK_AUTO_S8=-1   -> W2/W4 rescale blocksize = k, full-K S8 path
ARK_AUTO_S8=512  -> W2/W4 rescale blocksize = 512 when k % 512 == 0, K-block S8 path
ARK_AUTO_S8=0    -> disable W2/W4 rescale and use the original quantization path
## Validation

```bash
pytest -vvs auto_round_extension/ark/test/test_weightonly.py
442 passed, 24 skipped in 65.55s (0:01:05)

Perf

# CMD

python - <<'PY'
from test_weightonly import main_op
shapes = [
	(4096, 4096, 4096, 4096),
	(4096, 4096, 4096, 128),
	(4096, 4096, 4096, 64),
	(4096, 4096, 4096, 32),
]
print("=== W4A8 / A8W4 ===")
for m, n, k, block in shapes:
    main_op(m, n, k, block, "int8", "int4", "fp16", False, "xpu", False)
PY

SYCL-TLA on B70 ARK_AUTO_S8=-1

=== W4A8 / A8W4 ===

 m=4096, n=4096, k=4096, blocksize=4096, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.093750, Mean Diff: 0.009491, ref mean:7.507812, OUT mean:7.503906
[Performance] Time: 0.6834 ms
              GFLOPS: 201104.72
              Bandwidth: 110.48 GB/s

 m=4096, n=4096, k=4096, blocksize=128, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.070312, Mean Diff: 0.010101, ref mean:7.546875, OUT mean:7.546875
[Performance] Time: 0.6747 ms
              GFLOPS: 203716.00
              Bandwidth: 112.29 GB/s

 m=4096, n=4096, k=4096, blocksize=64, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.074219, Mean Diff: 0.009949, ref mean:7.546875, OUT mean:7.546875
[Performance] Time: 0.6856 ms
              GFLOPS: 200471.33
              Bandwidth: 110.89 GB/s

 m=4096, n=4096, k=4096, blocksize=32, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.078125, Mean Diff: 0.009880, ref mean:7.507812, OUT mean:7.507812
[Performance] Time: 0.6741 ms
              GFLOPS: 203884.94
              Bandwidth: 113.55 GB/s

SYCL-TLA on B70 ARK_AUTO_S8=0

=== W4A8 / A8W4 ===

 m=4096, n=4096, k=4096, blocksize=4096, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.085938, Mean Diff: 0.007950, ref mean:7.507812, OUT mean:7.507812
[Performance] Time: 0.6740 ms
              GFLOPS: 203901.36
              Bandwidth: 112.02 GB/s

 m=4096, n=4096, k=4096, blocksize=128, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.066406, Mean Diff: 0.008209, ref mean:7.546875, OUT mean:7.546875
[Performance] Time: 0.8733 ms
              GFLOPS: 157375.38
              Bandwidth: 86.75 GB/s

 m=4096, n=4096, k=4096, blocksize=64, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.062500, Mean Diff: 0.008217, ref mean:7.546875, OUT mean:7.546875
[Performance] Time: 0.9514 ms
              GFLOPS: 144456.15
              Bandwidth: 79.90 GB/s

 m=4096, n=4096, k=4096, blocksize=32, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.007812, Mean Diff: 0.001299, ref mean:7.507812, OUT mean:7.507812
[Performance] Time: 1.1420 ms
              GFLOPS: 120353.63
              Bandwidth: 67.03 GB/s

SYCL-TLA on B60

=== W4A8===
 m=4096, n=4096, k=4096, blocksize=4096, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.078125, Mean Diff: 0.007988, ref mean:7.476562, OUT mean:7.476562
[Performance] Time: 1.0585 ms
              GFLOPS: 129838.74
              Bandwidth: 71.33 GB/s

 m=4096, n=4096, k=4096, blocksize=128, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.062500, Mean Diff: 0.008240, ref mean:7.500000, OUT mean:7.500000
[Performance] Time: 1.4031 ms
              GFLOPS: 97953.18
              Bandwidth: 53.99 GB/s

 m=4096, n=4096, k=4096, blocksize=64, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.062500, Mean Diff: 0.008240, ref mean:7.480469, OUT mean:7.480469
[Performance] Time: 1.5445 ms
              GFLOPS: 88987.05
              Bandwidth: 49.22 GB/s

 m=4096, n=4096, k=4096, blocksize=32, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.007812, Mean Diff: 0.001297, ref mean:7.511719, OUT mean:7.511719
[Performance] Time: 1.7683 ms
              GFLOPS: 77723.89
              Bandwidth: 43.29 GB/s

Reference Joint_matrix on B60

=== W4A8 / A8W4 ===

 m=4096, n=4096, k=4096, blocksize=4096, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.078125, Mean Diff: 0.007988, ref mean:7.476562, OUT mean:7.476562
[Performance] Time: 1.1099 ms
              GFLOPS: 123827.48
              Bandwidth: 68.03 GB/s

 m=4096, n=4096, k=4096, blocksize=128, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.062500, Mean Diff: 0.008240, ref mean:7.500000, OUT mean:7.500000
[Performance] Time: 1.5408 ms
              GFLOPS: 89199.99
              Bandwidth: 49.17 GB/s

 m=4096, n=4096, k=4096, blocksize=64, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.062500, Mean Diff: 0.008240, ref mean:7.480469, OUT mean:7.480469
[Performance] Time: 1.9249 ms
              GFLOPS: 71400.29
              Bandwidth: 39.49 GB/s

 m=4096, n=4096, k=4096, blocksize=32, compute_type=int8, weight_type=int4, scale_type=fp16, asym=False, device=xpu
  Max Diff: 0.007812, Mean Diff: 0.001297, ref mean:7.511719, OUT mean:7.511719
[Performance] Time: 1.5898 ms
              GFLOPS: 86448.29
              Bandwidth: 48.15 GB/s

Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com>
Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com>
@Zhenzhong1 Zhenzhong1 changed the title [ARK] Optimize SYCL-TLA woqgemm_s8 [ARK] Optimize SYCL-TLA woqgemm_s8 on B60&B70 Aug 10, 2026
@Zhenzhong1 Zhenzhong1 changed the title [ARK] Optimize SYCL-TLA woqgemm_s8 on B60&B70 [ARK] Optimize SYCL-TLA woqgemm_s8 on B70 Aug 11, 2026
Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com>
@chensuyue chensuyue added this to the 0.15.0 milestone Aug 11, 2026
@Zhenzhong1
Zhenzhong1 marked this pull request as ready for review August 11, 2026 07:55

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Optimizes ARK’s SYCL-TLA int8 weight-only GEMM for B70 and adds device-specific build targeting.

Changes:

  • Adds B60/B70/PVC target detection and enables rescaling.
  • Fuses K-block accumulation, dequantization, and output finalization.
  • Adds G21/G31 compiler target mappings.

Reviewed changes

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

File Description
auto_round_extension/ark/setup.py Detects SYCL targets and configures rescaling.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp Implements optimized fused K-block GEMM.
auto_round_extension/ark/auto_round_kernel/CMakeLists.txt Maps GPU targets for the primary kernel build.
auto_round_extension/ark/auto_round_kernel/bestla/CMakeLists.txt Extends benchmark target mappings.

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

Comment thread auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp Outdated
Comment thread auto_round_extension/ark/auto_round_kernel/CMakeLists.txt

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Suppressed comments (2)

auto_round_extension/ark/auto_round_kernel/wrapper/include/xpu_wrapper.hpp:676

  • The positive ARK_AUTO_S8 path only checks that the requested new block size divides k, but packscale groups original scales using integer newblock_size / block. If the new size is not a multiple of the original block size (for example, block size 96 with ARK_AUTO_S8=512 and k=1536), an original scale spans two new blocks while scalext assigns it to only one, so rescaling can overflow and produce wrong results. Require aligned block boundaries or compute maxima for every intersecting original block.
    if (env_params::Instance()->auto_s8 > p->blocksize && p->k % env_params::Instance()->auto_s8 == 0) return true;

auto_round_extension/ark/setup.py:264

  • This enables rescaling even for oneAPI versions below 2025.3, where ARK_SYCL_TLA is disabled and the default build only enables DNNL. With the new ARK_AUTO_S8=-1 default, INT2/INT4 cases that previously fell back to dense DNNL (for example block size 32) are now routed through SyclS8Wrapper, whose build has neither the SYCL-TLA nor joint-matrix backend and throws at runtime. Only enable rescaling when one of those INT8 GEMM backends is enabled.
            "-DARK_RESCALE=ON",

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants