[Bugfix] Backport Inductor custom-op pattern matching fix (#58189)
Signed-off-by: mgoin <mgoin64@gmail.com>
Misha Goin committed
Sep 22, 2026 at 17:26 UTC
01e15bc94cd29cd1fcbbab484e96acb5b4f6bb14
2 files changed
+83
tests/compile/test_config.py
+31
@@ -70,6 +70,37 @@ def test_get_raw_stream_patch():
70
assert get_raw_stream is _cuda_getCurrentRawStream
71
72
73
+# Inductor can initialize CUDA even for CPU inputs.
74
+@pytest.mark.forked
75
+@pytest.mark.parametrize("use_v2", [False, True])
76
+def test_e8m0_custom_op_fullgraph(use_v2, use_fresh_inductor_cache):
77
+ """E8M0 inputs must not prevent mutable custom-op decomposition."""
78
+ with torch.library._scoped_library("test_e8m0", "FRAGMENT") as lib:
79
+ lib.define("scale_(Tensor(a!) out, Tensor scale) -> ()")
80
+
81
+ @torch.library.impl(lib, "scale_", "CPU")
82
+ def scale_impl(out, scale):
83
+ out.mul_(scale.float())
84
+
85
+ @torch.library.register_fake("test_e8m0::scale_", lib=lib)
86
+ def scale_fake(out, scale):
87
+ return None
88
+
89
+ def forward(x, scale):
90
+ out = x.clone()
91
+ torch.ops.test_e8m0.scale_(out, scale)
92
+ return out
93
+
94
+ x = torch.randn(4, 4)
95
+ scale = torch.full((4,), 128, dtype=torch.uint8).view(torch.float8_e8m0fnu)
96
+ compiled = torch.compile(
97
+ forward,
98
+ fullgraph=True,
99
+ options={"enable_auto_functionalized_v2": use_v2},
100
+ )
101
+ torch.testing.assert_close(compiled(x, scale), x * 2)
102
+
103
+
104
def test_copy_pass():
105
vllm_config = VllmConfig()
106
inductor_pass = FixFunctionalizationPass(vllm_config)
vllm/env_override.py
+52
@@ -916,6 +916,58 @@ def _patch_inductor_fallback_allow_list() -> None:
916
917
_patch_inductor_fallback_allow_list()
918
919
+
920
+def _apply_inductor_pattern_matcher_patch() -> None:
921
+ """Allow custom ops and functionalization wrappers with unsupported dtypes."""
922
+ from torch._inductor import pattern_matcher
923
+ from torch._inductor.lowering import fallback_node_due_to_unsupported_type
924
+
925
+ def fallback_for_builtin(node, allow_cpu_inputs=True):
926
+ return (
927
+ isinstance(node.target, torch._ops.OpOverload)
928
+ and torch._library.utils.is_builtin(node.target)
929
+ and fallback_node_due_to_unsupported_type(node, allow_cpu_inputs)
930
+ )
931
+
932
+ pattern_matcher.fallback_node_due_to_unsupported_type = fallback_for_builtin
933
+
934
+
935
+def _patch_inductor_pattern_matcher() -> None:
936
+ """Apply the backport when Inductor imports its pattern matcher."""
937
+ import sys
938
+
939
+ target_name = "torch._inductor.pattern_matcher"
940
+ if target_name in sys.modules:
941
+ _apply_inductor_pattern_matcher_patch()
942
+ return
943
+
944
+ import importlib.abc
945
+
946
+ class _PatternMatcherPatchFinder(importlib.abc.MetaPathFinder):
947
+ def find_spec(self, fullname, path, target=None):
948
+ if fullname != target_name:
949
+ return None
950
+ sys.meta_path.remove(self)
951
+ spec = importlib.util.find_spec(fullname)
952
+ if spec is None or spec.loader is None:
953
+ return None
954
+ original_exec = spec.loader.exec_module
955
+
956
+ def _exec_then_patch(module):
957
+ original_exec(module)
958
+ _apply_inductor_pattern_matcher_patch()
959
+
960
+ spec.loader.exec_module = _exec_then_patch # type: ignore[method-assign]
961
+ return spec
962
+
963
+ sys.meta_path.insert(0, _PatternMatcherPatchFinder())
964
+
965
+
966
+# Remove once the minimum supported torch includes
967
+# https://github.com/pytorch/pytorch/pull/196013.
968
+if not is_torch_equal_or_newer("2.16.0.dev"):
969
+ _patch_inductor_pattern_matcher()
970
+
971
# ============================================================
972
# Triton Autotuner determinism
973
# ============================================================