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 #72] [CLOSED] Add HIP backend with WMMA kernels for AMD RDNA4 (gfx12) #125

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/72
Author: @0xDELUXA
Created: 7/11/2026
Status: ❌ Closed

Base: main ← Head: amd/hip-rdna4-wmma

📝 Commits (1)

31e4bfe Add HIP backend with WMMA kernels for AMD RDNA4 (gfx12)

📊 Changes

24 files changed (+4771 additions, -40 deletions)

View changed files

📝 README.md (+102 -15)
📝 comfy_kitchen/__init__.py (+5 -0)
➕ comfy_kitchen/backends/hip/CMakeLists.txt (+113 -0)
➕ comfy_kitchen/backends/hip/__init__.py (+970 -0)
➕ comfy_kitchen/backends/hip/dlpack_bindings.cpp (+324 -0)
➕ comfy_kitchen/backends/hip/epilogue.h (+70 -0)
➕ comfy_kitchen/backends/hip/fp8_utils.h (+150 -0)
➕ comfy_kitchen/backends/hip/gemm_wmma.h (+238 -0)
➕ comfy_kitchen/backends/hip/hadamard.h (+163 -0)
➕ comfy_kitchen/backends/hip/ops/adaln.hip (+124 -0)
➕ comfy_kitchen/backends/hip/ops/apply_rope.hip (+119 -0)
➕ comfy_kitchen/backends/hip/ops/convrot_w4a4.hip (+121 -0)
➕ comfy_kitchen/backends/hip/ops/gemm_fp8.hip (+115 -0)
➕ comfy_kitchen/backends/hip/ops/gemm_int8.hip (+112 -0)
➕ comfy_kitchen/backends/hip/ops/gemv_awq.hip (+97 -0)
➕ comfy_kitchen/backends/hip/ops/per_tensor_fp8.hip (+120 -0)
➕ comfy_kitchen/backends/hip/ops/quantize_int8.hip (+136 -0)
➕ comfy_kitchen/backends/hip/ops/stochastic_round_fp8.hip (+129 -0)
➕ comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip (+302 -0)
➕ comfy_kitchen/backends/hip/wmma_gfx12.h (+75 -0)

...and 4 more files

📄 Description

Adds comfy_kitchen/backends/hip, a native backend for AMD RDNA4 (gfx12) built on the v_wmma_*_w32_gfx12 intrinsics. It does not link hipBLAS/hipBLASLt, and no linear op calls into them; every GEMM is compiled from source in the backend directory. Weight preparation and unsupported configurations fall back to eager.

Design

fp8, int8 and int4 all hold 8 bytes per lane per WMMA, so one K-step is 16 bytes of a row whether those bytes hold 16 fp8/int8 values or 32 int4 nibbles. The tile kernel (gemm_wmma.h) is byte-addressed and shared by the fp8, int8 and int4 GEMMs; the element type appears only in the Mma policy. It is software-pipelined at both the global (tile prefetch) and LDS (fragment) level, and uses grouped block ordering for L2 locality.

The lane/fragment layout for the gfx12 w32 WMMA instructions is documented in wmma_gfx12.h.

Ops

Op	Instruction
fp8 `scaled_mm`	`f32_16x16x16_fp8_fp8`
`int8_linear`	`i32_16x16x16_iu8`
`convrot_w4a4_linear`	`i32_16x16x32_iu4`, rotation fused into the quantizer; `linear_dtype="int8"` unpacks the weight onto the iu8 kernel
`scaled_mm_svdquant_w4a4`	`i32_16x16x32_iu4`, per-64 group scales, fused LoRA-up
`gemv_awq_w4a16`	scalar; bandwidth bound, no operand reuse for a tile

Plus fp8/int8 quantization, adaln, and the four apply_rope variants.

On ROCm, int8_linear fell through to eager, which routes torch._int_mm to hipBLASLt; it is now a WMMA kernel. NVFP4 and MXFP8 remain on eager: RDNA4 has neither fp4 WMMA nor microscaling hardware.

Dispatch

The backend registers only when every visible device is known to be gfx12 (the intrinsics have no gfx11 encoding); a device whose architecture cannot be read counts against it, since it cannot be shown to be gfx12. Registration is per-process while kernels launch on the tensor's own device, so one non-gfx12 GPU makes it decline and dispatch falls through to triton/eager. COMFY_KITCHEN_DISABLE_HIP=1 removes it from dispatch.

ROCm's fp8 scaled_mm is competitive on RDNA4 and is faster than the WMMA kernel on square, deep-K and small-M shapes, while the WMMA kernel is faster at large N. fp8 is routed to WMMA on every call the kernel supports, with no shape gate, to keep the backend free of BLAS calls. Calls outside its domain (swizzled operands, scaling other than tensor-wise, a K that is not a multiple of 16, a bias that is not 1-D of length N, operands not on the input's device) fall back to torch.

Build

setup.py builds the extension whenever a ROCm toolchain is present, on Linux and Windows alike:

pip install .

ROCm is discovered from ROCM_HOME/ROCM_PATH/HIP_PATH, then from the pip rocm-sdk wheel (via rocm-sdk path --root, which a ROCm PyTorch build already pulls in), then from /opt/rocm. No environment variables, no CC/CXX override and no Visual Studio developer shell are needed: the ROCm clang drives C, C++ and HIP alike, and locates the MSVC toolchain itself. The HIP extension pins the Ninja generator on every platform; CMake's default generator on Windows is Visual Studio, which cannot build the HIP language at all.

Target architectures default to the gfx12 GPUs the build machine can see, and to gfx1200;gfx1201 when it can see none (CI, cross-builds). To pick them explicitly:

COMFY_HIP_ARCHS=gfx1201 pip install .     # PYTORCH_ROCM_ARCH and GPU_ARCHS also honoured

Because the kernels are gfx12-only, the extension is skipped rather than built when the visible AMD GPUs are all of another architecture. COMFY_KITCHEN_BUILD_HIP=1 turns that skip (and a missing toolchain) into an error; COMFY_KITCHEN_BUILD_NO_HIP=1 suppresses the backend entirely.

Tests

tests/test_hip_wmma.py (95 tests) skips unless the backend registered, so it is inert on CI without an RDNA4 device. Kernels are validated against eager, and against an fp32 reference where the two legitimately differ (the HIP kernels quantize and reduce in fp32 where eager works in bf16, so they are not bitwise equal to eager, and are closer to the fp32 result). One test traps torch._scaled_mm and torch._int_mm and exercises every quantized path to assert neither is reached.

Tested on gfx1200 (RDNA4), ROCm 7.15 / PyTorch 2.12, Windows.

_{🔄 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/72 **Author:** [@0xDELUXA](https://github.com/0xDELUXA) **Created:** 7/11/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `amd/hip-rdna4-wmma` --- ### 📝 Commits (1) - [`31e4bfe`](https://github.com/Comfy-Org/comfy-kitchen/commit/31e4bfef6d67e0cf141d7d2751baea32ff4f6d7d) Add HIP backend with WMMA kernels for AMD RDNA4 (gfx12) ### 📊 Changes **24 files changed** (+4771 additions, -40 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+102 -15) 📝 `comfy_kitchen/__init__.py` (+5 -0) ➕ `comfy_kitchen/backends/hip/CMakeLists.txt` (+113 -0) ➕ `comfy_kitchen/backends/hip/__init__.py` (+970 -0) ➕ `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+324 -0) ➕ `comfy_kitchen/backends/hip/epilogue.h` (+70 -0) ➕ `comfy_kitchen/backends/hip/fp8_utils.h` (+150 -0) ➕ `comfy_kitchen/backends/hip/gemm_wmma.h` (+238 -0) ➕ `comfy_kitchen/backends/hip/hadamard.h` (+163 -0) ➕ `comfy_kitchen/backends/hip/ops/adaln.hip` (+124 -0) ➕ `comfy_kitchen/backends/hip/ops/apply_rope.hip` (+119 -0) ➕ `comfy_kitchen/backends/hip/ops/convrot_w4a4.hip` (+121 -0) ➕ `comfy_kitchen/backends/hip/ops/gemm_fp8.hip` (+115 -0) ➕ `comfy_kitchen/backends/hip/ops/gemm_int8.hip` (+112 -0) ➕ `comfy_kitchen/backends/hip/ops/gemv_awq.hip` (+97 -0) ➕ `comfy_kitchen/backends/hip/ops/per_tensor_fp8.hip` (+120 -0) ➕ `comfy_kitchen/backends/hip/ops/quantize_int8.hip` (+136 -0) ➕ `comfy_kitchen/backends/hip/ops/stochastic_round_fp8.hip` (+129 -0) ➕ `comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip` (+302 -0) ➕ `comfy_kitchen/backends/hip/wmma_gfx12.h` (+75 -0) _...and 4 more files_ </details> ### 📄 Description Adds `comfy_kitchen/backends/hip`, a native backend for AMD RDNA4 (gfx12) built on the `v_wmma_*_w32_gfx12` intrinsics. It does not link hipBLAS/hipBLASLt, and no linear op calls into them; every GEMM is compiled from source in the backend directory. Weight preparation and unsupported configurations fall back to eager. ### Design fp8, int8 and int4 all hold 8 bytes per lane per WMMA, so one K-step is 16 bytes of a row whether those bytes hold 16 fp8/int8 values or 32 int4 nibbles. The tile kernel (`gemm_wmma.h`) is byte-addressed and shared by the fp8, int8 and int4 GEMMs; the element type appears only in the `Mma` policy. It is software-pipelined at both the global (tile prefetch) and LDS (fragment) level, and uses grouped block ordering for L2 locality. The lane/fragment layout for the gfx12 w32 WMMA instructions is documented in `wmma_gfx12.h`. ### Ops | Op | Instruction | |----|-------------| | fp8 `scaled_mm` | `f32_16x16x16_fp8_fp8` | | `int8_linear` | `i32_16x16x16_iu8` | | `convrot_w4a4_linear` | `i32_16x16x32_iu4`, rotation fused into the quantizer; `linear_dtype="int8"` unpacks the weight onto the iu8 kernel | | `scaled_mm_svdquant_w4a4` | `i32_16x16x32_iu4`, per-64 group scales, fused LoRA-up | | `gemv_awq_w4a16` | scalar; bandwidth bound, no operand reuse for a tile | Plus fp8/int8 quantization, `adaln`, and the four `apply_rope` variants. On ROCm, `int8_linear` fell through to eager, which routes torch._int_mm to hipBLASLt; it is now a WMMA kernel. NVFP4 and MXFP8 remain on eager: RDNA4 has neither fp4 WMMA nor microscaling hardware. ### Dispatch The backend registers only when every visible device is known to be gfx12 (the intrinsics have no gfx11 encoding); a device whose architecture cannot be read counts against it, since it cannot be shown to be gfx12. Registration is per-process while kernels launch on the tensor's own device, so one non-gfx12 GPU makes it decline and dispatch falls through to triton/eager. `COMFY_KITCHEN_DISABLE_HIP=1` removes it from dispatch. ROCm's fp8 `scaled_mm` is competitive on RDNA4 and is faster than the WMMA kernel on square, deep-K and small-M shapes, while the WMMA kernel is faster at large N. fp8 is routed to WMMA on every call the kernel supports, with no shape gate, to keep the backend free of BLAS calls. Calls outside its domain (swizzled operands, scaling other than tensor-wise, a K that is not a multiple of 16, a bias that is not 1-D of length N, operands not on the input's device) fall back to torch. ### Build `setup.py` builds the extension whenever a ROCm toolchain is present, on Linux and Windows alike: ```bash pip install . ``` ROCm is discovered from `ROCM_HOME`/`ROCM_PATH`/`HIP_PATH`, then from the pip `rocm-sdk` wheel (via `rocm-sdk path --root`, which a ROCm PyTorch build already pulls in), then from `/opt/rocm`. No environment variables, no `CC`/`CXX` override and no Visual Studio developer shell are needed: the ROCm clang drives C, C++ and HIP alike, and locates the MSVC toolchain itself. The HIP extension pins the Ninja generator on every platform; CMake's default generator on Windows is Visual Studio, which cannot build the HIP language at all. Target architectures default to the gfx12 GPUs the build machine can see, and to `gfx1200`;`gfx1201` when it can see none (CI, cross-builds). To pick them explicitly: ```bash COMFY_HIP_ARCHS=gfx1201 pip install . # PYTORCH_ROCM_ARCH and GPU_ARCHS also honoured ``` Because the kernels are gfx12-only, the extension is skipped rather than built when the visible AMD GPUs are all of another architecture. `COMFY_KITCHEN_BUILD_HIP=1` turns that skip (and a missing toolchain) into an error; `COMFY_KITCHEN_BUILD_NO_HIP=1` suppresses the backend entirely. ### Tests `tests/test_hip_wmma.py` (95 tests) skips unless the backend registered, so it is inert on CI without an RDNA4 device. Kernels are validated against eager, and against an fp32 reference where the two legitimately differ (the HIP kernels quantize and reduce in fp32 where eager works in bf16, so they are not bitwise equal to eager, and are closer to the fp32 result). One test traps `torch._scaled_mm` and `torch._int_mm` and exercises every quantized path to assert neither is reached. Tested on gfx1200 (RDNA4), ROCm 7.15 / PyTorch 2.12, Windows. --- <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:36 +00:00

closed this issue
added the
pull-request
label

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

Comfy-Org/comfy-kitchen#125

No description provided.

Rows
Columns