[ROCm][Perf] Route the fused shared-expert gate GEMM through the platform dispatcher (#54185)
Signed-off-by: Mikko Tukiainen <Mikko.Tukiainen@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Mikko Tukiainen committed
Sep 21, 2026 at 20:17 UTC
17322375a2d7815ed3d7c9c3ffd718196564314e
1 file changed
+5
-2
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+5
-2
@@ -45,6 +45,7 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
45
SharedExperts,
46
SharedExpertsOrder,
47
)
48
+from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
49
from vllm.platforms import current_platform
50
from vllm.utils.torch_utils import (
51
_USE_LAYERNAME,
@@ -273,7 +274,7 @@ class MoERunner(MoERunnerInterface):
274
275
# When both gates are present and FSE is enabled, fuse their
276
# weight matrices into [num_experts + num_shared, hidden] so one
276
- # F.linear produces combined logits. The topk kernel can then
277
+ # GEMM produces combined logits. The topk kernel can then
278
# apply routing softmax and shared expert activation (sigmoid)
279
# in a single launch.
280
self._fse_fuse_gate = gate is not None and shared_expert_gate is not None
@@ -898,7 +899,9 @@ class MoERunner(MoERunnerInterface):
899
if self.gate is not None:
900
if self._fse_fuse_gate:
901
self._maybe_fuse_gate_weights()
901
- router_logits = F.linear(hidden_states, self._combined_gate_weight)
902
+ router_logits = dispatch_unquantized_gemm()(
903
+ self, hidden_states, self._combined_gate_weight, None
904
+ )
905
else:
906
router_logits, _ = self.gate(hidden_states)
907