[PR #122] [MERGED] Occupancy-aware HIP WMMA GEMM tile selection and vectorized fp8 elementwise kernels #160

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/122
Author: @0xDELUXA
Created: 8/18/2026
Status: ✅ Merged
Merged: 8/25/2026
Merged by: @comfyanonymous

Base: main ← Head: amd/hip-gemm-tile-selection

📝 Commits (2)

897aaba Occupancy-aware HIP WMMA GEMM tile selection and vectorized fp8 elementwise kernels
9192c77 Vectorize the stochastic-rounding fp8 quantize kernel

📊 Changes

8 files changed (+379 additions, -88 deletions)

View changed files

📝 comfy_kitchen/backends/hip/fp8_utils.h (+16 -0)
📝 comfy_kitchen/backends/hip/gemm_wmma.h (+78 -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 (+120 -12)
📝 comfy_kitchen/backends/hip/ops/stochastic_round_fp8.hip (+80 -37)
📝 tests/test_hip_wmma.py (+28 -2)
📝 tests/test_qdq.py (+55 -1)

📄 Description

This PR is a continuation of #94 by @crashingalexsan and carries two commits. The first cherry picks the author's commits with their permission, squashed into one commit that keeps their authorship, rebased onto current main. The second extends the same vectorization to stochastic_round_fp8, the one fp8 elementwise kernel #94 left scalar.

Performance-only change: no public API or numerics are changed, and all paths remain bit-identical.

GEMM tile selection

Replaces the previous M/N/K threshold-based WMMA tile selection with a shared launch_gemm_wmma() in gemm_wmma.h, used by both fp8 and int8 launchers. Selection now considers grid coverage, K depth, and warp grid.

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 selecting the same config stay within +/- 4%. Outputs are bit-identical.

Elementwise kernels

per_tensor_fp8 now has a 16-elements/thread path using 16-byte loads/stores when both pointers are aligned, with the original scalar path retained for misaligned views.

At 64M elements:

Vectorized: 290 GB/s
Scalar: 148 GB/s
Speedup: 1.96x
Output: bit-identical

Stochastic rounding

stochastic_round_fp8 receives the same vectorized path, reaching 297 GB/s at 64M elements vs 142 GB/s previously.

elements	scalar	vectorized	speedup
1M	41.1 us	34.8 us	1.18x
4M	107.5 us	81.3 us	1.32x
16M	435.5 us	231.5 us	1.88x
64M	1886.1 us	903.4 us	2.09x

The per-element logic is shared through stochastic_round_one() to guarantee identical rounding between paths. Vectorization helpers are shared through fp8_utils.h.

Scalar/vectorized paths were verified bit-identical across float32/float16/bfloat16, both e4m3fn and e5m2, multiple sizes, and edge cases including NaNs, infinities, +/- 0, +/- 448, 57344 and subnormals.

RDNA coverage

Run-verified: gfx1200
Compile-verified: gfx1030, gfx1100, gfx1151, gfx1201
Vectorized kernels use 128-bit loads/stores, spill nothing, and peak at 37 VGPRs.
On gfx11, fp8 scratch usage drops from 1440 B -> 804 B per output dtype.
The worst gfx11 spiller, the 256x128 BKB=64 tile, is removed.

Testing

gfx1200:

472 passed, 16 failed, 4 skipped
main: 416 passed, 16 failed, 4 skipped

The same 16 NVFP4 failures occur on both revisions and originate from comfy-kitchen's Triton backend on AMD, unrelated to this PR.

The 56 additional passing tests cover deep-K/skinny GEMMs and misaligned-view elementwise cases.

_{🔄 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/122 **Author:** [@0xDELUXA](https://github.com/0xDELUXA) **Created:** 8/18/2026 **Status:** ✅ Merged **Merged:** 8/25/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `amd/hip-gemm-tile-selection` --- ### 📝 Commits (2) - [`897aaba`](https://github.com/Comfy-Org/comfy-kitchen/commit/897aaba5621b94bf166f5107da957ae11dd6153b) Occupancy-aware HIP WMMA GEMM tile selection and vectorized fp8 elementwise kernels - [`9192c77`](https://github.com/Comfy-Org/comfy-kitchen/commit/9192c77751175c67ee76c7ef7ebb1327642320d9) Vectorize the stochastic-rounding fp8 quantize kernel ### 📊 Changes **8 files changed** (+379 additions, -88 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/hip/fp8_utils.h` (+16 -0) 📝 `comfy_kitchen/backends/hip/gemm_wmma.h` (+78 -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` (+120 -12) 📝 `comfy_kitchen/backends/hip/ops/stochastic_round_fp8.hip` (+80 -37) 📝 `tests/test_hip_wmma.py` (+28 -2) 📝 `tests/test_qdq.py` (+55 -1) </details> ### 📄 Description This PR is a continuation of #94 by @crashingalexsan and carries two commits. The first cherry picks the author's commits with their permission, squashed into one commit that keeps their authorship, rebased onto current `main`. The second extends the same vectorization to `stochastic_round_fp8`, the one fp8 elementwise kernel #94 left scalar. Performance-only change: no public API or numerics are changed, and all paths remain bit-identical. ## GEMM tile selection Replaces the previous `M`/`N`/`K` threshold-based WMMA tile selection with a shared `launch_gemm_wmma()` in `gemm_wmma.h`, used by both fp8 and int8 launchers. Selection now considers grid coverage, K depth, and warp grid. | 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 selecting the same config stay within +/- 4%. Outputs are bit-identical. ## Elementwise kernels `per_tensor_fp8` now has a 16-elements/thread path using 16-byte loads/stores when both pointers are aligned, with the original scalar path retained for misaligned views. At 64M elements: - Vectorized: 290 GB/s - Scalar: 148 GB/s - Speedup: 1.96x - Output: bit-identical ## Stochastic rounding `stochastic_round_fp8` receives the same vectorized path, reaching 297 GB/s at 64M elements vs 142 GB/s previously. | elements | scalar | vectorized | speedup | |---|---:|---:|---:| | 1M | 41.1 us | 34.8 us | 1.18x | | 4M | 107.5 us | 81.3 us | 1.32x | | 16M | 435.5 us | 231.5 us | 1.88x | | 64M | 1886.1 us | 903.4 us | 2.09x | The per-element logic is shared through `stochastic_round_one()` to guarantee identical rounding between paths. Vectorization helpers are shared through `fp8_utils.h`. Scalar/vectorized paths were verified bit-identical across `float32`/`float16`/`bfloat16`, both `e4m3fn` and `e5m2`, multiple sizes, and edge cases including NaNs, infinities, +/- 0, +/- 448, 57344 and subnormals. ## RDNA coverage - Run-verified: gfx1200 - Compile-verified: gfx1030, gfx1100, gfx1151, gfx1201 - Vectorized kernels use 128-bit loads/stores, spill nothing, and peak at 37 VGPRs. - On gfx11, fp8 scratch usage drops from 1440 B -> 804 B per output dtype. - The worst gfx11 spiller, the `256x128 BKB=64` tile, is removed. ## Testing gfx1200: - 472 passed, 16 failed, 4 skipped - `main`: 416 passed, 16 failed, 4 skipped The same 16 NVFP4 failures occur on both revisions and originate from comfy-kitchen's Triton backend on AMD, unrelated to this PR. The 56 additional passing tests cover deep-K/skinny GEMMs and misaligned-view elementwise cases. --- <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:59 +00:00