[diffusion] Measure warmup memory and layer usage per phase for residency calibration (1/4)

[mickqian]

First of four stacked PRs split out of #35335 (which becomes the fourth). This layer only measures: nothing about component placement changes yet.

What it records

During synthetic server warmup, one WarmupMemoryRecord per warmup request with everything a residency planner later needs:

• Per-phase device peaks. ComponentResidencyManager keys a warmup phase per stage setup / component use / prefetch / full-weight transition and records the allocator peak and the active components of each (WarmupPhasePeak).
• Iteration counts. A stage declares how many times its repeated unit runs as a formula of the step count (PipelineStage.default_workload_iterations); the base class evaluates it at the probe's steps (measured) and at the default workload's (target) after every warmup forward. DenoisingStage returns the step count, MiniMax-H3 steps − 1, DMD and the LTX-2 refiner their fixed distilled schedules. Only loops whose length is known inside them (block-wise causal DMD, Helios chunks, SANA-WM refiners) still call batch.record_stage_iterations explicitly. The perf metrics keep per-stage step timings.
• Per-layer uses. LayerwiseUsageTracker counts how often each layer of a layerwise component runs during the probe.

Types and estimators

auto_residency.py carries the record, target, workload and report types and the pure estimators that extrapolate a probe to the default workload (estimate_workload_phase_peaks, estimate_layerwise_layer_uses, estimate_default_workload_timing). No planning here.

Warmup plumbing

The warmup request builder gains the full-shape calibration probe (AUTO_RESIDENCY_TIMING_STEPS step floor), --warmup-sampling-params, and the auto_residency_args_skip_reason gate in auto_tune that says whether a deployment can be calibrated at all. The bounded warmup still runs first; the probe follows it, and before it runs the worker extrapolates the bounded measurement to the probe's shape and shrinks the probe (frames, then area) while that exceeds free memory minus the reserve, with all ranks agreeing on the most cautious shrink. A probe that still runs out of memory is retried smaller up to three times (the retry also recognizes cuBLAS/CUDA allocation failures); a warmup that fails altogether is fail-open unless --warmup-resolutions was given explicitly. Control-request structs move from entrypoints/utils.py to entrypoints/control_requests.py (a pure move the third layer builds on).

Peak-VRAM metric

The CI's runtime_peak_vram_mb was the worker's running maximum of peak reserved memory over every request, warmup included. This layer's calibration probe runs the default workload's full shape (1280×720×81f for Wan 2.2 T2V), which the CI requests never reach, so their reported runtime peak rose 7–25% without the requests using a byte more. Warmup requests now feed a separate warmup_peak (carried in the replica report, shown in the test summary as warmup_peak_vram_mb, not validated); the runtime figure is the serving requests' own peak again, and the fast_hunyuan_video baseline returns to main's value.

Stack

1. this PR — measurement and record types
2. [diffusion] Plan component residency from calibrated warmup records (2/4) #37917 — planner as a pure library (collect_residency_targets, plan_auto_residency, placement_budget)
3. [diffusion] Apply, validate and roll back calibrated residency after warmup (3/4) #37918 — apply / validate / rollback wiring (worker, scheduler, client, layerwise layouts, ServerArgs modes)
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. Tests: 748 planner/warmup unit tests pass on the top of the stack; this layer's own tests are test_performance_metrics, test_cfg_parallel_warmup, test_warmup_probe_degradation, test_component_residency and the estimator tests in test_auto_residency.

---

CI States

Latest PR Test (Base): ✅ Run #34026695033
Latest PR Test (Extra): ❌ Run #34026694906
Latest PR Test (AMD ROCm 7.2): ⏳ Run #34026695048

Follow-ups from CI (2026-09-05)

• Probe fitting reserve: the full-shape probe was shrunk against the planner's placement reserve (4 GiB / 10%), which cut a Wan 81f probe to 41f with 10 GiB to spare. The fit now keeps only an allocator margin (max(1 GiB, 2% of device memory)) and logs the requested shape's estimate.

• Warmup pool release: the probe runs shapes serving never sees, and its cached allocator blocks became the reserved floor of the first serving request (fast_hunyuan_video runtime peak 37.7–39.0 GiB vs the 35.9 GiB baseline). The worker now drops the pool once at the warmup → serving transition; runtime peaks measure the serving request again and warmup peaks stay in warmup_peak_vram_mb.

• Frame contract for lightened probes: halving frames bypassed pipeline_config.adjust_num_frames, so causal-block models (LongLive2: latent frames in blocks of 8) rejected the lightened probe. The halved count is now walked down to the largest value the contract accepts (61f → 29f), then the ladder moves to area.

• Non-PR failures seen on this head: AMD and NPU multimodal-gen jobs fail identically on main (NPU: minimax_h3_t2va_2npu / mova_360p_ti2va_2npu e2e and text-encoding stage); the 5090 minimax_h3_t2va_consumer_budget e2e is runner variance against the baseline [diffusion] fix: preserve mapped courier tensor lifetime #37965 tightened to 42.8 s (42.7 / 48.0 / 56.8 / 64.1 s across identical code).

• Re-warm after the probe: releasing the probe's pool alone still measured fast_hunyuan_video at 37.1 GiB reserved (main: 35.9 GiB, deterministic on both sides) because the serving request grows an empty pool differently from one pre-shaped by a small warmup. The bounded warmup now runs once more after the probe (bounded → probe → bounded), and the pool release moved to the request that follows the probe, so serving starts from the same allocator state as main. One extra bounded warmup request per model at startup; nothing changes when no probe runs.

• Baseline: longlive2_t2v h100 runtime peak 61190 → 62706 MiB. With the frame-contract fix the full-shape probe now runs for LongLive2, so the serving request starts from a different allocator pool history than on main; the reserved peak measured 62706 on seven consecutive attempts (main nightly: 61074) while placement and allocated usage are unchanged. Reserved peak carries a 1.5–3% dependence on pool history (same code measured 35926 / 36452 / 37134 / 37710 / 39026 for fast_hunyuan_video under different histories); a follow-up PR moves the guard to max_memory_allocated.

• IPC all-to-all staging after the probe (wan2_2_t2v_a14b_2gpu runtime peak +20–27%): the 2-rank IPC transport caches staging buffers keyed by message size and evicts by count only, so the 1280x720x81f probe left ~5 GiB of buffers per rank behind; the case then served at 20.3 GiB allocated against main's 15.1 GiB with identical placement (every other 2-GPU case matched main to 0.01 GiB). The staging cache is now released together with the allocator pool before the request that follows the probe. Probe lightening also no longer walks below the bounded warmup shape (it reached 16x16x1f on wan2_1_i2v_14b_480P_2gpu, which the model rejects).

• All-to-all staging caches are released after the probe (wan2_2_t2v_a14b_2gpu, Ulysses degree 2). Both ipc_a2a.get_staging and usp._A2A_STAGING_BUFFERS grow to the largest message seen and never shrink, so the 1280x720x81f probe left 1.6 + 4.3 GiB allocated per rank that the serving request never needed (allocated peak 20.3 GiB vs 15.1 on main, reserved 24–27 GiB vs the 21.5 GiB baseline). Attributed with torch.cuda.memory._record_memory_history on a 2×H200 devbox: before the first serving request 9.68 GiB allocated with the two Ulysses buffers at 2.16 GiB each; 5.80 GiB once both caches are dropped in the worker's post-probe release, serving peak allocated 19.15 → 15.45 GiB (main 15.11). The temporary instrumentation commit was reverted in the same push.