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, 02:15 UTC
[Refactor] Change the way of import triton (#20774)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 12, 2025 at 22:39 UTC
c1acd6d7d48505d070546b7afa922e4a93ac5447
5 files changed
+5
-9
tests/kernels/moe/test_batched_moe.py
+1
-1
@@ -6,7 +6,6 @@ from typing import Optional
6
7
import pytest
8
import torch
9
-import triton.language as tl
9
10
from tests.kernels.moe.utils import (batched_moe,
11
make_quantized_test_activations,
@@ -18,6 +17,7 @@ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
17
invoke_moe_batched_triton_kernel)
18
from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
19
from vllm.platforms import current_platform
20
+from vllm.triton_utils import tl
21
22
MNK_FACTORS = [
23
(1, 128, 128),
vllm/attention/ops/triton_unified_attention.py
+1
-2
@@ -8,10 +8,9 @@
8
# - Thomas Parnell <tpa@zurich.ibm.com>
9
10
import torch
11
-import triton
12
-import triton.language as tl
11
12
from vllm.logger import init_logger
13
+from vllm.triton_utils import tl, triton
14
15
logger = init_logger(__name__)
16
vllm/lora/ops/triton_ops/lora_expand_op.py
+1
-2
@@ -8,12 +8,11 @@ https://arxiv.org/abs/2310.18547
8
"""
9
10
import torch
11
-import triton
12
-import triton.language as tl
11
12
from vllm.lora.ops.triton_ops.kernel_utils import do_expand_kernel
13
from vllm.lora.ops.triton_ops.utils import _get_lora_b_ptr
14
from vllm.platforms import current_platform
15
+from vllm.triton_utils import tl, triton
16
from vllm.utils import direct_register_custom_op
17
18
vllm/lora/ops/triton_ops/lora_shrink_op.py
+1
-2
@@ -8,12 +8,11 @@ https://arxiv.org/abs/2310.18547
8
"""
9
10
import torch
11
-import triton
12
-import triton.language as tl
11
12
from vllm.lora.ops.triton_ops.kernel_utils import do_shrink_kernel
13
from vllm.lora.ops.triton_ops.utils import _get_lora_a_ptr
14
from vllm.platforms import current_platform
15
+from vllm.triton_utils import tl, triton
16
from vllm.utils import direct_register_custom_op
17
18
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+1
-2
@@ -4,8 +4,6 @@
4
from typing import Optional
5
6
import torch
7
-import triton
8
-import triton.language as tl
7
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
@@ -18,6 +16,7 @@ from vllm.model_executor.layers.fused_moe.utils import (
16
normalize_scales_shape)
17
from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
group_broadcast)
19
+from vllm.triton_utils import tl, triton
20
21
22
@triton.jit