diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 08bf0faf..0088830d 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -2205,10 +2205,92 @@ class DeepseekV4ForCausalLM(nn.Module): print(" Checkpoint loaded. Preparing NVFP4...", flush=True) self.model.finalize_mega_moe_weights() self.model._convert_nvfp4_post_load() + print(" Warming up tilelang kernels...", flush=True) + self._warmup_tilelang() print(" NVFP4 model ready ✓", flush=True) return loaded_params + def _warmup_tilelang(self) -> None: + """Force-compile all tilelang JIT kernels with dummy data. + + tilelang's @jit decorator compiles lazily on first call. In eager mode + (no cudagraphs), the HTTP server comes up before the first inference + triggers compilation — and any request hitting the model during + compilation crashes vLLM. This warmup ensures all kernels are compiled + before the server accepts traffic. + + We call the custom ops directly with 1-token dummy tensors to populate + the tilelang kernel cache. + """ + import torch + config = self.model.config + hc_mult = config.hc_mult + hidden_size = config.hidden_size + device = next(self.model.parameters()).device + hc_mult3 = hc_mult * (2 + hc_mult) + + # Warmup mhc_pre + residual = torch.randn(1, hc_mult, hidden_size, dtype=torch.bfloat16, + device=device) + fn = torch.randn(hc_mult3, hc_mult * hidden_size, dtype=torch.float32, + device=device) + hc_scale = torch.randn(3, dtype=torch.float32, device=device) + hc_base = torch.randn(hc_mult3, dtype=torch.float32, device=device) + + try: + torch.ops.vllm.mhc_pre( + residual=residual, + fn=fn, + hc_scale=hc_scale, + hc_base=hc_base, + rms_eps=config.rms_norm_eps, + hc_pre_eps=config.hc_eps, + hc_sinkhorn_eps=config.hc_eps, + hc_post_mult_value=2.0, + sinkhorn_repeat=config.hc_sinkhorn_iters, + ) + print(" mhc_pre ✓", flush=True) + except Exception as e: + print(f" mhc_pre warmup failed (non-fatal): {e}", flush=True) + + # Warmup mhc_post + x = torch.randn(1, hidden_size, dtype=torch.bfloat16, device=device) + post_mix = torch.randn(1, hc_mult, 1, dtype=torch.float32, device=device) + comb_mix = torch.randn(1, hc_mult, hc_mult, dtype=torch.float32, + device=device) + + try: + torch.ops.vllm.mhc_post(x, residual, post_mix, comb_mix) + print(" mhc_post ✓", flush=True) + except Exception as e: + print(f" mhc_post warmup failed (non-fatal): {e}", flush=True) + + # Warmup hc_head (also @torch.compile lazy) + try: + from vllm.platforms import current_platform + hc_head_fn = torch.randn( + hc_mult * (2 + hc_mult), hc_mult * hidden_size, + dtype=torch.float32, device=device) + hc_head_scale = torch.randn(3, dtype=torch.float32, device=device) + hc_head_base = torch.randn( + hc_mult * (2 + hc_mult), dtype=torch.float32, device=device) + hs = torch.randn(1, hc_mult, hidden_size, dtype=torch.bfloat16, + device=device) + hc_head(hs, hc_head_fn, hc_head_scale, hc_head_base, + config.rms_norm_eps, config.hc_eps) + print(" hc_head ✓", flush=True) + except Exception as e: + print(f" hc_head warmup failed (non-fatal): {e}", flush=True) + + # Free dummy tensors + del residual, fn, hc_scale, hc_base, x, post_mix, comb_mix + try: + del hc_head_fn, hc_head_scale, hc_head_base, hs + except NameError: + pass + torch.cuda.empty_cache() + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping()