diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py index 7f76b2616e..9e8b7df1c1 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py @@ -287,7 +287,10 @@ def fused_marlin_moe( dtype=hidden_states.dtype, ) # Marlin skips masked expert rows, so their shared cache must start at zero. - intermediate_cache13 = torch.zeros( + # Without an expert map every (token, expert) row is local and written by + # both GEMMs, so the ~M*topk*K zero-fill (0.8 GB per MoE layer for an + # 8K-token MiMo prefill chunk) is skipped. + intermediate_cache13 = (torch.zeros if expert_map is not None else torch.empty)( (M * topk_ids.shape[1] * max(gemm1_n, K),), device=hidden_states.device, dtype=hidden_states.dtype,