Triton, Flash-attension, Sage-attension и bitsandbytes с ROCm 7 в Windows / Хабр

В конце января 2026 вышел triton-windows 3.6.0.post25, который позволяет использовать flash-attention, sage-attention (v1) и другие библиотеки, использующие Triton, на картах AMD с поддержкой rocWMMA (то есть начиная с архитектуры gfx1100: RX 7900 XT/XTX и новее) в Windows.

Также, несмотря на то что в официальном репозитории bitsandbytes еще не приняли PR для поддержки ROCm 7, его все же можно собрать, внеся небольшие изменения в код. Эти изменения я уже сделал в своем форке.

В этой статье я расскажу, как установить все это себе, а также для примера запустим пару тестов в ComfyUI, в том числе со свежей LTX-2, и сделаем Qlora адаптер для модели Gemma 3 с использованием bitsandbytes.

Все скрипты используемые в статье можно найти в моем репозитории.

Установка компонентов

Средства сборки

Для сборки bitsandbytes и flash-attention2 нам понадобится установить некоторые средства разработки. Я использую Chocolatey.

choco install visualstudio2022buildtools -y --params "--add Microsoft.VisualStudio.Component.VC.Tools.x86.x64 --add Microsoft.VisualStudio.Component.VC.CMake.Project --add Microsoft.VisualStudio.Component.VC.ATL"
choco install git.install -y --params "'/GitAndUnixToolsOnPath'"
choco install cmake --version=3.31.0 -y
choco install ninja -y
choco install ccache -y
choco install python --version=3.12.0 -y

Примечание: Возможно, подойдет и другая версия Python, но я проверял на 3.12. Так же в моих репозитория пока собраны версии пакетов только для Python 3.12.

Дополнительно у меня установлена Visual Studio 2022 Community со следующим набором компонентов:

Workload: Desktop development with C++
Components: C++ CMake Tools for Windows, Git for Windows, C++ Clang Compiler for Windows, MSBuild Support for LLVM Toolset (clang)

Но я не уверен, что это необходимо.

ВАЖНО: Все дальнейшие действия выполняем в Developer PowerShell for VS 2022 или с активированными переменными окружения, выполнив в PowerShell:

cmd /c '"C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat" >nul 2>&1 && set' | ForEach-Object { if ($_ -match '^([^=]+)=(.*)$') { [System.Environment]::SetEnvironmentVariable($matches[1], $matches[2], 'Process') } }

Создание виртуального окружения (venv)

Все действия будем производить внутри виртуального окружения. Создаем и активируем его:

python -m venv .venv
.\.venv\Scripts\activate

Установка ROCm 7 и PyTorch

Первым делом установим ROCm 7. Если у вас одна видеокарта, просто устанавливаем последние пакеты этой командой:

pip install --index-url https://rocm.nightlies.amd.com/v2/gfx110X-all/ --pre torch torchaudio torchvision rocm[devel] --no-cache

Если во время установки возникают ошибки, связанные с невозможностью загрузить правильную версию torch, скачайте нужные .whl пакеты вручную и установите их командой вида:

pip install "C:\Path\To\rocm_sdk_core-7.12.xxxxxx-py3-none-win_amd64.whl" "C:\Path\To\torch-2.11.0a0+rocm7.12.xxxx-cp312-cp312-win_amd64.whl" ...

Примечание: Если у вас iGPU и dGPU от AMD и вы хотите, чтобы PyTorch работал на обеих карточках, вам необходимо собрать свою сборку ROCm. О том, как это сделать, я писал в прошлой статье.

Для дальнейших действий установим переменные окружения, необходимые для активации Triton под AMD. Сохраните код ниже в файл set_env.ps1 и выполните его (.\set_env.ps1).

set_env.ps1

# set_env.ps1
$ROCM_ROOT = (rocm-sdk path --root).Trim()
$ROCM_BIN = (rocm-sdk path --bin).Trim()
$env:ROCM_HOME = $ROCM_ROOT
$env:PATH = "$ROCM_ROOT\lib\llvm\bin;$ROCM_BIN;$env:PATH"
$env:CC = "clang-cl"
$env:CXX = "clang-cl"
$env:DISTUTILS_USE_SDK = "1"
$env:TRITON_PRINT_AUTOTUNING=1
$env:TRITON_CACHE_AUTOTUNING=1
$env:TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1

Установка Triton

Устанавливаем нужную версию Triton:

pip install triton-windows==3.6.0.post25

Примечание: Wheel-файлы Triton можно также скачать вручную со ссылке на GitHub Actions.

Проверка установки Triton:
Создайте файл test_triton.py со следующим содержимым:

test_triton.py

import torch
import triton
print(torch.zeros(3).cuda())
print("If you see tensor([0., 0., 0.], device='cuda:0'), then it works")

Запустите его:

python .\test_triton.py
# Вы должны увидеть это:
# tensor([0., 0., 0.], device='cuda:0')
# If you see tensor([0., 0., 0.], device='cuda:0'), then it works

Установка Flash-Attention

Клонируем репозиторий и собираем проект. Убедитесь, что выполнили set_env.ps1.

git clone https://github.com/Dao-AILab/flash-attention
cd flash-attention
python setup.py bdist_wheel
pip install .\dist\flash_attn-2.8.3-py3-none-any.whl

Важно: Если во время сборки начинают скачиваться composable_kernels, значит, вы забыли выполнить set_env.ps1.

Установка Sage-Attention

Для работы с ROCm нужна немного измененная версия 1.0.6 из репозитория ComfyUI-Zluda. Для удобства я вынес ее в отдельный репозиторий. Подсмотрено тут.

Установка из готового пакета:

pip install https://github.com/guinmoon/SageAttention-Rocm7/releases/download/v1.0.6_rocm7/sageattention-1.0.6-py3-none-any.whl

Или сборка вручную из оригинального репозитория + патчи ComfyUI-Zluda

git clone https://github.com/guinmoon/SageAttention-Rocm7
cd SageAttention-Rocm7
# Сборка и установка
python setup.py bdist_wheel
pip install .\dist\sageattention-1.0.6-py3-none-any.whl

Установка bitsandbytes

Способ 1: Из готового пакета:

pip install https://github.com/guinmoon/bitsandbytes_win_rocm/releases/download/v0.49.2_romc7.12/bitsandbytes-0.49.2.dev0-cp312-cp312-win_amd64.whl

Способ 2: Сборка из исходников

Клонируем форк и собираем проект

# Убедитесь, что set_env.ps1 выполнен!
git clone https://github.com/guinmoon/bitsandbytes_win_rocm
cd bitsandbytes_win_rocm
cmake -G Ninja -DCOMPUTE_BACKEND=hip -S . -B bitsandbytes -DCMAKE_BUILD_TYPE="Release" -DBNB_ROCM_ARCH="gfx1100;gfx1102;gfx1103" -DHIP_PLATFORM="amd" -DCMAKE_CXX_COMPILER="$env:ROCM_PATH/lib/llvm/bin/clang++.exe" -DCMAKE_PREFIX_PATH="$env:ROCM_PATH" -DCMAKE_HIP_COMPILER="$env:ROCM_PATH/lib/llvm/bin/clang++.exe" -DCMAKE_SHARED_LINKER_FLAGS="-L$env:ROCM_PATH/lib -lamdhip64 -lrocblas"
cmake --build bitsandbytes -j
python setup.py bdist_wheel
pip install .\dist\bitsandbytes-0.49.2.dev0-cp312-cp312-win_amd64.whl

После команды cmake --build должно появиться сообщение Linking CXX shared library libbitsandbytes_rocm72.dll.

Внимание: Если у вас не RX 7900 XT/XTX, замените gfx1100 на свою архитектуру.

Проверка установки bitsandbytes:
Создайте файл test_bnb.py

test_bnb.py

import bitsandbytes as bnb
print(f"BNB: {bnb.envs.hip_get_device_properties(0)['gcnArchName']}")
print(f"warpSize: {bnb.envs.hip_get_device_properties(0)['warpSize']}")
print("SUCCESS!")
print("Installation was successful!")

Запустите его:

python .\test_bnb.py
# BNB: gfx1100
# warpSize: 32
# SUCCESS!
# Installation was successful!

Установка ComfyUI

Теперь установим ComfyUI для тестирования:

git clone https://github.com/Comfy-Org/ComfyUI
cd ComfyUI
pip install -r .\requirements.txt

Запускаем и смотрим результаты

Я запущу ComfyUI с минимальным набором параметров. Мой скрипт запуска (run_comfy.ps1) выглядит так:

run_comfy.ps1

.\set_env.ps1 # Активируем переменные окружения
$env:HIP_VISIBLE_DEVICES = "1" # Только если у вас iGPU от AMD. Или используйте --cuda-device 1
$env:FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"
$env:FLASH_ATTENTION_TRITON_AMD_AUTOTUNE="FALSE"
$env:TORCH_BLAS_PREFER_CUBLASLT=0
$env:COMFYUI_ENABLE_MIOPEN=0

# Запускаем ComfyUI
python .\main.py --normalvram --use-pytorch-cross-attention
# или --use-flash-attention
# или --use-sage-attention

Тест 1: Flux1.Dev fp8, 1920x1080

Workflow
Стоит оговориться: я использую модель в fp8, однако карты RDNA 3 не поддерживают операции rocWMMA с fp8, и модель транслируется в BF16.

Результаты 3 генераций:

Pytorch-Cross-Attention (SDPA):
1 генерация: 116.05 seconds, 2 генерация: 87.76 seconds, 3 генерация: 88.27 seconds
Flash-Attention (--use-flash-attention):
1 генерация: 103.02 seconds, 2 генерация: 102.17 seconds, 3 генерация: 101.22 seconds
Sage-Attention (--use-sage-attention):
1 генерация: 89.82 seconds, 2 генерация: 73.97 seconds, 3 генерация: 74.43 seconds

Вывод: Flash-Attention незначительно выигрывает у SDPA, однако Sage-Attention дает ощутимый прирост производительности на последующих прогонах.

Тест 2: LTX-2

Workflow
Примечание: Файл taeltx_2.safetensors необходимо положить в директорию vae_approx.

Результаты для разрешения 960х720 (один прогон):

Pytorch-Cross-Attention (SDPA): 277.47 seconds
Flash-Attention: 286.93 seconds
Sage-Attention: 231.45 seconds

Вывод: В этом тесте Sage-Attention так же показал себя лучше всех.

Finetune Gemma 3: QLoRA с bitsandbytes (NF4)

Устанавливаем дополнительные зависимости для тренировки:

pip install jq transformers==4.57.5 peft==0.18.1 accelerate safetensors sentencepiece huggingface-hub trl==0.26.2 einops tqdm==4.67.1 traitlets==5.14.3

Запускаем дообучение Gemma 3 4B с квантованием скриптом Qlora.py. В скрипте можно менять реализацию attention:

С flash_attention_2

model = AutoModelForCausalLM.from_pretrained(
    MODEL, quantization_config=bnb_config, device_map="auto", attn_implementation="flash_attention_2"
)

Peak training memory qlora: 13.34 GB
--- 00:32:37 elapsed ---

С sdpa

model = AutoModelForCausalLM.from_pretrained(
    MODEL, quantization_config=bnb_config, device_map="auto", attn_implementation="sdpa"
)

Peak training memory qlora: 13.34 GB
--- 00:15:32 elapsed ---

Наблюдение: По моим замерам, sdpa в этом скрипте отрабатывает намного быстрее, чем flash_attention_2.

Для инференса дообученной модели с LoRA-адаптером запустите:

python inference_q.py

Заключение

Заработало: Теперь на Windows с картами AMD можно снова использовать Triton только теперь без Zluda.
Sage-Attention — лидер: В тестах генерации изображений Sage-Attention стабильно показывает лучшую производительность по сравнению с SDPA и Flash-Attention.
Не всегда лучше: Как показали тесты Flash-Attention 2, на мое удивление, отрабатывает хуже чем стандартная sdpa.
Экосистема растет, но: По прежнему остается множество программных продуктов которые нельзя завести под Windows с ROCm, например onnxruntime или nunchaku, однако сейчас дела с ROCm явно обстоят лучше чем пару лет назад.

Главные проблемы пока заключаются в необходимости ручной сборки многих компонентов и отсутствии официальной поддержки со стороны AMD и многих разработчиков библиотек.

Все скрипты используемые в статье можно найти в моем репозитории.

P.S. Я так понимаю, что с RX 7900 XT, в том же ценовом сегменте на б/у рынке, находится RTX 4070 super, однако у нее чуть меньше памяти. Если кто из читателей является счастливым владельцем оной, напишите пожалуйста в комментариях сильно ли проигрывает по скорости генерации и прочим параметрам карта от красных.