Fix DeepSeek V4 SM10 sparse MLA default

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,9 +1,6 @@
 # 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
@@ -36,40 +33,6 @@ 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
@@ -261,13 +224,16 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
             weight_layout = WeightLayout.BlockMajorK
             hidden_states_scale = a1q_scale.t().contiguous()
 
-        kwargs = dict(
+        flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe(
             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,
@@ -284,14 +250,6 @@ 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):
@@ -444,6 +402,9 @@ 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,
@@ -461,13 +422,6 @@ 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,8 +761,9 @@ 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 while the other
-    CUDA arches keep the FlashMLA path.
+    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.
     """
     backend = vllm_config.attention_config.backend
     device_capability = current_platform.get_device_capability()
@@ -785,6 +786,13 @@ 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