[diffusion] Apply, validate and roll back calibrated residency after warmup (3/4)

[mickqian]

Third of four stacked PRs split out of #35335. Base: the planner library. This layer is the runtime wiring that turns a plan into a placement.

Flow

server_warmup.maybe_apply_auto_residency runs between the synthetic warmup and the ready signal: apply → re-warm → validate → (rollback). GPUWorker.apply_auto_residency builds the rank report from the warmup records, plans jointly across ranks, applies the changes transactionally (apply_residency_changes / commit_residency_changes / rollback_residency_changes), and the validation re-warm rolls the round back on a memory shortfall or a materially slower probe. rollback_auto_residency restores the last calibrated placement. The scheduler and both scheduler clients fan the AutoResidencyReq control request out to every replica and merge the outcomes (a failed rollback wins over an adjusted replica).

Mechanisms it needs

• Layerwise offload: per-group resident layouts and pin frontiers, host-store repacking, disable/enable offload for promotions, and release of host stores once a resident placement is validated.
• ServerArgs: auto residency modes next to the explicit ones (configured_residency_mode vs runtime auto_residency_mode, the shared host_pin_budget); auto_tune excludes fixed loading paths (quantized, direct GPU, LTX-2 original) per component; quantized component offload capability is enforced at load (supports_quantized_component_offload).
• The post-request residency hint reuses the planner's candidate ranking.
• SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY is the kill switch (defined in the measurement layer, honored by the gate).

Tests

test_auto_residency_warmup (apply / validate / rollback rounds, replica merge), test_layerwise_offload, test_health_warmup_gate, test_server_args, test_dp_routing, test_transformer_quant, and the apply-side tests in test_auto_residency. 748 tests pass on this layer.

Stack

1. [diffusion] Measure warmup memory and layer usage per phase for residency calibration (1/4) #37916 — measurement and record types
2. [diffusion] Plan component residency from calibrated warmup records (2/4) #37917 — planner as a pure library
3. [diffusion] Apply, validate and roll back calibrated residency after warmup (3/4) #37918 — apply / validate / rollback wiring
4. [diffusion] Auto residency: per-pipeline containment fixes and CI baselines (4/4) #35335 — per-pipeline containment fixes and CI baselines

Every layer's tree is a prefix of #35335's; the union is byte-identical to the original branch merged with main.

---

CI States

Latest PR Test (Base): ❌ Run #34829636346
Latest PR Test (Extra): ❌ Run #34829635803
Latest PR Test (AMD ROCm 10): ❌ Run #34829636099

Regression fixes (2026-09-07)

Comparing this head against main's run 34015151838 on the same H100 runners showed three default-command cases serving slower: fast_hunyuan_video 9.2 s → 10.2–12.9 s, lingbot_video_moe_t2v 6.6 s → 8.0–8.4 s, hunyuanvideo_modelopt_fp8_t2v 2.6 s → 6.9–10.1 s (the last two had no e2e baseline; #38335 adds them). Commit 1e29438 fixes the causes, each reproduced on an H200 devbox with 80 GiB of free VRAM emulated:

case	main	this PR before	this PR after
fast_hunyuan_video	9.77 s	10.16 s	9.08 s
hunyuanvideo_modelopt_fp8_t2v	2.43 s	6.84 s	2.24 s
lingbot_video_moe_t2v	6.38 s	7.60 s	4.21 s

• ServerArgsAutoTuner no longer skips main's memory-threshold placement under server warmup. Calibration starts from that placement and refines it; it never serves below it when a promotion is declined.
• The probe fit subtracts the weight floor the planner already uses (resident components plus the largest streamed one), so a streamed 24 GiB DiT is no longer extrapolated as activations (77 GiB estimate for a 33 GiB request).
• finalize_resident_layerwise_placement only empties the host caching allocator when the freed pinned stores exceed 1 GiB. Dropping it after a 0.1 GiB VAE promotion discarded the pinned blocks component-offload text encoders recycle for their per-request D2H copy, so the next request re-pinned 8–14 GiB synchronously (3.5–4.8 s).
• TimestepPreparationStage logs the timesteps' shape instead of the tensor, which forced a device sync behind the DiT prefetch (158 ms vs 0.6 ms).

Follow-ups (2026-09-08). 5cb56be keys the MiniMax-H3 VAE decode-dtype store by module layout: the runtimes on shared CI runners expose different parameter sets (144 vs 288 tensors) and were deleting each other's store, which left the decoder weights in anonymous host memory and made the 2-GPU ref2va decode 12.9 s instead of 2.3 s. 5cb99db refreshes the allocated/reserved peak guards of the 29 H100 and 4 RTX 5090 cases whose text encoders or VAEs the planner now keeps resident, plus the warmup peak of wan2_2_t2v_a14b_2gpu and the sampled steps of two cases; no e2e expectation changed.

A full per-case comparison of this head against main's run 34015151838 (46 single-GPU, 28 two-GPU, 6 RTX 5090 and 5 B200 cases) shows every case equal or faster within main's own run-to-run noise (up to 1.08 on single samples). The B200 nvfp4 cases look 1.6-1.9x slower on some runs, but that tracks the runner node, not the branch: b200-cirrascale1-* and b200-di01-* run qwen_image_2512_modelopt_nvfp4_t2i at 331-359 ms/step for every branch including unrelated PRs, b200-fin03-* at 187-213 ms/step for main and for #37811; a B200 devbox gives 177 ms/step for both main and this head.

[mickqian]

/tag-and-rerun-ci