[PR #143] [MERGED] HIP port of #117: int8 Sol-Attn sparse attention #175

New issue

Closed

opened 2026-10-26 22:21:14 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:21:14 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/143
Author: @0xDELUXA
Created: 8/30/2026
Status: ✅ Merged
Merged: 9/1/2026
Merged by: @comfyanonymous

Base: main ← Head: amd/hip-sol-attn

📝 Commits (2)

f622acc Restore the registry's backend selection between tests
d9630b8 HIP port of #117: int8 Sol-Attn sparse attention

📊 Changes

18 files changed (+2432 additions, -114 deletions)

View changed files

📝 README.md (+1 -1)
📝 comfy_kitchen/backends/hip/CMakeLists.txt (+6 -0)
📝 comfy_kitchen/backends/hip/__init__.py (+381 -0)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+258 -0)
📝 comfy_kitchen/backends/hip/launchers.h (+22 -0)
📝 comfy_kitchen/backends/hip/sage_attention/int8_attn.hip (+0 -58)
📝 comfy_kitchen/backends/hip/sage_attention/sage_common.h (+111 -0)
➕ comfy_kitchen/backends/hip/sage_attention/sol_attn.hip (+234 -0)
➕ comfy_kitchen/backends/hip/sage_attention/sol_attn_exact.hip (+203 -0)
➕ comfy_kitchen/backends/hip/sage_attention/sol_attn_preprocess.hip (+268 -0)
➕ comfy_kitchen/backends/hip/sage_attention/sol_attn_producer.hip (+131 -0)
➕ comfy_kitchen/backends/hip/sage_attention/sol_attn_route.hip (+259 -0)
➕ comfy_kitchen/backends/hip/sage_attention/sol_attn_vtranspose.hip (+87 -0)
➕ comfy_kitchen/backends/hip/sage_attention/sol_common.h (+243 -0)
📝 tests/conftest.py (+23 -0)
📝 tests/test_constraints.py (+1 -1)
📝 tests/test_sol_attn.py (+189 -51)
📝 tests/test_tensor.py (+15 -3)

📄 Description

Brings the HIP backend back to parity with #117.

A capability port: sol_attn and sol_attn_chunked did not exist on HIP, so ROCm fell through to the O(T^2) eager reference, which refuses video length outright. Both entry points, the workspace plan, the chunked QKV producer, key_bias, block_len, tail=False and the VSA coarse gate are all covered, and tests/test_sol_attn.py now runs against whichever fused backend is built.

The kernels are a reimplementation. The CUDA carriers are stored under sol::perm_d and sol::perm_key, which exist only to turn an m16n8k32 lane's two operand words into one 8-byte load; RDNA's 16x16x16 WMMA already hands a lane a contiguous K slice, so both permutations are dropped and the scores are evaluated transposed the way int8_attn.hip does, which puts one query per lane and makes P^T the operand of the second MMA directly. Everything goes through the existing MmaInt8 / MmaBf16 policies, so gfx11 and gfx12 both fall out of the port.

On a gfx1200 card, (1, T, H, 128) bf16 at tau=1.4, against bf16 SDPA with the AOTriton flash backend enabled: 11.3x at T=32768 (113.3 ms to 10.0 ms), and 2.3x to 7.1x against this backend's own fused INT8 attention. Cosine against the eager Sol-Attn reference is 0.9999. Compile-verified on all 17 targets in architectures.json with no scratch spill; run-verified on RDNA4 only. RDNA2 is unsupported because it has no matrix cores.

One fix rides along in its own commit because it is what hid the rest. test_backend_priority changes registry priority without restoring it, leaving HIP registered but unreachable and causing two INT8 tests to run against eager incorrectly. tests/conftest.py now restores priority and the disabled set per test. Those INT8 tests now compare against the exact dequantized product instead of eager, since the fused kernels accumulate in INT32. ROCm failures drop from 113 to 101 - the same failures already present on main.

_{🔄 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/143 **Author:** [@0xDELUXA](https://github.com/0xDELUXA) **Created:** 8/30/2026 **Status:** ✅ Merged **Merged:** 9/1/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `amd/hip-sol-attn` --- ### 📝 Commits (2) - [`f622acc`](https://github.com/Comfy-Org/comfy-kitchen/commit/f622acc35738d056df867e766b9606ff52b8263c) Restore the registry's backend selection between tests - [`d9630b8`](https://github.com/Comfy-Org/comfy-kitchen/commit/d9630b8c6524c791671138520c84afa358ef08b3) HIP port of #117: int8 Sol-Attn sparse attention ### 📊 Changes **18 files changed** (+2432 additions, -114 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+1 -1) 📝 `comfy_kitchen/backends/hip/CMakeLists.txt` (+6 -0) 📝 `comfy_kitchen/backends/hip/__init__.py` (+381 -0) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+258 -0) 📝 `comfy_kitchen/backends/hip/launchers.h` (+22 -0) 📝 `comfy_kitchen/backends/hip/sage_attention/int8_attn.hip` (+0 -58) 📝 `comfy_kitchen/backends/hip/sage_attention/sage_common.h` (+111 -0) ➕ `comfy_kitchen/backends/hip/sage_attention/sol_attn.hip` (+234 -0) ➕ `comfy_kitchen/backends/hip/sage_attention/sol_attn_exact.hip` (+203 -0) ➕ `comfy_kitchen/backends/hip/sage_attention/sol_attn_preprocess.hip` (+268 -0) ➕ `comfy_kitchen/backends/hip/sage_attention/sol_attn_producer.hip` (+131 -0) ➕ `comfy_kitchen/backends/hip/sage_attention/sol_attn_route.hip` (+259 -0) ➕ `comfy_kitchen/backends/hip/sage_attention/sol_attn_vtranspose.hip` (+87 -0) ➕ `comfy_kitchen/backends/hip/sage_attention/sol_common.h` (+243 -0) 📝 `tests/conftest.py` (+23 -0) 📝 `tests/test_constraints.py` (+1 -1) 📝 `tests/test_sol_attn.py` (+189 -51) 📝 `tests/test_tensor.py` (+15 -3) </details> ### 📄 Description Brings the HIP backend back to parity with #117. A capability port: `sol_attn` and `sol_attn_chunked` did not exist on HIP, so ROCm fell through to the O(T^2) eager reference, which refuses video length outright. Both entry points, the workspace plan, the chunked QKV producer, `key_bias`, `block_len`, `tail=False` and the VSA coarse gate are all covered, and `tests/test_sol_attn.py` now runs against whichever fused backend is built. The kernels are a reimplementation. The CUDA carriers are stored under `sol::perm_d` and `sol::perm_key`, which exist only to turn an `m16n8k32` lane's two operand words into one 8-byte load; RDNA's `16x16x16` WMMA already hands a lane a contiguous K slice, so both permutations are dropped and the scores are evaluated transposed the way `int8_attn.hip` does, which puts one query per lane and makes `P^T` the operand of the second MMA directly. Everything goes through the existing `MmaInt8` / `MmaBf16` policies, so gfx11 and gfx12 both fall out of the port. On a gfx1200 card, `(1, T, H, 128)` bf16 at `tau=1.4`, against bf16 SDPA with the AOTriton flash backend enabled: 11.3x at T=32768 (113.3 ms to 10.0 ms), and 2.3x to 7.1x against this backend's own fused INT8 attention. Cosine against the eager Sol-Attn reference is 0.9999. Compile-verified on all 17 targets in `architectures.json` with no scratch spill; run-verified on RDNA4 only. RDNA2 is unsupported because it has no matrix cores. One fix rides along in its own commit because it is what hid the rest. `test_backend_priority` changes registry priority without restoring it, leaving HIP registered but unreachable and causing two INT8 tests to run against eager incorrectly. `tests/conftest.py` now restores priority and the disabled set per test. Those INT8 tests now compare against the exact dequantized product instead of eager, since the fused kernels accumulate in INT32. ROCm failures drop from 113 to 101 - the same failures already present on `main`. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 22:21:14 +00:00