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 #94] [CLOSED] HIP: fix RDNA4 GEMM tile selection, vectorize fp8 elementwise kernels #140

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/94
Author: @crashingalexsan
Created: 8/5/2026
Status: ❌ Closed

Base: main ← Head: rdna4-wmma-tuning

📝 Commits (10+)

e177160 HIP: fix RDNA4 GEMM tile selection and vectorize fp8 elementwise kernels
a9ee312 HIP: make the device CU cache atomic
d1e673c Rename device_cu_count to device_wgp_count, cover the BKB=128 tile paths
35df5ec Route skinny shapes (M or N <= 64) to the 64x64 tile
bfb2c0b Trigger CI
cfb2a59 Tighten selector and test-shape comments
90d9a97 Reword tuning provenance comment to RDNA4
a887137 Cover w4x4 deep-K on 40-48 WGP parts, unify float conversion helpers
4e3778f Note byte units on depth thresholds, cover the scalar elementwise fallback
690e9b4 Make misaligned-view tests compare real allocations and run the hip backend

📊 Changes

6 files changed (+295 additions, -50 deletions)

View changed files

📝 comfy_kitchen/backends/hip/gemm_wmma.h (+93 -0)
📝 comfy_kitchen/backends/hip/ops/gemm_fp8.hip (+1 -20)
📝 comfy_kitchen/backends/hip/ops/gemm_int8.hip (+1 -16)
📝 comfy_kitchen/backends/hip/ops/per_tensor_fp8.hip (+138 -12)
📝 tests/test_hip_wmma.py (+12 -2)
📝 tests/test_qdq.py (+50 -0)

📄 Description

HIP: fix RDNA4 GEMM tile selection, vectorize fp8 elementwise kernels

What this fixes

The WMMA GEMM launcher chose its tile from M/N/K thresholds that, on two
common shape classes, select the worst configuration available to it. Both are
pure selection bugs — the kernel itself is fine, it was just being handed the
wrong template parameters.

1. The BM=256 "deep-K" tile was gated on K > N, which is exactly where it
performs worst. On 4096x2048x8192 it runs at 152 TFLOPS where the same kernel
at 128x128 with BKB=128 reaches 188. Its 16 accumulator fragments per warp
cost more occupancy than the taller tile buys back. Removed entirely.

2. The M >= 96 && N >= 96 floor sent narrow shapes to the 128x128 tile even
when that yields fewer blocks than the device has CUs, leaving most of the GPU
idle. 2048x128x2048 ran at 18 TFLOPS; the 64x64 tile does 70.

What replaces them

A single launch_gemm_wmma() in gemm_wmma.h, shared by the fp8 and int8
launchers, selecting on three measured axes:

axis	rule	why
grid coverage	`64x64` when the `128x128` grid can't fill the device	a tile that leaves CUs idle loses more than its reuse gains
K depth	`BKB=128` from `K>=4096`	halves LDS round trips and `__syncthreads` pairs per K element
warp grid	at `K>=4096`, `w4x4/t2x2` while blocks/CU <= 4, else `w4x2/t2x4`	with few blocks per CU there is nothing to interleave across, so more waves within a block hide the latency instead

The K threshold is set by int8, not fp8: taking BKB=128 lower costs int8 0.81x
at K=640, where fp8 only breaks even.

The warp-grid crossover, consistent across both policies:

blocks/CU	1x	2x	3x	4x	5x	6x
`w4x4/t2x2` vs `w4x2/t2x4`	1.18x	1.24x	1.04x	1.05x	1.02x	0.97x

Results

Old policy vs new, same kernel, A/B interleaved with min-of-3 on an idle
gfx1201 (RX 9070 XT), pseudorandom operands:

shape	fp8	int8
2048x128x2048	3.85x	3.77x
256x4096x4096	1.66x	1.62x
512x512x512	1.48x	1.75x
4096x2048x8192	1.23x	1.24x
4096x4096x4096	1.14x	1.11x
overall (19 shapes)	1.06x	1.06x

Shapes where both policies select the same config land within ±4% (noise).

Measured through torch against hipBLASLt on the same shapes, the fp8 GEMM goes
from 1.21x slower than hipBLASLt to 1.11x — it closes 46% of the gap. Output
is bit-identical to hipBLASLt on every shape tested.

Elementwise kernels

per_tensor_fp8 quantize/dequantize moved one element per thread, so a wave
wrote 32 contiguous bytes. Added a 16-elements-per-thread path moving whole
16-byte chunks, gated on a 16-byte alignment check of both pointers — a
contiguous tensor can still be a mid-buffer view, whose base carries the offset.
Per-element arithmetic is unchanged, so both paths are bit-identical (verified on
raw bit patterns; note torch.equal reports False on fp8-derived data purely
because e4m3 0x7f/0xff decode to NaN).

elements	quantize	dequantize
4M	1.31x	1.31x
16M	1.78x	1.75x
64M	1.90x	1.69x
<=1M	neutral	neutral

Testing

gfx1201: test_hip_wmma, test_hip_dispatch, test_qdq, test_int8 —
338 passed, 0 failed (stable over repeated runs).
Compiles for gfx1100, gfx1151, gfx1201, gfx1030.
Pre-existing failures on this ROCm build are unaffected — triton-backend rope,
NVFP4/MXFP8 (no HIP kernels), and HIP graph capture all fail identically on
unmodified v0.2.26.

Limitations / for reviewers

All tuning was done on gfx12 (gfx1201). I have no gfx11 hardware. The
grid-coverage rule adapts itself via a runtime CU query and should carry over;
the K and warp-grid thresholds are gfx12 numbers.
gfx11 fp8 spills registers badly and this PR does not fix it. With no fp8
WMMA there, MmaFp8::Frag is a v16bf (8 VGPRs) rather than gfx12's v2i
(2), and compiler resource usage on gfx1100 shows every 128x128 fp8 tile
spilling: w4x2/t2x4 308 VGPRs spilled, 256x128 834, while 64x64x128 and
128x128 w4x4/t2x2 spill none. This PR helps gfx11 incidentally — it deletes
the 834-spill tile and routes narrow shapes to the spill-free 64x64 — but a
square deep-K shape moves from 308 to 369. Using w4x4/t2x2 throughout would
eliminate the spilling, but it is the slower grid on gfx12 once blocks are
plentiful, so a proper fix needs an arch-conditional grid and someone with an
RDNA3 card to validate it. Worth noting register pressure alone does not
predict throughput here.
Benchmarking caveat: operand data matters a lot. A cycling byte fill
switches far fewer transistors than realistic randn-derived fp8, so the GPU
holds a ~20% higher sustained clock — 218 vs 180 TFLOPS on the same kernel and
shape. All numbers here use pseudorandom operands. Timing two policies in
back-to-back blocks rather than interleaved also makes the second look ~30%
slower even when both select the identical config.
Where the win concentrates (thanks @0xDELUXA): sweeping K on a fixed
2048x2048, the speedup over the old selector is concentrated where K is a
multiple of 4096 (1.35x at K=4096, 1.98x at K=8192, vs ~1.03x at 3072/6144).
The likely mechanism is the deleted 256x128 tile hitting cache aliasing when
the row stride is a multiple of 4096 bytes. If you benchmark this PR at, say,
K=3072 only, it will look like a no-op — that is expected.

_{🔄 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/94 **Author:** [@crashingalexsan](https://github.com/crashingalexsan) **Created:** 8/5/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `rdna4-wmma-tuning` --- ### 📝 Commits (10+) - [`e177160`](https://github.com/Comfy-Org/comfy-kitchen/commit/e1771606d274fb27b8cbf7bc7c012b14a085e7f8) HIP: fix RDNA4 GEMM tile selection and vectorize fp8 elementwise kernels - [`a9ee312`](https://github.com/Comfy-Org/comfy-kitchen/commit/a9ee3128f09b823175e409f977a88f6d4b76b09d) HIP: make the device CU cache atomic - [`d1e673c`](https://github.com/Comfy-Org/comfy-kitchen/commit/d1e673c7edcd28892000733a3a05044c14a4630e) Rename device_cu_count to device_wgp_count, cover the BKB=128 tile paths - [`35df5ec`](https://github.com/Comfy-Org/comfy-kitchen/commit/35df5ec239a1a79a4e537ff224d1c2e9e53a0d1c) Route skinny shapes (M or N <= 64) to the 64x64 tile - [`bfb2c0b`](https://github.com/Comfy-Org/comfy-kitchen/commit/bfb2c0b0105c32249185f26057ae06eb967f0eee) Trigger CI - [`cfb2a59`](https://github.com/Comfy-Org/comfy-kitchen/commit/cfb2a597f567e94ac02eec5a64c336c3dc374658) Tighten selector and test-shape comments - [`90d9a97`](https://github.com/Comfy-Org/comfy-kitchen/commit/90d9a9767c3dce0125dfdf00e9a0c324929601ae) Reword tuning provenance comment to RDNA4 - [`a887137`](https://github.com/Comfy-Org/comfy-kitchen/commit/a88713764005c6c723534a05634b4810ca786917) Cover w4x4 deep-K on 40-48 WGP parts, unify float conversion helpers - [`4e3778f`](https://github.com/Comfy-Org/comfy-kitchen/commit/4e3778f10243649f6724b890bbea71c3452c222a) Note byte units on depth thresholds, cover the scalar elementwise fallback - [`690e9b4`](https://github.com/Comfy-Org/comfy-kitchen/commit/690e9b4897dcaa979ec47605971801f0e26a5d90) Make misaligned-view tests compare real allocations and run the hip backend ### 📊 Changes **6 files changed** (+295 additions, -50 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/hip/gemm_wmma.h` (+93 -0) 📝 `comfy_kitchen/backends/hip/ops/gemm_fp8.hip` (+1 -20) 📝 `comfy_kitchen/backends/hip/ops/gemm_int8.hip` (+1 -16) 📝 `comfy_kitchen/backends/hip/ops/per_tensor_fp8.hip` (+138 -12) 📝 `tests/test_hip_wmma.py` (+12 -2) 📝 `tests/test_qdq.py` (+50 -0) </details> ### 📄 Description # HIP: fix RDNA4 GEMM tile selection, vectorize fp8 elementwise kernels ## What this fixes The WMMA GEMM launcher chose its tile from `M`/`N`/`K` thresholds that, on two common shape classes, select the *worst* configuration available to it. Both are pure selection bugs — the kernel itself is fine, it was just being handed the wrong template parameters. **1. The `BM=256` "deep-K" tile was gated on `K > N`, which is exactly where it performs worst.** On `4096x2048x8192` it runs at 152 TFLOPS where the same kernel at `128x128` with `BKB=128` reaches 188. Its 16 accumulator fragments per warp cost more occupancy than the taller tile buys back. Removed entirely. **2. The `M >= 96 && N >= 96` floor sent narrow shapes to the `128x128` tile even when that yields fewer blocks than the device has CUs**, leaving most of the GPU idle. `2048x128x2048` ran at 18 TFLOPS; the `64x64` tile does 70. ## What replaces them A single `launch_gemm_wmma()` in `gemm_wmma.h`, shared by the fp8 and int8 launchers, selecting on three measured axes: | axis | rule | why | |---|---|---| | grid coverage | `64x64` when the `128x128` grid can't fill the device | a tile that leaves CUs idle loses more than its reuse gains | | K depth | `BKB=128` from `K>=4096` | halves LDS round trips and `__syncthreads` pairs per K element | | warp grid | at `K>=4096`, `w4x4/t2x2` while blocks/CU <= 4, else `w4x2/t2x4` | with few blocks per CU there is nothing to interleave *across*, so more waves *within* a block hide the latency instead | The K threshold is set by int8, not fp8: taking `BKB=128` lower costs int8 0.81x at `K=640`, where fp8 only breaks even. The warp-grid crossover, consistent across both policies: | blocks/CU | 1x | 2x | 3x | 4x | 5x | 6x | |---|---|---|---|---|---|---| | `w4x4/t2x2` vs `w4x2/t2x4` | 1.18x | 1.24x | 1.04x | 1.05x | 1.02x | 0.97x | ## Results Old policy vs new, same kernel, A/B interleaved with min-of-3 on an idle gfx1201 (RX 9070 XT), pseudorandom operands: | shape | fp8 | int8 | |---|---|---| | 2048x128x2048 | **3.85x** | **3.77x** | | 256x4096x4096 | 1.66x | 1.62x | | 512x512x512 | 1.48x | 1.75x | | 4096x2048x8192 | 1.23x | 1.24x | | 4096x4096x4096 | 1.14x | 1.11x | | **overall (19 shapes)** | **1.06x** | **1.06x** | Shapes where both policies select the same config land within ±4% (noise). Measured through torch against hipBLASLt on the same shapes, the fp8 GEMM goes from **1.21x slower than hipBLASLt to 1.11x** — it closes 46% of the gap. Output is bit-identical to hipBLASLt on every shape tested. ## Elementwise kernels `per_tensor_fp8` quantize/dequantize moved one element per thread, so a wave wrote 32 contiguous bytes. Added a 16-elements-per-thread path moving whole 16-byte chunks, gated on a 16-byte alignment check of both pointers — a contiguous tensor can still be a mid-buffer view, whose base carries the offset. Per-element arithmetic is unchanged, so both paths are bit-identical (verified on raw bit patterns; note `torch.equal` reports False on fp8-derived data purely because e4m3 `0x7f`/`0xff` decode to NaN). | elements | quantize | dequantize | |---|---|---| | 4M | 1.31x | 1.31x | | 16M | 1.78x | 1.75x | | 64M | 1.90x | 1.69x | | <=1M | neutral | neutral | ## Testing - gfx1201: `test_hip_wmma`, `test_hip_dispatch`, `test_qdq`, `test_int8` — **338 passed, 0 failed** (stable over repeated runs). - Compiles for `gfx1100`, `gfx1151`, `gfx1201`, `gfx1030`. - Pre-existing failures on this ROCm build are unaffected — triton-backend rope, NVFP4/MXFP8 (no HIP kernels), and HIP graph capture all fail identically on unmodified v0.2.26. ## Limitations / for reviewers - **All tuning was done on gfx12 (gfx1201). I have no gfx11 hardware.** The grid-coverage rule adapts itself via a runtime CU query and should carry over; the K and warp-grid thresholds are gfx12 numbers. - **gfx11 fp8 spills registers badly and this PR does not fix it.** With no fp8 WMMA there, `MmaFp8::Frag` is a `v16bf` (8 VGPRs) rather than gfx12's `v2i` (2), and compiler resource usage on gfx1100 shows every `128x128` fp8 tile spilling: `w4x2/t2x4` 308 VGPRs spilled, `256x128` 834, while `64x64x128` and `128x128 w4x4/t2x2` spill none. This PR helps gfx11 incidentally — it deletes the 834-spill tile and routes narrow shapes to the spill-free `64x64` — but a square deep-K shape moves from 308 to 369. Using `w4x4/t2x2` throughout would eliminate the spilling, but it is the slower grid on gfx12 once blocks are plentiful, so a proper fix needs an arch-conditional grid and someone with an RDNA3 card to validate it. Worth noting register pressure alone does not predict throughput here. - **Benchmarking caveat**: operand data matters a lot. A cycling byte fill switches far fewer transistors than realistic randn-derived fp8, so the GPU holds a ~20% higher sustained clock — 218 vs 180 TFLOPS on the same kernel and shape. All numbers here use pseudorandom operands. Timing two policies in back-to-back blocks rather than interleaved also makes the second look ~30% slower even when both select the identical config. - **Where the win concentrates** (thanks @0xDELUXA): sweeping K on a fixed 2048x2048, the speedup over the old selector is concentrated where K is a multiple of 4096 (1.35x at K=4096, 1.98x at K=8192, vs ~1.03x at 3072/6144). The likely mechanism is the deleted `256x128` tile hitting cache aliasing when the row stride is a multiple of 4096 bytes. If you benchmark this PR at, say, K=3072 only, it will look like a no-op — that is expected. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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

closed this issue
added the
pull-request
label

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

Comfy-Org/comfy-kitchen#140

No description provided.

Rows
Columns