[PR #195] [HIP] Speed up RMS+RoPE for 64-wide heads on wave32 #212

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/195
Author: @Yasei-no-otoko
Created: 9/24/2026
Status: 🔄 Open

Base: main ← Head: codex/h3-rms-rope-multirow

📝 Commits (4)

810527c Optimize HIP RMS RoPE for head dimension 64
fdd5833 Check HIP RMS RoPE accuracy per row
62b7050 Assert HIP RoPE on the rotated prefix
a91df93 Select RMS RoPE by stream wave size

📊 Changes

2 files changed (+187 additions, -34 deletions)

View changed files

📝 comfy_kitchen/backends/hip/ops/rms_rope.hip (+95 -34)
📝 tests/test_rms_rope.py (+92 -0)

📄 Description

For 64-wide heads on HIP devices reporting a 32-lane wavefront, process four rows per block with independent 32-lane shuffle reductions. This preserves reduction order, dtype rounding, partial rotary, and in-place/strided contracts. Dispatch queries the launch stream's device and its wave size. Other head dimensions, non-wave32 devices, and query failures use the existing one-row path. Batch, sequence length, row counts, and strides remain runtime arguments.

On Radeon 8060S (gfx1151), the four-row device kernel improved actual MiniMax H3 operand timing from 2.2002 ms to 1.4333 ms across seven rotating-order rounds. Five alternating warm BF16 video-VAE decode pairs at 864×480 / 124 frames improved the median from 53.3668 s to 52.7726 s (1.1%); the median paired reduction was 0.5942 s, and every decoded pixel was bitwise identical. Absolute times drifted, but every paired candidate run was faster. These measurements exclude the added host capability query and cover VAE decode only.

The final guarded build was compared with the same four-row kernel over 15 alternating real-operand pairs: synchronized wall medians were 1.5538 ms versus 1.5732 ms, with a median paired increase of 0.0188 ms and overlapping ranges. Q/K outputs were bitwise identical and V was unchanged on default and non-default streams. Full VAE decode was not remeasured after adding the host query.

Validation: 18 focused HIP tests passed on Windows / ROCm 10.2 / gfx1151, covering partial rotary, tail rows, packed QKV, FP16/FP32, strides, and in-place/out-of-place calls. The complete RMS+RoPE device translation unit compiled for gfx1010, gfx1030, gfx1100, gfx1151, gfx1200, gfx908, gfx90a, gfx942, and gfx950; metadata reports wave32 for the RDNA targets and wave64 for CDNA. Other GPUs were not run or benchmarked. The current backend package manifest excludes RDNA1 and CDNA; this PR does not enable those package targets.

_{🔄 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/195 **Author:** [@Yasei-no-otoko](https://github.com/Yasei-no-otoko) **Created:** 9/24/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `codex/h3-rms-rope-multirow` --- ### 📝 Commits (4) - [`810527c`](https://github.com/Comfy-Org/comfy-kitchen/commit/810527cf9b7eff1a1dfb4e195a15b33dc5f69a96) Optimize HIP RMS RoPE for head dimension 64 - [`fdd5833`](https://github.com/Comfy-Org/comfy-kitchen/commit/fdd5833917053c58a21500888e481631b9004591) Check HIP RMS RoPE accuracy per row - [`62b7050`](https://github.com/Comfy-Org/comfy-kitchen/commit/62b7050e96c223abc6eb793ff249aaebba955235) Assert HIP RoPE on the rotated prefix - [`a91df93`](https://github.com/Comfy-Org/comfy-kitchen/commit/a91df93e59fcecc0be15211233e9c94aee8e328d) Select RMS RoPE by stream wave size ### 📊 Changes **2 files changed** (+187 additions, -34 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/hip/ops/rms_rope.hip` (+95 -34) 📝 `tests/test_rms_rope.py` (+92 -0) </details> ### 📄 Description For 64-wide heads on HIP devices reporting a 32-lane wavefront, process four rows per block with independent 32-lane shuffle reductions. This preserves reduction order, dtype rounding, partial rotary, and in-place/strided contracts. Dispatch queries the launch stream's device and its wave size. Other head dimensions, non-wave32 devices, and query failures use the existing one-row path. Batch, sequence length, row counts, and strides remain runtime arguments. On Radeon 8060S (gfx1151), the four-row device kernel improved actual MiniMax H3 operand timing from 2.2002 ms to 1.4333 ms across seven rotating-order rounds. Five alternating warm BF16 video-VAE decode pairs at 864×480 / 124 frames improved the median from 53.3668 s to 52.7726 s (1.1%); the median paired reduction was 0.5942 s, and every decoded pixel was bitwise identical. Absolute times drifted, but every paired candidate run was faster. These measurements exclude the added host capability query and cover VAE decode only. The final guarded build was compared with the same four-row kernel over 15 alternating real-operand pairs: synchronized wall medians were 1.5538 ms versus 1.5732 ms, with a median paired increase of 0.0188 ms and overlapping ranges. Q/K outputs were bitwise identical and V was unchanged on default and non-default streams. Full VAE decode was not remeasured after adding the host query. Validation: 18 focused HIP tests passed on Windows / ROCm 10.2 / gfx1151, covering partial rotary, tail rows, packed QKV, FP16/FP32, strides, and in-place/out-of-place calls. The complete RMS+RoPE device translation unit compiled for gfx1010, gfx1030, gfx1100, gfx1151, gfx1200, gfx908, gfx90a, gfx942, and gfx950; metadata reports wave32 for the RDNA targets and wave64 for CDNA. Other GPUs were not run or benchmarked. The current backend package manifest excludes RDNA1 and CDNA; this PR does not enable those package targets. --- <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:37 +00:00

bsakaguchi referenced this issue

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

[PR #212] [Perf][Ascend] Reuse unpacked W4A4 weights within an explicit scope #224