[ARK] update README

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<sup>[1]</sup> |
 |------------------|:--------------------------------:|:-----------------:|:-----------------------:|
 | INT8             | INT8<sup>[2]</sup> / 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     |
 
 <sup>[1]</sup>: Quantization algorithms for integer types: symmetric or asymmetric.  
 <sup>[2]</sup>: 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.