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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user