[diffusion] feat: let every layerwise component be configurable (#35688)

python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py

@@ -1100,10 +1100,13 @@ class LayerwiseOffloadableModuleMixin:
         self.layerwise_offload_managers = []
         named_modules = dict(self.named_modules())
         configured_layer_names = []
-        # These legacy tuning knobs are explicitly DiT-scoped. Auxiliary
-        # components still support layerwise streaming, but their layers run
-        # once per component use and get no reuse benefit from DiT residency.
-        dit_tuning_enabled = self.layerwise_offload_dit_group_enabled
+        # `--dit-*` is the group default these fall back to, not a scope.
+        prefetch_value, resident_value, residency_policy = (
+            server_args.layerwise_tuning_for(
+                component_name,
+                dit_group=self.layerwise_offload_dit_group_enabled,
+            )
+        )
         for layer_name in self.layer_names:
             module_list = named_modules.get(layer_name)
             if not isinstance(module_list, (torch.nn.ModuleList, torch.nn.Sequential)):
@@ -1112,9 +1115,6 @@ class LayerwiseOffloadableModuleMixin:
                 continue
 
             num_layers = len(module_list)
-            prefetch_value = (
-                server_args.dit_offload_prefetch_size if dit_tuning_enabled else 0.0
-            )
             if current_platform.is_mps() and prefetch_value == 0.0:
                 prefetch_size = 0
             elif prefetch_value < 1.0:
@@ -1122,9 +1122,6 @@ class LayerwiseOffloadableModuleMixin:
             else:
                 prefetch_size = int(prefetch_value)
 
-            resident_value = (
-                server_args.dit_layerwise_resident_layers if dit_tuning_enabled else 0.0
-            )
             if resident_value <= 0:
                 resident_layers = 0
             elif resident_value < 1.0:
@@ -1151,11 +1148,7 @@ class LayerwiseOffloadableModuleMixin:
                 prefetch_size=prefetch_size,
                 resident_layers=resident_layers,
                 initialize=False,
-                residency_policy=(
-                    server_args.dit_layerwise_residency_policy
-                    if dit_tuning_enabled
-                    else RESIDENCY_POLICY_LEADING
-                ),
+                residency_policy=residency_policy,
             )
             self.layerwise_offload_managers.append(manager)
             configured_layer_names.append(layer_name)

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -340,6 +340,15 @@ class ServerArgs(DisaggServerArgsMixin):
     dit_layerwise_resident_layers: float = 0.0
     # Which layers those are: the leading ones, or spread evenly over the stack.
     dit_layerwise_residency_policy: str = RESIDENCY_POLICY_LEADING
+    # Per-component overrides of the three knobs above; an entry wins for that
+    # component.
+    layerwise_prefetch_size: dict[str, float] | str | None = field(default_factory=dict)
+    layerwise_resident_layers: dict[str, float] | str | None = field(
+        default_factory=dict
+    )
+    layerwise_residency_policy: dict[str, str] | str | None = field(
+        default_factory=dict
+    )
     offload_during_compile: bool = True
     text_encoder_cpu_offload: bool | None = None
     image_encoder_cpu_offload: bool | None = None
@@ -986,6 +995,71 @@ class ServerArgs(DisaggServerArgsMixin):
                 f"Available options are: {[e.name.lower() for e in AttentionBackendEnum]}"
             ) from None
 
+    @staticmethod
+    def _parse_component_value_map(
+        value: dict[str, Any] | str | None, *, option: str
+    ) -> dict[str, str]:
+        """Parse a ``component=value`` map, the same shape as component backends."""
+        if value is None or value == "":
+            return {}
+        if isinstance(value, dict):
+            return {str(k): str(v) for k, v in value.items()}
+        if not isinstance(value, str):
+            raise ValueError(
+                f"{option} must be a dict or a comma-separated component=value string"
+            )
+        try:
+            parsed = json.loads(value)
+            if isinstance(parsed, dict):
+                return {str(k): str(v) for k, v in parsed.items()}
+        except json.JSONDecodeError:
+            pass
+        result: dict[str, str] = {}
+        for pair in value.split(","):
+            pair = pair.strip()
+            if not pair:
+                continue
+            if "=" not in pair:
+                raise ValueError(f"{option} must use component=value entries")
+            component, entry = pair.split("=", 1)
+            result[component.strip()] = entry.strip()
+        return result
+
+    def layerwise_tuning_for(
+        self, component_name: str | None, *, dit_group: bool
+    ) -> tuple[float, float, str]:
+        """Prefetch size, resident layers and residency policy for one component."""
+        prefetch_map = self._parse_component_value_map(
+            self.layerwise_prefetch_size, option="--layerwise-prefetch-size"
+        )
+        resident_map = self._parse_component_value_map(
+            self.layerwise_resident_layers, option="--layerwise-resident-layers"
+        )
+        policy_map = self._parse_component_value_map(
+            self.layerwise_residency_policy, option="--layerwise-residency-policy"
+        )
+
+        def _pick(mapping: dict[str, str], group_default, aux_default):
+            if component_name is not None and component_name in mapping:
+                return mapping[component_name]
+            return group_default if dit_group else aux_default
+
+        prefetch = float(_pick(prefetch_map, self.dit_offload_prefetch_size, 0.0))
+        resident = float(_pick(resident_map, self.dit_layerwise_resident_layers, 0.0))
+        policy = str(
+            _pick(
+                policy_map,
+                self.dit_layerwise_residency_policy,
+                RESIDENCY_POLICY_LEADING,
+            )
+        )
+        if policy not in RESIDENCY_POLICIES:
+            raise ValueError(
+                f"unknown residency policy {policy!r} for component "
+                f"{component_name!r}, expected one of {RESIDENCY_POLICIES}"
+            )
+        return prefetch, resident, policy
+
     @staticmethod
     def _parse_component_attention_backend_map(
         value: dict[str, str] | str | None,
@@ -2153,6 +2227,37 @@ class ServerArgs(DisaggServerArgsMixin):
             "once (not re-streamed every step), so this trades VRAM for lower denoise "
             "latency when memory is available.",
         )
+        parser.add_argument(
+            "--layerwise-prefetch-size",
+            type=str,
+            default=None,
+            help="Per-component override of --dit-offload-prefetch-size, as "
+            "component=value entries, e.g. --layerwise-prefetch-size "
+            "text_encoder=2,vae=2. Same units as the DiT flag. Components with "
+            "no entry keep their group default. Prefetch overlaps a layer's "
+            "transfer with the previous layer's compute, which happens within a "
+            "single pass, so it is worth tuning on any streamed component.",
+        )
+        parser.add_argument(
+            "--layerwise-resident-layers",
+            type=str,
+            default=None,
+            help="Per-component override of --dit-layerwise-resident-layers, as "
+            "component=value entries, e.g. --layerwise-resident-layers "
+            "text_encoder=4. Resident layers are transferred once at startup "
+            "rather than streamed, so they cut the transfer of every pass "
+            "including the first -- an auxiliary component that runs once per "
+            "request still benefits, it just recovers the VRAM once per request "
+            "instead of once per denoising step.",
+        )
+        parser.add_argument(
+            "--layerwise-residency-policy",
+            type=str,
+            default=None,
+            help="Per-component override of --dit-layerwise-residency-policy, as "
+            "component=value entries, e.g. --layerwise-residency-policy "
+            "text_encoder=strided.",
+        )
         parser.add_argument(
             "--dit-layerwise-residency-policy",
             type=str,

python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py

@@ -190,6 +190,8 @@ class _TestServerArgs(SimpleNamespace):
     record_component_layerwise_capability = (
         ServerArgs.record_component_layerwise_capability
     )
+    _parse_component_value_map = staticmethod(ServerArgs._parse_component_value_map)
+    layerwise_tuning_for = ServerArgs.layerwise_tuning_for
 
 
 def _server_args(**kwargs):
@@ -210,6 +212,9 @@ def _server_args(**kwargs):
         dit_offload_prefetch_size=1,
         dit_layerwise_resident_layers=0.0,
         dit_layerwise_residency_policy=RESIDENCY_POLICY_LEADING,
+        layerwise_prefetch_size={},
+        layerwise_resident_layers={},
+        layerwise_residency_policy={},
         pin_cpu_memory=False,
         # the pin budget ranks candidates by bytes x steps, and reads the step
         # count off the pipeline's sampling defaults
@@ -1221,3 +1226,59 @@ def test_strided_forward_leaves_exactly_the_resident_set(monkeypatch):
     resident = set(range(8)) - set(manager._streamed_order)
     assert resident <= manager._gpu_layers
     assert len(manager._gpu_layers) <= len(resident) + manager.prefetch_size
+
+
+def test_layerwise_tuning_defaults_match_the_group():
+    """No per-component entry: the DiT group keeps its knobs, auxiliaries do not."""
+    args = _server_args(
+        dit_offload_prefetch_size=3,
+        dit_layerwise_resident_layers=20,
+        dit_layerwise_residency_policy=RESIDENCY_POLICY_STRIDED,
+    )
+    assert args.layerwise_tuning_for("transformer", dit_group=True) == (
+        3.0,
+        20.0,
+        RESIDENCY_POLICY_STRIDED,
+    )
+    assert args.layerwise_tuning_for("text_encoder", dit_group=False) == (
+        0.0,
+        0.0,
+        RESIDENCY_POLICY_LEADING,
+    )
+
+
+def test_layerwise_tuning_per_component_entry_wins():
+    """An auxiliary component can be tuned; its layers cost the same per pass."""
+    args = _server_args(
+        dit_offload_prefetch_size=3,
+        dit_layerwise_resident_layers=20,
+        layerwise_prefetch_size="text_encoder=2",
+        layerwise_resident_layers="text_encoder=4",
+        layerwise_residency_policy={"text_encoder": RESIDENCY_POLICY_STRIDED},
+    )
+    assert args.layerwise_tuning_for("text_encoder", dit_group=False) == (
+        2.0,
+        4.0,
+        RESIDENCY_POLICY_STRIDED,
+    )
+    # an entry for one component leaves every other component alone
+    assert args.layerwise_tuning_for("vae", dit_group=False) == (
+        0.0,
+        0.0,
+        RESIDENCY_POLICY_LEADING,
+    )
+    assert args.layerwise_tuning_for("transformer", dit_group=True)[:2] == (3.0, 20.0)
+
+
+def test_layerwise_tuning_rejects_unknown_policy():
+    args = _server_args(layerwise_residency_policy="vae=sideways")
+    with pytest.raises(ValueError, match="unknown residency policy"):
+        args.layerwise_tuning_for("vae", dit_group=False)
+
+
+def test_layerwise_tuning_accepts_json_and_pair_forms():
+    pair = _server_args(layerwise_resident_layers="vae=6,text_encoder=2")
+    assert pair.layerwise_tuning_for("vae", dit_group=False)[1] == 6.0
+    assert pair.layerwise_tuning_for("text_encoder", dit_group=False)[1] == 2.0
+    as_json = _server_args(layerwise_resident_layers='{"vae": 6}')
+    assert as_json.layerwise_tuning_for("vae", dit_group=False)[1] == 6.0