[PR #158] [MERGED] HIP port of #156: Sol-Attn token routing, plus a folded fp16 rounding and the test gate that hid it #185

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: amd/hip-sol-token-aug

📝 Commits (3)

3497a01 HIP port of #156: Sol-Attn token routing
bf180fd Keep the HIP RoPE fp16 rounding from being folded away by -ffast-math
b99d220 Run the shared test suites against the HIP backend

📊 Changes

20 files changed (+921 additions, -163 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+2 -2)
📝 comfy_kitchen/backends/cuda/__init__.py (+1 -2)
📝 comfy_kitchen/backends/hip/CMakeLists.txt (+1 -0)
📝 comfy_kitchen/backends/hip/__init__.py (+14 -21)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+34 -25)
📝 comfy_kitchen/backends/hip/launchers.h (+5 -5)
📝 comfy_kitchen/backends/hip/rope_math.h (+8 -3)
📝 comfy_kitchen/backends/hip/sage_attention/sol_attn.hip (+104 -33)
📝 comfy_kitchen/backends/hip/sage_attention/sol_attn_exact.hip (+75 -25)
📝 comfy_kitchen/backends/hip/sage_attention/sol_attn_producer.hip (+13 -3)
📝 comfy_kitchen/backends/hip/sage_attention/sol_attn_route.hip (+59 -7)
➕ comfy_kitchen/backends/hip/sage_attention/sol_attn_token.hip (+518 -0)
📝 comfy_kitchen/backends/hip/sage_attention/sol_attn_vtranspose.hip (+12 -5)
📝 comfy_kitchen/backends/hip/sage_attention/sol_common.h (+12 -0)
📝 tests/conftest.py (+2 -2)
📝 tests/test_adaln.py (+9 -9)
📝 tests/test_hip_wmma.py (+12 -3)
📝 tests/test_int8_input_act.py (+2 -2)
📝 tests/test_rope.py (+9 -9)
📝 tests/test_sol_attn.py (+29 -7)

📄 Description

Ports token routing from #156 to HIP, making token_aug work on RDNA instead of warning and running without it.

Adds sage_attention/sol_attn_token.hip with the five CUDA-equivalent kernels: centroids, tile list, two scoring passes, sort, and merge.
route now outputs the unrouted-block bitmap and best unrouted score.
vtranspose and the chunked producer write the row-major int8 V copy consumed by gathered token tiles.
The exact kernel walks each group's token list after routed blocks.
HIP workspace slots now match sol_attn.cu, so Python uses one plan for both backends.

This is a correctness port. RDNA WMMA already exposes a contiguous K slice, so token IDs index ki / ksb / vRow directly without CUDA's sol::perm_key permutation.

On a routing-adversarial synthetic case (T=16384, H=8, bf16), token_aug=256 improves dense-attention relative error from 0.648 -> 0.132 and DC bias from 0.165 -> 0.014, at 7.1 -> 8.1 ms. The selected tokens were independently re-derived from workspace centroid scores and verified sorted, within budget, candidate-only, and correctly cut at the histogram boundary.

The second commit fixes HIP round_fp16 under -ffast-math. Clang was folding the float -> half -> float round away, causing apply_rope to differ from eager by up to 9% of elements. Two inline-asm conversions restore exact rounding; fp16-freq cases are now bit-identical. The 2x24x4096x128 case goes 865 -> 1024 us; fp32/bf16 paths are unchanged and VGPR use drops by one. test_hip_wmma.py::test_rope_matches_eager now covers fp16 freqs and checks bit identity.

The third commit enables HIP in the test backend/device lists. This activates ~130 additional cases across AdaLN, RMS AdaLN, NA3D, FP8 quantization, int8 activations, and RoPE. All pass; non-HIP machines simply skip them.

test_rms_rope.py remains unchanged because its existing tolerance bands are already exceeded by eager on some seeds, while HIP tracks eager within 0.001 percentage points. Coverage remains in test_hip_wmma.py.

Suite result: 1483 passing vs 1349 before, with failures unchanged at 101 and none in HIP paths. Compile-verified on gfx1100/gfx1200 with no scratch spills; run-verified on gfx1200. All five token_aug tests in tests/test_sol_attn.py now run and pass.

_{🔄 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/158 **Author:** [@0xDELUXA](https://github.com/0xDELUXA) **Created:** 9/5/2026 **Status:** ✅ Merged **Merged:** 9/5/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `amd/hip-sol-token-aug` --- ### 📝 Commits (3) - [`3497a01`](https://github.com/Comfy-Org/comfy-kitchen/commit/3497a01e0a32ad8e14ae9c3b333c3664dd647b82) HIP port of #156: Sol-Attn token routing - [`bf180fd`](https://github.com/Comfy-Org/comfy-kitchen/commit/bf180fd8a169c7123a4b51cb38f7d2722b84b6de) Keep the HIP RoPE fp16 rounding from being folded away by -ffast-math - [`b99d220`](https://github.com/Comfy-Org/comfy-kitchen/commit/b99d2207b0865a3cc660ffb8265eadde3f891ead) Run the shared test suites against the HIP backend ### 📊 Changes **20 files changed** (+921 additions, -163 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+2 -2) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+1 -2) 📝 `comfy_kitchen/backends/hip/CMakeLists.txt` (+1 -0) 📝 `comfy_kitchen/backends/hip/__init__.py` (+14 -21) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+34 -25) 📝 `comfy_kitchen/backends/hip/launchers.h` (+5 -5) 📝 `comfy_kitchen/backends/hip/rope_math.h` (+8 -3) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_attn.hip` (+104 -33) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_attn_exact.hip` (+75 -25) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_attn_producer.hip` (+13 -3) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_attn_route.hip` (+59 -7) ➕ `comfy_kitchen/backends/hip/sage_attention/sol_attn_token.hip` (+518 -0) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_attn_vtranspose.hip` (+12 -5) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_common.h` (+12 -0) 📝 `tests/conftest.py` (+2 -2) 📝 `tests/test_adaln.py` (+9 -9) 📝 `tests/test_hip_wmma.py` (+12 -3) 📝 `tests/test_int8_input_act.py` (+2 -2) 📝 `tests/test_rope.py` (+9 -9) 📝 `tests/test_sol_attn.py` (+29 -7) </details> ### 📄 Description Ports token routing from #156 to HIP, making `token_aug` work on RDNA instead of warning and running without it. - Adds `sage_attention/sol_attn_token.hip` with the five CUDA-equivalent kernels: centroids, tile list, two scoring passes, sort, and merge. - `route` now outputs the unrouted-block bitmap and best unrouted score. - `vtranspose` and the chunked producer write the row-major int8 V copy consumed by gathered token tiles. - The exact kernel walks each group's token list after routed blocks. - HIP workspace slots now match `sol_attn.cu`, so Python uses one plan for both backends. This is a correctness port. RDNA WMMA already exposes a contiguous K slice, so token IDs index `ki` / `ksb` / `vRow` directly without CUDA's `sol::perm_key` permutation. On a routing-adversarial synthetic case (T=16384, H=8, bf16), `token_aug=256` improves dense-attention relative error from 0.648 -> 0.132 and DC bias from 0.165 -> 0.014, at 7.1 -> 8.1 ms. The selected tokens were independently re-derived from workspace centroid scores and verified sorted, within budget, candidate-only, and correctly cut at the histogram boundary. The second commit fixes HIP `round_fp16` under `-ffast-math`. Clang was folding the float -> half -> float round away, causing `apply_rope` to differ from eager by up to 9% of elements. Two inline-asm conversions restore exact rounding; fp16-freq cases are now bit-identical. The 2x24x4096x128 case goes 865 -> 1024 us; fp32/bf16 paths are unchanged and VGPR use drops by one. `test_hip_wmma.py::test_rope_matches_eager` now covers fp16 freqs and checks bit identity. The third commit enables HIP in the test backend/device lists. This activates ~130 additional cases across AdaLN, RMS AdaLN, NA3D, FP8 quantization, int8 activations, and RoPE. All pass; non-HIP machines simply skip them. `test_rms_rope.py` remains unchanged because its existing tolerance bands are already exceeded by eager on some seeds, while HIP tracks eager within 0.001 percentage points. Coverage remains in `test_hip_wmma.py`. Suite result: 1483 passing vs 1349 before, with failures unchanged at 101 and none in HIP paths. Compile-verified on gfx1100/gfx1200 with no scratch spills; run-verified on gfx1200. All five `token_aug` tests in `tests/test_sol_attn.py` now run and pass. --- <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:20 +00:00