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, 11:00 UTC
[Bugfix] Fix triton import with local TritonPlaceholder (#17446)
Signed-off-by: Mengqing Cao <cmq0113@163.com>
Mengqing Cao committed
May 6, 2025 at 17:53 UTC
f9bc5a0693d831467df738fce36c701f6a88a18d
30 files changed
+171
-81
benchmarks/kernels/benchmark_moe.py
+1
-1
@@ -10,12 +10,12 @@ from typing import Any, TypedDict
10
11
import ray
12
import torch
13
-import triton
13
from ray.experimental.tqdm_ray import tqdm
14
from transformers import AutoConfig
15
16
from vllm.model_executor.layers.fused_moe.fused_moe import *
17
from vllm.platforms import current_platform
18
+from vllm.triton_utils import triton
19
from vllm.utils import FlexibleArgumentParser
20
21
FP8_DTYPE = current_platform.fp8_dtype()
benchmarks/kernels/benchmark_rmsnorm.py
+1
-1
@@ -4,11 +4,11 @@ import itertools
4
from typing import Optional, Union
5
6
import torch
7
-import triton
7
from flashinfer.norm import fused_add_rmsnorm, rmsnorm
8
from torch import nn
9
10
from vllm import _custom_ops as vllm_ops
11
+from vllm.triton_utils import triton
12
13
14
class HuggingFaceRMSNorm(nn.Module):
benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py
+1
-1
@@ -6,13 +6,13 @@ import time
6
# Import DeepGEMM functions
7
import deep_gemm
8
import torch
9
-import triton
9
from deep_gemm import calc_diff, ceil_div, get_col_major_tma_aligned_tensor
10
11
# Import vLLM functions
12
from vllm import _custom_ops as ops
13
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
14
per_token_group_quant_fp8, w8a8_block_fp8_matmul)
15
+from vllm.triton_utils import triton
16
17
18
# Copied from
tests/kernels/attention/test_flashmla.py
+1
-1
@@ -5,11 +5,11 @@ import random
5
6
import pytest
7
import torch
8
-import triton
8
9
from vllm.attention.ops.flashmla import (flash_mla_with_kvcache,
10
get_mla_metadata,
11
is_flashmla_supported)
12
+from vllm.triton_utils import triton
13
14
15
def cal_diff(x: torch.Tensor, y: torch.Tensor, name: str) -> None:
tests/test_triton_utils.py
new
+92
@@ -0,0 +1,92 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import sys
4
+import types
5
+from unittest import mock
6
+
7
+from vllm.triton_utils.importing import (TritonLanguagePlaceholder,
8
+ TritonPlaceholder)
9
+
10
+
11
+def test_triton_placeholder_is_module():
12
+ triton = TritonPlaceholder()
13
+ assert isinstance(triton, types.ModuleType)
14
+ assert triton.__name__ == "triton"
15
+
16
+
17
+def test_triton_language_placeholder_is_module():
18
+ triton_language = TritonLanguagePlaceholder()
19
+ assert isinstance(triton_language, types.ModuleType)
20
+ assert triton_language.__name__ == "triton.language"
21
+
22
+
23
+def test_triton_placeholder_decorators():
24
+ triton = TritonPlaceholder()
25
+
26
+ @triton.jit
27
+ def foo(x):
28
+ return x
29
+
30
+ @triton.autotune
31
+ def bar(x):
32
+ return x
33
+
34
+ @triton.heuristics
35
+ def baz(x):
36
+ return x
37
+
38
+ assert foo(1) == 1
39
+ assert bar(2) == 2
40
+ assert baz(3) == 3
41
+
42
+
43
+def test_triton_placeholder_decorators_with_args():
44
+ triton = TritonPlaceholder()
45
+
46
+ @triton.jit(debug=True)
47
+ def foo(x):
48
+ return x
49
+
50
+ @triton.autotune(configs=[], key="x")
51
+ def bar(x):
52
+ return x
53
+
54
+ @triton.heuristics(
55
+ {"BLOCK_SIZE": lambda args: 128 if args["x"] > 1024 else 64})
56
+ def baz(x):
57
+ return x
58
+
59
+ assert foo(1) == 1
60
+ assert bar(2) == 2
61
+ assert baz(3) == 3
62
+
63
+
64
+def test_triton_placeholder_language():
65
+ lang = TritonLanguagePlaceholder()
66
+ assert isinstance(lang, types.ModuleType)
67
+ assert lang.__name__ == "triton.language"
68
+ assert lang.constexpr is None
69
+ assert lang.dtype is None
70
+ assert lang.int64 is None
71
+
72
+
73
+def test_triton_placeholder_language_from_parent():
74
+ triton = TritonPlaceholder()
75
+ lang = triton.language
76
+ assert isinstance(lang, TritonLanguagePlaceholder)
77
+
78
+
79
+def test_no_triton_fallback():
80
+ # clear existing triton modules
81
+ sys.modules.pop("triton", None)
82
+ sys.modules.pop("triton.language", None)
83
+ sys.modules.pop("vllm.triton_utils", None)
84
+ sys.modules.pop("vllm.triton_utils.importing", None)
85
+
86
+ # mock triton not being installed
87
+ with mock.patch.dict(sys.modules, {"triton": None}):
88
+ from vllm.triton_utils import HAS_TRITON, tl, triton
89
+ assert HAS_TRITON is False
90
+ assert triton.__class__.__name__ == "TritonPlaceholder"
91
+ assert triton.language.__class__.__name__ == "TritonLanguagePlaceholder"
92
+ assert tl.__class__.__name__ == "TritonLanguagePlaceholder"
vllm/attention/ops/blocksparse_attention/blocksparse_attention_kernel.py
+2
-2
@@ -1,8 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
import torch
4
-import triton
5
-import triton.language as tl
4
+
5
+from vllm.triton_utils import tl, triton
6
7
8
def blocksparse_flash_attn_varlen_fwd(
vllm/attention/ops/blocksparse_attention/utils.py
+2
-1
@@ -8,7 +8,8 @@ from functools import lru_cache
8
9
import numpy as np
10
import torch
11
-import triton
11
+
12
+from vllm.triton_utils import triton
13
14
15
class csr_matrix:
vllm/attention/ops/chunked_prefill_paged_decode.py
+1
-2
@@ -7,11 +7,10 @@
7
# - Thomas Parnell <tpa@zurich.ibm.com>
8
9
import torch
10
-import triton
11
-import triton.language as tl
10
11
from vllm import _custom_ops as ops
12
from vllm.platforms.rocm import use_rocm_custom_paged_attention
13
+from vllm.triton_utils import tl, triton
14
15
from .prefix_prefill import context_attention_fwd
16
vllm/attention/ops/prefix_prefill.py
+1
-2
@@ -4,10 +4,9 @@
4
# https://github.com/ModelTC/lightllm/blob/main/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py
5
6
import torch
7
-import triton
8
-import triton.language as tl
7
8
from vllm.platforms import current_platform
9
+from vllm.triton_utils import tl, triton
10
11
# Static kernels parameters
12
BASE_BLOCK = 128 if current_platform.has_device_capability(80) else 64
vllm/attention/ops/triton_decode_attention.py
+1
-3
@@ -30,10 +30,8 @@ It supports page size >= 1.
30
31
import logging
32
33
-import triton
34
-import triton.language as tl
35
-
33
from vllm.platforms import current_platform
34
+from vllm.triton_utils import tl, triton
35
36
is_hip_ = current_platform.is_rocm()
37
vllm/attention/ops/triton_flash_attention.py
+1
-2
@@ -25,11 +25,10 @@ Currently only the forward kernel is supported, and contains these features:
25
from typing import Optional
26
27
import torch
28
-import triton
29
-import triton.language as tl
28
29
from vllm import _custom_ops as ops
30
from vllm.platforms import current_platform
31
+from vllm.triton_utils import tl, triton
32
33
SUPPORTED_LAYOUTS = ['thd', 'bhsd', 'bshd']
34
vllm/attention/ops/triton_merge_attn_states.py
+2
-2
@@ -2,8 +2,8 @@
2
from typing import Optional
3
4
import torch
5
-import triton
6
-import triton.language as tl
5
+
6
+from vllm.triton_utils import tl, triton
7
8
9
# Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
vllm/lora/ops/triton_ops/kernel_utils.py
+1
-2
@@ -2,8 +2,7 @@
2
"""
3
Utilities for Punica kernel construction.
4
"""
5
-import triton
6
-import triton.language as tl
5
+from vllm.triton_utils import tl, triton
6
7
8
@triton.jit
vllm/model_executor/layers/fused_moe/fused_moe.py
+1
-2
@@ -6,8 +6,6 @@ import os
6
from typing import Any, Callable, Dict, List, Optional, Tuple
7
8
import torch
9
-import triton
10
-import triton.language as tl
9
10
import vllm.envs as envs
11
from vllm import _custom_ops as ops
@@ -21,6 +19,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
19
from vllm.model_executor.layers.quantization.utils.int8_utils import (
20
per_token_group_quant_int8, per_token_quant_int8)
21
from vllm.platforms import current_platform
22
+from vllm.triton_utils import tl, triton
23
from vllm.utils import direct_register_custom_op
24
25
from .rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled
vllm/model_executor/layers/fused_moe/moe_align_block_size.py
+1
-2
@@ -2,11 +2,10 @@
2
from typing import Optional, Tuple
3
4
import torch
5
-import triton
6
-import triton.language as tl
5
6
import vllm.envs as envs
7
from vllm import _custom_ops as ops
8
+from vllm.triton_utils import tl, triton
9
from vllm.utils import round_up
10
11
vllm/model_executor/layers/lightning_attn.py
+2
-2
@@ -1,9 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
import torch
3
-import triton
4
-import triton.language as tl
3
from einops import rearrange
4
5
+from vllm.triton_utils import tl, triton
6
+
7
8
@triton.jit
9
def _fwd_diag_kernel(Q, K, V, Out, S, b: tl.constexpr, h: tl.constexpr, n,
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+1
-3
@@ -4,13 +4,11 @@
4
# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/selective_state_update.py
5
6
import torch
7
-import triton
8
-import triton.language as tl
7
from packaging import version
8
9
from vllm import _custom_ops as ops
10
from vllm.attention.backends.utils import PAD_SLOT_ID
13
-from vllm.triton_utils import HAS_TRITON
11
+from vllm.triton_utils import HAS_TRITON, tl, triton
12
13
TRITON3 = HAS_TRITON and (version.parse(triton.__version__)
14
>= version.parse("3.0.0"))
vllm/model_executor/layers/mamba/ops/ssd_bmm.py
+2
-2
@@ -8,8 +8,8 @@
8
import math
9
10
import torch
11
-import triton
12
-import triton.language as tl
11
+
12
+from vllm.triton_utils import tl, triton
13
14
15
@triton.autotune(
vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py
+2
-2
@@ -6,10 +6,10 @@
6
# ruff: noqa: E501,SIM102
7
8
import torch
9
-import triton
10
-import triton.language as tl
9
from packaging import version
10
11
+from vllm.triton_utils import tl, triton
12
+
13
TRITON_22 = version.parse(triton.__version__) >= version.parse('2.2.0')
14
15
vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py
+2
-2
@@ -8,8 +8,8 @@
8
import math
9
10
import torch
11
-import triton
12
-import triton.language as tl
11
+
12
+from vllm.triton_utils import tl, triton
13
14
from .mamba_ssm import softplus
15
vllm/model_executor/layers/mamba/ops/ssd_combined.py
+2
-1
@@ -6,10 +6,11 @@
6
# ruff: noqa: E501
7
8
import torch
9
-import triton
9
from einops import rearrange
10
from packaging import version
11
12
+from vllm.triton_utils import triton
13
+
14
from .ssd_bmm import _bmm_chunk_fwd
15
from .ssd_chunk_scan import _chunk_scan_fwd
16
from .ssd_chunk_state import (_chunk_cumsum_fwd, _chunk_state_fwd,
vllm/model_executor/layers/mamba/ops/ssd_state_passing.py
+2
-2
@@ -6,8 +6,8 @@
6
# ruff: noqa: E501
7
8
import torch
9
-import triton
10
-import triton.language as tl
9
+
10
+from vllm.triton_utils import tl, triton
11
12
13
@triton.autotune(
vllm/model_executor/layers/quantization/awq_triton.py
+2
-2
@@ -1,8 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
import torch
4
-import triton
5
-import triton.language as tl
4
+
5
+from vllm.triton_utils import tl, triton
6
7
AWQ_TRITON_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]
8
vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py
+2
-2
@@ -3,8 +3,8 @@
3
from typing import Optional, Type
4
5
import torch
6
-import triton
7
-import triton.language as tl
6
+
7
+from vllm.triton_utils import tl, triton
8
9
10
def is_weak_contiguous(x: torch.Tensor):
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+1
-2
@@ -7,8 +7,6 @@ import os
7
from typing import Any, Dict, List, Optional, Tuple, Union
8
9
import torch
10
-import triton
11
-import triton.language as tl
10
11
from vllm import _custom_ops as ops
12
from vllm.logger import init_logger
@@ -17,6 +15,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
15
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
16
CUTLASS_BLOCK_FP8_SUPPORTED)
17
from vllm.platforms import current_platform
18
+from vllm.triton_utils import tl, triton
19
from vllm.utils import direct_register_custom_op
20
21
logger = init_logger(__name__)
vllm/model_executor/layers/quantization/utils/int8_utils.py
+1
-2
@@ -8,10 +8,9 @@ import os
8
from typing import Any, Dict, List, Optional, Tuple
9
10
import torch
11
-import triton
12
-import triton.language as tl
11
12
from vllm.platforms import current_platform
13
+from vllm.triton_utils import tl, triton
14
15
logger = logging.getLogger(__name__)
16
vllm/triton_utils/__init__.py
+10
-2
@@ -1,5 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-from vllm.triton_utils.importing import HAS_TRITON
3
+from vllm.triton_utils.importing import (HAS_TRITON, TritonLanguagePlaceholder,
4
+ TritonPlaceholder)
5
5
-__all__ = ["HAS_TRITON"]
6
+if HAS_TRITON:
7
+ import triton
8
+ import triton.language as tl
9
+else:
10
+ triton = TritonPlaceholder()
11
+ tl = TritonLanguagePlaceholder()
12
+
13
+__all__ = ["HAS_TRITON", "triton", "tl"]
vllm/triton_utils/importing.py
+31
-29
@@ -16,32 +16,34 @@ if not HAS_TRITON:
16
logger.info("Triton not installed or not compatible; certain GPU-related"
17
" functions will not be available.")
18
19
- class TritonPlaceholder(types.ModuleType):
20
-
21
- def __init__(self):
22
- super().__init__("triton")
23
- self.jit = self._dummy_decorator("jit")
24
- self.autotune = self._dummy_decorator("autotune")
25
- self.heuristics = self._dummy_decorator("heuristics")
26
- self.language = TritonLanguagePlaceholder()
27
- logger.warning_once(
28
- "Triton is not installed. Using dummy decorators. "
29
- "Install it via `pip install triton` to enable kernel"
30
- "compilation.")
31
-
32
- def _dummy_decorator(self, name):
33
-
34
- def decorator(func=None, **kwargs):
35
- if func is None:
36
- return lambda f: f
37
- return func
38
-
39
- return decorator
40
-
41
- class TritonLanguagePlaceholder(types.ModuleType):
42
-
43
- def __init__(self):
44
- super().__init__("triton.language")
45
- self.constexpr = None
46
- self.dtype = None
47
- self.int64 = None
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(*args, **kwargs):
36
+ if args and callable(args[0]):
37
+ return args[0]
38
+ return lambda f: f
39
+
40
+ return decorator
41
+
42
+
43
+class TritonLanguagePlaceholder(types.ModuleType):
44
+
45
+ def __init__(self):
46
+ super().__init__("triton.language")
47
+ self.constexpr = None
48
+ self.dtype = None
49
+ self.int64 = None
vllm/v1/sample/rejection_sampler.py
+1
-2
@@ -3,10 +3,9 @@ from typing import Optional
3
4
import torch
5
import torch.nn as nn
6
-import triton
7
-import triton.language as tl
6
7
from vllm.logger import init_logger
8
+from vllm.triton_utils import tl, triton
9
from vllm.v1.sample.metadata import SamplingMetadata
10
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
11
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
vllm/v1/spec_decode/eagle.py
+1
-2
@@ -1,8 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
import torch
3
import torch.nn as nn
4
-import triton
5
-import triton.language as tl
4
5
from vllm.config import CompilationLevel, VllmConfig, set_current_vllm_config
6
from vllm.forward_context import set_forward_context
@@ -11,6 +9,7 @@ from vllm.model_executor.model_loader.loader import get_model_loader
9
from vllm.model_executor.model_loader.utils import set_default_torch_dtype
10
from vllm.model_executor.models import ModelRegistry
11
from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
12
+from vllm.triton_utils import tl, triton
13
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
14
from vllm.v1.sample.metadata import SamplingMetadata
15