From 7f8f030000b628ea2cb033e7457a13dd0ac80f99 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 20 Aug 2026 22:38:05 +0800 Subject: [PATCH] [diffusion] feat: let every layerwise component be configurable (#35688) --- docs/docs/sglang-diffusion/api/cli.mdx | 16 +++ .../memory_managers/layerwise_offload.py | 23 ++-- .../runtime/server_args/server_args.py | 105 ++++++++++++++++++ .../test/unit/test_layerwise_offload.py | 61 ++++++++++ 4 files changed, 190 insertions(+), 15 deletions(-) diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 3498d730b..7be087b22 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -246,6 +246,22 @@ Values passed to the compatibility option `--layerwise-offload-components` must 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. +Those three set the default for every streamed component. To give one component its own values, use the `component=value` forms, which also accept JSON: + +```bash +sglang serve --model-path \ + --layerwise-offload-components dit,text_encoder \ + --layerwise-prefetch-size text_encoder=2 \ + --layerwise-resident-layers text_encoder=4 \ + --layerwise-residency-policy text_encoder=strided +``` + +- `--layerwise-prefetch-size`: how many layers to fetch ahead. Fractional values are a share of the stack, `>= 1` an absolute count. Deeper prefetch overlaps more of the transfer with compute, at the cost of staging buffers. +- `--layerwise-resident-layers`: how many layers stay on the GPU instead of being streamed. Resident layers are transferred once at startup, so they are removed from every pass. Fractional values are a share of the stack. +- `--layerwise-residency-policy`: `leading` keeps the first layers, `strided` spreads them across the stack so the transfers do not arrive as one burst. + +A component without an entry keeps the group default, so adding these options changes nothing until one is set. Both knobs trade VRAM for transfer, and the return differs by component: a DiT earns it back once per denoising step, a text encoder or VAE once per request. Measure before raising either. + ## Serve `sglang serve` starts the HTTP server and keeps the model loaded for repeated requests. diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py index 3bcf32398..e3f95ddf5 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 0238c570b..28dbcf149 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/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, diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index e9bd41974..d5f2a77bd 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/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