add AITER Triton RoPE as a registered ops with VLLM_USE_AITER_TRITON_ROPE=1

Author: k50112113

vllm/model_executor/layers/rotary_embedding/base.py

@@ -5,12 +5,11 @@
 
 import torch
 
-from vllm import envs
+import vllm.envs as envs
 from vllm.model_executor.custom_op import CustomOp
 
 from .common import apply_rotary_emb_dispatch, apply_rotary_emb_torch
-from .rocm_aiter_rope_ops import is_rocm_rotary_embedding_enabled
-
+from .rocm_aiter_rope_ops import is_rocm_rotary_embedding_enabled, is_rocm_triton_rotary_embedding_enabled
 
 @CustomOp.register("rotary_embedding")
 class RotaryEmbedding(CustomOp):
@@ -38,6 +37,7 @@ def __init__(
         self.cos_sin_cache: torch.Tensor
         self.register_buffer("cos_sin_cache", cache, persistent=False)
         self.is_rocm_aiter_enabled = is_rocm_rotary_embedding_enabled()
+        self.is_rocm_aiter_triton_enabled = is_rocm_triton_rotary_embedding_enabled()
 
     def _compute_inv_freq(self, base: float) -> torch.Tensor:
         """Compute the inverse frequency."""
@@ -110,8 +110,8 @@ def forward_cuda(
                                                        dtype=query.dtype)
 
         num_tokens = positions.numel()
-        if envs.VLLM_USE_AITER_TRITON_ROPE and num_tokens <= 128:
-            import aiter.ops.triton.rope as ops
+
+        if envs.VLLM_USE_AITER_TRITON_ROPE:
             assert key is not None
             cos, sin = self.cos_sin_cache.chunk(2, dim=-1)
             query_shape = query.shape
@@ -120,30 +120,20 @@ def forward_cuda(
             key = key.view(num_tokens, -1, self.head_size)
             query_ = query[..., :self.rotary_dim]
             key_ = key[..., :self.rotary_dim]
+            rotate_style = 0 if self.is_neox_style else 1
+            positions = positions.view(*query.shape[:1])
             if offsets is not None:
                 offsets = offsets.view(*query.shape[:1])
-                ops.rope_cached_thd_positions_offsets_2c_fwd_inplace(
-                    query_,
-                    key_,
-                    cos,
-                    sin,
-                    positions,
-                    offsets,
-                    0,
-                    True,
-                    False,
-                )
-            else:
-                ops.rope_cached_thd_positions_2c_fwd_inplace(
-                    query_,
-                    key_,
-                    cos,
-                    sin,
-                    positions,
-                    0,
-                    True,
-                    False,
-                )
+            torch.ops.vllm.rocm_aiter_rotary_emb_with_key_forward_triton(
+                positions,
+                sin,
+                cos,
+                query_,
+                key_,
+                offsets,
+                rotate_style,
+                False,
+            )
             query = query.view(query_shape)
             key = key.view(key_shape)
         else:
@@ -173,7 +163,9 @@ def forward_hip(
     ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         # currently only rotary embedding ops from AITER package are
         # supported for HiP forward.
-        if self.is_rocm_aiter_enabled:
+        if self.is_rocm_aiter_triton_enabled:
+            return self.forward_cuda(positions, query, key, offsets)
+        elif self.is_rocm_aiter_enabled:
             return self.forward_hip_rocm_aiter(positions, query, key, offsets,
                                                is_nope_first)
         return self.forward_native(positions, query, key, offsets)

vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py

@@ -89,7 +89,9 @@ def forward(
         offsets: Optional[torch.Tensor] = None,
     ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         """PyTorch-native implementation equivalent to forward()."""
-        if self.is_rocm_aiter_enabled:
+        if self.is_rocm_aiter_triton_enabled:
+            return self.forward_cuda(positions, query, key, offsets)
+        elif self.is_rocm_aiter_enabled:
             return self.forward_hip_rocm_aiter(positions, query, key, offsets)
 
         assert key is not None

vllm/model_executor/layers/rotary_embedding/rocm_aiter_rope_ops.py

@@ -13,6 +13,9 @@
 def is_rocm_rotary_embedding_enabled() -> bool:
     return (current_platform.is_rocm() and envs.VLLM_ROCM_USE_AITER)
 
+def is_rocm_triton_rotary_embedding_enabled() -> bool:
+    return (current_platform.is_rocm() and envs.VLLM_ROCM_USE_AITER and envs.VLLM_USE_AITER_TRITON_ROPE)
+
 
 def rocm_aiter_rotary_emb_without_key_forward_hip_impl(
     positions: torch.Tensor,
@@ -124,4 +127,65 @@ def rocm_aiter_rotary_emb_without_key_forward_hip_fake(
         mutates_args=["query"],
         fake_impl=rocm_aiter_rotary_emb_without_key_forward_hip_fake,
         dispatch_key=current_platform.dispatch_key,
+    )
+
+
+
+def rocm_aiter_rotary_emb_with_key_forward_triton_impl(
+    positions: torch.Tensor,
+    sin: torch.Tensor,
+    cos: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    offsets: Optional[torch.Tensor] = None,
+    rotate_style: int = 0,
+    is_nope_first: bool = False,
+) -> None:
+    import aiter.ops.triton.rope as ops
+    if offsets is None:
+        ops.rope_cached_thd_positions_2c_fwd_inplace(
+            query,
+            key,
+            cos,
+            sin,
+            positions,
+            rotate_style,
+            reuse_freqs_front_part=True,
+            nope_first=is_nope_first,
+        )
+    else:
+        ops.rope_cached_thd_positions_offsets_2c_fwd_inplace(
+            query,
+            key,
+            cos,
+            sin,
+            positions,
+            offsets,
+            rotate_style,
+            reuse_freqs_front_part=True,
+            nope_first=is_nope_first,
+        )
+
+
+def rocm_aiter_rotary_emb_with_key_forward_triton_fake(
+    positions: torch.Tensor,
+    sin: torch.Tensor,
+    cos: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    offsets: Optional[torch.Tensor] = None,
+    rotate_style: int = 0,
+    is_nope_first: bool = False,
+) -> None:
+    pass
+
+
+if is_rocm_triton_rotary_embedding_enabled():
+
+    direct_register_custom_op(
+        op_name="rocm_aiter_rotary_emb_with_key_forward_triton",
+        op_func=rocm_aiter_rotary_emb_with_key_forward_triton_impl,
+        mutates_args=["key", "query"],
+        fake_impl=rocm_aiter_rotary_emb_with_key_forward_triton_fake,
+        dispatch_key=current_platform.dispatch_key,
     )