[PR #146] Reuse Sage INT8 kernel for Sol exact attention #174

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/146
Author: @Zironic
Created: 8/31/2026
Status: 🔄 Open

Base: main ← Head: reuse-sage-int8-sol-exact

📝 Commits (1)

f94e2b1 Use Sage INT8 kernel for Sol exact attention

📊 Changes

7 files changed (+274 additions, -50 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+1 -0)
📝 comfy_kitchen/backends/cuda/sage_attention/qk_int_sv_i8_cuda.cuh (+96 -28)
📝 comfy_kitchen/backends/cuda/sage_attention/sage_attn_launcher.cu (+1 -1)
➕ comfy_kitchen/backends/cuda/sage_attention/sage_attn_sparse_launcher.cu (+97 -0)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu (+56 -15)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn_preprocess.cu (+12 -6)
📝 tests/test_sol_attn.py (+11 -0)

📄 Description

Summary

reuse Kitchen's pure-INT8 Sage consumer for compatible long-sequence Sol tail=False calls
reuse the existing Q/K/V quantizers and skip Sol's duplicate exact carriers on that path
preserve the current Sol exact fallback for tail approximation, short sequences, key bias, block lengths, nonzero sink starts, and chunked producer calls

Results

RTX 4070 (SM89), BF16 [1, 79567, 56, 128], public sol_attn:

Mode	`main`	This PR	Change
Top-K 15%, no tail	396.83 ms	364.83 ms	-8.1%
All blocks, no tail	2254.17 ms	2038.60 ms	-9.6%

At [1, 4096, 4, 128] with every block routed, relative L2 error versus FP32 dense attention improves from 1.975% to 1.545% (21.8% lower). Incremental Torch peak allocation was byte-identical in each paired benchmark.

Tests

CUDA build for SM89 and SM120
pytest tests/test_sol_attn.py tests/test_int8_attention.py -q: 146 passed, 2 skipped
CPU-only full suite: 520 passed, 1349 skipped; its one fused-QKV exact-equality failure reproduces unchanged on parent dae00a1

Sidenote

Kitchen's pure-INT8 Sage consumer could also be extended to replace the remaining Sol exact path, including resuming from the approximate-tail softmax state. I deliberately left that out of this PR for two reasons:

It would make an already fairly involved kernel change significantly larger.
In our compute-equivalence experiments, spending the same additional attention compute on higher exact density consistently retained more attention mass and produced lower L2 error than spending it on the Sol tail approximation.

For that reason, I don't think adding tail support belongs in this optimization PR, even though the underlying online-softmax representation makes it technically feasible.

_{🔄 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/146 **Author:** [@Zironic](https://github.com/Zironic) **Created:** 8/31/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `reuse-sage-int8-sol-exact` --- ### 📝 Commits (1) - [`f94e2b1`](https://github.com/Comfy-Org/comfy-kitchen/commit/f94e2b1619a66261dce76b4540763a0e6051576a) Use Sage INT8 kernel for Sol exact attention ### 📊 Changes **7 files changed** (+274 additions, -50 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+1 -0) 📝 `comfy_kitchen/backends/cuda/sage_attention/qk_int_sv_i8_cuda.cuh` (+96 -28) 📝 `comfy_kitchen/backends/cuda/sage_attention/sage_attn_launcher.cu` (+1 -1) ➕ `comfy_kitchen/backends/cuda/sage_attention/sage_attn_sparse_launcher.cu` (+97 -0) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu` (+56 -15) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn_preprocess.cu` (+12 -6) 📝 `tests/test_sol_attn.py` (+11 -0) </details> ### 📄 Description ## Summary - reuse Kitchen's pure-INT8 Sage consumer for compatible long-sequence Sol `tail=False` calls - reuse the existing Q/K/V quantizers and skip Sol's duplicate exact carriers on that path - preserve the current Sol exact fallback for tail approximation, short sequences, key bias, block lengths, nonzero sink starts, and chunked producer calls ## Results RTX 4070 (SM89), BF16 `[1, 79567, 56, 128]`, public `sol_attn`: | Mode | `main` | This PR | Change | | --- | ---: | ---: | ---: | | Top-K 15%, no tail | 396.83 ms | 364.83 ms | -8.1% | | All blocks, no tail | 2254.17 ms | 2038.60 ms | -9.6% | At `[1, 4096, 4, 128]` with every block routed, relative L2 error versus FP32 dense attention improves from 1.975% to 1.545% (21.8% lower). Incremental Torch peak allocation was byte-identical in each paired benchmark. ## Tests - CUDA build for SM89 and SM120 - `pytest tests/test_sol_attn.py tests/test_int8_attention.py -q`: 146 passed, 2 skipped - CPU-only full suite: 520 passed, 1349 skipped; its one fused-QKV exact-equality failure reproduces unchanged on parent `dae00a1` ## Sidenote Kitchen's pure-INT8 Sage consumer could also be extended to replace the remaining Sol exact path, including resuming from the approximate-tail softmax state. I deliberately left that out of this PR for two reasons: 1. It would make an already fairly involved kernel change significantly larger. 2. In our compute-equivalence experiments, spending the same additional attention compute on higher exact density consistently retained more attention mass and produced lower L2 error than spending it on the Sol tail approximation. For that reason, I don't think adding tail support belongs in this optimization PR, even though the underlying online-softmax representation makes it technically feasible. --- <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:12 +00:00

bsakaguchi referenced this issue

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

[PR #174] Optimize HIP WMMA INT8 GEMM staging, epilogue, and deep-K tiling #194