[Perf]Optimize rotary_emb implementation to use Triton operator for improved inference performance

[cynthieye]

Optimize Rotary Positional Embeddings with Triton Kernel in VLLM

This PR enhances rotary positional embedding computation by leveraging Triton-optimized kernels from flash_attn, addressing a significant performance bottleneck observed in models like Qwen2-VL.

Background
The original PyTorch-native rotary embedding implementation (rotary_emb) consumed 40-60% of total inference latency for Qwen2-VL, particularly scaling with output token count. Profiling revealed inefficiencies in tensor reshaping and element-wise operations.

Changes
Triton Kernel Integration:

Conditionally uses flash_attn.ops.triton.rotary.apply_rotary when flash-attn>=2.0 is available.

Falls back to the native PyTorch implementation otherwise.

Dynamic Implementation Selection:

Added _use_flash_attn flag to RotaryEmbedding classes, auto-detecting flash_attn availability during initialization.

Performance Gains
Qwen2-VL: Achieved 17% end-to-end speedup when generating 150 output tokens.

Generalization: Expected improvements for other mrotary-based models, pending further benchmarks.

Code Compatibility
Maintains backward compatibility with existing model architectures.

Requires no user-side changes—automatically prioritizes Triton kernel when dependencies are met.

This optimization significantly reduces rotary embedding overhead while preserving numerical equivalence, making it particularly impactful for long-context and long-generation scenarios.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[jeejeelee]

Perhaps using kernel from https://github.com/vllm-project/flash-attention would be more reasonable

[Isotr0py]

Shouldn't this triton kernel have been used for Qwen2-VL?

vllm/vllm/model_executor/models/qwen2_vl.py

Lines 233 to 243 in 87e067d

	def apply_rotary_pos_emb_vision(t: torch.Tensor,
	freqs: torch.Tensor,
	use_flash_attn=False) -> torch.Tensor:
	t_ = t.float()
	cos = freqs.cos()
	sin = freqs.sin()
	apply_rotary_emb = apply_rotary_emb_torch
	if use_flash_attn:
	from flash_attn.layers.rotary import apply_rotary_emb
	output = apply_rotary_emb(t_, cos, sin).type_as(t)
	return output

vllm/vllm/model_executor/models/qwen2_5_vl.py

Lines 276 to 282 in 87e067d

	use_flash_attn = self.attn_backend == _Backend.FLASH_ATTN
	q = apply_rotary_pos_emb_vision(q,
	rotary_pos_emb,
	use_flash_attn=use_flash_attn)
	k = apply_rotary_pos_emb_vision(k,
	rotary_pos_emb,
	use_flash_attn=use_flash_attn)

[cynthieye]

apply_rotary_pos_emb_vision

qwen2_vl.py

class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,
                                      SupportsLoRA, SupportsPP):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        ......

        self.language_model = init_vllm_registered_model(
            vllm_config=vllm_config,
            prefix=maybe_prefix(prefix, "language_model"),
            architectures=["Qwen2ForCausalLM"],
        )
        ......
       
            
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        intermediate_tensors: Optional[IntermediateTensors] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        **kwargs: object,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        ......

        hidden_states = self.language_model.model(
            input_ids=input_ids,
            positions=positions,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
        )
        return hidden_states

qwen2.py

from vllm.model_executor.layers.rotary_embedding import get_rope

class Qwen2Attention(nn.Module):
    def __init__(self,
                 hidden_size: int,
                 num_heads: int,
                 num_kv_heads: int,
                 max_position: int = 4096 * 32,
                 rope_theta: float = 10000,
                 cache_config: Optional[CacheConfig] = None,
                 quant_config: Optional[QuantizationConfig] = None,
                 rope_scaling: Optional[Tuple] = None) -> None:
        super().__init__()
        ......

        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position,
            base=self.rope_theta,
            rope_scaling=rope_scaling,
        )
        ......
    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
    ) -> torch.Tensor:
        ......
        q, k = self.rotary_emb(positions, q, k)
        ......

@Isotr0py In qwen2-vl, the Qwen2Model of qwen2 will be called, and the rotary-embedding in Qwen2Model does not use the Triton kernel. During the token ID generation stage, there will be a performance bottleneck, which accounts for 40-60% of the overall inference time

[Isotr0py]

Suggested change

	if use_flash_attn:
	return apply_rotary_emb(x.unsqueeze(0), cos, sin,
	not is_neox_style).squeeze(0)
	else:
	return _apply_rotary_emb_torch(x, cos, sin, is_neox_style)
	if current_platform.is_cuda_alike():
	from vllm_flash_attn.layers.rotary import apply_rotary_emb
	return apply_rotary_emb(x.unsqueeze(0), cos, sin,
	not is_neox_style).squeeze(0)
	else:
	return _apply_rotary_emb_torch(x, cos, sin, is_neox_style)

Since triton is only available for Nvidia and ROCm GPUs, we can simplify the implementation here.

[cynthieye]

hi，isotr0py，if I replace

        from flash_attn.layers.rotary import apply_rotary_emb
        return apply_rotary_emb(x.unsqueeze(0), cos, sin,
                                not is_neox_style).squeeze(0)

to

        from vllm_flash_attn.layers.rotary import apply_rotary_emb
        return apply_rotary_emb(x.unsqueeze(0), cos, sin,
                                not is_neox_style).squeeze(0)

Running the CI test will result in an error：
https://buildkite.com/vllm/ci/builds/18001#01964d1c-6acb-48b6-8b8e-6b296cdefd17

[2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/qwen2.py", line 243, in forward

  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] hidden_states = self.self_attn(
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] ^^^^^^^^^^^^^^^
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] return self._call_impl(*args, **kwargs)
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1750, in _call_impl
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] return forward_call(*args, **kwargs)
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/qwen2.py", line 176, in forward
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] q, k = self.rotary_emb(positions, q, k)
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] return self._call_impl(*args, **kwargs)
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1750, in _call_impl
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] return forward_call(*args, **kwargs)
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/rotary_embedding.py", line 992, in forward
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/rotary_embedding.py", line 81, in _apply_rotary_emb
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] from vllm_flash_attn.layers.rotary import apply_rotary_emb
  | [2025-04-19T08:53:45Z] ERROR 04-19 01:53:45 [core.py:386] ModuleNotFoundError: No module named 'vllm_flash_attn'

How should I solve it

[cynthieye]

I use flash.attn.layers.gross import
Mainly because qwen2-vl.py is used in this way, as follows

def apply_rotary_pos_emb_vision(t: torch.Tensor,
                                freqs: torch.Tensor,
                                use_flash_attn=False) -> torch.Tensor:
    t_ = t.float()
    cos = freqs.cos()
    sin = freqs.sin()
    apply_rotary_emb = apply_rotary_emb_torch
    if use_flash_attn:
        from flash_attn.layers.rotary import apply_rotary_emb
    output = apply_rotary_emb(t_, cos, sin).type_as(t)
    return output

[Isotr0py]

Oh, the import code should be vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb, otherwise it won't work outside the repo folder.

Qwen2-vl use original FA due to vllm_flash_attn compatibility issue for ViT, however, we shouldn't use it here since original FA is not a requirement.

[cynthieye]

Using vllm-vllm_flash_mattn.layers.total still encounters errors：
https://buildkite.com/vllm/ci/builds/18013#01964e8f-a7a0-4638-baa0-a57dfa5ac8b4

  | [2025-04-19T15:42:15Z] ERROR 04-19 08:42:15 [core.py:386] return forward_call(*args, **kwargs)
  | [2025-04-19T15:42:15Z] ERROR 04-19 08:42:15 [core.py:386] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | [2025-04-19T15:42:15Z] ERROR 04-19 08:42:15 [core.py:386] File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/rotary_embedding.py", line 992, in forward
  | [2025-04-19T15:42:15Z] ERROR 04-19 08:42:15 [core.py:386] query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
  | [2025-04-19T15:42:15Z] ERROR 04-19 08:42:15 [core.py:386] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | [2025-04-19T15:42:15Z] ERROR 04-19 08:42:15 [core.py:386] File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/rotary_embedding.py", line 81, in _apply_rotary_emb
  | [2025-04-19T15:42:15Z] ERROR 04-19 08:42:15 [core.py:386] from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
  | [2025-04-19T15:42:15Z] ERROR 04-19 08:42:15 [core.py:386] ModuleNotFoundError: No module named 'vllm.vllm_flash_attn.layers'

[cynthieye]

Are you saying that I need to make modifications https://github.com/vllm-project/flash-attention project ？
Can I modify it like this:

1. Create a 'layers' sub directory in the' vllm_flash-attn 'directory
2. Then copy 'flash_attn/layers/total. py' to the 'vllm_flash-attn/layers/' directory,
3. Import or export 'apply_rotary-emb' in 'vllm_flash-attn/layers/init. py'.
4. Ensure that the top-level 'init. py' of 'vllm_flash-attn' correctly exports the 'layers' module.

[Isotr0py]

Are you saying that I need to make modifications https://github.com/vllm-project/flash-attention project ？

Yes. We should not directly copy the source code from FA to vllm_flash_attn, because the copying should happen at the compilation during installation. So I recommend to make modifications in https://github.com/vllm-project/flash-attention to make rotary code copied during compilation.

[cynthieye]

The file I want to copy/flash attention/flash attn/layers/rotary.py contains the following code
from flash_attn.ops.triton.rotary import apply_rotary
If I copy /flash-attention/flash-attn/layers/rotary.py to the /flash-attention/vllm_flash-attn/ directory during compilation, the error will still occur during execution，
I need to modify the code to become
from vllm.vllm_flash_attn.ops.triton.rotary import apply_rotary

[cynthieye]

Are you saying that I need to make modifications https://github.com/vllm-project/flash-attention project ？

Yes. We should not directly copy the source code from FA to vllm_flash_attn, because the copying should happen at the compilation during installation. So I recommend to make modifications in https://github.com/vllm-project/flash-attention to make rotary code copied during compilation.

Can you help me merge this flash attn PR？
vllm-project/flash-attention#64

[cynthieye]

@Isotr0py hi，I have merged the changes to the flash attn repository. Can you help me review and merge this PR again

[Isotr0py]

Overall LGTM now, just leave some nits. PTAL!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @cynthieye.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Isotr0py]

LGTM now, thanks for your patience!

[vadiklyutiy]

Performance Gains
Qwen2-VL: Achieved 17% end-to-end speedup when generating 150 output tokens.

@cynthieye Could you clarify what command line args did you use for measurement?

[vadiklyutiy]

@cynthieye @Isotr0py
It seems to me that this PR caused some discrepancy with the design of the torch.compile integration.

V1 with torch.compile explicitly avoid using any custom implementation. For control what custom op implementation should be enable is used CompilationConfig.custom_ops and by default used native torch version (see CustomOp.dispatch_forward)

With this change for class MRotaryEmbedding was enabled using custom cuda implementation even if torch.compile is enabled.

@youkaichao Please correct me if I wrong.

[Isotr0py]

V1 with torch.compile explicitly avoid using any custom implementation. For control what custom op implementation should be enable is used CompilationConfig.custom_ops and by default used native torch version (see CustomOp.dispatch_forward)

But MRoPE implementation itself doesn't obey custom op implementation design very much, because it has overridden the forward method, and dispatch_forward won't take any effect.

[vadiklyutiy]

yes, correct. And I propose to fix it.