diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 5c674237..00f01236 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -524,6 +524,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): activation_clamp: float | None, fast_math: bool, ) -> None: + import os import nvfp4_megamoe_kernel as deep_gemm symm_buffer = self.get_symm_buffer() @@ -557,7 +558,6 @@ class DeepseekV4MegaMoEExperts(nn.Module): symm_buffer.topk_weights[:num_tokens].copy_(topk_weights) # Debug: check staging output - import os if int(os.environ.get('MEGA_MOE_DEBUG', '0')): print(f"[MEGA_MOE_DEBUG] After staging: x dtype={symm_buffer.x.dtype} shape={symm_buffer.x.shape}") print(f"[MEGA_MOE_DEBUG] x_sf dtype={symm_buffer.x_sf.dtype} shape={symm_buffer.x_sf.shape}") @@ -583,7 +583,6 @@ class DeepseekV4MegaMoEExperts(nn.Module): from nvfp4_megamoe_kernel import nvfp4_mega_moe_full as fp8_nvfp4_mega_moe # Debug: dump shapes before mega_moe - import os if int(os.environ.get('MEGA_MOE_DEBUG', '0')): l1_w, l1_sf = self._transformed_l1_weights l2_w, l2_sf = self._transformed_l2_weights