[PR #227] Tune unmasked SM120 D64 attention with smaller query tiles #236

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/227
Author: @Tokha233
Created: 10/1/2026
Status: 🔄 Open

Base: main ← Head: cuda/sm120-d64-query-tile

📝 Commits (5)

f395165 Tune unmasked SM120 D64 attention with smaller query tiles
07e3a00 Record current-main and combined attention test results
3c4e42c Add isolated-process D64 reproduction and decoder evidence
5ff2a31 Reject benchmark checkout and extension import fallbacks
fee9b94 Preserve SM120 D64 tuning alongside upstream Ampere tiles

📊 Changes

7 files changed (+2310 additions, -1 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/sage_attention/sage_attn_launcher.cu (+28 -1)
➕ docs/benchmarks/bench_sm120_d64_query_tile.py (+132 -0)
➕ docs/benchmarks/sm120-d64-independent-decode.json (+149 -0)
➕ docs/benchmarks/sm120-d64-independent-micro.json (+1326 -0)
➕ docs/benchmarks/sm120-d64-query-tile.json (+568 -0)
➕ docs/benchmarks/sm120-d64-query-tile.md (+76 -0)
➕ tests/test_int8_attention_d64_tile.py (+31 -0)

📄 Description

H3 video VAE attention on SM120 uses a 128-query/32-query-warp tile by default. This uses a 64-query/16-query-warp tile for the measured unmasked D64 workload, preserving the existing 128-row quantization scale layout and output bytes. Dispatch is restricted to SM120, 32 query/key heads, batch 1–8, equal sequence lengths 1024–2048, and positive attention scale. Quantizers, the cached-Q pipeline and public APIs are unchanged.

On RTX 5090 D v2, direct [4,32,1797,64] attention time falls by 4.56% (FP16) / 5.82% (BF16). A separate public prequantized-API benchmark in independent ABBA processes reproduces 4.55% / 4.80%; the smallest batch1/1024 shape is effectively neutral, and the 2049 fallback is unchanged within 0.09%.

Full INT8 VAE decode in independent ABBA processes improves from 9.962410 to 9.866619 seconds (0.9615%) on one 243-frame 768×1344latent. Each process loads only one extension, warms once and performs two formal decodes (four formal observations per arm); floating RGB and RGB8 hashes match across all 12 decodes. Two earlier attention-only integration checks measured 0.56% and 0.35%. These are decoder measurements, not end-to-end generation or online throughput claims.

Validation:

16 direct kernel shape/dtype cases and 18 public prequantized API cases match byte for byte, including 2049-token fallback.
The new split-head oracle tests the old 128-query tile against the new path with interleaved QKV and tail lengths.
The new reproducible independent-process benchmark passes 12 dtype/shape cases across four processes, with exact output signatures.
Combined with the existing nonpositive-scale fix (#224): 527 passed, 594 skipped. Current main plus the new launcher:429 passed, 594 skipped, two nonpositive-scale failures; the unmodified main launcher reproduces those same two failures. This positive-scale dispatch does not address #224.
Ruff and diff checks pass.

The benchmark, all raw timings and scope notes are in docs/benchmarks/sm120-d64-query-tile.md. Simultaneously loading/swapping multiple experimental _C modules caused an unrelated GEMM timing artifact; the independent-process reproduction avoids that harness issue.

_{🔄 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/227 **Author:** [@Tokha233](https://github.com/Tokha233) **Created:** 10/1/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `cuda/sm120-d64-query-tile` --- ### 📝 Commits (5) - [`f395165`](https://github.com/Comfy-Org/comfy-kitchen/commit/f3951658149c9c75774c874135cb7a0e9a45002e) Tune unmasked SM120 D64 attention with smaller query tiles - [`07e3a00`](https://github.com/Comfy-Org/comfy-kitchen/commit/07e3a00a3a15bfdc0a7ab5c9cbb6d8e82bc20724) Record current-main and combined attention test results - [`3c4e42c`](https://github.com/Comfy-Org/comfy-kitchen/commit/3c4e42cefb5de17b22d4491122ee31b85865d011) Add isolated-process D64 reproduction and decoder evidence - [`5ff2a31`](https://github.com/Comfy-Org/comfy-kitchen/commit/5ff2a31f5e5c5bbbaa1deb054d4a1e96f4f626ef) Reject benchmark checkout and extension import fallbacks - [`fee9b94`](https://github.com/Comfy-Org/comfy-kitchen/commit/fee9b942f0b03b1c4e24f2b537e9a730ee47d21d) Preserve SM120 D64 tuning alongside upstream Ampere tiles ### 📊 Changes **7 files changed** (+2310 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/sage_attention/sage_attn_launcher.cu` (+28 -1) ➕ `docs/benchmarks/bench_sm120_d64_query_tile.py` (+132 -0) ➕ `docs/benchmarks/sm120-d64-independent-decode.json` (+149 -0) ➕ `docs/benchmarks/sm120-d64-independent-micro.json` (+1326 -0) ➕ `docs/benchmarks/sm120-d64-query-tile.json` (+568 -0) ➕ `docs/benchmarks/sm120-d64-query-tile.md` (+76 -0) ➕ `tests/test_int8_attention_d64_tile.py` (+31 -0) </details> ### 📄 Description H3 video VAE attention on SM120 uses a 128-query/32-query-warp tile by default. This uses a 64-query/16-query-warp tile for the measured unmasked D64 workload, preserving the existing 128-row quantization scale layout and output bytes. Dispatch is restricted to SM120, 32 query/key heads, batch 1–8, equal sequence lengths 1024–2048, and positive attention scale. Quantizers, the cached-Q pipeline and public APIs are unchanged. On RTX 5090 D v2, direct `[4,32,1797,64]` attention time falls by 4.56% (FP16) / 5.82% (BF16). A separate public prequantized-API benchmark in independent ABBA processes reproduces 4.55% / 4.80%; the smallest batch1/1024 shape is effectively neutral, and the 2049 fallback is unchanged within 0.09%. Full INT8 VAE decode in independent ABBA processes improves from **9.962410 to 9.866619 seconds (0.9615%)** on one 243-frame 768×1344latent. Each process loads only one extension, warms once and performs two formal decodes (four formal observations per arm); floating RGB and RGB8 hashes match across all 12 decodes. Two earlier attention-only integration checks measured 0.56% and 0.35%. These are decoder measurements, not end-to-end generation or online throughput claims. Validation: - 16 direct kernel shape/dtype cases and 18 public prequantized API cases match byte for byte, including 2049-token fallback. - The new split-head oracle tests the old 128-query tile against the new path with interleaved QKV and tail lengths. - The new reproducible independent-process benchmark passes 12 dtype/shape cases across four processes, with exact output signatures. - Combined with the existing nonpositive-scale fix (#224): **527 passed, 594 skipped**. Current main plus the new launcher:429 passed, 594 skipped, two nonpositive-scale failures; the unmodified main launcher reproduces those same two failures. This positive-scale dispatch does not address #224. - Ruff and diff checks pass. The benchmark, all raw timings and scope notes are in `docs/benchmarks/sm120-d64-query-tile.md`. Simultaneously loading/swapping multiple experimental `_C` modules caused an unrelated GEMM timing artifact; the independent-process reproduction avoids that harness issue. --- <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:57 +00:00

[PR #236] Validate sol_attn binding arguments on CUDA like the HIP binding does #245