feat: Add int8 Sol-attention (CORE-391)

[kijai]

#" 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; attention-only 10.2 ms vs 71 ms at T=37k, 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 kernels, SM90/SM100 only):
  Sage-style smoothed-K INT8 QK and INT8 PV with u8 probabilities, targeting
  consumer GPUs (sm_80+, tuned on sm_120)
• Centroid-evaluated tail: the pooled branch is computed once per query
  block instead of per row (the routing decision already is), shrinking the
  routing pass 64x for ~5e-4 cosine, selectable back to per-row via
  centroid_tail=False
• Additions the original lacks: key_bias (key-only additive bias,
  SDPA-style shapes incl. bool, for example LTX guide strength / padding)

Four CUDA stages + an eager reference as test oracle, registry-dispatched
(bf16, head_dim 128). 52 tests covering the invariants that broke during
development: B>1, ragged tails, strided inputs, cap/sink/bias interactions,
stride agreement with register_fake.

Below ~12k tokens dense is usually faster; sm_89 compiles but is untuned.

Temporary custom node to test it:

sol_attn_minimax_v2.py

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 58148650-789d-4f54-9cbe-376d788f5969

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[hahajaja]

Build & test report on sm86 (RTX 3070 Laptop GPU, 8.6 GB)

Environment (local):

• OS: Windows; GPU: NVIDIA GeForce RTX 3070 Laptop GPU, sm_86, 8.6 GB
• CUDA: 13.0 (nvcc V13.0.88); Python 3.11.9; torch 2.10.0+cu130
• cmake 4.3.1; MSVC 2022 14.40; ninja 1.13
• Cloned the PR head (pull/117/head); initialized submodules third_party/cutlass and third_party/flash-attention with --depth 1.

1. Build failure (full build)

• The full extension did not build. The CUTLASS submodule headers failed to compile under the local CUDA 13.0 / MSVC 2022 toolchain. Errors appeared in third_party/cutlass/include/cute/atom/mma_traits_sm90_gmma.hpp, mma_traits_sm100.hpp, and cute/int_tuple.hpp, e.g.:
  • error C2672: "layout" / "stride" / "get": no matching overloaded function
  • error C2955, error C3203, error C2440
  • final: RuntimeError: CMake build failed for comfy_kitchen.backends.cuda._C
• Sources pulling in CUTLASS/cute: cutlass_gemm_int8.cu, convrot_w4a4.cu, turing_int4.cu, turing_int8.cu (#include <cutlass/...>) and flash_decode.cu (#include "flash_fwd_kernel.h" → cute/tensor.hpp).
• The project's CMakeLists.txt requires CUDA >= 12.8; the local toolchain is CUDA 13.0.

2. How I got a working build to test sol_attn

• Renamed third_party/cutlass/include → include_off so COMFY_HAVE_CUTLASS is not defined. cutlass_gemm_int8.cu then compiles its cuBLAS fallback path (no CUTLASS headers needed).
• Removed the 5 CUTLASS-dependent sources from CUDA_SOURCES in comfy_kitchen/backends/cuda/CMakeLists.txt: cutlass_gemm_int8.cu, turing_int4.cu, turing_int8.cu, convrot_w4a4.cu, flash_decode.cu.
• Added comfy_kitchen/backends/cuda/ops/sol_stubs.cu with empty extern "C" stubs for the symbols those files exported that dlpack_bindings.cpp links against. First attempt missed launch_cutlass_int4_dequant and launch_cutlass_int8_dequant_strided (link errors LNK2001: unresolved external symbol); after adding them the link succeeded.
• Built with COMFY_CUDA_ARCHS=86 (the default Windows arch list is 75-real;80-real;89;120f, which does not include 86). Build succeeded (exit 0); _C.cp311-win_amd64.pyd was produced.
• sol_attn itself does not depend on CUTLASS (it uses inline PTX mma.sync), so these exclusions do not affect the sol_attn kernels.