SageAttention: Thư viện Attention lượng tử hóa

SageAttention là gì?

SageAttention là thư viện CUDA/Triton ba thế hệ cho attention lượng tử hóa, tăng tốc suy luận transformer 2-5 lần so với FlashAttention trên GPU hiện đại. Nó áp INT8 smoothing cho tích QK^T và tích lũy FP8 hoặc FP16 cho tích PV, với kernel được điều chỉnh cho phần cứng Ampere, Ada, Hopper và Blackwell. API đơn giản qua hàm `sageattn` gọi được từ PyTorch, mỗi phiên bản exposes các trade-off precision/accuracy khác nhau.

Tính năng chính và kiến trúc GPU được hỗ trợ

• ✓Lượng tử hóa INT8 cho QK^T với granularity theo token, theo head, hoặc theo thread (SageAttention2)
• ✓Lượng tử hóa FP8 cho PV với accumulator FP16 — tránh suy giảm độ chính xác mà FP8 tích lũy nguyên thủy gây ra trên Hopper
• ✓Chiến lược tích lũy hai tầng cho PV trên Hopper WGMMA, giữ nguyên độ chính xác ở precision FP8
• ✓Làm mượt outliers trong attention logits trước khi lượng tử hóa, ngăn lỗi lớn mà INT8 rounding tiêu chuẩn gây ra
• ✓Backend Triton (v1 pip) và backend CUDA (v2++), tự chọn kernel qua `sageattn`
• ✓Biến thể sparse attention SpargeAttn mở rộng kernel lượng tử hóa cho arbitrary block-sparse patterns mà không cần training
• ✓Hỗ trợ sequence có độ dài biến đổi qua `sageattn_varlen` và native group-query attention
• ✓RTX 5090 đạt 560 TOPS trên attention kernel — nhanh hơn 2.7 lần so với FlashAttention2 trên cùng phần cứng

Dùng SageAttention trong code của bạn

import torch from sageattention import sageattn # Standard HND layout (batch, heads, seq, dim) q = torch.randn(1, 16, 512, 64, dtype=torch.float16, device='cuda') k = torch.randn(1, 16, 512, 64, dtype=torch.float16, device='cuda') v = torch.randn(1, 16, 512, 64, dtype=torch.float16, device='cuda') output = sageattn(q, k, v, tensor_layout='HND', is_causal=False) # Drop-in SDPA replacement for compatible models: # torch.nn.functional.scaled_dot_product_attention = sageattn Với video model kiểu CogVideoX, README cho thấy swap SDPA trực tiếp hoạt động. Với image và video DiT model, path được khuyên là sửa trực tiếp Attention class (ví dụ trong example/modify_mochi.py).

Cài đặt SageAttention

pip install sageattention==2.2.0 --no-build-isolation Để có feature v2++ mới nhất bao gồm hỗ trợ Blackwell: compile từ source bằng cách clone repo, rồi chạy python setup.py install. Đặt EXT_PARALLEL=4, NVCC_APPEND_FLAGS="--threads 8", và MAX_JOBS=32 để compile nhanh hơn. Lưu ý rằng sageattention v1 từ pip dùng kernel Triton chậm hơn — compile từ source mới kích hoạt được hiệu năng CUDA thực sự.

SageAttention so với FlashAttention

Ai nên dùng SageAttention?

Dùng SageAttention nếu bạn serving hoặc fine-tune LLM, model video generation (CogVideoX, Mochi), hoặc image model dạng DiT trên GPU Ampere, Ada, hoặc Hopper và kernel attention là nút thắt cổ chai. Nó tỏa sáng khi bạn có context dài hoặc batch size lớn. Bỏ qua nếu bạn dùng phần cứng cũ, không đáp ứng được CUDA 12+, hoặc cần giữ nguyên kết quả bit-exact — và với pipeline production nhạy cảm độ chính xác, mặc định dùng SageAttention2 thay vì v3.

Câu hỏi thường gặp

Bài toán mà nó giải quyết

Phép tính attention là nút thắt cổ chai chính trong transformer có context dài, và attention FP16/FP32 nguyên thủy scale kém theo sequence length. FlashAttention giải quyết được memory wall nhưng bỏ qua phần lượng tử hóa gây áp lực độ chính xác. SageAttention lấp khoảng trống đó bằng cách lượng tử hóa QK^T sang INT8 và PV sang FP8, ép thêm throughput từ tensor core mà không gây suy giảm độ chính xác như tính toán low-precision nguyên thủy.

Trường hợp sử dụng tốt nhất

• •Suy luận LLM context dài trên H100/H800/H20 với tích lũy FP8 PV
• •Tăng tốc suy luận video generation (CogVideoX, Mochi) qua thay attention DiT
• •Image transformer (DiT-based) với INT8 QK^T và FP16 PV cho tăng tốc không mất accuracy
• •Batch inference với sequence có độ dài khác nhau trong cùng một batch
• •Serving các mô hình lượng tử hóa mà attention computation là nút thắt cổ chai
• •Triển khai trên GPU Blackwell với SageAttention2++ nhắm vào RTX 5090

Repo liên quan