Intended Audience
- Science/Research
Programming Language
- Python :: 3
Topic
- Scientific/Engineering :: Artificial Intelligence
License
- OSI Approved :: Apache Software License
What is AutoRound Kernel (ARK)?
AutoRound Kernel (ARK) is a low-bit acceleration library for Intel platform, providing three categories of optimized operators for LLM inference.
| Operator Category | CPU | XPU (Battlemage) |
|---|---|---|
| Weight-Only Quantized Linear (INT4/INT8/FP8/FP4) | ✅ | ✅ |
| MoE Grouped GEMM | ❌ | ✅ |
| SageAttention (SDPA / SAGE) | ❌ | ✅ |
Validated CPU: Intel Xeon Scalable (Sapphire Rapids / Emerald Rapids), Intel Xeon 6 (Sierra Forest / Granite Rapids)
Validated GPU: Intel Arc B-Series / Arc Pro B-Series (Battlemage)
Highlights — Ecosystem Integration
ARK kernels are integrated into the following projects:
| Project | Integration | Description |
|---|---|---|
| vllm | inc_wna16_linear.py |
INCXPUARKLinearMethod — weight-only quantized linear on XPU via auto_round_kernel.qlinear.QuantLinear. The goal is to support all AutoRound-quantized models. The plan is tracked in Intel Quantization Support Roadmap |
| vllm-omni | sage_attn.py |
SageAttentionBackend — diffusion model attention on XPU via ARK.sagev1 |
| Transformers (via auto-round) | backend.py |
All models quantized by AutoRound automatically use ARK on CPU/XPU by default; no additional configuration required. 6 backends registered: auto_round_kernel[_xpu] (GPTQ no-zp), auto_round_kernel_zp[_xpu] (GPTQ +zp), auto_round_kernel_awq[_xpu] (AWQ).CPU: INT2/INT4/INT8; XPU: INT4/INT8 |
1. Linear (Weight-Only Quantized GEMM)
Low-bit weight-only linear for LLM inference. Both CPU and XPU are supported.
API
| API | Description | Platform |
|---|---|---|
QuantLinear (example ↓) |
Unified PyTorch module (GPTQ/AWQ/raw quantized checkpoint) | CPU / XPU |
QuantLinearGPTQ |
GPTQ-format checkpoint loader | CPU / XPU |
QuantLinearAWQ |
AWQ-format checkpoint loader | CPU / XPU |
QuantLinearFP8 |
FP8 weight-only linear | CPU / XPU |
woqgemm |
Low-level weight-only GEMM (packed format) | CPU / XPU |
woqgemm_s8 |
Low-level INT8-weight GEMM with scale | CPU / XPU |
_repack_quantized_weight |
Repack raw qweight/qzero/scale → ARK format | CPU / XPU |
_unpack_weight |
Unpack ARK-format weight back to full precision | CPU / XPU |
Key Features
W4A8 / W2A8 Rescale (QQQ-style) — On XPU, ARK supports ~QQQ-style compute: low-bit weights (INT2/INT4) are re-scaled to INT8 and computed via INT8 GEMM, avoiding FP16 dequantization for better throughput. Enabled automatically via environment variable
ARK_AUTO_S8; see xpu_wrapper.hpp.
Supported Data Types
CPU
| Weight dtype | Compute dtype | Scale dtype | Algorithm |
|---|---|---|---|
| INT1–INT8 | INT8[1] / BF16 / FP32 | BF16 / FP32 | sym / asym |
| FP8 (E4M3, E5M2) | BF16 / FP32 | FP32 / FP8 (E8M0) | NA |
| FP4 (E2M1) | BF16 / FP32 | BF16 / FP32 | NA |
XPU
| Weight dtype | Compute dtype | Scale dtype | Algorithm |
|---|---|---|---|
| INT4, INT8 | INT8 / FP16 | FP16 | sym |
| FP8 (E4M3, E5M2) | FP16 | FP16 / FP8 (E8M0) | NA |
[1] INT8 compute includes dynamic activation quantization; results are dequantized to floating-point.
Example
import auto_round_kernel as ark
# Prepare quantized weight: qweight [K, N] int4/int2, scale [K/G, N] fp16/fp32, zp [K/G, N] int4/int2
packw = ark.repack_quantized_weight(
qweight,
scale,
zp,
blocksize=128,
compute_type="fp16",
weight_type="int4",
scale_type="fp16",
asym=False,
)
# Run weight-only quantized GEMM: activation [M, K] → output [M, N]
output = ark.woqgemm(
activation, # [M, K] fp16/bf16
packw, # packed weight blob (INT8)
bias, # [1, N] optional bias
n, # output features
k, # input features
groupsize=128,
compute_type="fp16",
weight_type="int4",
scale_type="fp16",
asym=False,
)
# Decompose back to full precision for verification
decompressed = ark.unpack_weight(
packw,
dtype=torch.float16,
n=n,
k=k,
groupsize=128,
compute_type="fp16",
weight_type="int4",
scale_type="fp16",
asym=False,
)
See test_weightonly.py for an end-to-end example of weight repack, verification, and woqgemm execution on CPU and XPU.
2. MoE (Mixture-of-Experts Grouped GEMM)
Grouped GEMM for MoE layers where different experts process varying numbers of tokens.
API
| Function | Description | Platform | Activation Dtype | Weight Dtype |
|---|---|---|---|---|
ark.moe_gemm(...) (example ↓) |
Grouped GEMM across experts | XPU | FP16 / BF16 | FP16 / BF16 |
ark.moe_gemm(...) (WIP) |
Grouped GEMM with INT4 weight | XPU | FP16 / BF16 | INT4 🚧 |
ark.moe_gemm(...) (WIP) |
Grouped GEMM with INT2 weight | XPU | FP16 / BF16 | INT2 🚧 |
ark.moe_gemm(...) (WIP) |
Grouped GEMM with INT8 weight | XPU | FP16 / BF16 | INT8 🚧 |
🚧 INT2 / INT4 / INT8 weight support is under active development. See #PR.
Details
| Parameter | Shape | Dtype |
|---|---|---|
| activations | [total_tokens, K] |
FP16 / BF16 |
| weights | [num_experts, K, N] (row-major) |
FP16 / BF16 |
| num_tokens_per_expert | [num_experts] |
INT32 |
| scales (optional) | [num_experts, N] |
FP16 / BF16 |
| output | [total_tokens, N] |
same as activations |
Example
# FP16/BF16 MoE
output = ark.moe_gemm(activations, weights, num_tokens_per_expert)
# INT4 MoE (coming soon)
# output = ark.moe_gemm(activations, q4_weights, num_tokens_per_expert, scales=scales)
Build requirement: ARK_SYCL_TLA=ON. See test_moe.py.
3. SageAttention (XPU SDPA Acceleration)
ARK provides a full family of scaled dot-product attention kernels on XPU, ranging from vanilla FP16 SDPA to INT8-quantized SageAttention variants.
API Overview
| Function | Description | Q/K/V Input | PV Precision | Head Dim |
|---|---|---|---|---|
ark.sdpa (example ↓) |
FP16/BF16 SDPA (flash attention) | FP16 / BF16 | FP16 | 64, 96, 128, 192 |
ark.sage |
Low-level INT8 SAGE (pre-quantized Q/K) | INT8 (Q/K), FP16 (V) | FP16 | 64, 128 |
ark.sage_pvi8 |
Low-level INT8 SAGE (pre-quantized Q/K/V) | INT8 | INT8 | 64, 128 |
ark.sagev1 |
High-level FP16 → internal Q/K quant → SAGE | FP16 / BF16 | FP16 | 64, 128 |
ark.sagev1_pvi8 |
High-level FP16 → internal Q/K/V quant → SAGE PV INT8 | FP16 / BF16 | INT8 | 64, 128 |
ark.sageattn |
Dispatcher (sageattention-compatible API) | FP16 / BF16 | FP16 / INT8 | 64, 128 |
ark.sage_dynquant |
Dynamic INT8 block-wise Q/K quant → SAGE (drop-in SDPA replacement) | FP16 / BF16 | FP16 | 64, 128 |
Comparison
| Feature | sdpa |
sagev1 |
sagev1_pvi8 |
sage_dynquant |
|---|---|---|---|---|
| Q/K quantization | None | Internal INT8 | Internal INT8 | Internal INT8 |
| PV quantization | None | None | Internal INT8 | None |
| quant_block_size | N/A | 1 / ≥32 | 1 / ≥32 | 1 / ≥32 |
| Additive mask | ✅ [B,1,Sq,Skv] FP32 | ✅ | ✅ | ✅ |
| Causal mask | ✅ | ✅ | ✅ | ✅ |
| GQA | ✅ | ✅ | ✅ | ✅ |
| Tensor layout | HND / NHD | HND / NHD | HND / NHD | HND |
Drop-in SDPA Replacement
Replace torch.nn.functional.scaled_dot_product_attention globally for lm-eval:
cd /path/to/auto_round_extension/ark
PYTHONPATH=$PWD python tools/lm_eval_with_ark_sdpa.py \
--model hf \
--model_args pretrained=/path/to/model,trust_remote_code=True,dtype=bfloat16 \
--tasks hellaswag,piqa,winogrande \
--device xpu:0 --batch_size 1
The patching logic (in auto_round_kernel/torch_sdpa_patch.py) routes to ARK on XPU when constraints are met; otherwise falls back to PyTorch SDPA.
Constraints
| Constraint | sdpa |
sagev1 / sagev1_pvi8 / sage_dynquant |
|---|---|---|
| Q/K/V dtype | FP16, BF16 | FP16, BF16 |
| Head dim | 64, 96, 128, 192 | 64, 128 |
dropout_p |
must be 0.0 | must be 0.0 |
| Boolean mask | falls back to torch | falls back to torch |
| Additive mask shape | [B, 1, Sq, Skv] FP32 |
[B, 1, Sq, Skv] FP32 |
| quant_block_size | N/A | 1 (per-token) or ≥32 |
Installation
Install via pip
pip install auto-round-lib
Install from Source
pip install . --no-build-isolation
# or
python setup.py bdist_wheel; pip install dist/*
Build with MoE / SageAttention support requires ARK_SYCL_TLA=ON.
Tests
| Test | Description |
|---|---|
| test_weightonly.py | WOQ GEMM pack/unpack/run on CPU & XPU |
| test_moe.py | MoE grouped GEMM |
| test_flash_attn.py | SDPA (flash attention) prefill |
| test_sdpa.py | SDPA benchmark suite |
| test_sdpa_parity.py | SDPA vs PyTorch parity check |
| test_sage_dynquant.py | SageAttention dynamic INT8 quant benchmarks |
| test_bench_bmg.py | BMG SDPA / SageAttention benchmarking |
| test_matmul.py | Low-level matmul |
| test_packq.py | Weight packing utilities |
Wheel compatibility matrix
| Platform | CPython 3.10 | CPython 3.11 | CPython 3.12 | CPython 3.13 | CPython 3.14 |
|---|---|---|---|---|---|
| manylinux_2_27_x86_64 | |||||
| manylinux_2_28_x86_64 |