Comfy-Org/comfy-kitchen

Fork 0

mirror of https://github.com/Comfy-Org/comfy-kitchen.git synced 2026-11-05 07:47:40 +00:00

[PR #138] [CLOSED] HIP native BF16/INT8 attention and gfx12 INT8 inference stack for AMD ROCm #170

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/138
Author: @tvukovic-amd
Created: 8/25/2026
Status: ❌ Closed

Base: main ← Head: hip-gfx117x-gfx12x-native-attention

📝 Commits (6)

44fa296 Expose Kitchen HIP BF16/INT8 attention flags so AMD ROCm can select native gfx12 kernels without Sage or Flash
ec33fa2 Gate gfx12 long-BF16 SDPA off WMMA and pack INT8 attention into one workspace
38273d9 Enable gfx12 long-seq HIP BF16 attention without an SDPA support gate
6dcf2c2 Restore row-major INT8 projection packing policy.
95a0f93 Remove unused HIP attention launches and a dead INT8 tile_k arm.
96aeebc Keep branch INT8 GEMM tile selection for Navi 4 perf after upstream rebase

📊 Changes

29 files changed (+8149 additions, -122 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+273 -0)
📝 comfy_kitchen/backends/eager/__init__.py (+95 -0)
📝 comfy_kitchen/backends/eager/quantization.py (+282 -0)
➕ comfy_kitchen/backends/eager/residual.py (+18 -0)
📝 comfy_kitchen/backends/hip/CMakeLists.txt (+9 -0)
📝 comfy_kitchen/backends/hip/__init__.py (+1342 -23)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+722 -3)
📝 comfy_kitchen/backends/hip/epilogue.h (+45 -0)
📝 comfy_kitchen/backends/hip/gemm_wmma.h (+723 -23)
📝 comfy_kitchen/backends/hip/hadamard.h (+688 -8)
📝 comfy_kitchen/backends/hip/launchers.h (+95 -0)
📝 comfy_kitchen/backends/hip/mma.h (+5 -4)
➕ comfy_kitchen/backends/hip/ops/attention_bf16.hip (+499 -0)
📝 comfy_kitchen/backends/hip/ops/gemm_int8.hip (+358 -1)
📝 comfy_kitchen/backends/hip/ops/quantize_int8.hip (+157 -19)
➕ comfy_kitchen/backends/hip/ops/rms_convrot_quant.hip (+43 -0)
➕ comfy_kitchen/backends/hip/ops/rms_gated_residual.hip (+183 -0)
📝 comfy_kitchen/backends/hip/ops/rms_rope.hip (+215 -6)
📝 comfy_kitchen/backends/hip/sage_attention/int8_attn.hip (+413 -0)
📝 comfy_kitchen/backends/hip/sage_attention/quant_qk_int8.hip (+423 -16)

...and 9 more files

📄 Description

Summary

Adds native HIP attention and a gfx12 INT8+BF16 inference stack for AMD ROCm (RDNA4 / non-duplicated WMMA), so diffusion models can use Comfy Kitchen on gfx12 without SageAttention or FlashAttention.

Attention

New attention_bf16.hip and Python exports: hip_attention, hip_attention_is_supported
gfx12 INT8 attention: Q/K/V quant kernel updates, workspace packing, hip_int8_attention / hip_int8_attention_is_supported
Short or overhead-bound INT8 calls stay on HIP BF16 inside Kitchen, not PyTorch SDPA

INT8 linear / GEMM (Navi 4 perf)

gfx12-tuned launch_int8() with shape-specific paths (dual-M for Z-Image projection shapes, cooperative staging, BKB=64 defaults)
Rebased on upstream #122; INT8 does not use upstream launch_gemm_wmma() — that helper remains for fp8 only
gemm_wmma.h merges upstream #122 infra (device_wgp_count, fp8 launcher) with branch epilogues and dual-M kernels

Supporting inference ops

ConvRot / modulated quant, gated residual, SwiGLU split, WMMA weight packing (tensor/int8.py), sage Q/K/V quant updates, eager fallbacks

Tests

tests/test_int8_attention.py, updates to test_hip_wmma.py, test_int8.py, test_hip_dispatch.py

The diff is large because HIP kernels and the full gfx12 dispatch surface are verbose by nature; this is the perf-validated stack, not unrelated refactors.

Pair PR (ComfyUI): Adds --use-kitchen-bf16-attention and --use-kitchen-int8-attention routing to these APIs.

Target hardware

Primary: AMD ROCm gfx12xx (non-duplicated WMMA), e.g. RDNA4 / Navi 4
gfx11 (Navi 3): portable INT8 GEMM geometry unchanged; native attention APIs require non-duplicated WMMA (gfx12-class)

Test plan

Build HIP backend / wheels for gfx12 target(s)
pytest tests/test_int8_attention.py (gfx12 markers) on ROCm
pytest tests/test_hip_wmma.py -k gfx12 — dual-M, modulated QKV exactness
End-to-end with ComfyUI flags --use-kitchen-bf16-attention / --use-kitchen-int8-attention

Notes for reviewers

INT8 launch_int8() keeps explicit gfx12 tile selection (including dual-M for 512 × 3840 × {3840,11520,10240}) rather than delegating to upstream launch_gemm_wmma().
Rebase: Branch is on current main (includes #122); conflict resolution preserved Navi 4 INT8 paths while keeping fp8 on shared upstream launcher.
Scope: Full inference path for measured Z-Image-style workloads; splitting attention-only vs INT8-linear is possible if preferred, but this PR reflects the integrated perf-tested stack.

_{🔄 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/138 **Author:** [@tvukovic-amd](https://github.com/tvukovic-amd) **Created:** 8/25/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `hip-gfx117x-gfx12x-native-attention` --- ### 📝 Commits (6) - [`44fa296`](https://github.com/Comfy-Org/comfy-kitchen/commit/44fa296723122d52a454e09c2a201bbb94a5efcc) Expose Kitchen HIP BF16/INT8 attention flags so AMD ROCm can select native gfx12 kernels without Sage or Flash - [`ec33fa2`](https://github.com/Comfy-Org/comfy-kitchen/commit/ec33fa2fe7d01c7ed5cf7535de90f9de7e8b220e) Gate gfx12 long-BF16 SDPA off WMMA and pack INT8 attention into one workspace - [`38273d9`](https://github.com/Comfy-Org/comfy-kitchen/commit/38273d97b843e42a9ad06f0bb590b2e2a7ed21d1) Enable gfx12 long-seq HIP BF16 attention without an SDPA support gate - [`6dcf2c2`](https://github.com/Comfy-Org/comfy-kitchen/commit/6dcf2c25d936fa177ec593a20a9eb4317c16e9c7) Restore row-major INT8 projection packing policy. - [`95a0f93`](https://github.com/Comfy-Org/comfy-kitchen/commit/95a0f9366bf164653ce2b22d2fdc20dc8c87fb7b) Remove unused HIP attention launches and a dead INT8 tile_k arm. - [`96aeebc`](https://github.com/Comfy-Org/comfy-kitchen/commit/96aeebcc6ade5e23c758f0b8a3302dad94562c8d) Keep branch INT8 GEMM tile selection for Navi 4 perf after upstream rebase ### 📊 Changes **29 files changed** (+8149 additions, -122 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+273 -0) 📝 `comfy_kitchen/backends/eager/__init__.py` (+95 -0) 📝 `comfy_kitchen/backends/eager/quantization.py` (+282 -0) ➕ `comfy_kitchen/backends/eager/residual.py` (+18 -0) 📝 `comfy_kitchen/backends/hip/CMakeLists.txt` (+9 -0) 📝 `comfy_kitchen/backends/hip/__init__.py` (+1342 -23) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+722 -3) 📝 `comfy_kitchen/backends/hip/epilogue.h` (+45 -0) 📝 `comfy_kitchen/backends/hip/gemm_wmma.h` (+723 -23) 📝 `comfy_kitchen/backends/hip/hadamard.h` (+688 -8) 📝 `comfy_kitchen/backends/hip/launchers.h` (+95 -0) 📝 `comfy_kitchen/backends/hip/mma.h` (+5 -4) ➕ `comfy_kitchen/backends/hip/ops/attention_bf16.hip` (+499 -0) 📝 `comfy_kitchen/backends/hip/ops/gemm_int8.hip` (+358 -1) 📝 `comfy_kitchen/backends/hip/ops/quantize_int8.hip` (+157 -19) ➕ `comfy_kitchen/backends/hip/ops/rms_convrot_quant.hip` (+43 -0) ➕ `comfy_kitchen/backends/hip/ops/rms_gated_residual.hip` (+183 -0) 📝 `comfy_kitchen/backends/hip/ops/rms_rope.hip` (+215 -6) 📝 `comfy_kitchen/backends/hip/sage_attention/int8_attn.hip` (+413 -0) 📝 `comfy_kitchen/backends/hip/sage_attention/quant_qk_int8.hip` (+423 -16) _...and 9 more files_ </details> ### 📄 Description ## Summary Adds native HIP attention and a gfx12 INT8+BF16 inference stack for AMD ROCm (RDNA4 / non-duplicated WMMA), so diffusion models can use Comfy Kitchen on gfx12 without SageAttention or FlashAttention. ### Attention - New `attention_bf16.hip` and Python exports: `hip_attention`, `hip_attention_is_supported` - gfx12 INT8 attention: Q/K/V quant kernel updates, workspace packing, `hip_int8_attention` / `hip_int8_attention_is_supported` - Short or overhead-bound INT8 calls stay on **HIP BF16 inside Kitchen**, not PyTorch SDPA ### INT8 linear / GEMM (Navi 4 perf) - gfx12-tuned `launch_int8()` with shape-specific paths (dual-M for Z-Image projection shapes, cooperative staging, BKB=64 defaults) - Rebased on upstream `#122`; **INT8 does not** use upstream `launch_gemm_wmma()` — that helper remains for **fp8** only - `gemm_wmma.h` merges upstream `#122` infra (`device_wgp_count`, fp8 launcher) with branch epilogues and dual-M kernels ### Supporting inference ops - ConvRot / modulated quant, gated residual, SwiGLU split, WMMA weight packing (`tensor/int8.py`), sage Q/K/V quant updates, eager fallbacks ### Tests - `tests/test_int8_attention.py`, updates to `test_hip_wmma.py`, `test_int8.py`, `test_hip_dispatch.py` The diff is large because HIP kernels and the full gfx12 dispatch surface are verbose by nature; this is the perf-validated stack, not unrelated refactors. **Pair PR (ComfyUI):** Adds `--use-kitchen-bf16-attention` and `--use-kitchen-int8-attention` routing to these APIs. ## Target hardware - **Primary:** AMD ROCm **gfx12xx** (non-duplicated WMMA), e.g. RDNA4 / Navi 4 - **gfx11 (Navi 3):** portable INT8 GEMM geometry unchanged; native attention APIs require non-duplicated WMMA (gfx12-class) ## Test plan - [ ] Build HIP backend / wheels for gfx12 target(s) - [ ] `pytest tests/test_int8_attention.py` (gfx12 markers) on ROCm - [ ] `pytest tests/test_hip_wmma.py -k gfx12` — dual-M, modulated QKV exactness - [ ] End-to-end with ComfyUI flags `--use-kitchen-bf16-attention` / `--use-kitchen-int8-attention` ## Notes for reviewers - **INT8 `launch_int8()`** keeps explicit gfx12 tile selection (including dual-M for `512 × 3840 × {3840,11520,10240}`) rather than delegating to upstream `launch_gemm_wmma()`. - **Rebase:** Branch is on current `main` (includes `#122`); conflict resolution preserved Navi 4 INT8 paths while keeping fp8 on shared upstream launcher. - **Scope:** Full inference path for measured Z-Image-style workloads; splitting attention-only vs INT8-linear is possible if preferred, but this PR reflects the integrated perf-tested stack. --- <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:09 +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#170

No description provided.

Rows
Columns