A high-throughput and memory-efficient inference and serving engine for LLMs
[Perf] Fix dsv3_router_gemm heuristic (#44217)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Roberto L. Castro committed
Jun 10, 2026 at 15:08 UTC
6850839c6f5a33e2c856c7a49ba1dc2a5b42508b
1 file changed
+8
-6
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+8
-6
@@ -44,11 +44,10 @@ class GateLinear(ReplicatedLinear):
44
force_fp32_compute: bool = False,
45
prefix: str = "",
46
):
47
- is_hopper_or_blackwell = current_platform.is_device_capability(
48
- (9, 0)
49
- ) or current_platform.is_device_capability_family(100)
47
+ is_hopper = current_platform.is_device_capability((9, 0))
48
+ is_blackwell = current_platform.is_device_capability_family(100)
49
can_use_specialized_kernels = (
51
- current_platform.is_cuda() and is_hopper_or_blackwell and not bias
50
+ current_platform.is_cuda() and (is_hopper or is_blackwell) and not bias
51
)
52
53
# If fp32 compute is required and no specialized kernel is available,
@@ -73,13 +72,16 @@ class GateLinear(ReplicatedLinear):
72
and output_size in self.DSV3_SUPPORTED_NUM_EXPERTS
73
and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES
74
)
75
+ # See https://github.com/vllm-project/vllm/pull/44217
76
+ # for more details.
77
+ self._dsv3_max_batch = 16 if is_hopper else 8
78
79
# fp32 specialized kernel eligibility (SM90+, exact dims, fp32 weight)
80
self.allow_fp32_router_gemm = (
81
not bias
82
and self.weight.dtype == torch.float32
83
and current_platform.is_cuda()
82
- and is_hopper_or_blackwell
84
+ and (is_hopper or is_blackwell)
85
and output_size in self.FP32_SUPPORTED_NUM_EXPERTS
86
and input_size in self.FP32_SUPPORTED_HIDDEN_SIZES
87
)
@@ -112,7 +114,7 @@ class GateLinear(ReplicatedLinear):
114
self, x: torch.Tensor
115
) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
116
# Tier 1: DSV3 specialized kernel
115
- if self.allow_dsv3_router_gemm and x.shape[0] <= 16:
117
+ if self.allow_dsv3_router_gemm and x.shape[0] <= self._dsv3_max_batch:
118
output = ops.dsv3_router_gemm(
119
hidden_states=x,
120
router_weight=self.weight,