[PR #139] [ROCm] Add gfx12 INT8 WMMA GEMM schedule and fused INT8 ops for Z-Image #167

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/139
Author: @tvukovic-amd
Created: 8/27/2026
Status: 🔄 Open

Base: main ← Head: hip-gfx12x-native-attention

📝 Commits (10+)

6a94830 Add native gfx12 HIP attention and Z-Image INT8 fusions
f8ef6fc Address review feedback for gfx12 HIP kernels
eb6270f reslve code rabbit review comments
e70bc08 Merge branch 'main' into hip-gfx12x-native-attention
638e94c Use atomic WMMA state snapshots, scope ConvRot capability checks to the input device, and preserve float32 validation for HIP attention scales.
99ff431 Fix WMMA races and validate attention scales
8f0b518 Fix WMMA races and validate attention scales
ba44363 Fix Ruff lint
71c5a34 hip: enable native gfx1170 attention and GEMM schedules
9518b7e Merge branch 'main' into hip-gfx12x-native-attention

📊 Changes

25 files changed (+4133 additions, -96 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+228 -0)
📝 comfy_kitchen/backends/eager/__init__.py (+79 -0)
📝 comfy_kitchen/backends/eager/quantization.py (+150 -0)
➕ comfy_kitchen/backends/eager/residual.py (+18 -0)
📝 comfy_kitchen/backends/hip/CMakeLists.txt (+2 -0)
📝 comfy_kitchen/backends/hip/__init__.py (+727 -36)
➕ comfy_kitchen/backends/hip/device_arch.h (+45 -0)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+279 -14)
📝 comfy_kitchen/backends/hip/epilogue.h (+14 -0)
📝 comfy_kitchen/backends/hip/gemm_wmma.h (+518 -10)
📝 comfy_kitchen/backends/hip/hadamard.h (+491 -1)
📝 comfy_kitchen/backends/hip/launchers.h (+21 -0)
📝 comfy_kitchen/backends/hip/mma.h (+1 -1)
📝 comfy_kitchen/backends/hip/ops/gemm_int8.hip (+169 -1)
📝 comfy_kitchen/backends/hip/ops/quantize_int8.hip (+91 -10)
➕ comfy_kitchen/backends/hip/ops/rms_convrot_quant.hip (+35 -0)
➕ comfy_kitchen/backends/hip/ops/rms_gated_residual.hip (+180 -0)
📝 comfy_kitchen/backends/hip/ops/rms_rope.hip (+215 -6)
📝 comfy_kitchen/backends/hip/rope_math.h (+23 -0)
➕ comfy_kitchen/backends/hip/swiglu_bf16.h (+41 -0)

...and 5 more files

📄 Description

Summary

Adds INT8 WMMA GEMM optimizations for gfx12 and fused INT8 operations for accelerated Z-Image Turbo inference.

Key changes

Fused INT8 operations (new public APIs)

int8_linear_modulated and int8_linear_rms_modulated: ConvRot quantization fused with modulation (and RMSNorm), then an INT8 GEMM.
int8_linear_pair, int8_linear_pair_modulated and int8_linear_pair_rms_modulated: two projections that share one quantized activation.
int8_linear_swiglu_split: SwiGLU fused with quantization for the FFN down projection.
rms_gated_residual: RMSNorm, gate and residual add in one kernel.
fused_rms_modulated and fused_swiglu_ffn on the INT8 quantized-tensor layout, used by the ComfyUI integration.
A wave32 RMSNorm + RoPE kernel.
Eager reference implementations for every fused op.

INT8 GEMM and quantization

A new gfx12 WMMA INT8 GEMM schedule (tiled-A, dual-M).
A 512×2 ConvRot quantization schedule, with a spill path for large K.
Both schedules are gated by a shared cached device query (device_arch.h) that checks for the non-duplicated WMMA layout.

Compatibility

gfx11 keeps upstream's launch_gemm_wmma tile selection. Unsupported devices, batched modulation and modulation without a norm return NotImplemented and use the existing fallbacks.
Adds support for gfx1170 and gfx1171, which is enabled together with #141.

Required ComfyUI integration

This PR pairs with https://github.com/Comfy-Org/ComfyUI/pull/16674, which adds the Z-Image fusion integration, the attention command-line options and Dynamic VRAM handling. Both are needed to reproduce the benchmark below.

Testing

Tested on AMD Radeon RX 9070 XT (Navi 48, gfx1201, 16 GB), Ubuntu 24.04.2 LTS with https://github.com/Comfy-Org/ComfyUI/pull/16674's head.

python -m pytest tests/test_hip_dispatch.py tests/test_hip_wmma.py tests/test_int8.py -q

Result: 751 passed, 47 skipped, 3 failed. The three failures (test_rope_matches_eager[freqs_dtype1-True], test_rope_matches_eager[freqs_dtype2-True], test_int8_linear_correctness[hip]) also fail on the base branch, so this PR doesn't introduce them.

Z-Image Turbo benchmark (gfx1201)

Protocol: 1024×1024, 8 steps, batch size 1, Dynamic VRAM on. One warmup and five timed runs per row. Values are the median end-to-end workflow throughput over the five runs. The baseline is BF16 + AOTriton.

Configuration	Median	Speedup	Uplift	Loaded VRAM
BF16 + AOTriton (base)	1.63 it/s	1.000×	—	14.79 GiB
BF16 + Comfy Kitchen attention	1.81 it/s	1.109×	+10.9%	14.70 GiB
INT8 + AOTriton	2.17 it/s	1.332×	+33.2%	11.67 GiB
INT8 + Comfy Kitchen attention	2.55 it/s	1.561×	+56.1%	11.67 GiB

Big thanks to @jammm for contribution!

_{🔄 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/139 **Author:** [@tvukovic-amd](https://github.com/tvukovic-amd) **Created:** 8/27/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `hip-gfx12x-native-attention` --- ### 📝 Commits (10+) - [`6a94830`](https://github.com/Comfy-Org/comfy-kitchen/commit/6a94830bedf1c2cf59f8b2e1b304e6e05edc884a) Add native gfx12 HIP attention and Z-Image INT8 fusions - [`f8ef6fc`](https://github.com/Comfy-Org/comfy-kitchen/commit/f8ef6fc5958b0f6475927526d3c1b9f9d84184d1) Address review feedback for gfx12 HIP kernels - [`eb6270f`](https://github.com/Comfy-Org/comfy-kitchen/commit/eb6270f317fa5e7b5751a614447eccef52545c57) reslve code rabbit review comments - [`e70bc08`](https://github.com/Comfy-Org/comfy-kitchen/commit/e70bc0878928901e76ec0022523be98c2261c94f) Merge branch 'main' into hip-gfx12x-native-attention - [`638e94c`](https://github.com/Comfy-Org/comfy-kitchen/commit/638e94c247b727b6699e0e4ad0115882a3e4328e) Use atomic WMMA state snapshots, scope ConvRot capability checks to the input device, and preserve float32 validation for HIP attention scales. - [`99ff431`](https://github.com/Comfy-Org/comfy-kitchen/commit/99ff4310c78eec5d9cfd58359092a33e6055eee7) Fix WMMA races and validate attention scales - [`8f0b518`](https://github.com/Comfy-Org/comfy-kitchen/commit/8f0b518a2685fdcadba755f1f4ce821dc1463000) Fix WMMA races and validate attention scales - [`ba44363`](https://github.com/Comfy-Org/comfy-kitchen/commit/ba4436363892e5765cae7e7ef5e07f504eefb1ad) Fix Ruff lint - [`71c5a34`](https://github.com/Comfy-Org/comfy-kitchen/commit/71c5a34ed205dcca53299fd118222b3b74e8ad0c) hip: enable native gfx1170 attention and GEMM schedules - [`9518b7e`](https://github.com/Comfy-Org/comfy-kitchen/commit/9518b7e2b1ca1804d78352c7421f9fbc2383091c) Merge branch 'main' into hip-gfx12x-native-attention ### 📊 Changes **25 files changed** (+4133 additions, -96 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+228 -0) 📝 `comfy_kitchen/backends/eager/__init__.py` (+79 -0) 📝 `comfy_kitchen/backends/eager/quantization.py` (+150 -0) ➕ `comfy_kitchen/backends/eager/residual.py` (+18 -0) 📝 `comfy_kitchen/backends/hip/CMakeLists.txt` (+2 -0) 📝 `comfy_kitchen/backends/hip/__init__.py` (+727 -36) ➕ `comfy_kitchen/backends/hip/device_arch.h` (+45 -0) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+279 -14) 📝 `comfy_kitchen/backends/hip/epilogue.h` (+14 -0) 📝 `comfy_kitchen/backends/hip/gemm_wmma.h` (+518 -10) 📝 `comfy_kitchen/backends/hip/hadamard.h` (+491 -1) 📝 `comfy_kitchen/backends/hip/launchers.h` (+21 -0) 📝 `comfy_kitchen/backends/hip/mma.h` (+1 -1) 📝 `comfy_kitchen/backends/hip/ops/gemm_int8.hip` (+169 -1) 📝 `comfy_kitchen/backends/hip/ops/quantize_int8.hip` (+91 -10) ➕ `comfy_kitchen/backends/hip/ops/rms_convrot_quant.hip` (+35 -0) ➕ `comfy_kitchen/backends/hip/ops/rms_gated_residual.hip` (+180 -0) 📝 `comfy_kitchen/backends/hip/ops/rms_rope.hip` (+215 -6) 📝 `comfy_kitchen/backends/hip/rope_math.h` (+23 -0) ➕ `comfy_kitchen/backends/hip/swiglu_bf16.h` (+41 -0) _...and 5 more files_ </details> ### 📄 Description ## Summary Adds INT8 WMMA GEMM optimizations for gfx12 and fused INT8 operations for accelerated Z-Image Turbo inference. ## Key changes **Fused INT8 operations (new public APIs)** - `int8_linear_modulated` and `int8_linear_rms_modulated`: ConvRot quantization fused with modulation (and RMSNorm), then an INT8 GEMM. - `int8_linear_pair`, `int8_linear_pair_modulated` and `int8_linear_pair_rms_modulated`: two projections that share one quantized activation. - `int8_linear_swiglu_split`: SwiGLU fused with quantization for the FFN down projection. - `rms_gated_residual`: RMSNorm, gate and residual add in one kernel. - `fused_rms_modulated` and `fused_swiglu_ffn` on the INT8 quantized-tensor layout, used by the ComfyUI integration. - A wave32 RMSNorm + RoPE kernel. - Eager reference implementations for every fused op. **INT8 GEMM and quantization** - A new gfx12 WMMA INT8 GEMM schedule (tiled-A, dual-M). - A 512×2 ConvRot quantization schedule, with a spill path for large K. - Both schedules are gated by a shared cached device query (`device_arch.h`) that checks for the non-duplicated WMMA layout. **Compatibility** - gfx11 keeps upstream's `launch_gemm_wmma` tile selection. Unsupported devices, batched modulation and modulation without a norm return `NotImplemented` and use the existing fallbacks. - Adds support for gfx1170 and gfx1171, which is enabled together with #141. ## Required ComfyUI integration This PR pairs with https://github.com/Comfy-Org/ComfyUI/pull/16674, which adds the Z-Image fusion integration, the attention command-line options and Dynamic VRAM handling. Both are needed to reproduce the benchmark below. ## Testing Tested on AMD Radeon RX 9070 XT (Navi 48, gfx1201, 16 GB), Ubuntu 24.04.2 LTS with https://github.com/Comfy-Org/ComfyUI/pull/16674's head. ``` python -m pytest tests/test_hip_dispatch.py tests/test_hip_wmma.py tests/test_int8.py -q ``` Result: **751 passed, 47 skipped, 3 failed**. The three failures (`test_rope_matches_eager[freqs_dtype1-True]`, `test_rope_matches_eager[freqs_dtype2-True]`, `test_int8_linear_correctness[hip]`) also fail on the base branch, so this PR doesn't introduce them. ## Z-Image Turbo benchmark (gfx1201) Protocol: 1024×1024, 8 steps, batch size 1, Dynamic VRAM on. One warmup and five timed runs per row. Values are the median end-to-end workflow throughput over the five runs. The baseline is BF16 + AOTriton. | Configuration | Median | Speedup | Uplift | Loaded VRAM | |---|---|---|---|---| | BF16 + AOTriton (base) | 1.63 it/s | 1.000× | — | 14.79 GiB | | BF16 + Comfy Kitchen attention | 1.81 it/s | 1.109× | +10.9% | 14.70 GiB | | INT8 + AOTriton | 2.17 it/s | 1.332× | +33.2% | 11.67 GiB | | INT8 + Comfy Kitchen attention | 2.55 it/s | 1.561× | +56.1% | 11.67 GiB | Big thanks to @jammm for contribution! --- <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:05 +00:00

bsakaguchi referenced this issue

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

[PR #167] [MERGED] MiniMax H3 VAE kernels: fused encoder pad/norm, fp16-accumulate conv3d and GEMM, int8 residual epilogue #195

bsakaguchi referenced this issue

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

[PR #175] [MERGED] HIP port of #167: MiniMax H3 VAE kernels #197

bsakaguchi referenced this issue

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