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 3, 11:45 UTC
[Perf] Wire silu_and_mul_per_block_quant into TritonFP8MoE (MiniMax-M2) (#42497)
Signed-off-by: qianlihuang <yiliu.dong@qq.com> Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Co-authored-by: qianlihuang <yiliu.dong@qq.com>
Yiliu Dong committed
May 18, 2026 at 09:57 UTC
03ddc1c9bc5e448e0da6236268a611d7d001dbae
1 file changed
+31
-12
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+31
-12
@@ -31,6 +31,9 @@ from vllm.model_executor.layers.fused_moe.utils import (
31
_resize_cache,
32
moe_kernel_quantize_input,
33
)
34
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
35
+ is_deep_gemm_e8m0_used,
36
+)
37
from vllm.model_executor.layers.quantization.utils.quant_utils import (
38
QuantKey,
39
kFp8Dynamic128Sym,
@@ -283,20 +286,36 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
286
top_k_num=top_k_num,
287
)
288
286
- self.activation(
287
- activation, intermediate_cache2, intermediate_cache1.view(-1, N)
288
- )
289
-
289
a2q_scale: torch.Tensor | None = None
290
292
- qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
293
- intermediate_cache2,
294
- a2_scale,
295
- self.quant_dtype,
296
- self.per_act_token_quant,
297
- self.block_shape,
298
- quantization_emulation=self.quantization_emulation,
299
- )
291
+ # Fuse SiLU+Mul + FP8 block quantize into a single kernel
292
+ # when conditions permit (gated SiLU, fp8 block quant with
293
+ # group_size=128, no LoRA requiring the BF16 intermediate).
294
+ if (
295
+ activation == MoEActivation.SILU
296
+ and self.quant_config.use_fp8_w8a8
297
+ and self.block_shape == [128, 128]
298
+ and lora_context is None
299
+ and not is_deep_gemm_e8m0_used()
300
+ ):
301
+ qintermediate_cache2, a2q_scale = ops.silu_and_mul_per_block_quant(
302
+ intermediate_cache1.view(-1, N),
303
+ group_size=128,
304
+ quant_dtype=current_platform.fp8_dtype(),
305
+ )
306
+ else:
307
+ self.activation(
308
+ activation, intermediate_cache2, intermediate_cache1.view(-1, N)
309
+ )
310
+
311
+ qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
312
+ intermediate_cache2,
313
+ a2_scale,
314
+ self.quant_dtype,
315
+ self.per_act_token_quant,
316
+ self.block_shape,
317
+ quantization_emulation=self.quantization_emulation,
318
+ )
319
320
invoke_fused_moe_triton_kernel(
321
qintermediate_cache2,