From f191af7e298a5fbde729a9bdff3c331524594268 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 04:39:05 +0000 Subject: [PATCH] feat: warm up CuTeDSL kernel during model loading MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit JIT compiles the MLIR→PTX during finalize_weights instead of on the first inference request. Prevents vLLM's 5-min RPC timeout from killing the engine while workers are busy compiling. Warmup runs a single-token, single-expert forward pass — just enough to trigger compilation. Takes ~1-2 min, same as layertest. --- vllm/patches/deepseek_v4.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index c5e05f9d..86de2081 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -429,7 +429,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): l2_fp4, l2_sf, l2_gs = [], [], [] from tqdm import tqdm - for e in tqdm(range(self.num_local_experts), desc=" uint8→NVFP4 experts", unit="exp"): + for e in tqdm(range(self.num_local_experts), desc=" (byte view-cast)uint8→NVFP4 experts", unit="exp"): # ── L1: gate + up (fused) ── gate_w = self.w13_weight.data[e, :self.intermediate_size] # (intermediate, hidden//2) uint8 up_w = self.w13_weight.data[e, self.intermediate_size:] # (intermediate, hidden//2) uint8 @@ -505,6 +505,19 @@ class DeepseekV4MegaMoEExperts(nn.Module): self.w2_weight_scale_2 = None self.w2_input_scale = None + # Warm up the CuTeDSL kernel (JIT compiles MLIR→PTX on first call). + # This takes ~1-2 min but prevents the vLLM RPC timeout (5 min) from + # killing the engine when the first inference request triggers compilation. + print(" Warming up CuTeDSL NVFP4 kernel (first call compiles)...", flush=True) + dummy_hidden = torch.randn(1, self.hidden_size, dtype=torch.bfloat16, device=self._cutedsl_runner.l1_fp4[0].device) + dummy_ids = torch.zeros(1, 1, dtype=torch.int32, device=dummy_hidden.device) + dummy_weights = torch.ones(1, 1, dtype=torch.float32, device=dummy_hidden.device) + try: + self._cutedsl_runner.run(dummy_hidden, dummy_weights, dummy_ids, expert_indices=[0]) + print(" CuTeDSL warmup complete ✓", flush=True) + except Exception: + print(" CuTeDSL warmup failed (will compile on first inference)", flush=True) + def forward( self, hidden_states: torch.Tensor, @@ -1623,7 +1636,7 @@ class DeepseekV4Model(nn.Module): _shard_index = self._build_shard_index("/model") if os.path.isdir("/model") else None from tqdm import tqdm - for layer_idx, layer in tqdm(enumerate(self.layers), total=len(self.layers), desc=" NVFP4→FP8/BF16 convert", unit="layer"): + for layer_idx, layer in tqdm(enumerate(self.layers), total=len(self.layers), desc=" (upcast)NVFP4→FP8/BF16 convert", unit="layer"): attn = layer.attn # FP8 conversion: only wo_a