[PR #117] [MERGED] feat: Add int8 Sol-attention (CORE-391) #156

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/117
Author: @kijai
Created: 8/13/2026
Status: ✅ Merged
Merged: 8/29/2026
Merged by: @comfyanonymous

Base: main ← Head: sol_attn

📝 Commits (10+)

ad9a4a8 initial int8 sol-attn implemenation
75113d3 Merge remote-tracking branch 'upstream/main' into sol_attn
6716dfa Merge remote-tracking branch 'upstream/main' into sol_attn
c27d802 sol_attn: centroid-tail routing (1.38x), u8 P quantization, and key_bias masks
79a7532 Fixes
d7fc603 Merge remote-tracking branch 'upstream/main' into sol_attn
883e710 ruff
c04ef20 Add reuse_qkv_memory to reduce peak VRAM
85c4142 Merge remote-tracking branch 'upstream/main' into sol_attn
b17db5d Optimize routing

📊 Changes

18 files changed (+3337 additions, -29 deletions)

View changed files

📝 README.md (+1 -0)
📝 comfy_kitchen/__init__.py (+62 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+6 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+380 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+232 -0)
📝 comfy_kitchen/backends/cuda/sage_attention/attn_utils.cuh (+9 -29)
📝 comfy_kitchen/backends/cuda/sage_attention/mma.cuh (+20 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu (+250 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/sol_attn_exact.cu (+305 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/sol_attn_preprocess.cu (+262 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/sol_attn_producer.cu (+141 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/sol_attn_route.cu (+286 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/sol_attn_vtranspose.cu (+99 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/sol_layout.cuh (+309 -0)
📝 comfy_kitchen/backends/eager/__init__.py (+12 -0)
➕ comfy_kitchen/backends/eager/sol_attn.py (+275 -0)
📝 comfy_kitchen/constraints.py (+58 -0)
➕ tests/test_sol_attn.py (+630 -0)

📄 Description

Add sol_attn: training-free block-sparse attention (CUDA + eager)

Sol-Attn (arXiv 2607.24027): each 64-token query block attends a routed subset of key blocks exactly; all other blocks contribute one pooled term, so nothing leaves the softmax. Cost becomes O(T² · density) with a quality knob (tau).

MiniMax-H3 at T=80k/56 heads: 23.4 → 13.0 s/it (1.80x) end to end vs dense INT8 attention, at cos ≈ 0.97 vs dense (≥ 0.998 vs its own full-precision reference).

Additions to the original method:

All-INT8 where the original is bf16 (CuTe, SM90/SM100-only): Sage-style smoothed-K INT8 QK and INT8 PV with u8 probabilities, targeting consumer GPUs (sm_80+, tuned on sm_120). Routing pass runs on tensor cores.
Centroid-evaluated tail: the pooled branch is computed once per query block instead of per row, shrinking the routing pass 64x for ~5e-4 cosine.
Selection is either the paper's adaptive tau threshold or fixed per-block top-k (topk_ratio) — the latter reproduces lightx2v - SLA selection exactly (threshold computed through the kernel's own int8 quantization).
key_bias: key-only additive bias, SDPA-style shapes incl. bool (e.g. LTX guide strength / padding).
VSA mode (FastVideo / FastH3-VSA checkpoints): tail=False runs the softmax over the routed blocks only (the VSA/SLA fine stage), block_len marks live rows in zero-padded tiles so edge cubes route and pool on real tokens only, and coarse_gate adds VSA's gated coarse branch (gate · softmax(q̄ k̄ᵀ) v̄ per block) from block means the kernels already produce. With topk_ratio this is the full VSA recipe.
Chunked QKV producer (sol_attn_chunked): consumes the qkv projection in token chunks and emits int8 carriers directly, full bf16 Q/K/V never materializes (~5 GB peak saved at 113k tokens). Global quant stats carry over from the previous step, first call self-measures.
Four CUDA stages + an eager reference as test oracle, registry-dispatched (bf16, head_dim 128). 58 tests covering the invariants that broke during development: B>1, ragged tails, strided inputs, cap/sink/bias interactions, inference-mode, real-model rot_dim, activation-scale bootstrap.

Below ~12k tokens dense is usually faster.

Cost of the PR: +213 KB SASS (+0.96%), ~+1.5–2 MB per CI wheel, zero wall-time change, ~40 s extra CPU per wheel job.

Temporary custom node to test it:

sol_attn_minimax_v5.py

_{🔄 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/117 **Author:** [@kijai](https://github.com/kijai) **Created:** 8/13/2026 **Status:** ✅ Merged **Merged:** 8/29/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `sol_attn` --- ### 📝 Commits (10+) - [`ad9a4a8`](https://github.com/Comfy-Org/comfy-kitchen/commit/ad9a4a80029fe4e0bc787247b6beaaf29c4a564e) initial int8 sol-attn implemenation - [`75113d3`](https://github.com/Comfy-Org/comfy-kitchen/commit/75113d36fdf909280bab89355bb4cae508d45e3b) Merge remote-tracking branch 'upstream/main' into sol_attn - [`6716dfa`](https://github.com/Comfy-Org/comfy-kitchen/commit/6716dfa59cd5227e434c748c799a58cdf835b328) Merge remote-tracking branch 'upstream/main' into sol_attn - [`c27d802`](https://github.com/Comfy-Org/comfy-kitchen/commit/c27d802f1448e2bbce2721cab23ef0549b758ba2) sol_attn: centroid-tail routing (1.38x), u8 P quantization, and key_bias masks - [`79a7532`](https://github.com/Comfy-Org/comfy-kitchen/commit/79a75320aa71235dae671fee6904a18c24b08e90) Fixes - [`d7fc603`](https://github.com/Comfy-Org/comfy-kitchen/commit/d7fc60382a5381c3f174e27dbe810b1b05c38da7) Merge remote-tracking branch 'upstream/main' into sol_attn - [`883e710`](https://github.com/Comfy-Org/comfy-kitchen/commit/883e710f8a779c7766171ed6e2b07c1e8e032836) ruff - [`c04ef20`](https://github.com/Comfy-Org/comfy-kitchen/commit/c04ef206f3b84cedb7da26414547daa0fe789bea) Add reuse_qkv_memory to reduce peak VRAM - [`85c4142`](https://github.com/Comfy-Org/comfy-kitchen/commit/85c41421f710a2e4eb415ee2eb457ba11bf68956) Merge remote-tracking branch 'upstream/main' into sol_attn - [`b17db5d`](https://github.com/Comfy-Org/comfy-kitchen/commit/b17db5d7041d5c6fdccb501952fceaeb189c10a2) Optimize routing ### 📊 Changes **18 files changed** (+3337 additions, -29 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+1 -0) 📝 `comfy_kitchen/__init__.py` (+62 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+6 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+380 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+232 -0) 📝 `comfy_kitchen/backends/cuda/sage_attention/attn_utils.cuh` (+9 -29) 📝 `comfy_kitchen/backends/cuda/sage_attention/mma.cuh` (+20 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu` (+250 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/sol_attn_exact.cu` (+305 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/sol_attn_preprocess.cu` (+262 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/sol_attn_producer.cu` (+141 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/sol_attn_route.cu` (+286 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/sol_attn_vtranspose.cu` (+99 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/sol_layout.cuh` (+309 -0) 📝 `comfy_kitchen/backends/eager/__init__.py` (+12 -0) ➕ `comfy_kitchen/backends/eager/sol_attn.py` (+275 -0) 📝 `comfy_kitchen/constraints.py` (+58 -0) ➕ `tests/test_sol_attn.py` (+630 -0) </details> ### 📄 Description ## Add sol_attn: training-free block-sparse attention (CUDA + eager) Sol-Attn (arXiv 2607.24027): each 64-token query block attends a routed subset of key blocks exactly; all other blocks contribute one pooled term, so nothing leaves the softmax. Cost becomes O(T² · density) with a quality knob (tau). MiniMax-H3 at T=80k/56 heads: 23.4 → 13.0 s/it (1.80x) end to end vs dense INT8 attention, at cos ≈ 0.97 vs dense (≥ 0.998 vs its own full-precision reference). Additions to the original method: - All-INT8 where the original is bf16 (CuTe, SM90/SM100-only): Sage-style smoothed-K INT8 QK and INT8 PV with u8 probabilities, targeting consumer GPUs (sm_80+, tuned on sm_120). Routing pass runs on tensor cores. - Centroid-evaluated tail: the pooled branch is computed once per query block instead of per row, shrinking the routing pass 64x for ~5e-4 cosine. - Selection is either the paper's adaptive tau threshold or fixed per-block top-k (topk_ratio) — the latter reproduces lightx2v - SLA selection exactly (threshold computed through the kernel's own int8 quantization). key_bias: key-only additive bias, SDPA-style shapes incl. bool (e.g. LTX guide strength / padding). - VSA mode (FastVideo / FastH3-VSA checkpoints): tail=False runs the softmax over the routed blocks only (the VSA/SLA fine stage), block_len marks live rows in zero-padded tiles so edge cubes route and pool on real tokens only, and coarse_gate adds VSA's gated coarse branch (gate · softmax(q̄ k̄ᵀ) v̄ per block) from block means the kernels already produce. With topk_ratio this is the full VSA recipe. - Chunked QKV producer (sol_attn_chunked): consumes the qkv projection in token chunks and emits int8 carriers directly, full bf16 Q/K/V never materializes (~5 GB peak saved at 113k tokens). Global quant stats carry over from the previous step, first call self-measures. - Four CUDA stages + an eager reference as test oracle, registry-dispatched (bf16, head_dim 128). 58 tests covering the invariants that broke during development: B>1, ragged tails, strided inputs, cap/sink/bias interactions, inference-mode, real-model rot_dim, activation-scale bootstrap. Below ~12k tokens dense is usually faster. Cost of the PR: +213 KB SASS (+0.96%), ~+1.5–2 MB per CI wheel, zero wall-time change, ~40 s extra CPU per wheel job. Temporary custom node to test it: [sol_attn_minimax_v5.py](https://github.com/user-attachments/files/31576773/sol_attn_minimax_v5.py) --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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