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 12, 03:30 UTC
[Bugfix] Defunctionalize TRTLLM AR+Norm op for avoiding extra clone kernel before it (#29631)
Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
elvischenv committed
Dec 3, 2025 at 13:15 UTC
c719c40540a85c1e6aeee9af20f29db581da27f0
2 files changed
+14
-2
vllm/compilation/fix_functionalization.py
+12
@@ -103,6 +103,18 @@ class FixFunctionalizationPass(VllmInductorPass):
103
]:
104
mutated_args = {1: "result"}
105
self.defunctionalize(graph, node, mutated_args)
106
+ elif (
107
+ at_target
108
+ == torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default
109
+ ):
110
+ mutated_args = {
111
+ 1: "allreduce_in",
112
+ 2: "residual",
113
+ 3: "norm_out",
114
+ 4: "quant_out",
115
+ 5: "scale_out",
116
+ }
117
+ self.defunctionalize(graph, node, mutated_args)
118
# For some reason we need to specify the args for both
119
# silu_and_mul and silu_and_mul_quant. The kwargs
120
# pathway gets the wrong answer.
vllm/compilation/fx_utils.py
+2
-2
@@ -75,8 +75,8 @@ def find_op_nodes(
75
return
76
77
assert isinstance(op, OpOverload)
78
- if not op._schema.is_mutable:
79
- yield from graph.find_nodes(op="call_function", target=op)
78
+
79
+ yield from graph.find_nodes(op="call_function", target=op)
80
81
for n in graph.find_nodes(op="call_function", target=auto_functionalized):
82
if n.args[0] == op: