Форк официального SageAttention с полной оптимизацией под SM75 (Turing) — NVIDIA T4, RTX 2080, Quadro RTX
ВАЖНО (2026-07-03): Найдена и исправлена настоящая причина визуального шума:
ldmatrix_m8n8x4несовместим с INT8 MMA на Turing.ldmatrixраскладывает данные в FP16 tensor core layout, несовместимый с INT8 MMA fragment layout. Заменён на ручнуюuint32_tзагрузку.Дополнительно: K_GRAN теперь всегда
kPerBlock(как в официальном SageAttention), убранrepeat_interleaveкостыль.Требуется: пересборка (
python setup.py build_ext --inplace) и тест в ComfyUI.
Оригинальный SageAttention не поддерживает SM75 (Turing) — на T4 он просто не собирается. Мы исправили это и написали полноценный CUDA-кернел для архитектуры Turing, использующий:
- INT8 тензорные ядра для QK⊤ (m8n8k16 — 130 TOPS)
- FP16 тензорные ядра для PV (m8n8k16 — 65 TFLOPS)
- FP32 аккумуляцию для точности
- Полный FlashAttention-совместимый softmax (онлайн, warp-редукции)
- LSE (log-sum-exp) для Ring Attention / xDiT
| Компонент | Что сделано |
|---|---|
MMA-обёртки (csrc/mma.cuh) |
m8n8k16 INT8 + m8n8k4 FP16 для SM75 (m8n8k16 FP16 валиден только на SM80+). A=2,B=2 регистра для FP16 |
LDMATRIX-обёртки (csrc/mma.cuh) |
ldmatrix_m8n8x1(_trans), ldmatrix_m8n8x2(_trans), ldmatrix_m8n8x4(_trans) — полный набор для SM75 MMA-фрагментов |
| QK MMA | INT8 m8n8k16 с ldmatrix-загрузкой Q/K фрагментов |
| Онлайн-softmax | Полный FlashAttention: running max/denominator, warp-редукции через __shfl_xor_sync |
| P → smem → ldmatrix | Softmax-выход пишется в shared memory, перезагружается через ldmatrix для корректной MMA-раскладки |
| PV MMA | FP16 m8n8k4 с pv_k-циклом (K=4 → 2 прохода на 8×8 тайл). fk-цикл итерирует столбцы V (head_dim/8 суб-тайлов) |
| 2D RO_accum | Раздельные аккумуляторы для каждого mq-подтайла, онлайн-softmax через K-тайлы |
| Output store | Прямая запись в global O (без smem_O roundtrip) + smem_O staging вариант для бенчмарка |
| static_assert | Проверки на этапе компиляции: делимость head_dim, CTA_Q/WARP_Q, CTA_K/WARP_K |
| LSE | Log-sum-exp возврат с warp-редукциями для Ring Attention |
| Компонент | Что сделано |
|---|---|
| Авто-диспетчер | sageattn() → SM75 → sageattn_qk_int8_pv_fp16_cuda_sm75() |
| BF16→FP16 | Авто-каст для Turing (не поддерживает BF16) |
| Padding | Авто-подгон head_dim под 16 (INT8 MMA K) |
| Починены баги | try/except, stray-код, дублирующиеся блоки, экранированные docstring |
| Файл | Что проверяет |
|---|---|
test_sm75_kernel.py |
4 конфига head_dim, GQA, causal, LSE, edge cases, back-to-back переиспользование тензора |
bench/bench_sm75_output_store.py |
Микро-бенчмарк: smem_O staging vs прямой вывод (13 конфигов) |
| Файл | Назначение |
|---|---|
comfyui_nodes/ |
Custom Node — авто-установка через клонирование репо |
sageattn_t4_nodes.py |
Bridge-файл для авто-обнаружения ComfyUI (в корне репо) |
scripts/build_sm75_kaggle.sh |
Скрипт сборки для Kaggle T4×2 |
# Клонируй репо ПРЯМО в custom_nodes (НЕ копируй подпапку!):
cd ComfyUI/custom_nodes/
git clone https://github.com/THE-ANGEL-AI/SageAttention-SM75-path.git SageAttention-T4
cd SageAttention-T4
# Собери CUDA-кернелы под T4:
pip install -e .
# Перезапусти ComfyUIПосле перезапуска в меню нод появится категория 🧠 SageAttention-T4 с двумя нодами:
| Нода | Где использовать |
|---|---|
| 🧠 SageAttention-T4 Apply (INT8 Turbo) | Между загрузчиком модели и сэмплером. Параметры: smooth_k (on/off), enable (on/off) |
| 🧠 SageAttention-T4 Remove | Убирает патч, возвращает оригинальное внимание |
Схема workflow:
[Load LTX Model] → [🧠 SageAttention-T4 Apply] → [LTX Sampler] → [VAE Decode] → [Output]
Нода использует
model.add_object_patch()— патч действует только на эту модель, не ломая другие части workflow. Никаких глобальных monkey-patch'ей.
| Длина контекста | Ускорение attention | End-to-end (видео-модели) |
|---|---|---|
| 2K | ~1.3× | +10-15% |
| 8K | ~1.8× | +30-40% |
| 16K | ~2.5× | +60-75% |
При правильном распределении слоёв по GPU + INT8 внимании:
- LTX-Video 1280×720×15сек: с ~400 сек/итерацию → ~45-70 сек (5-9× общий прирост — в основном за счёт устранения PCIe-spill в системную RAM)
# Клонируем форк
git clone https://github.com/THE-ANGEL-AI/SageAttention-SM75-path.git
cd SageAttention-SM75-path
# Собираем SM75 расширение
python setup.py build_ext --inplace
# Или на Kaggle T4×2:
bash scripts/build_sm75_kaggle.sh
# Проверяем
python test_sm75_kernel.pyТребования: Python ≥3.9, PyTorch ≥2.3, CUDA ≥11.8, NVIDIA T4 (или другой SM75 GPU)
✅ Протестировано на Kaggle T4×2: CUDA 13.0, PyTorch 2.10.0, TorchVision 0.25.0, TorchAudio 2.10.0 — PyTorch cu130
from sageattention import sageattn
import torch.nn.functional as F
# Вариант 1: Прямой вызов
output = sageattn(q, k, v, tensor_layout="HND", is_causal=False)
# Вариант 2: Monkey-patch (все sdpa вызовы → SageAttention)
F.scaled_dot_product_attention = sageattnSageAttention-SM75-path/
├── csrc/qattn/
│ ├── attn_cuda_sm75.h # ← основной SM75-кернел
│ ├── pybind_sm75.cpp # ← pybind-регистрация
│ └── qk_int_sv_f16_cuda_sm75.cu
├── csrc/mma.cuh # ← SM75 MMA-обёртки (int8 m8n8k16, fp16 m8n8k4) + ldmatrix
├── sageattention/
│ ├── core.py # ← авто-диспетчер + SM75-функция
│ ├── quant.py # ← CUDA квантование
│ └── triton/ # ← Triton-кернелы
├── comfyui_nodes/
│ └── sageattention_node.py # ← ComfyUI Custom Node
├── test_sm75_kernel.py # ← юнит-тесты
├── bench/
│ └── bench_sm75_output_store.py
├── scripts/
│ └── build_sm75_kaggle.sh # ← сборка на Kaggle
└── setup.py # ← SM75-only сборка
Оригинальный SageAttention: thu-ml/SageAttention — Jintao Zhang, Jia Wei, Haofeng Huang, Pengle Zhang, Jun Zhu, Jianfei Chen.
Форк основан на XUANNISSAN/SageAttention-SM75-path.
⚡ SageAttention-SM75 · Предварительная версия · Тестирование в процессе