[PR #104] [MERGED] HIP port of #103: optimized int8 attention #149

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: amd/hip-int8-attention

📝 Commits (1)

4cd8ad2 HIP port of #103: int8 attention

📊 Changes

11 files changed (+2187 additions, -3 deletions)

View changed files

📝 README.md (+1 -1)
📝 comfy_kitchen/backends/hip/CMakeLists.txt (+4 -0)
📝 comfy_kitchen/backends/hip/__init__.py (+177 -0)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+347 -0)
📝 comfy_kitchen/backends/hip/mma.h (+29 -0)
➕ comfy_kitchen/backends/hip/sage_attention/int8_attn.hip (+584 -0)
➕ comfy_kitchen/backends/hip/sage_attention/quant_qk_int8.hip (+521 -0)
➕ comfy_kitchen/backends/hip/sage_attention/quant_v_int8.hip (+209 -0)
➕ comfy_kitchen/backends/hip/sage_attention/sage_common.h (+231 -0)
📝 comfy_kitchen/sage_attention.py (+65 -2)
📝 tests/test_int8_attention.py (+19 -0)

📄 Description

Overview

Correctness/performance port of #103 to the HIP backend, enabling int8_attention on AMD RDNA3, 3.5, 4. It declines on RDNA2, which has no matrix cores. The main performance gains are seen on RDNA4.

The kernel is a separate implementation, not a direct translation. RDNA's WMMA layout is transposed relative to mma.m16n8k16, so the kernel computes S^T = K @ Q^T and builds O^T = V^T @ P^T. This keeps the online softmax per-lane and, on gfx12, lets the score accumulator feed directly into the second MMA without going through shared memory. gfx11 needs two shuffles and four v_perm_b32 instructions instead.

The fragment and operand layout is inspired by "Add native gfx12 attention backend" by @jammm (thu-ml/SageAttention#368).

Both matmuls remain int8, as in #103, so the implementation does not depend on fp8 WMMA and RDNA3 is supported as well as RDNA4.

Quantization keeps the CUDA numerics while adapting scale granularity to the RDNA fragment layout: one scale per query row and one per 16 adjacent keys. V needs no key permutation.

Performance

Several additional optimizations were tested and reverted because they regressed performance or caused spills. A smaller performance change can be expected on RDNA3, while RDNA2 is unsupported because it has no matrix cores.

Measured on gfx1200 against torch.nn.functional.scaled_dot_product_attention with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1, best of seven alternating rounds:

1.35x at 4K keys, head_dim=128
1.55x at 8K
1.66x on 4K causal
2.10x at head_dim=256
1.14x on Wan-shaped 14040x512 cross attention

The attention kernel alone is 1.67-1.79x faster; quantization accounts for the remaining 15-20%. Below ~2K sequence length it is slower than SDPA due to host overhead and multiple kernel launches/allocations.

Validation

tests/test_int8_attention.py passes unchanged apart from the compute-capability gating test, which now has a HIP counterpart. A separate 71-case sweep covering different head dimensions, dtypes, masks, causal/GQA/ragged inputs, smooth_k, convrot and split prequantization stays within 0.015 NRMSE of fp32 SDPA, with fused and split paths bitwise identical.

The kernel compile-verifies for all 17 architectures in architectures.json as a single fat binary, with no scratch spill on gfx12. I runtime-verified and benchmarked RDNA4, while @tvukovic-amd validated RDNA3/3.5.

Trying it

The CI wheels include the HIP backend precompiled for every target in architectures.json, so no source build is required. Install the artifact for your platform and use it with ComfyUI. It additionally needs the --use-ck-attention flag.

_{🔄 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/104 **Author:** [@0xDELUXA](https://github.com/0xDELUXA) **Created:** 8/10/2026 **Status:** ✅ Merged **Merged:** 8/12/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `amd/hip-int8-attention` --- ### 📝 Commits (1) - [`4cd8ad2`](https://github.com/Comfy-Org/comfy-kitchen/commit/4cd8ad2603b611a71c88990fc274f1afa3652ca4) HIP port of #103: int8 attention ### 📊 Changes **11 files changed** (+2187 additions, -3 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+1 -1) 📝 `comfy_kitchen/backends/hip/CMakeLists.txt` (+4 -0) 📝 `comfy_kitchen/backends/hip/__init__.py` (+177 -0) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+347 -0) 📝 `comfy_kitchen/backends/hip/mma.h` (+29 -0) ➕ `comfy_kitchen/backends/hip/sage_attention/int8_attn.hip` (+584 -0) ➕ `comfy_kitchen/backends/hip/sage_attention/quant_qk_int8.hip` (+521 -0) ➕ `comfy_kitchen/backends/hip/sage_attention/quant_v_int8.hip` (+209 -0) ➕ `comfy_kitchen/backends/hip/sage_attention/sage_common.h` (+231 -0) 📝 `comfy_kitchen/sage_attention.py` (+65 -2) 📝 `tests/test_int8_attention.py` (+19 -0) </details> ### 📄 Description ## Overview Correctness/performance port of #103 to the HIP backend, enabling `int8_attention` on AMD RDNA3, 3.5, 4. It declines on RDNA2, which has no matrix cores. The main performance gains are seen on RDNA4. The kernel is a separate implementation, not a direct translation. RDNA's WMMA layout is transposed relative to `mma.m16n8k16`, so the kernel computes `S^T = K @ Q^T` and builds `O^T = V^T @ P^T`. This keeps the online softmax per-lane and, on gfx12, lets the score accumulator feed directly into the second MMA without going through shared memory. gfx11 needs two shuffles and four `v_perm_b32` instructions instead. The fragment and operand layout is inspired by "Add native gfx12 attention backend" by @jammm (thu-ml/SageAttention#368). Both matmuls remain int8, as in #103, so the implementation does not depend on fp8 WMMA and RDNA3 is supported as well as RDNA4. Quantization keeps the CUDA numerics while adapting scale granularity to the RDNA fragment layout: one scale per query row and one per 16 adjacent keys. V needs no key permutation. ## Performance Several additional optimizations were tested and reverted because they regressed performance or caused spills. A smaller performance change can be expected on RDNA3, while RDNA2 is unsupported because it has no matrix cores. Measured on gfx1200 against `torch.nn.functional.scaled_dot_product_attention` with `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1`, best of seven alternating rounds: - 1.35x at 4K keys, `head_dim=128` - 1.55x at 8K - 1.66x on 4K causal - 2.10x at `head_dim=256` - 1.14x on Wan-shaped 14040x512 cross attention The attention kernel alone is 1.67-1.79x faster; quantization accounts for the remaining 15-20%. Below ~2K sequence length it is slower than SDPA due to host overhead and multiple kernel launches/allocations. ## Validation `tests/test_int8_attention.py` passes unchanged apart from the compute-capability gating test, which now has a HIP counterpart. A separate 71-case sweep covering different head dimensions, dtypes, masks, causal/GQA/ragged inputs, `smooth_k`, `convrot` and split prequantization stays within 0.015 NRMSE of fp32 SDPA, with fused and split paths bitwise identical. The kernel compile-verifies for all 17 architectures in `architectures.json` as a single fat binary, with no scratch spill on gfx12. I runtime-verified and benchmarked RDNA4, while @tvukovic-amd validated RDNA3/3.5. ## Trying it The CI wheels include the HIP backend precompiled for every target in `architectures.json`, so no source build is required. Install the artifact for your platform and use it with ComfyUI. It additionally needs the `--use-ck-attention` flag. --- <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:52 +00:00