diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py index 484b8aa9d107c..86cdd25f2c873 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py @@ -215,7 +215,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute): ) assert M_sum % block_m == 0 - workspace1 = (M_sum, max(N, K)) + workspace1 = (M_sum, N) workspace2 = (M_sum, max(N // 2, K)) output = (M, K) return (workspace1, workspace2, output)