[PR #254] Vectorize HIP rms_adaln and split-half rms_rope #259

New issue

Open

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

bsakaguchi commented

2026-10-26 22:22:11 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/254
Author: @qweqweewqe7-create
Created: 10/17/2026
Status: 🔄 Open

Base: main ← Head: moss-hip-adaln-rope-vectorize

📝 Commits (2)

7c01dd2 Vectorize HIP rms_adaln and split-half rms_rope
8a3e749 Copy rms_adaln/rms_rope vector packs with memcpy

📊 Changes

2 files changed (+318 additions, -0 deletions)

View changed files

📝 comfy_kitchen/backends/hip/ops/adaln.hip (+159 -0)
📝 comfy_kitchen/backends/hip/ops/rms_rope.hip (+159 -0)

📄 Description

Summary

Vectorized HIP fast paths for two memory-bound kernels used by modulated DiT blocks (MOSS-Image, among others):

rms_adaln / adaln: half/bf16 rows with D % 8 == 0 and D <= 4096. Each thread keeps up to four 16-byte chunks of the row in registers, so x is read once instead of twice; scale/shift are loaded 8 at a time; the reductions use wave shuffles.
split-half rms_rope: half/bf16 q/k with a contiguous head dim of 64, 128 or 256, full rotation and contiguous fp32 [.., rot/2, 2, 2] freqs. head_dim / 8 lanes share a row and each holds four (p, p + head_dim/2) pairs of q and k, so both are read once; the norm reduces with shuffles. The in-place variant stays safe, since every lane loads its elements before storing them.

Every other dtype, layout and size keeps the existing kernels. The per-element math and rounding are the existing kernels'; only the order of the fp32 sums differs.

Performance

RX 9070 XT (gfx1201), PyTorch 2.16.0a0+rocm10.2.0a20261015, HIP event timing, 50 iterations after 10 warmups:

case	before	after	after (GB/s)
`rms_adaln` bf16 6304x3072, broadcast scale/shift	0.278 ms	0.125 ms	619
`rms_adaln` fp16 1000x3072, per-row fp32 scale/shift	0.079 ms	0.045 ms	274
`rms_adaln` bf16 4097x1536	0.123 ms	0.055 ms	462
`rms_adaln` bf16 2048x4096	0.140 ms	0.053 ms	637
`rms_rope_split_half` bf16 qkv slice, 6304 x 24 x 128	0.664 ms	0.298 ms	542
`rms_rope_split_half` fp16 3001 x 16 x 64	0.224 ms	0.064 ms	407
`rms_rope_split_half` bf16 1003 x 8 x 256, fp32 scale	0.098 ms	0.051 ms	365
`rms_rope_split_half_` bf16 in place, 2051 x 24 x 128	0.254 ms	0.088 ms	596

A device-to-device copy measures 618 GB/s on the same card. MOSS-Image 1024x1536 DiT forward: 584 -> 554 ms.

Correctness

Against the old kernels, outputs match except where the fp32 sum order flips a rounding: at most one bf16/fp16 ulp at the value's magnitude, in about 1e-5 of the elements.
Against eager, the mismatch rate is equal or lower than the old kernels' (for example 1.9e-5 -> 1.1e-5 for the MOSS rope case).
tests/test_adaln.py tests/test_rms_rope.py and the HIP dispatch/WMMA/INT8 files on RX 9070 XT (gfx1201) and RX 9060 XT (gfx1200): same results as unmodified main. The only failures (18 triton cases and test_int8_linear_correctness[hip]) fail identically on main.

Update: rebased onto v0.2.40, compared with #259's wave32 rms_rope

Rebased onto main (f8b7e97, v0.2.40). #259 added rms_rope_wave32_128_kernel, which now handles every head_dim=128 call on wave32, split-half included. The vectorized split-half path here is checked first and only takes the cases it supports. Everything else goes through #259's dispatch. The interleaved rms_rope is unchanged and gives the same output hash as main.

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 50 calls, averaged over two alternating runs. MOSS-Image 1024x1536 shapes (6304 tokens, D=3072, 24 heads x 128, BHND):

op	9070 XT main	9070 XT PR	9060 XT main	9060 XT PR
`rms_adaln` bf16 6304 x 3072	0.296 ms	0.151 ms	0.571 ms	0.311 ms
`adaln` bf16 6304 x 3072	0.361	0.154	0.694	0.303
`rms_rope_split_half` bf16	0.502	0.307	0.922	0.615
`rms_rope` (interleaved, unchanged)	0.353	0.335	0.632	0.644

The split-half path is still about 1.5-1.6x faster than #259's wave32 kernel. rms_adaln and adaln are about 2x faster.

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) show the same failures as 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/254 **Author:** [@qweqweewqe7-create](https://github.com/qweqweewqe7-create) **Created:** 10/17/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `moss-hip-adaln-rope-vectorize` --- ### 📝 Commits (2) - [`7c01dd2`](https://github.com/Comfy-Org/comfy-kitchen/commit/7c01dd2220cb8849463c3177b935dd274744e00e) Vectorize HIP rms_adaln and split-half rms_rope - [`8a3e749`](https://github.com/Comfy-Org/comfy-kitchen/commit/8a3e749b2356cce93ee74bfeffa9e42d1e1a160e) Copy rms_adaln/rms_rope vector packs with memcpy ### 📊 Changes **2 files changed** (+318 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/hip/ops/adaln.hip` (+159 -0) 📝 `comfy_kitchen/backends/hip/ops/rms_rope.hip` (+159 -0) </details> ### 📄 Description ## Summary Vectorized HIP fast paths for two memory-bound kernels used by modulated DiT blocks (MOSS-Image, among others): - **`rms_adaln` / `adaln`**: half/bf16 rows with `D % 8 == 0` and `D <= 4096`. Each thread keeps up to four 16-byte chunks of the row in registers, so `x` is read once instead of twice; scale/shift are loaded 8 at a time; the reductions use wave shuffles. - **split-half `rms_rope`**: half/bf16 `q`/`k` with a contiguous head dim of 64, 128 or 256, full rotation and contiguous fp32 `[.., rot/2, 2, 2]` freqs. `head_dim / 8` lanes share a row and each holds four `(p, p + head_dim/2)` pairs of `q` and `k`, so both are read once; the norm reduces with shuffles. The in-place variant stays safe, since every lane loads its elements before storing them. Every other dtype, layout and size keeps the existing kernels. The per-element math and rounding are the existing kernels'; only the order of the fp32 sums differs. ## Performance RX 9070 XT (gfx1201), PyTorch `2.16.0a0+rocm10.2.0a20261015`, HIP event timing, 50 iterations after 10 warmups: | case | before | after | after (GB/s) | | --- | ---: | ---: | ---: | | `rms_adaln` bf16 6304x3072, broadcast scale/shift | 0.278 ms | 0.125 ms | 619 | | `rms_adaln` fp16 1000x3072, per-row fp32 scale/shift | 0.079 ms | 0.045 ms | 274 | | `rms_adaln` bf16 4097x1536 | 0.123 ms | 0.055 ms | 462 | | `rms_adaln` bf16 2048x4096 | 0.140 ms | 0.053 ms | 637 | | `rms_rope_split_half` bf16 qkv slice, 6304 x 24 x 128 | 0.664 ms | 0.298 ms | 542 | | `rms_rope_split_half` fp16 3001 x 16 x 64 | 0.224 ms | 0.064 ms | 407 | | `rms_rope_split_half` bf16 1003 x 8 x 256, fp32 scale | 0.098 ms | 0.051 ms | 365 | | `rms_rope_split_half_` bf16 in place, 2051 x 24 x 128 | 0.254 ms | 0.088 ms | 596 | A device-to-device copy measures 618 GB/s on the same card. MOSS-Image 1024x1536 DiT forward: 584 -> 554 ms. ## Correctness - Against the old kernels, outputs match except where the fp32 sum order flips a rounding: at most one bf16/fp16 ulp at the value's magnitude, in about 1e-5 of the elements. - Against eager, the mismatch rate is equal or lower than the old kernels' (for example 1.9e-5 -> 1.1e-5 for the MOSS rope case). - `tests/test_adaln.py tests/test_rms_rope.py` and the HIP dispatch/WMMA/INT8 files on RX 9070 XT (gfx1201) and RX 9060 XT (gfx1200): same results as unmodified `main`. The only failures (18 triton cases and `test_int8_linear_correctness[hip]`) fail identically on `main`. ## Update: rebased onto v0.2.40, compared with #259's wave32 rms_rope Rebased onto `main` (`f8b7e97`, v0.2.40). #259 added `rms_rope_wave32_128_kernel`, which now handles every head_dim=128 call on wave32, split-half included. The vectorized split-half path here is checked first and only takes the cases it supports. Everything else goes through #259's dispatch. The interleaved `rms_rope` is unchanged and gives the same output hash as `main`. 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 50 calls, averaged over two alternating runs. MOSS-Image 1024x1536 shapes (6304 tokens, D=3072, 24 heads x 128, BHND): | op | 9070 XT main | 9070 XT PR | 9060 XT main | 9060 XT PR | | --- | ---: | ---: | ---: | ---: | | `rms_adaln` bf16 6304 x 3072 | 0.296 ms | 0.151 ms | 0.571 ms | 0.311 ms | | `adaln` bf16 6304 x 3072 | 0.361 | 0.154 | 0.694 | 0.303 | | `rms_rope_split_half` bf16 | 0.502 | 0.307 | 0.922 | 0.615 | | `rms_rope` (interleaved, unchanged) | 0.353 | 0.335 | 0.632 | 0.644 | The split-half path is still about 1.5-1.6x faster than #259's wave32 kernel. `rms_adaln` and `adaln` are about 2x faster. 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`) show the same failures as 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:22:11 +00:00

bsakaguchi referenced this issue

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

[PR #259] [MERGED] AMD int8 performance #264