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 #174] Optimize HIP WMMA INT8 GEMM staging, epilogue, and deep-K tiling #194

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: optimize-hip-int8-wmma-gemm

📝 Commits (5)

6f26009 Optimize HIP WMMA INT8 GEMM
4711125 Preserve signed zero without INT8 bias
2ab6a7c Use the 128x256 BKB=128 INT8 WMMA tile on gfx12 at every K
21c422d Keep the cached INT8 epilogue scale order exact
06c5ee1 Keep gfx12 INT8 linear on the row-major route

📊 Changes

4 files changed (+218 additions, -14 deletions)

View changed files

📝 comfy_kitchen/backends/hip/__init__.py (+8 -0)
📝 comfy_kitchen/backends/hip/epilogue.h (+25 -0)
📝 comfy_kitchen/backends/hip/gemm_wmma.h (+164 -14)
📝 tests/test_hip_wmma.py (+21 -0)

📄 Description

Summary

skip per-chunk bounds checks for block-uniform full operand tiles
cache row/column scales and bias values in the rowwise epilogue, with a bounds-check-free full output tile path
make grouped-M ordering a kernel template parameter while preserving GroupM=4 for existing instantiations
add a balanced 128x256, 2x4-wave, TM4xTN4, GroupM=8 INT8 tile on gfx12, with BKB=128, used at every K (see the update below; it was BKB=64 and K >= 4096 only)
preserve the no-bias epilogue's signed-zero behavior and cover it with a raw WMMA GEMM regression test

The new tile is selected with a cached host-side HIP device property query. The generated COMFY_MMA_GFX12 macro is intentionally not used in the launcher because the __gfx*__ conditions that define it are device-pass-only. If the runtime architecture query fails, or the current device is not gfx12, selection falls back to the existing cross-architecture tile policy.

FP8 and INT4 tile selection is unchanged. The common staging optimization and defaulted GROUP_M template parameter are shared infrastructure.

Performance

Raw HIP event timing on an RX 9070 XT (gfx1201), PyTorch 2.15.0a0+rocm10.1.0a20260822, HIP 7.16.26332, BF16 output, M=30000, N=21504, K=5376. Each ABBA run used 5 warmups and 15 measured iterations with deterministic inputs.

Build/order	Median
upstream A1	39.347 ms
optimized B1	33.049 ms
optimized B2	33.055 ms
upstream A2	39.340 ms

The mean of the two run medians improves from 39.3435 ms to 33.052 ms, a cumulative 19.04% speedup (176.3 to 209.9 TOPS). This comparison is the complete stack in this PR against current upstream, not the isolated effect of the new geometry.

With the other optimizations held constant, the new tile was also validated across QKV M sweeps with no observed regression: about 3.0-3.5% over the previous canonical tile on gfx1201 and 5.08% at M=30000 on gfx1200. The GroupM=8 ablation contributed 1.59% at gfx1201 QKV M=30000 and was neutral-to-positive across the tested projection shapes.

Static resource inspection of the winning specialization on gfx1200 and gfx1201 reported 23 SGPRs, 221 VGPRs, 27,648 bytes LDS, zero scratch, and no spills.

Validation

built a gfx1200;gfx1201 fat HIP extension
deterministic INT8 correctness on RX 9060 XT/gfx1200 and RX 9070 XT/gfx1201: exact match, max absolute error 0.0
tests/test_hip_wmma.py tests/test_hip_dispatch.py tests/test_gemm_bias_contract.py: 452 passed, 11 skipped on each GPU
compile-checked gfx1100;gfx1151
git diff --check

Update: rebased onto v0.2.38, BKB=128 at every K

Rebased onto current main (42081bf, v0.2.38) without conflicts, plus one commit, Use the 128x256 BKB=128 INT8 WMMA tile on gfx12 at every K. The gfx12 INT8 tile now uses BKB=128, which halves the K-loop barriers, and covers shallow K too. Before this commit, K < 4096 used the cross-architecture 128x128x64 tile. Outputs are bit-identical (int32 accumulation).

RX 9070 XT (gfx1201), PyTorch 2.16.0a0+rocm10.2.0a20261015, HIP event timing, int8_linear with ConvRot (activation quant included). "Before" is this PR's previous tile policy.

MOSS-Image 1024x1536 DiT projections, M=6304:

projection (K -> N)	before	after
QKV 3072 -> 9216	1.86 ms	1.81 ms
out 3072 -> 3072	0.82 ms	0.77 ms
FFN up/gate 3072 -> 16384	3.53 ms	3.28 ms
FFN down (SiTU) 8192 -> 3072	2.34 ms	2.13 ms

The deep-K GEMM alone at 8192 -> 3072 goes from about 2.5 ms (BKB=64) to about 1.5 ms (212 TOPS). The full MOSS-Image DiT forward goes from 535 to 508 ms.

M sweep (256 to 4096), speedup of int8_linear: K=3072/N=3072 1.01-1.07x, K=3072/N=9216 1.00-1.09x, K=8192/N=3072 1.03-2.02x, K=4096/N=4096 1.07-1.36x. M=128 stays below the WGP threshold and keeps the existing small tile.

H3 projections (already on the previous deep-K tile), M=15000/30000, alternating runs:

projection (K -> N)	before	after
QKV 5376 -> 16128, M=30000	25.81 ms	25.38 ms
out 5376 -> 5376, M=30000	9.25 ms	9.24 ms
MLP up 5376 -> 14336, M=30000	23.08 ms	22.71 ms
MLP down 14336 -> 5376, M=30000	24.59 ms	24.33 ms
5376 -> 21504, M=30000 (shape above)	34.10 ms	33.50 ms

That is +1.0-1.8% across the H3 shapes, and the same at M=15000.

Resources of the new specialization (bf16, fp16 and fp32 outputs alike, gfx1200 and gfx1201): 256 VGPRs, 3 VGPRs spilled (16 bytes/lane of scratch), 52,224 bytes LDS, occupancy 4 waves/SIMD. Unlike the previous BKB=64 tile, this one spills a little. It is still the fastest of the tiles swept here (15 configurations) at every shape above, and the numbers above include the spill.

Re-validation after the rebase, 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_w4a8_requant.py tests/test_convrot_unpacked_quantization.py tests/test_adaln.py tests/test_rms_rope.py tests/test_hip_convrot_perf.py tests/test_gemm_bias_contract.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). The conflict with #259 in gemm_wmma.h is resolved by keeping both sides:

GROUP_M now follows #259's COALESCED_STAGING and TILED_A template parameters.
The full-tile fast path in TileStager::load uses the same chunk index as store, so it is also correct for coalesced staging.

Two commits are added:

Keep the cached INT8 epilogue scale order exact. Fast-math folded the cached scales into acc * (sa * sb). That rounds differently from (acc * sa) * sb in rare bf16 cases (1-26 of 0.6M-5.9M outputs), and the new test_gfx12_int8_gemm_is_exact caught it. apply_cached now pins the association, and the test passes.
Keep gfx12 INT8 linear on the row-major route. #259 sends gfx12/gfx117 int8_linear through the tiled-A producer and int8_gemm_a_tiled128. With this PR's 128x256 BKB=128 tile, the row-major route is faster at every MOSS-Image/H3 projection, and its outputs are bit-identical to the tiled-A route. gfx117 keeps tiled-A, and COMFY_KITCHEN_INT8_TILED_A=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:

shape (M x K -> N)	9070 XT main	9070 XT PR	9060 XT main	9060 XT PR
MOSS qkv 12544 x 3072 -> 9216	3.985 ms	3.733 ms	7.447 ms	7.049 ms
MOSS out 12544 x 3072 -> 3072	1.609	1.511	2.766	2.667
MOSS up/gate 12544 x 3072 -> 16384	6.805	6.380	13.067	12.177
MOSS down 12544 x 8192 -> 3072	3.771	3.579	7.028	6.815
H3 qkv 38339 x 5376 -> 21504	46.17	42.31	89.38	82.43
H3 out 38339 x 7168 -> 5376	16.15	15.13	31.45	29.34
H3 fc1 38339 x 5376 -> 28672	61.20	56.07	117.57	108.16
H3 fc2 38339 x 14336 -> 5376	32.67	30.21	64.45	59.69

With COMFY_KITCHEN_INT8_TILED_A=1, the PR build matches main within about 1%. The gain therefore comes from the row-major tile.

On the 9070 XT, the raw int8_gemm (no activation quant) goes from 4.09 to 3.14 ms at MOSS qkv and from 68.6 to 55.0 ms at H3 fc1. Its outputs are identical to 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 signed-zero test added here.
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/174 **Author:** [@qweqweewqe7-create](https://github.com/qweqweewqe7-create) **Created:** 9/13/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `optimize-hip-int8-wmma-gemm` --- ### 📝 Commits (5) - [`6f26009`](https://github.com/Comfy-Org/comfy-kitchen/commit/6f26009e44d7e37d05ff162599e03a6d6b49827f) Optimize HIP WMMA INT8 GEMM - [`4711125`](https://github.com/Comfy-Org/comfy-kitchen/commit/47111256e5539b0bfa7ad35e227ee5acc20df65d) Preserve signed zero without INT8 bias - [`2ab6a7c`](https://github.com/Comfy-Org/comfy-kitchen/commit/2ab6a7cbc610214a662088ba774c3b6ed59ee50e) Use the 128x256 BKB=128 INT8 WMMA tile on gfx12 at every K - [`21c422d`](https://github.com/Comfy-Org/comfy-kitchen/commit/21c422deaa858b8177ae088cb99902be8af0d60a) Keep the cached INT8 epilogue scale order exact - [`06c5ee1`](https://github.com/Comfy-Org/comfy-kitchen/commit/06c5ee157c3174050c7f1f646899003c750720ff) Keep gfx12 INT8 linear on the row-major route ### 📊 Changes **4 files changed** (+218 additions, -14 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/hip/__init__.py` (+8 -0) 📝 `comfy_kitchen/backends/hip/epilogue.h` (+25 -0) 📝 `comfy_kitchen/backends/hip/gemm_wmma.h` (+164 -14) 📝 `tests/test_hip_wmma.py` (+21 -0) </details> ### 📄 Description ## Summary - skip per-chunk bounds checks for block-uniform full operand tiles - cache row/column scales and bias values in the rowwise epilogue, with a bounds-check-free full output tile path - make grouped-M ordering a kernel template parameter while preserving `GroupM=4` for existing instantiations - add a balanced `128x256`, 2x4-wave, `TM4xTN4`, `GroupM=8` INT8 tile on gfx12, with `BKB=128`, used at every K (see the update below; it was `BKB=64` and K >= 4096 only) - preserve the no-bias epilogue's signed-zero behavior and cover it with a raw WMMA GEMM regression test The new tile is selected with a cached host-side HIP device property query. The generated `COMFY_MMA_GFX12` macro is intentionally not used in the launcher because the `__gfx*__` conditions that define it are device-pass-only. If the runtime architecture query fails, or the current device is not gfx12, selection falls back to the existing cross-architecture tile policy. FP8 and INT4 tile selection is unchanged. The common staging optimization and defaulted `GROUP_M` template parameter are shared infrastructure. ## Performance Raw HIP event timing on an RX 9070 XT (`gfx1201`), PyTorch `2.15.0a0+rocm10.1.0a20260822`, HIP `7.16.26332`, BF16 output, `M=30000, N=21504, K=5376`. Each ABBA run used 5 warmups and 15 measured iterations with deterministic inputs. | Build/order | Median | | --- | ---: | | upstream A1 | 39.347 ms | | optimized B1 | 33.049 ms | | optimized B2 | 33.055 ms | | upstream A2 | 39.340 ms | The mean of the two run medians improves from 39.3435 ms to 33.052 ms, a cumulative 19.04% speedup (176.3 to 209.9 TOPS). This comparison is the complete stack in this PR against current upstream, not the isolated effect of the new geometry. With the other optimizations held constant, the new tile was also validated across QKV M sweeps with no observed regression: about 3.0-3.5% over the previous canonical tile on gfx1201 and 5.08% at M=30000 on gfx1200. The `GroupM=8` ablation contributed 1.59% at gfx1201 QKV M=30000 and was neutral-to-positive across the tested projection shapes. Static resource inspection of the winning specialization on gfx1200 and gfx1201 reported 23 SGPRs, 221 VGPRs, 27,648 bytes LDS, zero scratch, and no spills. ## Validation - built a `gfx1200;gfx1201` fat HIP extension - deterministic INT8 correctness on RX 9060 XT/gfx1200 and RX 9070 XT/gfx1201: exact match, max absolute error 0.0 - `tests/test_hip_wmma.py tests/test_hip_dispatch.py tests/test_gemm_bias_contract.py`: 452 passed, 11 skipped on each GPU - compile-checked `gfx1100;gfx1151` - `git diff --check` ## Update: rebased onto v0.2.38, BKB=128 at every K Rebased onto current `main` (`42081bf`, v0.2.38) without conflicts, plus one commit, *Use the 128x256 BKB=128 INT8 WMMA tile on gfx12 at every K*. The gfx12 INT8 tile now uses `BKB=128`, which halves the K-loop barriers, and covers shallow K too. Before this commit, `K < 4096` used the cross-architecture 128x128x64 tile. Outputs are bit-identical (int32 accumulation). RX 9070 XT (gfx1201), PyTorch `2.16.0a0+rocm10.2.0a20261015`, HIP event timing, `int8_linear` with ConvRot (activation quant included). "Before" is this PR's previous tile policy. MOSS-Image 1024x1536 DiT projections, M=6304: | projection (K -> N) | before | after | | --- | ---: | ---: | | QKV 3072 -> 9216 | 1.86 ms | 1.81 ms | | out 3072 -> 3072 | 0.82 ms | 0.77 ms | | FFN up/gate 3072 -> 16384 | 3.53 ms | 3.28 ms | | FFN down (SiTU) 8192 -> 3072 | 2.34 ms | 2.13 ms | The deep-K GEMM alone at 8192 -> 3072 goes from about 2.5 ms (`BKB=64`) to about 1.5 ms (212 TOPS). The full MOSS-Image DiT forward goes from 535 to 508 ms. M sweep (256 to 4096), speedup of `int8_linear`: K=3072/N=3072 1.01-1.07x, K=3072/N=9216 1.00-1.09x, K=8192/N=3072 1.03-2.02x, K=4096/N=4096 1.07-1.36x. M=128 stays below the WGP threshold and keeps the existing small tile. H3 projections (already on the previous deep-K tile), M=15000/30000, alternating runs: | projection (K -> N) | before | after | | --- | ---: | ---: | | QKV 5376 -> 16128, M=30000 | 25.81 ms | 25.38 ms | | out 5376 -> 5376, M=30000 | 9.25 ms | 9.24 ms | | MLP up 5376 -> 14336, M=30000 | 23.08 ms | 22.71 ms | | MLP down 14336 -> 5376, M=30000 | 24.59 ms | 24.33 ms | | 5376 -> 21504, M=30000 (shape above) | 34.10 ms | 33.50 ms | That is +1.0-1.8% across the H3 shapes, and the same at M=15000. Resources of the new specialization (bf16, fp16 and fp32 outputs alike, gfx1200 and gfx1201): 256 VGPRs, 3 VGPRs spilled (16 bytes/lane of scratch), 52,224 bytes LDS, occupancy 4 waves/SIMD. Unlike the previous `BKB=64` tile, this one spills a little. It is still the fastest of the tiles swept here (15 configurations) at every shape above, and the numbers above include the spill. Re-validation after the rebase, 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_w4a8_requant.py tests/test_convrot_unpacked_quantization.py tests/test_adaln.py tests/test_rms_rope.py tests/test_hip_convrot_perf.py tests/test_gemm_bias_contract.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). The conflict with #259 in `gemm_wmma.h` is resolved by keeping both sides: - `GROUP_M` now follows #259's `COALESCED_STAGING` and `TILED_A` template parameters. - The full-tile fast path in `TileStager::load` uses the same chunk index as `store`, so it is also correct for coalesced staging. Two commits are added: - *Keep the cached INT8 epilogue scale order exact.* Fast-math folded the cached scales into `acc * (sa * sb)`. That rounds differently from `(acc * sa) * sb` in rare bf16 cases (1-26 of 0.6M-5.9M outputs), and the new `test_gfx12_int8_gemm_is_exact` caught it. `apply_cached` now pins the association, and the test passes. - *Keep gfx12 INT8 linear on the row-major route.* #259 sends gfx12/gfx117 `int8_linear` through the tiled-A producer and `int8_gemm_a_tiled128`. With this PR's 128x256 BKB=128 tile, the row-major route is faster at every MOSS-Image/H3 projection, and its outputs are bit-identical to the tiled-A route. gfx117 keeps tiled-A, and `COMFY_KITCHEN_INT8_TILED_A=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: | shape (M x K -> N) | 9070 XT main | 9070 XT PR | 9060 XT main | 9060 XT PR | | --- | ---: | ---: | ---: | ---: | | MOSS qkv 12544 x 3072 -> 9216 | 3.985 ms | 3.733 ms | 7.447 ms | 7.049 ms | | MOSS out 12544 x 3072 -> 3072 | 1.609 | 1.511 | 2.766 | 2.667 | | MOSS up/gate 12544 x 3072 -> 16384 | 6.805 | 6.380 | 13.067 | 12.177 | | MOSS down 12544 x 8192 -> 3072 | 3.771 | 3.579 | 7.028 | 6.815 | | H3 qkv 38339 x 5376 -> 21504 | 46.17 | 42.31 | 89.38 | 82.43 | | H3 out 38339 x 7168 -> 5376 | 16.15 | 15.13 | 31.45 | 29.34 | | H3 fc1 38339 x 5376 -> 28672 | 61.20 | 56.07 | 117.57 | 108.16 | | H3 fc2 38339 x 14336 -> 5376 | 32.67 | 30.21 | 64.45 | 59.69 | With `COMFY_KITCHEN_INT8_TILED_A=1`, the PR build matches `main` within about 1%. The gain therefore comes from the row-major tile. On the 9070 XT, the raw `int8_gemm` (no activation quant) goes from 4.09 to 3.14 ms at MOSS qkv and from 68.6 to 55.0 ms at H3 fc1. Its outputs are identical to `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 signed-zero test added here. - 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:26 +00:00

bsakaguchi referenced this issue

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

[PR #194] [MERGED] comfy kitchen int8 attention speedups on HIP. #213