Add key bias to chunked Sol-Attn producers

[xmarre]

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: xmarre/MiniMax-H3-Flow-Aligned-Regenerate#29

Core contract: Comfy-Org/ComfyUI#16239

Sol-H3 provider: xmarre/ComfyUI-Sol-H3#9

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

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

[xmarre]

Closing this draft because its key_bias extension was introduced specifically for the weighted MiniMax-H3 Mixed-Grid stack, and that downstream execution path has now been retired from the production plan. Later matched Flow/Continuum evidence explicitly retired mixed_grid_low_suffix for exact Native Masked continuation; production now uses the conservative target-grid Target Input path.

The branch is preserved on the fork at checkpoint/pre-mixed-grid-retirement-171-20260916. No compiled weighted-Mixed-Grid acceptance work is required for the current production topology.