From 6850839c6f5a33e2c856c7a49ba1dc2a5b42508b Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Wed, 10 Jun 2026 15:08:41 +0200 Subject: [PATCH] [Perf] Fix dsv3_router_gemm heuristic (#44217) Signed-off-by: LopezCastroRoberto --- .../layers/fused_moe/router/gate_linear.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index 0a57a6f4df..5867ce3e9a 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -44,11 +44,10 @@ class GateLinear(ReplicatedLinear): force_fp32_compute: bool = False, prefix: str = "", ): - is_hopper_or_blackwell = current_platform.is_device_capability( - (9, 0) - ) or current_platform.is_device_capability_family(100) + is_hopper = current_platform.is_device_capability((9, 0)) + is_blackwell = current_platform.is_device_capability_family(100) can_use_specialized_kernels = ( - current_platform.is_cuda() and is_hopper_or_blackwell and not bias + current_platform.is_cuda() and (is_hopper or is_blackwell) and not bias ) # If fp32 compute is required and no specialized kernel is available, @@ -73,13 +72,16 @@ class GateLinear(ReplicatedLinear): and output_size in self.DSV3_SUPPORTED_NUM_EXPERTS and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES ) + # See https://github.com/vllm-project/vllm/pull/44217 + # for more details. + self._dsv3_max_batch = 16 if is_hopper else 8 # fp32 specialized kernel eligibility (SM90+, exact dims, fp32 weight) self.allow_fp32_router_gemm = ( not bias and self.weight.dtype == torch.float32 and current_platform.is_cuda() - and is_hopper_or_blackwell + and (is_hopper or is_blackwell) and output_size in self.FP32_SUPPORTED_NUM_EXPERTS and input_size in self.FP32_SUPPORTED_HIDDEN_SIZES ) @@ -112,7 +114,7 @@ class GateLinear(ReplicatedLinear): self, x: torch.Tensor ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: # Tier 1: DSV3 specialized kernel - if self.allow_dsv3_router_gemm and x.shape[0] <= 16: + if self.allow_dsv3_router_gemm and x.shape[0] <= self._dsv3_max_batch: output = ops.dsv3_router_gemm( hidden_states=x, router_weight=self.weight,