[diffusion] refactor: unify component residency controls (#34736)

docs/cookbook/diffusion/Krea/Krea-2.mdx

@@ -245,16 +245,18 @@ SGLANG_CACHE_DIT_MC=3 \
 sglang serve --model-path krea/Krea-2-Raw
 ```
 
-#### 4.2.2 Memory & CPU Offload
+#### 4.2.2 Memory and Component Residency
 
 Krea-2's DiT is ~24 GB in bf16 (the bulk of the model). On memory-constrained GPUs you can keep less of it resident:
 
-- `--dit-layerwise-offload`: stream the DiT's transformer blocks layer-by-layer with async host-to-device prefetch overlap, so only a small working set stays on the GPU. This is the primary way to fit Krea-2 on a single consumer / 32 GB-class card, at a modest latency cost. Tune the memory/latency trade-off with `--dit-offload-prefetch-size` (`0.0` prefetches one layer for the lowest memory; larger values prefetch more layers -- faster but more memory).
-- `--dit-cpu-offload`: keep the whole DiT in host memory. Combine it with `--dit-layerwise-offload` for the lowest peak GPU memory (weights stay on host and only the layers needed for the current step are brought on-device).
-- `--text-encoder-cpu-offload`: offload the Qwen3-VL text encoder (it is idle during the denoise loop).
-- `--vae-cpu-offload`: offload the VAE.
+- `--component-residency dit=layerwise-offload`: stream the DiT's transformer blocks layer-by-layer with async host-to-device prefetch overlap, so only a small working set stays on the GPU. This is the primary way to fit Krea-2 on a single consumer / 32 GB-class card, at a modest latency cost. Tune the memory/latency trade-off with `--dit-offload-prefetch-size` (`0.0` prefetches one layer for the lowest memory; larger values prefetch more layers -- faster but more memory).
+- `--component-residency dit=component-offload`: keep the complete DiT on CPU between denoising uses. This and layerwise offload are distinct modes; do not combine them for the same component.
+- `--component-residency text_encoder=component-offload`: offload the Qwen3-VL text encoder while it is idle during denoising.
+- `--component-residency vae=component-offload`: offload the VAE between uses.
 - `--pin-cpu-memory`: pin host memory for offload. Add only as a temporary workaround if you hit `CUDA error: invalid argument`.
 
+The legacy `--dit-layerwise-offload`, `--dit-cpu-offload`, `--text-encoder-cpu-offload`, and `--vae-cpu-offload` forms remain accepted. If both legacy DiT offload flags are enabled, layerwise offload is the effective DiT mode.
+
 On large-VRAM GPUs (e.g. H200), keep everything resident (offloads off) for the fastest latency.
 
 ## 5. Benchmark

docs/cookbook/diffusion/README.mdx

@@ -32,6 +32,14 @@ Each recipe provides step-by-step instructions to help you quickly implement SGL
 3. Adapt configurations to your specific hardware and requirements
 4. Join our community to share feedback and improvements
 
+For memory placement, prefer the unified
+[`--component-residency`](/docs/sglang-diffusion/api/cli#component-residency)
+selector. Each component resolves to exactly one of `resident`,
+`component-offload`, or `layerwise-offload`. Existing options such as
+`--dit-cpu-offload`, `--text-encoder-cpu-offload`,
+`--image-encoder-cpu-offload`, and `--vae-cpu-offload` remain supported by all
+recipes that already use them.
+
 The sglang diffusion cookbook directory structure are shown below:
 
 ```text Example

docs/docs/sglang-diffusion/api/cli.mdx

@@ -200,33 +200,42 @@ For supported native pipelines, set `SGLANG_CACHE_DIT_ENABLED=true` to enable Ca
 
 For supported image pipelines, breakable CUDA graph can be enabled with `--enable-breakable-cuda-graph`, but you must declare every served resolution in `--warmup-resolutions` so warmup captures matching graph signatures.
 
-### Component CPU Offload
+### Component Residency
 
-Use `--cpu-offload-components` to explicitly select components for coarse CPU offload:
+Use `--component-residency COMPONENT=MODE` to assign one runtime residency mode to each native pipeline component:
 
 ```bash Command
 sglang generate \
   --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
-  --cpu-offload-components dit text_encoder \
+  --component-residency all=resident text_encoder=layerwise-offload vae=component-offload \
   --prompt "A quiet city street after rain"
 ```
 
-Component names are matched against the loaded pipeline's component keys from `model_index.json`, so names such as `transformer_2`, `audio_vae`, and `connectors` are supported without adding them to a registry. The group aliases `dit`, `text_encoder`, `image_encoder`, and `vae` remain available. Use `all` to offload every loaded `torch.nn.Module`, or `none` to disable coarse component CPU offload. Non-module components such as tokenizers and schedulers are unaffected. This unified option cannot be combined with the legacy per-component CPU offload flags. Layerwise offload remains independently controlled by `--layerwise-offload-components`.
+The available modes are:
 
-### Layerwise Offload
+- `resident`: keep the complete component on the accelerator.
+- `component-offload`: keep the complete component on CPU between uses, moving it to the accelerator before each declared use and back to CPU afterward.
+- `layerwise-offload`: keep component weights on CPU and stream its declared layers during execution.
 
-Use layerwise offload when a large component does not fit comfortably in GPU memory. By default, `--dit-layerwise-offload` only applies to legacy DiT components. Use `--layerwise-offload-components` to select pipeline component names explicitly (`--layerwise-offload-modules` is accepted as an alias):
+Selectors match exact loaded component keys from `model_index.json`, including names such as `transformer_2`, `audio_vae`, and `connectors`. The group selectors `dit`, `text_encoder`, `image_encoder`, and `vae` are also available, together with `all`. An exact key overrides a matching group, and a group overrides `all`. Components without a matching canonical selector retain their explicit legacy setting or automatic/model default.
+
+The existing `--dit-cpu-offload`, `--text-encoder-cpu-offload`, `--image-encoder-cpu-offload`, `--vae-cpu-offload`, and `--cpu-offload-components` options remain supported. New and legacy options may be mixed: `--component-residency` wins only for components it matches, while unmatched legacy settings remain effective. Legacy layerwise selectors take precedence over legacy component-offload selectors for the same component. Explicit `--dit-layerwise-offload false` makes the DiT resident unless another explicit DiT selector, such as `--dit-cpu-offload true` or `--component-residency dit=component-offload`, selects a different mode.
+
+Layerwise selection is strict. A native weighted component selected for `layerwise-offload` must declare its layer structure; otherwise startup fails with the unsupported component name instead of silently changing modes. FSDP applies only to resident components. The Diffusers backend supports only pipeline-wide `all=resident` and `all=component-offload`.
+
+### Layerwise Offload Tuning
+
+Use layerwise offload when a component does not fit comfortably in GPU memory. The compatibility options `--dit-layerwise-offload` and `--layerwise-offload-components` remain available (`--layerwise-offload-modules` is an alias), while new deployments can select the mode directly:
 
 ```bash Command
 sglang generate \
   --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \
-  --dit-layerwise-offload \
-  --layerwise-offload-components transformer text_encoder \
+  --component-residency transformer=layerwise-offload text_encoder=layerwise-offload \
   --dit-offload-prefetch-size 0 \
   --prompt "A quiet city street after rain"
 ```
 
-The values must match keys in the selected pipeline's `pipeline.modules`, such as `transformer`, `text_encoder`, `image_encoder`, `vae`, `condition_image_encoder`, `spatial_upsampler`, or `vocoder`. Use `all` to select every layerwise-offloadable component. Prefer the smallest component set that solves the memory issue because layerwise offload can increase latency.
+Layerwise tuning options such as `--dit-offload-prefetch-size`, `--dit-layerwise-resident-layers`, and `--dit-layerwise-residency-policy` continue to control the streamed layer working set. Prefer the smallest component set that solves the memory issue because layerwise offload can increase latency.
 
 ## Serve
 

docs/docs/sglang-diffusion/deployment_cookbook.mdx

@@ -1,10 +1,10 @@
 ---
 title: "Deployment and Performance Modes"
-description: "Choose CPU offload, FSDP, CFG parallelism, SP, TP, and performance-mode presets in SGLang Diffusion."
+description: "Choose component residency, FSDP, CFG parallelism, SP, TP, and performance-mode presets in SGLang Diffusion."
 tag: "preserve"
 ---
 
-This page gives practical defaults for choosing `--performance-mode`, CPU offload, FSDP, CFG parallelism, SP, and TP.
+This page gives practical defaults for choosing `--performance-mode`, component residency, FSDP, CFG parallelism, SP, and TP.
 
 ## Quick Rule
 
@@ -24,15 +24,15 @@ Use the simplest setting that fits your memory target:
   <tbody>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Fastest single-GPU run when the model fits</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Disable CPU offload and do not use FSDP.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use resident components and do not use FSDP.</td>
     </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Lower single-GPU memory usage</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use component CPU offload, or layerwise DiT offload for supported Wan/MOVA models.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use component offload, then layerwise offload when a complete component still does not fit comfortably.</td>
     </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Faster multi-GPU Qwen/Wan CFG generation</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use FSDP with CFG parallelism and disable CPU offload.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use FSDP with CFG parallelism and keep the sharded component resident.</td>
     </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Sequence length or video-shape scaling</td>
@@ -48,7 +48,7 @@ Use the simplest setting that fits your memory target:
 Base the decision on available memory on the selected GPU(s).
 
 - For multi-GPU deployment: the least-free selected GPU is the bottleneck. A busy 80GiB GPU can behave like a much smaller GPU.
-- For single-GPU deployment: FSDP shards DiT weights across multiple GPUs. It is not useful for keeping a single-GPU deployment on one GPU; for that case use CPU offload.
+- For single-GPU deployment: FSDP shards weights across multiple GPUs. It is not useful for keeping a single-GPU deployment on one GPU; use component or layerwise offload instead.
 
 ## Health Probes
 
@@ -130,7 +130,17 @@ See [OpenAI API: Served model name](/docs/sglang-diffusion/api/openai_api#served
 
 `auto` checks selected GPU memory before applying FSDP. In multi-GPU runs it uses the least available memory across selected GPUs, and only turns on FSDP automatically when doing so can replace DiT offload. For image workloads with at least 45 GiB available per selected GPU, it keeps the repeatedly reused DiT resident and uses layerwise offload for large auxiliary encoders; below that threshold it keeps the DiT offloaded. Video DiT residency remains model- and workload-specific because frame count and resolution change its peak memory substantially. When the model default uses CFG and the user did not set a parallelism policy, `auto` may also enable CFG parallelism. `speed` intentionally does not check memory; it is the mode for users who prefer latency/throughput and accept OOM risk. It keeps `torch.compile` disabled by default because its effect varies by model and workload. A model-specific deployment config may enable a validated compile path, and `--enable-torch-compile true` always opts in explicitly.
 
-The modes tune residency for native pipeline components declared to the component residency manager. Today this covers the major DiT, text/image encoder, VAE, vocoder, and upsampler components; DiT can use layerwise offload when supported, while text encoders use either resident execution or component CPU offload. Do not assume text-encoder layerwise offload unless a model implements and validates it.
+The modes tune native pipeline components declared to the component residency manager. DiTs, text/image encoders, VAEs, vocoders, adapters, and upsamplers can use layerwise offload when their native module declares its executable layer structure. Explicitly selecting an unsupported component fails at startup instead of falling back to another residency mode.
+
+For direct control, assign one of `resident`, `component-offload`, or `layerwise-offload` with `--component-residency COMPONENT=MODE`:
+
+```bash
+sglang generate \
+  --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --component-residency dit=layerwise-offload text_encoder=component-offload vae=resident
+```
+
+Existing per-component CPU-offload and layerwise flags remain supported. Canonical selectors override matching legacy settings only; unmatched legacy settings and automatic defaults remain effective. See [Component Residency](/docs/sglang-diffusion/api/cli#component-residency) for the complete precedence rules.
 
 When `torch.compile` is enabled, `--offload-during-compile` stays on by default. During compile warmup it temporarily offloads the DiT and evicts resident non-DiT components so `max-autotune` fits on tighter-memory GPUs, then restores the configured serving residency before real traffic.
 
@@ -168,11 +178,11 @@ In this example, `auto` will not re-enable FSDP. The same applies to parallelism
 
 ## Interpreting The Levers
 
-**No offload** keeps model components resident on GPU. It is usually fastest when memory is sufficient.
+**Resident** keeps the complete component on the accelerator. It is usually fastest when memory is sufficient.
 
-**Component CPU offload** lowers GPU memory by moving large components to CPU. It is simple and robust, but it usually trades latency for memory.
+**Component offload** keeps a complete component on CPU between declared uses. It is simple and robust, but each use pays a whole-component transfer.
 
-**Layerwise DiT offload** lowers DiT memory further for supported Wan/MOVA models by moving DiT layers between CPU and GPU. It can be the best single-GPU memory mode, but may increase latency and lower throughput.
+**Layerwise offload** streams the declared layers of any supported native weighted component. It lowers peak accelerator memory further, but may increase latency and lower throughput.
 
 **FSDP** shards DiT weights across multiple GPUs and all-gathers weights during forward. It can reduce DiT CPU offload cost on multi-GPU deployments, especially for validated Wan I2V workloads.