[Refactor] Change the way of import triton (#20774)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Author: yewentao256

tests/kernels/moe/test_batched_moe.py

@@ -6,7 +6,6 @@
 
 import pytest
 import torch
-import triton.language as tl
 
 from tests.kernels.moe.utils import (batched_moe,
                                      make_quantized_test_activations,
@@ -18,6 +17,7 @@
     invoke_moe_batched_triton_kernel)
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
 from vllm.platforms import current_platform
+from vllm.triton_utils import tl
 
 MNK_FACTORS = [
     (1, 128, 128),

vllm/attention/ops/triton_unified_attention.py

@@ -8,10 +8,9 @@
 #  - Thomas Parnell <tpa@zurich.ibm.com>
 
 import torch
-import triton
-import triton.language as tl
 
 from vllm.logger import init_logger
+from vllm.triton_utils import tl, triton
 
 logger = init_logger(__name__)
 

vllm/lora/ops/triton_ops/lora_expand_op.py

@@ -8,12 +8,11 @@
 """
 
 import torch
-import triton
-import triton.language as tl
 
 from vllm.lora.ops.triton_ops.kernel_utils import do_expand_kernel
 from vllm.lora.ops.triton_ops.utils import _get_lora_b_ptr
 from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
 from vllm.utils import direct_register_custom_op
 
 

vllm/lora/ops/triton_ops/lora_shrink_op.py

@@ -8,12 +8,11 @@
 """
 
 import torch
-import triton
-import triton.language as tl
 
 from vllm.lora.ops.triton_ops.kernel_utils import do_shrink_kernel
 from vllm.lora.ops.triton_ops.utils import _get_lora_a_ptr
 from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
 from vllm.utils import direct_register_custom_op
 
 

vllm/model_executor/layers/fused_moe/fused_batched_moe.py

@@ -4,8 +4,6 @@
 from typing import Optional
 
 import torch
-import triton
-import triton.language as tl
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
@@ -18,6 +16,7 @@
     normalize_scales_shape)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     group_broadcast)
+from vllm.triton_utils import tl, triton
 
 
 @triton.jit