[diffusion] kernel: support qk rotary_embedding in one triton kernel #15480
Conversation
Summary of ChangesHello @triple-Mu, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the efficiency of Rotary Positional Embeddings (RoPE) within the system. By developing a specialized Triton kernel that processes both query and key tensors concurrently, the change minimizes computational overhead and accelerates the application of positional encodings across various diffusion models. This optimization directly contributes to faster model inference and improved overall performance. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new Triton kernel to apply rotary embeddings to query and key tensors simultaneously, which is a great optimization to reduce kernel launch overhead. The changes are well-contained and applied across several models to leverage this new fused kernel. I have identified a couple of issues: a redundant conditional check in _apply_qk_rotary_emb that includes unreachable code, and an unused interleaved parameter in the new Triton kernel _qk_rotary_embedding_kernel which could lead to incorrect behavior. Addressing these points will improve code correctness and clarity.
| if is_neox_style: | ||
| q1, q2 = torch.chunk(q, 2, dim=-1) | ||
| k1, k2 = torch.chunk(k, 2, dim=-1) | ||
| else: | ||
| q1 = q[..., ::2] | ||
| q2 = q[..., 1::2] | ||
| k1 = k[..., ::2] | ||
| k2 = k[..., 1::2] |
There was a problem hiding this comment.
The inner if is_neox_style: check is redundant because it is already inside an if is_neox_style: block. This makes the else branch with slicing logic unreachable. For neox_style, only torch.chunk should be used. The slicing logic is for GPT-J style, which is handled by the Triton kernel in the outer else block.
q1, q2 = torch.chunk(q, 2, dim=-1)
k1, k2 = torch.chunk(k, 2, dim=-1)| stride_k_row, | ||
| stride_cos_row, | ||
| stride_sin_row, | ||
| interleaved: tl.constexpr, |
There was a problem hiding this comment.
The interleaved parameter is defined as a tl.constexpr but is not used within the _qk_rotary_embedding_kernel. The kernel's implementation currently assumes an interleaved memory layout for q and k tensors due to the calculation of offsets_x1 and offsets_x2. If non-interleaved tensors are passed, this will lead to incorrect results. Consider implementing the logic to handle both interleaved and non-interleaved cases based on this parameter, or remove it if only interleaved is supported.
1 participant
Motivation
support qk rotary_embedding in one triton kernel.
Profiler:
before:
after:
result comparison:
before:
A_cat_and_a_dog_baking_a_cake_together_in_a_kitchen._The_cat_is_carefully_measuring_flour_while_the_20251219-220049_47059d65.mp4
after:
A_cat_and_a_dog_baking_a_cake_together_in_a_kitchen._The_cat_is_carefully_measuring_flour_while_the_20251219-224628_47059d65.mp4
Modifications
Accuracy Tests
Benchmarking and Profiling
Checklist