[PR #156] [MERGED] Sparse attention: Add token routing for improved quality #184

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/156
Author: @kijai
Created: 9/4/2026
Status: ✅ Merged
Merged: 9/5/2026
Merged by: @comfyanonymous

Base: main ← Head: sol_token_aug_main

📝 Commits (2)

0e45b9a Add token routing for improved quality
1128df6 Fixes and adjustments

📊 Changes

15 files changed (+1051 additions, -190 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+25 -1)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+1 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+18 -7)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+30 -30)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu (+83 -26)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn_exact.cu (+69 -110)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn_producer.cu (+4 -3)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn_route.cu (+46 -4)
➕ comfy_kitchen/backends/cuda/sage_attention/sol_attn_token.cu (+482 -0)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn_vtranspose.cu (+6 -4)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_layout.cuh (+157 -0)
📝 comfy_kitchen/backends/eager/sol_attn.py (+5 -1)
📝 comfy_kitchen/backends/hip/__init__.py (+16 -0)
📝 comfy_kitchen/constraints.py (+4 -0)
📝 tests/test_sol_attn.py (+105 -4)

📄 Description

Adds token_aug (0, or a multiple of 64 up to 256): each query block additionally attends up to that many top-scoring tokens outside its routed blocks, with the remaining tail made exact for the block centroid.

For now CUDA only (HIP warns and runs without it), needs someone with AMD to support and test.

On MiniMax-H3 captures it roughly halves the sparse-pass error and DC bias for about a third more attention time.

The DC error manifests as brightness/detail "pulsing" every 5 latent frames (around 0.7s)

sol-attn test, before:

https://github.com/user-attachments/assets/d461b1e5-310e-45fc-858f-8d1459bd323c

after:

https://github.com/user-attachments/assets/852e08e4-e299-4f09-9246-2b74f00b5f9a

_{🔄 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/156 **Author:** [@kijai](https://github.com/kijai) **Created:** 9/4/2026 **Status:** ✅ Merged **Merged:** 9/5/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `sol_token_aug_main` --- ### 📝 Commits (2) - [`0e45b9a`](https://github.com/Comfy-Org/comfy-kitchen/commit/0e45b9a7c06461e64c66a7460573fa45d945a57f) Add token routing for improved quality - [`1128df6`](https://github.com/Comfy-Org/comfy-kitchen/commit/1128df65f7695f84b947f3b79c59e88b125681b4) Fixes and adjustments ### 📊 Changes **15 files changed** (+1051 additions, -190 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+25 -1) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+1 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+18 -7) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+30 -30) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu` (+83 -26) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn_exact.cu` (+69 -110) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn_producer.cu` (+4 -3) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn_route.cu` (+46 -4) ➕ `comfy_kitchen/backends/cuda/sage_attention/sol_attn_token.cu` (+482 -0) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn_vtranspose.cu` (+6 -4) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_layout.cuh` (+157 -0) 📝 `comfy_kitchen/backends/eager/sol_attn.py` (+5 -1) 📝 `comfy_kitchen/backends/hip/__init__.py` (+16 -0) 📝 `comfy_kitchen/constraints.py` (+4 -0) 📝 `tests/test_sol_attn.py` (+105 -4) </details> ### 📄 Description Adds token_aug (0, or a multiple of 64 up to 256): each query block additionally attends up to that many top-scoring tokens outside its routed blocks, with the remaining tail made exact for the block centroid. For now CUDA only (HIP warns and runs without it), needs someone with AMD to support and test. On MiniMax-H3 captures it roughly halves the sparse-pass error and DC bias for about a third more attention time. The DC error manifests as brightness/detail "pulsing" every 5 latent frames (around 0.7s) sol-attn test, before: https://github.com/user-attachments/assets/d461b1e5-310e-45fc-858f-8d1459bd323c after: https://github.com/user-attachments/assets/852e08e4-e299-4f09-9246-2b74f00b5f9a --- <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:20 +00:00