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 1, 23:45 UTC
[Inductor] Fast-path Inductor fallback for vllm::*/vllm_aiter::* custom ops (#42129)
Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Oxana Korzh committed
Jun 3, 2026 at 23:03 UTC
b4b4aaa70ee992e96504c27fad4535d085a5448f
2 files changed
+348
tests/compile/test_inductor_fallback_allow_list_patch.py
new
+250
@@ -0,0 +1,250 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the Inductor FALLBACK_ALLOW_LIST patch in env_override.py.
4
+
5
+The patch wraps ``torch._inductor.lowering.FALLBACK_ALLOW_LIST`` in a thin
6
+proxy that auto-allows any custom op in the ``vllm::`` or ``vllm_aiter::``
7
+namespaces. This routes those ops through Inductor's fast-path
8
+``make_fallback(target, warn=False, override_decomp=True)`` and avoids the
9
+expensive ``error.operator_str(target, args, kwargs)`` formatting that
10
+recursively stringifies every input ``TensorBox``.
11
+
12
+The slow path is what made ``torch.compile`` effectively hang on Kimi-K2.6
13
+TP=8 (deep MoE/TP IR provenance trees). These tests cover both the proxy's
14
+semantics in isolation and the membership-check fast-path that Inductor's
15
+``GraphLowering.call_function`` actually performs, so we can validate the
16
+optimization without needing a full GPU compile.
17
+"""
18
+
19
+import time
20
+
21
+import pytest
22
+
23
+from vllm.env_override import (
24
+ _patch_inductor_fallback_allow_list,
25
+ _VllmFallbackAllowList,
26
+)
27
+
28
+
29
+class TestVllmFallbackAllowListProxy:
30
+ """Unit tests for the membership-proxy semantics."""
31
+
32
+ def test_vllm_namespace_auto_allowed(self):
33
+ proxy = _VllmFallbackAllowList(set())
34
+ assert "vllm::all_reduce" in proxy
35
+ assert "vllm::fused_add_rms_norm" in proxy
36
+ assert "vllm::all_reduce.default" in proxy
37
+
38
+ def test_vllm_aiter_namespace_auto_allowed(self):
39
+ proxy = _VllmFallbackAllowList(set())
40
+ assert "vllm_aiter::fused_add_rms_norm" in proxy
41
+ assert "vllm_aiter::rocm_aiter_fused_moe" in proxy
42
+
43
+ def test_unknown_namespace_falls_through(self):
44
+ proxy = _VllmFallbackAllowList({"torchvision::roi_align"})
45
+ assert "torchvision::roi_align" in proxy
46
+ assert "made_up_ns::nonexistent_op" not in proxy
47
+
48
+ def test_non_string_falls_through_to_inner(self):
49
+ sentinel = object()
50
+ inner = {sentinel}
51
+ proxy = _VllmFallbackAllowList(inner)
52
+ assert sentinel in proxy
53
+ assert object() not in proxy
54
+
55
+ def test_prefix_only_match_not_substring(self):
56
+ proxy = _VllmFallbackAllowList(set())
57
+ assert "not_vllm::something" not in proxy
58
+ assert " vllm::space_prefixed" not in proxy
59
+
60
+ def test_standard_entries_preserved(self):
61
+ base = {"torchvision::roi_align", "aten::index_add"}
62
+ proxy = _VllmFallbackAllowList(base)
63
+ assert "torchvision::roi_align" in proxy
64
+ assert "aten::index_add" in proxy
65
+ assert "aten::__not_present__" not in proxy
66
+
67
+ def test_add_and_discard_delegate_to_inner(self):
68
+ inner: set[str] = set()
69
+ proxy = _VllmFallbackAllowList(inner)
70
+ proxy.add("custom::op")
71
+ assert "custom::op" in inner
72
+ proxy.discard("custom::op")
73
+ assert "custom::op" not in inner
74
+
75
+ def test_iter_len_repr(self):
76
+ base = {"torchvision::roi_align", "aten::index_add"}
77
+ proxy = _VllmFallbackAllowList(base)
78
+ assert set(iter(proxy)) == base
79
+ assert len(proxy) == len(base)
80
+ assert "torchvision::roi_align" in repr(proxy)
81
+
82
+ def test_getattr_delegates_to_inner(self):
83
+ class _Inner:
84
+ sentinel = "i_am_inner"
85
+
86
+ def some_method(self):
87
+ return 42
88
+
89
+ inner = _Inner()
90
+ proxy = _VllmFallbackAllowList(inner)
91
+ assert proxy.sentinel == "i_am_inner"
92
+ assert proxy.some_method() == 42
93
+
94
+ def test_sentinel_attribute(self):
95
+ proxy = _VllmFallbackAllowList(set())
96
+ assert proxy._vllm_patched is True
97
+
98
+
99
+class TestPatchApplication:
100
+ """Integration tests verifying the patch reaches ``torch._inductor``."""
101
+
102
+ def test_patch_applied_to_lowering(self):
103
+ import torch._inductor.lowering as _lowering
104
+
105
+ assert getattr(_lowering.FALLBACK_ALLOW_LIST, "_vllm_patched", False), (
106
+ "env_override._patch_inductor_fallback_allow_list did not run"
107
+ )
108
+
109
+ def test_graph_module_local_binding_rebound(self):
110
+ # ``torch/_inductor/graph.py`` does:
111
+ # from torch._inductor.lowering import FALLBACK_ALLOW_LIST
112
+ # so the patch has to overwrite the graph module's local binding too,
113
+ # otherwise the fast-path check in GraphLowering.call_function still
114
+ # sees the original (unwrapped) OrderedSet.
115
+ import torch._inductor.graph as _graph
116
+ import torch._inductor.lowering as _lowering
117
+
118
+ if not hasattr(_graph, "FALLBACK_ALLOW_LIST"):
119
+ pytest.skip(
120
+ "torch._inductor.graph no longer imports FALLBACK_ALLOW_LIST "
121
+ "as a module-level symbol; nothing to rebind."
122
+ )
123
+
124
+ assert _graph.FALLBACK_ALLOW_LIST is _lowering.FALLBACK_ALLOW_LIST
125
+
126
+ def test_patch_is_idempotent(self):
127
+ import torch._inductor.lowering as _lowering
128
+
129
+ first = _lowering.FALLBACK_ALLOW_LIST
130
+ _patch_inductor_fallback_allow_list()
131
+ _patch_inductor_fallback_allow_list()
132
+ assert _lowering.FALLBACK_ALLOW_LIST is first
133
+
134
+ def test_real_vllm_ops_in_real_allow_list(self):
135
+ # End-to-end membership check using the live (already-patched) object.
136
+ import torch._inductor.lowering as _lowering
137
+
138
+ allow_list = _lowering.FALLBACK_ALLOW_LIST
139
+ assert "vllm::all_reduce" in allow_list
140
+ assert "vllm::fused_add_rms_norm" in allow_list
141
+ assert "vllm_aiter::fused_add_rms_norm" in allow_list
142
+
143
+
144
+class TestInductorFallbackFastPath:
145
+ """Emulates ``GraphLowering.call_function``'s FALLBACK_ALLOW_LIST check.
146
+
147
+ The relevant snippet in ``torch/_inductor/graph.py`` is roughly::
148
+
149
+ base_name = target.name()
150
+ if base_name not in FALLBACK_ALLOW_LIST:
151
+ log.info(
152
+ "Creating implicit fallback for:\\n%s",
153
+ error.operator_str(target, args, kwargs),
154
+ )
155
+ out = make_fallback(target, ...)
156
+
157
+ On a deep MoE/TP graph (Kimi-K2.6 at TP=4/8) ``operator_str`` recurses
158
+ through every input ``TensorBox.__str__`` and ends up taking many minutes
159
+ of CPU per encountered op. The patch ensures the membership test
160
+ short-circuits for ``vllm::*``/``vllm_aiter::*`` ops so the slow path is
161
+ never entered. These tests pin that behaviour without needing a real
162
+ GPU compile.
163
+ """
164
+
165
+ def _simulate_graph_lowering(self, target_names: list[str]):
166
+ """Returns the set of target names that would have hit the slow
167
+ operator_str() path under the patched FALLBACK_ALLOW_LIST.
168
+ """
169
+ import torch._inductor.lowering as _lowering
170
+
171
+ allow_list = _lowering.FALLBACK_ALLOW_LIST
172
+ slow_path_hits: list[str] = []
173
+ for name in target_names:
174
+ if name not in allow_list:
175
+ slow_path_hits.append(name)
176
+ return slow_path_hits
177
+
178
+ def test_vllm_ops_skip_slow_path(self):
179
+ slow = self._simulate_graph_lowering(
180
+ [
181
+ "vllm::all_reduce",
182
+ "vllm::fused_add_rms_norm",
183
+ "vllm_aiter::rocm_aiter_fused_moe",
184
+ "vllm_aiter::asm_moe",
185
+ ]
186
+ )
187
+ assert slow == [], (
188
+ "Patched FALLBACK_ALLOW_LIST must short-circuit for all "
189
+ f"vllm::*/vllm_aiter::* ops; got slow-path hits: {slow}"
190
+ )
191
+
192
+ def test_non_vllm_ops_still_hit_slow_path(self):
193
+ # Without the patch this is also what would happen; with the patch
194
+ # the behaviour for non-vllm namespaces must be unchanged.
195
+ slow = self._simulate_graph_lowering(
196
+ ["my_user_ns::custom_op", "fancy_ns::something_else"]
197
+ )
198
+ assert "my_user_ns::custom_op" in slow
199
+ assert "fancy_ns::something_else" in slow
200
+
201
+ def test_kimi_k2_6_style_op_stream(self):
202
+ """Emulates one decoder layer's worth of fallback hits.
203
+
204
+ Kimi-K2.6 at TP=4 lowers a stream of ``vllm::all_reduce`` +
205
+ ``vllm_aiter::fused_add_rms_norm`` calls (one per residual block)
206
+ plus a handful of fused-MoE ops. Pre-patch every one of these would
207
+ invoke ``operator_str`` and stringify a hundreds-deep IR provenance
208
+ tree; post-patch they must all short-circuit.
209
+ """
210
+ n_layers = 64 # Kimi-K2.6 has ~64 decoder layers per replica
211
+ op_stream: list[str] = []
212
+ for _ in range(n_layers):
213
+ op_stream.extend(
214
+ [
215
+ "vllm::all_reduce",
216
+ "vllm_aiter::fused_add_rms_norm",
217
+ "vllm_aiter::rocm_aiter_fused_moe",
218
+ ]
219
+ )
220
+
221
+ start = time.perf_counter()
222
+ slow = self._simulate_graph_lowering(op_stream)
223
+ elapsed_s = time.perf_counter() - start
224
+
225
+ assert slow == [], (
226
+ f"Expected all {len(op_stream)} vllm/vllm_aiter ops to take "
227
+ f"the fast path; got {len(slow)} slow-path hits."
228
+ )
229
+ # ``__contains__`` is O(1) per call, so a Kimi-sized stream should
230
+ # complete in well under a second even on a slow runner. The
231
+ # pre-patch slow path took many minutes per op on Kimi-K2.6 TP=8.
232
+ assert elapsed_s < 1.0, (
233
+ f"FALLBACK_ALLOW_LIST membership check is unexpectedly slow: "
234
+ f"{elapsed_s:.3f}s for {len(op_stream)} ops"
235
+ )
236
+
237
+ def test_inner_set_membership_still_works_for_standard_ops(self):
238
+ """The patch must not break Inductor's existing fallback decisions
239
+ for non-vllm ops such as ``torchvision::roi_align``."""
240
+ import torch._inductor.lowering as _lowering
241
+
242
+ allow_list = _lowering.FALLBACK_ALLOW_LIST
243
+ # ``torchvision::roi_align`` has been a member of the upstream
244
+ # FALLBACK_ALLOW_LIST since the original Inductor implementation.
245
+ # If the proxy ever broke pass-through, this would regress.
246
+ if "torchvision::roi_align" not in allow_list:
247
+ pytest.skip(
248
+ "Upstream FALLBACK_ALLOW_LIST no longer ships "
249
+ "torchvision::roi_align; nothing to verify."
250
+ )
vllm/env_override.py
+98
@@ -758,3 +758,101 @@ def _patch_cpp_indirect_assert_if_needed():
758
759
760
_patch_cpp_indirect_assert_if_needed()
761
+
762
+# ============================================================
763
+# Inductor FALLBACK_ALLOW_LIST fast-path for vllm::*/vllm_aiter::* ops
764
+# ============================================================
765
+# When Inductor encounters a custom op without a registered lowering or
766
+# decomposition (e.g. vllm::all_reduce, vllm_aiter::fused_add_rms_norm) it
767
+# correctly creates an implicit fallback that calls into the eager Python
768
+# impl. However, unless `base_name` (e.g. "vllm::all_reduce") is in
769
+# torch._inductor.lowering.FALLBACK_ALLOW_LIST, GraphLowering.call_function
770
+# (torch/_inductor/graph.py:~1283) takes the slow path that emits
771
+# log.info("Creating implicit fallback for:\n%s",
772
+# error.operator_str(target, args, kwargs))
773
+# `operator_str` eagerly recurses through __str__ on every input TensorBox;
774
+# for deep MoE/TP graphs (e.g. Kimi-K2.6 at TP=8) the IR provenance tree
775
+# behind a TP all-reduce input or a residual-fed RMSNorm input is hundreds
776
+# of layers deep, and stringifying it consumes many minutes of CPU per call,
777
+# effectively hanging compilation.
778
+#
779
+# Patching FALLBACK_ALLOW_LIST membership to also match any "vllm::*" or
780
+# "vllm_aiter::*" base_name routes our custom ops through the fast path
781
+# `make_fallback(target, warn=False, override_decomp=True)` instead. This
782
+# preserves all downstream behaviour (allreduce_rms_fusion still pattern-
783
+# matches them, partitioning still works, fallback semantics identical) but
784
+# skips the expensive log formatting on the FIRST encounter of each op.
785
+#
786
+# We wrap the OrderedSet in a thin proxy that:
787
+# - Returns True from __contains__ for any vllm::*/vllm_aiter::* op
788
+# - Otherwise delegates to the underlying set (preserving membership of
789
+# the standard entries like "torchvision::roi_align", "aten::index_add")
790
+# - Forwards add()/__iter__()/__len__()/etc. so other Inductor code paths
791
+# that mutate or iterate the set keep working.
792
+
793
+_VLLM_FALLBACK_NAMESPACE_PREFIXES = ("vllm::", "vllm_aiter::")
794
+
795
+
796
+class _VllmFallbackAllowList:
797
+ """Membership proxy that auto-allows vllm::*/vllm_aiter::* base_names."""
798
+
799
+ _vllm_patched = True
800
+
801
+ def __init__(self, inner):
802
+ self._inner = inner
803
+
804
+ def __contains__(self, item):
805
+ if isinstance(item, str) and item.startswith(_VLLM_FALLBACK_NAMESPACE_PREFIXES):
806
+ return True
807
+ return item in self._inner
808
+
809
+ def add(self, item):
810
+ self._inner.add(item)
811
+
812
+ def discard(self, item):
813
+ self._inner.discard(item)
814
+
815
+ def __iter__(self):
816
+ return iter(self._inner)
817
+
818
+ def __len__(self):
819
+ return len(self._inner)
820
+
821
+ def __repr__(self):
822
+ return f"_VllmFallbackAllowList({self._inner!r})"
823
+
824
+ def __getattr__(self, name):
825
+ return getattr(self._inner, name)
826
+
827
+
828
+def _patch_inductor_fallback_allow_list() -> None:
829
+ """Wrap torch._inductor.lowering.FALLBACK_ALLOW_LIST so any custom op in
830
+ the ``vllm::`` or ``vllm_aiter::`` namespaces is treated as a member.
831
+
832
+ Idempotent: a sentinel attribute on the proxy prevents re-wrapping.
833
+ """
834
+ try:
835
+ from torch._inductor import lowering as _lowering
836
+ except ImportError:
837
+ return
838
+
839
+ base = getattr(_lowering, "FALLBACK_ALLOW_LIST", None)
840
+ if base is None or getattr(base, "_vllm_patched", False):
841
+ return
842
+
843
+ _lowering.FALLBACK_ALLOW_LIST = _VllmFallbackAllowList(base)
844
+
845
+ # torch/_inductor/graph.py imports the symbol at module load time:
846
+ # from torch._inductor.lowering import FALLBACK_ALLOW_LIST
847
+ # so we also need to overwrite the local binding in the graph module if
848
+ # it has already been imported.
849
+ try:
850
+ from torch._inductor import graph as _graph
851
+
852
+ if hasattr(_graph, "FALLBACK_ALLOW_LIST"):
853
+ _graph.FALLBACK_ALLOW_LIST = _lowering.FALLBACK_ALLOW_LIST
854
+ except ImportError:
855
+ pass
856
+
857
+
858
+_patch_inductor_fallback_allow_list()