A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[gpt-oss] guard import when triton kernel is not installed (#22529)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Yongye Zhu committed
Aug 8, 2025 at 11:18 UTC
f756a682d96ba1824b6a759017f9d27a7f5f0182
1 file changed
+19
-10
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+19
-10
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any, Optional
3
+from typing import TYPE_CHECKING, Any, Optional
4
5
import torch
6
@@ -8,13 +8,16 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
9
TopKWeightAndReduceDelegate)
10
from vllm.model_executor.layers.fused_moe.utils import extract_required_args
11
+from vllm.utils import has_triton_kernels
12
12
-if True:
13
+if has_triton_kernels():
14
import triton_kernels.swiglu
14
- from triton_kernels.matmul_ogs import (FnSpecs, FusedActivation,
15
- PrecisionConfig, matmul_ogs)
15
+ from triton_kernels.matmul_ogs import FnSpecs, FusedActivation, matmul_ogs
16
from triton_kernels.routing import routing
17
18
+if TYPE_CHECKING:
19
+ from triton_kernels.matmul_ogs import PrecisionConfig
20
+
21
22
def triton_kernel_moe_forward(
23
hidden_states: torch.Tensor,
@@ -33,8 +36,8 @@ def triton_kernel_moe_forward(
36
w2_scale: Optional[torch.Tensor] = None,
37
w1_bias: Optional[torch.Tensor] = None,
38
w2_bias: Optional[torch.Tensor] = None,
36
- w1_precision=None, # PrecisionConfig or None
37
- w2_precision=None, # PrecisionConfig or None
39
+ w1_precision: Optional["PrecisionConfig"] = None,
40
+ w2_precision: Optional["PrecisionConfig"] = None,
41
a1_scale: Optional[torch.Tensor] = None,
42
a2_scale: Optional[torch.Tensor] = None,
43
block_shape: Optional[list[int]] = None,
@@ -90,8 +93,8 @@ def triton_kernel_fused_experts(
93
w2_scale: Optional[torch.Tensor] = None,
94
w1_bias: Optional[torch.Tensor] = None,
95
w2_bias: Optional[torch.Tensor] = None,
93
- w1_precision=None, # PrecisionConfig or None
94
- w2_precision=None, # PrecisionConfig or None
96
+ w1_precision: Optional["PrecisionConfig"] = None,
97
+ w2_precision: Optional["PrecisionConfig"] = None,
98
a1_scale: Optional[torch.Tensor] = None,
99
a2_scale: Optional[torch.Tensor] = None,
100
block_shape: Optional[list[int]] = None,
@@ -141,8 +144,14 @@ def triton_kernel_fused_experts(
144
145
class BatchedOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
146
144
- def __init__(self, quant_config, max_num_tokens: int, num_dispatchers: int,
145
- w1_precision: PrecisionConfig, w2_precision: PrecisionConfig):
147
+ def __init__(
148
+ self,
149
+ quant_config,
150
+ max_num_tokens: int,
151
+ num_dispatchers: int,
152
+ w1_precision: "PrecisionConfig",
153
+ w2_precision: "PrecisionConfig",
154
+ ):
155
super().__init__(quant_config)
156
self.max_num_tokens = max_num_tokens
157
self.num_dispatchers = num_dispatchers