fix: remove CuTeDSL warmup — OOM with 175GB model loaded

The warmup allocated 1GB of dummy tensors but the model already
uses 175.7GB of the 178.35GB per GPU. No room.

With FULL_AND_PIEWISE CUDA graph mode, the kernel compiles during
the graph capture phase (which manages memory properly). The warmup
was a band-aid for eager mode and is now redundant.
This commit is contained in:
2026-05-16 07:32:17 +00:00
parent e0814eb54e
commit a2cac7a7fe

View File

@@ -517,20 +517,9 @@ 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).
# CuTeDSL caches by (M, N, K) shape — different shapes trigger new
# compiles. Running warmup on every layer ensures all kernels are
# compiled before inference, preventing vLLM RPC timeouts.
# The MMA tiler needs >= 128 tokens; using 128 for the warmup.
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)
dummy_weights = torch.ones(128, 1, dtype=torch.float32, device=device)
try:
self._cutedsl_runner.run(dummy_hidden, dummy_weights, dummy_ids, expert_indices=[0])
except Exception as exc:
print(f" CuTeDSL warmup FAILED: {exc}", flush=True)
print(" The CUDA context may be corrupted. Check kernel alignment/tiling.", flush=True)
# Note: No explicit CuTeDSL warmup here. With FULL_AND_PIECEWISE
# CUDA graph mode, the kernel compiles during graph capture (startup).
# In eager mode, the first inference triggers JIT compilation.
def forward(
self,