From bdea2777eaeab65937dffdbdcb372b03e2cad441 Mon Sep 17 00:00:00 2001 From: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:33:56 +0300 Subject: [PATCH] [Bugfix][MoE] Fix batched CUTLASS workspace overallocation (#55579) Signed-off-by: Itay Alroy Co-authored-by: Tyler Michael Smith --- vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py index 80c6926fd4..402077b017 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py @@ -477,14 +477,12 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base): expert_tokens_meta: mk.ExpertTokensMetadata | None, activation: MoEActivation, ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - num_dp = self.num_dispatchers - assert num_dp is not None experts_per_worker = self.moe_config.num_local_experts activation_out_dim = self.adjust_N_for_activation(N, activation) - workspace1 = (experts_per_worker, M * num_dp, max(N, K)) + workspace1 = (experts_per_worker, M, max(N, K)) workspace2 = ( experts_per_worker, - M * num_dp, + M, max(activation_out_dim, K), ) output = (experts_per_worker, M, K)