iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Bugfix][MoE] Fix batched CUTLASS workspace overallocation (#55579)

Browse Source 
Signed-off-by: Itay Alroy <ialroy@nvidia.com>
Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
This commit is contained in:
Itay Alroy
2026-09-10 15:33:56 +00:00
committed by GitHub
co-authored by Tyler Michael Smith
parent db723d245a
commit bdea2777ea
1 changed files with 2 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+2 -4
View File
@@ -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)