auto-round-lib 0.14.2


pip install auto-round-lib

  Latest version

Released: Jul 14, 2026

Project Links

Meta

Classifiers

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
Extras: None
Dependencies:
torch (>=2.13.0)