diff --git a/vllm/patches/fused_moe/experts/cutedsl_moe.py b/vllm/patches/fused_moe/experts/cutedsl_moe.py index 77ee0f9b..5c900e18 100644 --- a/vllm/patches/fused_moe/experts/cutedsl_moe.py +++ b/vllm/patches/fused_moe/experts/cutedsl_moe.py @@ -271,13 +271,9 @@ class CuTeDSLMoEExperts(mk.FusedMoEExpertsModular): # Our runner manages its own workspace internally (pre-allocated buffers) workspace1 = (0,) workspace2 = (0,) - # K is the packed dimension from w1.shape[-1]. - # For NVFP4 uint8 packed weights, K_packed = K_logical // 2. - # The output of the L2 GEMM is hidden_dim (unpacked). - # If K == hidden_dim, weights are BF16 (not packed). - # If K == hidden_dim // 2, weights are NVFP4 packed. - output_dim = max(K * 2, self.hidden_dim) - output = (M, output_dim) + # The output of the full 2-stage MoE pipeline is hidden_dim. + # K comes from hidden_states.size(-1) (full BF16 dimension, not packed). + output = (M, self.hidden_dim) return (workspace1, workspace2, output) def apply( @@ -313,9 +309,4 @@ class CuTeDSLMoEExperts(mk.FusedMoEExpertsModular): ) # Copy result into output tensor - if result.shape != output.shape: - import sys - print(f"[CuTeDSL MoE] SHAPE MISMATCH: result={result.shape} output={output.shape} " - f"hidden_dim={self.hidden_dim} w1={w1.shape if w1 is not None else None} " - f"hs={hidden_states.shape}", file=sys.stderr, flush=True) output.copy_(result)