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 6, 04:30 UTC
[Bug][MoE] Strengthen _supports_current_device() checks in the TRTLLM FP8, NVFP4, and FlashInfer CuteDSL MoE experts (#36728)
Signed-off-by: Yifan Zong <yzong@redhat.com>
yzong-rh committed
Mar 23, 2026 at 17:02 UTC
e85f8f0932f65c17d64c9bb8e85a57bb007b369d
4 files changed
+19
-5
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
+6
-1
@@ -23,6 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
from vllm.platforms import current_platform
24
from vllm.utils.flashinfer import (
25
flashinfer_cutedsl_grouped_gemm_nt_masked,
26
+ has_flashinfer_cutedsl_grouped_gemm_nt_masked,
27
scaled_fp4_grouped_quantize,
28
silu_and_mul_scaled_nvfp4_experts_quantize,
29
)
@@ -60,7 +61,11 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
61
@staticmethod
62
def _supports_current_device() -> bool:
63
p = current_platform
63
- return p.is_cuda() and p.is_device_capability_family(100)
64
+ return (
65
+ p.is_cuda()
66
+ and p.is_device_capability_family(100)
67
+ and has_flashinfer_cutedsl_grouped_gemm_nt_masked()
68
+ )
69
70
@staticmethod
71
def _supports_no_act_and_mul() -> bool:
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+6
-2
@@ -27,6 +27,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
27
kMxfp8Static,
28
)
29
from vllm.platforms import current_platform
30
+from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
31
32
logger = init_logger(__name__)
33
@@ -61,8 +62,11 @@ class TrtLlmFp8ExpertsBase:
62
def _supports_current_device() -> bool:
63
"""Supports only Blackwell-family GPUs."""
64
p = current_platform
64
- # Add check flashinfer trtllm is available
65
- return p.is_cuda() and p.is_device_capability_family(100)
65
+ return (
66
+ p.is_cuda()
67
+ and p.is_device_capability_family(100)
68
+ and has_flashinfer_trtllm_fused_moe()
69
+ )
70
71
@staticmethod
72
def _supports_no_act_and_mul() -> bool:
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+6
-1
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
kNvfp4Static,
25
)
26
from vllm.platforms import current_platform
27
+from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
28
29
30
class TrtLlmNvFp4ExpertsBase:
@@ -80,7 +81,11 @@ class TrtLlmNvFp4ExpertsBase:
81
def _supports_current_device() -> bool:
82
"""Supports only Blackwell-family GPUs."""
83
p = current_platform
83
- return p.is_cuda() and p.is_device_capability_family(100)
84
+ return (
85
+ p.is_cuda()
86
+ and p.is_device_capability_family(100)
87
+ and has_flashinfer_trtllm_fused_moe()
88
+ )
89
90
@staticmethod
91
def _supports_no_act_and_mul() -> bool:
vllm/utils/flashinfer.py
+1
-1
@@ -244,7 +244,7 @@ def has_flashinfer_cutedsl_grouped_gemm_nt_masked() -> bool:
244
required_functions = [
245
("flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked"),
246
("flashinfer", "scaled_fp4_grouped_quantize"),
247
- ("flashinfer", "silu_and_scaled_nvfp4_experts_quantize"),
247
+ ("flashinfer", "silu_and_mul_scaled_nvfp4_experts_quantize"),
248
]
249
250
for module_name, attr_name in required_functions: