[Bugfix][Misc] Add a defensive check before importing triton

Signed-off-by: Mengqing Cao <cmq0113@163.com>

Author: MengqingCao

benchmarks/kernels/benchmark_lora.py

@@ -17,8 +17,14 @@
 from utils import ArgPool, Bench, CudaGraphBenchParams
 from weight_shapes import WEIGHT_SHAPES
 
-from vllm.lora.ops.triton_ops import LoRAKernelMeta, lora_expand, lora_shrink
-from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
+from vllm.triton_utils import HAS_TRITON
+
+if HAS_TRITON:
+    from vllm.lora.ops.triton_ops import (LoRAKernelMeta, lora_expand,
+                                          lora_shrink)
+    from vllm.lora.ops.triton_ops.utils import (_LORA_A_PTR_DICT,
+                                                _LORA_B_PTR_DICT)
+
 from vllm.utils import FlexibleArgumentParser
 
 DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())

benchmarks/kernels/benchmark_moe.py

@@ -10,7 +10,12 @@
 
 import ray
 import torch
-import triton
+
+from vllm.triton_utils import HAS_TRITON
+
+if HAS_TRITON:
+    import triton
+
 from ray.experimental.tqdm_ray import tqdm
 from transformers import AutoConfig
 

benchmarks/kernels/benchmark_rmsnorm.py

@@ -4,7 +4,12 @@
 from typing import Optional, Union
 
 import torch
-import triton
+
+from vllm.triton_utils import HAS_TRITON
+
+if HAS_TRITON:
+    import triton
+
 from flashinfer.norm import fused_add_rmsnorm, rmsnorm
 from torch import nn
 

benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py

@@ -6,7 +6,12 @@
 # Import DeepGEMM functions
 import deep_gemm
 import torch
-import triton
+
+from vllm.triton_utils import HAS_TRITON
+
+if HAS_TRITON:
+    import triton
+
 from deep_gemm import calc_diff, ceil_div, get_col_major_tma_aligned_tensor
 
 # Import vLLM functions

tests/kernels/test_flashmla.py

@@ -5,7 +5,11 @@
 
 import pytest
 import torch
-import triton
+
+from vllm.triton_utils.importing import HAS_TRITON
+
+if HAS_TRITON:
+    import triton
 
 from vllm.attention.ops.flashmla import (flash_mla_with_kvcache,
                                          get_mla_metadata,

vllm/attention/ops/blocksparse_attention/blocksparse_attention_kernel.py

@@ -1,8 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import torch
-import triton
-import triton.language as tl
+
+from vllm.triton_utils import HAS_TRITON
+
+if HAS_TRITON:
+    import triton
+    import triton.language as tl
+
+from vllm.triton_utils import triton_heuristics_decorator, triton_jit_decorator
 
 
 def blocksparse_flash_attn_varlen_fwd(
@@ -122,7 +128,7 @@ def blocksparse_flash_attn_varlen_fwd(
     return out
 
 
-@triton.jit
+@triton_jit_decorator
 def _fwd_kernel_inner(
     acc,
     l_i,
@@ -227,11 +233,11 @@ def _fwd_kernel_inner(
     return acc, l_i, m_i
 
 
-@triton.heuristics({
+@triton_heuristics_decorator({
     "M_LT_N":
     lambda kwargs: kwargs["BLOCK_M"] < kwargs["BLOCK_N"],
 })
-@triton.jit
+@triton_jit_decorator
 def _fwd_kernel_batch_inference(
     Q,
     K,

vllm/attention/ops/blocksparse_attention/utils.py

@@ -8,7 +8,11 @@
 
 import numpy as np
 import torch
-import triton
+
+from vllm.triton_utils import HAS_TRITON
+
+if HAS_TRITON:
+    import triton
 
 
 class csr_matrix:

vllm/attention/ops/chunked_prefill_paged_decode.py

@@ -7,18 +7,23 @@
 #  - Thomas Parnell <tpa@zurich.ibm.com>
 
 import torch
-import triton
-import triton.language as tl
+
+from vllm.triton_utils import HAS_TRITON
+
+if HAS_TRITON:
+    import triton
+    import triton.language as tl
+from vllm.triton_utils import triton_jit_decorator
 
 from .prefix_prefill import context_attention_fwd
 
 
-@triton.jit
+@triton_jit_decorator
 def cdiv_fn(x, y):
     return (x + y - 1) // y
 
 
-@triton.jit
+@triton_jit_decorator
 def kernel_paged_attention_2d(
         output_ptr,  # [num_tokens, num_query_heads, head_size]
         query_ptr,  # [num_tokens, num_query_heads, head_size]

vllm/attention/ops/prefix_prefill.py

@@ -4,8 +4,12 @@
 # https://github.com/ModelTC/lightllm/blob/main/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py
 
 import torch
-import triton
-import triton.language as tl
+
+from vllm.triton_utils import HAS_TRITON
+
+if HAS_TRITON:
+    import triton
+    import triton.language as tl
 
 from vllm.platforms import current_platform
 
@@ -16,7 +20,7 @@
 # To check compatibility
 IS_TURING = current_platform.get_device_capability() == (7, 5)
 
-if triton.__version__ >= "2.1.0":
+if HAS_TRITON and triton.__version__ >= "2.1.0":
 
     @triton.jit
     def _fwd_kernel(

vllm/attention/ops/triton_decode_attention.py

@@ -30,10 +30,14 @@
 
 import logging
 
-import triton
-import triton.language as tl
+from vllm.triton_utils import HAS_TRITON
+
+if HAS_TRITON:
+    import triton
+    import triton.language as tl
 
 from vllm.platforms import current_platform
+from vllm.triton_utils import triton_jit_decorator
 
 is_hip_ = current_platform.is_rocm()
 
@@ -46,13 +50,13 @@
     "can be ignored.")
 
 
-@triton.jit
+@triton_jit_decorator
 def tanh(x):
     # Tanh is just a scaled sigmoid
     return 2 * tl.sigmoid(2 * x) - 1
 
 
-@triton.jit
+@triton_jit_decorator
 def _fwd_kernel_stage1(
     Q,
     K_Buffer,
@@ -228,7 +232,7 @@ def _decode_att_m_fwd(
     )
 
 
-@triton.jit
+@triton_jit_decorator
 def _fwd_grouped_kernel_stage1(
     Q,
     K_Buffer,
@@ -468,7 +472,7 @@ def _decode_grouped_att_m_fwd(
     )
 
 
-@triton.jit
+@triton_jit_decorator
 def _fwd_kernel_stage2(
     Mid_O,
     o,

vllm/attention/ops/triton_flash_attention.py

@@ -22,46 +22,52 @@
 """
 
 import torch
-import triton
-import triton.language as tl
+
+from vllm.triton_utils import HAS_TRITON
+
+if HAS_TRITON:
+    import triton
+    import triton.language as tl
+
+from vllm.triton_utils import triton_autotune_decorator, triton_jit_decorator
 
 torch_dtype: tl.constexpr = torch.float16
 
 
-@triton.jit
+@triton_jit_decorator
 def cdiv_fn(x, y):
     return (x + y - 1) // y
 
 
-@triton.jit
+@triton_jit_decorator
 def max_fn(x, y):
     return tl.math.max(x, y)
 
 
-@triton.jit
+@triton_jit_decorator
 def dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride):
     ms = tl.arange(0, m)
     ns = tl.arange(0, n)
     return philox_offset + ms[:, None] * stride + ns[None, :]
 
 
-@triton.jit
+@triton_jit_decorator
 def dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride):
     rng_offsets = dropout_offsets(philox_seed, philox_offset, dropout_p, m, n,
                                   stride).to(tl.uint32)
     # TODO: use tl.randint for better performance
     return tl.rand(philox_seed, rng_offsets)
 
 
-@triton.jit
+@triton_jit_decorator
 def dropout_mask(philox_seed, philox_offset, dropout_p, m, n, stride):
     rng_output = dropout_rng(philox_seed, philox_offset, dropout_p, m, n,
                              stride)
     rng_keep = rng_output > dropout_p
     return rng_keep
 
 
-@triton.jit
+@triton_jit_decorator
 def load_fn(block_ptr, first, second, pad):
     if first and second:
         tensor = tl.load(block_ptr, boundary_check=(0, 1), padding_option=pad)
@@ -74,7 +80,7 @@ def load_fn(block_ptr, first, second, pad):
     return tensor
 
 
-@triton.jit
+@triton_jit_decorator
 def _attn_fwd_inner(
     acc,
     l_i,
@@ -208,7 +214,7 @@ def _attn_fwd_inner(
     return acc, l_i, m_i
 
 
-@triton.autotune(
+@triton_autotune_decorator(
     configs=[
         triton.Config(
             {
@@ -306,7 +312,7 @@ def _attn_fwd_inner(
     ],
     key=['IS_CAUSAL', 'dropout_p', 'BLOCK_DMODEL'],
 )
-@triton.jit
+@triton_jit_decorator
 def attn_fwd(
     Q,
     K,

vllm/attention/ops/triton_merge_attn_states.py

@@ -2,8 +2,14 @@
 from typing import Optional
 
 import torch
-import triton
-import triton.language as tl
+
+from vllm.triton_utils import HAS_TRITON
+
+if HAS_TRITON:
+    import triton
+    import triton.language as tl
+
+from vllm.triton_utils import triton_jit_decorator
 
 
 # Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
@@ -35,7 +41,7 @@ def merge_attn_states(
     )
 
 
-@triton.jit
+@triton_jit_decorator
 def merge_attn_states_kernel(
     output,  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
     output_lse,  # [NUM_HEADS, NUM_TOKENS]