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 12, 02:15 UTC
[Kernels] Remove BatchedTritonOrDeepGemmExperts and default fallback to Triton (#29929)
Signed-off-by: Bill Nell <bnell@redhat.com> Signed-off-by: bnellnm <49004751+bnellnm@users.noreply.github.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>
bnellnm committed
Dec 3, 2025 at 15:49 UTC
2902c348265639de300c95cbcae1c26486f57ac7
5 files changed
+46
-217
docs/design/moe_kernel_features.md
+1
-2
@@ -90,7 +90,6 @@ To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels
90
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`cutlass_moe_fp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.cutlass_moe_fp8],</br>[`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
91
| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`flashinfer_cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.flashinfer_cutlass_moe_fp4],</br>[`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
92
| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
93
-| deep gemm+triton<sup>2</sup> | standard,</br>batched | all<sup>1</sup> | G(128),A,T | silu, gelu | <sup>6</sup> | Y | [`TritonOrDeepGemmExperts`][vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe.TritonOrDeepGemmExperts],</br>[`BatchedTritonOrDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe.BatchedTritonOrDeepGemmExperts] |
93
| marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
94
| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts] |
95
| pallas | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_pallas.fused_moe] |
@@ -114,5 +113,5 @@ The following table shows "families" of modular kernels that are intended to wor
113
| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEPermuteExpertsUnpermute` subclasses |
114
|---------|-----------------------------------------|----------------------------------------------|
115
| deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,</br>`TritonExperts`,</br>`TritonOrDeepGemmExperts`,</br>`CutlassExpertsFp8`, </br>`MarlinExperts` |
117
-| deepep_low_latency,</br>pplx | `DeepEPLLPrepareAndFinalize`,</br>`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,</br>`BatchedTritonExperts`,</br>`BatchedTritonOrDeepGemmExperts`,</br>`CutlassBatchedExpertsFp8`,</br>`BatchedMarlinExperts` |
116
+| deepep_low_latency,</br>pplx | `DeepEPLLPrepareAndFinalize`,</br>`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,</br>`BatchedTritonExperts`,</br>`CutlassBatchedExpertsFp8`,</br>`BatchedMarlinExperts` |
117
| flashinfer | `FlashInferCutlassMoEPrepareAndFinalize` | `FlashInferExperts` |
tests/kernels/moe/modular_kernel_tools/mk_objects.py
-17
@@ -13,9 +13,6 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
13
from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
14
BatchedDeepGemmExperts,
15
)
16
-from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import (
17
- BatchedTritonOrDeepGemmExperts,
18
-)
16
from vllm.model_executor.layers.fused_moe.config import (
17
FusedMoEConfig,
18
FusedMoEQuantConfig,
@@ -286,16 +283,6 @@ if has_deep_gemm() and is_deep_gemm_supported():
283
needs_matching_quant=False,
284
needs_deep_gemm=True,
285
)
289
- register_experts(
290
- BatchedTritonOrDeepGemmExperts,
291
- batched_format,
292
- common_float_and_int_types,
293
- blocked_quantization_support=True,
294
- supports_chunking=False,
295
- supports_expert_map=False,
296
- needs_matching_quant=True,
297
- needs_deep_gemm=True,
298
- )
286
register_experts(
287
TritonOrDeepGemmExperts,
288
standard_format,
@@ -457,10 +444,6 @@ def make_fused_experts(
444
kwargs = batch_kwargs | quant_kwargs
445
print(f"Making BatchedTritonExperts {kwargs} ...")
446
experts = BatchedTritonExperts(**kwargs)
460
- elif fused_experts_type == BatchedTritonOrDeepGemmExperts:
461
- kwargs = batch_kwargs | quant_kwargs | deepgemm_kwargs
462
- print(f"Making BatchedTritonOrDeepGemmExperts {kwargs} ...")
463
- experts = BatchedTritonOrDeepGemmExperts(**kwargs)
447
elif fused_experts_type == DeepGemmExperts:
448
print(f"Making DeepGemmExperts {quant_config} ...")
449
experts = DeepGemmExperts(quant_config)
vllm/model_executor/layers/fused_moe/__init__.py
-4
@@ -60,9 +60,6 @@ if HAS_TRITON:
60
from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
61
BatchedDeepGemmExperts,
62
)
63
- from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501
64
- BatchedTritonOrDeepGemmExperts,
65
- )
63
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
64
CutlassBatchedExpertsFp8,
65
CutlassExpertsFp8,
@@ -98,7 +95,6 @@ if HAS_TRITON:
95
"DeepGemmExperts",
96
"BatchedDeepGemmExperts",
97
"TritonOrDeepGemmExperts",
101
- "BatchedTritonOrDeepGemmExperts",
98
]
99
else:
100
# Some model classes directly use the custom ops. Add placeholders
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
deleted
-180
@@ -1,180 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import torch
5
-
6
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
-from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
8
- BatchedDeepGemmExperts,
9
-)
10
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11
-from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts
12
-from vllm.utils.deep_gemm import get_mk_alignment_for_contiguous_layout
13
-
14
-
15
-class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
16
- def __init__(
17
- self,
18
- max_num_tokens: int,
19
- num_dispatchers: int,
20
- quant_config: FusedMoEQuantConfig,
21
- allow_deep_gemm: bool = False,
22
- ):
23
- super().__init__(quant_config)
24
-
25
- self.batched_triton_experts = BatchedTritonExperts(
26
- max_num_tokens=max_num_tokens,
27
- num_dispatchers=num_dispatchers,
28
- quant_config=self.quant_config,
29
- )
30
-
31
- self.allow_deep_gemm = (
32
- allow_deep_gemm
33
- and self.quant_config.use_fp8_w8a8
34
- and self.block_shape == get_mk_alignment_for_contiguous_layout()
35
- )
36
-
37
- self.batched_deep_gemm_experts = (
38
- BatchedDeepGemmExperts(
39
- max_num_tokens=max_num_tokens,
40
- num_dispatchers=num_dispatchers,
41
- quant_config=self.quant_config,
42
- )
43
- if self.allow_deep_gemm
44
- else None
45
- )
46
-
47
- assert (
48
- self.batched_deep_gemm_experts is not None
49
- or self.batched_triton_experts is not None
50
- )
51
-
52
- @property
53
- def activation_formats(
54
- self,
55
- ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
56
- if self.batched_triton_experts is not None:
57
- assert (
58
- self.batched_deep_gemm_experts is None
59
- or self.batched_deep_gemm_experts.activation_formats
60
- == self.batched_triton_experts.activation_formats
61
- )
62
- return self.batched_triton_experts.activation_formats
63
- else:
64
- assert self.batched_deep_gemm_experts is not None
65
- return self.batched_deep_gemm_experts.activation_formats
66
-
67
- def supports_chunking(self) -> bool:
68
- bdge = self.batched_deep_gemm_experts
69
- bte = self.batched_triton_experts
70
- return (bdge is None or bdge.supports_chunking()) and (
71
- bte is None or bte.supports_chunking()
72
- )
73
-
74
- def supports_expert_map(self) -> bool:
75
- bdge = self.batched_deep_gemm_experts
76
- bte = self.batched_triton_experts
77
- return (bdge is None or bdge.supports_expert_map()) and (
78
- bte is None or bte.supports_expert_map()
79
- )
80
-
81
- def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
82
- bdge = self.batched_deep_gemm_experts
83
- bte = self.batched_triton_experts
84
- bdge_war = bdge.finalize_weight_and_reduce_impl() if bdge else None
85
- bte_war = bte.finalize_weight_and_reduce_impl() if bte else None
86
- is_bdge_war = bdge_war is not None
87
- is_bte_war = bte_war is not None
88
-
89
- if is_bdge_war and is_bte_war:
90
- assert bdge_war == bte_war, (
91
- "Both implementations should agree on WeightAndReduce impls. "
92
- f"Got bdge_war: {bdge_war}, and bte_war: {bte_war}"
93
- )
94
-
95
- if bdge_war is not None:
96
- return bdge_war
97
-
98
- assert bte_war is not None
99
- return bte_war
100
-
101
- def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
102
- return act_dtype
103
-
104
- def workspace_shapes(
105
- self,
106
- M: int,
107
- N: int,
108
- K: int,
109
- topk: int,
110
- global_num_experts: int,
111
- local_num_experts: int,
112
- expert_tokens_metadata: mk.ExpertTokensMetadata | None,
113
- ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
114
- # Note: the deep gemm workspaces are strictly larger than the triton
115
- # workspaces so we can be pessimistic here and allocate for DeepGemm
116
- # even if we fall back to triton later, e.g. if expert maps are set.
117
- if self.allow_deep_gemm:
118
- assert self.batched_deep_gemm_experts is not None
119
- return self.batched_deep_gemm_experts.workspace_shapes(
120
- M,
121
- N,
122
- K,
123
- topk,
124
- global_num_experts,
125
- local_num_experts,
126
- expert_tokens_metadata,
127
- )
128
- else:
129
- assert self.batched_triton_experts is not None
130
- return self.batched_triton_experts.workspace_shapes(
131
- M,
132
- N,
133
- K,
134
- topk,
135
- global_num_experts,
136
- local_num_experts,
137
- expert_tokens_metadata,
138
- )
139
-
140
- def apply(
141
- self,
142
- output: torch.Tensor,
143
- hidden_states: torch.Tensor,
144
- w1: torch.Tensor,
145
- w2: torch.Tensor,
146
- topk_weights: torch.Tensor,
147
- topk_ids: torch.Tensor,
148
- activation: str,
149
- global_num_experts: int,
150
- expert_map: torch.Tensor | None,
151
- a1q_scale: torch.Tensor | None,
152
- a2_scale: torch.Tensor | None,
153
- workspace13: torch.Tensor,
154
- workspace2: torch.Tensor,
155
- expert_tokens_meta: mk.ExpertTokensMetadata | None,
156
- apply_router_weight_on_input: bool,
157
- ):
158
- experts = (
159
- self.batched_deep_gemm_experts
160
- if self.allow_deep_gemm
161
- else self.batched_triton_experts
162
- )
163
- assert experts is not None
164
- experts.apply(
165
- output,
166
- hidden_states,
167
- w1,
168
- w2,
169
- topk_weights,
170
- topk_ids,
171
- activation,
172
- global_num_experts,
173
- expert_map,
174
- a1q_scale,
175
- a2_scale,
176
- workspace13,
177
- workspace2,
178
- expert_tokens_meta,
179
- apply_router_weight_on_input,
180
- )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+45
-14
@@ -90,8 +90,10 @@ from vllm.platforms import CpuArchEnum, current_platform
90
from vllm.scalar_type import scalar_types
91
from vllm.utils.deep_gemm import (
92
get_col_major_tma_aligned_tensor,
93
+ get_mk_alignment_for_contiguous_layout,
94
is_deep_gemm_e8m0_used,
95
)
96
+from vllm.utils.import_utils import has_deep_gemm
97
98
logger = init_logger(__name__)
99
@@ -1088,9 +1090,11 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1090
1091
return experts
1092
1091
- # triton path
1092
- from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501
1093
- BatchedTritonOrDeepGemmExperts,
1093
+ from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
1094
+ BatchedDeepGemmExperts,
1095
+ )
1096
+ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
1097
+ BatchedTritonExperts,
1098
)
1099
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
1100
TritonOrDeepGemmExperts,
@@ -1098,6 +1102,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1102
1103
assert not self.rocm_aiter_moe_enabled and not self.use_marlin
1104
1105
+ use_deep_gemm = envs.VLLM_USE_DEEP_GEMM and envs.VLLM_MOE_USE_DEEP_GEMM
1106
+
1107
if (
1108
prepare_finalize.activation_format
1109
== FusedMoEActivationFormat.BatchedExperts
@@ -1105,22 +1111,47 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1111
max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
1112
assert max_num_tokens_per_rank is not None
1113
1108
- logger.debug("BatchedTritonExperts(%s)", self.__class__.__name__)
1109
- return BatchedTritonOrDeepGemmExperts(
1110
- max_num_tokens=max_num_tokens_per_rank,
1111
- num_dispatchers=prepare_finalize.num_dispatchers(),
1112
- quant_config=self.moe_quant_config,
1113
- allow_deep_gemm=(
1114
- envs.VLLM_USE_DEEP_GEMM and envs.VLLM_MOE_USE_DEEP_GEMM
1115
- ),
1114
+ if use_deep_gemm and not has_deep_gemm():
1115
+ raise RuntimeError(
1116
+ "DeepGEMM requested for MoE layer but not installed."
1117
+ )
1118
+
1119
+ compatible_with_deep_gemm = (
1120
+ self.moe_quant_config.use_fp8_w8a8
1121
+ and self.moe_quant_config.block_shape
1122
+ == get_mk_alignment_for_contiguous_layout()
1123
)
1124
+
1125
+ # If this MoE layer is compatible with DeepGEMM, the proper env
1126
+ # vars are set and DeepGEMM is not installed, throw an error.
1127
+ if use_deep_gemm and compatible_with_deep_gemm and not has_deep_gemm():
1128
+ raise RuntimeError(
1129
+ f"MoE layer incompatible with DeepGEMM, expected "
1130
+ f"fp8==True, got {self.moe_quant_config.use_fp8_w8a8}"
1131
+ f"or block_shape {self.moe_quant_config.block_shape}"
1132
+ f"=={get_mk_alignment_for_contiguous_layout()}."
1133
+ )
1134
+
1135
+ if use_deep_gemm and compatible_with_deep_gemm and has_deep_gemm():
1136
+ logger.debug("BatchedDeepGemmExperts(%s)", self.__class__.__name__)
1137
+ return BatchedDeepGemmExperts(
1138
+ max_num_tokens=max_num_tokens_per_rank,
1139
+ num_dispatchers=prepare_finalize.num_dispatchers(),
1140
+ quant_config=self.moe_quant_config,
1141
+ )
1142
+ else:
1143
+ logger.debug("BatchedTritonExperts(%s)", self.__class__.__name__)
1144
+ return BatchedTritonExperts(
1145
+ max_num_tokens=max_num_tokens_per_rank,
1146
+ num_dispatchers=prepare_finalize.num_dispatchers(),
1147
+ quant_config=self.moe_quant_config,
1148
+ )
1149
+
1150
else:
1151
logger.debug("TritonOrDeepGemmExperts(%s)", self.__class__.__name__)
1152
return TritonOrDeepGemmExperts(
1153
self.moe_quant_config,
1121
- allow_deep_gemm=(
1122
- envs.VLLM_USE_DEEP_GEMM and envs.VLLM_MOE_USE_DEEP_GEMM
1123
- ),
1154
+ allow_deep_gemm=use_deep_gemm,
1155
)
1156
1157
def get_fused_moe_quant_config(