Comfy-Org/comfy-kitchen

Fork 0

mirror of https://github.com/Comfy-Org/comfy-kitchen.git synced 2026-11-04 23:37:42 +00:00

[PR #183] Optimize HIP ConvRot register transpose and block geometry #202

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/183
Author: @qweqweewqe7-create
Created: 9/16/2026
Status: 🔄 Open

Base: main ← Head: optimize-hip-convrot-register-transpose

📝 Commits (5)

22d7127 Optimize HIP ConvRot register transpose and block geometry
8b94a26 Handle misaligned BF16 SwiGLU ConvRot inputs
aa5258a Vectorize ConvRot INT8 SiTU input loads, tune MOSS-Image widths
76f5a7e Align two v0.2.40 HIP tests with the ConvRot changes
da6565e Keep the fused ConvRot INT8 producer on gfx12

📊 Changes

3 files changed (+372 additions, -23 deletions)

View changed files

📝 comfy_kitchen/backends/hip/hadamard.h (+314 -22)
📝 comfy_kitchen/backends/hip/ops/quantize_int8.hip (+21 -0)
📝 tests/test_hip_wmma.py (+37 -1)

📄 Description

Summary

Optimize the HIP INT8 ConvRot path by reducing wave-local LDS transpose traffic
and selecting block geometries that better match the dominant H3 projection
depths on RDNA4.

The optimized path:

performs the S4 transpose in wave-local registers;
performs the S16 transpose in registers while preserving the floating-point
evaluation order needed for exact quantization under -ffast-math;
uses LDS only once the transform becomes cross-wave at the S16 -> S64 handoff;
uses a single LDS transpose buffer;
uses validated K-dependent block geometries;
preserves the newer upstream fused RMSNorm ConvRot path.

Performance

Broad H3 projection sweeps over approximately M=15k through 110k:

RX 9070 XT / gfx1201: weighted improvement 19.418%, 18/18 wins
RX 9060 XT / gfx1200: weighted improvement 17.914%

The sweep covered QKV, attention output, MLP up, and MLP down projection
geometries rather than tuning for one fixed production M.

Correctness

Historical optimization validation:

deterministic ConvRot fingerprint validation: 155 / 155 exact
relevant validation suite: 398 passed
validated on gfx1200 and gfx1201

After rebasing onto current upstream main, the existing optimized ConvRot
paths were compared bit-for-bit against the previously validated authoritative
winner over representative BF16/FP16/FP32, GELU, SwiGLU, fused and spill-path
cases. Both gfx1200 and gfx1201 were 13 / 13 bitwise exact.

The current upstream ConvRot/input-activation tests, including the newer fused
RMSNorm path, were also re-run on both RDNA4 GPUs.

Notes

Changing block geometry changes the parallel reduction order, so a direct
bit-for-bit comparison between the optimized block sizes and the older
1024-thread upstream geometry is not used as the rebase-regression criterion.
Rebase equivalence is checked against the already validated optimized winner,
while the upstream test suite covers numerical/API correctness.

Update: rebased onto v0.2.38 (SiTU port #252), SiTU vector loads, MOSS-Image widths

Rebased onto current main (42081bf, v0.2.38), which resolves the conflict with the HIP SiTU ConvRot port (#252). The existing two commits are unchanged in substance. One commit is added, Vectorize ConvRot INT8 SiTU input loads, tune MOSS-Image widths:

the fused quantizer loads a bf16 [hidden | gate] SiTU row with two 8-byte loads, as SwiGLU already does, instead of eight dtype-dispatched scalar loads. The scalar and vector loaders share one SiTU formula (situ_value);
K=3072 uses 128-thread and K=8192 256-thread blocks when M >= 1024 (MOSS-Image attention/FFN-in and SiTU FFN-down widths). Short M keeps the wide block.

RX 9070 XT (gfx1201), PyTorch 2.16.0a0+rocm10.2.0a20261015, activation rotate+quant, alternating runs. "Before" is this PR's previous state.

case	before	after
MOSS-Image K=3072, M=6304	0.188 ms	0.118 ms
MOSS-Image SiTU K=8192 (16384-wide input), M=6304	0.805 ms	0.582 ms
H3 K=5376, M=30000	0.95 ms	0.95 ms
H3 K=14336, M=30000	2.65 ms	2.64 ms
H3 SwiGLU K=14336, M=30000	3.36 ms	3.35 ms

The MOSS-Image DiT forward goes from 553 to about 537 ms. The H3 widths are unchanged.

A packed four-column int8 store was also tried. It made the H3 widths 15-18% slower and K=3072 slower as well, so it is not included.

Correctness: outputs are bit-identical to the previous state, except the bf16 SiTU path, where 101 of 51.6M int8 values move by one step at rounding boundaries. That path's accuracy against an fp64 reference is the same: at the differing positions the old and new kernels match the reference 49 and 52 times. Block sizes do not change the result of the max reduction, so the new widths are bit-identical to the old ones.

Re-validation on RX 9070 XT (gfx1201) and RX 9060 XT (gfx1200), gfx1200;gfx1201 build: tests/test_int8*.py tests/test_hip_wmma.py tests/test_hip_dispatch.py tests/test_convrot_unpacked_quantization.py tests/test_hip_convrot_perf.py tests/test_w4a8_requant.py tests/test_adaln.py tests/test_rms_rope.py give the same results as unmodified main (1176 passed). The only failures, 18 triton cases and test_int8_linear_correctness[hip], fail identically on main.

Update: rebased onto v0.2.40 (#259)

Rebased onto main (f8b7e97, v0.2.40). One conflict, in launch_convrot_quant_fused_impl: it now combines this PR's single-buffer shared-memory size with #259's TILED_QOUT kernel template. The fused kernel body already uses one buffer per group, and convrot_fused_lds_fits uses the same size.

Two commits are added:

Align two v0.2.40 HIP tests with the ConvRot changes.
- #259's binding rejects a misaligned x. The offset-input SwiGLU test now checks two things: the binding refuses the misaligned input, and the public wrapper, which realigns, matches the aligned result.
- The single-buffer producer fits K=24576 in LDS at 768 threads. The global-vs-fused exactness test therefore uses K=28672, the first width that still spills at 96 rows.
Keep the fused ConvRot INT8 producer on gfx12. #259 sends bf16 ConvRot with M >= 1024 and K in 2560..16384 to the transposed 512x2 kernel on gfx12/gfx117. With this PR's producer, the fused kernel is as fast or faster on gfx12 and yields the same int8 values and scales. gfx117 keeps the transposed kernel, and COMFY_KITCHEN_CONVROT_TRANSPOSED=1 selects it on gfx12.

RX 9070 XT (gfx1201) and RX 9060 XT (gfx1200), PyTorch 2.16.0a0+rocm10.3.0a20261023, HIP 7.17.26423. Timings are wall-clock medians over windows of calls, averaged over two alternating runs. int8_linear with ConvRot, bf16 out, this PR with the transposed vs the fused producer:

shape (M x K -> N)	9070 XT transposed	9070 XT fused	9060 XT transposed	9060 XT fused
MOSS qkv 12544 x 3072 -> 9216	4.002 ms	3.918 ms	7.718 ms	7.633 ms
MOSS out 12544 x 3072 -> 3072	1.616	1.505	2.835	2.718
MOSS up/gate 12544 x 3072 -> 16384	6.826	6.719	13.299	13.203
MOSS down 12544 x 8192 -> 3072	3.807	3.758	7.123	7.158
H3 qkv / out / fc1 / fc2 (M=38339)		within +-0.6%		within +-0.6%

At every shape, the two producers give identical outputs, and both match main.

Tests on both GPUs (test_hip_wmma, test_hip_dispatch, test_gemm_bias_contract, test_int8*, test_adaln, test_rms_rope, test_int8_attention, test_sol_attn, test_convrot_unpacked_quantization, test_w4a8_mma, test_w6a8):

2401 passed: main's 2400 plus the offset-input test from this PR.
The failures are the same as on unmodified main: 18 triton rope cases and test_int8_linear_correctness[hip].

🤖 Generated with Claude Code

_{🔄 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/183 **Author:** [@qweqweewqe7-create](https://github.com/qweqweewqe7-create) **Created:** 9/16/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `optimize-hip-convrot-register-transpose` --- ### 📝 Commits (5) - [`22d7127`](https://github.com/Comfy-Org/comfy-kitchen/commit/22d7127cdc97ccd922824088b9e38f0831e6f537) Optimize HIP ConvRot register transpose and block geometry - [`8b94a26`](https://github.com/Comfy-Org/comfy-kitchen/commit/8b94a26694908e22dfcb3277c355dac6272f5f5e) Handle misaligned BF16 SwiGLU ConvRot inputs - [`aa5258a`](https://github.com/Comfy-Org/comfy-kitchen/commit/aa5258a16c875c6e69e6c8696a9709e17b4ae171) Vectorize ConvRot INT8 SiTU input loads, tune MOSS-Image widths - [`76f5a7e`](https://github.com/Comfy-Org/comfy-kitchen/commit/76f5a7ef2aae063dee4c68422d8203845169a434) Align two v0.2.40 HIP tests with the ConvRot changes - [`da6565e`](https://github.com/Comfy-Org/comfy-kitchen/commit/da6565eb3b7708d3ff8c126a90e0c9198c310047) Keep the fused ConvRot INT8 producer on gfx12 ### 📊 Changes **3 files changed** (+372 additions, -23 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/hip/hadamard.h` (+314 -22) 📝 `comfy_kitchen/backends/hip/ops/quantize_int8.hip` (+21 -0) 📝 `tests/test_hip_wmma.py` (+37 -1) </details> ### 📄 Description ## Summary Optimize the HIP INT8 ConvRot path by reducing wave-local LDS transpose traffic and selecting block geometries that better match the dominant H3 projection depths on RDNA4. The optimized path: - performs the S4 transpose in wave-local registers; - performs the S16 transpose in registers while preserving the floating-point evaluation order needed for exact quantization under `-ffast-math`; - uses LDS only once the transform becomes cross-wave at the S16 -> S64 handoff; - uses a single LDS transpose buffer; - uses validated K-dependent block geometries; - preserves the newer upstream fused RMSNorm ConvRot path. ## Performance Broad H3 projection sweeps over approximately M=15k through 110k: - RX 9070 XT / gfx1201: weighted improvement **19.418%**, 18/18 wins - RX 9060 XT / gfx1200: weighted improvement **17.914%** The sweep covered QKV, attention output, MLP up, and MLP down projection geometries rather than tuning for one fixed production M. ## Correctness Historical optimization validation: - deterministic ConvRot fingerprint validation: **155 / 155 exact** - relevant validation suite: **398 passed** - validated on gfx1200 and gfx1201 After rebasing onto current upstream `main`, the existing optimized ConvRot paths were compared bit-for-bit against the previously validated authoritative winner over representative BF16/FP16/FP32, GELU, SwiGLU, fused and spill-path cases. Both gfx1200 and gfx1201 were **13 / 13 bitwise exact**. The current upstream ConvRot/input-activation tests, including the newer fused RMSNorm path, were also re-run on both RDNA4 GPUs. ## Notes Changing block geometry changes the parallel reduction order, so a direct bit-for-bit comparison between the optimized block sizes and the older 1024-thread upstream geometry is not used as the rebase-regression criterion. Rebase equivalence is checked against the already validated optimized winner, while the upstream test suite covers numerical/API correctness. ## Update: rebased onto v0.2.38 (SiTU port #252), SiTU vector loads, MOSS-Image widths Rebased onto current `main` (`42081bf`, v0.2.38), which resolves the conflict with the HIP SiTU ConvRot port (#252). The existing two commits are unchanged in substance. One commit is added, *Vectorize ConvRot INT8 SiTU input loads, tune MOSS-Image widths*: - the fused quantizer loads a bf16 `[hidden | gate]` SiTU row with two 8-byte loads, as SwiGLU already does, instead of eight dtype-dispatched scalar loads. The scalar and vector loaders share one SiTU formula (`situ_value`); - `K=3072` uses 128-thread and `K=8192` 256-thread blocks when `M >= 1024` (MOSS-Image attention/FFN-in and SiTU FFN-down widths). Short M keeps the wide block. RX 9070 XT (gfx1201), PyTorch `2.16.0a0+rocm10.2.0a20261015`, activation rotate+quant, alternating runs. "Before" is this PR's previous state. | case | before | after | | --- | ---: | ---: | | MOSS-Image K=3072, M=6304 | 0.188 ms | 0.118 ms | | MOSS-Image SiTU K=8192 (16384-wide input), M=6304 | 0.805 ms | 0.582 ms | | H3 K=5376, M=30000 | 0.95 ms | 0.95 ms | | H3 K=14336, M=30000 | 2.65 ms | 2.64 ms | | H3 SwiGLU K=14336, M=30000 | 3.36 ms | 3.35 ms | The MOSS-Image DiT forward goes from 553 to about 537 ms. The H3 widths are unchanged. A packed four-column int8 store was also tried. It made the H3 widths 15-18% slower and K=3072 slower as well, so it is not included. Correctness: outputs are bit-identical to the previous state, except the bf16 SiTU path, where 101 of 51.6M int8 values move by one step at rounding boundaries. That path's accuracy against an fp64 reference is the same: at the differing positions the old and new kernels match the reference 49 and 52 times. Block sizes do not change the result of the max reduction, so the new widths are bit-identical to the old ones. Re-validation on RX 9070 XT (gfx1201) and RX 9060 XT (gfx1200), `gfx1200;gfx1201` build: `tests/test_int8*.py tests/test_hip_wmma.py tests/test_hip_dispatch.py tests/test_convrot_unpacked_quantization.py tests/test_hip_convrot_perf.py tests/test_w4a8_requant.py tests/test_adaln.py tests/test_rms_rope.py` give the same results as unmodified `main` (1176 passed). The only failures, 18 triton cases and `test_int8_linear_correctness[hip]`, fail identically on `main`. ## Update: rebased onto v0.2.40 (#259) Rebased onto `main` (`f8b7e97`, v0.2.40). One conflict, in `launch_convrot_quant_fused_impl`: it now combines this PR's single-buffer shared-memory size with #259's `TILED_QOUT` kernel template. The fused kernel body already uses one buffer per group, and `convrot_fused_lds_fits` uses the same size. Two commits are added: - *Align two v0.2.40 HIP tests with the ConvRot changes.* - #259's binding rejects a misaligned `x`. The offset-input SwiGLU test now checks two things: the binding refuses the misaligned input, and the public wrapper, which realigns, matches the aligned result. - The single-buffer producer fits K=24576 in LDS at 768 threads. The global-vs-fused exactness test therefore uses K=28672, the first width that still spills at 96 rows. - *Keep the fused ConvRot INT8 producer on gfx12.* #259 sends bf16 ConvRot with M >= 1024 and K in 2560..16384 to the transposed 512x2 kernel on gfx12/gfx117. With this PR's producer, the fused kernel is as fast or faster on gfx12 and yields the same int8 values and scales. gfx117 keeps the transposed kernel, and `COMFY_KITCHEN_CONVROT_TRANSPOSED=1` selects it on gfx12. RX 9070 XT (gfx1201) and RX 9060 XT (gfx1200), PyTorch `2.16.0a0+rocm10.3.0a20261023`, HIP 7.17.26423. Timings are wall-clock medians over windows of calls, averaged over two alternating runs. `int8_linear` with ConvRot, bf16 out, this PR with the transposed vs the fused producer: | shape (M x K -> N) | 9070 XT transposed | 9070 XT fused | 9060 XT transposed | 9060 XT fused | | --- | ---: | ---: | ---: | ---: | | MOSS qkv 12544 x 3072 -> 9216 | 4.002 ms | 3.918 ms | 7.718 ms | 7.633 ms | | MOSS out 12544 x 3072 -> 3072 | 1.616 | 1.505 | 2.835 | 2.718 | | MOSS up/gate 12544 x 3072 -> 16384 | 6.826 | 6.719 | 13.299 | 13.203 | | MOSS down 12544 x 8192 -> 3072 | 3.807 | 3.758 | 7.123 | 7.158 | | H3 qkv / out / fc1 / fc2 (M=38339) | | within +-0.6% | | within +-0.6% | At every shape, the two producers give identical outputs, and both match `main`. Tests on both GPUs (`test_hip_wmma`, `test_hip_dispatch`, `test_gemm_bias_contract`, `test_int8*`, `test_adaln`, `test_rms_rope`, `test_int8_attention`, `test_sol_attn`, `test_convrot_unpacked_quantization`, `test_w4a8_mma`, `test_w6a8`): - 2401 passed: `main`'s 2400 plus the offset-input test from this PR. - The failures are the same as on unmodified `main`: 18 triton rope cases and `test_int8_linear_correctness[hip]`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

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

bsakaguchi referenced this issue

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

[PR #202] [MERGED] HIP port of #192: zero_pad and strided NDHWC views for conv3d and GroupNorm #216