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

fix: compile CuTeDSL kernel once per process, not per MoE layer

Browse Source 
The warmup was running for every MoE layer (61 layers × 8 ranks = 488
compile attempts). The kernel is cached after the first compile —
subsequent calls are instant. But the print spam was insane.

Now uses a class-level flag to compile exactly once per process.
All 61 layers on a rank share the same compiled kernel.
This commit is contained in:
biondizzle
2026-05-16 05:16:53 +00:00
parent 936982c5aa
commit f19932d8db
1 changed files with 18 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
vllm/patches/deepseek_v4.py
View File

@@ -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,