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 #276] Issue HIP INT8 attention K/V staging loads together #276

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/276
Author: @qweqweewqe7-create
Created: 10/23/2026
Status: 🔄 Open

Base: main ← Head: optimize-hip-int8-attn-staging

📝 Commits (1)

370901d Issue INT8 attention K/V staging loads together

📊 Changes

1 file changed (+19 additions, -7 deletions)

View changed files

📝 comfy_kitchen/backends/hip/sage_attention/int8_attn.hip (+19 -7)

📄 Description

Summary

The INT8 attention body copies each K/V tile with for (idx = tid; idx < N; idx += kThreads) loops. The compiler cannot tell that every trip runs. So each global load sits behind its own exec branch and is waited on before the next one: load -> s_wait_loadcnt 0 -> ds_store, four times per tile.

This change counts trips from zero with a fixed trip count, so all four loads issue before the first store. It also gives each thread one K base pointer per tile with constant row offsets.

D128 interior loop: staging goes from 58 to 16 instructions, VALU from 266 to 256, and VGPR from 178 to 174, with no spills. Outputs are bit-identical.

Performance

RX 9070 XT (gfx1201) and RX 9060 XT (gfx1200), PyTorch 2.16.0a0+rocm10.3.0a20261023, HIP 7.17.26423. Timings are wall-clock medians over windows of calls, averaged over two alternating runs. Measured at the shapes ComfyUI issues for each model, main -> this PR:

body: int8_attention_from_prequantized
full: the call including prequantization

shape	9070 XT body	9070 XT full	9060 XT body	9060 XT full
Anima self, 2 x 16 x 4096 x 128	1.560 -> 1.490 ms (-4.5%)	2.074 -> 1.984 (-4.3%)	2.984 -> 2.899 (-2.8%)	3.968 -> 3.830 (-3.5%)
Anima cross, kv 512	0.250 -> 0.245 (-2.0%)	0.425 -> 0.416 (-2.1%)	0.496 -> 0.485 (-2.2%)	0.885 -> 0.861 (-2.8%)
Krea2, 1 x 48 x 4138 x 128	2.445 -> 2.351 (-3.9%)	3.226 -> 3.135 (-2.8%)	4.590 -> 4.409 (-3.9%)	6.167 -> 5.998 (-2.7%)
MOSS-Image, 1 x 24 x 6191 x 128	2.661 -> 2.576 (-3.2%)	3.314 -> 3.221 (-2.8%)	5.033 -> 4.896 (-2.7%)	5.979 -> 5.825 (-2.6%)
MOSS-Image, 1 x 24 x 6177 x 128	2.727 -> 2.615 (-4.1%)	3.319 -> 3.220 (-3.0%)	4.992 -> 4.857 (-2.7%)	5.975 -> 5.796 (-3.0%)
H3 fp16 D64 packed qkv, B4 x 32 x 1797		1.169 -> 1.160 (-0.8%)		2.334 -> 2.349 (+0.6%)
H3 fp16 D64 packed qkv, B1		0.366 -> 0.358 (-2.0%)		0.685 -> 0.673 (-1.8%)

On both GPUs, output hashes are identical to main at every shape.

Validation

Tests on both GPUs (test_hip_wmma, test_hip_dispatch, test_gemm_bias_contract, test_int8*, test_adaln, test_rms_rope, test_int8_attention, test_sol_attn, test_convrot_unpacked_quantization, test_w4a8_mma, test_w6a8) show the same failures as unmodified main: 18 triton rope cases and test_int8_linear_correctness[hip].

🤖 Generated with Claude Code

_{🔄 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/276 **Author:** [@qweqweewqe7-create](https://github.com/qweqweewqe7-create) **Created:** 10/23/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `optimize-hip-int8-attn-staging` --- ### 📝 Commits (1) - [`370901d`](https://github.com/Comfy-Org/comfy-kitchen/commit/370901dde3bc35f108f705aa81cf988acec20ca1) Issue INT8 attention K/V staging loads together ### 📊 Changes **1 file changed** (+19 additions, -7 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/hip/sage_attention/int8_attn.hip` (+19 -7) </details> ### 📄 Description ## Summary The INT8 attention body copies each K/V tile with `for (idx = tid; idx < N; idx += kThreads)` loops. The compiler cannot tell that every trip runs. So each global load sits behind its own exec branch and is waited on before the next one: load -> `s_wait_loadcnt 0` -> `ds_store`, four times per tile. This change counts trips from zero with a fixed trip count, so all four loads issue before the first store. It also gives each thread one K base pointer per tile with constant row offsets. D128 interior loop: staging goes from 58 to 16 instructions, VALU from 266 to 256, and VGPR from 178 to 174, with no spills. Outputs are bit-identical. ## Performance RX 9070 XT (gfx1201) and RX 9060 XT (gfx1200), PyTorch `2.16.0a0+rocm10.3.0a20261023`, HIP 7.17.26423. Timings are wall-clock medians over windows of calls, averaged over two alternating runs. Measured at the shapes ComfyUI issues for each model, `main` -> this PR: - **body**: `int8_attention_from_prequantized` - **full**: the call including prequantization | shape | 9070 XT body | 9070 XT full | 9060 XT body | 9060 XT full | | --- | ---: | ---: | ---: | ---: | | Anima self, 2 x 16 x 4096 x 128 | 1.560 -> 1.490 ms (-4.5%) | 2.074 -> 1.984 (-4.3%) | 2.984 -> 2.899 (-2.8%) | 3.968 -> 3.830 (-3.5%) | | Anima cross, kv 512 | 0.250 -> 0.245 (-2.0%) | 0.425 -> 0.416 (-2.1%) | 0.496 -> 0.485 (-2.2%) | 0.885 -> 0.861 (-2.8%) | | Krea2, 1 x 48 x 4138 x 128 | 2.445 -> 2.351 (-3.9%) | 3.226 -> 3.135 (-2.8%) | 4.590 -> 4.409 (-3.9%) | 6.167 -> 5.998 (-2.7%) | | MOSS-Image, 1 x 24 x 6191 x 128 | 2.661 -> 2.576 (-3.2%) | 3.314 -> 3.221 (-2.8%) | 5.033 -> 4.896 (-2.7%) | 5.979 -> 5.825 (-2.6%) | | MOSS-Image, 1 x 24 x 6177 x 128 | 2.727 -> 2.615 (-4.1%) | 3.319 -> 3.220 (-3.0%) | 4.992 -> 4.857 (-2.7%) | 5.975 -> 5.796 (-3.0%) | | H3 fp16 D64 packed qkv, B4 x 32 x 1797 | | 1.169 -> 1.160 (-0.8%) | | 2.334 -> 2.349 (+0.6%) | | H3 fp16 D64 packed qkv, B1 | | 0.366 -> 0.358 (-2.0%) | | 0.685 -> 0.673 (-1.8%) | On both GPUs, output hashes are identical to `main` at every shape. ## Validation Tests on both GPUs (`test_hip_wmma`, `test_hip_dispatch`, `test_gemm_bias_contract`, `test_int8*`, `test_adaln`, `test_rms_rope`, `test_int8_attention`, `test_sol_attn`, `test_convrot_unpacked_quantization`, `test_w4a8_mma`, `test_w6a8`) show the same failures as unmodified `main`: 18 triton rope cases and `test_int8_linear_correctness[hip]`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

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

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#276

No description provided.

Rows
Columns