[PR #134] [CLOSED] HIP: accelerate AWQ W4A16 with WMMA #168

New issue

Closed

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/134
Author: @tangzzycc
Created: 8/24/2026
Status: ❌ Closed

Base: main ← Head: opt/hip-awq-w4a16-mma

📝 Commits (4)

031236d HIP: accelerate AWQ W4A16 with WMMA
dbe90f4 fix(hip): keep mixed-dtype AWQ inputs on scalar path
2d55d1b fix(hip): refine AWQ W4A16 dispatch
a50e241 fix(hip): refine AWQ W4A16 large-M fallback

📊 Changes

5 files changed (+313 additions, -20 deletions)

View changed files

📝 comfy_kitchen/backends/hip/__init__.py (+54 -8)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+3 -3)
📝 comfy_kitchen/backends/hip/ops/gemv_awq.hip (+125 -3)
📝 tests/test_hip_dispatch.py (+8 -0)
📝 tests/test_hip_wmma.py (+123 -6)

📄 Description

Summary

This PR improves the HIP implementation of gemv_awq_w4a16 with size-based dispatch:

The tiered design follows the existing CUDA backend: GEMV for small M, fused MMA for medium M, and dequantization followed by PyTorch matmul for very large M. The HIP implementation uses the project's existing RDNA WMMA policies.

Keep the existing wave-reduction GEMV path for M <= 8.
Use a fused FP16/BF16 WMMA kernel for medium-M workloads.
Dequantize the weight and use PyTorch matmul for M > 2048.
Keep unsupported layouts and mixed-dtype inputs on the existing scalar path.

The public API, packed weight layout, output layout, and bias behavior remain unchanged.

Implementation

The fused kernel uses a 16 x 128 x 64 tile. Each workgroup stages one activation tile and one dequantized weight tile in LDS. The dequantized weights are shared across 16 activation rows and accumulated with the existing MmaF16 or MmaBf16 policy.

The WMMA path requires group size 64, matching activation/scale/output dtypes, and a 16-byte-aligned activation pointer. Other inputs continue to use the original kernel. For very large M, materializing the dequantized weight is amortized by the tuned PyTorch matrix multiplication path.

Performance

Measured on an AMD Radeon 8060S (gfx1151) with ROCm 7.2.

BF16, N=K=4096, medium-M comparison against the previous HIP kernel:

M	Before	After	Speedup
1	0.119 ms	0.119 ms	1.00x
8	0.757 ms	0.757 ms	1.00x
16	1.493 ms	0.149 ms	10.05x
32	2.972 ms	0.302 ms	9.84x
64	5.947 ms	0.609 ms	9.76x
128	11.943 ms	0.889 ms	13.43x
256	24.117 ms	1.228 ms	19.64x

BF16, N=K=4096, routing comparison:

M	Fused WMMA	Dequant + matmul	Selected path	Selected-path advantage
1024	4.514 ms	5.327 ms	WMMA	1.18x
1536	5.445 ms	5.944 ms	WMMA	1.09x
3072	10.284 ms	8.043 ms	PyTorch matmul	1.28x
4096	13.517 ms	9.853 ms	PyTorch matmul	1.37x

The crossover near the dispatch limit varies with dtype and projection shape, so the fallback is intentionally limited to very large M.

Architecture verification

The extension was successfully compiled for every currently supported gfx11/gfx12 target

Code-object inspection confirmed FP16 and BF16 WMMA instructions on every target. The optimized kernels use 20,736 bytes of LDS and have no private-memory, VGPR, or SGPR spills. Runtime performance was measured on gfx1151; the remaining targets were validated through compilation and code-object inspection.

Validation

CMAKE_BUILD_PARALLEL_LEVEL=16 .venv/bin/python setup.py \
  build_ext --inplace --no-cuda --hip \
  --hip-archs=gfx1100,gfx1101,gfx1102,gfx1103,gfx1150,gfx1151,gfx1152,gfx1153,gfx1200,gfx1201

PYTHONPATH=. .venv/bin/pytest \
  tests/test_hip_wmma.py tests/test_hip_dispatch.py tests/test_dlpack.py -q

PYTHONPATH=. .venv/bin/python \
  build/awq_w4a16_gfx11/benchmark.py --large \
  --warmup 3 --iterations 10 --rounds 3 --n 4096 --k 4096

Test result:

408 passed, 9 skipped

The AWQ tests cover FP16/BF16 execution, small and non-tile-aligned shapes, scalar fallback, misaligned contiguous inputs, bias, and large-M routing.

_{🔄 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/134 **Author:** [@tangzzycc](https://github.com/tangzzycc) **Created:** 8/24/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `opt/hip-awq-w4a16-mma` --- ### 📝 Commits (4) - [`031236d`](https://github.com/Comfy-Org/comfy-kitchen/commit/031236db2b4e988dae5236c9778508b71dba1a2e) HIP: accelerate AWQ W4A16 with WMMA - [`dbe90f4`](https://github.com/Comfy-Org/comfy-kitchen/commit/dbe90f467716942dc19937bffb11042e425cb9ec) fix(hip): keep mixed-dtype AWQ inputs on scalar path - [`2d55d1b`](https://github.com/Comfy-Org/comfy-kitchen/commit/2d55d1b47749571250b686b0545e030ea494a09e) fix(hip): refine AWQ W4A16 dispatch - [`a50e241`](https://github.com/Comfy-Org/comfy-kitchen/commit/a50e24175609410477fcd99f8f8d1099e029b625) fix(hip): refine AWQ W4A16 large-M fallback ### 📊 Changes **5 files changed** (+313 additions, -20 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/hip/__init__.py` (+54 -8) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+3 -3) 📝 `comfy_kitchen/backends/hip/ops/gemv_awq.hip` (+125 -3) 📝 `tests/test_hip_dispatch.py` (+8 -0) 📝 `tests/test_hip_wmma.py` (+123 -6) </details> ### 📄 Description ## Summary This PR improves the HIP implementation of `gemv_awq_w4a16` with size-based dispatch: The tiered design follows the existing CUDA backend: GEMV for small M, fused MMA for medium M, and dequantization followed by PyTorch matmul for very large M. The HIP implementation uses the project's existing RDNA WMMA policies. - Keep the existing wave-reduction GEMV path for `M <= 8`. - Use a fused FP16/BF16 WMMA kernel for medium-M workloads. - Dequantize the weight and use PyTorch matmul for `M > 2048`. - Keep unsupported layouts and mixed-dtype inputs on the existing scalar path. The public API, packed weight layout, output layout, and bias behavior remain unchanged. ## Implementation The fused kernel uses a `16 x 128 x 64` tile. Each workgroup stages one activation tile and one dequantized weight tile in LDS. The dequantized weights are shared across 16 activation rows and accumulated with the existing `MmaF16` or `MmaBf16` policy. The WMMA path requires group size 64, matching activation/scale/output dtypes, and a 16-byte-aligned activation pointer. Other inputs continue to use the original kernel. For very large M, materializing the dequantized weight is amortized by the tuned PyTorch matrix multiplication path. ## Performance Measured on an AMD Radeon 8060S (`gfx1151`) with ROCm 7.2. BF16, `N=K=4096`, medium-M comparison against the previous HIP kernel: | M | Before | After | Speedup | |---:|---:|---:|---:| | 1 | 0.119 ms | 0.119 ms | 1.00x | | 8 | 0.757 ms | 0.757 ms | 1.00x | | 16 | 1.493 ms | 0.149 ms | 10.05x | | 32 | 2.972 ms | 0.302 ms | 9.84x | | 64 | 5.947 ms | 0.609 ms | 9.76x | | 128 | 11.943 ms | 0.889 ms | 13.43x | | 256 | 24.117 ms | 1.228 ms | 19.64x | BF16, `N=K=4096`, routing comparison: | M | Fused WMMA | Dequant + matmul | Selected path | Selected-path advantage | |---:|---:|---:|---|---:| | 1024 | 4.514 ms | 5.327 ms | WMMA | 1.18x | | 1536 | 5.445 ms | 5.944 ms | WMMA | 1.09x | | 3072 | 10.284 ms | 8.043 ms | PyTorch matmul | 1.28x | | 4096 | 13.517 ms | 9.853 ms | PyTorch matmul | 1.37x | The crossover near the dispatch limit varies with dtype and projection shape, so the fallback is intentionally limited to very large M. ## Architecture verification The extension was successfully compiled for every currently supported gfx11/gfx12 target Code-object inspection confirmed FP16 and BF16 WMMA instructions on every target. The optimized kernels use 20,736 bytes of LDS and have no private-memory, VGPR, or SGPR spills. Runtime performance was measured on gfx1151; the remaining targets were validated through compilation and code-object inspection. ## Validation ```bash CMAKE_BUILD_PARALLEL_LEVEL=16 .venv/bin/python setup.py \ build_ext --inplace --no-cuda --hip \ --hip-archs=gfx1100,gfx1101,gfx1102,gfx1103,gfx1150,gfx1151,gfx1152,gfx1153,gfx1200,gfx1201 PYTHONPATH=. .venv/bin/pytest \ tests/test_hip_wmma.py tests/test_hip_dispatch.py tests/test_dlpack.py -q PYTHONPATH=. .venv/bin/python \ build/awq_w4a16_gfx11/benchmark.py --large \ --warmup 3 --iterations 10 --rounds 3 --n 4096 --k 4096 ``` Test result: ```text 408 passed, 9 skipped ``` The AWQ tests cover FP16/BF16 execution, small and non-tile-aligned shapes, scalar fallback, misaligned contiguous inputs, bias, and large-M routing. --- <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:05 +00:00