[PR #184] Optimize HIP Sol exact attention with paired query blocks #205

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/184
Author: @qweqweewqe7-create
Created: 9/17/2026
Status: 🔄 Open

Base: main ← Head: optimize-hip-sol-pair128

📝 Commits (1)

53e545b Optimize HIP Sol exact attention with paired query blocks

📊 Changes

1 file changed (+238 additions, -0 deletions)

View changed files

📝 comfy_kitchen/backends/hip/sage_attention/sol_attn_exact.hip (+238 -0)

📄 Description

Summary

Pair adjacent 64-query blocks in the HIP Sol exact-attention consumer so they can share K / V^T / KSB staging when their routed key-block lists overlap.

The two routed block lists are already sorted. The new 256-thread CTA merges those lists with a two-pointer walk while preserving each query block's original traversal order:

cohort 0 handles the first 64Q block
cohort 1 handles the second 64Q block
shared routed blocks are staged once and consumed by both cohorts
blocks unique to one route are consumed only by that cohort
odd final query blocks are handled by leaving the second cohort inactive

Token-augmented calls (n_tok != 0) keep using the existing exact kernel.

Performance

Measured on RX 9070 XT / gfx1201, BF16 MiniMax H3-style Sol exact workloads. Timings are median exact-core latency using ABBA ordering.

Sequence length	Existing	Pair128	Improvement
15,000	16.416 ms	13.076 ms	20.35%
24,576	26.677 ms	21.799 ms	18.29%
38,987	45.075 ms	36.038 ms	20.05%
49,152	56.083 ms	45.437 ms	18.98%
76,596	93.558 ms	74.259 ms	20.63%
110,000	137.367 ms	109.394 ms	20.36%

The gain stays roughly flat across the tested H3 sequence-length range instead of depending on one specific shape.

Correctness

For the BF16 sweep above, Pair128 matched the existing exact kernel bit-for-bit on sampled outputs and kmean_next:

sampled output exact: true
max absolute difference: 0
relative L2 error: 0
kmean_next exact: true

FP16 was also compared directly against the existing exact kernel at ragged / odd-query-block lengths:

T=3,137: bit-exact, max abs 0, rel L2 0
T=15,000: bit-exact, max abs 0, rel L2 0

Tests

Built for both gfx1200 and gfx1201 and ran the Sol attention test suite with Pair128 as the default non-token exact path:

RX 9070 XT / gfx1201: 104 passed, 1 skipped
RX 9060 XT / gfx1200: 104 passed, 1 skipped

The suite covers ragged tails, batch > 1, FP16/BF16, strided inputs, chunked producer behavior, sinks, top-k routing, and token augmentation / fallback behavior.

_{🔄 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/184 **Author:** [@qweqweewqe7-create](https://github.com/qweqweewqe7-create) **Created:** 9/17/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `optimize-hip-sol-pair128` --- ### 📝 Commits (1) - [`53e545b`](https://github.com/Comfy-Org/comfy-kitchen/commit/53e545b87ddac48cf7bb0f3e6b4d8a2c0f616022) Optimize HIP Sol exact attention with paired query blocks ### 📊 Changes **1 file changed** (+238 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/hip/sage_attention/sol_attn_exact.hip` (+238 -0) </details> ### 📄 Description ## Summary Pair adjacent 64-query blocks in the HIP Sol exact-attention consumer so they can share K / V^T / KSB staging when their routed key-block lists overlap. The two routed block lists are already sorted. The new 256-thread CTA merges those lists with a two-pointer walk while preserving each query block's original traversal order: - cohort 0 handles the first 64Q block - cohort 1 handles the second 64Q block - shared routed blocks are staged once and consumed by both cohorts - blocks unique to one route are consumed only by that cohort - odd final query blocks are handled by leaving the second cohort inactive Token-augmented calls (`n_tok != 0`) keep using the existing exact kernel. ## Performance Measured on RX 9070 XT / gfx1201, BF16 MiniMax H3-style Sol exact workloads. Timings are median exact-core latency using ABBA ordering. | Sequence length | Existing | Pair128 | Improvement | |---:|---:|---:|---:| | 15,000 | 16.416 ms | 13.076 ms | 20.35% | | 24,576 | 26.677 ms | 21.799 ms | 18.29% | | 38,987 | 45.075 ms | 36.038 ms | 20.05% | | 49,152 | 56.083 ms | 45.437 ms | 18.98% | | 76,596 | 93.558 ms | 74.259 ms | 20.63% | | 110,000 | 137.367 ms | 109.394 ms | 20.36% | The gain stays roughly flat across the tested H3 sequence-length range instead of depending on one specific shape. ## Correctness For the BF16 sweep above, Pair128 matched the existing exact kernel bit-for-bit on sampled outputs and `kmean_next`: - sampled output exact: true - max absolute difference: 0 - relative L2 error: 0 - `kmean_next` exact: true FP16 was also compared directly against the existing exact kernel at ragged / odd-query-block lengths: - T=3,137: bit-exact, max abs 0, rel L2 0 - T=15,000: bit-exact, max abs 0, rel L2 0 ## Tests Built for both gfx1200 and gfx1201 and ran the Sol attention test suite with Pair128 as the default non-token exact path: - RX 9070 XT / gfx1201: `104 passed, 1 skipped` - RX 9060 XT / gfx1200: `104 passed, 1 skipped` The suite covers ragged tails, batch > 1, FP16/BF16, strided inputs, chunked producer behavior, sinks, top-k routing, and token augmentation / fallback behavior. --- <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:21:33 +00:00

bsakaguchi referenced this issue

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

[PR #205] Fall back to eager when a one-element INT8 scale adds leading dims #219