diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index 7465ea442c..2d4ef13e33 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -70,6 +70,37 @@ def test_get_raw_stream_patch(): assert get_raw_stream is _cuda_getCurrentRawStream +# Inductor can initialize CUDA even for CPU inputs. +@pytest.mark.forked +@pytest.mark.parametrize("use_v2", [False, True]) +def test_e8m0_custom_op_fullgraph(use_v2, use_fresh_inductor_cache): + """E8M0 inputs must not prevent mutable custom-op decomposition.""" + with torch.library._scoped_library("test_e8m0", "FRAGMENT") as lib: + lib.define("scale_(Tensor(a!) out, Tensor scale) -> ()") + + @torch.library.impl(lib, "scale_", "CPU") + def scale_impl(out, scale): + out.mul_(scale.float()) + + @torch.library.register_fake("test_e8m0::scale_", lib=lib) + def scale_fake(out, scale): + return None + + def forward(x, scale): + out = x.clone() + torch.ops.test_e8m0.scale_(out, scale) + return out + + x = torch.randn(4, 4) + scale = torch.full((4,), 128, dtype=torch.uint8).view(torch.float8_e8m0fnu) + compiled = torch.compile( + forward, + fullgraph=True, + options={"enable_auto_functionalized_v2": use_v2}, + ) + torch.testing.assert_close(compiled(x, scale), x * 2) + + def test_copy_pass(): vllm_config = VllmConfig() inductor_pass = FixFunctionalizationPass(vllm_config) diff --git a/vllm/env_override.py b/vllm/env_override.py index 87501eb469..bfcdfa7af2 100644 --- a/vllm/env_override.py +++ b/vllm/env_override.py @@ -916,6 +916,58 @@ def _patch_inductor_fallback_allow_list() -> None: _patch_inductor_fallback_allow_list() + +def _apply_inductor_pattern_matcher_patch() -> None: + """Allow custom ops and functionalization wrappers with unsupported dtypes.""" + from torch._inductor import pattern_matcher + from torch._inductor.lowering import fallback_node_due_to_unsupported_type + + def fallback_for_builtin(node, allow_cpu_inputs=True): + return ( + isinstance(node.target, torch._ops.OpOverload) + and torch._library.utils.is_builtin(node.target) + and fallback_node_due_to_unsupported_type(node, allow_cpu_inputs) + ) + + pattern_matcher.fallback_node_due_to_unsupported_type = fallback_for_builtin + + +def _patch_inductor_pattern_matcher() -> None: + """Apply the backport when Inductor imports its pattern matcher.""" + import sys + + target_name = "torch._inductor.pattern_matcher" + if target_name in sys.modules: + _apply_inductor_pattern_matcher_patch() + return + + import importlib.abc + + class _PatternMatcherPatchFinder(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname != target_name: + return None + sys.meta_path.remove(self) + spec = importlib.util.find_spec(fullname) + if spec is None or spec.loader is None: + return None + original_exec = spec.loader.exec_module + + def _exec_then_patch(module): + original_exec(module) + _apply_inductor_pattern_matcher_patch() + + spec.loader.exec_module = _exec_then_patch # type: ignore[method-assign] + return spec + + sys.meta_path.insert(0, _PatternMatcherPatchFinder()) + + +# Remove once the minimum supported torch includes +# https://github.com/pytorch/pytorch/pull/196013. +if not is_torch_equal_or_newer("2.16.0.dev"): + _patch_inductor_pattern_matcher() + # ============================================================ # Triton Autotuner determinism # ============================================================