diff --git a/auto_round_extension/ark/README.md b/auto_round_extension/ark/README.md index 5b66ccc634..a9d8eb63a9 100644 --- a/auto_round_extension/ark/README.md +++ b/auto_round_extension/ark/README.md @@ -1,5 +1,6 @@ ## What is AutoRound Kernel? -AutoRound Kernel is a low-bit acceleration library for Intel platform. + +AutoRound Kernel (ARK) is a low-bit acceleration library for Intel platform, providing high-performance kernels for LLM inference including weight-only quantization, flash attention (with [SageAttention](https://github.com/thu-ml/SageAttention) v1 support), and Mixture-of-Experts (MOE) computation. The kernels are optimized for the following CPUs: * Intel Xeon Scalable processor (formerly Sapphire Rapids, and Emerald Rapids) @@ -10,8 +11,10 @@ The kernels are optimized for the following GPUs: (formerly Battlemage) ## Key Features -AutoRound Kernel provides weight-only linear computational capabilities for LLM inference. Specifically, the weight-only-quantization configs we support are given in the table below: -### CPU + +AutoRound Kernel provides the following computational capabilities for LLM inference: + +### Weight-Only Quantization (WOQ) Linear | Weight dtype | Compute dtype | Scale dtype | Algorithm[1] | |------------------|:--------------------------------:|:-----------------:|:-----------------------:| | INT8 | INT8[2] / BF16 / FP32 | BF16 / FP32 | sym / asym | @@ -25,44 +28,76 @@ AutoRound Kernel provides weight-only linear computational capabilities for LLM | FP8 (E4M3, E5M2) | BF16 / FP32 | FP32 / FP8 (E8M0) | NA | | FP4 (E2M1) | BF16 / FP32 | BF16 / FP32 | NA | -### XPU +### XPU Weight-Only Quantization | Weight dtype | Compute dtype | Scale dtype | Algorithm | |------------------|:--------------:|:-----------------:|:----------:| | INT8 | INT8 / FP16 | FP16 | sym | | INT4 | INT8 / FP16 | FP16 | sym | +| INT2 | INT8 / FP16 | FP16 | sym | | FP8 (E4M3, E5M2) | FP16 | FP16 / FP8 (E8M0) | NA | [1]: Quantization algorithms for integer types: symmetric or asymmetric. [2]: Includes dynamic activation quantization; results are dequantized to floating-point formats. - + +### Flash Attention (XPU only) + +ARK provides multiple attention backends for prefill and decode, with [SageAttention](https://github.com/thu-ml/SageAttention) v1 support — offering INT8-quantized attention variants for higher throughput on Intel Arc GPUs. (SageAttention v2/v3 are under development.) + +| Backend | Description | Q/K/V dtype | Head dim | Features | +|---------|-------------|-------------|----------|----------| +| `sdpa` | Standard flash attention (FP16/BF16) | FP16 / BF16 | 64, 96, 128, 192 | Causal mask, additive mask, GQA | +| `sagev1` | **SageAttention v1** with INT8 Q/K quantization | FP16 / BF16 | 64, 128 | Block-wise INT8 QK, PV in half | +| `sagev1_pvi8` | **SageAttention v1** with INT8 Q/K/V quantization | FP16 / BF16 | 64, 128 | Block-wise INT8 QK + INT8 PV | +| `sage` | Low-level SageAttention with pre-quantized INT8 Q/K | INT8 (Q/K), FP16 (V) | 64, 128 | External Q/K scales | +| `sage_pvi8` | Low-level SageAttention with pre-quantized INT8 Q/K/V | INT8 (Q/K/V) | 64, 128 | External Q/K/V scales | +| `sage_dynquant` | SageAttention with fused dynamic INT8 quantization | FP16 | 64, 128 | Auto-quantizes Q/K internally | + +All attention backends support both `HND` (`[B, H, N, D]`) and `NHD` (`[B, N, H, D]`) tensor layouts, as well as non-contiguous (sliced) input tensors. + +### MOE GEMM (XPU only) + +Grouped GEMM for Mixture-of-Experts layers, supporting FP16/BF16 with variable token counts per expert. ## Installation + ### 1. Install via pip ```bash pip install auto-round-lib ``` ### 2. Install from Source + +Requires a sourced oneAPI environment (2025.3+ recommended for SYCL TLA support). + ```bash +# Source oneAPI environment +source /opt/intel/oneapi/setvars.sh + +# Build and install pip install . --no-build-isolation # or -python setup.py bdist_wheel;pip install dist/* +python setup.py bdist_wheel; pip install dist/* ``` +The build system automatically detects the oneAPI version. SYCL TLA (Tensor Linear Algebra) kernels are enabled when oneAPI >= 2025.3. + ### Validated Hardware Environment + #### CPU based on [Intel 64 architecture or compatible processors](https://en.wikipedia.org/wiki/X86-64): * Intel Xeon Scalable processor (Granite Rapids) + #### GPU built on Intel's Xe architecture: * Intel Arc B-Series Graphics (Battlemage) -### Resources +## API Reference -#### QuantLinear API - ARK exposes a unified weight-only linear interface through QuantLinear, QuantLinearGPTQ, QuantLinearAWQ, and QuantLinearFP8. Please refer to the [QLinear](auto_round_kernel/qlinear.py) for more integration details. +### QuantLinear (Weight-Only Quantization) - The expected lifecycle is: create the module, load quantized tensors from the checkpoint, call post_init() once to repack weights into the ARK-friendly layout, and then call forward() during inference. +ARK exposes a unified weight-only linear interface through `QuantLinear`, `QuantLinearGPTQ`, `QuantLinearAWQ`, and `QuantLinearFP8`. Please refer to the [QLinear](auto_round_kernel/qlinear.py) for more integration details. - Minimal usage: +The expected lifecycle is: create the module, load quantized tensors from the checkpoint, call `post_init()` once to repack weights into the ARK-friendly layout, and then call `forward()` during inference. + +Minimal usage: ```python from auto_round_kernel.qlinear import QuantLinear @@ -82,13 +117,131 @@ qlinear.post_init() y = qlinear(x) ``` -#### A Weight-Only Example - A runnable end-to-end example is available in [test_weightonly.py](test/test_weightonly.py). It demonstrates how to prepare quantized weights and scales, call repack_quantized_weight to build ARK-packed weights, verify correctness with unpack_weight, and run woqgemm on CPU and XPU. +### Attention APIs + +#### `ark.sdpa` — Standard Flash Attention + +```python +import auto_round_kernel as ark + +output = ark.sdpa( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + scale=None, + tensor_layout="HND", # "HND" or "NHD" +) +``` + +- Q/K/V dtype: FP16 or BF16 +- Supported head dims: 64, 96, 128, 192 +- Supports GQA (grouped query attention) via different Hq/Hkv + +#### `ark.sagev1` — SAGE v1 Attention + +```python +output = ark.sagev1( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + scale=None, + enable_gqa=False, + quant_block_size=64, # block size for INT8 QK quantization + tensor_layout="HND", +) +``` + +- Q/K/V dtype: FP16 or BF16 +- Supported head dims: 64, 128 +- Internally quantizes Q/K to INT8 per block; PV computed in half precision +- Falls back to `ark.sdpa` when `quant_block_size <= 0` + +#### `ark.sagev1_pvi8` — SAGE v1 with INT8 PV + +Same interface as `sagev1` but also quantizes V to INT8 internally for higher throughput. -#### Replace torch SDPA and run lm-eval - ARK also exposes an XPU SDPA kernel through `ARK.sdpa(...)`. If you want to replace `torch.nn.functional.scaled_dot_product_attention` globally for evaluation without editing model code, use the helper launcher in [tools/lm_eval_with_ark_sdpa.py](tools/lm_eval_with_ark_sdpa.py). +#### `ark.sage_dynquant` — SAGE with Fused Dynamic Quantization + +```python +output = ark.sage_dynquant( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + scale=None, + enable_gqa=False, + quant_block_size=64, +) +``` + +- Takes FP16 Q/K/V, performs fused block-wise INT8 quantization of Q/K via SYCL kernel +- Supports quant_block_size: 1 (per-token), 32, 64, 128, 256 +- Auto-pads sequence lengths for block alignment + +#### `ark.sageattn` — SageAttention-Compatible Dispatcher + +ARK provides a drop-in replacement for the `sageattention.sageattn` API, enabling seamless integration with existing SageAttention workflows: + +```python +import auto_round_kernel as ark + +# Drop-in replacement for sageattention.sageattn +output = ark.sageattn( + q, + k, + v, + tensor_layout="HND", + is_causal=False, + sm_scale=None, + return_lse=False, + kernel="v1_pvhalf", # or "v1_pvi8" + **kwargs, +) +``` + +This mirrors the [SageAttention](https://github.com/thu-ml/SageAttention) interface for drop-in compatibility, allowing models using SageAttention to run on Intel Arc GPUs without code changes. + +### MOE GEMM + +```python +output = ark.moe_gemm( + activations, # [total_tokens, K] FP16/BF16 + weights, # [num_experts, K, N] FP16/BF16 + num_tokens_per_expert, # [num_experts] int32 + scales=None, # optional [num_experts, N] FP16/BF16 +) +``` + +### Patching torch SDPA + +ARK can globally replace `torch.nn.functional.scaled_dot_product_attention` for evaluation, including SageAttention backends: + +```python +import auto_round_kernel as ark + +# Patch with standard SDPA backend +ark.patch_torch_sdpa(backend="sdpa") + +# Patch with SageAttention v1 backend (INT8 QK, PV half) +ark.patch_torch_sdpa(backend="sagev1", quant_block_size=64) + +# Patch with SageAttention v1 + INT8 PV backend +ark.patch_torch_sdpa(backend="sagev1_pvi8", quant_block_size=64) + +# Restore original +ark.unpatch_torch_sdpa() +``` + +Or use the helper launcher for lm-eval: - Example: ```bash cd /path/to/auto_round_extension/ark PYTHONPATH=$PWD python tools/lm_eval_with_ark_sdpa.py \ @@ -99,9 +252,42 @@ PYTHONPATH=$PWD python tools/lm_eval_with_ark_sdpa.py \ --batch_size 1 ``` - Notes: - * The patch only routes calls to ARK on XPU when the inputs match ARK kernel constraints; otherwise it falls back to the original torch SDPA. - * Supported Q/K/V dtypes are FP16 and BF16. - * Supported head dimensions are 64, 96, 128, and 192. - * `dropout_p` must be 0.0 for the ARK path. - * Additive masks are supported when they can be normalized to `[B, 1, Sq, Skv]`; boolean masks fall back to torch. +### Low-Level Matrix Operations + +```python +# FP16/BF16 matrix multiply with bias +C = ark.matmul(A, B, bias) + +# INT8 matrix multiply (s8s8s32) +C = ark.igemm_s8s8s32(A, B) + +# Weight-only quantized GEMM with INT8 weights +C = ark.woqgemm_s8(A, B, scaleB, bias) + +# General weight-only quantized GEMM +C = ark.woqgemm(A, B, bias, n, k, groupsize, compute_type, weight_type, scale_type, asym) +``` + +## Testing + +Unit tests are available in the [test](test/) directory: + +| Test file | Description | +|-----------|-------------| +| `test_weightonly.py` | Weight-only quantized GEMM (CPU + XPU) | +| `test_flash_attn.py` | Flash attention (sdpa) correctness | +| `test_sdpa.py` | SDPA benchmark suite | +| `test_sdpa_parity.py` | SDPA parity with non-contiguous inputs and layouts | +| `test_sage_dynquant.py` | SageDynQuant block-wise benchmark | +| `test_bench_bmg.py` | BMG comparison benchmark | +| `test_matmul.py` | General matrix multiply | +| `test_packq.py` | Weight packing/unpacking | +| `test_moe.py` | MOE GEMM correctness | + +## Notes + +- The SDPA patch only routes calls to ARK on XPU when inputs match kernel constraints; otherwise it falls back to the original torch SDPA. +- Supported Q/K/V dtypes for attention are FP16 and BF16 (except SAGE variants which may use INT8 internally). +- `dropout_p` must be 0.0 for all ARK attention paths. +- Additive masks are supported when they can be normalized to `[B, 1, Sq, Skv]`; boolean masks fall back to torch. +- Non-contiguous (sliced) input tensors are supported for all attention backends.