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

Signed-off-by: Mengqing Cao <cmq0113@163.com>
2025-05-06 17:53:09 +08:00
parent 05e1f96419
commit f9bc5a0693
30 changed files with 165 additions and 75 deletions
--- a/vllm/model_executor/layers/lightning_attn.py
+++ b/vllm/model_executor/layers/lightning_attn.py
@@ -1,9 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 import torch
-import triton
-import triton.language as tl
 from einops import rearrange

+from vllm.triton_utils import tl, triton
+

@triton.jit
 def _fwd_diag_kernel(Q, K, V, Out, S, b: tl.constexpr, h: tl.constexpr, n,