[Diffusion] Add a permanent lifetime for layerwise resident layers

[mickqian]

Follows #40590 / #40592 / #40593 (all merged), which name and document the
lifetime the resident set has today; this one adds the other lifetime, which is
what the flag was designed to mean. main is merged in (a5e7f87), so the diff
is exactly this change.

What the flag does today, and what it was meant to do

--layerwise-resident-layers keeps a set of a component's layers on the GPU
instead of streaming them. Today that set lives for the component's run: it is
pinned by the layer-0 pre-hook when the component starts its forwards for a
request, released by finish_use → release_all() when it finishes, and a
pinned host copy of every resident layer is kept for the next request.
finish_request forces preferred = False for any module that holds
residents, so nothing keeps the set past the request.

That is the right shape when several components' resident sets have to share
one budget: only the running component holds its set, so the peak is the
largest single set rather than their sum. It is the wrong shape when the sum
fits -- the layers are transferred again on every request, and host memory is
pinned for a reload that never had to happen. The original intent of the flag
was the second case: place the layers once, at load, keep no host copy, never
release them.

Both shapes are real, on adjacent consumer cards, for the same model and the
same flag value. Qwen-Image-2.1 with DiT+VAE resident (14.5 GiB) and the
text encoder's resident set at 0.8 (~9-11 GiB): the sum fits on an RTX 5090
(31.4 GiB) and keeping the set was measured faster (14.136 → 13.737 s); it does
not fit on an RTX 4090 (23.5 GiB, ~6.9 GiB free once DiT+VAE are resident), and
keeping it left the server unable to start. So the lifetime cannot be a
constant, and it is not the same knob as which layers are resident.

The change

A lifetime axis beside the existing policy axis, in the same two-flag shape:

--dit-layerwise-residency-lifetime {forward,permanent}       # group default
--layerwise-residency-lifetime transformer=permanent,video_vae=forward

lifetime	placed	host copy	released	transfers
forward (default, today's behaviour)	when the component starts running its layers for a request	kept, pinned	when it finishes	once per request
permanent	at load, before any forward or warmup	none	never	once, for the life of the server

forward stays the default, so no existing command changes behaviour. Whether
permanent becomes the default when the sum fits is a placement decision; it
belongs with the residency planner (#37917 already carries a
permanent_residency field on ResidencyTarget and splits its option keys
into permanent / stage), not in a constant here.

Manager

• Permanent layers are split out of layer_groups before the hosting plan
  runs, so they take no host store and spend none of the pin budget. They are
  placed with a layout-keeping device copy -- a transposed FP8 view has to stay
  transposed, the same reason the strided host store exists -- shared by the
  CUDA and MPS load paths.
• Residency is armed from construction, so _head_of_stream starts at the
  first streamed layer rather than layer 0.
• release_layer returns early for a permanent layer regardless of force.
  That single check covers release_all, the disable_offload → enable_offload
  round trip the compile-warmup and LoRA paths take, skip-gap and wraparound
  releases alike. sync_layer_to_cpu has nothing to sync;
  iter_materialized_weights reads them from the module as it does any
  non-offloaded parameter. prepare_for_next_req no longer visits them.
• A permanent lifetime with zero resident layers warns and streams everything.
• The readiness log names the lifetime:
  resident=53/66 (permanent), policy=leading, and placement logs what it
  put on the device: placed 29 permanent resident layers (10.42 GiB) on the device at load; no host copy is kept.

Server args and help

One group default plus one per-component map, validated in the CLI (choices)
and in the Python API (ValueError), with the same "has no effect" warnings the
policy knob gives. The two resident-layer help strings now describe the default
lifetime instead of the permanent one they used to promise, and point at the
new knob; the CLI reference and the deployment cookbook gain the lifetime and
the rule for choosing it (sum fits → permanent; otherwise forward).

Wording is the user's, not the runtime's: "while the component runs for a
request", "released when it finishes", "for the life of the server". No
ComponentUse vocabulary on any user-facing surface.

Validation

Unit tests, on an H200 devbox at this head (2bb100d):
test_layerwise_offload.py 170 passed (158 on main + 12 new);
test_server_args.py -k "lifetime or residency" 21 passed + 13 subtests.
The new tests pin: placement at load without a host store and with layout kept;
survival through release_layer(force=True), release_all,
sync_all_layers_to_cpu, and a disable_offload → enable_offload round trip;
prepare_for_next_req not touching the set (same data_ptr, no prefetch
issued); the hosting plan receiving only the streamed layers;
iter_materialized_weights yielding every parameter exactly once; the
zero-resident warning; rejection of unknown values in both the manager and
layerwise_tuning_for; the readiness log; and the help text's claims.

End to end, Qwen-Image-2.1 on one H200, the 5090 recipe's shape
(--component-residency dit=resident text_encoder=layerwise-offload vae=resident --layerwise-resident-layers text_encoder=0.8), only the lifetime
flag differing; 1024², 40 steps, seed 42, generator_device=cpu; one warm
request then five measured:

	forward	permanent	Δ
e2e p50	4.778 s [4.752 4.767 4.778 4.799 4.786]	4.547 s [4.546 4.547 4.543 4.556 4.555]	−0.231 s (−4.8%), ranges disjoint
worker RSS after load	18.19 GiB	6.87 GiB	−11.3 GiB: the host copies are gone
placed at load	—	11.19 GiB (29/36 + 22/27 + 2/3 layers)	= 53/66
VRAM peak during requests	36447 MiB	36485 MiB	+38 MiB
output PNG md5	1db98827…	1db98827…	bitwise identical

The −0.231 s is the per-request transfer it removes: 11.19 GiB at this card's
~48 GB/s host-to-device is ~0.23 s. The peak is unchanged because the DiT is
resident in both runs, so the text encoder's set sits on top of the same
footprint either way -- on a card where the sum fits, permanent costs nothing
in peak. nvidia-smi cannot show the forward release at all: the caching
allocator keeps the freed blocks reserved, which is why "VRAM after load" reads
the same in both modes (36437 vs 36473 MiB) and why host RSS, not nvidia-smi,
is the number that moves.

Run again in the opposite order (permanent first, then forward, each a
fresh server) to rule out order effects: permanent p50 4.547 s
[4.545 4.545 4.547 4.554 4.549], forward p50 4.785 s
[4.789 4.785 4.789 4.764 4.772], Δ −0.238 s (−5.0%); RSS 6.87 vs 18.12 GiB;
the same md5 from all four servers. Across both orders every one of the ten
permanent samples is below every one of the ten forward samples (max 4.556 s
vs min 4.752 s). The one thing order did move was the first request of the
first server on the box (6.9 s vs ~4.6-4.8 s everywhere else): a cold page cache
on a freshly downloaded checkpoint, not a lifetime effect, so it is not claimed
above.

CI guard

wan2_1_t2v_1.3b_encoder_permanent_residents (1 GPU): the existing 1.3B case
with --layerwise-resident-layers text_encoder=0.8 --layerwise-residency-lifetime text_encoder=permanent. Baseline seeded by CI on
the H100 pool (b2fefd5822):

	streamed sibling wan2_1_t2v_1.3b	this case
TextEncodingStage	534.89 ms	171.56 ms
load_peak_vram	7814 MiB	17792 MiB
avg denoise step	148 ms	148 ms

The guard is the stage: with the set on the device the stage is compute plus the
streamed 20%; if the lifetime ever regresses to per-request transfer, the stage
lands near the sibling's 535 ms, 3.1× this baseline, against a
non_denoise_stage tolerance of 1.8×. load_peak_vram carrying the set pins
that placement happens at load. Removing the transfer from the measured stage
also makes this case less sensitive to runner disk speed than its forward
siblings.

CI's first run of the seeded case reproduced the baseline (load_peak_vram
17792 MiB to the MiB) and failed only on the consistency check: no ground truth
exists under the new case id. The case now sets run_consistency_check=False
(bd8caa64d6): the lifetime moves weights, not math, the output is the base case's,
wan2_1_t2v_1.3b already guards it, and the H200 runs above showed it bitwise
identical across both lifetimes. A GT copy under a second name would turn this
into a ci-data revision bump for the whole suite.

Not in this PR

• Making the planner choose permanent when the sum fits ([diffusion] Plan component residency from calibrated warmup records (2/4) #37917 / [diffusion] Apply, validate and roll back calibrated residency after warmup (3/4) #37918).
• Loading permanent layers straight from the checkpoint to the device; they
  still pass through host memory once at load, as every layer does today.
• _h2d_bytes_a_pin_would_save ranks components for the pin budget by whole
  module bytes; permanent layers never move, so it slightly over-ranks a
  component that has them. A priority heuristic, not a correctness issue.
• MPS: the path is symmetric with CUDA but was not run here.
• Warning when a component=value map names a component the pipeline does
  not have. Today --layerwise-resident-layers dit=4 is silently ignored
  (the maps take transformer, not the group alias); the CLI reference now
  says so, but a startup warning from the module-level configure, which sees
  every component, would catch it at the source. Separate change.

Relationship to the merged PRs

#40590 (release_after_use / release_all split) and #40592 (retain_resident_layers)
change what finish_use calls; this PR puts its invariant one level lower, in
release_layer, so it composes with both unchanged. #40593 rewrote the same two
help strings and the readiness-log line; the merge keeps its user vocabulary and
adds the lifetime on top, and the log now names the lifetime ((forward) /
(permanent)) where #40593 wrote (per request).

Churn: existing production files +220/−25 over 3 files (layerwise_offload.py
+122/−10, server_args.py +91/−15, layerwise_offload_components.py +7); new
production files 0; tests +308/−1; docs +10/−6.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ✅ Run #35750453665
Latest PR Test (Extra): ✅ Run #35750453457
Latest PR Test (AMD ROCm 10): ❌ Run #35750453551

[mickqian]

CI status note (from the author's triage, so reviewers need not re-derive it).

Every red on this PR falls in one of three buckets, none of them from the diff:

1. NVIDIA multimodal-gen-test-2-gpu (0) — Load Latency (excluding warmup) perf validations only. All 15 validation failures in the latest run are that one metric; no correctness or e2e/denoise metric fails. Measured/expected load on runner h100-novita9-gpu-23: wan2_2_t2v_a14b_2gpu 1.48× on the first attempt and 1.26× on the in-job retry, wan2_1_i2v_14b_480P_2gpu 1.58× → 1.09×, wan2_2_i2v_a14b 1.24×, while the small checkpoints pass (flux2_modelopt_fp8_tp2 0.97×, qwen_image_t2i_2_gpus 1.12×). Second read of the same weights on the same machine is faster and the biggest checkpoints are worst: that is a cold page cache on the runner, not code. The control: [Diffusion] Correct the resident-layer help text to match its scope #40593, whose diff is help text and a doc, fails the same job with the same metric on the same runner pool.
2. AMD extra-a-*, multimodal-gen-test-*-amd — extra-a-test-1-gpu-small-amd fails identically on [Diffusion] Separate a use-scoped layerwise release from release_all #40590, [diffusion] feat: allow a component use retain its layerwise resident set #40592, [Diffusion] Correct the resident-layer help text to match its scope #40593 and [Diffusion] Add a permanent lifetime for layerwise resident layers #40599 (job logs are not uploaded by the runner); multimodal-gen-test-1-gpu-amd (0) is the whole unit-test suite timing out at the 30-minute budget with test_layerwise_offload.py fully PASSED before the cut. The sibling extra-a-test-1-gpu-large-amd passes on every PR.
3. NPU multimodal-gen-test-4-npu-a3 (0) — wan2_2_t2v_14b_w8a8_2npu: [consistency] GT image not found (404 on the pinned ci-data revision), with e2e 191197 ms vs baseline 191226 ms.

diffusion-coverage-check and the *-finish jobs are aggregators of the above (coverage saw 9 of 28 2-GPU cases because the cancelled partitions uploaded no report).

[mickqian]

Added a lightweight e2e guard for the permanent lifetime (wan2_1_t2v_1.3b_encoder_permanent_residents, 1 GPU).

It is the existing wan2_1_t2v_1.3b case with --layerwise-resident-layers text_encoder=0.8 --layerwise-residency-lifetime text_encoder=permanent. The guard is the shape of TextEncodingStage: with the set on the device the stage is compute plus the streamed 20%; if the layers ever go back to being transferred per request, the stage regains the whole ~9 GB and lands near the streamed sibling's 535 ms — about 3× a permanent baseline against the 1.8× non_denoise_stage tolerance. load_peak_vram carries the placed set, so it also pins that placement happens at load, not on the first forward. Removing the transfer from the measured stage also makes this case less sensitive to runner disk speed than its forward siblings.

The baseline is seeded the usual way: this run's multimodal-gen-test-1-gpu partition will fail with Testcase … not found in h100.json and log the add this baseline block; the next commit adds it to perf_baselines/h100.json.

[mickqian]

CI note, B200 lane (multimodal-gen-test-1-b200) on bd8caa6: the four NVFP4 cases (flux1, flux2, ideogram4, qwen_image_2512) failed E2E Latency on runner b200-cirrascale4-4567, 04:57–05:27 UTC. Not this diff, and not main:

• Within one job the same case swings request to request — flux1_modelopt_nvfp4_t2i at 1.75×, 3.41×, 1.75×, 4.01× its baseline; qwen_image_2512 at 1.95×, 4.54×, 4.29×, 1.98×. A code regression is a shift, not a 2–4× random walk; this is a contended or throttled GPU.
• The same lane passed on [Diffusion] Separate a use-scoped layerwise release from release_all #40590, [diffusion] feat: allow a component use retain its layerwise resident set #40592 (04:00 UTC, an hour earlier) and [Diffusion] Correct the resident-layer help text to match its scope #40593.
• git log 50ec9702d0..14e9c40a72 -- python/sglang/multimodal_gen is empty: this PR's base carries no multimodal_gen change the sibling PRs' base lacks.
• Nothing in this PR touches NVFP4, ModelOpt or the B200 suite; the new guard case is in the H100 1-GPU suite only.

One rerun of that job once the run completes. The seeded guard case reproduced its baseline on the previous head (load_peak_vram 17792 MiB to the MiB) and failed only on the missing consistency GT, which bd8caa6 addresses.