PR #39488: [vLLM IR][Rope] Port RotaryEmbedding and DeepseekScalingRotaryEmbedding to IR Ops

• Repository: vllm-project/vllm
• Author: wxsIcey
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/39488

Description (problem / solution / changelog)

Purpose

Port RotaryEmbedding and DeepseekScalingRotaryEmbedding to vLLM IR, following https://github.com/vllm-project/vllm/issues/38756. For the base RotaryEmbedding, keep the key=None path out of the IR maybe_inplace call and fall back to the existing static implementation.

This preserves cross-layer KV sharing behavior and avoids torch.library schema inference failures on optional returns.

Test Plan

To be tested

Test Result

---

<details> <summary> Essential Elements of an Effective PR Description Checklist </summary>

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

</details>

Changed files

• docs/design/vllm_ir.md (added, +626/-0)
• tests/compile/backend.py (modified, +15/-2)
• tests/compile/passes/distributed/test_sequence_parallelism.py (modified, +13/-16)
• tests/compile/passes/ir/test_clone_cleanup.py (added, +370/-0)
• tests/compile/passes/ir/test_inplace_functionalization.py (added, +403/-0)
• tests/compile/passes/test_functionalization.py (modified, +2/-2)
• tests/compile/passes/test_fusion.py (modified, +4/-15)
• tests/ir/test_inplace_op.py (added, +102/-0)
• tests/test_config.py (modified, +6/-3)
• vllm/_aiter_ops.py (modified, +1/-99)
• vllm/compilation/backends.py (modified, +19/-0)
• vllm/compilation/passes/fusion/allreduce_rms_fusion.py (modified, +16/-9)
• vllm/compilation/passes/fusion/matcher_utils.py (modified, +0/-149)
• vllm/compilation/passes/fusion/qk_norm_rope_fusion.py (modified, +27/-21)
• vllm/compilation/passes/fusion/rms_quant_fusion.py (modified, +26/-12)
• vllm/compilation/passes/fusion/rocm_aiter_fusion.py (modified, +30/-16)
• vllm/compilation/passes/fusion/rope_kvcache_fusion.py (modified, +12/-11)
• vllm/compilation/passes/fusion/sequence_parallelism.py (modified, +11/-9)
• vllm/compilation/passes/inductor_pass.py (modified, +4/-0)
• vllm/compilation/passes/ir/clone_elimination.py (added, +117/-0)
• vllm/compilation/passes/ir/inplace_functionalization.py (added, +98/-0)
• vllm/compilation/passes/ir/lowering_pass.py (modified, +5/-35)
• vllm/compilation/passes/ir/utils.py (added, +40/-0)
• vllm/compilation/passes/pass_manager.py (modified, +7/-1)
• vllm/config/kernel.py (modified, +6/-0)
• vllm/config/vllm.py (modified, +1/-2)
• vllm/envs.py (modified, +3/-7)
• vllm/ir/op.py (modified, +114/-5)
• vllm/ir/ops/__init__.py (modified, +3/-2)
• vllm/ir/ops/layernorm.py (modified, +22/-0)
• vllm/ir/ops/rotary_embedding.py (added, +162/-0)
• vllm/kernels/aiter_ops.py (modified, +117/-0)
• vllm/kernels/oink_ops.py (modified, +46/-5)
• vllm/kernels/triton/layernorm_batch_invariant.py (modified, +30/-0)
• vllm/kernels/vllm_c.py (modified, +69/-0)
• vllm/kernels/xpu_ops.py (modified, +41/-0)
• vllm/model_executor/layers/layernorm.py (modified, +10/-202)
• vllm/model_executor/layers/rotary_embedding/base.py (modified, +20/-70)
• vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py (modified, +16/-43)
• vllm/platforms/cpu.py (modified, +14/-0)
• vllm/platforms/cuda.py (modified, +15/-1)
• vllm/platforms/rocm.py (modified, +16/-1)