[PR #168] sol_attn: optional blk_cnt out-parameter for the routed-block count #191

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/168
Author: @fblissjr
Created: 9/8/2026
Status: 🔄 Open

Base: main ← Head: sol-blk-cnt-pr

📝 Commits (3)

06abeed sol_attn: optional blk_cnt out-parameter for the routed-block count
1a3475d sol_attn: the eager blk_cnt check is its own, not the fused backends'
956fc58 sol_attn_chunked: optional keyword-only blk_cnt out-parameter

📊 Changes

6 files changed (+345 additions, -4 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+10 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+30 -2)
📝 comfy_kitchen/backends/eager/sol_attn.py (+30 -2)
📝 comfy_kitchen/backends/hip/__init__.py (+15 -0)
📝 comfy_kitchen/constraints.py (+12 -0)
📝 tests/test_sol_attn.py (+248 -0)

📄 Description

TLDR changes:
Sol's route stage already works out how many key blocks each query block attends exactly and writes it into the scratch workspace that sol_attn allocates for the call. that workspace is freed the moment the call returns and the count is gone before the caller sees any of it - this copies it out first to a consumer (if there is one, since it's opt-in and off by default)**

**Note: It's opt-in and off by default (blk_cnt=None), pass nothing and nothing changes: no allocation, no copy, no sync, and no kernel reads the destination so you end up with the same output either way

so what?

tau's effect is invisible. You can't tell whether tau 1.0 routed 5% or 50% of blocks without reimplementing the route stage in python, sinks and forced diagonal and ties included, and that still wouldn't be the code that ran... figured since I leveraged it for some solid data gathering research, I'd share back in case useful to others

other stuff:

CUDA copies the plan's slice after the launch, eager fills from its own route mask, HIP raises if given one. Same slot is there but i have no AMD card to verify that on
tests: closed form at both tau extremes, monotone in tau, the top-k lower bound, the tie case with no upper bound, batch slices, wrong-buffer rejection.

happy to move it off the main signature if you'd rather not grow the args

_{🔄 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/168 **Author:** [@fblissjr](https://github.com/fblissjr) **Created:** 9/8/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `sol-blk-cnt-pr` --- ### 📝 Commits (3) - [`06abeed`](https://github.com/Comfy-Org/comfy-kitchen/commit/06abeedd85af9d7151332ef9d61902ebede73f1c) sol_attn: optional blk_cnt out-parameter for the routed-block count - [`1a3475d`](https://github.com/Comfy-Org/comfy-kitchen/commit/1a3475d9f2d04d25c54a0694eed35b025e8c545d) sol_attn: the eager blk_cnt check is its own, not the fused backends' - [`956fc58`](https://github.com/Comfy-Org/comfy-kitchen/commit/956fc582873080ff19f4ef763844e200c8d22e24) sol_attn_chunked: optional keyword-only blk_cnt out-parameter ### 📊 Changes **6 files changed** (+345 additions, -4 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+10 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+30 -2) 📝 `comfy_kitchen/backends/eager/sol_attn.py` (+30 -2) 📝 `comfy_kitchen/backends/hip/__init__.py` (+15 -0) 📝 `comfy_kitchen/constraints.py` (+12 -0) 📝 `tests/test_sol_attn.py` (+248 -0) </details> ### 📄 Description TLDR changes: Sol's route stage already works out how many key blocks each query block attends exactly and writes it into the scratch workspace that `sol_attn` allocates for the call. that workspace is freed the moment the call returns and the count is gone before the caller sees any of it - this copies it out first to a consumer (if there is one, since it's opt-in and off by default)** **Note: It's opt-in and off by default (`blk_cnt=None`), pass nothing and nothing changes: no allocation, no copy, no sync, and no kernel reads the destination so you end up with the same output either way so what? - tau's effect is invisible. You can't tell whether tau 1.0 routed 5% or 50% of blocks without reimplementing the route stage in python, sinks and forced diagonal and ties included, and that still wouldn't be the code that ran... figured since I leveraged it for some solid data gathering research, I'd share back in case useful to others other stuff: - CUDA copies the plan's slice after the launch, eager fills from its own route mask, HIP raises if given one. Same slot is there but i have no AMD card to verify that on - tests: closed form at both tau extremes, monotone in tau, the top-k lower bound, the tie case with no upper bound, batch slices, wrong-buffer rejection. happy to move it off the main signature if you'd rather not grow the args --- <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:25 +00:00

bsakaguchi referenced this issue

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

[PR #191] [MERGED] Add w6a8 quantization support #208

bsakaguchi referenced this issue

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

[PR #193] [MERGED] HIP port of #191: W6A8 decode in the W4A8 kernels #210