biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: clearer warmup message — 'Compiling CuTeDSL NVFP4 MegaMoE kernel'

Browse Source
This commit is contained in:
biondizzle
2026-05-16 04:40:31 +00:00
parent f191af7e29
commit a70d2d3984
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/patches/deepseek_v4.py
View File

@@ -508,13 +508,13 @@ class DeepseekV4MegaMoEExperts(nn.Module):
# 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)
print(" Compiling CuTeDSL NVFP4 MegaMoE kernel (one-time JIT, ~1-2 min)...", 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)
print(" CuTeDSL NVFP4 MegaMoE kernel compiled ✓", flush=True)
except Exception:
print(" CuTeDSL warmup failed (will compile on first inference)", flush=True)