[PR #271] sol_attn: caller-supplied block lists (block-sparse attention) #275

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/271
Author: @kijai
Created: 10/22/2026
Status: 🔄 Open

Base: main ← Head: flashvsr_block_sparse

📝 Commits (1)

16e4b27 sol_attn: caller-supplied block lists

📊 Changes

8 files changed (+163 additions, -42 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+14 -3)
📝 comfy_kitchen/backends/cuda/__init__.py (+9 -2)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+31 -11)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu (+36 -20)
📝 comfy_kitchen/backends/eager/sol_attn.py (+18 -6)
📝 comfy_kitchen/backends/hip/__init__.py (+4 -0)
📝 comfy_kitchen/constraints.py (+11 -0)
📝 tests/test_sol_attn.py (+40 -0)

📄 Description

Adds optional kv_indices / kv_counts to sol_attn. When given, they replace routing. Each 64-token query block attends exactly its listed 64-token key blocks through the existing exact stage, starting from an empty softmax state. Sinks, the forced diagonal, the pooled tail and token routing don't apply, and a query block with no ids returns zeros.

This is used with FlashVSR+ (core support to follow) was trained with its own block mask: a local window intersected with a top-k over pooled window scores, chunk-causal. It needs exact attention over exactly those blocks. Without this, the only option is masked SDPA with a per-token mask.

Results: FlashVSR+ Turbo on an RTX 4090, outputs within bf16 noise of the upstream reference:

	block lists	masked SDPA
1024x1920, 33 frames (full DiT)	2.7 s	9.0 s
2048x3840, one 16-frame chunk	3.65 s	28.9 s

_{🔄 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/271 **Author:** [@kijai](https://github.com/kijai) **Created:** 10/22/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `flashvsr_block_sparse` --- ### 📝 Commits (1) - [`16e4b27`](https://github.com/Comfy-Org/comfy-kitchen/commit/16e4b2797bbc82fed1b80c0225b6cf2aaf2abc2f) sol_attn: caller-supplied block lists ### 📊 Changes **8 files changed** (+163 additions, -42 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+14 -3) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+9 -2) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+31 -11) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu` (+36 -20) 📝 `comfy_kitchen/backends/eager/sol_attn.py` (+18 -6) 📝 `comfy_kitchen/backends/hip/__init__.py` (+4 -0) 📝 `comfy_kitchen/constraints.py` (+11 -0) 📝 `tests/test_sol_attn.py` (+40 -0) </details> ### 📄 Description Adds optional `kv_indices` / `kv_counts` to `sol_attn`. When given, they replace routing. Each 64-token query block attends exactly its listed 64-token key blocks through the existing exact stage, starting from an empty softmax state. Sinks, the forced diagonal, the pooled tail and token routing don't apply, and a query block with no ids returns zeros. This is used with FlashVSR+ (core support to follow) was trained with its own block mask: a local window intersected with a top-k over pooled window scores, chunk-causal. It needs exact attention over exactly those blocks. Without this, the only option is masked SDPA with a per-token mask. **Results:** FlashVSR+ Turbo on an RTX 4090, outputs within bf16 noise of the upstream reference: | | block lists | masked SDPA | |---|---|---| | 1024x1920, 33 frames (full DiT) | 2.7 s | 9.0 s | | 2048x3840, one 16-frame chunk | 3.65 s | 28.9 s | --- <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:22:21 +00:00

bsakaguchi referenced this issue

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

[PR #275] Add ck.input_act with Metal swiglu/situ for Apple GPUs (depends on #270), based on main @ f8b7e97 (v0.2.40) #277