[PR #79] [MERGED] Add rms_rope to Triton backend. #130

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/79
Author: @bedovyy
Created: 7/19/2026
Status: ✅ Merged
Merged: 7/20/2026
Merged by: @comfyanonymous

Base: main ← Head: triton_rms_rope

📝 Commits (1)

c8e14d6 Add rms_rope to Triton backend.

📊 Changes

2 files changed (+249 additions, -0 deletions)

View changed files

📝 comfy_kitchen/backends/triton/__init__.py (+41 -0)
➕ comfy_kitchen/backends/triton/rms_rope.py (+208 -0)

📄 Description

This PR adds the rms_rope implementation to the Triton backend.

Code formatting follows the conventions in rope.py.
I placed it in a separate rms_rope.py file since I wasn't entirely sure of the best location for it within the current directory structure.

Testing:
test_rms_rope_cuda.txt
test_rms_rope_rocm.txt

CUDA (RTX5090): Passed all 72 Triton-related items in test_rms_rope.
- Note: A warning occurs in the FLUX-freqs_fp32-bf16-triton-rms_rope case when the input is bf16 and the weight is fp32. I decided to keep the current logic as it provides a good balance between speed and accuracy. - means no idea how to resolve.
ROCm (AI PRO R9700): 20 items related to freqs=torch.bfloat16 failed. However, this behavior is consistent with the current test_rope results, so it's not a new regression.

Resolves:
https://github.com/Comfy-Org/ComfyUI/issues/14968

Overall generation speed remains similar or shows a slight improvement.

Benchmark Logs - R9700 Power limit 260

Edit: Because freqs_batch is changed to runtime scalar, I redo the generation test on Anima.

after PR + no torch.compile
100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:15<00:00,  1.88it/s]
[INFO] Prompt executed in 16.35 seconds
100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:16<00:00,  1.87it/s]
[INFO] Prompt executed in 16.41 seconds
100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:16<00:00,  1.86it/s]
[INFO] Prompt executed in 16.49 seconds
100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:16<00:00,  1.85it/s]
[INFO] Prompt executed in 16.55 seconds

after PR + torch.compile
100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:13<00:00,  2.23it/s]
[INFO] Prompt executed in 13.82 seconds
100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:13<00:00,  2.22it/s]
[INFO] Prompt executed in 13.92 seconds
100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:13<00:00,  2.21it/s]
[INFO] Prompt executed in 13.97 seconds
100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:13<00:00,  2.20it/s]
[INFO] Prompt executed in 14.00 seconds

_{🔄 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/79 **Author:** [@bedovyy](https://github.com/bedovyy) **Created:** 7/19/2026 **Status:** ✅ Merged **Merged:** 7/20/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `triton_rms_rope` --- ### 📝 Commits (1) - [`c8e14d6`](https://github.com/Comfy-Org/comfy-kitchen/commit/c8e14d63d161e09c79f7354f0e1df04593c62dd5) Add rms_rope to Triton backend. ### 📊 Changes **2 files changed** (+249 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/triton/__init__.py` (+41 -0) ➕ `comfy_kitchen/backends/triton/rms_rope.py` (+208 -0) </details> ### 📄 Description This PR adds the `rms_rope` implementation to the Triton backend. - Code formatting follows the conventions in `rope.py`. - I placed it in a separate `rms_rope.py` file since I wasn't entirely sure of the best location for it within the current directory structure. **Testing:** [test_rms_rope_cuda.txt](https://github.com/user-attachments/files/30165891/test_rms_rope_cuda.txt) [test_rms_rope_rocm.txt](https://github.com/user-attachments/files/30165892/test_rms_rope_rocm.txt) - **CUDA (RTX5090):** Passed all 72 Triton-related items in `test_rms_rope`. - *Note:* A warning occurs in the `FLUX-freqs_fp32-bf16-triton-rms_rope` case when the input is `bf16` and the weight is `fp32`. I decided to keep the current logic as it provides a good balance between speed and accuracy. *- means no idea how to resolve.* - **ROCm (AI PRO R9700):** 20 items related to `freqs=torch.bfloat16` failed. However, this behavior is consistent with the current `test_rope` results, so it's not a new regression. **Resolves:** https://github.com/Comfy-Org/ComfyUI/issues/14968 Overall generation speed remains similar or shows a slight improvement. <details><summary>Benchmark Logs - R9700 Power limit 260</summary> *Edit: Because `freqs_batch` is changed to runtime scalar, I redo the generation test on Anima.* ```text after PR + no torch.compile 100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:15<00:00, 1.88it/s] [INFO] Prompt executed in 16.35 seconds 100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:16<00:00, 1.87it/s] [INFO] Prompt executed in 16.41 seconds 100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:16<00:00, 1.86it/s] [INFO] Prompt executed in 16.49 seconds 100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:16<00:00, 1.85it/s] [INFO] Prompt executed in 16.55 seconds after PR + torch.compile 100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:13<00:00, 2.23it/s] [INFO] Prompt executed in 13.82 seconds 100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:13<00:00, 2.22it/s] [INFO] Prompt executed in 13.92 seconds 100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:13<00:00, 2.21it/s] [INFO] Prompt executed in 13.97 seconds 100%|██████████████████████████████████████████████████████████████████████████████████| 30/30 [00:13<00:00, 2.20it/s] [INFO] Prompt executed in 14.00 seconds ``` </details> --- <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:39 +00:00