[PR #100] [MERGED] HIP ports of #99 (fused W4A8 requantize) and #95 (na3d / na2d) #145

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: amd/hip-w4a8-requant-na3d

📝 Commits (2)

f3810f4 HIP port of #99: fused W4A8 requantize with stochastic rounding
b035b98 HIP port of #95: na3d / na2d neighborhood attention

📊 Changes

10 files changed (+1386 additions, -2 deletions)

View changed files

📝 README.md (+2 -2)
📝 comfy_kitchen/backends/hip/CMakeLists.txt (+2 -0)
📝 comfy_kitchen/backends/hip/__init__.py (+206 -0)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+74 -0)
📝 comfy_kitchen/backends/hip/launchers.h (+17 -0)
📝 comfy_kitchen/backends/hip/mma.h (+97 -0)
➕ comfy_kitchen/backends/hip/ops/na3d.hip (+324 -0)
➕ comfy_kitchen/backends/hip/ops/w4a8_requant.hip (+282 -0)
📝 tests/test_hip_dispatch.py (+5 -0)
📝 tests/test_hip_wmma.py (+377 -0)

📄 Description

Catch-up for the two kernels added since #97, one commit each. All numbers on gfx1200, ROCm 10.1, PyTorch 2.14.

HIP port of #99: fused W4A8 requantize

Performance port. stochastic_rounding is backend-agnostic and already reached HIP through the shared eager packer, so nothing was broken. What was missing is the kernel: HIP registered no quantize_w4a8_int8_weight, so every weight prepare and LoRA-merge requantize ran eager torch.

ops/w4a8_requant.hip mirrors ops/w4a8_gemm.cu, with the same PCG draw keyed by the element's global index, so a HIP requantize and a CUDA one produce identical stochastic results. The row max is an LDS tree instead of a shuffle reduction, so it
does not depend on the wave width. Two deliberate differences:

ConvRot rotation stays on eager on purpose. Eager rotates with a batched matmul against the 256x256 Hadamard, which runs on the matrix cores: 0.42 ms at 5120x5120. HIP's butterfly ConvRot is LDS-bound and does rotate plus an int8 quantize in 1.03 ms, so a rotate-only kernel carved out of it would be slower than what it replaces, and its rounding differs from the matmul's (about 8 percent of int8 codes land one level apart), which would cost the eager parity below.
The K bound is read from the device rather than assumed: _C.w4a8_requant_max_k() reports the LDS budget of the card the weight lives on, and the wrapper falls back to the chunked eager packer past it. Asymmetric, uniform, fp32-scale and non-16 group sizes reach that packer unchanged.

bf16 weights, ms. SR is the same call with stochastic_rounding set, which is the LoRA-merge path #99 added:

shape	eager	HIP		eager, SR	HIP, SR
3072 x 3072	27.3	0.98	28x	27.2	1.09	25x
5120 x 5120	77.4	2.52	31x	77.0	2.68	29x
13824 x 5120	211	5.59	38x	210	6.02	35x

The gap is memory traffic, not arithmetic. Eager materialises the weight in fp32 and runs the codebook assign, two ALS iterations and the level assign as separate full-tensor elementwise passes; the kernel keeps a 16-wide group in registers and
reads the rotated weight twice in total.

Not bit-exact with eager, by one ulp. The kernel cannot reproduce the summation order of eager's 16-wide reductions, which is torch's and not a sequential accumulate. That moves an ALS group scale by an ulp, visible only where the scale sat on an e4m3 rounding boundary: over 2.6M nibbles at 2560 x 2048, 4 group scales one e4m3 code apart and 12 nibbles one int4 level apart. The tests bound both, require every mismatch to be an adjacent code, and require a decode no less faithful to the weight than eager's. The CUDA kernel has the same property.

HIP port of #95: na3d / na2d

Coverage port. na3d landed with no HIP entry, so on ROCm it fell through to triton or eager. Triton is not a fallback the backend can lean on: ComfyUI gates it behind --enable-triton-backend and some installs have no triton, which left na3d on the eager SDPA-with-mask path there.

ops/na3d.hip follows the flash-attention-2 shape of ops/na3d.cu: one wave covers 16 queries along W at a fixed (t, h) and sweeps that block's key span 16 keys at a time. Two things differ, both forced by the instruction. RDNA's 16x16x16 WMMA hands each lane one accumulator column over 8 rows while the A operand wants one row of the K-step, so the probability tile round-trips through LDS rather than being repacked in registers the way mma.m16n8k16 allows; and both matmuls contract
A @ B^T, so V is staged transposed.

mma.h gains MmaBf16/MmaF16 for both matrix-core generations, so the kernel addresses fragments through the existing frag_row/acc_row/acc_col abstraction. The gfx12 layout it assumes was checked against a CPU reference with a standalone
probe rather than read off a table: exact match. LDS rows are padded to 9 dwords, so with 32 four-byte banks and gcd(9, 32) = 1 the 16 row-strided lanes of a fragment read reach 16 distinct banks; 8 or 12 dwords collapse them onto 4 or 8 and cost 12
to 20 percent.

fp32 has no 16-bit matrix path and is left to triton/eager, matching the CUDA entry. na2d needs nothing of its own.

bf16, ms. Eager is measured with TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1, which lets its SDPA use Flash via AOTriton on gfx1200 and is worth about 7x; without it the eager column is 277 / 127 / 26.0 ms:

shape	kernel	HIP	triton	eager
(1, 21, 30, 52, 12, 64)	(3, 7, 7)	8.38	13.67 (1.63x)	36.2 (4.3x)
(1, 16, 32, 32, 8, 64)	(5, 5, 5)	2.88	4.52 (1.57x)	16.6 (5.8x)
(1, 1, 64, 64, 8, 64)	(1, 13, 13)	0.41	0.75 (1.84x)	4.80 (11.7x)

The triton column is what a ROCm user gets only with triton installed and enabled; the eager column is what everyone else was getting, and only if they set that flag on architectures that are only experimentally supported by AOTriton.

Tests

83 added, in tests/test_hip_wmma.py and tests/test_hip_dispatch.py (tests/test_na.py's get_capable_backends enumerates only cuda/triton/eager, so it does not reach HIP). na3d is checked against a per-query loop over the NATTEN window in fp32 across the same 15 cases #95 uses, crossed with fp16/bf16 and both scale modes. Requantize covers eager parity across three dtypes, row-block independence, the four layouts that must decline, the LDS-budget fallback, and the draw's reproducibility and unbiasedness.

Full suite against an origin/main worktree, comparing per-test outcome maps rather than counts: 83 added and passing, 0 removed, 0 changed outcomes. The 102 pre-existing failures (triton-on-AMD compiler bugs, bf16 tolerance in triton/eager rope, torch ROCm MX gemm) are on origin/main too.

Architecture coverage

All 17 targets in architectures.json compile; gfx1200 is run-verified.

	RDNA2	RDNA3 / 3.5	RDNA4
W4A8 requantize	registers, elementwise	compile-verified	run-verified
na3d / na2d	withheld, no matrix cores	compile-verified	run-verified

gfx1100 and gfx1151 each emit 40 v_wmma for na3d with no traps and no scratch, so they take the real path and not the stub. na3d joins _WMMA_ONLY_OPS, pinned by an assertion rather than left implicit.

_{🔄 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/100 **Author:** [@0xDELUXA](https://github.com/0xDELUXA) **Created:** 8/7/2026 **Status:** ✅ Merged **Merged:** 8/9/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `amd/hip-w4a8-requant-na3d` --- ### 📝 Commits (2) - [`f3810f4`](https://github.com/Comfy-Org/comfy-kitchen/commit/f3810f4072c86d9a86fdb81c7becf43a1318ee5d) HIP port of #99: fused W4A8 requantize with stochastic rounding - [`b035b98`](https://github.com/Comfy-Org/comfy-kitchen/commit/b035b9824f19b0ec66f8d0035c5788226f1f0395) HIP port of #95: na3d / na2d neighborhood attention ### 📊 Changes **10 files changed** (+1386 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+2 -2) 📝 `comfy_kitchen/backends/hip/CMakeLists.txt` (+2 -0) 📝 `comfy_kitchen/backends/hip/__init__.py` (+206 -0) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+74 -0) 📝 `comfy_kitchen/backends/hip/launchers.h` (+17 -0) 📝 `comfy_kitchen/backends/hip/mma.h` (+97 -0) ➕ `comfy_kitchen/backends/hip/ops/na3d.hip` (+324 -0) ➕ `comfy_kitchen/backends/hip/ops/w4a8_requant.hip` (+282 -0) 📝 `tests/test_hip_dispatch.py` (+5 -0) 📝 `tests/test_hip_wmma.py` (+377 -0) </details> ### 📄 Description Catch-up for the two kernels added since #97, one commit each. All numbers on gfx1200, ROCm 10.1, PyTorch 2.14. ## HIP port of #99: fused W4A8 requantize Performance port. `stochastic_rounding` is backend-agnostic and already reached HIP through the shared eager packer, so nothing was broken. What was missing is the kernel: HIP registered no `quantize_w4a8_int8_weight`, so every weight prepare and LoRA-merge requantize ran eager torch. `ops/w4a8_requant.hip` mirrors `ops/w4a8_gemm.cu`, with the same PCG draw keyed by the element's global index, so a HIP requantize and a CUDA one produce identical stochastic results. The row max is an LDS tree instead of a shuffle reduction, so it does not depend on the wave width. Two deliberate differences: - ConvRot rotation stays on eager on purpose. Eager rotates with a batched matmul against the 256x256 Hadamard, which runs on the matrix cores: 0.42 ms at 5120x5120. HIP's butterfly ConvRot is LDS-bound and does rotate plus an int8 quantize in 1.03 ms, so a rotate-only kernel carved out of it would be slower than what it replaces, and its rounding differs from the matmul's (about 8 percent of int8 codes land one level apart), which would cost the eager parity below. - The K bound is read from the device rather than assumed: `_C.w4a8_requant_max_k()` reports the LDS budget of the card the weight lives on, and the wrapper falls back to the chunked eager packer past it. Asymmetric, uniform, fp32-scale and non-16 group sizes reach that packer unchanged. bf16 weights, ms. SR is the same call with `stochastic_rounding` set, which is the LoRA-merge path #99 added: | shape | eager | HIP | | eager, SR | HIP, SR | | |---|---|---|---|---|---|---| | 3072 x 3072 | 27.3 | 0.98 | 28x | 27.2 | 1.09 | 25x | | 5120 x 5120 | 77.4 | 2.52 | 31x | 77.0 | 2.68 | 29x | | 13824 x 5120 | 211 | 5.59 | 38x | 210 | 6.02 | 35x | The gap is memory traffic, not arithmetic. Eager materialises the weight in fp32 and runs the codebook assign, two ALS iterations and the level assign as separate full-tensor elementwise passes; the kernel keeps a 16-wide group in registers and reads the rotated weight twice in total. Not bit-exact with eager, by one ulp. The kernel cannot reproduce the summation order of eager's 16-wide reductions, which is torch's and not a sequential accumulate. That moves an ALS group scale by an ulp, visible only where the scale sat on an e4m3 rounding boundary: over 2.6M nibbles at 2560 x 2048, 4 group scales one e4m3 code apart and 12 nibbles one int4 level apart. The tests bound both, require every mismatch to be an adjacent code, and require a decode no less faithful to the weight than eager's. The CUDA kernel has the same property. ## HIP port of #95: na3d / na2d Coverage port. na3d landed with no HIP entry, so on ROCm it fell through to triton or eager. Triton is not a fallback the backend can lean on: ComfyUI gates it behind `--enable-triton-backend` and some installs have no triton, which left na3d on the eager SDPA-with-mask path there. `ops/na3d.hip` follows the flash-attention-2 shape of `ops/na3d.cu`: one wave covers 16 queries along W at a fixed (t, h) and sweeps that block's key span 16 keys at a time. Two things differ, both forced by the instruction. RDNA's `16x16x16` WMMA hands each lane one accumulator column over 8 rows while the A operand wants one row of the K-step, so the probability tile round-trips through LDS rather than being repacked in registers the way `mma.m16n8k16` allows; and both matmuls contract `A @ B^T`, so V is staged transposed. `mma.h` gains `MmaBf16`/`MmaF16` for both matrix-core generations, so the kernel addresses fragments through the existing `frag_row`/`acc_row`/`acc_col` abstraction. The gfx12 layout it assumes was checked against a CPU reference with a standalone probe rather than read off a table: exact match. LDS rows are padded to 9 dwords, so with 32 four-byte banks and gcd(9, 32) = 1 the 16 row-strided lanes of a fragment read reach 16 distinct banks; 8 or 12 dwords collapse them onto 4 or 8 and cost 12 to 20 percent. fp32 has no 16-bit matrix path and is left to triton/eager, matching the CUDA entry. na2d needs nothing of its own. bf16, ms. Eager is measured with `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1`, which lets its SDPA use Flash via AOTriton on gfx1200 and is worth about 7x; without it the eager column is 277 / 127 / 26.0 ms: | shape | kernel | HIP | triton | eager | |---|---|---|---|---| | (1, 21, 30, 52, 12, 64) | (3, 7, 7) | 8.38 | 13.67 (1.63x) | 36.2 (4.3x) | | (1, 16, 32, 32, 8, 64) | (5, 5, 5) | 2.88 | 4.52 (1.57x) | 16.6 (5.8x) | | (1, 1, 64, 64, 8, 64) | (1, 13, 13) | 0.41 | 0.75 (1.84x) | 4.80 (11.7x) | The triton column is what a ROCm user gets only with triton installed and enabled; the eager column is what everyone else was getting, and only if they set that flag on architectures that are only experimentally supported by AOTriton. ## Tests 83 added, in `tests/test_hip_wmma.py` and `tests/test_hip_dispatch.py` (`tests/test_na.py`'s `get_capable_backends` enumerates only cuda/triton/eager, so it does not reach HIP). na3d is checked against a per-query loop over the NATTEN window in fp32 across the same 15 cases #95 uses, crossed with fp16/bf16 and both scale modes. Requantize covers eager parity across three dtypes, row-block independence, the four layouts that must decline, the LDS-budget fallback, and the draw's reproducibility and unbiasedness. Full suite against an `origin/main` worktree, comparing per-test outcome maps rather than counts: 83 added and passing, 0 removed, 0 changed outcomes. The 102 pre-existing failures (triton-on-AMD compiler bugs, bf16 tolerance in triton/eager rope, torch ROCm MX gemm) are on `origin/main` too. ## Architecture coverage All 17 targets in `architectures.json` compile; gfx1200 is run-verified. | | RDNA2 | RDNA3 / 3.5 | RDNA4 | |---|---|---|---| | W4A8 requantize | registers, elementwise | compile-verified | run-verified | | na3d / na2d | withheld, no matrix cores | compile-verified | run-verified | gfx1100 and gfx1151 each emit 40 `v_wmma` for na3d with no traps and no scratch, so they take the real path and not the stub. na3d joins `_WMMA_ONLY_OPS`, pinned by an assertion rather than left implicit. --- <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:49 +00:00