[Bugfix][MoE] Fix batched CUTLASS workspace overallocation (#55579)
Signed-off-by: Itay Alroy <ialroy@nvidia.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Itay Alroy committed
Sep 10, 2026 at 18:33 UTC
bdea2777eaeab65937dffdbdcb372b03e2cad441
1 file changed
+2
-4
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+2
-4
@@ -477,14 +477,12 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):
477
expert_tokens_meta: mk.ExpertTokensMetadata | None,
478
activation: MoEActivation,
479
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
480
- num_dp = self.num_dispatchers
481
- assert num_dp is not None
480
experts_per_worker = self.moe_config.num_local_experts
481
activation_out_dim = self.adjust_N_for_activation(N, activation)
484
- workspace1 = (experts_per_worker, M * num_dp, max(N, K))
482
+ workspace1 = (experts_per_worker, M, max(N, K))
483
workspace2 = (
484
experts_per_worker,
487
- M * num_dp,
485
+ M,
486
max(activation_out_dim, K),
487
)
488
output = (experts_per_worker, M, K)