Triton, Flash-attension, Sage-attension и bitsandbytes с Rocm7 в 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).

Установка Triton

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

pip install triton-windows==3.6.0.post25

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

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

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

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

Установка 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: Сборка из исходников

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

Внимание: Если у вас не RX 7900 XT/XTX, замените gfx1100 на свою архитектуру (например, gfx1030 для RX 6700 XT).

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

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

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) выглядит так:

Тест 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:

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

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

python inference_q.py

Заключение

1. Заработало: Теперь на Windows с картами AMD можно использовать снова использовать Triton только теперь без Zluda.

2. Sage-Attention — лидер: В тестах генерации изображений Sage-Attention стабильно показывает лучшую производительность по сравнению с SDPA и Flash-Attention.

3. Не всегда лучше: Как показал тест с QLoRA, Flash-Attention 2 не всегда является серебряной пулей, и в некоторых задачах стоится проверять разные бэкенды.

4. Экосистема растет, но: По прежнему остается множество программных продуктов которые нельзя завести под Windows с ROCm, например onnxruntime или nunchaku, однако сейчас дела с ROCm явно обстоят лучше чем пару лет назад.

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

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

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