[Bugfix] Fix triton import with local TritonPlaceholder (#17446)

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

Author: MengqingCao

benchmarks/kernels/benchmark_moe.py

@@ -10,12 +10,12 @@
 
 import ray
 import torch
-import triton
 from ray.experimental.tqdm_ray import tqdm
 from transformers import AutoConfig
 
 from vllm.model_executor.layers.fused_moe.fused_moe import *
 from vllm.platforms import current_platform
+from vllm.triton_utils import triton
 from vllm.utils import FlexibleArgumentParser
 
 FP8_DTYPE = current_platform.fp8_dtype()

benchmarks/kernels/benchmark_rmsnorm.py

@@ -4,11 +4,11 @@
 from typing import Optional, Union
 
 import torch
-import triton
 from flashinfer.norm import fused_add_rmsnorm, rmsnorm
 from torch import nn
 
 from vllm import _custom_ops as vllm_ops
+from vllm.triton_utils import triton
 
 
 class HuggingFaceRMSNorm(nn.Module):

benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py

@@ -6,13 +6,13 @@
 # Import DeepGEMM functions
 import deep_gemm
 import torch
-import triton
 from deep_gemm import calc_diff, ceil_div, get_col_major_tma_aligned_tensor
 
 # Import vLLM functions
 from vllm import _custom_ops as ops
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     per_token_group_quant_fp8, w8a8_block_fp8_matmul)
+from vllm.triton_utils import triton
 
 
 # Copied from

tests/kernels/attention/test_flashmla.py

@@ -5,11 +5,11 @@
 
 import pytest
 import torch
-import triton
 
 from vllm.attention.ops.flashmla import (flash_mla_with_kvcache,
                                          get_mla_metadata,
                                          is_flashmla_supported)
+from vllm.triton_utils import triton
 
 
 def cal_diff(x: torch.Tensor, y: torch.Tensor, name: str) -> None:

tests/test_triton_utils.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+import types
+from unittest import mock
+
+from vllm.triton_utils.importing import (TritonLanguagePlaceholder,
+                                         TritonPlaceholder)
+
+
+def test_triton_placeholder_is_module():
+    triton = TritonPlaceholder()
+    assert isinstance(triton, types.ModuleType)
+    assert triton.__name__ == "triton"
+
+
+def test_triton_language_placeholder_is_module():
+    triton_language = TritonLanguagePlaceholder()
+    assert isinstance(triton_language, types.ModuleType)
+    assert triton_language.__name__ == "triton.language"
+
+
+def test_triton_placeholder_decorators():
+    triton = TritonPlaceholder()
+
+    @triton.jit
+    def foo(x):
+        return x
+
+    @triton.autotune
+    def bar(x):
+        return x
+
+    @triton.heuristics
+    def baz(x):
+        return x
+
+    assert foo(1) == 1
+    assert bar(2) == 2
+    assert baz(3) == 3
+
+
+def test_triton_placeholder_decorators_with_args():
+    triton = TritonPlaceholder()
+
+    @triton.jit(debug=True)
+    def foo(x):
+        return x
+
+    @triton.autotune(configs=[], key="x")
+    def bar(x):
+        return x
+
+    @triton.heuristics(
+        {"BLOCK_SIZE": lambda args: 128 if args["x"] > 1024 else 64})
+    def baz(x):
+        return x
+
+    assert foo(1) == 1
+    assert bar(2) == 2
+    assert baz(3) == 3
+
+
+def test_triton_placeholder_language():
+    lang = TritonLanguagePlaceholder()
+    assert isinstance(lang, types.ModuleType)
+    assert lang.__name__ == "triton.language"
+    assert lang.constexpr is None
+    assert lang.dtype is None
+    assert lang.int64 is None
+
+
+def test_triton_placeholder_language_from_parent():
+    triton = TritonPlaceholder()
+    lang = triton.language
+    assert isinstance(lang, TritonLanguagePlaceholder)
+
+
+def test_no_triton_fallback():
+    # clear existing triton modules
+    sys.modules.pop("triton", None)
+    sys.modules.pop("triton.language", None)
+    sys.modules.pop("vllm.triton_utils", None)
+    sys.modules.pop("vllm.triton_utils.importing", None)
+
+    # mock triton not being installed
+    with mock.patch.dict(sys.modules, {"triton": None}):
+        from vllm.triton_utils import HAS_TRITON, tl, triton
+        assert HAS_TRITON is False
+        assert triton.__class__.__name__ == "TritonPlaceholder"
+        assert triton.language.__class__.__name__ == "TritonLanguagePlaceholder"
+        assert tl.__class__.__name__ == "TritonLanguagePlaceholder"

vllm/attention/ops/blocksparse_attention/blocksparse_attention_kernel.py

@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import torch
-import triton
-import triton.language as tl
+
+from vllm.triton_utils import tl, triton
 
 
 def blocksparse_flash_attn_varlen_fwd(

vllm/attention/ops/blocksparse_attention/utils.py

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

vllm/attention/ops/chunked_prefill_paged_decode.py

@@ -7,11 +7,10 @@
 #  - Thomas Parnell <[email protected]>
 
 import torch
-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 tl, triton
 
 from .prefix_prefill import context_attention_fwd
 

vllm/attention/ops/prefix_prefill.py

@@ -4,10 +4,9 @@
 # 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.platforms import current_platform
+from vllm.triton_utils import tl, triton
 
 # Static kernels parameters
 BASE_BLOCK = 128 if current_platform.has_device_capability(80) else 64

vllm/attention/ops/triton_decode_attention.py

@@ -30,10 +30,8 @@
 
 import logging
 
-import triton
-import triton.language as tl
-
 from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
 
 is_hip_ = current_platform.is_rocm()