[diffusion] Plan reversible residency before warmup

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

@@ -88,7 +88,7 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis
 - PEFT `adapter_config.json` semantics are applied automatically for named adapter slots, RSLoRA, and per-layer `alpha_pattern`. Adapters that require unsupported auxiliary parameters or runtime behavior, such as DoRA, fail before weight injection instead of silently using ordinary LoRA math.
 - `--lora-merge-mode {auto|merge|dynamic}`: choose how LoRA is applied. `auto` statically merges regular weights and uses dynamic LoRA for FSDP-sharded weights to avoid full-gather peaks.
 - `--num-gpus {N}`: number of GPUs to use
-- `--performance-mode {manual|auto|speed|memory}` / `--mode`: preset for latency/throughput and memory defaults. `auto` is the default. With an eligible synthetic server warmup, it starts from a load-safe placement, measures the representative workload, and jointly selects component/layer residency against every rank and execution-phase VRAM constraint plus HostPin limits. These are simultaneous constraints on one placement decision, not independently spendable budgets. Warmup also observes each declared layer group's call count, so repeated groups can receive independent resident-layer allocations while one-shot or unused groups use HostPin or whole-component residency instead. Multi-process HostPin planning stays inside non-overlapping per-worker shares. Explicit/online quantization and direct GPU loading fix only their affected components; eligible encoders, VAEs, and other components in the same pipeline still participate. The selected placement targets the model-default (or `--warmup-sampling-params`) workload; use representative warmup parameters or explicit residency for larger production shapes, batches, or output counts. Candidate planning uses live allocated memory, while the single post-placement validation also checks allocator-reserved memory so the selected layout retains request headroom. Pipelines that cannot preserve their numerical path across a runtime residency change keep their static placement. Layerwise prefetch depth and residency policy remain explicit because their latency effect depends on runtime overlap rather than weight size alone. Offline, realtime, and no-warmup paths retain the model/card fallback policy. It uses FSDP only for validated DiT-offload replacement paths. `speed` keeps `torch.compile` disabled unless a model-specific deployment config opts in after validation; pass `--enable-torch-compile true` to enable it explicitly. Use `manual` to keep performance-related server args under explicit user control. Explicit offload, FSDP, and parallelism flags take precedence in all modes.
+- `--performance-mode {manual|auto|speed|memory}` / `--mode`: preset for latency/throughput and memory defaults. `auto` is the default. With an eligible synthetic server warmup, it starts from a load-safe checkpoint-weight placement, measures the representative workload, and jointly selects component/layer residency against every rank and execution-phase VRAM constraint plus node-wide HostPin limits. These are simultaneous constraints on one placement decision, not independently spendable budgets. Warmup also observes each declared layer group's call count, so repeated groups can receive independent resident-layer allocations and implicit residency policies while one-shot or unused groups use HostPin or whole-component residency instead. Multi-process workers begin with non-overlapping HostPin shares; the joint solve may reassign them asymmetrically without exceeding the node allowance. Explicit/online quantization and direct GPU loading fix only their affected components; eligible encoders, VAEs, and other components in the same pipeline still participate. The selected placement targets the model-default (or `--warmup-sampling-params`) workload with one output per prompt; use explicit residency or representative warmup parameters for larger production shapes, batches, or output counts. Candidate planning uses live allocated memory, while the single post-placement validation also checks allocator-reserved memory so the selected layout retains request headroom. Pipelines that cannot preserve their numerical path across a runtime residency change keep their static placement. Layerwise prefetch depth remains explicit because its latency effect depends on runtime overlap rather than weight size alone. Explicit residency policies also remain authoritative; otherwise native policy candidates may be tuned. Offline, realtime, and no-warmup paths retain the static fallback policy. It uses FSDP only for validated DiT-offload replacement paths. `speed` keeps `torch.compile` disabled unless a model-specific deployment config opts in after validation; pass `--enable-torch-compile true` to enable it explicitly. Use `manual` to keep performance-related server args under explicit user control. Explicit offload, FSDP, and parallelism flags take precedence in all modes.
 - `--direct-gpu-weight-loading {true|false}`: opt into direct GPU loading for an unquantized, GPU-resident, TP=1 DiT by materializing its complete checkpoint state dict on GPU. Compatible tensors become parameter storage without another GPU copy, while tensors requiring transforms can still need temporary allocations. It is disabled by default; startup time and peak GPU memory remain model-dependent, so benchmark the target model before deployment. It is incompatible with DiT CPU/layerwise offload and FSDP.
 - `--tp-size {N}`: tensor parallelism size. Depending on the pipeline, it can shard the DiT, one or more encoders, or both.
 - `--sp-degree {N}`: sequence parallelism size

docs/docs/sglang-diffusion/deployment_cookbook.mdx

@@ -128,11 +128,11 @@ See [OpenAI API: Served model name](/docs/sglang-diffusion/api/openai_api#served
   </tbody>
 </table>
 
-With an eligible `sglang serve` warmup, `auto` records allocated-memory peaks for the representative request, active component layouts, and temporary per-layer call counts for every declared layer group. One placement decision then satisfies every rank and execution-phase VRAM constraint together with HostPin and HostPin-transition limits. These are simultaneous constraints on one decision vector, not separate planners or independently spendable pools. The target frontier covers component offload, resident components, independent resident counts for repeatedly executed layer groups, and pinned layer-store reassignment; one-shot or unused groups do not consume stage-resident VRAM. Explicit/online quantization and direct GPU loading keep only their affected components fixed, so unrelated encoders and VAEs remain eligible. A post-adjustment warmup realizes the selected state and validates its memory peak before the server becomes ready. A round containing only moves to resident uses one full-shape denoising step; demotions and layerwise or HostPin rebalancing retain the longer timing sample. Validation may keep or roll back the one selected placement, but never starts another search round.
+With an eligible `sglang serve` warmup, `auto` records allocated-memory peaks for the representative request, active component layouts, and temporary per-layer call counts for every declared layer group. One serving-placement decision then satisfies every rank and execution-phase VRAM constraint together with node-wide HostPin and HostPin-transition limits. These are simultaneous constraints on one decision vector, not separate planners or independently spendable pools. The earlier preload plan is only a replaceable load-feasibility seed: after exact modules exist, the serving planner rebuilds the complete frontier and may replace every seeded choice. This preserves the optimum of the supplied serving frontier, but does not claim to jointly minimize startup and serving latency over every possible lifecycle. The target frontier covers component offload, resident components, independent resident counts and implicit residency policies for repeatedly executed layer groups, and pinned layer-store reassignment; one-shot or unused groups do not consume stage-resident VRAM. Explicit/online quantization and direct GPU loading keep only their affected components fixed, so unrelated encoders and VAEs remain eligible. A post-adjustment warmup realizes the selected state and validates its memory peak before the server becomes ready. A round containing only moves to resident uses one full-shape denoising step; demotions and layerwise or HostPin rebalancing retain the longer timing sample. Validation may keep or roll back the one selected placement, but never starts another search round.
 
-This optimization targets the model-default workload, batch size 1, unless `--warmup-sampling-params` supplies a representative request. Larger production shapes or batches should set that override or use explicit `--component-residency`. Layerwise prefetch depth and residency policy remain user-controlled because their latency effect depends on measured copy/compute overlap; the planner does not infer them from weight bytes. Multi-process runs use non-overlapping per-worker HostPin shares, which prevents node overcommit but can be more conservative than a node-global asymmetric allocation.
+This optimization targets the model-default workload and one output per prompt unless `--warmup-sampling-params` supplies a representative request. Larger production shapes, batches, or output counts should set that override or use explicit `--component-residency`. Layerwise prefetch depth remains user-controlled because its latency effect depends on measured copy/compute overlap. An explicitly selected residency policy also remains authoritative; otherwise the planner may choose among the native policy candidates. Multi-process runs begin with non-overlapping per-worker HostPin shares. The joint placement solve may then assign asymmetric per-rank quotas while keeping their sum within the same node-wide allowance, so HostPin is optimized with VRAM rather than by a separate planner.
 
-Before warmup, `auto` still uses load-safe model/card defaults. It checks selected GPU memory before applying FSDP, uses the least available memory across selected GPUs, and only enables FSDP automatically on validated paths where it can replace DiT offload. Offline, realtime, no-warmup, and unsupported serving paths retain those static defaults. When the model default uses CFG and the user did not set a parallelism policy, `auto` may also enable CFG parallelism. `memory` favors lower memory, while `speed` favors latency/throughput and accepts OOM risk. `speed` keeps `torch.compile` disabled by default because its effect varies by model and workload; `--enable-torch-compile true` always opts in explicitly.
+Before warmup, `auto` inventories the selected checkpoint weights on every rank and may seed native DiTs as resident while keeping auxiliary components on their configured loading path. This seed reserves activation headroom and is not a second serving optimizer: weight size alone cannot safely promote encoders or VAEs before the target shape is measured. Each rank reads its checkpoint independently; neither this seed nor layerwise initialization introduces a load-time weight all-gather. Auto parallelism still uses the least available memory across selected GPUs and only enables FSDP on validated paths where it can replace DiT offload. Offline, realtime, no-warmup, and unsupported serving paths retain static defaults. When the model default uses CFG and the user did not set a parallelism policy, `auto` may also enable CFG parallelism. `memory` favors lower memory, while `speed` favors latency/throughput and accepts OOM risk. `speed` keeps `torch.compile` disabled by default because its effect varies by model and workload; `--enable-torch-compile true` always opts in explicitly.
 
 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.
 
@@ -213,6 +213,8 @@ A 1.3B video DiT moves about 2.6 GB a step, on the order of a tenth of a one-sec
 
 **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.
 
+Checkpoint loading does not use that forward-time collective as a loading primitive. Each rank reads the checkpoint independently and constructs its final local DTensor storage directly; the fallback for a full checkpoint tensor slices locally instead of redistributing from another rank. Layerwise offload follows the same rule: its CPU layer stores contain that rank's local DTensor shards and wrap them back with `DTensor.from_local`, without a load-time all-gather. FSDP forward execution still performs the collectives required by its sharding policy.
+
 FSDP sharding granularity matters. SGLang Diffusion prefers sharding direct repeated transformer block entries such as `transformer_blocks.0` or `blocks.0`. Coarser sharding lowers wrapper count but can increase all-gather peak memory; finer sharding can reduce transient memory but adds communication and scheduling overhead. If a model does not define an explicit sharding rule, the loader falls back to repeated block class names and common direct numbered block paths.
 
 **CFG parallelism** splits positive and negative CFG branches across GPUs. For Qwen/Wan workloads with normal step counts, this is the most reliable multi-GPU speedup observed so far.

python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py

@@ -51,10 +51,10 @@ class GetDisaggStatsReq(msgspec.Struct):
 
 
 class AutoResidencyReq(msgspec.Struct, frozen=True):
-    """Apply or roll back the warmup-calibrated residency promotion.
+    """Plan, apply, or roll back automatic component residency.
 
     Sent by the server warmup orchestrator after the synthetic warmup
     requests finish and before the server reports ready.
     """
 
-    action: str = "apply"  # "apply" | "validate" | "rollback"
+    action: str = "apply"  # "apply" | "apply_static" | "validate" | "rollback"

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -40,7 +40,10 @@
 from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
 from sglang.multimodal_gen.runtime.server_warmup import (
     run_sync_client_warmup,
+    run_sync_startup_warmup,
+    should_apply_pre_warmup_auto_residency,
     should_run_explicit_client_warmup,
+    should_run_synthetic_server_warmup,
 )
 from sglang.multimodal_gen.runtime.utils.logging_utils import (
     GREEN,
@@ -191,6 +194,11 @@ def _start_local_server_if_needed(
         return processes
 
     def _run_client_warmup_if_needed(self) -> None:
+        if should_apply_pre_warmup_auto_residency(
+            self.server_args
+        ) or should_run_synthetic_server_warmup(self.server_args):
+            run_sync_startup_warmup(self.server_args, sync_scheduler_client.forward)
+            return
         if not should_run_explicit_client_warmup(self.server_args):
             return
 

python/sglang/multimodal_gen/runtime/entrypoints/http_server.py

@@ -36,7 +36,9 @@
 from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
 from sglang.multimodal_gen.runtime.server_warmup import (
     maybe_apply_auto_residency,
+    maybe_apply_pre_warmup_auto_residency,
     run_async_client_warmup,
+    should_apply_pre_warmup_auto_residency,
     should_run_synthetic_server_warmup,
 )
 from sglang.multimodal_gen.runtime.utils.logging_utils import (
@@ -84,21 +86,31 @@ async def _run_server_warmup_after_http_live(
     server_args: ServerArgs, warmup_done: asyncio.Event
 ) -> None:
     try:
-        if not should_run_synthetic_server_warmup(server_args):
+        run_static_planner = should_apply_pre_warmup_auto_residency(server_args)
+        run_warmup = should_run_synthetic_server_warmup(server_args)
+        if not run_static_planner and not run_warmup:
             warmup_done.set()
             return
 
         await _wait_until_http_live(server_args)
 
-        await run_async_client_warmup(
-            server_args,
-            async_scheduler_client.forward,
-            fail_open=server_args.warmup_resolutions is None,
-        )
-        # Freeze the auto residency decision before the server reports ready.
-        # Raises only when a rollback failed, which aborts startup below.
-        await maybe_apply_auto_residency(server_args, async_scheduler_client.forward)
-        logger.info("The server is fired up and ready to roll!")
+        if run_static_planner:
+            await maybe_apply_pre_warmup_auto_residency(
+                server_args, async_scheduler_client.forward
+            )
+
+        if run_warmup:
+            await run_async_client_warmup(
+                server_args,
+                async_scheduler_client.forward,
+                fail_open=server_args.warmup_resolutions is None,
+            )
+            # Freeze the auto residency decision before the server reports ready.
+            # Raises only when a rollback failed, which aborts startup below.
+            await maybe_apply_auto_residency(
+                server_args, async_scheduler_client.forward
+            )
+            logger.info("The server is fired up and ready to roll!")
         warmup_done.set()
     except asyncio.CancelledError:
         raise

python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py

@@ -7,5 +7,9 @@
 
 
 class BridgeLoader(PlainStateDictComponentLoader):
+    """Loader for MOVA dual tower bridge with FSDP support."""
+
+    supports_fsdp_loading = True
+
     component_names = ["dual_tower_bridge"]
     config_classes = {"dual_tower_bridge": MOVADualTowerConfig}

python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py

@@ -149,6 +149,18 @@ class ComponentLoader(ABC):
     # diffusers or transformers
     expected_library: str = ""
 
+    # --attention-backend primarily selects the DiT backend. Auxiliary
+    # components may fall back when that global choice is incompatible; an
+    # explicit --component-attention-backends entry remains strict.
+    allow_global_attention_backend_fallback = True
+    # Gates only --component-quantizations.<name>. Quantization declared by a
+    # checkpoint is discovered and admitted by the component's normal loader.
+    supports_online_quantization_override = False
+    supports_fsdp_loading = False
+    # Gates only --component-direct-gpu-weight-loading.<name>. The checkpoint
+    # source stays component-specific because its streaming ABI is loader-owned.
+    supports_direct_gpu_weight_loading = False
+    supports_fsdp_inference = False
     _loaders_registered = False
 
     def __init_subclass__(cls, **kwargs):
@@ -1005,6 +1017,19 @@ class PipelineComponentLoader:
     Utility class for loading the components in a pipeline.
     """
 
+    @staticmethod
+    def supports_fsdp_loading(
+        component_name: str,
+        transformers_or_diffusers: str,
+        component_architecture: str | None = None,
+    ) -> bool:
+        loader = ComponentLoader.for_component_type(
+            component_name,
+            transformers_or_diffusers,
+            component_architecture,
+        )
+        return loader.supports_fsdp_loading
+
     @staticmethod
     def load_component(
         component_name: str,

python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py

@@ -155,6 +155,10 @@ def _server_args_for_transformer_component(
 class TransformerLoader(OnlineQuantizationComponentLoader):
     """Shared loader for (video/audio) DiT transformers."""
 
+    allow_global_attention_backend_fallback = False
+    supports_online_quantization_override = True
+    supports_fsdp_loading = True
+    supports_fsdp_inference = True
     component_names = [
         "transformer",
         "unconditional_transformer",