biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: handle input_scale as 1D or 2D (EP splits change the shape)

Browse Source
This commit is contained in:
biondizzle
2026-05-16 22:49:30 +00:00
parent 139c9c37cd
commit b382a7a528
1 changed files with 13 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
vllm/patches/deepseek_v4.py
View File

@@ -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()