diff --git a/tests/compile/test_inductor_fallback_allow_list_patch.py b/tests/compile/test_inductor_fallback_allow_list_patch.py new file mode 100644 index 00000000000..29fe9962e34 --- /dev/null +++ b/tests/compile/test_inductor_fallback_allow_list_patch.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the Inductor FALLBACK_ALLOW_LIST patch in env_override.py. + +The patch wraps ``torch._inductor.lowering.FALLBACK_ALLOW_LIST`` in a thin +proxy that auto-allows any custom op in the ``vllm::`` or ``vllm_aiter::`` +namespaces. This routes those ops through Inductor's fast-path +``make_fallback(target, warn=False, override_decomp=True)`` and avoids the +expensive ``error.operator_str(target, args, kwargs)`` formatting that +recursively stringifies every input ``TensorBox``. + +The slow path is what made ``torch.compile`` effectively hang on Kimi-K2.6 +TP=8 (deep MoE/TP IR provenance trees). These tests cover both the proxy's +semantics in isolation and the membership-check fast-path that Inductor's +``GraphLowering.call_function`` actually performs, so we can validate the +optimization without needing a full GPU compile. +""" + +import time + +import pytest + +from vllm.env_override import ( + _patch_inductor_fallback_allow_list, + _VllmFallbackAllowList, +) + + +class TestVllmFallbackAllowListProxy: + """Unit tests for the membership-proxy semantics.""" + + def test_vllm_namespace_auto_allowed(self): + proxy = _VllmFallbackAllowList(set()) + assert "vllm::all_reduce" in proxy + assert "vllm::fused_add_rms_norm" in proxy + assert "vllm::all_reduce.default" in proxy + + def test_vllm_aiter_namespace_auto_allowed(self): + proxy = _VllmFallbackAllowList(set()) + assert "vllm_aiter::fused_add_rms_norm" in proxy + assert "vllm_aiter::rocm_aiter_fused_moe" in proxy + + def test_unknown_namespace_falls_through(self): + proxy = _VllmFallbackAllowList({"torchvision::roi_align"}) + assert "torchvision::roi_align" in proxy + assert "made_up_ns::nonexistent_op" not in proxy + + def test_non_string_falls_through_to_inner(self): + sentinel = object() + inner = {sentinel} + proxy = _VllmFallbackAllowList(inner) + assert sentinel in proxy + assert object() not in proxy + + def test_prefix_only_match_not_substring(self): + proxy = _VllmFallbackAllowList(set()) + assert "not_vllm::something" not in proxy + assert " vllm::space_prefixed" not in proxy + + def test_standard_entries_preserved(self): + base = {"torchvision::roi_align", "aten::index_add"} + proxy = _VllmFallbackAllowList(base) + assert "torchvision::roi_align" in proxy + assert "aten::index_add" in proxy + assert "aten::__not_present__" not in proxy + + def test_add_and_discard_delegate_to_inner(self): + inner: set[str] = set() + proxy = _VllmFallbackAllowList(inner) + proxy.add("custom::op") + assert "custom::op" in inner + proxy.discard("custom::op") + assert "custom::op" not in inner + + def test_iter_len_repr(self): + base = {"torchvision::roi_align", "aten::index_add"} + proxy = _VllmFallbackAllowList(base) + assert set(iter(proxy)) == base + assert len(proxy) == len(base) + assert "torchvision::roi_align" in repr(proxy) + + def test_getattr_delegates_to_inner(self): + class _Inner: + sentinel = "i_am_inner" + + def some_method(self): + return 42 + + inner = _Inner() + proxy = _VllmFallbackAllowList(inner) + assert proxy.sentinel == "i_am_inner" + assert proxy.some_method() == 42 + + def test_sentinel_attribute(self): + proxy = _VllmFallbackAllowList(set()) + assert proxy._vllm_patched is True + + +class TestPatchApplication: + """Integration tests verifying the patch reaches ``torch._inductor``.""" + + def test_patch_applied_to_lowering(self): + import torch._inductor.lowering as _lowering + + assert getattr(_lowering.FALLBACK_ALLOW_LIST, "_vllm_patched", False), ( + "env_override._patch_inductor_fallback_allow_list did not run" + ) + + def test_graph_module_local_binding_rebound(self): + # ``torch/_inductor/graph.py`` does: + # from torch._inductor.lowering import FALLBACK_ALLOW_LIST + # so the patch has to overwrite the graph module's local binding too, + # otherwise the fast-path check in GraphLowering.call_function still + # sees the original (unwrapped) OrderedSet. + import torch._inductor.graph as _graph + import torch._inductor.lowering as _lowering + + if not hasattr(_graph, "FALLBACK_ALLOW_LIST"): + pytest.skip( + "torch._inductor.graph no longer imports FALLBACK_ALLOW_LIST " + "as a module-level symbol; nothing to rebind." + ) + + assert _graph.FALLBACK_ALLOW_LIST is _lowering.FALLBACK_ALLOW_LIST + + def test_patch_is_idempotent(self): + import torch._inductor.lowering as _lowering + + first = _lowering.FALLBACK_ALLOW_LIST + _patch_inductor_fallback_allow_list() + _patch_inductor_fallback_allow_list() + assert _lowering.FALLBACK_ALLOW_LIST is first + + def test_real_vllm_ops_in_real_allow_list(self): + # End-to-end membership check using the live (already-patched) object. + import torch._inductor.lowering as _lowering + + allow_list = _lowering.FALLBACK_ALLOW_LIST + assert "vllm::all_reduce" in allow_list + assert "vllm::fused_add_rms_norm" in allow_list + assert "vllm_aiter::fused_add_rms_norm" in allow_list + + +class TestInductorFallbackFastPath: + """Emulates ``GraphLowering.call_function``'s FALLBACK_ALLOW_LIST check. + + The relevant snippet in ``torch/_inductor/graph.py`` is roughly:: + + base_name = target.name() + if base_name not in FALLBACK_ALLOW_LIST: + log.info( + "Creating implicit fallback for:\\n%s", + error.operator_str(target, args, kwargs), + ) + out = make_fallback(target, ...) + + On a deep MoE/TP graph (Kimi-K2.6 at TP=4/8) ``operator_str`` recurses + through every input ``TensorBox.__str__`` and ends up taking many minutes + of CPU per encountered op. The patch ensures the membership test + short-circuits for ``vllm::*``/``vllm_aiter::*`` ops so the slow path is + never entered. These tests pin that behaviour without needing a real + GPU compile. + """ + + def _simulate_graph_lowering(self, target_names: list[str]): + """Returns the set of target names that would have hit the slow + operator_str() path under the patched FALLBACK_ALLOW_LIST. + """ + import torch._inductor.lowering as _lowering + + allow_list = _lowering.FALLBACK_ALLOW_LIST + slow_path_hits: list[str] = [] + for name in target_names: + if name not in allow_list: + slow_path_hits.append(name) + return slow_path_hits + + def test_vllm_ops_skip_slow_path(self): + slow = self._simulate_graph_lowering( + [ + "vllm::all_reduce", + "vllm::fused_add_rms_norm", + "vllm_aiter::rocm_aiter_fused_moe", + "vllm_aiter::asm_moe", + ] + ) + assert slow == [], ( + "Patched FALLBACK_ALLOW_LIST must short-circuit for all " + f"vllm::*/vllm_aiter::* ops; got slow-path hits: {slow}" + ) + + def test_non_vllm_ops_still_hit_slow_path(self): + # Without the patch this is also what would happen; with the patch + # the behaviour for non-vllm namespaces must be unchanged. + slow = self._simulate_graph_lowering( + ["my_user_ns::custom_op", "fancy_ns::something_else"] + ) + assert "my_user_ns::custom_op" in slow + assert "fancy_ns::something_else" in slow + + def test_kimi_k2_6_style_op_stream(self): + """Emulates one decoder layer's worth of fallback hits. + + Kimi-K2.6 at TP=4 lowers a stream of ``vllm::all_reduce`` + + ``vllm_aiter::fused_add_rms_norm`` calls (one per residual block) + plus a handful of fused-MoE ops. Pre-patch every one of these would + invoke ``operator_str`` and stringify a hundreds-deep IR provenance + tree; post-patch they must all short-circuit. + """ + n_layers = 64 # Kimi-K2.6 has ~64 decoder layers per replica + op_stream: list[str] = [] + for _ in range(n_layers): + op_stream.extend( + [ + "vllm::all_reduce", + "vllm_aiter::fused_add_rms_norm", + "vllm_aiter::rocm_aiter_fused_moe", + ] + ) + + start = time.perf_counter() + slow = self._simulate_graph_lowering(op_stream) + elapsed_s = time.perf_counter() - start + + assert slow == [], ( + f"Expected all {len(op_stream)} vllm/vllm_aiter ops to take " + f"the fast path; got {len(slow)} slow-path hits." + ) + # ``__contains__`` is O(1) per call, so a Kimi-sized stream should + # complete in well under a second even on a slow runner. The + # pre-patch slow path took many minutes per op on Kimi-K2.6 TP=8. + assert elapsed_s < 1.0, ( + f"FALLBACK_ALLOW_LIST membership check is unexpectedly slow: " + f"{elapsed_s:.3f}s for {len(op_stream)} ops" + ) + + def test_inner_set_membership_still_works_for_standard_ops(self): + """The patch must not break Inductor's existing fallback decisions + for non-vllm ops such as ``torchvision::roi_align``.""" + import torch._inductor.lowering as _lowering + + allow_list = _lowering.FALLBACK_ALLOW_LIST + # ``torchvision::roi_align`` has been a member of the upstream + # FALLBACK_ALLOW_LIST since the original Inductor implementation. + # If the proxy ever broke pass-through, this would regress. + if "torchvision::roi_align" not in allow_list: + pytest.skip( + "Upstream FALLBACK_ALLOW_LIST no longer ships " + "torchvision::roi_align; nothing to verify." + ) diff --git a/vllm/env_override.py b/vllm/env_override.py index 78270c2bee3..fe084ce5149 100644 --- a/vllm/env_override.py +++ b/vllm/env_override.py @@ -758,3 +758,101 @@ def _patch_cpp_indirect_assert_if_needed(): _patch_cpp_indirect_assert_if_needed() + +# ============================================================ +# Inductor FALLBACK_ALLOW_LIST fast-path for vllm::*/vllm_aiter::* ops +# ============================================================ +# When Inductor encounters a custom op without a registered lowering or +# decomposition (e.g. vllm::all_reduce, vllm_aiter::fused_add_rms_norm) it +# correctly creates an implicit fallback that calls into the eager Python +# impl. However, unless `base_name` (e.g. "vllm::all_reduce") is in +# torch._inductor.lowering.FALLBACK_ALLOW_LIST, GraphLowering.call_function +# (torch/_inductor/graph.py:~1283) takes the slow path that emits +# log.info("Creating implicit fallback for:\n%s", +# error.operator_str(target, args, kwargs)) +# `operator_str` eagerly recurses through __str__ on every input TensorBox; +# for deep MoE/TP graphs (e.g. Kimi-K2.6 at TP=8) the IR provenance tree +# behind a TP all-reduce input or a residual-fed RMSNorm input is hundreds +# of layers deep, and stringifying it consumes many minutes of CPU per call, +# effectively hanging compilation. +# +# Patching FALLBACK_ALLOW_LIST membership to also match any "vllm::*" or +# "vllm_aiter::*" base_name routes our custom ops through the fast path +# `make_fallback(target, warn=False, override_decomp=True)` instead. This +# preserves all downstream behaviour (allreduce_rms_fusion still pattern- +# matches them, partitioning still works, fallback semantics identical) but +# skips the expensive log formatting on the FIRST encounter of each op. +# +# We wrap the OrderedSet in a thin proxy that: +# - Returns True from __contains__ for any vllm::*/vllm_aiter::* op +# - Otherwise delegates to the underlying set (preserving membership of +# the standard entries like "torchvision::roi_align", "aten::index_add") +# - Forwards add()/__iter__()/__len__()/etc. so other Inductor code paths +# that mutate or iterate the set keep working. + +_VLLM_FALLBACK_NAMESPACE_PREFIXES = ("vllm::", "vllm_aiter::") + + +class _VllmFallbackAllowList: + """Membership proxy that auto-allows vllm::*/vllm_aiter::* base_names.""" + + _vllm_patched = True + + def __init__(self, inner): + self._inner = inner + + def __contains__(self, item): + if isinstance(item, str) and item.startswith(_VLLM_FALLBACK_NAMESPACE_PREFIXES): + return True + return item in self._inner + + def add(self, item): + self._inner.add(item) + + def discard(self, item): + self._inner.discard(item) + + def __iter__(self): + return iter(self._inner) + + def __len__(self): + return len(self._inner) + + def __repr__(self): + return f"_VllmFallbackAllowList({self._inner!r})" + + def __getattr__(self, name): + return getattr(self._inner, name) + + +def _patch_inductor_fallback_allow_list() -> None: + """Wrap torch._inductor.lowering.FALLBACK_ALLOW_LIST so any custom op in + the ``vllm::`` or ``vllm_aiter::`` namespaces is treated as a member. + + Idempotent: a sentinel attribute on the proxy prevents re-wrapping. + """ + try: + from torch._inductor import lowering as _lowering + except ImportError: + return + + base = getattr(_lowering, "FALLBACK_ALLOW_LIST", None) + if base is None or getattr(base, "_vllm_patched", False): + return + + _lowering.FALLBACK_ALLOW_LIST = _VllmFallbackAllowList(base) + + # torch/_inductor/graph.py imports the symbol at module load time: + # from torch._inductor.lowering import FALLBACK_ALLOW_LIST + # so we also need to overwrite the local binding in the graph module if + # it has already been imported. + try: + from torch._inductor import graph as _graph + + if hasattr(_graph, "FALLBACK_ALLOW_LIST"): + _graph.FALLBACK_ALLOW_LIST = _lowering.FALLBACK_ALLOW_LIST + except ImportError: + pass + + +_patch_inductor_fallback_allow_list()