debug more2

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