diff --git a/src/nvfp4_megamoe_kernel/weight_transform.py b/src/nvfp4_megamoe_kernel/weight_transform.py index 4e2d4696..0e01033e 100644 --- a/src/nvfp4_megamoe_kernel/weight_transform.py +++ b/src/nvfp4_megamoe_kernel/weight_transform.py @@ -79,11 +79,6 @@ def transform_nvfp4_weights_for_mega_moe( l1_weight, l1_weight_scale = l1_tuple l2_weight, l2_weight_scale = l2_tuple - # DEBUG: check weights BEFORE transform - print(f"[WT-IN] l1_w shape={l1_weight.shape} absmax={l1_weight.view(torch.int8).abs().max().item()} " - f"l1_sf shape={l1_weight_scale.shape} sf_absmax={l1_weight_scale.view(torch.uint8).abs().max().item()} " - f"l2_w shape={l2_weight.shape} absmax={l2_weight.view(torch.int8).abs().max().item()}") - # Fold global scales into block scales # The logical_widths branch was wrong: it treated gs as per-projection # scalars and only used experts 0 and 1's scales for ALL experts. @@ -112,8 +107,4 @@ def transform_nvfp4_weights_for_mega_moe( l1_sf_out = l1_sf_out.transpose(-2, -1).contiguous() l2_sf_out = l2_sf_out.transpose(-2, -1).contiguous() - # DEBUG: check weights AFTER transform - print(f"[WT-OUT] l1_w shape={l1_weight_out.shape} absmax={l1_weight_out.view(torch.int8).abs().max().item()} " - f"l1_sf shape={l1_sf_out.shape} sf_absmax={l1_sf_out.view(torch.uint8).abs().max().item()}") - return (l1_weight_out, l1_sf_out), (l2_weight_out, l2_sf_out) diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 1b66350c..2c54fcc8 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -371,6 +371,13 @@ class DeepseekV4MegaMoEExperts(nn.Module): if local_expert_id == -1: return False + # DEBUG: log weight loads for expert params + if "w13_weight" in weight_name and local_expert_id < 2: + print(f"[WT-LOAD] {weight_name} expert={expert_id}→local={local_expert_id} " + f"shard={shard_id} loaded_shape={tuple(loaded_weight.shape)} " + f"param_shape={tuple(param.data[local_expert_id].shape)} " + f"loaded_absmax={loaded_weight.view(torch.int8).abs().max().item()}") + # Scalar params (weight_scale_2, input_scale): 1D per-expert if "weight_scale_2" in weight_name or "input_scale" in weight_name: param.data[local_expert_id].copy_(loaded_weight)