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
[Bugfix] Fix accuracy issue of TRTLLM FP8 MOE and improve logging (#25895)
Signed-off-by: Pavani Majety <pmajety@nvidia.com>
Pavani Majety committed
Sep 30, 2025 at 07:51 UTC
ef283548f75122b0e8ce49ce2548fbb49446d7c7
2 files changed
+29
-17
vllm/model_executor/layers/quantization/fp8.py
+23
-16
@@ -434,14 +434,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
434
self.weight_block_size = self.quant_config.weight_block_size
435
self.block_quant = self.weight_block_size is not None
436
437
- self.flashinfer_moe_backend: Optional[FlashinferMoeBackend] = None
437
self.fused_experts: Optional[
438
mk.FusedMoEModularKernel] = None # type: ignore
440
- if envs.VLLM_USE_FLASHINFER_MOE_FP8 and has_flashinfer_moe():
441
- self.flashinfer_moe_backend = get_flashinfer_moe_backend()
442
- logger.info_once(
443
- f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"
444
- )
439
+
440
# For GPUs that lack FP8 hardware support, we can leverage the Marlin
441
# kernel for fast weight-only FP8 quantization
442
self.use_marlin = (not current_platform.has_device_capability(89)
@@ -450,14 +445,27 @@ class Fp8MoEMethod(FusedMoEMethodBase):
445
if current_platform.is_rocm():
446
self.use_marlin = False
447
448
+ # First check for Flashinfer MOE on Blackwell GPUs
449
+ self.flashinfer_moe_backend: Optional[FlashinferMoeBackend] = None
450
+ if (current_platform.is_cuda()
451
+ and current_platform.is_device_capability(100)
452
+ and envs.VLLM_USE_FLASHINFER_MOE_FP8 and has_flashinfer_moe()):
453
+ self.flashinfer_moe_backend = get_flashinfer_moe_backend()
454
+ logger.info_once(
455
+ f"Detected Blackwell GPUs, using FlashInfer "
456
+ f"{self.flashinfer_moe_backend.value} kernels for FP8 MOE.")
457
+
458
# Check for DeepGemm support.
459
self.allow_deep_gemm = False
460
if envs.VLLM_USE_DEEP_GEMM:
461
if not has_deep_gemm():
462
logger.warning_once("Failed to import DeepGemm kernels.")
463
elif not self.block_quant:
459
- logger.warning_once("Model is not block quantized. Not using "
460
- "DeepGemm kernels")
464
+ logger.warning_once("Model is not block quantized. Not using"
465
+ " DeepGemm kernels")
466
+ elif self.flashinfer_moe_backend:
467
+ logger.info_once("DeepGemm disabled: FlashInfer MOE is"
468
+ " enabled.")
469
elif (is_deep_gemm_supported()):
470
logger.info_once("Using DeepGemm kernels for Fp8MoEMethod.")
471
self.allow_deep_gemm = True
@@ -471,15 +479,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):
479
logger.debug_once("Model is not block quantized. Not using "
480
"CutlassBlockScaledGroupedGemm kernels")
481
elif (current_platform.is_cuda()
474
- and current_platform.is_device_capability(100)):
482
+ and current_platform.is_device_capability(100)
483
+ and not self.flashinfer_moe_backend):
484
logger.info_once(
476
- "Using CutlassBlockScaledGroupedGemm kernels for Fp8MoEMethod."
477
- )
485
+ "Using CutlassBlockScaledGroupedGemm kernels for Fp8 MOE "
486
+ "on SM100.")
487
self.allow_cutlass_block_scaled_grouped_gemm = True
479
- else:
480
- logger.warning_once(
481
- "CutlassBlockScaledGroupedGemm not supported on the current "
482
- "platform.")
488
489
def create_weights(self, layer: Module, num_experts: int, hidden_size: int,
490
intermediate_size_per_partition: int,
@@ -934,7 +939,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
939
import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
940
assert (renormalize and use_grouped_topk
941
and custom_routing_function is None)
937
- result = torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
942
+ e_score_correction_bias = (e_score_correction_bias.to(
943
+ x.dtype) if e_score_correction_bias is not None else None)
944
+ return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
945
routing_logits=router_logits.to(torch.float32),
946
routing_bias=e_score_correction_bias,
947
x=x,
vllm/utils/deep_gemm.py
+6
-1
@@ -27,7 +27,8 @@ def is_deep_gemm_supported() -> bool:
27
is_supported_arch = current_platform.is_cuda() and (
28
current_platform.is_device_capability(90)
29
or current_platform.is_device_capability(100))
30
- return envs.VLLM_USE_DEEP_GEMM and has_deep_gemm() and is_supported_arch
30
+ return (envs.VLLM_USE_DEEP_GEMM and has_deep_gemm() and is_supported_arch
31
+ and not envs.VLLM_USE_FLASHINFER_MOE_FP8)
32
33
34
@functools.cache
@@ -46,6 +47,10 @@ def is_deep_gemm_e8m0_used() -> bool:
47
logger.info_once("DeepGEMM E8M0 disabled: _fp8_gemm_nt_impl not found")
48
return False
49
50
+ if envs.VLLM_USE_FLASHINFER_MOE_FP8:
51
+ logger.info_once("DeepGEMM E8M0 disabled: FlashInfer MOE is enabled.")
52
+ return False
53
+
54
if current_platform.is_device_capability(100) and \
55
envs.VLLM_USE_DEEP_GEMM_E8M0:
56
logger.info_once("DeepGEMM E8M0 enabled on Blackwell GPU.")