[PR #242] CUDA Sol-Attn: deterministic overflow retry with exact-tail preservation #249

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/242
Author: @xmarre
Created: 10/12/2026
Status: 🔄 Open

Base: main ← Head: fix/sol-token-overflow-deterministic-20261012

📝 Commits (1)

b2346b8 Make CUDA Sol token overflow admission deterministic without losing tail mass

📊 Changes

3 files changed (+90 additions, -11 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu (+6 -5)
📝 comfy_kitchen/backends/cuda/sage_attention/sol_attn_token.cu (+58 -6)
📝 tests/test_sol_attn.py (+26 -0)

📄 Description

Summary

Fixes #241: in CUDA Sol-Attn with token_aug > 0, multiple split CTAs used atomicAdd to reserve more exact tokens than the fixed capacity. The first n_tok arrivals depended on scheduling, so sorting the surviving IDs did not make membership deterministic. This caused the real MiniMax-H3 A/A′ low-source replay to diverge despite equal histogram/count hashes.

Numerical change only (3 files): sol_attn.cu, sol_attn_token.cu, tests/test_sol_attn.py. No feature flag, attention-backend gate, host sampler change or change to Flow/VDN/Spectrum defaults.

Algorithm

Preserve the existing parallel pass 1/2, exact-token admission and approximate softmax tail for ordinary groups.
When there are multiple CTA splits, flag only query groups with tokCnt > n_tok. For each flagged group, reset its count and neutralize earlier split tail states.
Re-execute pass 2 only for flagged groups with one CTA per group, traversing key tiles in a deterministic order. Regenerate both capped exact-token IDs and the complementary approximate tail before sorting and merging.
Add a tokOverflow workspace slot, exposed by sol_attn_plan. Existing C++/Python wrapper ABI signatures are otherwise unchanged. This affects both direct and chunked Sol-Attn.

Reproduction and evidence

The source-of-truth H3 T=36445, H=56, token_aug=256 Core trace has 15,960 groups, with 12 groups exceeding the 256-token capacity. Previously (run 01906), A/A′ had matching reservation/histogram hashes but different retained-ID hashes, first native prediction RMS 0.018553374, and low-clean suffix RMS 0.110349394.

With this patch in the owner’s installed GPU stack (run 01907), all 12 real overflowing A/A′ groups now have equal retained-ID, histogram and reservation hashes. A/A′ has 103/103 identical observed native/sampler boundary hashes, a bitwise-identical first actual low prediction and bitwise-identical final clean suffix (RMS and max error 0). B's deliberately changed prefix retains a meaningful different trajectory, so the A/B intervention is not accidentally masked. Frozen archive hashes were independently verified. Full H3 evidence.

GPU tests: RTX PRO 6000 Blackwell, PyTorch 2.10.0+cu130, CUDA 13, actual installed _C.abi3.so: 7/7 targeted numerical tests passed, no skips, including repeated saturated top-bin admission, retained surplus/tail fidelity, dense-reference accuracy, direct/chunked parity, and token-budget/no-tail validation. Invocation used --import-mode=importlib from outside the checkout, to import the installed native wheel rather than uncompiled source. Test report. Linux/Windows wheel build in fork CI passed; additional combined overlay wheel test in run 38201552241 passed.

Dependencies and scope

This PR is independently reviewable. Companion build-system fix #244 addresses source-wheel CMake cache/toolchain reuse; #246 independently addresses a separate native CUDA/PyTorch resource-handle failure on the owner's CUDA13 environment. The owner qualified the full H3 result with all three overlays, so results do not imply this numerical PR alone fixes cross-runtime linkage.

For reviewers: inspect fallback tail reset and one-CTA replay semantics (no double-count/lost mass), overflow-only scope/launch overhead, tokOverflow workspace ABI, and compatibility/performance on other SM architectures. SM120 functional replay is validated, but dedicated SM80 numerical performance, compute-sanitizer/racecheck and non-overflow microbenchmarking have not yet been reported. Do not infer final audiovisual seam continuity from deterministic low-source replay.

Patch structure: one commit on upstream main be9985bf2977133fb5a48aac4a9df8b6760194c8; no merge/release performed.

_{🔄 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/242 **Author:** [@xmarre](https://github.com/xmarre) **Created:** 10/12/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/sol-token-overflow-deterministic-20261012` --- ### 📝 Commits (1) - [`b2346b8`](https://github.com/Comfy-Org/comfy-kitchen/commit/b2346b8ed5b30497b7493508d60c9f740ac16bb8) Make CUDA Sol token overflow admission deterministic without losing tail mass ### 📊 Changes **3 files changed** (+90 additions, -11 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn.cu` (+6 -5) 📝 `comfy_kitchen/backends/cuda/sage_attention/sol_attn_token.cu` (+58 -6) 📝 `tests/test_sol_attn.py` (+26 -0) </details> ### 📄 Description ## Summary Fixes [#241](https://github.com/Comfy-Org/comfy-kitchen/issues/241): in CUDA Sol-Attn with `token_aug > 0`, multiple split CTAs used `atomicAdd` to reserve more exact tokens than the fixed capacity. The first `n_tok` arrivals depended on scheduling, so sorting the surviving IDs did **not** make membership deterministic. This caused the real MiniMax-H3 A/A′ low-source replay to diverge despite equal histogram/count hashes. **Numerical change only** (3 files): `sol_attn.cu`, `sol_attn_token.cu`, `tests/test_sol_attn.py`. No feature flag, attention-backend gate, host sampler change or change to Flow/VDN/Spectrum defaults. ### Algorithm 1. Preserve the existing parallel pass 1/2, exact-token admission and approximate softmax tail for ordinary groups. 2. When there are multiple CTA splits, flag only query groups with `tokCnt > n_tok`. For each flagged group, reset its count and neutralize earlier split tail states. 3. Re-execute pass 2 only for flagged groups with **one CTA per group**, traversing key tiles in a deterministic order. Regenerate **both** capped exact-token IDs and the complementary approximate tail before sorting and merging. 4. Add a `tokOverflow` workspace slot, exposed by `sol_attn_plan`. Existing C++/Python wrapper ABI signatures are otherwise unchanged. This affects both direct and chunked Sol-Attn. ### Reproduction and evidence The source-of-truth H3 `T=36445, H=56, token_aug=256` Core trace has 15,960 groups, with **12 groups exceeding the 256-token capacity**. Previously (run **01906**), A/A′ had matching reservation/histogram hashes but *different retained-ID hashes*, first native prediction RMS **0.018553374**, and low-clean suffix RMS **0.110349394**. With this patch in the owner’s installed GPU stack (run **01907**), **all 12 real overflowing A/A′ groups now have equal retained-ID, histogram and reservation hashes**. A/A′ has **103/103 identical observed native/sampler boundary hashes**, a **bitwise-identical first actual low prediction** and **bitwise-identical final clean suffix** (RMS and max error **0**). B's deliberately changed prefix retains a meaningful different trajectory, so the A/B intervention is not accidentally masked. Frozen archive hashes were independently verified. [Full H3 evidence](https://github.com/Comfy-Org/comfy-kitchen/pull/242#issuecomment-6118741574). **GPU tests:** RTX PRO 6000 Blackwell, PyTorch `2.10.0+cu130`, CUDA 13, actual installed `_C.abi3.so`: **7/7 targeted numerical tests passed**, no skips, including repeated saturated top-bin admission, retained surplus/tail fidelity, dense-reference accuracy, direct/chunked parity, and token-budget/no-tail validation. Invocation used `--import-mode=importlib` from outside the checkout, to import the installed native wheel rather than uncompiled source. [Test report](https://github.com/Comfy-Org/comfy-kitchen/pull/242#issuecomment-6118440878). Linux/Windows wheel build in [fork CI](https://github.com/xmarre/comfy-kitchen/actions/runs/38196943588) passed; additional combined overlay wheel test in [run 38201552241](https://github.com/xmarre/comfy-kitchen/actions/runs/38201552241) passed. ### Dependencies and scope This PR is independently reviewable. Companion build-system fix [#244](https://github.com/Comfy-Org/comfy-kitchen/pull/244) addresses source-wheel CMake cache/toolchain reuse; [#246](https://github.com/Comfy-Org/comfy-kitchen/pull/246) independently addresses a separate native CUDA/PyTorch resource-handle failure on the owner's CUDA13 environment. The owner qualified the full H3 result with **all three** overlays, so results do **not** imply this numerical PR alone fixes cross-runtime linkage. **For reviewers:** inspect fallback tail reset and one-CTA replay semantics (no double-count/lost mass), overflow-only scope/launch overhead, `tokOverflow` workspace ABI, and compatibility/performance on other SM architectures. SM120 functional replay is validated, but dedicated SM80 numerical performance, compute-sanitizer/racecheck and non-overflow microbenchmarking have not yet been reported. **Do not infer final audiovisual seam continuity from deterministic low-source replay.** **Patch structure:** one commit on upstream `main` `be9985bf2977133fb5a48aac4a9df8b6760194c8`; no merge/release performed. --- <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:59 +00:00

bsakaguchi referenced this issue

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

[PR #249] Add an mps backend: weight-only INT8 and W4A8/W6A8 linears for Apple GPUs #254

bsakaguchi referenced this issue

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

[PR #270] Add Metal RoPE and AdaLN kernels for Apple GPUs (mps backend), based on main @ f8b7e97 (v0.2.40) #273