remove torch module patch

Signed-off-by: Isotr0py <[email protected]>

Author: Isotr0py

vllm/triton_utils/importing.py

@@ -1,10 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-
 import importlib
-import pkgutil
 import sys
 import types
 from importlib.util import find_spec
+from typing import Optional
 
 from vllm.logger import init_logger
 
@@ -19,17 +18,37 @@
     logger.info("Triton not installed or not compatible; certain GPU-related"
                 " functions will not be available.")
 
-    class TritonPlaceholder(types.ModuleType):
+    class TritonModulePlaceholder(types.ModuleType):
+
+        def __init__(
+            self,
+            name: str,
+            dummy_objects: Optional[list[str]] = None,
+        ):
+            super().__init__(name)
+
+            if dummy_objects is not None:
+                for obj_name in dummy_objects:
+                    setattr(self, obj_name, object)
+
+    class TritonLanguagePlaceholder(TritonModulePlaceholder):
+
+        def __init__(self):
+            super().__init__("triton.language")
+            self.constexpr = None
+            self.dtype = None
+
+    class TritonPlaceholder(TritonModulePlaceholder):
 
         def __init__(self):
-            super().__init__("triton")
+            super().__init__("triton", dummy_objects=["Config"])
             self.jit = self._dummy_decorator("jit")
             self.autotune = self._dummy_decorator("autotune")
             self.heuristics = self._dummy_decorator("heuristics")
             self.language = TritonLanguagePlaceholder()
             logger.warning_once(
                 "Triton is not installed. Using dummy decorators. "
-                "Install it via `pip install triton` to enable kernel"
+                "Install it via `pip install triton` to enable kernel "
                 "compilation.")
 
         def _dummy_decorator(self, name):
@@ -41,34 +60,27 @@ def decorator(func=None, **kwargs):
 
             return decorator
 
-    class TritonLanguagePlaceholder(types.ModuleType):
-
-        def __init__(self):
-            super().__init__("triton.language")
-            self.constexpr = None
-            self.dtype = None
+    # Hack `_is_triton_available` in torch to return False
+    torch_hints = importlib.import_module("torch._inductor.runtime.hints")
+    torch_hints._is_triton_available = lambda: False  # type: ignore[attr-defined]
 
-    def init_torch_inductor_runtime():
-        name = "torch._inductor.runtime"
-        torch_runtime = importlib.import_module(name)
-        path = torch_runtime.__path__
-        for module_info in pkgutil.iter_modules(path, name + "."):
-            if not module_info.ispkg:
-                try:
-                    importlib.import_module(module_info.name)
-                except Exception as e:
-                    logger.warning(
-                        "Ignore import error when loading " \
-                        "%s: %s", module_info.name, e)
-                    continue
-
-    # initialize torch inductor without triton placeholder
-    # FIXME(Isotr0py): See if we can remove this after bumping torch version
-    # to 2.7.0, because torch 2.7.0 has a better triton check.
-    init_torch_inductor_runtime()
-    # Replace the triton module in sys.modules with the placeholder
+    # Replace the triton module and torch triton helpers in sys.modules
+    # with the placeholder
     sys.modules['triton'] = TritonPlaceholder()
     sys.modules['triton.language'] = TritonLanguagePlaceholder()
+    sys.modules['torch._inductor.runtime.triton_helpers'] = types.ModuleType(
+        "triton_helpers")
+
+    # Replace triton submodules with dummy objects to keep compatibility with
+    # torch triton check
+    triton_modules_with_objects = {
+        "triton.compiler": ["CompiledKernel"],
+        "triton.runtime.autotuner": ["OutOfResources"],
+        "triton.runtime.jit": ["KernelInterface"],
+    }
+    for module_name, dummy_objects in triton_modules_with_objects.items():
+        sys.modules[module_name] = TritonModulePlaceholder(
+            module_name, dummy_objects=dummy_objects)
 
 if 'triton' in sys.modules:
     logger.info("Triton module has been replaced with a placeholder.")