[PR #70] [MERGED] Add RMS fusion for RoPE #122

New issue

Closed

opened 2026-10-26 22:20:36 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:20:36 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/70
Author: @contentis
Created: 7/10/2026
Status: ✅ Merged
Merged: 7/13/2026
Merged by: @comfyanonymous

Base: main ← Head: rmsnorm_rope

📝 Commits (4)

7a57d66 Add RMS fusion to RoPE
1a6190e Add RMS fusion to RoPE
26cfae6 Merge remote-tracking branch 'origin/rmsnorm_rope' into rmsnorm_rope
87a673c remove merge duplicates

📊 Changes

9 files changed (+1728 additions, -1 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+68 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+1 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+267 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+199 -0)
📝 comfy_kitchen/backends/cuda/float_utils.cuh (+1 -0)
➕ comfy_kitchen/backends/cuda/ops/rms_rope.cu (+568 -0)
📝 comfy_kitchen/backends/eager/__init__.py (+82 -1)
📝 comfy_kitchen/backends/eager/rope.py (+180 -0)
➕ tests/test_rms_rope.py (+362 -0)

📄 Description

RMSNorm + RoPE

Measured BF16 Q/K/frequencies, FP32 RMS scales, eps=1e-6, 10 warmups, 8 samples, and 10 grouped iterations. BW is API-boundary bytes (inputs + outputs + frequencies + scales); Rated peak is that API BW divided by the GPU's nominal DRAM bandwidth; speedup is unfused / fused. The unfused baseline is F.rms_norm for Q/K followed by the matching ck.apply_rope* call.

Reference peak bandwidths: RTX 5000 Ada 470 GB/s; RTX PRO 6000 Blackwell 1790 GB/s.

GPU	Config	Variant	Fused us	Unfused us	Speedup	Fused BW GB/s	Rated peak	Unfused BW GB/s
RTX 5000 Ada	FLUX BHND (1,24,4352,128)	interleaved	229.734	2070.259	9.01x	475.3	101.1%	52.7
RTX 5000 Ada	LTX BHND (2,32,4996,64)	interleaved	352.896	3605.120	10.22x	471.2	100.3%	46.1
RTX 5000 Ada	ZIMAGE BNHD (1,4096,30,128)	interleaved	271.168	2609.459	9.62x	471.8	100.4%	49.0
RTX 5000 Ada	WAN BNHD (2,12288,16,128)	split-half	894.003	9101.312	10.18x	464.5	98.8%	45.6
RTX 5000 Ada	FLUX BHND (1,24,4352,128)	split-half	232.614	2069.747	8.90x	469.4	99.9%	52.8
RTX 5000 Ada	LTX BHND (2,32,4996,64)	split-half	355.187	3599.002	10.13x	468.1	99.6%	46.2
RTX PRO 6000 Blackwell	FLUX BHND (1,24,4352,128)	interleaved	37.280	380.850	10.22x	2928.8	163.6%	286.7
RTX PRO 6000 Blackwell	LTX BHND (2,32,4996,64)	interleaved	98.642	818.253	8.30x	1685.6	94.2%	203.2
RTX PRO 6000 Blackwell	ZIMAGE BNHD (1,4096,30,128)	interleaved	38.834	548.502	14.12x	3294.2	184.0%	233.2
RTX PRO 6000 Blackwell	WAN BNHD (2,12288,16,128)	split-half	291.855	2902.382	9.94x	1422.8	79.5%	143.1
RTX PRO 6000 Blackwell	FLUX BHND (1,24,4352,128)	split-half	36.536	373.450	10.22x	2988.4	167.0%	292.4
RTX PRO 6000 Blackwell	LTX BHND (2,32,4996,64)	split-half	101.569	814.978	8.02x	1637.0	91.5%	204.0

Some bandwidths are above the limit due to L2 cache hits that skew the exact result.

_{🔄 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/70 **Author:** [@contentis](https://github.com/contentis) **Created:** 7/10/2026 **Status:** ✅ Merged **Merged:** 7/13/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `rmsnorm_rope` --- ### 📝 Commits (4) - [`7a57d66`](https://github.com/Comfy-Org/comfy-kitchen/commit/7a57d66c8d6eb53d7eb6b50b03bbc1bc198c0732) Add RMS fusion to RoPE - [`1a6190e`](https://github.com/Comfy-Org/comfy-kitchen/commit/1a6190e5200d07ca6a53203d03228d5ef90a526e) Add RMS fusion to RoPE - [`26cfae6`](https://github.com/Comfy-Org/comfy-kitchen/commit/26cfae62528b2619779ecef2ad20cadde9854a01) Merge remote-tracking branch 'origin/rmsnorm_rope' into rmsnorm_rope - [`87a673c`](https://github.com/Comfy-Org/comfy-kitchen/commit/87a673c836feef7ce57d745d97527e843cc4f6c4) remove merge duplicates ### 📊 Changes **9 files changed** (+1728 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+68 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+1 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+267 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+199 -0) 📝 `comfy_kitchen/backends/cuda/float_utils.cuh` (+1 -0) ➕ `comfy_kitchen/backends/cuda/ops/rms_rope.cu` (+568 -0) 📝 `comfy_kitchen/backends/eager/__init__.py` (+82 -1) 📝 `comfy_kitchen/backends/eager/rope.py` (+180 -0) ➕ `tests/test_rms_rope.py` (+362 -0) </details> ### 📄 Description # RMSNorm + RoPE Measured BF16 Q/K/frequencies, FP32 RMS scales, `eps=1e-6`, 10 warmups, 8 samples, and 10 grouped iterations. `BW` is API-boundary bytes (inputs + outputs + frequencies + scales); `Rated peak` is that API BW divided by the GPU's nominal DRAM bandwidth; speedup is `unfused / fused`. The unfused baseline is `F.rms_norm` for Q/K followed by the matching `ck.apply_rope*` call. Reference peak bandwidths: RTX 5000 Ada 470 GB/s; RTX PRO 6000 Blackwell 1790 GB/s. | GPU | Config | Variant | Fused us | Unfused us | Speedup | Fused BW GB/s | Rated peak | Unfused BW GB/s | |---|---|---|---:|---:|---:|---:|---:|---:| | RTX 5000 Ada | FLUX BHND (1,24,4352,128) | interleaved | 229.734 | 2070.259 | 9.01x | 475.3 | 101.1% | 52.7 | | RTX 5000 Ada | LTX BHND (2,32,4996,64) | interleaved | 352.896 | 3605.120 | 10.22x | 471.2 | 100.3% | 46.1 | | RTX 5000 Ada | ZIMAGE BNHD (1,4096,30,128) | interleaved | 271.168 | 2609.459 | 9.62x | 471.8 | 100.4% | 49.0 | | RTX 5000 Ada | WAN BNHD (2,12288,16,128) | split-half | 894.003 | 9101.312 | 10.18x | 464.5 | 98.8% | 45.6 | | RTX 5000 Ada | FLUX BHND (1,24,4352,128) | split-half | 232.614 | 2069.747 | 8.90x | 469.4 | 99.9% | 52.8 | | RTX 5000 Ada | LTX BHND (2,32,4996,64) | split-half | 355.187 | 3599.002 | 10.13x | 468.1 | 99.6% | 46.2 | | RTX PRO 6000 Blackwell | FLUX BHND (1,24,4352,128) | interleaved | 37.280 | 380.850 | 10.22x | 2928.8 | 163.6% | 286.7 | | RTX PRO 6000 Blackwell | LTX BHND (2,32,4996,64) | interleaved | 98.642 | 818.253 | 8.30x | 1685.6 | 94.2% | 203.2 | | RTX PRO 6000 Blackwell | ZIMAGE BNHD (1,4096,30,128) | interleaved | 38.834 | 548.502 | 14.12x | 3294.2 | 184.0% | 233.2 | | RTX PRO 6000 Blackwell | WAN BNHD (2,12288,16,128) | split-half | 291.855 | 2902.382 | 9.94x | 1422.8 | 79.5% | 143.1 | | RTX PRO 6000 Blackwell | FLUX BHND (1,24,4352,128) | split-half | 36.536 | 373.450 | 10.22x | 2988.4 | 167.0% | 292.4 | | RTX PRO 6000 Blackwell | LTX BHND (2,32,4996,64) | split-half | 101.569 | 814.978 | 8.02x | 1637.0 | 91.5% | 204.0 | Some bandwidths are above the limit due to L2 cache hits that skew the exact result. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 22:20:36 +00:00