@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f38f8c9742

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 19:00 UTC

[ROCm]: Enable customop and rope+kvcache fusion for AITER RoPE (#35180)

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

Rohan Potdar committed Feb 24, 2026 at 22:36 UTC f38f8c974277a2fbd4bcb915416fd907fc023366

9 files changed +139 -67

tests/compile/passes/test_rope_kvcache_fusion.py

+9 -1

@@ -177,7 +177,10 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):

177 def ops_in_model_before(self) -> list[torch._ops.OpOverload]:

178 ops = []

179 if self.enable_rope_custom_op:

180 - ops.append(ROTARY_OP)

180 + if rocm_aiter_ops.is_triton_rotary_embed_enabled():

181 + ops.append(torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default)

182 + else:

183 + ops.append(ROTARY_OP)

184 else:

185 ops.append(INDEX_SELECT_OP)

186 ops.append(torch.ops.vllm.unified_kv_cache_update.default)

@@ -196,6 +199,7 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):

199 ],

200 )

201 @pytest.mark.parametrize("enable_rope_custom_op", [True]) # [True, False])

202 +@pytest.mark.parametrize("enable_aiter_triton_rope", [True, False])

203 @pytest.mark.parametrize("num_heads", [64])

204 @pytest.mark.parametrize("num_kv_heads", [8])

205 @pytest.mark.parametrize("head_size", [64])

@@ -210,6 +214,7 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):

214 def test_rope_kvcache_fusion(

215 attn_backend: AttentionBackendEnum,

216 enable_rope_custom_op: bool,

217 + enable_aiter_triton_rope: bool,

218 num_heads: int,

219 num_kv_heads: int,

220 head_size: int,

@@ -245,6 +250,9 @@ def test_rope_kvcache_fusion(

250

251 with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:

252 m.setenv("VLLM_ROCM_USE_AITER", "1")

253 + m.setenv(

254 + "VLLM_ROCM_USE_AITER_TRITON_ROPE", "1" if enable_aiter_triton_rope else "0"

255 + )

256 rocm_aiter_ops.refresh_env_variables()

257

258 model = QKRoPEKVCacheTestModel(

vllm/_aiter_ops.py

+65 -36

@@ -831,6 +831,59 @@ def _rocm_aiter_triton_add_rmsnorm_pad_fake(

831 return out, residual_out

832

833

834 +def _triton_rotary_embedding_impl(

835 + positions: torch.Tensor,

836 + query: torch.Tensor,

837 + key: torch.Tensor,

838 + head_size: int,

839 + cos_sin_cache: torch.Tensor,

840 + is_neox: bool,

841 + offsets: torch.Tensor | None = None,

842 +) -> None:

843 + # Modifies query and key in-place

844 + from aiter.ops.triton.rope.rope import (

845 + rope_cached_thd_positions_offsets_2c_fwd_inplace,

846 + )

847 +

848 + num_tokens = positions.numel()

849 + cos, sin = cos_sin_cache.chunk(2, dim=-1)

850 + query_shape = query.shape

851 + key_shape = key.shape

852 + rotate_style = 0 if is_neox else 1

853 + rotary_dim = head_size

854 +

855 + query = query.view(num_tokens, -1, head_size)

856 + key = key.view(num_tokens, -1, head_size)

857 + query_ = query[..., :rotary_dim]

858 + key_ = key[..., :rotary_dim]

859 + positions = positions.view(*query.shape[:1])

860 + rope_cached_thd_positions_offsets_2c_fwd_inplace(

861 + query_,

862 + key_,

863 + cos,

864 + sin,

865 + positions,

866 + offsets,

867 + rotate_style,

868 + reuse_freqs_front_part=True,

869 + nope_first=False,

870 + )

871 + query = query.view(query_shape)

872 + key = key.view(key_shape)

873 +

874 +

875 +def _triton_rotary_embedding_fake(

876 + positions: torch.Tensor,

877 + query: torch.Tensor,

878 + key: torch.Tensor,

879 + head_size: int,

880 + cos_sin_cache: torch.Tensor,

881 + is_neox_style: bool,

882 + offsets: torch.Tensor | None = None,

883 +) -> None:

884 + return

885 +

886 +

887 # Global flag to ensure ops are registered only once

888 _OPS_REGISTERED = False

889

@@ -1178,6 +1231,14 @@ class rocm_aiter_ops:

1231 dispatch_key=current_platform.dispatch_key,

1232 )

1233

1234 + # Register rocm aiter rotary embedding custom op

1235 + direct_register_custom_op(

1236 + op_name="rocm_aiter_triton_rotary_embedding",

1237 + op_func=_triton_rotary_embedding_impl,

1238 + mutates_args=["query", "key"], # These tensors are modified in-place

1239 + fake_impl=_triton_rotary_embedding_fake,

1240 + )

1241 +

1242 _OPS_REGISTERED = True

1243

1244 @staticmethod

@@ -1220,6 +1281,10 @@ class rocm_aiter_ops:

1281 def get_triton_add_rmsnorm_pad_op() -> OpOverload:

1282 return torch.ops.vllm.rocm_aiter_triton_add_rmsnorm_pad.default

1283

1284 + @staticmethod

1285 + def get_triton_rotary_embedding_op() -> OpOverload:

1286 + return torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default

1287 +

1288 @staticmethod

1289 def rms_norm(

1290 x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float

@@ -1482,42 +1547,6 @@ class rocm_aiter_ops:

1547 gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y)

1548 return y

1549

1485 - @staticmethod

1486 - def triton_rotary_embed(

1487 - positions: torch.Tensor,

1488 - query: torch.Tensor,

1489 - key: torch.Tensor,

1490 - cos_sin_cache: torch.Tensor,

1491 - head_size: int,

1492 - rotary_dim: int,

1493 - is_neox_style: bool,

1494 - ):

1495 - from aiter.ops.triton.rope import rope_cached_thd_positions_2c_fwd_inplace

1496 -

1497 - num_tokens = positions.numel()

1498 - cos, sin = cos_sin_cache.chunk(2, dim=-1)

1499 - query_shape = query.shape

1500 - key_shape = key.shape

1501 - rotate_style = 0 if is_neox_style else 1

1502 -

1503 - query = query.view(num_tokens, -1, head_size)

1504 - key = key.view(num_tokens, -1, head_size)

1505 - query_ = query[..., :rotary_dim]

1506 - key_ = key[..., :rotary_dim]

1507 - positions = positions.view(*query.shape[:1])

1508 - rope_cached_thd_positions_2c_fwd_inplace(

1509 - query_,

1510 - key_,

1511 - cos,

1512 - sin,

1513 - positions,

1514 - rotate_style,

1515 - reuse_freqs_front_part=True,

1516 - nope_first=False,

1517 - )

1518 - query = query.view(query_shape)

1519 - key = key.view(key_shape)

1520 -

1550 @staticmethod

1551 def triton_rope_and_cache(

1552 query: torch.Tensor,

vllm/compilation/passes/fusion/matcher_utils.py

@@ -89,10 +89,13 @@ class MatcherRotaryEmbedding(MatcherCustomOp):

89 num_heads: int,

90 num_kv_heads: int,

91 use_flashinfer: bool = False,

92 + match_rocm_aiter: bool | None = None,

93 enabled: bool | None = None,

94 ) -> None:

95 if enabled is None:

96 enabled = RotaryEmbedding.enabled()

97 + if match_rocm_aiter is None:

98 + match_rocm_aiter = rocm_aiter_ops.is_triton_rotary_embed_enabled()

100 super().__init__(enabled)

101 self.is_neox = is_neox

@@ -104,6 +107,8 @@ class MatcherRotaryEmbedding(MatcherCustomOp):

107 self.rotary_dim = head_size

108 if use_flashinfer:

109 self.rotary_op = FLASHINFER_ROTARY_OP

110 + elif match_rocm_aiter:

111 + self.rotary_op = rocm_aiter_ops.get_triton_rotary_embedding_op()

112 else:

113 self.rotary_op = ROTARY_OP

114

vllm/compilation/passes/utility/scatter_split_replace.py

+5 -1

@@ -60,6 +60,10 @@ class ScatterSplitReplacementPass(VllmInductorPass):

60 def __call__(self, graph: fx.Graph) -> None:

61 count = 0

63 + target_ops = [torch.ops._C.rotary_embedding.default]

64 + if hasattr(torch.ops.vllm, "rocm_aiter_triton_rotary_embedding"):

65 + target_ops.append(torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default)

66 +

67 for node in graph.nodes:

68 if not is_func(node, auto_functionalized):

69 continue

@@ -67,7 +71,7 @@ class ScatterSplitReplacementPass(VllmInductorPass):

71 kwargs = node.kwargs

72 at_target = node.args[0]

70 - if at_target == torch.ops._C.rotary_embedding.default:

74 + if at_target in target_ops:

75 query = kwargs["query"]

76 key = kwargs["key"]

77 getitem_nodes = {}

vllm/config/compilation.py

+13 -16

@@ -123,6 +123,8 @@ class PassConfig:

123 """Enable async TP."""

124 fuse_allreduce_rms: bool = Field(default=None)

125 """Enable flashinfer allreduce fusion."""

126 + enable_qk_norm_rope_fusion: bool = False

127 + """Enable fused Q/K RMSNorm + RoPE pass."""

128

129 # ROCm/AITER specific fusions

130 fuse_act_padding: bool = Field(default=None)

@@ -153,8 +155,6 @@ class PassConfig:

155 8: 1, # 1MB

156 },

157 }, where key is the device capability"""

156 - enable_qk_norm_rope_fusion: bool = False

157 - """Enable fused Q/K RMSNorm + RoPE pass."""

158

159 # TODO(luka) better pass enabling system.

160

@@ -834,23 +834,20 @@ class CompilationConfig:

834 func if isinstance(func, InductorPass) else CallableInductorPass(func)

835 )

836

837 - if self.pass_config.enable_qk_norm_rope_fusion:

837 + if (

838 + self.pass_config.enable_qk_norm_rope_fusion

839 + and "+rotary_embedding" not in self.custom_ops

840 + ):

841 # TODO(zhuhaoran): support rope native forward match and remove this.

842 # Linked issue: https://github.com/vllm-project/vllm/issues/28042

843 self.custom_ops.append("+rotary_embedding")

841 - if self.pass_config.fuse_rope_kvcache:

842 - from vllm._aiter_ops import rocm_aiter_ops

843 -

844 - if rocm_aiter_ops.is_triton_rotary_embed_enabled():

845 - logger.warning(

846 - "Cannot use VLLM_ROCM_USE_AITER_TRITON_ROPE with "

847 - "fuse_rope_kvcache. Disabling fuse_rope_kvcache."

848 - )

849 - self.pass_config.fuse_rope_kvcache = False

850 - else:

851 - # TODO(Rohan138): support rope native forward match and remove this.

852 - # Linked issue: https://github.com/vllm-project/vllm/issues/28042

853 - self.custom_ops.append("+rotary_embedding")

844 + if (

845 + self.pass_config.fuse_rope_kvcache

846 + and "+rotary_embedding" not in self.custom_ops

847 + ):

848 + # TODO(Rohan138): support rope native forward match and remove this.

849 + # Linked issue: https://github.com/vllm-project/vllm/issues/28042

850 + self.custom_ops.append("+rotary_embedding")

851

852 if (

853 is_torch_equal_or_newer("2.9.0.dev")

vllm/config/vllm.py

+20 -3

@@ -126,14 +126,27 @@ def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:

126 )

127

128

129 +def enable_rope_kvcache_fusion(cfg: "VllmConfig") -> bool:

130 + """Enable if rotary embedding custom op is active and

131 + use_inductor_graph_partition is enabled.

132 + """

133 + from vllm._aiter_ops import rocm_aiter_ops

134 +

135 + return (

136 + rocm_aiter_ops.is_enabled()

137 + and cfg.compilation_config.is_custom_op_enabled("rotary_embedding")

138 + and cfg.compilation_config.use_inductor_graph_partition

139 + )

140 +

141 +

142 def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:

143 """Enable if using AITER RMSNorm and AITER Triton GEMMs

144 and hidden size is 2880 i.e. gpt-oss; otherwise Inductor handles fusion."""

145 + from vllm._aiter_ops import rocm_aiter_ops

146

147 return (

134 - envs.VLLM_ROCM_USE_AITER

135 - and envs.VLLM_ROCM_USE_AITER_RMSNORM

136 - and envs.VLLM_ROCM_USE_AITER_TRITON_GEMM

148 + rocm_aiter_ops.is_rmsnorm_enabled()

149 + and not rocm_aiter_ops.is_triton_gemm_enabled()

150 and cfg.model_config is not None

151 and cfg.model_config.get_hidden_size() == 2880

152 )

@@ -149,6 +162,7 @@ OPTIMIZATION_LEVEL_00 = {

162 "enable_sp": False,

163 "fuse_gemm_comms": False,

164 "fuse_act_padding": False,

165 + "fuse_rope_kvcache": False,

166 },

167 "cudagraph_mode": CUDAGraphMode.NONE,

168 "use_inductor_graph_partition": False,

@@ -167,6 +181,7 @@ OPTIMIZATION_LEVEL_01 = {

181 "enable_sp": False,

182 "fuse_gemm_comms": False,

183 "fuse_act_padding": enable_norm_pad_fusion,

184 + "fuse_rope_kvcache": enable_rope_kvcache_fusion,

185 },

186 "cudagraph_mode": CUDAGraphMode.PIECEWISE,

187 "use_inductor_graph_partition": False,

@@ -185,6 +200,7 @@ OPTIMIZATION_LEVEL_02 = {

200 "enable_sp": IS_DENSE,

201 "fuse_gemm_comms": IS_DENSE,

202 "fuse_act_padding": enable_norm_pad_fusion,

203 + "fuse_rope_kvcache": enable_rope_kvcache_fusion,

204 },

205 "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,

206 "use_inductor_graph_partition": False,

@@ -203,6 +219,7 @@ OPTIMIZATION_LEVEL_03 = {

219 "enable_sp": IS_DENSE,

220 "fuse_gemm_comms": IS_DENSE,

221 "fuse_act_padding": enable_norm_pad_fusion,

222 + "fuse_rope_kvcache": enable_rope_kvcache_fusion,

223 },

224 "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,

225 "use_inductor_graph_partition": False,

vllm/envs.py

+3 -3

@@ -105,7 +105,7 @@ if TYPE_CHECKING:

105 VLLM_ROCM_USE_AITER_MLA: bool = True

106 VLLM_ROCM_USE_AITER_MHA: bool = True

107 VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False

108 - VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False

108 + VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = True

109 VLLM_ROCM_USE_AITER_FP8BMM: bool = True

110 VLLM_ROCM_USE_AITER_FP4BMM: bool = True

111 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False

@@ -937,9 +937,9 @@ environment_variables: dict[str, Callable[[], Any]] = {

937 os.getenv("VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", "False").lower() in ("true", "1")

938 ),

939 # Whether to use aiter rope.

940 - # By default is disabled.

940 + # By default is enabled.

941 "VLLM_ROCM_USE_AITER_TRITON_ROPE": lambda: (

942 - os.getenv("VLLM_ROCM_USE_AITER_TRITON_ROPE", "False").lower() in ("true", "1")

942 + os.getenv("VLLM_ROCM_USE_AITER_TRITON_ROPE", "True").lower() in ("true", "1")

943 ),

944 # Whether to use aiter triton fp8 bmm kernel

945 # By default is enabled.

vllm/model_executor/layers/rotary_embedding/base.py

+11 -7

@@ -47,15 +47,20 @@ class RotaryEmbeddingBase(CustomOp):

47 if not hasattr(self, "use_flashinfer"):

48 self.use_flashinfer = False

50 + self.use_aiter = (

51 + self.enabled() and rocm_aiter_ops.is_triton_rotary_embed_enabled()

52 + )

53 + if self.use_aiter:

54 + self.rocm_aiter_triton_rotary_embedding = (

55 + rocm_aiter_ops.get_triton_rotary_embedding_op()

56 + )

57 +

58 if init_cache:

59 cache = self._compute_cos_sin_cache()

60 if not self.use_flashinfer:

61 cache = cache.to(dtype)

62 self.cos_sin_cache: torch.Tensor

63 self.register_buffer("cos_sin_cache", cache, persistent=False)

56 - self.is_rocm_triton_rotary_embed_enabled = (

57 - rocm_aiter_ops.is_triton_rotary_embed_enabled()

58 - )

65 self.apply_rotary_emb = ApplyRotaryEmb(

66 is_neox_style=self.is_neox_style,

@@ -231,15 +236,14 @@ class RotaryEmbedding(RotaryEmbeddingBase):

236 query: torch.Tensor,

237 key: torch.Tensor | None = None,

238 ) -> tuple[torch.Tensor, torch.Tensor | None]:

234 - if self.is_rocm_triton_rotary_embed_enabled:

239 + if self.use_aiter:

240 cos_sin_cache = self._match_cos_sin_cache_dtype(query)

236 - rocm_aiter_ops.triton_rotary_embed(

241 + self.rocm_aiter_triton_rotary_embedding(

242 positions,

243 query,

244 key,

240 - cos_sin_cache,

245 self.head_size,

242 - self.rotary_dim,

246 + cos_sin_cache,

247 self.is_neox_style,

248 )

249 return query, key

vllm/platforms/rocm.py

@@ -494,6 +494,7 @@ class RocmPlatform(Platform):

494 use_aiter_rms_norm = rocm_aiter_ops.is_rmsnorm_enabled()

495 use_aiter_fp8_linear = rocm_aiter_ops.is_linear_fp8_enabled()

496 use_aiter_fused_se = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()

497 + use_aiter_triton_rope = rocm_aiter_ops.is_triton_rotary_embed_enabled()

498

499 if compilation_config.cudagraph_mode.has_full_cudagraphs():

500 # decode context parallel does not support full cudagraphs

@@ -558,6 +559,13 @@ class RocmPlatform(Platform):

559 and "-grouped_topk" not in compilation_config.custom_ops

560 ):

561 compilation_config.custom_ops.append("+grouped_topk")

562 + # Enable rotary embedding when using AITER if its not disabled by user

563 + if (

564 + use_aiter_triton_rope

565 + and "+rotary_embedding" not in compilation_config.custom_ops

566 + and "-rotary_embedding" not in compilation_config.custom_ops

567 + ):

568 + compilation_config.custom_ops.append("+rotary_embedding")

569

570 # Default dispatch to rocm's sparse_attn_indexer implementation

571 compilation_config.custom_ops.append("+sparse_attn_indexer")