[PR #224] Fix nonpositive scale ordering in CUDA INT8 attention #239

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/224
Author: @Tokha233
Created: 9/30/2026
Status: 🔄 Open

Base: main ← Head: cuda/nonpositive-attention-scale

📝 Commits (3)

96b7062 Fix nonpositive scale ordering in CUDA INT8 attention
556fb21 Merge main and separate attention scale and decode template flags
c664dae Sync nonpositive attention scale fix with Kitchen 0.2.40

📊 Changes

3 files changed (+63 additions, -15 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/sage_attention/qk_int_sv_i8_cuda.cuh (+20 -14)
📝 comfy_kitchen/backends/cuda/sage_attention/sage_attn_launcher.cu (+11 -0)
📝 tests/test_int8_attention.py (+32 -1)

📄 Description

The ordinary unmasked CUDA INT8 attention kernel can return all-zero output for a negative attention scale when the last KV tile is partial. This patch applies nonpositive scaling before the maximum reduction and before padding masks, while retaining the existing positive-scale specialization.

Cause and change

For a negative scale, max(QK) * scale is not max(QK * scale). The raw-score padding sentinel also becomes a large positive logit after multiplication by a negative scale; at zero scale padding becomes a valid zero logit. Select a separate specialization for scale <= 0, pre-scale the logits, and apply the padding mask in the scaled domain. Masked/prepared routes and positive-scale arithmetic remain unchanged.

The existing long/partial-tile test also used an automatically selected BF16 SDPA oracle that returns NaNs for zero/negative scales on the tested build. Use explicit FP32 SDPBackend.MATH and assert its finiteness, without relaxing the existing NRMSE < 0.03 threshold. Add 96 regression combinations covering FP16/BF16, D64/128/256, GQA, single/full/partial KV tiles, both CTA_K sizes and both probability paths; each checks direct and prequantized entry points.

RTX 5090 D v2 validation

CUDA 13.0.88, PyTorch 2.12.0+cu130, baseline 19ea55b, patch 96b7062. Recompiled the affected launcher TU (the only include site of the changed header) with the baseline CMake flags and unchanged baseline objects.

Six existing INT8/RMS/GEMM test files plus the new cases: 999 passed, 598 skipped, 0 failed. Skips are retained backend/capability-specific tests.
Reproducer B=1, H=4, Lq=129, Lkv=8193, D=128, BF16, seed 31:

attention scale	baseline NRMSE vs FP32 math	patched NRMSE
0	0.016976889	0.008582708
-1/sqrt(128)	1.000000000 (all zeros)	0.016013842

48 positive-scale shape/dtype cases and four benchmark shapes retain identical output SHA-256 hashes across AB/BA runs.
Positive dense attention at 14,850 tokens / 56 heads: 11.90499 -> 11.90943 ms; 32,700 tokens / 42 heads: 42.82232 -> 42.82554 ms (pooled median of 16 samples per arm; each sample is 10 CUDA Graph replays). These are noise-level differences, not speedup claims.
Full H3 Larry v4 INT8, 8 steps, 14,850 tokens: video and audio latent hashes remain identical. Two post-warmup calls per arm average 15.73348 -> 15.73964 s; this is a correctness/smoke comparison, not a serving-throughput benchmark.

This isolates the inherited failure reported during #218 validation; it does not depend on that performance PR. Other GPU architectures were not runtime-tested here.

_{🔄 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/224 **Author:** [@Tokha233](https://github.com/Tokha233) **Created:** 9/30/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `cuda/nonpositive-attention-scale` --- ### 📝 Commits (3) - [`96b7062`](https://github.com/Comfy-Org/comfy-kitchen/commit/96b70626b4bdde4c6ad8827d51d1a125650d9abd) Fix nonpositive scale ordering in CUDA INT8 attention - [`556fb21`](https://github.com/Comfy-Org/comfy-kitchen/commit/556fb2134577d40985bb6d3c31df79c8b10de2b8) Merge main and separate attention scale and decode template flags - [`c664dae`](https://github.com/Comfy-Org/comfy-kitchen/commit/c664daefd1bda70cbe2cb7993297181f7e451e73) Sync nonpositive attention scale fix with Kitchen 0.2.40 ### 📊 Changes **3 files changed** (+63 additions, -15 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/sage_attention/qk_int_sv_i8_cuda.cuh` (+20 -14) 📝 `comfy_kitchen/backends/cuda/sage_attention/sage_attn_launcher.cu` (+11 -0) 📝 `tests/test_int8_attention.py` (+32 -1) </details> ### 📄 Description The ordinary unmasked CUDA INT8 attention kernel can return all-zero output for a negative attention scale when the last KV tile is partial. This patch applies nonpositive scaling before the maximum reduction and before padding masks, while retaining the existing positive-scale specialization. ### Cause and change For a negative `scale`, `max(QK) * scale` is not `max(QK * scale)`. The raw-score padding sentinel also becomes a large positive logit after multiplication by a negative scale; at zero scale padding becomes a valid zero logit. Select a separate specialization for `scale <= 0`, pre-scale the logits, and apply the padding mask in the scaled domain. Masked/prepared routes and positive-scale arithmetic remain unchanged. The existing long/partial-tile test also used an automatically selected BF16 SDPA oracle that returns NaNs for zero/negative scales on the tested build. Use explicit FP32 `SDPBackend.MATH` and assert its finiteness, without relaxing the existing NRMSE < 0.03 threshold. Add 96 regression combinations covering FP16/BF16, D64/128/256, GQA, single/full/partial KV tiles, both CTA_K sizes and both probability paths; each checks direct and prequantized entry points. ### RTX 5090 D v2 validation CUDA 13.0.88, PyTorch 2.12.0+cu130, baseline 19ea55b, patch 96b7062. Recompiled the affected launcher TU (the only include site of the changed header) with the baseline CMake flags and unchanged baseline objects. - Six existing INT8/RMS/GEMM test files plus the new cases: **999 passed, 598 skipped, 0 failed**. Skips are retained backend/capability-specific tests. - Reproducer B=1, H=4, Lq=129, Lkv=8193, D=128, BF16, seed 31: | attention scale | baseline NRMSE vs FP32 math | patched NRMSE | | --- | ---: | ---: | | 0 | 0.016976889 | 0.008582708 | | -1/sqrt(128) | 1.000000000 (all zeros) | 0.016013842 | - 48 positive-scale shape/dtype cases and four benchmark shapes retain identical output SHA-256 hashes across AB/BA runs. - Positive dense attention at 14,850 tokens / 56 heads: 11.90499 -> 11.90943 ms; 32,700 tokens / 42 heads: 42.82232 -> 42.82554 ms (pooled median of 16 samples per arm; each sample is 10 CUDA Graph replays). These are noise-level differences, not speedup claims. - Full H3 Larry v4 INT8, 8 steps, 14,850 tokens: video and audio latent hashes remain identical. Two post-warmup calls per arm average 15.73348 -> 15.73964 s; this is a correctness/smoke comparison, not a serving-throughput benchmark. This isolates the inherited failure reported during #218 validation; it does not depend on that performance PR. Other GPU architectures were not runtime-tested here. --- <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:53 +00:00

bsakaguchi referenced this issue

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

[PR #239] [MERGED] Performance improvements for qwen 3.x LLM (Part 1) #246