RFC: Unify Rotary Position Embedding Implementations Across Models

Background

PR #2393 optimized Wan2.2's apply_rotary_emb_wan function by replacing the inefficient torch.empty_like + strided slice assignment pattern with torch.stack + flatten. This achieved a 25% speedup (44s → 33s on NPU, 55s → 44s on A100).

During code review, we discovered that multiple models in vllm-omni have their own rotary embedding implementations with varying efficiency levels.

Current State

Model	Implementation	Pattern	Status
wan_transformer	apply_rotary_emb_wan	torch.stack + flatten	✅ Optimized (PR #2393)
helios_transformer	apply_rotary_emb	empty_like + strided slice	❌ Needs optimization
stable_audio	apply_rotary_embed	torch.cat	✅ Optimized
qwen_image	apply_rotary_emb_qwen	torch.stack + flatten	✅ Optimized
omnigen2	apply_rotary_emb	torch.stack + flatten	✅ Optimized

Meanwhile, vLLM core already provides optimized implementations:

• vllm.model_executor.layers.rotary_embedding.common.ApplyRotaryEmb - Uses efficient patterns for both Neox-style (torch.cat) and interleaved/GPT-J style (torch.stack + flatten)
• vllm.vllm_flash_attn.layers.rotary - Compiled CUDA kernel, optimized at kernel level

Problem

1. Code duplication: Each model implements its own RoPE function
2. Inconsistent performance: Some implementations use inefficient patterns (strided slice assignment causes InplaceCopy/ViewCopy overhead on NPU)
3. Maintenance burden: Optimizations need to be applied to each model individually
4. Missing hardware optimizations: Custom ops and compiled kernels from vLLM core are not being leveraged

Proposal

We propose to unify rotary embedding implementations across all models in vllm-omni:

Phase 1: Immediate Fixes

• Apply the torch.stack + flatten pattern to helios_transformer.py (same fix as PR #2393)
• This is a low-risk change that maintains the same mathematical semantics

Phase 2: Centralized Implementation

• Create a unified apply_rotary_emb utility in a shared location (e.g., vllm_omni/layers/rotary_embedding.py)
• Support all required variations:
  • Neox-style (even/odd split)
  • Interleaved/GPT-J style
  • Optional: different RoPE scaling methods (linear, dynamic, yarn, etc.)

Phase 3: Leverage vLLM Core

• Investigate using ApplyRotaryEmb from vLLM core directly
• For NPU: Consider implementing a custom NPU op for RoPE if profiling shows further optimization opportunities
• This aligns with vLLM's design philosophy of hardware-agnostic model code with platform-specific kernels

Benefits

1. Performance: Ensure all models benefit from optimized implementations
2. Maintainability: Single source of truth for RoPE logic
3. Consistency: Same behavior and performance characteristics across models
4. Future-proof: Easier to add new hardware optimizations (NPU custom ops, etc.)

Call for Discussion

1. Should we prioritize reusing vLLM core's ApplyRotaryEmb or maintain our own implementation?
2. Are there any model-specific RoPE variations that cannot be unified?
3. What's the preferred location for the shared implementation?

---

cc @fhfuih @wtomin @hsliuustc0106