[PR #218] Pipeline SM120 D128 dense INT8 attention with TMA and cached Q #230

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: cuda/sm120-dense-tma

📝 Commits (2)

e931351 Pipeline SM120 D128 dense INT8 attention with TMA and cached Q
2eb97e0 Validate TMA against math SDPA and fix nonpositive scales

📊 Changes

15 files changed (+1661 additions, -0 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+14 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/dense_tma_sm120.cu (+94 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/dense_tma_sm120.cuh (+363 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/dense_tma_sm120.h (+16 -0)
📝 comfy_kitchen/backends/cuda/sage_attention/sage_attn_launcher.cu (+19 -0)
➕ comfy_kitchen/backends/cuda/sage_attention/tma_pipeline.cuh (+44 -0)
➕ docs/benchmarks/attention-post-review-benchmark.json (+143 -0)
➕ docs/benchmarks/attention-public-benchmark.json (+142 -0)
➕ docs/benchmarks/attention-reference-regression.json (+467 -0)
➕ docs/benchmarks/dense-tma-sm120.md (+91 -0)
➕ docs/benchmarks/full-dit-post-review.json (+39 -0)
➕ docs/benchmarks/full-dit.json (+82 -0)
➕ docs/benchmarks/source-sha256.json (+12 -0)
➕ tests/benchmark_attention_tma.py (+59 -0)
➕ tests/test_int8_attention_tma.py (+76 -0)

📄 Description

This adds a narrow SM120 path for long unmasked D128 INT8 attention: a three-slot TMA K/V ring and producer warpgroup overlap operand movement with cached-Q computation. For positive scales, it reuses main's quantization, QK/PV MMA order, online-softmax helpers and last-tile masking. A separate nonpositive-scale specialization applies scaling before the maximum and padding mask; the independent reference tests uncovered that the initial implementation inherited main's incorrect negative-scale behavior.

Main's cached-Q path is currently enabled on Ada. This is a separate SM120 implementation, selected only for contiguous BHSD, BF16 output, CTA128, Lq/Lkv>=8192 and uint32-safe offsets. Other inputs retain the ordinary launcher. The specialized object is compiled for SM120a; TMA entry points are resolved through the runtime without a new libcuda link dependency.

RTX 5090 D v2 measurements

Baseline 19ea55b (Kitchen 0.2.36), CUDA13.0.88, PyTorch2.12.0+cu130. Eight groups of ten warmed CUDA Graph replays, separate baseline/PR extensions. All default-positive-scale benchmark output tensors match exactly.

B,Hq,Hkv,Lq,Lkv	Main ms	This PR ms	Less operator time
1,8,2,8192,8193	0.602890	0.579014	3.96%
1,8,8,16384,16384	2.144539	2.092317	2.44%
1,42,42,32700,32700	42.892517	42.624614	0.62%
2,4,2,8197,9217	0.712371	0.683634	4.03%

Raw samples, exact A/B output procedure and benchmark script are in docs/benchmarks/dense-tma-sm120.md. The large H3 shape has only a small gain; no broad deployment-speed claim is intended. A three-run public H3 DiT smoke test averaged 15.860485 → 15.839752 s (0.131%, too small for a reliable end-to-end claim), with identical video/audio latent SHA256.

Validation and known baseline issues

Built the full extension and the dedicated CMake SM120a target.
Added independent FP32 math SDPA references for default, zero and negative scale, with repeated GQA heads and query-only chunking for bounded memory. All 9 new accuracy/repeat/stream/graph tests pass; NRMSE is 0.009147–0.016305, below 0.03.
The new reference caught all-zero negative-scale results in both unmodified main and the initial PR (NRMSE 1.0 for the two partial-tile cases). The TMA route now scales before max/masking; negative-scale NRMSE is 0.016199–0.016305. Default-positive arithmetic remains unchanged. Raw before/after data: docs/benchmarks/attention-reference-regression.json.
Complete attention suite after the fix: 428 passed, 594 skipped, 2 failed. The two existing Lq129/Lkv8193 failures remain outside the TMA route and reproduce on unmodified main: optimized SDPA supplies NaN references for zero/negative scale, and the ordinary kernel also has the inherited negative-scale issue. These tests were not skipped or weakened; this PR does not fix the ordinary kernel.
Fresh public-API benchmark after the fix: operator time reductions 3.71%, 2.20%, 0.74%, 4.35% in table order, all default-positive BF16 outputs still exact against main. Separate warmed processes, 8×10 graph replays; raw samples in attention-post-review-benchmark.json. Small timing differences remain subject to clock/run-order noise.
Complete H3 DiT rerun after the fix: three warmed runs average 15.834297 s; both audio/video latent hashes still match the original baseline in every run. This is a correctness smoke check, not a fresh paired end-to-end speedup claim. Raw: full-dit-post-review.json.
Ruff and whitespace checks pass.

This preserves dense semantics. No sparsity, cache approximation or probability precision change is introduced. DCO signed off. Prepared with AI assistance and measured on the stated hardware. Broader SM120 cards and Windows are not yet benchmarked; feedback on the initial dispatch range is welcome.

_{🔄 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/218 **Author:** [@Tokha233](https://github.com/Tokha233) **Created:** 9/30/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `cuda/sm120-dense-tma` --- ### 📝 Commits (2) - [`e931351`](https://github.com/Comfy-Org/comfy-kitchen/commit/e931351942d2ed2f2eecec653ff2e2621a9dd22a) Pipeline SM120 D128 dense INT8 attention with TMA and cached Q - [`2eb97e0`](https://github.com/Comfy-Org/comfy-kitchen/commit/2eb97e0948180e568d722dab65ca8c2db1fa0790) Validate TMA against math SDPA and fix nonpositive scales ### 📊 Changes **15 files changed** (+1661 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+14 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/dense_tma_sm120.cu` (+94 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/dense_tma_sm120.cuh` (+363 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/dense_tma_sm120.h` (+16 -0) 📝 `comfy_kitchen/backends/cuda/sage_attention/sage_attn_launcher.cu` (+19 -0) ➕ `comfy_kitchen/backends/cuda/sage_attention/tma_pipeline.cuh` (+44 -0) ➕ `docs/benchmarks/attention-post-review-benchmark.json` (+143 -0) ➕ `docs/benchmarks/attention-public-benchmark.json` (+142 -0) ➕ `docs/benchmarks/attention-reference-regression.json` (+467 -0) ➕ `docs/benchmarks/dense-tma-sm120.md` (+91 -0) ➕ `docs/benchmarks/full-dit-post-review.json` (+39 -0) ➕ `docs/benchmarks/full-dit.json` (+82 -0) ➕ `docs/benchmarks/source-sha256.json` (+12 -0) ➕ `tests/benchmark_attention_tma.py` (+59 -0) ➕ `tests/test_int8_attention_tma.py` (+76 -0) </details> ### 📄 Description This adds a narrow SM120 path for long unmasked D128 INT8 attention: a three-slot TMA K/V ring and producer warpgroup overlap operand movement with cached-Q computation. For positive scales, it reuses main's quantization, QK/PV MMA order, online-softmax helpers and last-tile masking. A separate nonpositive-scale specialization applies scaling before the maximum and padding mask; the independent reference tests uncovered that the initial implementation inherited main's incorrect negative-scale behavior. Main's cached-Q path is currently enabled on Ada. This is a separate SM120 implementation, selected only for contiguous BHSD, BF16 output, CTA128, Lq/Lkv>=8192 and uint32-safe offsets. Other inputs retain the ordinary launcher. The specialized object is compiled for SM120a; TMA entry points are resolved through the runtime without a new libcuda link dependency. ### RTX 5090 D v2 measurements Baseline `19ea55b` (Kitchen 0.2.36), CUDA13.0.88, PyTorch2.12.0+cu130. Eight groups of ten warmed CUDA Graph replays, separate baseline/PR extensions. All default-positive-scale benchmark output tensors match exactly. | B,Hq,Hkv,Lq,Lkv | Main ms | This PR ms | Less operator time | |---|---:|---:|---:| | 1,8,2,8192,8193 | 0.602890 | 0.579014 | 3.96% | | 1,8,8,16384,16384 | 2.144539 | 2.092317 | 2.44% | | 1,42,42,32700,32700 | 42.892517 | 42.624614 | 0.62% | | 2,4,2,8197,9217 | 0.712371 | 0.683634 | 4.03% | Raw samples, exact A/B output procedure and benchmark script are in `docs/benchmarks/dense-tma-sm120.md`. The large H3 shape has only a small gain; no broad deployment-speed claim is intended. A three-run public H3 DiT smoke test averaged 15.860485 → 15.839752 s (0.131%, too small for a reliable end-to-end claim), with identical video/audio latent SHA256. ### Validation and known baseline issues - Built the full extension and the dedicated CMake SM120a target. - Added independent FP32 math SDPA references for default, zero and negative scale, with repeated GQA heads and query-only chunking for bounded memory. All 9 new accuracy/repeat/stream/graph tests pass; NRMSE is 0.009147–0.016305, below 0.03. - The new reference caught all-zero negative-scale results in both unmodified main and the initial PR (NRMSE 1.0 for the two partial-tile cases). The TMA route now scales before max/masking; negative-scale NRMSE is 0.016199–0.016305. Default-positive arithmetic remains unchanged. Raw before/after data: `docs/benchmarks/attention-reference-regression.json`. - Complete attention suite after the fix: 428 passed, 594 skipped, 2 failed. The two existing Lq129/Lkv8193 failures remain outside the TMA route and reproduce on unmodified main: optimized SDPA supplies NaN references for zero/negative scale, and the ordinary kernel also has the inherited negative-scale issue. These tests were not skipped or weakened; this PR does not fix the ordinary kernel. - Fresh public-API benchmark after the fix: operator time reductions 3.71%, 2.20%, 0.74%, 4.35% in table order, all default-positive BF16 outputs still exact against main. Separate warmed processes, 8×10 graph replays; raw samples in `attention-post-review-benchmark.json`. Small timing differences remain subject to clock/run-order noise. - Complete H3 DiT rerun after the fix: three warmed runs average 15.834297 s; both audio/video latent hashes still match the original baseline in every run. This is a correctness smoke check, not a fresh paired end-to-end speedup claim. Raw: `full-dit-post-review.json`. - Ruff and whitespace checks pass. This preserves dense semantics. No sparsity, cache approximation or probability precision change is introduced. DCO signed off. Prepared with AI assistance and measured on the stated hardware. Broader SM120 cards and Windows are not yet benchmarked; feedback on the initial dispatch range is welcome. --- <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:48 +00:00

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

Comfy-Org/comfy-kitchen#230

No description provided.

Rows
Columns