[PR #193] [MERGED] HIP port of #191: W6A8 decode in the W4A8 kernels #210

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/193
Author: @0xDELUXA
Created: 9/23/2026
Status: ✅ Merged
Merged: 9/23/2026
Merged by: @comfyanonymous

Base: main ← Head: amd/hip-w6a8

📝 Commits (1)

fcbf2fb HIP port of #191: W6A8 decode in the W4A8 kernels

📊 Changes

7 files changed (+323 additions, -103 deletions)

View changed files

📝 comfy_kitchen/backends/hip/__init__.py (+7 -13)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+41 -7)
📝 comfy_kitchen/backends/hip/launchers.h (+5 -4)
📝 comfy_kitchen/backends/hip/ops/w4a8_dequant.hip (+71 -39)
📝 comfy_kitchen/backends/hip/ops/w4a8_gemv.hip (+46 -22)
📝 comfy_kitchen/backends/hip/w4a8_codec.h (+37 -18)
📝 tests/test_hip_wmma.py (+116 -0)

📄 Description

Performance port. #191 left 6-bit weights on HIP routed to the eager decode and _eager.int8_linear, so w6a8_int8 checkpoints worked but ran about 7x slower than W4A8 at prefill sizes and about 200x slower at decode. This templates the shared HIP decode helper in w4a8_codec.h on the bit width, following the CUDA kernels: the grouped dequant, the chunked GEMM loop and the decode GEMV all read the 2-bit high plane and decode (code - 32) * s_rel onto the same int8 grid. The eager route in w4a8_int8_linear is gone.

The bindings infer the width from the packed row, as CUDA's w4a8_bits_from_width does, and reject the same layouts: a row that is neither K/2 nor 3K/4 bytes, 6-bit with K % 32 != 0, a group that is not a multiple of 16, or a codebook. The GEMV keeps its paired 16-byte loads only when the 3K/4-byte row stays 16-byte aligned, and otherwise uses its one-vector loop.

H3 out_proj shape (N=5376, K=7168, bf16) on an RX 9060 XT (gfx1200), W6A8 before / after, with W4A8 for reference: M=1 20.7 ms / 0.099 ms (W4A8 0.105 ms), M=4096 28.1 ms / 3.85 ms (W4A8 3.80 ms). This matches the "same speed as W4A8" result from #191 on CUDA.

The new tests check that the 6-bit decode is bit-exact with eager across group sizes, both scale dtypes and rows that are not 16-byte aligned. They also check that every linear route (GEMV, chunked, fp32 scales) is bit-identical to int8_linear on the eagerly decoded weight. The other tests cover binding rejections and layout dispatch to HIP. Run-verified on gfx1200 only. Compile-verified for all 17 targets in architectures.json; the new kernel instances compile without scratch on gfx1030, gfx1100, gfx1151, gfx1200 and gfx1201.

_{🔄 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/193 **Author:** [@0xDELUXA](https://github.com/0xDELUXA) **Created:** 9/23/2026 **Status:** ✅ Merged **Merged:** 9/23/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `amd/hip-w6a8` --- ### 📝 Commits (1) - [`fcbf2fb`](https://github.com/Comfy-Org/comfy-kitchen/commit/fcbf2fb655e5cf2f31995be9bdda9aab39bad2b8) HIP port of #191: W6A8 decode in the W4A8 kernels ### 📊 Changes **7 files changed** (+323 additions, -103 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/hip/__init__.py` (+7 -13) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+41 -7) 📝 `comfy_kitchen/backends/hip/launchers.h` (+5 -4) 📝 `comfy_kitchen/backends/hip/ops/w4a8_dequant.hip` (+71 -39) 📝 `comfy_kitchen/backends/hip/ops/w4a8_gemv.hip` (+46 -22) 📝 `comfy_kitchen/backends/hip/w4a8_codec.h` (+37 -18) 📝 `tests/test_hip_wmma.py` (+116 -0) </details> ### 📄 Description Performance port. #191 left 6-bit weights on HIP routed to the eager decode and `_eager.int8_linear`, so `w6a8_int8` checkpoints worked but ran about 7x slower than W4A8 at prefill sizes and about 200x slower at decode. This templates the shared HIP decode helper in `w4a8_codec.h` on the bit width, following the CUDA kernels: the grouped dequant, the chunked GEMM loop and the decode GEMV all read the 2-bit high plane and decode `(code - 32) * s_rel` onto the same int8 grid. The eager route in `w4a8_int8_linear` is gone. The bindings infer the width from the packed row, as CUDA's `w4a8_bits_from_width` does, and reject the same layouts: a row that is neither `K/2` nor `3K/4` bytes, 6-bit with `K % 32 != 0`, a group that is not a multiple of 16, or a codebook. The GEMV keeps its paired 16-byte loads only when the `3K/4`-byte row stays 16-byte aligned, and otherwise uses its one-vector loop. H3 `out_proj` shape (N=5376, K=7168, bf16) on an RX 9060 XT (gfx1200), W6A8 before / after, with W4A8 for reference: M=1 20.7 ms / 0.099 ms (W4A8 0.105 ms), M=4096 28.1 ms / 3.85 ms (W4A8 3.80 ms). This matches the "same speed as W4A8" result from #191 on CUDA. The new tests check that the 6-bit decode is bit-exact with eager across group sizes, both scale dtypes and rows that are not 16-byte aligned. They also check that every linear route (GEMV, chunked, fp32 scales) is bit-identical to `int8_linear` on the eagerly decoded weight. The other tests cover binding rejections and layout dispatch to HIP. Run-verified on gfx1200 only. Compile-verified for all 17 targets in `architectures.json`; the new kernel instances compile without scratch on gfx1030, gfx1100, gfx1151, gfx1200 and gfx1201. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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