[PR #85] [MERGED] Refactor RoPE kernels and add in-place variants #131

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: rope_refactor

📝 Commits (4)

4e70824 Refactor RoPE kernels and add in-place variants
b1f515e remove eager fallback
b8ce472 address code rabbit issues
0bf232e full validation

📊 Changes

16 files changed (+1647 additions, -1395 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+79 -1)
➕ comfy_kitchen/_rope_utils.py (+77 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+0 -1)
📝 comfy_kitchen/backends/cuda/__init__.py (+196 -136)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+83 -77)
➖ comfy_kitchen/backends/cuda/ops/apply_rope.cu (+0 -407)
📝 comfy_kitchen/backends/cuda/ops/rms_rope.cu (+274 -491)
➕ comfy_kitchen/backends/cuda/ops/rope_device.cuh (+99 -0)
📝 comfy_kitchen/backends/eager/__init__.py (+27 -0)
📝 comfy_kitchen/backends/eager/rope.py (+199 -8)
📝 comfy_kitchen/backends/triton/__init__.py (+41 -3)
📝 comfy_kitchen/backends/triton/rms_rope.py (+79 -26)
📝 comfy_kitchen/backends/triton/rope.py (+77 -43)
📝 comfy_kitchen/constraints.py (+8 -4)
📝 tests/test_rms_rope.py (+165 -108)
📝 tests/test_rope.py (+243 -90)

📄 Description

Add in-place APIs for interleaved and split-half RoPE, including fused RMS-RoPE variants.
Refactor CUDA RoPE implementation around shared device helpers; consolidate the former apply_rope kernel into the RMS-RoPE source.
Extend CUDA and Triton paths to support strided views, frequency broadcasting, GQA query/key head-count differences, and paired tensors with different strides.
Expand RoPE/RMS-RoPE coverage, in-place, broadcast, stride, GQA, and safety tests.

This should also address: https://github.com/Comfy-Org/comfy-kitchen/pull/77

_{🔄 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/85 **Author:** [@contentis](https://github.com/contentis) **Created:** 7/27/2026 **Status:** ✅ Merged **Merged:** 7/27/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `rope_refactor` --- ### 📝 Commits (4) - [`4e70824`](https://github.com/Comfy-Org/comfy-kitchen/commit/4e708249e911c47f45caf8958ba60ee4de0a0035) Refactor RoPE kernels and add in-place variants - [`b1f515e`](https://github.com/Comfy-Org/comfy-kitchen/commit/b1f515eb240ebe06c1ee93f8787e9fa3060541bf) remove eager fallback - [`b8ce472`](https://github.com/Comfy-Org/comfy-kitchen/commit/b8ce472c41d4170b84eb06b011ac5e10e40eac27) address code rabbit issues - [`0bf232e`](https://github.com/Comfy-Org/comfy-kitchen/commit/0bf232ef662b72c1662b36c32b5bf0c631fe9565) full validation ### 📊 Changes **16 files changed** (+1647 additions, -1395 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+79 -1) ➕ `comfy_kitchen/_rope_utils.py` (+77 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+0 -1) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+196 -136) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+83 -77) ➖ `comfy_kitchen/backends/cuda/ops/apply_rope.cu` (+0 -407) 📝 `comfy_kitchen/backends/cuda/ops/rms_rope.cu` (+274 -491) ➕ `comfy_kitchen/backends/cuda/ops/rope_device.cuh` (+99 -0) 📝 `comfy_kitchen/backends/eager/__init__.py` (+27 -0) 📝 `comfy_kitchen/backends/eager/rope.py` (+199 -8) 📝 `comfy_kitchen/backends/triton/__init__.py` (+41 -3) 📝 `comfy_kitchen/backends/triton/rms_rope.py` (+79 -26) 📝 `comfy_kitchen/backends/triton/rope.py` (+77 -43) 📝 `comfy_kitchen/constraints.py` (+8 -4) 📝 `tests/test_rms_rope.py` (+165 -108) 📝 `tests/test_rope.py` (+243 -90) </details> ### 📄 Description - Add in-place APIs for interleaved and split-half RoPE, including fused RMS-RoPE variants. - Refactor CUDA RoPE implementation around shared device helpers; consolidate the former apply_rope kernel into the RMS-RoPE source. - Extend CUDA and Triton paths to support strided views, frequency broadcasting, GQA query/key head-count differences, and paired tensors with different strides. - Expand RoPE/RMS-RoPE coverage, in-place, broadcast, stride, GQA, and safety tests. This should also address: https://github.com/Comfy-Org/comfy-kitchen/pull/77 --- <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:41 +00:00

closed this issue
added the
pull-request
label

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#131

No description provided.

Rows
Columns