[Bugfix] Fix TritonPlaceholder conflicts with torch.compile

Signed-off-by: Mengqing Cao <[email protected]>

Author: MengqingCao

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/attention/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,21 +7,26 @@
 #  - Thomas Parnell <[email protected]>
 
 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 import _custom_ops as ops
 from vllm.platforms.rocm import use_rocm_custom_paged_attention
+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,10 +4,15 @@
 # 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
+from vllm.triton_utils import triton_jit_decorator
 
 # Static kernels parameters
 BASE_BLOCK = 128 if current_platform.has_device_capability(80) else 64
@@ -32,7 +37,7 @@
 #     ],
 #     key=["BLOCK_SIZE", "MAX_Q_LEN", "MAX_CTX_LEN"]
 # )
-@triton.jit
+@triton_jit_decorator
 def _fwd_kernel(Q,
                 K,
                 V,
@@ -280,7 +285,7 @@ def _fwd_kernel(Q,
     return
 
 
-@triton.jit
+@triton_jit_decorator
 def _fwd_kernel_flash_attn_v2(
     Q,
     K,
@@ -466,7 +471,7 @@ def _fwd_kernel_flash_attn_v2(
     return
 
 
-@triton.jit
+@triton_jit_decorator
 def _fwd_kernel_alibi(
     Q,
     K,

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()
 
@@ -47,13 +51,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,
@@ -229,7 +233,7 @@ def _decode_att_m_fwd(
     )
 
 
-@triton.jit
+@triton_jit_decorator
 def _fwd_grouped_kernel_stage1(
     Q,
     K_Buffer,
@@ -469,7 +473,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

@@ -25,11 +25,16 @@
 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 import _custom_ops as ops
 from vllm.platforms import current_platform
+from vllm.triton_utils import triton_jit_decorator
 
 SUPPORTED_LAYOUTS = ['thd', 'bhsd', 'bshd']
 
@@ -234,19 +239,19 @@ def check_args(self, q, k, v, o):
         assert self.layout == 'thd' or not self.varlen
 
 
-@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)
 
 
 # Convenience function to load with optional boundary checks.
 # "First" is the major dim, "second" is the minor dim.
-@triton.jit
+@triton_jit_decorator
 def masked_load(ptrs, offset_first, offset_second, boundary_first,
                 boundary_second):
     if offset_first is not None and offset_second is not None:
@@ -264,7 +269,7 @@ def masked_load(ptrs, offset_first, offset_second, boundary_first,
     return tensor
 
 
-@triton.jit
+@triton_jit_decorator
 def compute_alibi_block(alibi_slope,
                         seqlen_q,
                         seqlen_k,
@@ -318,14 +323,14 @@ def compute_alibi_tensor(alibi_slopes, seqlen_q, seqlen_k):
         -1) * relative_pos  # (Z, H, N_CTX_Q, N_CTX_K)
 
 
-@triton.jit
+@triton_jit_decorator
 def quant_fp8(x, scale):
     x *= scale
     x = tl.clamp(x, FP8_MIN, FP8_MAX)
     return x
 
 
-@triton.jit
+@triton_jit_decorator
 def _attn_fwd_inner(
     acc,
     l_i,
@@ -676,7 +681,7 @@ def get_autotune_configs():
     key=autotune_keys,
     use_cuda_graph=True,
 )
-@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]

vllm/lora/ops/triton_ops/kernel_utils.py

@@ -2,11 +2,15 @@
 """
 Utilities for Punica kernel construction.
 """
-import triton
-import triton.language as tl
+from vllm.triton_utils import HAS_TRITON
 
+if HAS_TRITON:
+    import triton.language as tl
 
-@triton.jit
+from vllm.triton_utils import triton_jit_decorator
+
+
+@triton_jit_decorator
 def mm_k(a_ptr, b_ptr, ak_stride, bk_stride, offset_k, K: tl.constexpr,
          BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
          EVEN_K: tl.constexpr, SPLIT_K: tl.constexpr, CAST_TYPE: tl.constexpr,
@@ -59,7 +63,7 @@ def mm_k(a_ptr, b_ptr, ak_stride, bk_stride, offset_k, K: tl.constexpr,
     return accumulator
 
 
-@triton.jit
+@triton_jit_decorator
 def do_expand_kernel(
     pid_n,
     lora_index,
@@ -161,7 +165,7 @@ def do_expand_kernel(
     tl.store(c_ptr, tiled_c, mask=c_mask)
 
 
-@triton.jit
+@triton_jit_decorator
 def do_shrink_kernel(
     pid_n,
     pid_sk,

vllm/lora/ops/triton_ops/lora_expand.py

@@ -14,10 +14,11 @@
 
 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.triton_utils import triton_jit_decorator
 from vllm.utils import direct_register_custom_op
 
 
-@triton.jit
+@triton_jit_decorator
 def _lora_expand_kernel(
         input_ptr,
         lora_ptr,