Triton Backend - Ensure different dim1/dim2 matched against freqs_cis independently#77
Triton Backend - Ensure different dim1/dim2 matched against freqs_cis independently#77sleppyrobot wants to merge 1 commit into
Conversation
|
✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged. |
|
Important Review skippedNo new commits to review since the last review. ⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughChangesThe Triton RoPE backend now prepares and launches the kernel separately for each tensor, using tensor-specific shapes and strides. Frequency cropping and contiguous input handling are performed in the new single-tensor helper; public entry points retain their signatures. RoPE geometry handling
🚥 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: 2
🤖 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/triton/rope.py`:
- Around line 162-164: Remove the forced x.contiguous() conversion and update
the Triton RoPE kernel to accept and use separate stride_x_out_* parameters for
output indexing. Allocate x_out with contiguous format while preserving x’s
original strides for reads, and pass both stride sets through the kernel launch
to avoid the intermediate input copy.
- Around line 169-170: Remove the freqs_cis.is_contiguous() conditional and the
corresponding contiguous() conversion after the slice, leaving freqs_cis as the
strided view passed to the kernel.
🪄 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: 3152d3ad-00b4-4171-b31c-b48d18ecc2f6
📒 Files selected for processing (1)
comfy_kitchen/backends/triton/rope.py
| # Ensure inputs are contiguous | ||
| if not x1.is_contiguous(): | ||
| x1 = x1.contiguous() | ||
| if not x.is_contiguous(): | ||
| x = x.contiguous() |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy lift
Avoid copying x by supporting separate output strides in the kernel.
Forcing x to be contiguous prevents illegal memory accesses by ensuring x and x_out share the exact same strides (since torch.empty_like on a view often falls back to standard dense strides). However, this incurs a full memory copy of the input.
To save some time and keep it fast, don't copy the tensor, let the strides do the heavy cast!
Consider updating the Triton kernel to accept separate stride_x_out_* arguments. This would allow you to allocate a contiguous x_out (torch.empty_like(x, memory_format=torch.contiguous_format)) and read from a non-contiguous x directly without any intermediate allocations.
🤖 Prompt for 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.
In `@comfy_kitchen/backends/triton/rope.py` around lines 162 - 164, Remove the
forced x.contiguous() conversion and update the Triton RoPE kernel to accept and
use separate stride_x_out_* parameters for output indexing. Allocate x_out with
contiguous format while preserving x’s original strides for reads, and pass both
stride sets through the kernel launch to avoid the intermediate input copy.
| if not freqs_cis.is_contiguous(): | ||
| freqs_cis = freqs_cis.contiguous() |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick win
Remove unnecessary contiguous check for freqs_cis.
Since the kernel only reads from freqs_cis and explicitly computes offsets using its strides (stride_freqs_*), there is no need to force it to be contiguous. When freqs_cis is sliced on line 168, it becomes a view, and this check forces an expensive memory copy on every forward pass.
Let's slice the cost and drop the copy—leaving freqs_cis as a view won't make things floppy!
💨 Proposed fix
- if not freqs_cis.is_contiguous():
- freqs_cis = freqs_cis.contiguous()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if not freqs_cis.is_contiguous(): | |
| freqs_cis = freqs_cis.contiguous() |
🤖 Prompt for 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.
In `@comfy_kitchen/backends/triton/rope.py` around lines 169 - 170, Remove the
freqs_cis.is_contiguous() conditional and the corresponding contiguous()
conversion after the slice, leaving freqs_cis as the strided view passed to the
kernel.
|
@coderabbitai recheck |
|
✅ Action performedReview finished.
|
|
I have read and agree to the Contributor License Agreement |
sleppyrobot
changed the title
1 participant
Since Krea2 launched i have not been able to get it running with --enable-triton-backend on Linux on my 7900xtx, it triggers ; "CUDA error: an illegal memory access was encountered"
Comparing the eager rope.py to the triton revealed an issue that may occur due to the shapes\stride generated by Krea2.
This change mimics the eager logic to allow the model to run, also give a small speed up as triton rope is faster than eager rope.