ROCm 7 on Windows: Triton, FlashAttention, SageAttention & bitsandbytes

A guide to enabling advanced PyTorch extensions (FlashAttention, SageAttention, bitsandbytes) for AMD GPUs (RDNA 3/gfx1100+) on Windows using ROCm 7.

🚧 Status: Experimental | Requires manual compilation and specific versions.

✨ Features

• Triton Support: triton-windows==3.6.0.post25 with ROCm support.
• Accelerated Attention:
  • flash-attention (Dao-AILab)
  • sage-attention==1.0.6 (patched for ROCm)
• 4-bit Quantization: bitsandbytes compiled for ROCm 7 (NF4, QLoRA).
• Tested Frameworks: Works with ComfyUI and standard PyTorch training scripts.

📋 Prerequisites

• GPU: AMD RDNA 3 Architecture (gfx1100+, e.g., RX 7900 XT/XTX).
• OS: Windows 10/11
• Python: 3.12 (recommended)
• ROCm: 7.12 nightly builds.

Install Build Tools (using 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 python --version=3.12.0 -y

⚙️ Installation

1. Install PyTorch with ROCm 7:

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

2. Set up build environment: Save the following as set_env.ps1 and run it before any compilation.

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

3. Install Triton:

   pip install triton-windows==3.6.0.post25

4. Install FlashAttention:

   git clone https://github.com/Dao-AILab/flash-attention
   cd flash-attention
   .\..\set_env.ps1 # Source the env script
   python setup.py bdist_wheel
   pip install dist\flash_attn-*.whl

5. Install SageAttention (patched for ROCm):

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

   Form manual compilation:

   git clone https://github.com/thu-ml/SageAttention -b sageattention-1
   # Скачиваем необходимые патчи
   curl -o .\sageattention\attn_qk_int8_per_block.py https://raw.githubusercontent.com/patientx/ComfyUI-Zluda/refs/heads/master/comfy/customzluda/sa/attn_qk_int8_per_block.py
   curl -o .\sageattention\attn_qk_int8_per_block_causal.py https://raw.githubusercontent.com/patientx/ComfyUI-Zluda/refs/heads/master/comfy/customzluda/sa/attn_qk_int8_per_block_causal.py
   curl -o .\sageattention\quant_per_block.py https://raw.githubusercontent.com/patientx/ComfyUI-Zluda/refs/heads/master/comfy/customzluda/sa/quant_per_block.py
   # Сборка и установка
   python setup.py bdist_wheel
   pip install .\dist\sageattention-1.0.6-py3-none-any.whl
   

6. Install bitsandbytes (Pre-built wheel):

   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

   For manual compilation:

   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
   

🚀 Usage Example (ComfyUI)

1. Install ComfyUI:

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

2. Run with different attention backends:

   # Activate environment vars
   .\path\to\set_env.ps1
   # Run ComfyUI
   python main.py --normalvram --use-sage-attention # or --use-flash-attention

📊 Performance (ComfyUI)

Tested on an AMD RX 7900 XTX.

Model / Task	SDPA (s)	FlashAttention (s)	SageAttention (s)
Flux1.Dev (1080p)	88.3	101.2	74.4
LTX-2	277.5	286.9	231.5

SageAttention provides a significant speedup in stable diffusion workflows.

🔧 Fine-tuning Example (QLoRA)

Install extra dependencies:

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

Use Qlora.py for training and inference_q.py for inference. Create secret.py with hf_token = 'hf_**********'

⚠️ Important Notes

• This is an unofficial, community-supported setup.
• The architecture (gfx1100) must be specified correctly in build commands.
• Performance can vary significantly between different models and tasks.
• Always source the set_env.ps1 script before compiling any component.

📚 References & Credits

• Original Detailed Guide (Habr): Russian Article
• Patched bitsandbytes fork: guinmoon/bitsandbytes_win_rocm
• Triton Windows Wheels: Comfy-Org/wheels