[PR #220] Add projection-ready BSHD output to prequantized INT8 attention #232

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: cuda/prequantized-bshd-output

📝 Commits (5)

e931351 Pipeline SM120 D128 dense INT8 attention with TMA and cached Q
2eb97e0 Validate TMA against math SDPA and fix nonpositive scales
2a1650b Add contiguous BSHD output to prequantized attention
9f70304 Bind benchmark operands per case
13bbea6 Cover BSHD TMA streams graphs and nonpositive scales

📊 Changes

21 files changed (+2032 additions, -15 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+14 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+13 -6)
➕ comfy_kitchen/backends/cuda/sage_attention/dense_tma_sm120.cu (+97 -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)
📝 comfy_kitchen/sage_attention.py (+21 -9)
➕ 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/layout-benchmark.json (+186 -0)
➕ docs/benchmarks/layout.md (+30 -0)
➕ docs/benchmarks/source-sha256.json (+12 -0)
➕ tests/benchmark_attention_tma.py (+59 -0)
➕ tests/benchmark_layout.py (+59 -0)
➕ tests/test_int8_attention_layout.py (+54 -0)

...and 1 more files

📄 Description

Adds an explicit output_layout="BSHD" option to prequantized INT8 attention, so output projections can consume contiguous BSHD storage without a separate transpose/copy. BHSD remains the default. Native CUDA uses output strides; HIP retains a conversion fallback.

Depends on #218. This branch includes #218; the additional layout-only diff is https://github.com/Tokha233/comfy-kitchen/compare/cuda/sm120-dense-tma...cuda/prequantized-bshd-output . Please review the layout-only comparison as the incremental contribution; #218 must merge first.

Validation on RTX 5090 D v2 / Torch 2.12.0+cu130: 37 tests passed (layout, batch 2, GQA, masks, D64/96/128/256, BF16/FP16/FP32, streams and CUDA Graphs, plus #218 accuracy regressions). All benchmark outputs are bitwise equal.

B,H,S,D	BHSD + reorder	Direct BSHD	Time reduction
1,56,14850,128	12.16595 ms	11.82173 ms	2.83%
1,56,32700,128	57.38058 ms	56.60735 ms	1.35%
1,56,87142,128	402.39087 ms	400.22939 ms	0.54%
1,56,90461,128	433.01341 ms	430.83040 ms	0.50%

Eight alternating groups, four graph replays per sample. These are attention-plus-layout timings, not whole-model/serving speedups. Raw samples and reproduction: docs/benchmarks/layout.md, tests/benchmark_layout.py. The separate ComfyUI consumer in Comfy-Org/ComfyUI#16678 completed four alternating full-sampler pairs: 15.778404 → 15.713747 s (−0.410%), with equal video/audio latent hashes. Both arms used the same TMA build. Other GPUs/HIP are not tested here.

Prepared with AI assistance; tests and measurements were executed on the stated hardware. DCO signed off.

Latest combined validation and per-PR fixes: https://github.com/Tokha233/ComfyUI-H3-SpeedKit/blob/main/docs/UPSTREAM-RECOVERY-TESTS-20261001.zh-CN.md . Combined gains overlap and must not be added to the single-operation timings above.

_{🔄 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/220 **Author:** [@Tokha233](https://github.com/Tokha233) **Created:** 9/30/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `cuda/prequantized-bshd-output` --- ### 📝 Commits (5) - [`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 - [`2a1650b`](https://github.com/Comfy-Org/comfy-kitchen/commit/2a1650b44a74f0f9d0e3993c4b882bab62b51c01) Add contiguous BSHD output to prequantized attention - [`9f70304`](https://github.com/Comfy-Org/comfy-kitchen/commit/9f70304669543c7f691f46234519d9b06fbaad39) Bind benchmark operands per case - [`13bbea6`](https://github.com/Comfy-Org/comfy-kitchen/commit/13bbea660c11c20718336ff3bc791667b235b105) Cover BSHD TMA streams graphs and nonpositive scales ### 📊 Changes **21 files changed** (+2032 additions, -15 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+14 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+13 -6) ➕ `comfy_kitchen/backends/cuda/sage_attention/dense_tma_sm120.cu` (+97 -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) 📝 `comfy_kitchen/sage_attention.py` (+21 -9) ➕ `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/layout-benchmark.json` (+186 -0) ➕ `docs/benchmarks/layout.md` (+30 -0) ➕ `docs/benchmarks/source-sha256.json` (+12 -0) ➕ `tests/benchmark_attention_tma.py` (+59 -0) ➕ `tests/benchmark_layout.py` (+59 -0) ➕ `tests/test_int8_attention_layout.py` (+54 -0) _...and 1 more files_ </details> ### 📄 Description Adds an explicit `output_layout="BSHD"` option to prequantized INT8 attention, so output projections can consume contiguous BSHD storage without a separate transpose/copy. BHSD remains the default. Native CUDA uses output strides; HIP retains a conversion fallback. **Depends on #218.** This branch includes #218; the additional layout-only diff is https://github.com/Tokha233/comfy-kitchen/compare/cuda/sm120-dense-tma...cuda/prequantized-bshd-output . Please review the layout-only comparison as the incremental contribution; #218 must merge first. Validation on RTX 5090 D v2 / Torch 2.12.0+cu130: 37 tests passed (layout, batch 2, GQA, masks, D64/96/128/256, BF16/FP16/FP32, streams and CUDA Graphs, plus #218 accuracy regressions). All benchmark outputs are bitwise equal. | B,H,S,D | BHSD + reorder | Direct BSHD | Time reduction | |---|---:|---:|---:| | 1,56,14850,128 | 12.16595 ms | 11.82173 ms | 2.83% | | 1,56,32700,128 | 57.38058 ms | 56.60735 ms | 1.35% | | 1,56,87142,128 | 402.39087 ms | 400.22939 ms | 0.54% | | 1,56,90461,128 | 433.01341 ms | 430.83040 ms | 0.50% | Eight alternating groups, four graph replays per sample. These are attention-plus-layout timings, not whole-model/serving speedups. Raw samples and reproduction: `docs/benchmarks/layout.md`, `tests/benchmark_layout.py`. The separate ComfyUI consumer in Comfy-Org/ComfyUI#16678 completed four alternating full-sampler pairs: 15.778404 → 15.713747 s (−0.410%), with equal video/audio latent hashes. Both arms used the same TMA build. Other GPUs/HIP are not tested here. Prepared with AI assistance; tests and measurements were executed on the stated hardware. DCO signed off. Latest combined validation and per-PR fixes: https://github.com/Tokha233/ComfyUI-H3-SpeedKit/blob/main/docs/UPSTREAM-RECOVERY-TESTS-20261001.zh-CN.md . Combined gains overlap and must not be added to the single-operation timings above. --- <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:50 +00:00

bsakaguchi referenced this issue

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

[PR #232] [MERGED] Add 3D mesh operations for remeshing (CORE-560) #241