Handle older FlashInfer FP8 MoE signatures

Co-authored-by: OpenAI Codex <codex@openai.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import functools
+import inspect
+
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -33,6 +36,40 @@ from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
 logger = init_logger(__name__)
 
 
+@functools.cache
+def _flashinfer_moe_supports_swiglu_params(fn_name: str) -> bool:
+    try:
+        import flashinfer
+
+        fn = getattr(flashinfer.fused_moe, fn_name)
+        return "gemm1_alpha" in inspect.signature(fn).parameters
+    except (TypeError, ValueError):
+        # If FlashInfer stops exposing an inspectable Python signature, keep
+        # passing the parameters so an incompatible install fails explicitly.
+        return True
+
+
+def _add_swiglu_params_if_supported(
+    kwargs: dict[str, object],
+    fn_name: str,
+    gemm1_alpha: torch.Tensor | None,
+    gemm1_beta: torch.Tensor | None,
+    gemm1_clamp_limit: torch.Tensor | None,
+) -> None:
+    swiglu_params = {
+        "gemm1_alpha": gemm1_alpha,
+        "gemm1_beta": gemm1_beta,
+        "gemm1_clamp_limit": gemm1_clamp_limit,
+    }
+    if _flashinfer_moe_supports_swiglu_params(fn_name):
+        kwargs.update(swiglu_params)
+    elif any(param is not None for param in swiglu_params.values()):
+        raise RuntimeError(
+            "The installed FlashInfer TRTLLM FP8 MoE kernel does not support "
+            "per-expert SwiGLU parameters. Please upgrade FlashInfer."
+        )
+
+
 class TrtLlmFp8ExpertsBase:
     """
     Fp8 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic
@@ -224,16 +261,13 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
             weight_layout = WeightLayout.BlockMajorK
             hidden_states_scale = a1q_scale.t().contiguous()
 
-        flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe(
+        kwargs = dict(
             topk_ids=packed_topk_ids,
             routing_bias=None,
             hidden_states=hidden_states,
             hidden_states_scale=hidden_states_scale,
             gemm1_weights=w1,
             gemm1_weights_scale=self.quant_config.w1_scale,
-            gemm1_alpha=self.gemm1_alpha,
-            gemm1_beta=self.gemm1_beta,
-            gemm1_clamp_limit=self.gemm1_clamp_limit,
             gemm2_weights=w2,
             gemm2_weights_scale=self.quant_config.w2_scale,
             num_experts=global_num_experts,
@@ -250,6 +284,14 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
             fp8_quantization_type=fp8_quant_type,
             output=output,
         )
+        _add_swiglu_params_if_supported(
+            kwargs,
+            "trtllm_fp8_block_scale_routed_moe",
+            self.gemm1_alpha,
+            self.gemm1_beta,
+            self.gemm1_clamp_limit,
+        )
+        flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe(**kwargs)
 
 
 class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolithic):
@@ -402,9 +444,6 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
             hidden_states_scale=hidden_states_scale,
             gemm1_weights=w1,
             gemm1_weights_scale=self.quant_config.w1_scale,
-            gemm1_alpha=self.gemm1_alpha,
-            gemm1_beta=self.gemm1_beta,
-            gemm1_clamp_limit=self.gemm1_clamp_limit,
             gemm2_weights=w2,
             gemm2_weights_scale=self.quant_config.w2_scale,
             num_experts=global_num_experts,
@@ -422,6 +461,13 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
         )
         if is_mxfp8 or activation == MoEActivation.RELU2_NO_MUL:
             kwargs["activation_type"] = activation_type
+        _add_swiglu_params_if_supported(
+            kwargs,
+            "trtllm_fp8_block_scale_moe",
+            self.gemm1_alpha,
+            self.gemm1_beta,
+            self.gemm1_clamp_limit,
+        )
         return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(**kwargs)
 
     def _apply_per_tensor(

vllm/models/deepseek_v4/nvidia/model.py

@@ -761,9 +761,8 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
 
     The generic CUDA backend selector does not instantiate DSv4 layers directly,
     so map generic sparse-MLA choices to the DSv4-specialized attention class.
-    Without an explicit backend, SM12 defaults to FlashInfer. SM10x also uses
-    FlashInfer for plain FP8 KV cache because the FlashMLA sparse kernel only
-    accepts the older 656-byte row layout, not DeepSeek V4's 584-byte row.
+    Without an explicit backend, SM12 defaults to FlashInfer while the other
+    CUDA arches keep the FlashMLA path.
     """
     backend = vllm_config.attention_config.backend
     device_capability = current_platform.get_device_capability()
@@ -786,13 +785,6 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
     ):
         return DeepseekV4FlashMLAAttention
 
-    cache_dtype = vllm_config.cache_config.cache_dtype
-    if (
-        device_capability is not None
-        and device_capability.major == 10
-        and cache_dtype in ("fp8", "fp8_e4m3", "fp8_e5m2")
-    ):
-        return DeepseekV4FlashInferMLAAttention
     if device_capability is not None and device_capability.major == 12:
         return DeepseekV4FlashInferSM120Attention
     return DeepseekV4FlashMLAAttention