Add RMS fusion for RoPE - #70
Merged
Merged
Conversation
# Conflicts: # comfy_kitchen/backends/cuda/__init__.py
|
✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
💤 Files with no reviewable changes (1)
📝 WalkthroughWalkthroughChangesAdds four RMSNorm-plus-RoPE operators with eager and CUDA implementations, public exports, custom-op registration, CUDA kernel bindings, layout and dtype constraints, and cross-backend numerical tests. RMS RoPE support
Sequence Diagram(s)sequenceDiagram
participant Caller
participant ComfyKitchen
participant CUDABackend
participant CUDAKernel
Caller->>ComfyKitchen: call rms_rope variant
ComfyKitchen->>CUDABackend: dispatch tensors, scales, frequencies
CUDABackend->>CUDAKernel: validate metadata and launch kernel
CUDAKernel-->>CUDABackend: write normalized and rotated outputs
CUDABackend-->>Caller: return q and k outputs
🚥 Pre-merge checks | ✅ 2✅ Passed checks (2 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
✨ Simplify code
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@comfy_kitchen/backends/cuda/__init__.py`:
- Around line 2699-2782: Remove the duplicated FunctionConstraints entries for
rms_rope, rms_rope1, rms_rope_split_half, and rms_rope_split_half1 from the
later portion of the constraints dictionary, keeping the original definitions
and eliminating duplicate-key overwrites.
- Around line 1948-2122: Remove the duplicate RMS-RoPE helper and wrapper
definitions, keeping only one complete copy of _native_rms_rope_layout,
_has_vectorizable_rms_rope_rows, _rms_rope1_cuda, _rms_rope_cuda, rms_rope1,
rms_rope, rms_rope_split_half1, and rms_rope_split_half.
In `@tests/test_rms_rope.py`:
- Around line 237-306: Update test_split_half_cross_backend to skip when "eager"
is absent from backends, before populating or accessing results; retain the
existing minimum-backend check and add an explicit "eager" membership guard so
both results["eager"] references remain safe.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 35016161-9930-48e5-b4c0-9bf6c5b7a591
📒 Files selected for processing (9)
comfy_kitchen/__init__.pycomfy_kitchen/backends/cuda/CMakeLists.txtcomfy_kitchen/backends/cuda/__init__.pycomfy_kitchen/backends/cuda/dlpack_bindings.cppcomfy_kitchen/backends/cuda/float_utils.cuhcomfy_kitchen/backends/cuda/ops/rms_rope.cucomfy_kitchen/backends/eager/__init__.pycomfy_kitchen/backends/eager/rope.pytests/test_rms_rope.py
Contributor
Author
|
I have read and agree to the Contributor License Agreement |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
RMSNorm + RoPE
Measured BF16 Q/K/frequencies, FP32 RMS scales,
eps=1e-6, 10 warmups, 8 samples, and 10 grouped iterations.BWis API-boundary bytes (inputs + outputs + frequencies + scales);Rated peakis that API BW divided by the GPU's nominal DRAM bandwidth; speedup isunfused / fused. The unfused baseline isF.rms_normfor Q/K followed by the matchingck.apply_rope*call.Reference peak bandwidths: RTX 5000 Ada 470 GB/s; RTX PRO 6000 Blackwell 1790 GB/s.
Some bandwidths are above the limit due to L2 cache hits that skew the exact result.