[PR #171] [CLOSED] Add key bias to chunked Sol-Attn producers #193

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/171
Author: @xmarre
Created: 9/10/2026
Status: ❌ Closed

Base: main ← Head: feature/mixed-grid-key-bias

📝 Commits (1)

663bfa3 Add key bias to chunked Sol-Attn producers

📊 Changes

12 files changed (+200 additions, -29 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/__init__.py (+11 -1)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+12 -3)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu (+3 -3)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn_producer.cu (+5 -3)
📝 comfy_kitchen/backends/hip/__init__.py (+10 -1)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+5 -3)
📝 comfy_kitchen/backends/hip/launchers.h (+4 -3)
📝 comfy_kitchen/backends/hip/sage_attention/sol_attn.hip (+9 -8)
📝 comfy_kitchen/backends/hip/sage_attention/sol_attn_producer.hip (+7 -4)
📝 tests/test_sol_attn.py (+35 -0)
➕ tests/test_sol_attn_chunked_api.py (+42 -0)
➕ tests/test_sol_attn_weighted_token_aug.py (+57 -0)

📄 Description

Summary

Extend the fused-QKV sol_attn_chunked path with the same per-key logit-bias carrier already supported by direct sol_attn, for both CUDA and HIP.

The chunked API accepts natural-log key_bias, normalizes it to the existing float32 [B,T] contract, converts it once to log2 units, and passes it into K-row quantization. The producer continues to compute pooled K sums and V calibration statistics from the unweighted activations; only the exact routed branch receives the score bias.

This preserves the existing sparse-routing/calibration behavior while allowing callers that stream fused QKV to apply exact per-key weighting without materializing full Q/K/V tensors.

Contract

Existing positional calls are preserved. key_bias is a new optional keyword-only argument after the legacy positional tail.
Biased key blocks must be sink-routed, matching direct sol_attn semantics where key bias is consumed by the exact branch.
CUDA and HIP expose the same public behavior.
Cold/bootstrap and primed producer paths both carry the bias.
kmean / vscale calibration remains based on unweighted activations; the measure changes exact score semantics, not calibration values.
Token augmentation remains confined to blocks that routing left non-exact. Weighted sink blocks cannot enter the token candidate bitmap or pooled tail.

Tests and source audit

Adds shared fused-backend coverage for:

biased chunked output vs. direct sol_attn on the same post-RoPE Q/K/V;
cold/bootstrap and primed chunked paths;
proof that the bias changes the result;
invalid key_bias shape rejection before producer execution;
unchanged no-bias behavior;
an explicit token_aug=64 weighted interaction case that compares cold and primed chunked execution against direct fused Sol-Attn, checks finite output, verifies weighted output differs from unweighted output, and requires kmean / vscale to remain bit-identical to the unweighted calibration;
source-level CUDA/HIP API regression coverage proving that the pre-existing rope_eps, tail, block_len, coarse_gate, and token_aug parameters remain positional while key_bias is optional and keyword-only.

The reviewed CUDA and HIP producer implementations both compute post-RoPE K sums and V statistics before key-bias-bearing K quantization. The bias therefore cannot alter the kmean / vscale calibration values. Both low-level bindings require the producer bias carrier to be packed contiguous float32 with exactly B*T elements before a raw pointer reaches the kernel.

The token-routing audit also closes the weighted exact-prefix interaction mechanically on both backends: route initialization inserts sink blocks directly into the exact list; sink blocks are excluded from the normal routed-candidate calculation, excluded from the token candidate bitmap, and excluded from the pooled tail. Token augmentation therefore only refines unrelated non-sink blocks and cannot consume or duplicate a biased exact sink block.

The current clean head 663bfa3ab5c29e9cff6b2ab68d7f287ea54ae6bb is one implementation commit directly on current upstream main (21003fa97bf3b180393446d729ae630ceb6c2a52) and is mergeable. A pre-test checkpoint is retained at checkpoint/pr-171-before-token-aug-test-20260911.

For this exact head, CLA run 34563733264 passed. Build Wheels run 34563734866 concluded action_required before build execution because fork workflow authorization is required. That status is not a compiler failure and provides no CUDA/HIP numerical evidence.

This draft therefore still requires a compiled supported-GPU run before review. Minimum CUDA acceptance is:

pytest -q tests/test_sol_attn.py tests/test_sol_attn_weighted_token_aug.py -k 'chunked_key_bias'

That covers zero-token-budget weighted cold/primed parity plus the enabled-token-routing interaction. HIP has structural/source parity but remains behaviorally unvalidated unless run on supported HIP hardware.

Architecture companion: https://github.com/xmarre/MiniMax-H3-Flow-Aligned-Regenerate/pull/29

Core contract: https://github.com/Comfy-Org/ComfyUI/pull/16239

Sol-H3 provider: https://github.com/xmarre/ComfyUI-Sol-H3/pull/9

_{🔄 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/171 **Author:** [@xmarre](https://github.com/xmarre) **Created:** 9/10/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `feature/mixed-grid-key-bias` --- ### 📝 Commits (1) - [`663bfa3`](https://github.com/Comfy-Org/comfy-kitchen/commit/663bfa3ab5c29e9cff6b2ab68d7f287ea54ae6bb) Add key bias to chunked Sol-Attn producers ### 📊 Changes **12 files changed** (+200 additions, -29 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/__init__.py` (+11 -1) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+12 -3) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu` (+3 -3) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn_producer.cu` (+5 -3) 📝 `comfy_kitchen/backends/hip/__init__.py` (+10 -1) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+5 -3) 📝 `comfy_kitchen/backends/hip/launchers.h` (+4 -3) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_attn.hip` (+9 -8) 📝 `comfy_kitchen/backends/hip/sage_attention/sol_attn_producer.hip` (+7 -4) 📝 `tests/test_sol_attn.py` (+35 -0) ➕ `tests/test_sol_attn_chunked_api.py` (+42 -0) ➕ `tests/test_sol_attn_weighted_token_aug.py` (+57 -0) </details> ### 📄 Description ## Summary Extend the fused-QKV `sol_attn_chunked` path with the same per-key logit-bias carrier already supported by direct `sol_attn`, for both CUDA and HIP. The chunked API accepts natural-log `key_bias`, normalizes it to the existing float32 `[B,T]` contract, converts it once to log2 units, and passes it into K-row quantization. The producer continues to compute pooled K sums and V calibration statistics from the unweighted activations; only the exact routed branch receives the score bias. This preserves the existing sparse-routing/calibration behavior while allowing callers that stream fused QKV to apply exact per-key weighting without materializing full Q/K/V tensors. ## Contract - Existing positional calls are preserved. `key_bias` is a new optional **keyword-only** argument after the legacy positional tail. - Biased key blocks must be sink-routed, matching direct `sol_attn` semantics where key bias is consumed by the exact branch. - CUDA and HIP expose the same public behavior. - Cold/bootstrap and primed producer paths both carry the bias. - `kmean` / `vscale` calibration remains based on unweighted activations; the measure changes exact score semantics, not calibration values. - Token augmentation remains confined to blocks that routing left non-exact. Weighted sink blocks cannot enter the token candidate bitmap or pooled tail. ## Tests and source audit Adds shared fused-backend coverage for: - biased chunked output vs. direct `sol_attn` on the same post-RoPE Q/K/V; - cold/bootstrap and primed chunked paths; - proof that the bias changes the result; - invalid `key_bias` shape rejection before producer execution; - unchanged no-bias behavior; - an explicit `token_aug=64` weighted interaction case that compares cold and primed chunked execution against direct fused Sol-Attn, checks finite output, verifies weighted output differs from unweighted output, and requires `kmean` / `vscale` to remain bit-identical to the unweighted calibration; - source-level CUDA/HIP API regression coverage proving that the pre-existing `rope_eps`, `tail`, `block_len`, `coarse_gate`, and `token_aug` parameters remain positional while `key_bias` is optional and keyword-only. The reviewed CUDA and HIP producer implementations both compute post-RoPE K sums and V statistics before key-bias-bearing K quantization. The bias therefore cannot alter the `kmean` / `vscale` calibration values. Both low-level bindings require the producer bias carrier to be packed contiguous float32 with exactly `B*T` elements before a raw pointer reaches the kernel. The token-routing audit also closes the weighted exact-prefix interaction mechanically on both backends: route initialization inserts sink blocks directly into the exact list; sink blocks are excluded from the normal routed-candidate calculation, excluded from the token candidate bitmap, and excluded from the pooled tail. Token augmentation therefore only refines unrelated non-sink blocks and cannot consume or duplicate a biased exact sink block. The current clean head `663bfa3ab5c29e9cff6b2ab68d7f287ea54ae6bb` is one implementation commit directly on current upstream `main` (`21003fa97bf3b180393446d729ae630ceb6c2a52`) and is mergeable. A pre-test checkpoint is retained at `checkpoint/pr-171-before-token-aug-test-20260911`. For this exact head, CLA run `34563733264` passed. `Build Wheels` run `34563734866` concluded `action_required` before build execution because fork workflow authorization is required. That status is **not a compiler failure** and provides no CUDA/HIP numerical evidence. This draft therefore still requires a compiled supported-GPU run before review. Minimum CUDA acceptance is: ```bash pytest -q tests/test_sol_attn.py tests/test_sol_attn_weighted_token_aug.py -k 'chunked_key_bias' ``` That covers zero-token-budget weighted cold/primed parity plus the enabled-token-routing interaction. HIP has structural/source parity but remains behaviorally unvalidated unless run on supported HIP hardware. Architecture companion: https://github.com/xmarre/MiniMax-H3-Flow-Aligned-Regenerate/pull/29 Core contract: https://github.com/Comfy-Org/ComfyUI/pull/16239 Sol-H3 provider: https://github.com/xmarre/ComfyUI-Sol-H3/pull/9 --- <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:26 +00:00