From b382a7a52822b7e2949c01c158b1c56fc24a44c5 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 22:49:30 +0000 Subject: [PATCH] fix: handle input_scale as 1D or 2D (EP splits change the shape) --- vllm/patches/deepseek_v4.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 4c5096cc..abf2a2ad 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -513,12 +513,19 @@ class DeepseekV4MegaMoEExperts(nn.Module): self._cutedsl_runner.l2_gs = l2_gs # Set activation global scales from checkpoint input_scale - # The input_scale is the pre-computed activation normalization factor - # for each projection. Using hardcoded 1/2688 gives wrong scale for most layers. - # Read from the nn.Parameters BEFORE they're freed below. - # input_scale shape: (num_experts, 2) for w13 (gate, up), (num_experts, 1) for w2 - l1_igs = self.w13_input_scale.data[:, 0] # gate input_scale - l2_igs = self.w2_input_scale.data[:, 0] # down input_scale + # The input_scale is the pre-computed activation normalization factor. + # w13_input_scale shape: (num_experts, 2) for gate+up, but may be (num_experts,) after EP split + # w2_input_scale shape: (num_experts, 1) or (num_experts,) + w13_igs = self.w13_input_scale.data + w2_igs = self.w2_input_scale.data + if w13_igs.dim() == 2: + l1_igs = w13_igs[:, 0] # gate input_scale + else: + l1_igs = w13_igs # already 1D per expert + if w2_igs.dim() == 2: + l2_igs = w2_igs[:, 0] + else: + l2_igs = w2_igs self._cutedsl_runner.l1_activation_global_scale = l1_igs.mean().item() self._cutedsl_runner.l2_activation_global_scale = l2_igs.mean().item()