### Install SageAttention3 from Source

Source: https://github.com/thu-ml/sageattention/blob/main/sageattention3_blackwell/README.md

Clone the repository, navigate to the directory, and install the package using setup.py. Ensure you have the base environment requirements met.

```bash
git clone https://github.com/thu-ml/SageAttention
cd SageAttention/sageattention3_blackwell 
python setup.py install
```

--------------------------------

### Install SageAttention3 for Blackwell GPUs

Source: https://context7.com/thu-ml/sageattention/llms.txt

Install SageAttention3 specifically for Blackwell GPUs. This requires Python 3.13+, PyTorch 2.8.0+, and CUDA 12.8+.

```bash
cd SageAttention/sageattention3_blackwell
python setup.py install
```

--------------------------------

### Install FlashAttention3

Source: https://github.com/thu-ml/sageattention/blob/main/bench/README.md

Clone the flash-attention repository and install the package. Ensure you checkout the correct commit for the 2.7.2 release.

```bash
git clone https://github.com/Dao-AILab/flash-attention.git --recursive
git checkout b7d29fb3b79f0b78b1c369a52aaa6628dabfb0d7 # 2.7.2 release
cd hopper
python setup.py install
```

--------------------------------

### Install SageAttention V1

Source: https://github.com/thu-ml/sageattention/blob/main/README.md

Use this command to install SageAttention V1 if you need the Triton version, which is slower than V2/V2++/V3. Refer to the SageAttention-1 branch for details.

```bash
pip install sageattention==1.0.6
```

--------------------------------

### Install SageAttention 2.2.0

Source: https://github.com/thu-ml/sageattention/blob/main/README.md

Install SageAttention 2.2.0, which includes SageAttention2++, using pip. This command installs the package without build isolation.

```bash
pip install sageattention==2.2.0 --no-build-isolation
```

--------------------------------

### Install Dependencies for Parallel Inference

Source: https://github.com/thu-ml/sageattention/blob/main/example/README.md

Before running parallel SageAttention inference, ensure you have the latest versions of xDiT (xfuser) and diffusers installed from their source repositories. This involves installing xfuser with flash_attn support and building/installing the diffusers wheel.

```bash
# install latest xDiT(xfuser).
pip install "xfuser[flash_attn]"
# install latest diffusers (>=0.32.0.dev0), need by latest xDiT.
git clone https://github.com/huggingface/diffusers.git
cd diffusers && python3 setup.py bdist_wheel && cd dist && python3 -m pip install *.whl
```

--------------------------------

### SageAttention Triton INT8/FP16 Kernel Examples

Source: https://context7.com/thu-ml/sageattention/llms.txt

Demonstrates using the Triton implementation of SageAttention with INT8 quantization for Q/K and FP16 for PV. Supports attention masks and returning log-sum-exp.

```python
import torch
from sageattention import sageattn_qk_int8_pv_fp16_triton

batch, heads, seq_len, head_dim = 2, 32, 8192, 128
q = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")
k = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")
v = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")

# With attention mask (bool tensor broadcastable to [batch, heads, q_len, k_len])
mask = torch.ones(batch, 1, seq_len, seq_len, dtype=torch.bool, device="cuda")
mask[:, :, :, seq_len // 2:] = False  # mask out second half of keys
out_masked = sageattn_qk_int8_pv_fp16_triton(
    q, k, v,
    tensor_layout="HND",
    is_causal=False,
    attn_mask=mask,
    smooth_k=True,
)

# Use CUDA quantization backend for better performance via kernel fusion
out_cuda_quant = sageattn_qk_int8_pv_fp16_triton(
    q, k, v,
    tensor_layout="HND",
    quantization_backend="cuda",
    is_causal=False,
)

# Disable smooth_k (saves slight overhead, may reduce accuracy)
out_no_smooth = sageattn_qk_int8_pv_fp16_triton(q, k, v, smooth_k=False)

# Return log-sum-exp
out, lse = sageattn_qk_int8_pv_fp16_triton(q, k, v, return_lse=True)
# lse shape: (1, 24, 2048)
```

--------------------------------

### Run CogVideoX Inference with SageAttention

Source: https://github.com/thu-ml/sageattention/blob/main/example/README.md

Navigate to the example directory and run the cogvideox_infer.py script with the --attention_type sage flag. This will generate a video using SageAttention, which is typically faster than using SDPA.

```bash
cd example
python cogvideox_infer.py --model cogvideox-2b --compile --attention_type sage
```

--------------------------------

### Benchmark Kernel with Custom Arguments

Source: https://github.com/thu-ml/sageattention/blob/main/bench/README.md

Execute kernel benchmarks using specific arguments for quantization granularity and accumulation data types. The examples show configurations for RTX 4090 and H100 GPUs.

```bash
# on RTX 4090
python bench_qk_int8_pv_fp8_cuda.py --pv_accum_dtype fp32+fp16 --quant_gran per_warp
```

```bash
# on H100
python bench_qk_int8_pv_fp8_cuda_sm90.py --pv_accum_dtype fp32+fp32 --quant_gran per_thread
```

--------------------------------

### Execute Parallel SageAttention Inference

Source: https://github.com/thu-ml/sageattention/blob/main/example/README.md

After installing the necessary dependencies, execute the run_parallel.sh script to perform parallel SageAttention inference.

```bash
./run_parallel.sh
```

--------------------------------

### SageAttention Usage Example

Source: https://github.com/thu-ml/sageattention/blob/main/README.md

Basic usage of the SageAttention module. Ensure q, k, and v are FP16/BF16 dtype. The tensor_layout can be 'HND' or 'NHD', and is_causal controls causal masking.

```python
from sageattention import sageattn
attn_output = sageattn(q, k, v, tensor_layout="HND", is_causal=False)
```

--------------------------------

### Kernel Benchmarking with SageAttention

Source: https://context7.com/thu-ml/sageattention/llms.txt

Utilize benchmark scripts in the `bench/` directory to compare SageAttention performance against FlashAttention2 and FlashAttention3. Key arguments include `--quant_gran` for quantization granularity and `--pv_accum_dtype` for precision of two-level accumulation.

```bash
# RTX 4090 (sm89) — SageAttention2++ mode
python bench/bench_qk_int8_pv_fp8_cuda.py --pv_accum_dtype fp32+fp16 --quant_gran per_warp
```

```bash
# H100 (sm90) — Hopper-optimized two-level FP32 accumulation
python bench/bench_qk_int8_pv_fp8_cuda_sm90.py --pv_accum_dtype fp32+fp32 --quant_gran per_thread
```

```bash
# Ampere (sm80, A100) — INT8/FP16 CUDA kernel
python bench/bench_qk_int8_pv_fp16_cuda.py
```

```bash
# Triton INT8/FP16 baseline (all architectures)
python bench/bench_qk_int8_pv_fp16_triton.py
```

--------------------------------

### Compile SageAttention from Source

Source: https://github.com/thu-ml/sageattention/blob/main/README.md

Compile SageAttention from source. Optional environment variables can be set for parallel compilation and CUDA compilation flags.

```bash
git clone https://github.com/thu-ml/SageAttention.git
cd SageAttention 
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32 # Optional
python setup.py install
```

--------------------------------

### Compile SageAttention from Source

Source: https://context7.com/thu-ml/sageattention/llms.txt

Compile SageAttention from source. Ensure TORCH_CUDA_ARCH_LIST is set to target your specific GPU architecture. Adjust EXT_PARALLEL, NVCC_APPEND_FLAGS, and MAX_JOBS as needed for your build environment.

```bash
git clone https://github.com/thu-ml/SageAttention.git
cd SageAttention
export TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0"  # adjust to your GPU
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
python setup.py install
```

--------------------------------

### Run FlashAttention3 Benchmarks

Source: https://context7.com/thu-ml/sageattention/llms.txt

Execute benchmarks for FlashAttention3, including standard and FP8 precision variants. Ensure FlashAttention3 is compiled from source before running.

```python
python bench/bench_fa3.py
```

```python
python bench/bench_fa3_fp8.py
```

--------------------------------

### sageattn_qk_int8_pv_fp16_triton — Triton INT8/FP16 kernel (Ampere/Ada)

Source: https://context7.com/thu-ml/sageattention/llms.txt

SageAttention with per-block INT8 quantization for Q and K, and FP16 PV accumulation implemented using Triton. Suitable for RTX 3090 (sm86) and any Ampere/Ada GPU. Supports an optional boolean or float32 attention mask, key smoothing (`smooth_k`), and switchable quantization backend (`triton` or `cuda`).

```APIDOC
## sageattn_qk_int8_pv_fp16_triton — Triton INT8/FP16 kernel (Ampere/Ada)

### Description
SageAttention with per-block INT8 quantization for Q and K, and FP16 PV accumulation implemented using Triton. Suitable for RTX 3090 (sm86) and any Ampere/Ada GPU. Supports an optional boolean or float32 attention mask, key smoothing (`smooth_k`), and switchable quantization backend (`triton` or `cuda`). The FP16 accumulator is periodically promoted to a FP32 buffer each iteration to limit numerical error.

### Method
`sageattn_qk_int8_pv_fp16_triton`

### Parameters
- **q** (Tensor) - Query tensor.
- **k** (Tensor) - Key tensor.
- **v** (Tensor) - Value tensor.
- **tensor_layout** (str) - Layout of the input tensors, either "HND" (batch, heads, seq, dim) or "NHD" (batch, seq, heads, dim). Defaults to "HND".
- **is_causal** (bool) - Whether to apply causal masking. Defaults to False.
- **attention_mask** (Optional[Union[bool, torch.Tensor]]) - Optional attention mask.
- **smooth_k** (bool) - Whether to apply key smoothing. Defaults to False.

### Request Example
```python
import torch
from sageattention import sageattn_qk_int8_pv_fp16_triton

batch, heads, seq_len, head_dim = 1, 24, 2048, 64
q = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
k = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
v = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")

# Basic non-causal call with default settings
out = sageattn_qk_int8_pv_fp16_triton(q, k, v, tensor_layout="HND", is_causal=False)
```

### Response
#### Success Response (200)
- **out** (Tensor) - The attention output.
```

--------------------------------

### Replace scaled_dot_product_attention with SageAttention

Source: https://github.com/thu-ml/sageattention/blob/main/example/README.md

To use SageAttention, import it and assign it to F.scaled_dot_product_attention. This allows any model that uses this function to leverage SageAttention.

```python
from sageattention import sageattn
import torch.nn.functional as F

F.scaled_dot_product_attention = sageattn
```

--------------------------------

### Use sageattn_qk_int8_pv_fp16_triton for Triton INT8/FP16 Kernel

Source: https://context7.com/thu-ml/sageattention/llms.txt

This Triton-implemented kernel performs per-block INT8 quantization for Q and K, with FP16 PV accumulation. It's suitable for Ampere/Ada GPUs (e.g., RTX 3090) and supports optional attention masks, key smoothing (`smooth_k`), and switchable quantization backends. The FP16 accumulator is periodically promoted to FP32 to mitigate numerical errors.

```python
import torch
from sageattention import sageattn_qk_int8_pv_fp16_triton

batch, heads, seq_len, head_dim = 1, 24, 2048, 64
q = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
k = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
v = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")

# Basic non-causal call with default settings
out = sageattn_qk_int8_pv_fp16_triton(q, k, v, tensor_layout="HND", is_causal=False)

```

--------------------------------

### Integrate SageAttention Globally

Source: https://context7.com/thu-ml/sageattention/llms.txt

Replace the default scaled dot product attention with SageAttention for a 2-3x speedup in attention kernel performance. This global replacement requires no changes to model weights or training and does not measurably affect output quality.

```python
F.scaled_dot_product_attention = sageattn
```

--------------------------------

### Replace Scaled Dot Product Attention with SageAttention

Source: https://github.com/thu-ml/sageattention/blob/main/README.md

This snippet demonstrates how to replace the default `scaled_dot_product_attention` with SageAttention for potential performance gains. Ensure compatibility with your specific model.

```python
import torch.nn.functional as F

from sageattention import sageattn
F.scaled_dot_product_attention = sageattn
```

--------------------------------

### FP4 Attention for Blackwell GPUs

Source: https://context7.com/thu-ml/sageattention/llms.txt

SageAttention3 with FP4 quantization for Q/K and V on Blackwell GPUs (sm120, RTX 5090). Requires Python >= 3.13, torch >= 2.8.0, CUDA >= 12.8. Inputs must be in HND layout. Falls back to SDPA for head_dim >= 256.

```python
import torch
from sageattn3 import sageattn3_blackwell  # separate install from sageattention3_blackwell/

# Requires RTX 5090 / Blackwell GPU (sm120)
batch, heads, seq_len, head_dim = 1, 24, 4096, 128
q = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")
k = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")
v = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")

# Non-causal with per-block mean smoothing (default)
out = sageattn3_blackwell(q, k, v, is_causal=False, per_block_mean=True)
# out shape: (1, 24, 4096, 128)

# Causal attention
out_causal = sageattn3_blackwell(q, k, v, is_causal=True)

# Disable per_block_mean (uses global mean instead; slightly faster, slightly less accurate)
out_global = sageattn3_blackwell(q, k, v, is_causal=False, per_block_mean=False)

# Drop-in replacement for SDPA (only replaces image/video self-attention layers)
import torch.nn.functional as F
F.scaled_dot_product_attention = sageattn3_blackwell
```

--------------------------------

### sageattn_varlen

Source: https://context7.com/thu-ml/sageattention/llms.txt

SageAttention kernel designed for handling batches of variable-length sequences efficiently. It utilizes packed tensors and cumulative sequence length arrays (`cu_seqlens`) for optimized processing. This implementation, built with Triton, features per-block INT8 quantization for QK and FP16 for PV. It offers similar functionality to `flash_attn_varlen_func` from FlashAttention but incorporates SageAttention's specific quantization strategy.

```APIDOC
## `sageattn_varlen` — Variable-length sequence attention

SageAttention for batches with variable-length sequences, using packed (unpadded) tensors and cumulative sequence length arrays (`cu_seqlens`). Implemented via Triton with per-block INT8 quantization for QK and FP16 PV. Equivalent to `flash_attn_varlen_func` in FlashAttention but with SageAttention's quantization strategy.

```python
import torch
from sageattention import sageattn_varlen

# Batch of 3 sequences with lengths [512, 1024, 768], packed into flat tensors
seq_lens_q = [512, 1024, 768]
seq_lens_k = [512, 1024, 768]
total_q = sum(seq_lens_q)
total_k = sum(seq_lens_k)
num_heads, head_dim = 16, 128

q = torch.randn(total_q, num_heads, head_dim, dtype=torch.float16, device="cuda")
k = torch.randn(total_k, num_heads, head_dim, dtype=torch.float16, device="cuda")
v = torch.randn(total_k, num_heads, head_dim, dtype=torch.float16, device="cuda")

# Build cumulative sequence length arrays (must start with 0)
cu_seqlens_q = torch.tensor([0, 512, 1536, 2304], dtype=torch.int32, device="cuda")
cu_seqlens_k = torch.tensor([0, 512, 1536, 2304], dtype=torch.int32, device="cuda")

out = sageattn_varlen(
    q, k, v,
    cu_seqlens_q=cu_seqlens_q,
    cu_seqlens_k=cu_seqlens_k,
    max_seqlen_q=1024,
    max_seqlen_k=1024,
    is_causal=False,
    smooth_k=True,
)
# out shape: (2304, 16, 128)

# Causal variable-length (e.g., autoregressive language model)
out_causal = sageattn_varlen(
    q, k, v,
    cu_seqlens_q=cu_seqlens_q,
    cu_seqlens_k=cu_seqlens_k,
    max_seqlen_q=1024,
    max_seqlen_k=1024,
    is_causal=True,
)
```
```

--------------------------------

### SageAttention CUDA INT8/FP16 Kernel (Ampere sm80)

Source: https://context7.com/thu-ml/sageattention/llms.txt

Utilizes CUDA kernels for INT8 Q/K and FP16 PV on Ampere GPUs (sm80). Offers fine-grained control over quantization granularity and PV accumulation data types for speed/accuracy tuning.

```python
import torch
from sageattention import sageattn_qk_int8_pv_fp16_cuda

# Requires GPU with compute capability 8.0 (A100, A800, RTX 3090 sm86 uses triton variant)
batch, heads, seq_len, head_dim = 2, 32, 8192, 128
q = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")
k = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")
v = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")

# SageAttention2++ mode: per-thread quant + fp16+fp32 two-level accumulation (highest speed)
out_pp = sageattn_qk_int8_pv_fp16_cuda(
    q, k, v,
    tensor_layout="HND",
    is_causal=False,
    qk_quant_gran="per_thread",
    pv_accum_dtype="fp16+fp32",  # Two-level accumulation = SageAttention2++
)

# Most accurate mode: fp32 PV accumulation
out_fp32 = sageattn_qk_int8_pv_fp16_cuda(
    q, k, v,
    pv_accum_dtype="fp32",
    smooth_k=True,
)

# fp16 PV with smooth_v to handle biased value tensors (e.g., CogVideoX)
out_smooth_v = sageattn_qk_int8_pv_fp16_cuda(
    q, k, v,
    pv_accum_dtype="fp16",
    smooth_k=True,
    smooth_v=True,  # Subtract per-channel V mean; only active when pv_accum_dtype="fp16"
)

# Return LSE for distributed attention
out, lse = sageattn_qk_int8_pv_fp16_cuda(q, k, v, return_lse=True)
```

--------------------------------

### Basic SageAttention3 Usage

Source: https://github.com/thu-ml/sageattention/blob/main/sageattention3_blackwell/README.md

Import the sageattn3_blackwell function and use it with query, key, and value tensors. The is_causal parameter controls the use of a causal mask. Input tensors should be FP16/BF16.

```python
from sageattn3 import sageattn3_blackwell
attn_output = sageattn3_blackwell(q, k, v, is_causal=False)
```

--------------------------------

### Utilize SageAttention for Distributed Inference

Source: https://context7.com/thu-ml/sageattention/llms.txt

Enable SageAttention for distributed inference scenarios, such as sequence parallelism or Ring Attention, by using the `return_lse=True` flag. This flag exposes the log-sum-exp normalization factor required by frameworks like xDiT.

```python
return_lse=True
```

--------------------------------

### Global Replacement of scaled_dot_product_attention

Source: https://context7.com/thu-ml/sageattention/llms.txt

Globally replace `F.scaled_dot_product_attention` with `sageattn` before loading any model. This method is suitable for most DiT-based image and video generation models. Ensure this is done before the model's forward pass is called.

```python
import torch
import torch.nn.functional as F
from sageattention import sageattn

# Global replacement — must be done before model is loaded or forward() is called
F.scaled_dot_product_attention = sageattn

# --- CogVideoX example ---
from diffusers import CogVideoXPipeline
from diffusers.utils import export_to_video

pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-2b", torch_dtype=torch.float16)
pipe = torch.compile(pipe.transformer, mode="max-autotune-no-cudagraphs")  # optional
pipe.enable_model_cpu_offload()
pipe.vae.enable_slicing()
pipe.vae.enable_tiling()

video = pipe(
    prompt="A panda eating bamboo in a lush forest",
    num_videos_per_prompt=1,
    num_inference_steps=50,
    num_frames=49,
    guidance_scale=6,
    generator=torch.Generator(device="cuda").manual_seed(42),
).frames[0]
export_to_video(video, "output.mp4", fps=8)
```

--------------------------------

### sageattn3_blackwell

Source: https://context7.com/thu-ml/sageattention/llms.txt

SageAttention3 kernel featuring FP4 quantization for Q/K and V, optimized for NVIDIA Blackwell GPUs (sm120, RTX 5090). This kernel achieves over 1000 TOPS throughput. It is located in the `sageattention3_blackwell/` directory and requires Python 3.13+, PyTorch 2.8.0+, and CUDA 12.8+. Input tensors must use the HND layout (`batch, heads, seq, dim`). The kernel automatically falls back to SDPA when `head_dim` is 256 or greater.

```APIDOC
## `sageattn3_blackwell` — FP4 attention for Blackwell GPUs (RTX 5090)

SageAttention3 uses microscaling FP4 quantization for both Q/K and V on NVIDIA Blackwell (sm120, RTX 5090) GPUs, achieving throughput exceeding 1000 TOPS. Located in `sageattention3_blackwell/`, requires Python ≥ 3.13, torch ≥ 2.8.0, and CUDA ≥ 12.8. Inputs must be in HND layout (`batch, heads, seq, dim`). Falls back to SDPA for head_dim ≥ 256.

```python
import torch
from sageattn3 import sageattn3_blackwell  # separate install from sageattention3_blackwell/

# Requires RTX 5090 / Blackwell GPU (sm120)
batch, heads, seq_len, head_dim = 1, 24, 4096, 128
q = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")
k = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")
v = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")

# Non-causal with per-block mean smoothing (default)
out = sageattn3_blackwell(q, k, v, is_causal=False, per_block_mean=True)
# out shape: (1, 24, 4096, 128)

# Causal attention
out_causal = sageattn3_blackwell(q, k, v, is_causal=True)

# Disable per_block_mean (uses global mean instead; slightly faster, slightly less accurate)
out_global = sageattn3_blackwell(q, k, v, is_causal=False, per_block_mean=False)

# Drop-in replacement for SDPA (only replaces image/video self-attention layers)
import torch.nn.functional as F
F.scaled_dot_product_attention = sageattn3_blackwell
```
```

--------------------------------

### Apply SageAttention3 Blackwell for FP4 Microscaling

Source: https://context7.com/thu-ml/sageattention/llms.txt

Leverage `sageattn3_blackwell` for FP4 microscaling attention on the latest Blackwell hardware. This optimization achieves high TOPS performance on GPUs like the RTX 5090.

```python
sageattn3_blackwell
```

--------------------------------

### SageAttention CUDA INT8/FP8 Kernel (Ada sm89, Blackwell sm120)

Source: https://context7.com/thu-ml/sageattention/llms.txt

Employs CUDA kernels for INT8 Q/K and FP8 V, targeting Ada Lovelace and Blackwell architectures. Supports SageAttention2++ two-level accumulation for maximum throughput and requires specific CUDA versions.

```python
import torch
from sageattention import sageattn_qk_int8_pv_fp8_cuda

# Requires GPU with compute capability 8.9 (RTX 4090, L40, L20) or 12.0 (RTX 5090)
batch, heads, seq_len, head_dim = 2, 24, 4096, 128
q = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
k = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
v = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")

# SageAttention2++ (recommended for RTX 4090): fp32+fp16 two-level accumulation
out = sageattn_qk_int8_pv_fp8_cuda(
    q, k, v,
    tensor_layout="HND",
    is_causal=False,
    qk_quant_gran="per_thread",
    pv_accum_dtype="fp32+fp16",   # SageAttention2++ mode
)

# More accurate two-level fp32+fp32 accumulation
out_accurate = sageattn_qk_int8_pv_fp8_cuda(
    q, k, v,
    pv_accum_dtype="fp32+fp32",
    smooth_k=True,
)

# per_warp quantization granularity (used on Blackwell by sageattn auto-dispatch)
out_warp = sageattn_qk_int8_pv_fp8_cuda(
    q, k, v,
    qk_quant_gran="per_warp",
    pv_accum_dtype="fp32+fp16",
)

# Causal attention with LSE output
out_causal, lse = sageattn_qk_int8_pv_fp8_cuda(
    q, k, v,
    is_causal=True,
    pv_accum_dtype="fp32+fp16",
    return_lse=True,
)
# lse shape: (2, 24, 4096)
```

--------------------------------

### sageattn — Auto-selecting attention kernel

Source: https://context7.com/thu-ml/sageattention/llms.txt

The top-level entry point that automatically selects the best SageAttention kernel based on the detected GPU compute capability. It accepts FP16 or BF16 Q/K/V tensors in either HND or NHD layout and optionally returns the log-sum-exp for distributed attention patterns.

```APIDOC
## sageattn — Auto-selecting attention kernel

### Description
The top-level entry point that automatically selects the best SageAttention kernel based on the detected GPU compute capability (sm80→CUDA INT8/FP16, sm86→Triton INT8/FP16, sm89→CUDA INT8/FP8 SageAttention2++, sm90→Hopper-optimized CUDA INT8/FP8, sm120/sm121→CUDA INT8/FP8 per-warp). Accepts FP16 or BF16 Q/K/V tensors in either HND or NHD layout and optionally returns the log-sum-exp for distributed attention patterns like Ring Attention.

### Method
`sageattn`

### Parameters
- **q** (Tensor) - Query tensor.
- **k** (Tensor) - Key tensor.
- **v** (Tensor) - Value tensor.
- **tensor_layout** (str) - Layout of the input tensors, either "HND" (batch, heads, seq, dim) or "NHD" (batch, seq, heads, dim). Defaults to "HND".
- **is_causal** (bool) - Whether to apply causal masking. Defaults to False.
- **return_lse** (bool) - Whether to return the log-sum-exp. Defaults to False.
- **sm_scale** (float) - Optional softmax scale. Defaults to None.

### Request Example
```python
import torch
from sageattention import sageattn

# Standard non-causal attention, HND layout (batch, heads, seq, dim)
batch, heads, seq_len, head_dim = 2, 16, 4096, 128
q = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
k = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
v = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")

out = sageattn(q, k, v, tensor_layout="HND", is_causal=False)
# out shape: (2, 16, 4096, 128)

# Causal attention, NHD layout (batch, seq, heads, dim)
q_nhd = q.permute(0, 2, 1, 3)  # (2, 4096, 16, 128)
k_nhd = k.permute(0, 2, 1, 3)
v_nhd = v.permute(0, 2, 1, 3)
out_nhd = sageattn(q_nhd, k_nhd, v_nhd, tensor_layout="NHD", is_causal=True)
# out_nhd shape: (2, 4096, 16, 128)

# Return log-sum-exp for Ring Attention / distributed attention
out, lse = sageattn(q, k, v, tensor_layout="HND", is_causal=False, return_lse=True)
# lse shape: (2, 16, 4096)

# Group-query attention: 16 Q heads, 4 KV heads
k_gqa = torch.randn(batch, 4, seq_len, head_dim, dtype=torch.float16, device="cuda")
v_gqa = torch.randn(batch, 4, seq_len, head_dim, dtype=torch.float16, device="cuda")
out_gqa = sageattn(q, k_gqa, v_gqa, tensor_layout="HND", is_causal=False)
# out_gqa shape: (2, 16, 4096, 128)

# Custom softmax scale
out_scaled = sageattn(q, k, v, sm_scale=0.08, tensor_layout="HND")
```

### Response
#### Success Response (200)
- **out** (Tensor) - The attention output.
- **lse** (Tensor, optional) - The log-sum-exp, returned if `return_lse` is True.
```

--------------------------------

### SageAttention Usage

Source: https://github.com/thu-ml/sageattention/blob/main/README.md

This snippet demonstrates how to use the sageattn function from the sageattention library to compute attention outputs. It specifies the input tensor layout and whether to use a causal mask.

```APIDOC
## sageattn

### Description
Computes attention outputs using the SageAttention mechanism.

### Parameters
- **q** (tensor) - Query tensor.
- **k** (tensor) - Key tensor.
- **v** (tensor) - Value tensor.
- **tensor_layout** (string) - Optional. Specifies the layout of the input tensors. Defaults to "HND". Can be "HND" or "NHD".
- **is_causal** (boolean) - Optional. Determines whether to use a causal mask. Defaults to False.

### Request Example
```python
from sageattention import sageattn

attn_output = sageattn(q, k, v, tensor_layout="HND", is_causal=False)
```

### Notes
- Input tensors `q`, `k`, `v` should be of FP16/BF16 dtype.
- Default shape for `tensor_layout="HND"` is `(batch_size, head_num, seq_len, head_dim)`.
- For shape `(batch_size, seq_len, head_num, head_dim)`, set `tensor_layout="NHD"`.
```

--------------------------------

### sageattn_qk_int8_pv_fp16_cuda

Source: https://context7.com/thu-ml/sageattention/llms.txt

CUDA INT8/FP16 kernel optimized for Ampere sm80 GPUs (A100, A800). Supports per-thread or per-warp INT8 quantization for Q and K, and flexible FP16/FP32 accumulation for PV. Includes options for `smooth_k` and `smooth_v` to manage accuracy and handle biased value tensors.

```APIDOC
## sageattn_qk_int8_pv_fp16_cuda — CUDA INT8/FP16 kernel (Ampere sm80)

SageAttention with per-block or per-thread INT8 quantization for Q and K, FP16 PV implemented in CUDA — specifically optimized for A100/A800 (sm80). Exposes `qk_quant_gran` (`"per_thread"` or `"per_warp"`) and `pv_accum_dtype` (`"fp16"`, `"fp16+fp32"`, or `"fp32"`) for fine-grained control over the speed/accuracy trade-off. The `smooth_v` option additionally subtracts the per-channel V mean when `pv_accum_dtype="fp16"` to correct for biased value tensors (e.g., CogVideoX-2b).

```python
import torch
from sageattention import sageattn_qk_int8_pv_fp16_cuda

# Requires GPU with compute capability 8.0 (A100, A800, RTX 3090 sm86 uses triton variant)
batch, heads, seq_len, head_dim = 2, 32, 8192, 128
q = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")
k = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")
v = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")

# SageAttention2++ mode: per-thread quant + fp16+fp32 two-level accumulation (highest speed)
out_pp = sageattn_qk_int8_pv_fp16_cuda(
    q, k, v,
    tensor_layout="HND",
    is_causal=False,
    qk_quant_gran="per_thread",
    pv_accum_dtype="fp16+fp32",  # Two-level accumulation = SageAttention2++
)

# Most accurate mode: fp32 PV accumulation
out_fp32 = sageattn_qk_int8_pv_fp16_cuda(
    q, k, v,
    pv_accum_dtype="fp32",
    smooth_k=True,
)

# fp16 PV with smooth_v to handle biased value tensors (e.g., CogVideoX)
out_smooth_v = sageattn_qk_int8_pv_fp16_cuda(
    q, k, v,
    pv_accum_dtype="fp16",
    smooth_k=True,
    smooth_v=True,  # Subtract per-channel V mean; only active when pv_accum_dtype="fp16"
)

# Return LSE for distributed attention
out, lse = sageattn_qk_int8_pv_fp16_cuda(q, k, v, return_lse=True)
```
```

--------------------------------

### Variable-Length Sequence Attention

Source: https://context7.com/thu-ml/sageattention/llms.txt

Handles batches with variable-length sequences using packed tensors and `cu_seqlens`. Implemented via Triton with INT8 QK and FP16 PV quantization. Equivalent to `flash_attn_varlen_func` but with SageAttention's quantization.

```python
import torch
from sageattention import sageattn_varlen

# Batch of 3 sequences with lengths [512, 1024, 768], packed into flat tensors
seq_lens_q = [512, 1024, 768]
seq_lens_k = [512, 1024, 768]
total_q = sum(seq_lens_q)
total_k = sum(seq_lens_k)
num_heads, head_dim = 16, 128

q = torch.randn(total_q, num_heads, head_dim, dtype=torch.float16, device="cuda")
k = torch.randn(total_k, num_heads, head_dim, dtype=torch.float16, device="cuda")
v = torch.randn(total_k, num_heads, head_dim, dtype=torch.float16, device="cuda")

# Build cumulative sequence length arrays (must start with 0)
cu_seqlens_q = torch.tensor([0, 512, 1536, 2304], dtype=torch.int32, device="cuda")
cu_seqlens_k = torch.tensor([0, 512, 1536, 2304], dtype=torch.int32, device="cuda")

out = sageattn_varlen(
    q, k, v,
    cu_seqlens_q=cu_seqlens_q,
    cu_seqlens_k=cu_seqlens_k,
    max_seqlen_q=1024,
    max_seqlen_k=1024,
    is_causal=False,
    smooth_k=True,
)
# out shape: (2304, 16, 128)

# Causal variable-length (e.g., autoregressive language model)
out_causal = sageattn_varlen(
    q, k, v,
    cu_seqlens_q=cu_seqlens_q,
    cu_seqlens_k=cu_seqlens_k,
    max_seqlen_q=1024,
    max_seqlen_k=1024,
    is_causal=True,
)
```

--------------------------------

### sageattn_qk_int8_pv_fp8_cuda

Source: https://context7.com/thu-ml/sageattention/llms.txt

CUDA INT8/FP8 kernel targeting Ada Lovelace (sm89) and Blackwell (sm120) GPUs. Features INT8 quantization for Q/K and per-channel FP8 for V. Supports SageAttention2++ (fp32+fp16 accumulation) for maximum throughput. Requires CUDA version 12.4+ for sm89 and 12.8+ for sm120.

```APIDOC
## sageattn_qk_int8_pv_fp8_cuda — CUDA INT8/FP8 kernel (Ada sm89, Blackwell sm120)

SageAttention with INT8 quantization for Q and K and per-channel FP8 quantization for V, using CUDA kernels targeting Ada Lovelace (RTX 4090, L40, sm89) and Blackwell (sm120/sm121). Setting `pv_accum_dtype="fp32+fp16"` activates the **SageAttention2++** two-level accumulation strategy for maximum throughput. Requires CUDA ≥ 12.4 for sm89, CUDA ≥ 12.8 for sm120.

```python
import torch
from sageattention import sageattn_qk_int8_pv_fp8_cuda

# Requires GPU with compute capability 8.9 (RTX 4090, L40, L20) or 12.0 (RTX 5090)
batch, heads, seq_len, head_dim = 2, 24, 4096, 128
q = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
k = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
v = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")

# SageAttention2++ (recommended for RTX 4090): fp32+fp16 two-level accumulation
out = sageattn_qk_int8_pv_fp8_cuda(
    q, k, v,
    tensor_layout="HND",
    is_causal=False,
    qk_quant_gran="per_thread",
    pv_accum_dtype="fp32+fp16",   # SageAttention2++ mode
)

# More accurate two-level fp32+fp32 accumulation
out_accurate = sageattn_qk_int8_pv_fp8_cuda(
    q, k, v,
    pv_accum_dtype="fp32+fp32",
    smooth_k=True,
)

# per_warp quantization granularity (used on Blackwell by sageattn auto-dispatch)
out_warp = sageattn_qk_int8_pv_fp8_cuda(
    q, k, v,
    qk_quant_gran="per_warp",
    pv_accum_dtype="fp32+fp16",
)

# Causal attention with LSE output
out_causal, lse = sageattn_qk_int8_pv_fp8_cuda(
    q, k, v,
    is_causal=True,
    pv_accum_dtype="fp32+fp16",
    return_lse=True,
)
# lse shape: (2, 24, 4096)
```
```

--------------------------------

### Use sageattn for Auto-selected Attention Kernel

Source: https://context7.com/thu-ml/sageattention/llms.txt

The `sageattn` function automatically selects the optimal SageAttention kernel based on the detected GPU. It supports HND and NHD tensor layouts, causal masking, group-query attention, and custom softmax scales. Use `return_lse=True` for distributed attention patterns.

```python
import torch
from sageattention import sageattn

# Standard non-causal attention, HND layout (batch, heads, seq, dim)
batch, heads, seq_len, head_dim = 2, 16, 4096, 128
q = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
k = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")
v = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float16, device="cuda")

out = sageattn(q, k, v, tensor_layout="HND", is_causal=False)
# out shape: (2, 16, 4096, 128)

# Causal attention, NHD layout (batch, seq, heads, dim)
q_nhd = q.permute(0, 2, 1, 3)  # (2, 4096, 16, 128)
k_nhd = k.permute(0, 2, 1, 3)
v_nhd = v.permute(0, 2, 1, 3)
out_nhd = sageattn(q_nhd, k_nhd, v_nhd, tensor_layout="NHD", is_causal=True)
# out_nhd shape: (2, 4096, 16, 128)

# Return log-sum-exp for Ring Attention / distributed attention
out, lse = sageattn(q, k, v, tensor_layout="HND", is_causal=False, return_lse=True)
# lse shape: (2, 16, 4096)

# Group-query attention: 16 Q heads, 4 KV heads
k_gqa = torch.randn(batch, 4, seq_len, head_dim, dtype=torch.float16, device="cuda")
v_gqa = torch.randn(batch, 4, seq_len, head_dim, dtype=torch.float16, device="cuda")
out_gqa = sageattn(q, k_gqa, v_gqa, tensor_layout="HND", is_causal=False)
# out_gqa shape: (2, 16, 4096, 128)

# Custom softmax scale
out_scaled = sageattn(q, k, v, sm_scale=0.08, tensor_layout="HND")
```

--------------------------------

### sageattn_qk_int8_pv_fp8_cuda_sm90

Source: https://context7.com/thu-ml/sageattention/llms.txt

Hopper-optimized CUDA INT8/FP8 kernel for attention mechanisms. It is specifically designed for Hopper GPUs (H100, H800, H20, sm90) and utilizes WGMMA instructions and TMA-based memory access. This kernel offers throughput comparable to FlashAttention3-FP8 while maintaining superior numerical accuracy. By default, it employs `pv_accum_dtype="fp32+fp32"` for two-level FP32 accumulation, which is recommended for Hopper's FP32 MMA hardware behavior. Requires CUDA version 12.3 or later.

```APIDOC
## `sageattn_qk_int8_pv_fp8_cuda_sm90` — Hopper-optimized CUDA INT8/FP8 kernel

SageAttention with INT8/FP8 quantization specifically optimized for Hopper GPUs (H100, H800, H20, sm90) using WGMMA instructions and TMA-based memory access. Delivers throughput matching FlashAttention3-FP8 while offering significantly better numerical accuracy. Uses `pv_accum_dtype="fp32+fp32"` (two-level FP32 accumulation) as default to account for Hopper's FP32 MMA hardware behavior. Requires CUDA ≥ 12.3.

```python
import torch
from sageattention import sageattn_qk_int8_pv_fp8_cuda_sm90

# Requires GPU with compute capability 9.0 (H100, H800, H20)
batch, heads_q, heads_kv, seq_len, head_dim = 1, 32, 8, 8192, 128
q = torch.randn(batch, heads_q, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")
k = torch.randn(batch, heads_kv, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")
v = torch.randn(batch, heads_kv, seq_len, head_dim, dtype=torch.bfloat16, device="cuda")

# Default two-level fp32+fp32 accumulation (recommended for Hopper)
out = sageattn_qk_int8_pv_fp8_cuda_sm90(
    q, k, v,
    tensor_layout="HND",
    is_causal=False,
    pv_accum_dtype="fp32+fp32",
    smooth_k=True,
)

# Causal decoding with per-thread quantization
out_causal = sageattn_qk_int8_pv_fp8_cuda_sm90(
    q, k, v,
    is_causal=True,
    qk_quant_gran="per_thread",
    pv_accum_dtype="fp32+fp32",
)

# Return LSE for Ring Attention / Ulysses sequence parallelism
out, lse = sageattn_qk_int8_pv_fp8_cuda_sm90(
    q, k, v,
    pv_accum_dtype="fp32+fp32",
    return_lse=True,
)
# lse shape: (1, 32, 8192)
```
```