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 2, 17:01 UTC
[Bugfix][Kernel] TRTLLM NVFP4 MoE chunking (#43599)
Signed-off-by: amitz-nv <203509407+amitz-nv@users.noreply.github.com>
amitz-nv committed
May 28, 2026 at 03:36 UTC
381edde1b9bf6e4efce1ef2008d68ecdd47f415e
4 files changed
+74
-21
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
-3
@@ -99,9 +99,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
99
) -> bool:
100
return True
101
102
- def supports_chunking(self) -> bool:
103
- return False
104
-
102
def supports_expert_map(self) -> bool:
103
return False
104
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
-3
@@ -88,9 +88,6 @@ class TrtLlmFp8ExpertsBase:
88
or moe_parallel_config.use_ag_rs_all2all_kernels
89
) and not moe_parallel_config.enable_eplb
90
91
- def supports_chunking(self) -> bool:
92
- return False
93
-
91
def supports_expert_map(self) -> bool:
92
return False
93
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
-3
@@ -113,9 +113,6 @@ class TrtLlmMxfp4ExpertsBase:
113
def activation_format() -> mk.FusedMoEActivationFormat:
114
return mk.FusedMoEActivationFormat.Standard
115
116
- def supports_chunking(self) -> bool:
117
- return False
118
-
116
def supports_expert_map(self) -> bool:
117
return False
118
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+74
-12
@@ -157,8 +157,27 @@ class TrtLlmNvFp4ExpertsBase:
157
def activation_format() -> mk.FusedMoEActivationFormat:
158
return mk.FusedMoEActivationFormat.Standard
159
160
- def supports_chunking(self) -> bool:
161
- return False
160
+ def _get_chunk_size(self) -> int:
161
+ MAX_GRID_Y = 65535
162
+ MAX_TILE_TOKENS_DIM = 128
163
+
164
+ def _calc_max_supported_tokens(top_k: int, num_experts: int) -> int:
165
+ """Calculates the max number of supported tokens, so the CUDA grid.Y limit
166
+ won't be reached.
167
+ Based on getMaxNumCtasInBatchDim function in flashinfer's TRTLLM MoE runner:
168
+ https://github.com/flashinfer-ai/flashinfer/blob/719ee23fd82cb220d51ad118ca60198718f6c9d1/include/flashinfer/trtllm/fused_moe/runner.h#L97
169
+ Which given numTokens, topK, numExperts, tileTokensDim calculates maxNumCtas
170
+ which is used as the CUDA grid.Y dimension, which we want to
171
+ be <= MAX_GRID_Y. Solving for numTokens gives the formula below.
172
+ """
173
+ return (
174
+ num_experts + (MAX_GRID_Y - num_experts + 1) * MAX_TILE_TOKENS_DIM - 1
175
+ ) // top_k
176
+
177
+ # Using 305k or more causes IMA error in the kernel, so limit to 300k.
178
+ return min(
179
+ 300000, _calc_max_supported_tokens(self.topk, self.moe_config.num_experts)
180
+ )
181
182
def supports_expert_map(self) -> bool:
183
return False
@@ -199,7 +218,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
218
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
219
return TopKWeightAndReduceNoOP()
220
202
- def apply(
221
+ def _invoke_kernel(
222
self,
223
output: torch.Tensor,
224
hidden_states: torch.Tensor,
@@ -209,18 +228,10 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
228
topk_ids: torch.Tensor,
229
activation: MoEActivation,
230
global_num_experts: int,
212
- expert_map: torch.Tensor | None,
213
- a1q_scale: torch.Tensor | None,
214
- a2_scale: torch.Tensor | None,
215
- workspace13: torch.Tensor,
216
- workspace2: torch.Tensor,
217
- expert_tokens_meta: mk.ExpertTokensMetadata | None,
218
- apply_router_weight_on_input: bool,
231
+ a1q_scale: torch.Tensor,
232
):
233
import flashinfer
234
222
- assert self._supports_activation(activation)
223
- assert a1q_scale is not None
235
assert self.quant_config.w1_scale is not None
236
assert self.quant_config.w2_scale is not None
237
@@ -262,6 +273,57 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
273
output=output,
274
)
275
276
+ def apply(
277
+ self,
278
+ output: torch.Tensor,
279
+ hidden_states: torch.Tensor,
280
+ w1: torch.Tensor,
281
+ w2: torch.Tensor,
282
+ topk_weights: torch.Tensor,
283
+ topk_ids: torch.Tensor,
284
+ activation: MoEActivation,
285
+ global_num_experts: int,
286
+ expert_map: torch.Tensor | None,
287
+ a1q_scale: torch.Tensor | None,
288
+ a2_scale: torch.Tensor | None,
289
+ workspace13: torch.Tensor,
290
+ workspace2: torch.Tensor,
291
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
292
+ apply_router_weight_on_input: bool,
293
+ ):
294
+ assert self._supports_activation(activation)
295
+ assert a1q_scale is not None
296
+
297
+ M = hidden_states.shape[0]
298
+ chunk_size = self._get_chunk_size()
299
+
300
+ if chunk_size >= M:
301
+ self._invoke_kernel(
302
+ output,
303
+ hidden_states,
304
+ w1,
305
+ w2,
306
+ topk_weights,
307
+ topk_ids,
308
+ activation,
309
+ global_num_experts,
310
+ a1q_scale,
311
+ )
312
+ else:
313
+ for start in range(0, M, chunk_size):
314
+ end = min(start + chunk_size, M)
315
+ self._invoke_kernel(
316
+ output[start:end],
317
+ hidden_states[start:end],
318
+ w1,
319
+ w2,
320
+ topk_weights[start:end],
321
+ topk_ids[start:end],
322
+ activation,
323
+ global_num_experts,
324
+ a1q_scale[start:end],
325
+ )
326
+
327
328
class TrtLlmNvFp4ExpertsMonolithic(
329
TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsMonolithic