[PR #115] [MERGED] HIP port of #111: bf16 flash attention decode #154

New issue

Closed

opened 2026-10-26 22:20:55 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:20:55 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/115
Author: @0xDELUXA
Created: 8/13/2026
Status: ✅ Merged
Merged: 8/18/2026
Merged by: @comfyanonymous

Base: main ← Head: amd/hip-flash-attention-decode

📝 Commits (1)

66d2ac2 HIP port of #111: bf16 flash attention decode

📊 Changes

7 files changed (+698 additions, -10 deletions)

View changed files

📝 comfy_kitchen/backends/hip/CMakeLists.txt (+1 -0)
📝 comfy_kitchen/backends/hip/__init__.py (+45 -0)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+133 -0)
📝 comfy_kitchen/backends/hip/launchers.h (+11 -0)
➕ comfy_kitchen/backends/hip/ops/flash_decode.hip (+366 -0)
📝 comfy_kitchen/flash_attention.py (+34 -10)
📝 tests/test_flash_attention.py (+108 -0)

📄 Description

Adds flash_attention_decode to the HIP backend as an independent performance port of #111. It keeps the same launcher signature, split-then-combine structure, and Python API as CUDA, with a native HIP kernel in ops/flash_decode.hip.

comfy_kitchen/flash_attention.py now selects the backend like sage_attention.py. The GQA reshape, split heuristic, and workspace allocation remain shared.

Kernel

The kernel targets GQA decode, which is primarily bandwidth- and latency-bound.

Key optimizations:

8 waves per block: Improves latency hiding. Benchmarks across 11 shapes showed 8 waves are 1.24x faster than 4 on average. 16 waves exceed the LDS limit at the widest tile.
Branch-free prefetch: Uses a measured 4-16 key prefetch depth depending on tile size.
DPP wave reduction: Replaces the LDS-based ds_bpermute reduction with a bit-identical DPP sequence, improving GQA performance by 1.7-2.2x.
No WMMA: plain FMA, no matrix cores. Gated to RDNA 3 and RDNA 4, since RDNA 2 has no bf16.

Performance

Compared with scaled_dot_product_attention(enable_gqa), BF16, head_dim 128:

batch	KV heads	groups	length	kernel	SDPA	speedup
1	8	4	4096	0.120 ms	0.160 ms	1.34x
1	8	4	16384	0.291 ms	0.661 ms	2.28x
1	8	4	32768	0.506 ms	1.270 ms	2.51x
4	8	4	4096	0.275 ms	0.387 ms	1.41x
4	8	4	16384	0.943 ms	1.439 ms	1.53x
8	8	4	8192	0.912 ms	1.398 ms	1.53x
1	32	1	8192	0.532 ms	0.451 ms	0.85x
4	32	1	8192	1.872 ms	1.737 ms	0.93x
1	4	8	16384	0.266 ms	0.558 ms	2.10x

The main gains come from GQA reuse and avoiding unnecessary cache reads.

Variable-length decode shows larger gains because SDPA must process the full cache capacity:

Fill	Kernel	SDPA	Speedup
256	0.198 ms	2.831 ms	14.3x
1024	0.202 ms	2.831 ms	14.0x
4096	0.266 ms	2.831 ms	10.6x
16384	0.939 ms	2.831 ms	3.0x
32768	1.842 ms	2.831 ms	1.5x

Correctness

Upstream flash-attention tests now run on ROCm, including graph capture.
54 shapes x 5 split configurations = 270 configurations, with zero failures.
Worst relative max-abs error: 0.0035, consistent with BF16 rounding.
All 17 targets build into the fat binary.
Compile-verified on gfx1030, gfx1100, gfx1151, and gfx1200.
Run-verified on 1200.
Zero scratch usage and zero ds_bpermute.
Full-suite failure set is unchanged: 107 pre-existing ROCm failures.
Installed wheel tests: 64 passed, 9 skipped, up from 62 passed and 11 skipped.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/Comfy-Org/comfy-kitchen/pull/115 **Author:** [@0xDELUXA](https://github.com/0xDELUXA) **Created:** 8/13/2026 **Status:** ✅ Merged **Merged:** 8/18/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `amd/hip-flash-attention-decode` --- ### 📝 Commits (1) - [`66d2ac2`](https://github.com/Comfy-Org/comfy-kitchen/commit/66d2ac298e783e88a0d708a0b06240b585b8de1a) HIP port of #111: bf16 flash attention decode ### 📊 Changes **7 files changed** (+698 additions, -10 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/hip/CMakeLists.txt` (+1 -0) 📝 `comfy_kitchen/backends/hip/__init__.py` (+45 -0) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+133 -0) 📝 `comfy_kitchen/backends/hip/launchers.h` (+11 -0) ➕ `comfy_kitchen/backends/hip/ops/flash_decode.hip` (+366 -0) 📝 `comfy_kitchen/flash_attention.py` (+34 -10) 📝 `tests/test_flash_attention.py` (+108 -0) </details> ### 📄 Description Adds `flash_attention_decode` to the HIP backend as an independent performance port of #111. It keeps the same launcher signature, split-then-combine structure, and Python API as CUDA, with a native HIP kernel in `ops/flash_decode.hip`. `comfy_kitchen/flash_attention.py` now selects the backend like `sage_attention.py`. The GQA reshape, split heuristic, and workspace allocation remain shared. ## Kernel The kernel targets GQA decode, which is primarily bandwidth- and latency-bound. Key optimizations: - 8 waves per block: Improves latency hiding. Benchmarks across 11 shapes showed 8 waves are 1.24x faster than 4 on average. 16 waves exceed the LDS limit at the widest tile. - Branch-free prefetch: Uses a measured 4-16 key prefetch depth depending on tile size. - DPP wave reduction: Replaces the LDS-based `ds_bpermute` reduction with a bit-identical DPP sequence, improving GQA performance by 1.7-2.2x. - No WMMA: plain FMA, no matrix cores. Gated to RDNA 3 and RDNA 4, since RDNA 2 has no bf16. ## Performance Compared with `scaled_dot_product_attention(enable_gqa)`, BF16, head_dim 128: | batch | KV heads | groups | length | kernel | SDPA | speedup | | --- | --- | --- | --- | --- | --- | --- | | 1 | 8 | 4 | 4096 | 0.120 ms | 0.160 ms | 1.34x | | 1 | 8 | 4 | 16384 | 0.291 ms | 0.661 ms | 2.28x | | 1 | 8 | 4 | 32768 | 0.506 ms | 1.270 ms | 2.51x | | 4 | 8 | 4 | 4096 | 0.275 ms | 0.387 ms | 1.41x | | 4 | 8 | 4 | 16384 | 0.943 ms | 1.439 ms | 1.53x | | 8 | 8 | 4 | 8192 | 0.912 ms | 1.398 ms | 1.53x | | 1 | 32 | 1 | 8192 | 0.532 ms | 0.451 ms | 0.85x | | 4 | 32 | 1 | 8192 | 1.872 ms | 1.737 ms | 0.93x | | 1 | 4 | 8 | 16384 | 0.266 ms | 0.558 ms | 2.10x | The main gains come from GQA reuse and avoiding unnecessary cache reads. Variable-length decode shows larger gains because SDPA must process the full cache capacity: | Fill | Kernel | SDPA | Speedup | | --- | --- | --- | --- | | 256 | 0.198 ms | 2.831 ms | 14.3x | | 1024 | 0.202 ms | 2.831 ms | 14.0x | | 4096 | 0.266 ms | 2.831 ms | 10.6x | | 16384 | 0.939 ms | 2.831 ms | 3.0x | | 32768 | 1.842 ms | 2.831 ms | 1.5x | ## Correctness - Upstream flash-attention tests now run on ROCm, including graph capture. - 54 shapes x 5 split configurations = 270 configurations, with zero failures. - Worst relative max-abs error: 0.0035, consistent with BF16 rounding. - All 17 targets build into the fat binary. - Compile-verified on gfx1030, gfx1100, gfx1151, and gfx1200. - Run-verified on 1200. - Zero scratch usage and zero `ds_bpermute`. - Full-suite failure set is unchanged: 107 pre-existing ROCm failures. - Installed wheel tests: 64 passed, 9 skipped, up from 62 passed and 11 skipped. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 22:20:55 +00:00