flash-attention-v2-RDNA3-minimal

a simple Flash Attention v2 implementation with ROCM (RDNA3 GPU, roc wmma), mainly used for stable diffusion(ComfyUI) in Windows ZLUDA environments.

Build and Test

minimum integration:

──rocwmma_fattn
   │  FlashAttn.py
   │  host.cpp
   │  kernel_bf16.cu
   │  kernel_fp16.cu
   └─ zluda_hijack_torch_hip_ext.py

Linux with rocm:

run test: python bench_with_sdpa.py

Windows with zluda

Need MSVC Compiler, AMD HIP SDK and rocWMMA Library.

Install rocwmma library: https://github.com/ROCm/rocWMMA

clone it and copy library/include/rocwmma to HIP SDK installation path of include folder

In cmd.exe, run vcvars64.bat to active MSVC Environment, then run zluda -- python bench_with_sdpa.py

Pre-build Extension

Tested work with PyTorch 2.2.1 + cu118 windows zluda, gfx1100 GPU

comfyui: https://github.com/Repeerc/ComfyUI-flash-attention-rdna3-win-zluda

webui: https://github.com/Repeerc/sd-webui-flash-attention-zluda-win

To do

• [x] backward pass
• [x] causal mask (need more optimization)
• [ ] unaligned 32x seqlen padding optimization
• [ ] Load tile into LDS
• [ ] attention bias
• [ ] matrix multiplication optimization
• [x] fix poor performance in BF16
• [ ] ...

Benchmark

OS: Windows 11

GPU: 7900xtx (gfx1100)

PyTorch 2.2.1 + CU118 ZLUDA, Python 3.10, HIP 5.7.1

FP16, causal = False

Triton build from: https://github.com/triton-lang/triton

git hash: 47fc046ff29c9ea2ee90e987c39628a540603c8f

test use Triton windows pre-build version: https://github.com/Repeerc/triton-windows-amdgpu

Triton offcial version use 06-fused-attention.py

CK-based(Composable Kernel) flash attention version compiled from: https://github.com/ROCm/flash-attention/tree/howiejay/navi_support

CK-based flash attention windows porting: https://github.com/Repeerc/flash-attn-composable-kernel-gfx110x-windows-port

seqlen with 32x aligened

[B N H D] format rearrange and contiguous to [B H N D]

seqlen without 32x aligened

[B N H D] format rearrange and contiguous to [B H N D]

fwd+bwd

FP16, causal = True

Performance in Stable Diffusion (ComfyUI)

OS: Windows 11

GPU: 7900xtx (gfx1100)

PyTorch 2.2.1 + CU118 ZLUDA, Python 3.10

Sampler: Euler

SDXL U-Net Lora training

unet_lr = 0.0001
lr_scheduler = "constant"
lr_warmup_steps = 0
optimizer_type = "AdamW"
network_dim = 32
network_alpha = 32
seed = 1337
mixed_precision = "fp16"
full_fp16 = false
full_bf16 = false
fp8_base = true
no_half_vae = false