[mickqian]

Motivation

--offload-during-compile (#29862, on by default) layerwise-offloads the DiT for the torch.compile warmup, so max-autotune compiles each layer with only that layer resident. Residency is restored on the first non-warmup request. That is the real-path prewarm, which still runs during startup.

The restore does not bring back whole-graph compile. The warmup traces the offload hooks, and their torch.compiler.disable calls graph-break inside the block loop. Dynamo then skips the DiT frame, and it applies that skip to the code object for every later call: ConvertFrameReturn defaults to FrameExecStrategy(SKIP, DEFAULT) with apply_to_code=True. Removing the hooks changes nothing, so every request keeps running the per-block graphs the warmup left behind.

A toy reproduction is in the new test. It is a 3-block loop with a disabled pre-hook, compiled with CompileCounter. After the hooks are removed, serving compiles 0 frames. With this PR it compiles one frame with all 6 ops.

Modifications

DenoisingStage._offload_for_torch_compile_warmup calls torch._dynamo.reset() right after disable_offload(). The next request compiles whole graphs from the autotune results the warmup already cached, so the compile peak that #29862 caps stays capped. That request is the real-path prewarm, at startup.

Accuracy Tests

No numerical change beyond graph partitioning: the same kernels, fused across block boundaries instead of within each block.

Speed Tests and Profiling

4xGB300, TP2, served, 1024x1024, --enable-torch-compile true with default flags (offload during compile on). Base: main 3831e7e, the then-open GB300 fixes, #43601 and #43602. Each cell is the median of 5 requests after warmup. Two interleaved rounds (ABAB), each arm on a fresh server.

case	before	this PR	change
Qwen-Image-2.1, round 1 / 2	2.150 / 2.167 s	1.672 / 1.556 s	−22% / −28%
FLUX.1-dev, round 1 / 2	2.631 / 2.688 s	2.531 / 2.561 s	−4% / −5%
Z-Image-Turbo, round 1 / 2	0.945 / 0.912 s	0.948 / 0.904 s	—

Qwen-Image-2.1 comes close to the run with --offload-during-compile false (1.50 s). Z-Image does not move: with the flag off it serves in 0.57 s, so something else in its offloaded warmup still costs it.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests: test_offload_during_compile.py. It runs the real context manager around a compiled toy DiT. Warmup keeps the DiT offloaded. Serving compiles one whole graph (6 ops); without this PR it compiles none (AssertionError: 0 != 1). On GB300: 2 passed, against 1 failed / 1 passed on the base tree; with test_layerwise_offload.py on this branch rebased onto main 5730345: 199 passed.
• Update documentation: N/A.
• Provide accuracy and speed benchmark results.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ✅ Run #38101028112
Latest PR Test (Extra): ❌ Run #38101028042
Latest PR Test (AMD ROCm 10): ⏳ Run #38101028175

[mickqian]

/tag-and-rerun-ci