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, 14:00 UTC
[ROCm]: fix aiter rope functionalization (#35533)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Rohan Potdar committed
Feb 27, 2026 at 16:42 UTC
e3691988d0bf7c915d544d9204b53506815464ca
1 file changed
+9
-1
vllm/compilation/passes/utility/fix_functionalization.py
+9
-1
@@ -37,6 +37,14 @@ class FixFunctionalizationPass(VllmInductorPass):
37
38
self.nodes_to_remove: list[torch.fx.Node] = []
39
count = 0
40
+
41
+ rope_targets = [torch.ops._C.rotary_embedding.default]
42
+
43
+ if hasattr(torch.ops.vllm, "rocm_aiter_triton_rotary_embedding"):
44
+ rope_targets.append(
45
+ torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default
46
+ )
47
+
48
for node in graph.nodes:
49
if not is_func(node, auto_functionalized):
50
continue # Avoid deep if-elif nesting
@@ -44,7 +52,7 @@ class FixFunctionalizationPass(VllmInductorPass):
52
kwargs = node.kwargs
53
at_target = node.args[0]
54
47
- if at_target == torch.ops._C.rotary_embedding.default:
55
+ if at_target in rope_targets:
56
query = kwargs["query"]
57
key = kwargs["key"]
58
getitem_nodes = self.getitem_users(node)