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 15, 13:30 UTC
[Bugfix][Misc] Use TritonPlaceholderModule to defensively import triton (#15099)
Signed-off-by: Mengqing Cao <cmq0113@163.com>
Mengqing Cao committed
Apr 25, 2025 at 13:51 UTC
2f54045508930a741f3638cc3c1174fdd672087f
5 files changed
+53
-6
benchmarks/kernels/benchmark_lora.py
+8
-2
@@ -17,8 +17,14 @@ from torch.utils.benchmark import Measurement as TMeasurement
17
from utils import ArgPool, Bench, CudaGraphBenchParams
18
from weight_shapes import WEIGHT_SHAPES
19
20
-from vllm.lora.ops.triton_ops import LoRAKernelMeta, lora_expand, lora_shrink
21
-from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
20
+from vllm.triton_utils import HAS_TRITON
21
+
22
+if HAS_TRITON:
23
+ from vllm.lora.ops.triton_ops import (LoRAKernelMeta, lora_expand,
24
+ lora_shrink)
25
+ from vllm.lora.ops.triton_ops.utils import (_LORA_A_PTR_DICT,
26
+ _LORA_B_PTR_DICT)
27
+
28
from vllm.utils import FlexibleArgumentParser
29
30
DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+3
-1
@@ -10,8 +10,10 @@ from packaging import version
10
11
from vllm import _custom_ops as ops
12
from vllm.attention.backends.utils import PAD_SLOT_ID
13
+from vllm.triton_utils import HAS_TRITON
14
14
-TRITON3 = version.parse(triton.__version__) >= version.parse("3.0.0")
15
+TRITON3 = HAS_TRITON and (version.parse(triton.__version__)
16
+ >= version.parse("3.0.0"))
17
18
if TRITON3:
19
vllm/triton_utils/__init__.py
+1
-1
@@ -2,4 +2,4 @@
2
3
from vllm.triton_utils.importing import HAS_TRITON
4
5
-__all__ = ["HAS_TRITON"]
\ No newline at end of file
5
+__all__ = ["HAS_TRITON"]
vllm/triton_utils/importing.py
+38
-2
@@ -1,17 +1,53 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+import sys
4
+import types
5
from importlib.util import find_spec
6
7
from vllm.logger import init_logger
6
-from vllm.platforms import current_platform
8
9
logger = init_logger(__name__)
10
11
HAS_TRITON = (
12
find_spec("triton") is not None
12
- and not current_platform.is_xpu() # Not compatible
13
+ or find_spec("pytorch-triton-xpu") is not None # Not compatible
14
)
15
16
if not HAS_TRITON:
17
logger.info("Triton not installed or not compatible; certain GPU-related"
18
" functions will not be available.")
19
+
20
+ class TritonPlaceholder(types.ModuleType):
21
+
22
+ def __init__(self):
23
+ super().__init__("triton")
24
+ self.jit = self._dummy_decorator("jit")
25
+ self.autotune = self._dummy_decorator("autotune")
26
+ self.heuristics = self._dummy_decorator("heuristics")
27
+ self.language = TritonLanguagePlaceholder()
28
+ logger.warning_once(
29
+ "Triton is not installed. Using dummy decorators. "
30
+ "Install it via `pip install triton` to enable kernel"
31
+ "compilation.")
32
+
33
+ def _dummy_decorator(self, name):
34
+
35
+ def decorator(func=None, **kwargs):
36
+ if func is None:
37
+ return lambda f: f
38
+ return func
39
+
40
+ return decorator
41
+
42
+ class TritonLanguagePlaceholder(types.ModuleType):
43
+
44
+ def __init__(self):
45
+ super().__init__("triton.language")
46
+ self.constexpr = None
47
+ self.dtype = None
48
+
49
+ sys.modules['triton'] = TritonPlaceholder()
50
+ sys.modules['triton.language'] = TritonLanguagePlaceholder()
51
+
52
+if 'triton' in sys.modules:
53
+ logger.info("Triton module has been replaced with a placeholder.")
vllm/utils.py
+3
@@ -63,6 +63,9 @@ from torch.library import Library
63
from typing_extensions import Never, ParamSpec, TypeIs, assert_never
64
65
import vllm.envs as envs
66
+# NOTE: import triton_utils to make TritonPlaceholderModule work
67
+# if triton is unavailable
68
+import vllm.triton_utils # noqa: F401
69
from vllm.logger import enable_trace_function_call, init_logger
70
71
if TYPE_CHECKING: