diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 634372f6..939a4737 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -221,7 +221,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): It handles NVFP4 natively with full Blackwell pipeline overlap (TMA → MMA → Epilogue). This replaces the broken C++ CUTLASS kernel (see README.md for the full story). """ - _cutedsl_runner: 'CuTeDSLMoERunner | None' = None + _cutedsl_compiled: bool = False # NVFP4 E2M1 lookup table (positive values, sign from bit 3) E2M1_LUT = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] @@ -330,6 +330,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): set_weight_attrs(self.w2_input_scale, weight_attrs) self._cutedsl_runner = None + self._cutedsl_compiled = False # Register in the static forward context so the custom-op wrapper # can look up this module by name from within a torch.compile graph. @@ -505,23 +506,26 @@ class DeepseekV4MegaMoEExperts(nn.Module): self.w2_input_scale = None # Warm up the CuTeDSL kernel (JIT compiles MLIR→PTX on first call). + # Only compile once per process — the kernel is cached after that. # This takes ~1-2 min but prevents the vLLM RPC timeout (5 min) from # killing the engine when the first inference request triggers compilation. # Note: The MMA tiler is (128,128,256) — we need >= 128 tokens to fill # a tile. Using 128 tokens, 1 expert for the warmup. - print(" Compiling CuTeDSL NVFP4 MegaMoE kernel (one-time JIT, ~1-2 min)...", flush=True) - try: - device = self._cutedsl_runner.l1_fp4[0].device - dummy_hidden = torch.randn(128, self.hidden_size, dtype=torch.bfloat16, device=device) - dummy_ids = torch.zeros(128, 1, dtype=torch.int32, device=device) # all to expert 0 - dummy_weights = torch.ones(128, 1, dtype=torch.float32, device=device) - self._cutedsl_runner.run(dummy_hidden, dummy_weights, dummy_ids, expert_indices=[0]) - print(" CuTeDSL NVFP4 MegaMoE kernel compiled ✓", flush=True) - except Exception as exc: - # CUDA illegal memory access corrupts the context — can't recover. - # Log the error clearly so the user knows to check the kernel. - print(f" CuTeDSL warmup FAILED: {exc}", flush=True) - print(" The CUDA context may be corrupted. Check kernel alignment/tiling.", flush=True) + if not DeepseekV4MegaMoEExperts._cutedsl_compiled: + DeepseekV4MegaMoEExperts._cutedsl_compiled = True + print(" Compiling CuTeDSL NVFP4 MegaMoE kernel (one-time JIT, ~1-2 min)...", flush=True) + try: + device = self._cutedsl_runner.l1_fp4[0].device + dummy_hidden = torch.randn(128, self.hidden_size, dtype=torch.bfloat16, device=device) + dummy_ids = torch.zeros(128, 1, dtype=torch.int32, device=device) # all to expert 0 + dummy_weights = torch.ones(128, 1, dtype=torch.float32, device=device) + self._cutedsl_runner.run(dummy_hidden, dummy_weights, dummy_ids, expert_indices=[0]) + print(" CuTeDSL NVFP4 MegaMoE kernel compiled ✓", flush=True) + except Exception as exc: + # CUDA illegal memory access corrupts the context — can't recover. + # Log the error clearly so the user knows to check the kernel. + print(f" CuTeDSL warmup FAILED: {exc}", flush=True) + print(" The CUDA context may be corrupted. Check kernel alignment/tiling.", flush=True) def forward( self,