[PR #97] [MERGED] HIP port of #90: grouped W4A8 over the INT8 GEMM #143

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: amd/hip-w4a8-int8

📝 Commits (1)

e66f734 HIP port of #90: grouped W4A8 over the INT8 GEMM

📊 Changes

10 files changed (+839 additions, -24 deletions)

View changed files

📝 comfy_kitchen/backends/hip/CMakeLists.txt (+1 -0)
📝 comfy_kitchen/backends/hip/__init__.py (+288 -0)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+84 -4)
📝 comfy_kitchen/backends/hip/gemm_wmma.h (+5 -3)
➕ comfy_kitchen/backends/hip/launchers.h (+36 -0)
📝 comfy_kitchen/backends/hip/ops/convrot_w4a4.hip (+3 -3)
📝 comfy_kitchen/backends/hip/ops/gemm_fp8.hip (+3 -3)
📝 comfy_kitchen/backends/hip/ops/gemm_int8.hip (+16 -11)
➕ comfy_kitchen/backends/hip/ops/w4a8_dequant.hip (+182 -0)
📝 tests/test_hip_wmma.py (+221 -0)

📄 Description

Ports the AsymW4A8Int8 path from #90 to the HIP backend, which registered none of it and so fell through to eager on ROCm. Rebased on #96; that follow-up needed no HIP work (see below).

What this adds

ops/w4a8_dequant.hip: grouped int4 -> int8 weight decode. One thread reads 8 packed bytes and writes 16 int8 through a uint4 store, fp32 and e4m3 group scales via one templated kernel, optional codebook staged in LDS. Reuses fp8_utils.h for the e4m3 decode.
dequantize_w4a8_int8_weight and w4a8_int8_linear plus constraints. The linear feeds the decoded weight into the existing WMMA int8_linear with convrot=True, so the activation rotation stays in the fused quantizer. w4a8_int8_linear is in _WMMA_ONLY_OPS, so RDNA2 routes it to triton/eager; the decode is elementwise and stays registered there.
A chunked path, matching what #90 does on CUDA. The GEMM writeback gained an ldc row stride so a chunk can write an N-column slice of the output; the fp8 and int4 GEMMs pass N and are unchanged. The weight is decoded a chunk at a time so it is still cached when the GEMM reads it back, instead of the whole [N, K] int8 weight round-tripping through global memory.

quantize_w4a8_int8_weight is deliberately not registered. CUDA's version exists to use its native ConvRot rotate kernel, which HIP does not have, and that rotate measures at 0.1-0.2% of quantize time, so eager's packing math is what runs either way.

Chunk sizing

The chunk width targets the device's own L2, capped at the CUDA backend's 4096 columns so a shallow K does not widen a chunk pointlessly and floored at 1024 so a deep K does not starve the GEMM's N grid. Chunking is limited to M <= 64, the conservative end of where it won on every shape tried: above that the weight is decoded in one pass, which is what this path did before chunking, so a part whose crossover sits higher loses the extra win rather than regressing. Both knobs are overridable via COMFY_KITCHEN_W4A8_CHUNK_COLS and COMFY_KITCHEN_W4A8_CHUNK_MAX_ROWS.

Measurements (gfx1200, chunked vs single pass, min of 5x100 iters)

shape	M=1	M=8	M=16	M=64	M>64
N=12288 K=3072	1.77x	3.08x	1.43x	1.40x	off
N=3072 K=12288	1.49x	1.65x	1.27x	1.30x	off
N=5120 K=5120	1.80x	1.21x	1.32x	1.26x	off
N=3072 K=3072	1.13x	1.40x	1.02x	1.00x	off

Against the eager fallback this replaces, the whole path is 4.6x to 23x faster.

On #96

The frozen codebook LUT, the row-chunked quantize and the codebook_tensor kwarg all live in the shared eager layer, which the HIP backend inherits because it does not register quantize_w4a8_int8_weight. #96 does change what the packer emits, so there is a test that the HIP decode is bit exact against eager on a row-chunked quantize and on a requantize with a pinned codebook.

Testing

The int4 -> int8 decode is bit exact against eager across group sizes 4/8/16/32/64, both scale dtypes, and codebook on/off. The dequantized weight is bit identical. Linear matches eager to ~4e-3 relative, the activation quantizer rounding difference; eager's own error against fp is 0.02-0.055 on the same tensors.
Chunk boundaries including a short trailing chunk (N=1152 with 384-column chunks) produce output identical to the single pass, on both the GEMV and WMMA paths.
303 passed, 2 skipped on test_hip_wmma.py and test_hip_dispatch.py.

Only ran on gfx1200. The 17-target fat binary builds clean, and the decode kernel compiles to identical ISA on gfx1030, gfx1100, gfx1151 and gfx1201 (31 VGPRs, no scratch spills, same vectorized load/store), but RDNA2/3/3.5 are compile and disassemble only here. The ldc change touches gemm_wmma.h, shared by the fp8, int8 and int4 GEMMs on every family, so a check on a 7900 or 6800 would be welcome.

Required

Core PR:
https://github.com/Comfy-Org/ComfyUI/pull/15308

Test model:
https://huggingface.co/Kijai/MiniMax-H3-experimental/blob/main/minimax_h3_fl2va_pruned_w4a8_mixed.safetensors

_{🔄 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/97 **Author:** [@0xDELUXA](https://github.com/0xDELUXA) **Created:** 8/6/2026 **Status:** ✅ Merged **Merged:** 8/7/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `amd/hip-w4a8-int8` --- ### 📝 Commits (1) - [`e66f734`](https://github.com/Comfy-Org/comfy-kitchen/commit/e66f7342f0aba1a613db6909cabe49f997fe418e) HIP port of #90: grouped W4A8 over the INT8 GEMM ### 📊 Changes **10 files changed** (+839 additions, -24 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/hip/CMakeLists.txt` (+1 -0) 📝 `comfy_kitchen/backends/hip/__init__.py` (+288 -0) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+84 -4) 📝 `comfy_kitchen/backends/hip/gemm_wmma.h` (+5 -3) ➕ `comfy_kitchen/backends/hip/launchers.h` (+36 -0) 📝 `comfy_kitchen/backends/hip/ops/convrot_w4a4.hip` (+3 -3) 📝 `comfy_kitchen/backends/hip/ops/gemm_fp8.hip` (+3 -3) 📝 `comfy_kitchen/backends/hip/ops/gemm_int8.hip` (+16 -11) ➕ `comfy_kitchen/backends/hip/ops/w4a8_dequant.hip` (+182 -0) 📝 `tests/test_hip_wmma.py` (+221 -0) </details> ### 📄 Description Ports the AsymW4A8Int8 path from #90 to the HIP backend, which registered none of it and so fell through to eager on ROCm. Rebased on #96; that follow-up needed no HIP work (see below). ### What this adds - `ops/w4a8_dequant.hip`: grouped int4 -> int8 weight decode. One thread reads 8 packed bytes and writes 16 int8 through a uint4 store, fp32 and e4m3 group scales via one templated kernel, optional codebook staged in LDS. Reuses `fp8_utils.h` for the e4m3 decode. - `dequantize_w4a8_int8_weight` and `w4a8_int8_linear` plus constraints. The linear feeds the decoded weight into the existing WMMA `int8_linear` with `convrot=True`, so the activation rotation stays in the fused quantizer. `w4a8_int8_linear` is in `_WMMA_ONLY_OPS`, so RDNA2 routes it to triton/eager; the decode is elementwise and stays registered there. - A chunked path, matching what #90 does on CUDA. The GEMM writeback gained an `ldc` row stride so a chunk can write an N-column slice of the output; the fp8 and int4 GEMMs pass N and are unchanged. The weight is decoded a chunk at a time so it is still cached when the GEMM reads it back, instead of the whole `[N, K]` int8 weight round-tripping through global memory. `quantize_w4a8_int8_weight` is deliberately not registered. CUDA's version exists to use its native ConvRot rotate kernel, which HIP does not have, and that rotate measures at 0.1-0.2% of quantize time, so eager's packing math is what runs either way. ### Chunk sizing The chunk width targets the device's own L2, capped at the CUDA backend's 4096 columns so a shallow K does not widen a chunk pointlessly and floored at 1024 so a deep K does not starve the GEMM's N grid. Chunking is limited to M <= 64, the conservative end of where it won on every shape tried: above that the weight is decoded in one pass, which is what this path did before chunking, so a part whose crossover sits higher loses the extra win rather than regressing. Both knobs are overridable via `COMFY_KITCHEN_W4A8_CHUNK_COLS` and `COMFY_KITCHEN_W4A8_CHUNK_MAX_ROWS`. ### Measurements (gfx1200, chunked vs single pass, min of 5x100 iters) | shape | M=1 | M=8 | M=16 | M=64 | M>64 | | -------------- | ----- | ----- | ----- | ----- | ---- | | N=12288 K=3072 | 1.77x | 3.08x | 1.43x | 1.40x | off | | N=3072 K=12288 | 1.49x | 1.65x | 1.27x | 1.30x | off | | N=5120 K=5120 | 1.80x | 1.21x | 1.32x | 1.26x | off | | N=3072 K=3072 | 1.13x | 1.40x | 1.02x | 1.00x | off | Against the eager fallback this replaces, the whole path is 4.6x to 23x faster. ### On #96 The frozen codebook LUT, the row-chunked quantize and the `codebook_tensor` kwarg all live in the shared eager layer, which the HIP backend inherits because it does not register `quantize_w4a8_int8_weight`. #96 does change what the packer emits, so there is a test that the HIP decode is bit exact against eager on a row-chunked quantize and on a requantize with a pinned codebook. ### Testing - The int4 -> int8 decode is bit exact against eager across group sizes 4/8/16/32/64, both scale dtypes, and codebook on/off. The dequantized weight is bit identical. Linear matches eager to ~4e-3 relative, the activation quantizer rounding difference; eager's own error against fp is 0.02-0.055 on the same tensors. - Chunk boundaries including a short trailing chunk (N=1152 with 384-column chunks) produce output identical to the single pass, on both the GEMV and WMMA paths. - 303 passed, 2 skipped on `test_hip_wmma.py` and `test_hip_dispatch.py`. Only ran on gfx1200. The 17-target fat binary builds clean, and the decode kernel compiles to identical ISA on gfx1030, gfx1100, gfx1151 and gfx1201 (31 VGPRs, no scratch spills, same vectorized load/store), but RDNA2/3/3.5 are compile and disassemble only here. The `ldc` change touches `gemm_wmma.h`, shared by the fp8, int8 and int4 GEMMs on every family, so a check on a 7900 or 6800 would be welcome. ### Required Core PR: https://github.com/Comfy-Org/ComfyUI/pull/15308 Test model: https://huggingface.co/Kijai/MiniMax-H3-experimental/blob/main/minimax_h3_fl2va_pruned_w4a8_mixed.safetensors --- <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:48 +00:00