[PR #217] Add exact BF16 register ConvRot256 fast path for SM120 #227

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: cuda/convrot-register-sm120

📝 Commits (1)

6e2006d Add exact BF16 register ConvRot256 fast path for SM120

📊 Changes

9 files changed (+1291 additions, -0 deletions)

View changed files

📝 NOTICE (+7 -0)
➕ comfy_kitchen/backends/cuda/ops/convrot_register.cuh (+143 -0)
📝 comfy_kitchen/backends/cuda/ops/int8_linear.cu (+42 -0)
➕ docs/benchmarks/convrot-graph-benchmark.json (+837 -0)
➕ docs/benchmarks/convrot-register-sm120.md (+47 -0)
➕ docs/benchmarks/full-dit.json (+82 -0)
➕ docs/benchmarks/source-sha256.json (+7 -0)
➕ tests/benchmark_convrot_register.py (+52 -0)
➕ tests/test_convrot_register.py (+74 -0)

📄 Description

This adds an exact BF16 register implementation of ConvRot256 for SM120 at K=5376/7168. It keeps Kitchen's FP32 operation association, per-stage 0.5 scaling, BF16 quantization barriers and rounding. The existing implementation handles other widths, architectures, dtypes, stochastic rounding and fused input activations.

Eight-lane groups keep rotation values in registers and exchange the required values with warp shuffles, replacing shared-memory rotation scratch. Aligned vector loads/stores are guarded in the launcher. There is no new quantization format or approximation.

RTX 5090 D v2 measurements

Baseline 19ea55b (Kitchen 0.2.36), CUDA 13.0.88, PyTorch 2.12.0+cu130. CUDA events, warmed graph replays, eight alternating groups; all raw samples and a reproduction script are included in docs/benchmarks/convrot-register-sm120.md.

M×K	Main	This PR	Less operator time
4096×5376	0.049215 ms	0.023842 ms	51.55%
32700×5376	0.489247 ms	0.444126 ms	9.22%
65536×5376	0.976363 ms	0.891171 ms	8.73%
32700×7168	0.647534 ms	0.592686 ms	8.47%
90461×7168	1.781528 ms	1.643211 ms	7.76%

Full public H3 DiT smoke test (Larry v4, 8 steps, 768×512/124 frames, fixed seed/condition): three warmed runs averaged 15.860485 → 15.815302 s; both video/audio latent SHA256 match. The 0.285% full-DiT difference is small and sequentially measured, so I am not presenting it as a robust end-to-end speedup.

Validation

CUDA extension built against the pinned upstream dependencies.
80 passed, 3 skipped: new ConvRot tests plus existing input-activation tests.
INT8 output and FP32 scale bits match the retained generic kernel for both widths, including zero/negative zero, tiny/huge values, NaN/Inf, random BF16 bit patterns, nondefault stream and graph replay.
Ruff and git diff --check pass.

Layout provenance is attributed to SGLang PR #38040; Kitchen arithmetic is retained. Commit is DCO signed off. Prepared with AI assistance; the measurements above were run on the stated hardware. Other GPU architectures have not been benchmarked for this specialization.

_{🔄 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/217 **Author:** [@Tokha233](https://github.com/Tokha233) **Created:** 9/30/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `cuda/convrot-register-sm120` --- ### 📝 Commits (1) - [`6e2006d`](https://github.com/Comfy-Org/comfy-kitchen/commit/6e2006d4a715939847ff553cb46949d18c5f7171) Add exact BF16 register ConvRot256 fast path for SM120 ### 📊 Changes **9 files changed** (+1291 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `NOTICE` (+7 -0) ➕ `comfy_kitchen/backends/cuda/ops/convrot_register.cuh` (+143 -0) 📝 `comfy_kitchen/backends/cuda/ops/int8_linear.cu` (+42 -0) ➕ `docs/benchmarks/convrot-graph-benchmark.json` (+837 -0) ➕ `docs/benchmarks/convrot-register-sm120.md` (+47 -0) ➕ `docs/benchmarks/full-dit.json` (+82 -0) ➕ `docs/benchmarks/source-sha256.json` (+7 -0) ➕ `tests/benchmark_convrot_register.py` (+52 -0) ➕ `tests/test_convrot_register.py` (+74 -0) </details> ### 📄 Description This adds an exact BF16 register implementation of ConvRot256 for SM120 at K=5376/7168. It keeps Kitchen's FP32 operation association, per-stage 0.5 scaling, BF16 quantization barriers and rounding. The existing implementation handles other widths, architectures, dtypes, stochastic rounding and fused input activations. Eight-lane groups keep rotation values in registers and exchange the required values with warp shuffles, replacing shared-memory rotation scratch. Aligned vector loads/stores are guarded in the launcher. There is no new quantization format or approximation. ### RTX 5090 D v2 measurements Baseline `19ea55b` (Kitchen 0.2.36), CUDA 13.0.88, PyTorch 2.12.0+cu130. CUDA events, warmed graph replays, eight alternating groups; all raw samples and a reproduction script are included in `docs/benchmarks/convrot-register-sm120.md`. | M×K | Main | This PR | Less operator time | |---|---:|---:|---:| | 4096×5376 | 0.049215 ms | 0.023842 ms | 51.55% | | 32700×5376 | 0.489247 ms | 0.444126 ms | 9.22% | | 65536×5376 | 0.976363 ms | 0.891171 ms | 8.73% | | 32700×7168 | 0.647534 ms | 0.592686 ms | 8.47% | | 90461×7168 | 1.781528 ms | 1.643211 ms | 7.76% | Full public H3 DiT smoke test (Larry v4, 8 steps, 768×512/124 frames, fixed seed/condition): three warmed runs averaged 15.860485 → 15.815302 s; both video/audio latent SHA256 match. The 0.285% full-DiT difference is small and sequentially measured, so I am not presenting it as a robust end-to-end speedup. ### Validation - CUDA extension built against the pinned upstream dependencies. - 80 passed, 3 skipped: new ConvRot tests plus existing input-activation tests. - INT8 output and FP32 scale bits match the retained generic kernel for both widths, including zero/negative zero, tiny/huge values, NaN/Inf, random BF16 bit patterns, nondefault stream and graph replay. - Ruff and `git diff --check` pass. Layout provenance is attributed to SGLang PR #38040; Kitchen arithmetic is retained. Commit is DCO signed off. Prepared with AI assistance; the measurements above were run on the stated hardware. Other GPU architectures have not been benchmarked for this specialization. --- <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:47 +00:00

bsakaguchi referenced this issue

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

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

bsakaguchi referenced this issue

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

[PR #231] Tune SM120 V quantization for large D128 inputs #240