From 7f66f1cc4ea8f4081799e37fa98cbc59879a8079 Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 16 Aug 2026 21:22:10 +0800 Subject: [PATCH 01/90] [diffusion] feat: warmup-calibrated auto residency promotion Under --performance-mode auto with server warmup, measure per-rank peak GPU memory during the synthetic warmup, extrapolate it to the model's default workload (scaling only the activation part above the pre-forward allocated baseline), and promote implicitly offloaded components (component offload -> resident, layerwise offload -> fully loaded) when the estimate plus promoted weights fits under max(10% VRAM, 4 GiB) reserve. The decision is computed from all-gathered rank reports so every rank promotes identically; any rank failure rolls all ranks back, and the residency freezes before /health turns ready. Orchestration: warmup (measure) -> AutoResidencyReq(apply) fanned out through the existing control-req path -> re-run the synthetic warmup under the final residency (physically moves promoted components with per-use dtype semantics, rebuilds compile caches, and proves the layout fits) -> ready. A failed re-warm rolls back and re-warms the original strategy; only a failed rollback aborts startup. Excluded paths: explicit placement, FSDP, diffusers backend, BCG, cache-dit, dynamic batching, dp>1, realtime/disagg, and quantized checkpoints (residency shifts measurably moved fp8 DiT outputs before). Kill switch: SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY. Also fixes ridden-along bugs: - server warmup frame caps now re-apply the model frame contract (LongLive2's capped 17 frames -> 5 latent frames broke its 8-frame causal block, so every server warmup failed silently under fail-open; now re-aligned to 29) - get_can_stay_resident_components sized components from the loader's GPU-load delta, which is ~0 for exactly the offloaded components it reports on; it now sizes live modules through layerwise CPU buffers - enable_offload() re-registered forward hooks on managers that were never disabled, double-firing every layer hook Co-Authored-By: Claude Fable 5 --- python/sglang/multimodal_gen/envs.py | 7 + .../runtime/entrypoints/http_server.py | 4 + .../runtime/entrypoints/utils.py | 11 + .../runtime/managers/gpu_worker.py | 248 ++++++++- .../memory_managers/auto_residency.py | 458 ++++++++++++++++ .../memory_managers/component_manager.py | 17 +- .../memory_managers/layerwise_offload.py | 3 + .../runtime/managers/scheduler.py | 10 + .../runtime/scheduler_client.py | 2 + .../runtime/server_args/server_args.py | 4 + .../multimodal_gen/runtime/server_warmup.py | 114 ++++ .../test/unit/test_auto_residency.py | 508 ++++++++++++++++++ 12 files changed, 1368 insertions(+), 18 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_auto_residency.py diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index d67500b72877..56e6de4f935e 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -37,6 +37,7 @@ SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB: float | None = None SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB: float | None = None SGLANG_DIFFUSION_STAGE_LOGGING: bool = False + SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY: bool = False SGLANG_DIFFUSION_CFG_GATE_STEP: float = 1.0 # cache-dit env vars (primary transformer) # on by default; engages only on 2 ranks with peer-to-peer access and falls @@ -260,6 +261,12 @@ def _getter(): # If set, sgl_diffusion will enable stage logging, which will print the time # taken for each stage "SGLANG_DIFFUSION_STAGE_LOGGING": _lazy_bool("SGLANG_DIFFUSION_STAGE_LOGGING"), + # Kill-switch for the warmup-calibrated auto residency promotion that runs + # under `--performance-mode auto` with server warmup. Set to disable the + # promotion without giving up the rest of the auto performance policy. + "SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY": _lazy_bool( + "SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY" + ), # Fraction of denoising steps that run both CFG branches before reusing the # last conditional-minus-unconditional residual. Keep 1.0 to disable. "SGLANG_DIFFUSION_CFG_GATE_STEP": _lazy_float( diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index 1b5098acb184..bd4da6eee808 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -35,6 +35,7 @@ from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client 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, run_async_client_warmup, should_run_synthetic_server_warmup, ) @@ -94,6 +95,9 @@ async def _run_server_warmup_after_http_live( 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: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index e964b91a133c..73c21225c18a 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -23,6 +23,7 @@ from typing import Any, Callable, List, Optional, Sequence, Union import imageio +import msgspec import numpy as np import torch from PIL import Image @@ -202,6 +203,16 @@ class GetDisaggStatsReq: pass +class AutoResidencyReq(msgspec.Struct, frozen=True): + """Apply or roll back the warmup-calibrated residency promotion. + + Sent by the server warmup orchestrator after the synthetic warmup + requests finish and before the server reports ready. + """ + + action: str = "apply" # "apply" | "rollback" + + def format_lora_message( lora_nickname: Union[str, List[str]], target: Union[str, List[str]], diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 14a2375633f7..796d01c4219b 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -13,6 +13,7 @@ import numpy as np import torch +import torch.nn as nn from setproctitle import setproctitle from sglang.multimodal_gen.runtime.utils.logging_utils import ( # isort: skip @@ -43,12 +44,40 @@ get_tp_group, get_ulysses_parallel_rank, get_ulysses_parallel_world_size, + get_world_group, ) from sglang.multimodal_gen.runtime.entrypoints.utils import ( materialize_output_sample, post_process_sample, save_outputs, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( + GIB_BYTES, + PROMOTION_STATUS_PROMOTED, + PROMOTION_STATUS_ROLLBACK_FAILED, + PROMOTION_STATUS_ROLLED_BACK, + PROMOTION_STATUS_SKIPPED, + AppliedPromotion, + DefaultWorkload, + RankResidencyReport, + WarmupMemoryRecord, + apply_promotions, + collect_promotion_candidates, + component_resident_size_bytes, + estimate_default_workload_peak_bytes, + format_plan_summary, + plan_auto_residency, + plan_summary_payload, + resolve_default_workload, + rollback_promotions, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( + peek_global_component_residency_manager, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( + COMPONENT_OFFLOAD, + LAYERWISE_OFFLOAD, +) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( configure_layerwise_offload_modules, ) @@ -167,6 +196,10 @@ def __init__( self.cfg_cpu_group = self.cfg_group.cpu_group self._realtime_sessions = RealtimeSessionCache(max_sessions=1) self.memory_occupation: MemoryOccupationController | None = None + # per-rank memory measurements of server warmup forwards; consumed by + # the auto-residency promotion decision before the server turns ready + self._auto_residency_warmup_records: list[WarmupMemoryRecord] = [] + self._auto_residency_applied: list[AppliedPromotion] = [] def release_realtime_session(self, session_id: str) -> OutputBatch: """release the session of a realtime connection""" @@ -491,9 +524,17 @@ def _execute_forward_common( """ output_batch = None forward_failed = False + measure_server_warmup = ( + bool(req.extra.get("server_based_warmup")) and not current_platform.is_cpu() + ) + warmup_baseline_allocated_bytes = 0 try: if not current_platform.is_cpu() and not current_platform.is_mps(): torch.get_device_module().reset_peak_memory_stats() + if measure_server_warmup: + warmup_baseline_allocated_bytes = ( + torch.get_device_module().memory_allocated() + ) start_time = ( execution_start_time @@ -605,8 +646,30 @@ def _execute_forward_common( # clean cache if OOM if not current_platform.is_cpu(): torch.get_device_module().empty_cache() + if measure_server_warmup and output_batch is not None: + self._record_server_warmup_memory( + req=req, + baseline_allocated_bytes=warmup_baseline_allocated_bytes, + succeeded=output_batch.error is None, + ) return output_batch + def _record_server_warmup_memory( + self, *, req: Req, baseline_allocated_bytes: int, succeeded: bool + ) -> None: + self._auto_residency_warmup_records.append( + WarmupMemoryRecord( + width=int(req.width or 0), + height=int(req.height or 0), + num_frames=int(req.num_frames or 1), + baseline_allocated_bytes=int(baseline_allocated_bytes), + peak_reserved_bytes=int( + torch.get_device_module().max_memory_reserved() + ), + succeeded=succeeded, + ) + ) + def _materialize_output_transport( self, output_batch: OutputBatch, @@ -1023,34 +1086,185 @@ def get_can_stay_resident_components( if not self.pipeline: return can_stay_resident - memory_usages = self.pipeline.memory_usages + # Size from the live modules, not pipeline.memory_usages: the loader + # records the GPU-load delta, which is ~0 for exactly the offloaded + # components this recommendation is about. + modules = self.pipeline.modules ordered_names = [ - name - for name in OFFLOAD_DISABLE_RECOMMENDATION_ORDER - if name in memory_usages + name for name in OFFLOAD_DISABLE_RECOMMENDATION_ORDER if name in modules ] ordered_names.extend( - name - for name in memory_usages - if name not in OFFLOAD_DISABLE_RECOMMENDATION_ORDER + name for name in modules if name not in OFFLOAD_DISABLE_RECOMMENDATION_ORDER ) for name in ordered_names: - usage = memory_usages[name] - if not ( - self.server_args.should_cpu_offload_component(name) - or self.server_args.should_configure_layerwise_offload_for_lazy_component( - name - ) - ): + module = modules[name] + if not isinstance(module, nn.Module): continue - if usage is None: + residency_mode = self.server_args.residency_mode(name) + if residency_mode not in (COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD): continue - if usage <= remaining_gpu_mem_gb: + usage_gb = component_resident_size_bytes(module, residency_mode) / GIB_BYTES + if 0 < usage_gb <= remaining_gpu_mem_gb: can_stay_resident.append(name) - remaining_gpu_mem_gb -= usage + remaining_gpu_mem_gb -= usage_gb return can_stay_resident + def apply_auto_residency(self) -> OutputBatch: + """Apply the warmup-calibrated residency promotion on this replica. + + Every rank executes this handler at the same queue position; the plan + is computed from all-gathered rank reports so each rank reaches the + same decision. A failure on any rank rolls every rank back to the + original strategy. + """ + workload = resolve_default_workload(self.server_args) + local_report = self._build_auto_residency_report(workload=workload) + reports = self._auto_residency_all_gather(local_report) + plan = plan_auto_residency(reports=reports) + summary = format_plan_summary(plan=plan, workload=workload) + if plan.skip_reason is not None or not plan.promotions: + if self.is_output_rank: + logger.info("%s", summary) + return OutputBatch( + output=plan_summary_payload( + plan=plan, workload=workload, status=PROMOTION_STATUS_SKIPPED + ) + ) + + apply_error: str | None = None + try: + self._auto_residency_applied = apply_promotions( + plan=plan, + modules=self.pipeline.modules, + server_args=self.server_args, + ) + except Exception as e: # this rank already rolled itself back + logger.error( + "Auto residency promotion failed on rank %d: %s", + self.rank, + e, + exc_info=True, + ) + self._auto_residency_applied = [] + apply_error = str(e) + + rank_errors = [ + error for error in self._auto_residency_all_gather(apply_error) if error + ] + if rank_errors: + return self._rollback_auto_residency_after_failure(cause=rank_errors[0]) + + self._invalidate_component_strategies( + [candidate.component_name for candidate in plan.promotions] + ) + if self.is_output_rank: + logger.info("%s", summary) + return OutputBatch( + output=plan_summary_payload( + plan=plan, workload=workload, status=PROMOTION_STATUS_PROMOTED + ) + ) + + def rollback_auto_residency(self) -> OutputBatch: + """Restore the pre-promotion residency (post-promotion re-warm failed).""" + rollback_error = self._rollback_applied_promotions() + rank_errors = [ + error for error in self._auto_residency_all_gather(rollback_error) if error + ] + if rank_errors: + return OutputBatch( + error=f"auto residency rollback failed: {rank_errors[0]}", + output={"status": PROMOTION_STATUS_ROLLBACK_FAILED}, + ) + return OutputBatch(output={"status": PROMOTION_STATUS_ROLLED_BACK}) + + def _build_auto_residency_report( + self, *, workload: DefaultWorkload + ) -> RankResidencyReport: + records = list(self._auto_residency_warmup_records) + skip_reason = None if records else "no server warmup measurements" + estimated_peak_bytes = estimate_default_workload_peak_bytes( + records=records, target_units=workload.workload_units() + ) + # What this process may still grow into: free VRAM plus what its + # allocator already reserved. Unlike the raw device total, this + # excludes memory held by other processes on a shared GPU. + free_bytes, _ = torch.get_device_module().mem_get_info() + budget_bytes = int(free_bytes) + int( + torch.get_device_module().memory_reserved() + ) + candidates = collect_promotion_candidates( + modules=self.pipeline.modules, + residency_mode_of=self.server_args.residency_mode, + explicit_residency_mode_of=self.server_args.explicit_residency_mode, + custom_strategy_names=self.pipeline.component_residency_strategies, + num_inference_steps=workload.num_inference_steps, + ) + return RankResidencyReport( + rank=self.rank, + budget_bytes=budget_bytes, + estimated_peak_bytes=estimated_peak_bytes, + candidates=candidates, + skip_reason=skip_reason, + ) + + def _auto_residency_all_gather(self, obj: Any) -> list[Any]: + world_group = get_world_group() + if world_group.world_size == 1: + return [obj] + gathered: list[Any] = [None] * world_group.world_size + torch.distributed.all_gather_object(gathered, obj, group=world_group.cpu_group) + return gathered + + def _rollback_applied_promotions(self) -> str | None: + if not self._auto_residency_applied: + return None + applied = self._auto_residency_applied + try: + rollback_promotions( + applied=applied, + modules=self.pipeline.modules, + server_args=self.server_args, + ) + except Exception as e: + logger.error( + "Auto residency rollback failed on rank %d: %s", + self.rank, + e, + exc_info=True, + ) + return str(e) + self._invalidate_component_strategies( + [promotion.component_name for promotion in applied] + ) + self._auto_residency_applied = [] + return None + + def _rollback_auto_residency_after_failure(self, *, cause: str) -> OutputBatch: + rollback_error = self._rollback_applied_promotions() + rank_errors = [ + error for error in self._auto_residency_all_gather(rollback_error) if error + ] + if rank_errors: + return OutputBatch( + error=( + f"auto residency promotion failed ({cause}) and rollback " + f"failed: {rank_errors[0]}" + ), + output={"status": PROMOTION_STATUS_ROLLBACK_FAILED}, + ) + return OutputBatch( + error=f"auto residency promotion failed; original strategy restored: {cause}", + output={"status": PROMOTION_STATUS_ROLLED_BACK}, + ) + + @staticmethod + def _invalidate_component_strategies(component_names: List[str]) -> None: + manager = peek_global_component_residency_manager() + if manager is not None: + manager.invalidate_component_strategies(component_names) + def set_lora( self, lora_nickname: Union[str, List[str]], diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py new file mode 100644 index 000000000000..6e01fb6d19e0 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -0,0 +1,458 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Warmup-calibrated automatic component residency promotion. + +Under ``--performance-mode auto`` with server warmup, each rank measures the +peak GPU memory of the synthetic warmup requests, extrapolates it to the +model's default workload, and promotes implicitly offloaded components +(component offload -> resident, layerwise offload -> fully loaded) when the +estimate plus the promoted weights still fits under a safety reserve. + +The estimate splits the measured peak into a persistent part (weights and +buffers alive before the forward) and an activation part (everything above +it); only the activation part scales with the workload ratio. Scaling the +whole peak would multiply resident weights by the video frame/area cap ratio +(~16x for Wan-class defaults) and promotion would never trigger. + +Promotion targets the model default workload only (default resolution, +default frames, batch=1). Larger shapes, batches, or multi-image inputs need +explicit ``--component-residency``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, Iterable, Mapping + +import msgspec +import torch +import torch.nn as nn + +from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( + COMPONENT_OFFLOAD, + LAYERWISE_OFFLOAD, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import ( + is_fsdp_managed_module, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + LayerwiseOffloadableModuleMixin, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( + is_dit_component_name, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.server_args import ServerArgs + +logger = init_logger(__name__) + +GIB_BYTES = 1024**3 + +# Activation memory rarely scales perfectly linearly with workload units; +# pad the extrapolated activation part before checking the budget. +ACTIVATION_EXTRAPOLATION_MARGIN = 1.2 +# Always keep this much VRAM free after promotion for allocator slack, +# shape variance, and CUDA graph or compile pools. +VRAM_RESERVE_FRACTION = 0.10 +MIN_VRAM_RESERVE_BYTES = 4 * GIB_BYTES + +AUTO_RESIDENCY_FEATURE_NAME = "auto residency promotion" + +PROMOTION_STATUS_SKIPPED = "skipped" +PROMOTION_STATUS_PROMOTED = "promoted" +PROMOTION_STATUS_ROLLED_BACK = "rolled_back" +PROMOTION_STATUS_ROLLBACK_FAILED = "rollback_failed" + + +class WarmupMemoryRecord(msgspec.Struct, frozen=True): + """Per-rank memory measurement of one server warmup forward.""" + + width: int + height: int + num_frames: int + baseline_allocated_bytes: int + peak_reserved_bytes: int + succeeded: bool + + def workload_units(self) -> int: + return max(1, self.width) * max(1, self.height) * max(1, self.num_frames) + + +class PromotionCandidate(msgspec.Struct, frozen=True): + """An implicitly offloaded component that could become resident.""" + + component_name: str + residency_mode: str + promoted_weight_bytes: int + # Estimated per-request host-to-device traffic this promotion removes. + h2d_bytes_per_request: int + + +class DefaultWorkload(msgspec.Struct, frozen=True): + """The model-default request shape the promotion is calibrated for.""" + + width: int | None + height: int | None + num_frames: int + num_inference_steps: int + + def workload_units(self) -> int | None: + if self.width is None or self.height is None: + return None + return max(1, self.width) * max(1, self.height) * max(1, self.num_frames) + + def describe(self) -> str: + if self.width is None or self.height is None: + return "model-default" + return f"{self.width}x{self.height}x{self.num_frames}f" + + +class RankResidencyReport(msgspec.Struct, frozen=True): + """One rank's inputs to the replica-wide promotion decision.""" + + rank: int + budget_bytes: int + estimated_peak_bytes: int | None + candidates: list[PromotionCandidate] = [] + skip_reason: str | None = None + + +class AutoResidencyPlan(msgspec.Struct, frozen=True): + """Deterministic promotion decision shared by every rank.""" + + estimated_peak_bytes: int = 0 + reserve_bytes: int = 0 + budget_bytes: int = 0 + promotions: list[PromotionCandidate] = [] + skip_reason: str | None = None + + +class AppliedPromotion(msgspec.Struct, frozen=True): + component_name: str + residency_mode: str + + +def resolve_default_workload(server_args: ServerArgs) -> DefaultWorkload: + """Resolve the default request shape promotion is optimized for.""" + from sglang.multimodal_gen.runtime.warmup_request_builder import ( + get_model_sampling_defaults, + ) + + defaults = get_model_sampling_defaults(server_args) + width = defaults.width + height = defaults.height + if (width is None or height is None) and defaults.supported_resolutions: + width, height = defaults.supported_resolutions[0] + num_frames = defaults.num_frames or 1 + if num_frames > 1: + num_frames = server_args.pipeline_config.adjust_num_frames(num_frames) + return DefaultWorkload( + width=width, + height=height, + num_frames=num_frames, + num_inference_steps=defaults.num_inference_steps or 1, + ) + + +def estimate_default_workload_peak_bytes( + *, + records: Iterable[WarmupMemoryRecord], + target_units: int | None, +) -> int | None: + """Extrapolate warmup peaks to the default workload. + + Returns None when there is no usable measurement (no records, or any + warmup forward failed -- a failed warmup means the measured peak does not + cover the full request path). + """ + records = list(records) + if not records or any(not record.succeeded for record in records): + return None + + estimates = [] + for record in records: + peak = record.peak_reserved_bytes + ratio = 1.0 if target_units is None else target_units / record.workload_units() + if ratio <= 1.0: + estimates.append(peak) + continue + baseline = min(record.baseline_allocated_bytes, peak) + activation = peak - baseline + estimates.append( + baseline + int(activation * ratio * ACTIVATION_EXTRAPOLATION_MARGIN) + ) + return max(estimates) + + +def _module_weight_bytes(module: nn.Module) -> int: + """Full weight+buffer footprint, reading through layerwise CPU buffers.""" + from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + iter_materialized_weights, + ) + + param_bytes = sum( + tensor.numel() * tensor.element_size() + for _, tensor in iter_materialized_weights(module) + ) + buffer_bytes = sum( + tensor.numel() * tensor.element_size() for tensor in module.buffers() + ) + return param_bytes + buffer_bytes + + +def _layerwise_offloaded_bytes(module: LayerwiseOffloadableModuleMixin) -> int: + total = 0 + for manager in module.layerwise_offload_managers: + if not manager.enabled: + continue + total += sum( + tensor.numel() * tensor.element_size() + for _, tensor in manager.iter_cpu_weights() + ) + return total + + +def component_resident_size_bytes(module: nn.Module, residency_mode: str) -> int: + """Extra GPU bytes a promotion of this component would pin resident.""" + if residency_mode == LAYERWISE_OFFLOAD: + if not isinstance(module, LayerwiseOffloadableModuleMixin): + return 0 + return _layerwise_offloaded_bytes(module) + return _module_weight_bytes(module) + + +def collect_promotion_candidates( + *, + modules: Mapping[str, object], + residency_mode_of: Callable[[str], str], + explicit_residency_mode_of: Callable[[str], str | None], + custom_strategy_names: Iterable[str], + num_inference_steps: int, +) -> list[PromotionCandidate]: + """List implicitly offloaded native components eligible for promotion. + + Only components whose residency was chosen by the auto policy qualify: + explicit placements, FSDP-managed modules, and components driven by a + pipeline-custom residency strategy are never touched. + """ + custom_names = set(custom_strategy_names) + candidates = [] + for name in sorted(modules): + module = modules[name] + if not isinstance(module, nn.Module): + continue + if name in custom_names: + continue + mode = residency_mode_of(name) + if mode not in (COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD): + continue + if explicit_residency_mode_of(name) is not None: + continue + if is_fsdp_managed_module(module): + continue + promoted_bytes = component_resident_size_bytes(module, mode) + if promoted_bytes <= 0: + continue + # A layerwise DiT re-streams its layers once per denoise forward; + # every other offloaded component transfers once per request. + uses_per_request = ( + max(1, num_inference_steps) + if mode == LAYERWISE_OFFLOAD and is_dit_component_name(name) + else 1 + ) + candidates.append( + PromotionCandidate( + component_name=name, + residency_mode=mode, + promoted_weight_bytes=promoted_bytes, + h2d_bytes_per_request=promoted_bytes * uses_per_request, + ) + ) + return candidates + + +def _skip_plan(reason: str) -> AutoResidencyPlan: + return AutoResidencyPlan(skip_reason=reason) + + +def _consensus_candidates( + reports: list[RankResidencyReport], +) -> list[PromotionCandidate]: + """Merge per-rank candidates: keep components every rank agrees on. + + Sizes are worst-cased with the per-rank maximum so a promotion never + fits on one rank but overflows another. + """ + per_rank_maps = [ + {candidate.component_name: candidate for candidate in report.candidates} + for report in reports + ] + common_names = set(per_rank_maps[0]) + for candidate_map in per_rank_maps[1:]: + common_names &= set(candidate_map) + + merged = [] + for name in sorted(common_names): + rank_candidates = [candidate_map[name] for candidate_map in per_rank_maps] + modes = {candidate.residency_mode for candidate in rank_candidates} + if len(modes) != 1: + continue + merged.append( + PromotionCandidate( + component_name=name, + residency_mode=modes.pop(), + promoted_weight_bytes=max( + candidate.promoted_weight_bytes for candidate in rank_candidates + ), + h2d_bytes_per_request=max( + candidate.h2d_bytes_per_request for candidate in rank_candidates + ), + ) + ) + return merged + + +def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyPlan: + """Turn the gathered rank reports into one deterministic promotion plan.""" + if not reports: + return _skip_plan("no rank reports") + for report in reports: + if report.skip_reason is not None: + return _skip_plan(f"rank {report.rank}: {report.skip_reason}") + if report.estimated_peak_bytes is None: + return _skip_plan(f"rank {report.rank}: no usable warmup measurement") + + estimated_peak = max(report.estimated_peak_bytes for report in reports) + budget = min(report.budget_bytes for report in reports) + reserve = max(int(budget * VRAM_RESERVE_FRACTION), MIN_VRAM_RESERVE_BYTES) + + candidates = _consensus_candidates(reports) + if not candidates: + return _skip_plan("no implicitly offloaded components to promote") + + # Biggest per-request H2D savings first; name breaks ties deterministically. + ordered = sorted( + candidates, + key=lambda candidate: ( + -candidate.h2d_bytes_per_request, + candidate.component_name, + ), + ) + promotions = [] + promoted_bytes = 0 + for candidate in ordered: + next_bytes = promoted_bytes + candidate.promoted_weight_bytes + if estimated_peak + next_bytes + reserve <= budget: + promotions.append(candidate) + promoted_bytes = next_bytes + + return AutoResidencyPlan( + estimated_peak_bytes=estimated_peak, + reserve_bytes=reserve, + budget_bytes=budget, + promotions=promotions, + ) + + +def apply_promotions( + *, + plan: AutoResidencyPlan, + modules: Mapping[str, object], + server_args: ServerArgs, +) -> list[AppliedPromotion]: + """Flip the planned components to resident on this rank. + + Component-offload modules are only re-marked resident; the physical move + happens on their next use (the post-promotion re-warmup) through + ``ResidentStrategy``, which also applies any per-use target dtype. + Layerwise modules load all layers now and drop their hooks. + + Raises on failure after undoing the promotions already applied, so a + caller observing an exception knows this rank is back on the original + strategy. + """ + applied: list[AppliedPromotion] = [] + try: + for candidate in plan.promotions: + module = modules.get(candidate.component_name) + if not isinstance(module, nn.Module): + raise RuntimeError( + f"promotion target {candidate.component_name!r} is missing" + ) + # Record before acting so a mid-promotion failure rolls back the + # partially promoted component as well. + applied.append( + AppliedPromotion( + component_name=candidate.component_name, + residency_mode=candidate.residency_mode, + ) + ) + server_args.require_component_resident( + candidate.component_name, feature_name=AUTO_RESIDENCY_FEATURE_NAME + ) + if candidate.residency_mode == LAYERWISE_OFFLOAD: + assert isinstance(module, LayerwiseOffloadableModuleMixin) + module.disable_offload() + except Exception: + rollback_promotions(applied=applied, modules=modules, server_args=server_args) + raise + return applied + + +def rollback_promotions( + *, + applied: Iterable[AppliedPromotion], + modules: Mapping[str, object], + server_args: ServerArgs, +) -> None: + """Restore the original residency for previously applied promotions.""" + for promotion in reversed(list(applied)): + server_args.release_required_component_residency(promotion.component_name) + module = modules.get(promotion.component_name) + if not isinstance(module, nn.Module): + continue + if promotion.residency_mode == LAYERWISE_OFFLOAD: + assert isinstance(module, LayerwiseOffloadableModuleMixin) + module.enable_offload() + else: + module.to("cpu") + if torch.get_device_module().is_available(): + torch.get_device_module().empty_cache() + + +def format_plan_summary(*, plan: AutoResidencyPlan, workload: DefaultWorkload) -> str: + """One-line decision summary for the startup log.""" + if plan.skip_reason is not None: + return f"Auto residency: skipped ({plan.skip_reason})" + promoted = ( + ", ".join( + f"{candidate.component_name}" + f"({candidate.residency_mode}, " + f"{candidate.promoted_weight_bytes / GIB_BYTES:.1f} GiB)" + for candidate in plan.promotions + ) + or "none" + ) + return ( + f"Auto residency: target={workload.describe()} " + f"steps={workload.num_inference_steps}, " + f"estimated_peak={plan.estimated_peak_bytes / GIB_BYTES:.1f} GiB, " + f"reserve={plan.reserve_bytes / GIB_BYTES:.1f} GiB, " + f"budget={plan.budget_bytes / GIB_BYTES:.1f} GiB, " + f"promoted=[{promoted}]" + ) + + +def plan_summary_payload( + *, plan: AutoResidencyPlan, workload: DefaultWorkload, status: str +) -> dict: + """Structured decision summary returned to the warmup orchestrator.""" + return { + "status": status, + "skip_reason": plan.skip_reason, + "target_workload": workload.describe(), + "estimated_peak_bytes": plan.estimated_peak_bytes, + "reserve_bytes": plan.reserve_bytes, + "budget_bytes": plan.budget_bytes, + "promoted": [candidate.component_name for candidate in plan.promotions], + } diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py index d2e5e448966d..8bdcf4b9d458 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py @@ -1,4 +1,4 @@ -from collections.abc import Iterator +from collections.abc import Iterable, Iterator from contextlib import contextmanager from dataclasses import dataclass from typing import Mapping, MutableMapping, Protocol, Sequence @@ -153,6 +153,16 @@ def refresh_server_args(self, server_args: ServerArgs) -> None: self._strategy_cache.clear() self.server_args = server_args + def invalidate_component_strategies(self, component_names: Iterable[str]) -> None: + """Drop cached strategies after an in-place residency change. + + The cache only auto-invalidates when the ``server_args`` object + identity changes; runtime promotions mutate residency on the same + object and must invalidate explicitly. + """ + for component_name in component_names: + self._strategy_cache.pop(component_name, None) + def begin_request( self, stages: Sequence[ComponentResidencyStage], @@ -674,6 +684,11 @@ def _empty_cache_after_large_release( _GLOBAL_COMPONENT_RESIDENCY_MANAGER: ComponentResidencyManager | None = None +def peek_global_component_residency_manager() -> ComponentResidencyManager | None: + """Return the process-global manager without creating one.""" + return _GLOBAL_COMPONENT_RESIDENCY_MANAGER + + def get_global_component_residency_manager( pipeline: ComponentResidencyPipeline, server_args: ServerArgs, 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 fb1e37831c91..3e391a3c46ca 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 @@ -1836,6 +1836,9 @@ def enable_offload(self) -> None: if self.layerwise_offload_managers is None: return for manager in self.layerwise_offload_managers: + if manager.enabled: + # already armed; re-registering would double the forward hooks + continue if manager._configured: manager.enabled = True manager.sync_all_layers_to_cpu() diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index c85af07e94ee..0d1f0e5f9d17 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -25,6 +25,7 @@ UpdateWeightFromTensorReqInput, ) from sglang.multimodal_gen.runtime.entrypoints.utils import ( + AutoResidencyReq, GetDisaggStatsReq, ListLorasReq, MergeLoraWeightsReq, @@ -140,6 +141,7 @@ def __init__( self._running = True self.request_handlers = { + AutoResidencyReq: self._handle_auto_residency, SetLoraReq: self._handle_set_lora, MergeLoraWeightsReq: self._handle_merge_lora, UnmergeLoraWeightsReq: self._handle_unmerge_lora, @@ -201,6 +203,14 @@ def _handle_get_disagg_stats(self, _reqs: List[Any]) -> OutputBatch: output=stats or {"role": "monolithic", "message": "not in disagg mode"} ) + def _handle_auto_residency(self, reqs: List[Any]) -> OutputBatch: + req = reqs[0] + if not isinstance(self.worker, GPUWorker): + return OutputBatch(error="auto residency requires a GPU worker") + if req.action == "rollback": + return self.worker.rollback_auto_residency() + return self.worker.apply_auto_residency() + def _handle_set_lora(self, reqs: List[Any]) -> OutputBatch: # TODO: return set status # TODO: return with SetLoRAResponse or something more appropriate diff --git a/python/sglang/multimodal_gen/runtime/scheduler_client.py b/python/sglang/multimodal_gen/runtime/scheduler_client.py index aa0ef7c1e8fb..423c68d7bf25 100644 --- a/python/sglang/multimodal_gen/runtime/scheduler_client.py +++ b/python/sglang/multimodal_gen/runtime/scheduler_client.py @@ -16,6 +16,7 @@ UpdateWeightFromTensorReqInput, ) from sglang.multimodal_gen.runtime.entrypoints.utils import ( + AutoResidencyReq, ListLorasReq, MergeLoraWeightsReq, SetLoraReq, @@ -39,6 +40,7 @@ # Control ops mutate replica state (weights, LoRA, memory, shutdown), so with # DP they must reach every replica rather than one. _CONTROL_REQ_TYPES = ( + AutoResidencyReq, SetLoraReq, MergeLoraWeightsReq, UnmergeLoraWeightsReq, 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 73d1b7bf7cb2..cb0415426f72 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -1543,6 +1543,10 @@ def require_component_resident( ) self._required_resident_components.add(component_name) + def release_required_component_residency(self, component_name: str) -> None: + """Undo ``require_component_resident`` (auto-residency rollback).""" + self._required_resident_components.discard(component_name) + def should_use_fsdp_for_component(self, component_name: str) -> bool: return bool( self.use_fsdp_inference diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index db5102b387ca..438b5ee0b0ae 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -101,6 +101,120 @@ def should_run_explicit_client_warmup(server_args: ServerArgs) -> bool: ) +def auto_residency_skip_reason(server_args: ServerArgs) -> str | None: + """Server-args-only gate for the warmup-calibrated residency promotion. + + Only rules out paths the promotion was not designed for; the workers + re-check per component (explicit placement, FSDP modules, custom + strategies, missing sizes) and per measurement. + """ + from sglang.multimodal_gen import envs + from sglang.multimodal_gen.configs.quantization.nunchaku import ( + NunchakuSVDQuantArgs, + ) + from sglang.multimodal_gen.runtime.platforms import current_platform + + if envs.SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY: + return "disabled via SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY" + if server_args.performance_mode != "auto": + return f"performance_mode={server_args.performance_mode}" + if not should_run_synthetic_server_warmup(server_args): + return "no synthetic server warmup to calibrate from" + if server_args.backend == "diffusers": + return "diffusers backend" + if server_args.enable_breakable_cuda_graph: + return "breakable CUDA graph captures during warmup" + if envs.SGLANG_CACHE_DIT_ENABLED: + return "cache-dit enabled" + if server_args.batching_max_size > 1: + return "dynamic batching enabled" + if (server_args.dp_size or 1) > 1: + return "dp replicas warm up independently" + if server_args.use_fsdp_inference: + return "FSDP inference" + nunchaku = server_args.nunchaku_config + svdquant_enabled = nunchaku is not None and ( + not isinstance(nunchaku, NunchakuSVDQuantArgs) or nunchaku.enable_svdquant + ) + if ( + server_args.quantization is not None + or server_args.transformer_weights_path + or svdquant_enabled + ): + # Residency changes shift the memory layout, which measurably moved + # fp8-quantized DiT outputs in the past (see PR #29649); keep + # quantized checkpoints on their configured strategy. + return "quantized checkpoint" + if not current_platform.is_cuda(): + return "requires CUDA" + return None + + +def _auto_residency_status(response: OutputBatch) -> str | None: + if isinstance(response.output, dict): + return response.output.get("status") + return None + + +async def maybe_apply_auto_residency( + server_args: ServerArgs, + forward: Callable[[Req], Awaitable[OutputBatch]], +) -> None: + """Promote implicitly offloaded components after warmup, then re-warm. + + Runs between the synthetic warmup and the ready signal, so the residency + is frozen before ``/health`` turns 200. If a promotion (or the + post-promotion re-warmup) fails, the workers roll back and startup + continues on the original strategy; only a failed rollback raises, which + aborts startup. + """ + from sglang.multimodal_gen.runtime.entrypoints.utils import AutoResidencyReq + from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( + PROMOTION_STATUS_PROMOTED, + PROMOTION_STATUS_ROLLBACK_FAILED, + PROMOTION_STATUS_ROLLED_BACK, + ) + + skip_reason = auto_residency_skip_reason(server_args) + if skip_reason is not None: + logger.debug("Auto residency: skipped (%s)", skip_reason) + return + + response = await forward(AutoResidencyReq(action="apply")) + status = _auto_residency_status(response) + if status == PROMOTION_STATUS_ROLLBACK_FAILED: + raise RuntimeError(f"auto residency rollback failed: {response.error}") + if response.error is not None: + logger.warning( + "Auto residency promotion not applied; continuing on the original " + "strategy: %s", + response.error, + ) + return + if status != PROMOTION_STATUS_PROMOTED: + return + + # Re-run the synthetic warmup under the final residency: it physically + # moves promoted components, rebuilds compile caches invalidated by the + # placement change, and proves the promoted layout fits before ready. + try: + await run_async_client_warmup(server_args, forward, fail_open=False) + except Exception as e: + logger.warning( + "Post-promotion re-warmup failed (%s); rolling back auto residency", e + ) + rollback = await forward(AutoResidencyReq(action="rollback")) + if ( + rollback.error is not None + or _auto_residency_status(rollback) != PROMOTION_STATUS_ROLLED_BACK + ): + raise RuntimeError( + f"auto residency rollback failed: {rollback.error}" + ) from e + # restore warm caches for the original strategy before turning ready + await run_async_client_warmup(server_args, forward, fail_open=True) + + def format_warmup_req(req_or_group: Any) -> str: req = get_first_generation_req(req_or_group) prefix = ( diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py new file mode 100644 index 000000000000..14c1a8e2c2c4 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -0,0 +1,508 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for warmup-calibrated auto residency promotion.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType +from sglang.multimodal_gen.configs.pipeline_configs.longlive2 import LongLive2T2VConfig +from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( + ACTIVATION_EXTRAPOLATION_MARGIN, + GIB_BYTES, + MIN_VRAM_RESERVE_BYTES, + AutoResidencyPlan, + PromotionCandidate, + RankResidencyReport, + WarmupMemoryRecord, + apply_promotions, + collect_promotion_candidates, + component_resident_size_bytes, + estimate_default_workload_peak_bytes, + plan_auto_residency, + rollback_promotions, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( + COMPONENT_OFFLOAD, + LAYERWISE_OFFLOAD, + RESIDENT, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + LayerwiseOffloadableModuleMixin, +) +from sglang.multimodal_gen.runtime.warmup_request_builder import ( + SERVER_WARMUP_MAX_VIDEO_FRAMES, + _resolve_warmup_num_frames, +) + + +def _record( + *, + width=832, + height=480, + num_frames=17, + baseline_gib=10, + peak_gib=12, + succeeded=True, +) -> WarmupMemoryRecord: + return WarmupMemoryRecord( + width=width, + height=height, + num_frames=num_frames, + baseline_allocated_bytes=baseline_gib * GIB_BYTES, + peak_reserved_bytes=peak_gib * GIB_BYTES, + succeeded=succeeded, + ) + + +class TestEstimateDefaultWorkloadPeak: + def test_same_shape_uses_measured_peak(self): + record = _record() + estimate = estimate_default_workload_peak_bytes( + records=[record], target_units=record.workload_units() + ) + assert estimate == record.peak_reserved_bytes + + def test_unknown_target_uses_measured_peak(self): + record = _record() + estimate = estimate_default_workload_peak_bytes( + records=[record], target_units=None + ) + assert estimate == record.peak_reserved_bytes + + def test_capped_video_scales_only_the_activation_part(self): + # Warmup capped to 832x480x17; Wan-class default is 704x1280x121. + record = _record() + target_units = 704 * 1280 * 121 + ratio = target_units / record.workload_units() + assert ratio > 10 # the cap ratio this formula exists for + + estimate = estimate_default_workload_peak_bytes( + records=[record], target_units=target_units + ) + activation = record.peak_reserved_bytes - record.baseline_allocated_bytes + expected = record.baseline_allocated_bytes + int( + activation * ratio * ACTIVATION_EXTRAPOLATION_MARGIN + ) + assert estimate == expected + # Scaling the whole peak would inflate the estimate by the resident + # weights times the cap ratio and promotion would never trigger. + naive = int(record.peak_reserved_bytes * ratio) + assert estimate < naive + + def test_multiple_records_take_the_max(self): + low = _record(peak_gib=12) + high = _record(peak_gib=20) + estimate = estimate_default_workload_peak_bytes( + records=[low, high], target_units=low.workload_units() + ) + assert estimate == high.peak_reserved_bytes + + def test_failed_warmup_disables_estimation(self): + records = [_record(), _record(succeeded=False)] + assert ( + estimate_default_workload_peak_bytes(records=records, target_units=None) + is None + ) + + def test_no_records_disables_estimation(self): + assert ( + estimate_default_workload_peak_bytes(records=[], target_units=None) is None + ) + + +def _candidate( + name: str, + *, + mode: str = COMPONENT_OFFLOAD, + weight_gib: int = 10, + h2d_gib: int | None = None, +) -> PromotionCandidate: + return PromotionCandidate( + component_name=name, + residency_mode=mode, + promoted_weight_bytes=weight_gib * GIB_BYTES, + h2d_bytes_per_request=(h2d_gib if h2d_gib is not None else weight_gib) + * GIB_BYTES, + ) + + +def _report( + *, + rank: int = 0, + budget_gib: int = 100, + estimated_gib: int | None = 50, + candidates: list[PromotionCandidate] | None = None, + skip_reason: str | None = None, +) -> RankResidencyReport: + return RankResidencyReport( + rank=rank, + budget_bytes=budget_gib * GIB_BYTES, + estimated_peak_bytes=( + None if estimated_gib is None else estimated_gib * GIB_BYTES + ), + candidates=candidates if candidates is not None else [_candidate("vae")], + skip_reason=skip_reason, + ) + + +class TestPlanAutoResidency: + def test_rank_skip_reason_propagates(self): + plan = plan_auto_residency( + reports=[_report(), _report(rank=1, skip_reason="no measurements")] + ) + assert plan.skip_reason == "rank 1: no measurements" + assert not plan.promotions + + def test_missing_estimate_skips(self): + plan = plan_auto_residency(reports=[_report(estimated_gib=None)]) + assert plan.skip_reason is not None + + def test_worst_case_aggregation_across_ranks(self): + plan = plan_auto_residency( + reports=[ + _report(rank=0, budget_gib=100, estimated_gib=40), + _report(rank=1, budget_gib=80, estimated_gib=60), + ] + ) + assert plan.budget_bytes == 80 * GIB_BYTES + assert plan.estimated_peak_bytes == 60 * GIB_BYTES + + def test_reserve_has_absolute_floor(self): + # 10% of a 20 GiB budget is 2 GiB; the floor must lift it to 4 GiB. + plan = plan_auto_residency(reports=[_report(budget_gib=20, estimated_gib=10)]) + assert plan.reserve_bytes == MIN_VRAM_RESERVE_BYTES + + def test_greedy_promotion_by_h2d_savings(self): + # dit saves 25 GiB x 50 steps per request; text_encoder saves 30 GiB + # once. dit must be promoted first, after which the text encoder no + # longer fits (50 + 25 + 30 + 10 > 100). + candidates = [ + _candidate("text_encoder", weight_gib=30), + _candidate("dit", mode=LAYERWISE_OFFLOAD, weight_gib=25, h2d_gib=25 * 50), + ] + plan = plan_auto_residency( + reports=[_report(budget_gib=100, estimated_gib=50, candidates=candidates)] + ) + assert [c.component_name for c in plan.promotions] == ["dit"] + + def test_smaller_component_still_fits_after_skipping_a_big_one(self): + candidates = [ + _candidate("text_encoder", weight_gib=45, h2d_gib=100), + _candidate("vae", weight_gib=5, h2d_gib=5), + ] + plan = plan_auto_residency( + reports=[_report(budget_gib=100, estimated_gib=50, candidates=candidates)] + ) + assert [c.component_name for c in plan.promotions] == ["vae"] + + def test_consensus_requires_component_on_every_rank(self): + plan = plan_auto_residency( + reports=[ + _report( + rank=0, + candidates=[_candidate("vae"), _candidate("text_encoder")], + ), + _report(rank=1, candidates=[_candidate("vae", weight_gib=12)]), + ] + ) + names = [c.component_name for c in plan.promotions] + assert names == ["vae"] + # worst-case size across ranks + assert plan.promotions[0].promoted_weight_bytes == 12 * GIB_BYTES + + def test_no_candidates_skips(self): + plan = plan_auto_residency(reports=[_report(candidates=[])]) + assert plan.skip_reason is not None + + +class _FakeLayerwiseManager: + def __init__(self, tensors: dict[str, torch.Tensor]): + self.enabled = True + self._configured = True + self._tensors = tensors + self.load_all_layers_calls = 0 + self.remove_hooks_calls = 0 + self.register_hooks_calls = 0 + self.sync_to_cpu_calls = 0 + self.release_all_calls = 0 + + def iter_cpu_weights(self): + yield from self._tensors.items() + + def load_all_layers(self): + self.load_all_layers_calls += 1 + + def remove_forward_hooks(self): + self.remove_hooks_calls += 1 + + def register_forward_hooks(self): + self.register_hooks_calls += 1 + + def sync_all_layers_to_cpu(self): + self.sync_to_cpu_calls += 1 + + def release_all(self): + self.release_all_calls += 1 + + +class _FakeLayerwiseDit(LayerwiseOffloadableModuleMixin, nn.Module): + def __init__(self, managers: list[_FakeLayerwiseManager]): + nn.Module.__init__(self) + self.layerwise_offload_managers = managers + + +class _StubResidencyArgs: + """Duck-typed stand-in for the two ServerArgs hooks promotions use.""" + + def __init__(self): + self.required: set[str] = set() + + def require_component_resident(self, component_name, *, feature_name): + self.required.add(component_name) + + def release_required_component_residency(self, component_name): + self.required.discard(component_name) + + +def _plan_for(candidates: list[PromotionCandidate]) -> AutoResidencyPlan: + return AutoResidencyPlan( + estimated_peak_bytes=GIB_BYTES, + reserve_bytes=MIN_VRAM_RESERVE_BYTES, + budget_bytes=100 * GIB_BYTES, + promotions=candidates, + ) + + +class TestSizeAccounting: + def test_component_offload_counts_params_and_buffers(self): + module = nn.Linear(8, 8) + module.register_buffer("stats", torch.zeros(4)) + expected = sum(t.numel() * t.element_size() for t in module.parameters()) + sum( + t.numel() * t.element_size() for t in module.buffers() + ) + assert component_resident_size_bytes(module, COMPONENT_OFFLOAD) == expected + + def test_layerwise_counts_manager_cpu_buffers_not_placeholders(self): + cpu_weights = {"layers.0.w": torch.zeros(1024), "layers.1.w": torch.zeros(1024)} + module = _FakeLayerwiseDit([_FakeLayerwiseManager(cpu_weights)]) + expected = sum(t.numel() * t.element_size() for t in cpu_weights.values()) + assert component_resident_size_bytes(module, LAYERWISE_OFFLOAD) == expected + + +class TestCollectPromotionCandidates: + def _modes(self, mapping): + return lambda name: mapping.get(name, RESIDENT) + + def test_filters_and_sizes(self): + te = nn.Linear(8, 8) + vae = nn.Linear(4, 4) + dit = _FakeLayerwiseDit( + [_FakeLayerwiseManager({"layers.0.w": torch.zeros(1024)})] + ) + modules = { + "text_encoder": te, + "vae": vae, + "transformer": dit, + "explicit_encoder": nn.Linear(2, 2), + "custom": nn.Linear(2, 2), + "scheduler": object(), + } + modes = self._modes( + { + "text_encoder": COMPONENT_OFFLOAD, + "transformer": LAYERWISE_OFFLOAD, + "explicit_encoder": COMPONENT_OFFLOAD, + "custom": COMPONENT_OFFLOAD, + } + ) + candidates = collect_promotion_candidates( + modules=modules, + residency_mode_of=modes, + explicit_residency_mode_of=lambda name: ( + COMPONENT_OFFLOAD if name == "explicit_encoder" else None + ), + custom_strategy_names={"custom"}, + num_inference_steps=50, + ) + by_name = {c.component_name: c for c in candidates} + # resident vae, explicit placement, custom strategy, non-module: out + assert set(by_name) == {"text_encoder", "transformer"} + assert by_name["text_encoder"].h2d_bytes_per_request == ( + by_name["text_encoder"].promoted_weight_bytes + ) + # layerwise DiT re-streams its layers once per denoise step + assert by_name["transformer"].h2d_bytes_per_request == ( + by_name["transformer"].promoted_weight_bytes * 50 + ) + + +class TestApplyAndRollback: + def test_component_offload_promotion_marks_resident_without_moving(self): + module = nn.Linear(4, 4) + args = _StubResidencyArgs() + applied = apply_promotions( + plan=_plan_for([_candidate("text_encoder", weight_gib=1)]), + modules={"text_encoder": module}, + server_args=args, + ) + assert args.required == {"text_encoder"} + assert [p.component_name for p in applied] == ["text_encoder"] + + rollback_promotions( + applied=applied, modules={"text_encoder": module}, server_args=args + ) + assert args.required == set() + + def test_layerwise_promotion_loads_all_layers_and_rollback_rearms(self): + manager = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) + module = _FakeLayerwiseDit([manager]) + args = _StubResidencyArgs() + candidate = _candidate("transformer", mode=LAYERWISE_OFFLOAD, weight_gib=1) + + applied = apply_promotions( + plan=_plan_for([candidate]), + modules={"transformer": module}, + server_args=args, + ) + assert manager.load_all_layers_calls == 1 + assert manager.remove_hooks_calls == 1 + assert manager.enabled is False + + rollback_promotions( + applied=applied, modules={"transformer": module}, server_args=args + ) + assert manager.enabled is True + assert manager.register_hooks_calls == 1 + assert args.required == set() + + def test_mid_failure_rolls_back_already_applied_promotions(self): + manager = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) + module = _FakeLayerwiseDit([manager]) + args = _StubResidencyArgs() + candidates = [ + _candidate("transformer", mode=LAYERWISE_OFFLOAD, weight_gib=1), + _candidate("missing_component", weight_gib=1), + ] + with pytest.raises(RuntimeError, match="missing_component"): + apply_promotions( + plan=_plan_for(candidates), + modules={"transformer": module}, + server_args=args, + ) + assert args.required == set() + assert manager.enabled is True + assert manager.register_hooks_calls == 1 + + def test_enable_offload_does_not_double_register_hooks(self): + manager = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) + module = _FakeLayerwiseDit([manager]) + # never disabled: re-enabling must not stack a second set of hooks + module.enable_offload() + assert manager.register_hooks_calls == 0 + + +class TestWarmupFrameAdjustment: + def _server_args(self, *, bcg: bool = False) -> SimpleNamespace: + return SimpleNamespace( + pipeline_config=LongLive2T2VConfig(), + enable_breakable_cuda_graph=bcg, + ) + + def _defaults(self) -> SimpleNamespace: + return SimpleNamespace(num_frames=61) + + def test_capped_frames_keep_the_model_frame_contract(self): + # LongLive2 default 61 frames is capped to 17, whose 5 latent frames + # break the 8-frame causal block; the builder must re-align to 29. + assert SERVER_WARMUP_MAX_VIDEO_FRAMES == 17 + num_frames = _resolve_warmup_num_frames( + self._server_args(), self._defaults(), server_based_warmup=True + ) + assert num_frames == 29 + + def test_bcg_keeps_full_serving_frames(self): + num_frames = _resolve_warmup_num_frames( + self._server_args(bcg=True), self._defaults(), server_based_warmup=True + ) + assert num_frames == 61 + + def test_non_server_warmup_keeps_default_frames(self): + num_frames = _resolve_warmup_num_frames( + self._server_args(), self._defaults(), server_based_warmup=False + ) + assert num_frames == 61 + + +class TestAutoResidencySkipReason: + def _base_args(self, **overrides) -> SimpleNamespace: + args = SimpleNamespace( + performance_mode="auto", + warmup_mode="server", + warmup_resolutions=None, + backend="sglang", + enable_breakable_cuda_graph=False, + batching_max_size=1, + dp_size=1, + use_fsdp_inference=False, + quantization=None, + transformer_weights_path=None, + nunchaku_config=None, + pipeline_config=SimpleNamespace(task_type=ModelTaskType.T2V), + ) + for key, value in overrides.items(): + setattr(args, key, value) + return args + + def _skip_reason(self, args): + from sglang.multimodal_gen.runtime.server_warmup import ( + auto_residency_skip_reason, + ) + + return auto_residency_skip_reason(args) + + def test_env_kill_switch(self, monkeypatch): + monkeypatch.setenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", "1") + reason = self._skip_reason(self._base_args()) + assert ( + reason is not None and "SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY" in reason + ) + + def test_manual_performance_mode(self, monkeypatch): + monkeypatch.delenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", raising=False) + reason = self._skip_reason(self._base_args(performance_mode="manual")) + assert reason == "performance_mode=manual" + + @pytest.mark.parametrize( + "overrides, expected_fragment", + [ + ({"warmup_mode": "request"}, "server warmup"), + ({"backend": "diffusers"}, "diffusers"), + ({"enable_breakable_cuda_graph": True}, "CUDA graph"), + ({"batching_max_size": 4}, "batching"), + ({"dp_size": 2}, "dp replicas"), + ({"use_fsdp_inference": True}, "FSDP"), + ({"quantization": "fp8"}, "quantized"), + ({"transformer_weights_path": "/x.safetensors"}, "quantized"), + ], + ) + def test_excluded_paths(self, monkeypatch, overrides, expected_fragment): + monkeypatch.delenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", raising=False) + monkeypatch.delenv("SGLANG_CACHE_DIT_ENABLED", raising=False) + reason = self._skip_reason(self._base_args(**overrides)) + assert reason is not None and expected_fragment in reason + + def test_cache_dit_excluded(self, monkeypatch): + monkeypatch.delenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", raising=False) + monkeypatch.setenv("SGLANG_CACHE_DIT_ENABLED", "true") + reason = self._skip_reason(self._base_args()) + assert reason is not None and "cache-dit" in reason + + def test_eligible_path_reaches_platform_gate(self, monkeypatch): + monkeypatch.delenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", raising=False) + monkeypatch.delenv("SGLANG_CACHE_DIT_ENABLED", raising=False) + reason = self._skip_reason(self._base_args()) + # on a CUDA host everything passes; CPU CI stops at the platform gate + assert reason is None or reason == "requires CUDA" From 0c0f117b0183433658b103b375b46d2ea276a627 Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 16 Aug 2026 22:36:18 +0800 Subject: [PATCH 02/90] [diffusion] fix auto residency video estimation with two-point calibration Single-point extrapolation cannot separate constant peak costs (streamed layer weights, attention workspace, tiled VAE decode) from workload-linear activations when offload leaves the pre-forward baseline nearly empty: measured on Wan2.1-14B/H100, a ~30 GiB default-workload peak was estimated as 182.9 GiB and promotion never fired. Frame-capped video server warmup now adds one smaller calibration size (9 frames, auto mode only) so the estimator can fit peak = constant + slope * units and extrapolate only the measured linear part; a measurement at or above the target bounds the peak directly, and a single usable size keeps the conservative baseline split. Co-Authored-By: Claude Fable 5 --- .../runtime/managers/gpu_worker.py | 6 +- .../memory_managers/auto_residency.py | 62 ++++++++++- .../runtime/warmup_request_builder.py | 62 +++++++++++ .../test/unit/test_auto_residency.py | 104 +++++++++++++++++- 4 files changed, 227 insertions(+), 7 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 796d01c4219b..8c072c493ff8 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1122,7 +1122,11 @@ def apply_auto_residency(self) -> OutputBatch: local_report = self._build_auto_residency_report(workload=workload) reports = self._auto_residency_all_gather(local_report) plan = plan_auto_residency(reports=reports) - summary = format_plan_summary(plan=plan, workload=workload) + summary = format_plan_summary( + plan=plan, + workload=workload, + records=self._auto_residency_warmup_records, + ) if plan.skip_reason is not None or not plan.promotions: if self.is_output_rank: logger.info("%s", summary) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 6e01fb6d19e0..8621e2ef3713 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -161,6 +161,18 @@ def estimate_default_workload_peak_bytes( ) -> int | None: """Extrapolate warmup peaks to the default workload. + Preference order: + 1. A measurement at or above the target workload bounds the peak directly. + 2. Two distinct measured workload sizes fit ``peak = constant + slope * + units``; only the fitted linear part is extrapolated (and padded). + Under offload the pre-forward baseline is nearly empty, so a + single-point split cannot separate constant costs (streamed weights, + attention workspace, tiled VAE decode) from workload-linear + activations -- measured on Wan2.1-14B, single-point scaling + overestimated a ~30 GiB peak as ~183 GiB. + 3. One usable size: scale everything above the pre-forward allocated + baseline (conservative; may block promotion but never over-promotes). + Returns None when there is no usable measurement (no records, or any warmup forward failed -- a failed warmup means the measured peak does not cover the full request path). @@ -169,15 +181,42 @@ def estimate_default_workload_peak_bytes( if not records or any(not record.succeeded for record in records): return None + peak_by_units: dict[int, int] = {} + for record in records: + units = record.workload_units() + peak_by_units[units] = max( + peak_by_units.get(units, 0), record.peak_reserved_bytes + ) + + if target_units is None: + return max(peak_by_units.values()) + + covering_peaks = [ + peak for units, peak in peak_by_units.items() if units >= target_units + ] + if covering_peaks: + return max(covering_peaks) + + if len(peak_by_units) >= 2: + # fit on the two largest sizes: closest to the target, best local slope + (large_units, large_peak), (small_units, small_peak) = sorted( + peak_by_units.items(), reverse=True + )[:2] + slope = (large_peak - small_peak) / (large_units - small_units) + if slope >= 0: + constant = large_peak - slope * large_units + return int( + constant + slope * target_units * ACTIVATION_EXTRAPOLATION_MARGIN + ) + # negative slope is measurement noise; fall through to the + # conservative single-point estimate + estimates = [] for record in records: peak = record.peak_reserved_bytes - ratio = 1.0 if target_units is None else target_units / record.workload_units() - if ratio <= 1.0: - estimates.append(peak) - continue baseline = min(record.baseline_allocated_bytes, peak) activation = peak - baseline + ratio = target_units / record.workload_units() estimates.append( baseline + int(activation * ratio * ACTIVATION_EXTRAPOLATION_MARGIN) ) @@ -420,7 +459,12 @@ def rollback_promotions( torch.get_device_module().empty_cache() -def format_plan_summary(*, plan: AutoResidencyPlan, workload: DefaultWorkload) -> str: +def format_plan_summary( + *, + plan: AutoResidencyPlan, + workload: DefaultWorkload, + records: Iterable[WarmupMemoryRecord] = (), +) -> str: """One-line decision summary for the startup log.""" if plan.skip_reason is not None: return f"Auto residency: skipped ({plan.skip_reason})" @@ -433,9 +477,17 @@ def format_plan_summary(*, plan: AutoResidencyPlan, workload: DefaultWorkload) - ) or "none" ) + measured = ", ".join( + f"{record.width}x{record.height}x{record.num_frames}f=" + f"{record.peak_reserved_bytes / GIB_BYTES:.1f}GiB" + for record in records + if record.succeeded + ) + measured_part = f"measured=[{measured}], " if measured else "" return ( f"Auto residency: target={workload.describe()} " f"steps={workload.num_inference_steps}, " + f"{measured_part}" f"estimated_peak={plan.estimated_peak_bytes / GIB_BYTES:.1f} GiB, " f"reserve={plan.reserve_bytes / GIB_BYTES:.1f} GiB, " f"budget={plan.budget_bytes / GIB_BYTES:.1f} GiB, " diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py index b5f1be1d16bc..dfe0c8b35baf 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -40,6 +40,10 @@ SERVER_WARMUP_VIDEO_MAX_AREA = 832 * 480 SERVER_WARMUP_MAX_VIDEO_FRAMES = 17 SERVER_WARMUP_LTX2_TWO_STAGE_MAX_VIDEO_FRAMES = 25 +# Second, smaller frame count used to calibrate auto residency: two workload +# sizes let the estimator fit peak = constant + slope * units instead of +# scaling one measured peak (whose constant part dominates under offload). +SERVER_WARMUP_CALIBRATION_VIDEO_FRAMES = 9 SERVER_WARMUP_IMAGE_STEPS = 2 SERVER_WARMUP_VIDEO_STEPS = 2 @@ -259,6 +263,37 @@ def _resolve_warmup_num_frames( return server_args.pipeline_config.adjust_num_frames(warmup_num_frames) +def _resolve_calibration_num_frames( + server_args: ServerArgs, + sampling_defaults: SamplingParams, + warmup_num_frames: int, + *, + server_based_warmup: bool, +) -> int | None: + """A second, smaller warmup frame count for auto-residency calibration. + + Only produced when the main video warmup was frame-capped under + ``--performance-mode auto``: the auto-residency estimator then has two + workload sizes to fit the constant/linear split of the peak. Returns None + when a second measurement adds no information (same adjusted frames). + """ + if not server_based_warmup or not _is_video_warmup_task(server_args): + return None + if server_args.performance_mode != "auto": + return None + if getattr(server_args, "enable_breakable_cuda_graph", False) is True: + return None + default_num_frames = getattr(sampling_defaults, "num_frames", 1) or 1 + if warmup_num_frames >= default_num_frames: + return None + calibration_num_frames = server_args.pipeline_config.adjust_num_frames( + SERVER_WARMUP_CALIBRATION_VIDEO_FRAMES + ) + if calibration_num_frames == warmup_num_frames or calibration_num_frames <= 0: + return None + return calibration_num_frames + + def _effective_cfg_scale(sampling_defaults: SamplingParams) -> float | None: if getattr(sampling_defaults, "true_cfg_scale", None) is not None: return sampling_defaults.true_cfg_scale @@ -369,6 +404,16 @@ def build_warmup_reqs( sampling_defaults, server_based_warmup=server_based_warmup, ) + calibration_num_frames = ( + _resolve_calibration_num_frames( + server_args, + sampling_defaults, + warmup_num_frames, + server_based_warmup=server_based_warmup, + ) + if warmup_resolutions is None + else None + ) # build warmup reqs warmup_reqs = [] @@ -403,6 +448,23 @@ def build_warmup_reqs( elif negative_prompt is not None and cfg_scale is not None and cfg_scale > 1.0: req_kwargs["do_classifier_free_guidance"] = True + if calibration_num_frames is not None and (width, height) == resolutions[0]: + # smaller-workload measurement first; it also pre-warms kernels + # so the main warmup measurement is cleaner + calibration_kwargs = req_kwargs.copy() + calibration_kwargs["sampling_params"] = copy(req_kwargs["sampling_params"]) + calibration_kwargs["num_frames"] = calibration_num_frames + calibration_req = Req(**calibration_kwargs) + calibration_req.set_as_warmup(warmup_steps) + calibration_req.sampling_params.prepare_synthetic_warmup_request_for_queue( + calibration_req, server_args + ) + if return_warmup_result: + calibration_req.extra["return_warmup_result"] = True + if server_based_warmup: + calibration_req.extra["server_based_warmup"] = True + warmup_reqs.append(calibration_req) + run_real_path_prewarm = server_based_warmup and server_args.enable_torch_compile prompts = ( (DEFAULT_PLACEHOLDER_PROMPT,) + TORCH_COMPILE_REAL_PATH_PREWARM_PROMPTS diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 14c1a8e2c2c4..3728feca401f 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -34,6 +34,7 @@ ) from sglang.multimodal_gen.runtime.warmup_request_builder import ( SERVER_WARMUP_MAX_VIDEO_FRAMES, + _resolve_calibration_num_frames, _resolve_warmup_num_frames, ) @@ -72,7 +73,7 @@ def test_unknown_target_uses_measured_peak(self): ) assert estimate == record.peak_reserved_bytes - def test_capped_video_scales_only_the_activation_part(self): + def test_single_capped_record_scales_only_the_activation_part(self): # Warmup capped to 832x480x17; Wan-class default is 704x1280x121. record = _record() target_units = 704 * 1280 * 121 @@ -92,6 +93,55 @@ def test_capped_video_scales_only_the_activation_part(self): naive = int(record.peak_reserved_bytes * ratio) assert estimate < naive + def test_two_point_fit_separates_constant_from_linear(self): + # Two calibration sizes let the estimator measure the slope instead + # of assuming everything above the baseline scales. Under offload the + # baseline is nearly empty, so the single-point formula would scale + # the whole peak (~x13 here); the fit extrapolates only the measured + # linear part. + small = _record(num_frames=9, baseline_gib=1, peak_gib=14) + large = _record(num_frames=17, baseline_gib=1, peak_gib=16) + target_units = 1280 * 720 * 81 + + estimate = estimate_default_workload_peak_bytes( + records=[small, large], target_units=target_units + ) + slope = (large.peak_reserved_bytes - small.peak_reserved_bytes) / ( + large.workload_units() - small.workload_units() + ) + constant = large.peak_reserved_bytes - slope * large.workload_units() + expected = int( + constant + slope * target_units * ACTIVATION_EXTRAPOLATION_MARGIN + ) + assert estimate == expected + single_point = estimate_default_workload_peak_bytes( + records=[large], target_units=target_units + ) + assert estimate < single_point + + def test_negative_slope_falls_back_to_single_point_formula(self): + small = _record(num_frames=9, peak_gib=16) + large = _record(num_frames=17, peak_gib=14) + target_units = 1280 * 720 * 81 + estimate = estimate_default_workload_peak_bytes( + records=[small, large], target_units=target_units + ) + fallback = max( + estimate_default_workload_peak_bytes( + records=[record], target_units=target_units + ) + for record in (small, large) + ) + assert estimate == fallback + + def test_covering_measurement_bounds_the_target(self): + capped = _record(num_frames=17, peak_gib=12) + full = _record(width=1280, height=720, num_frames=81, peak_gib=30) + estimate = estimate_default_workload_peak_bytes( + records=[capped, full], target_units=1280 * 720 * 81 + ) + assert estimate == full.peak_reserved_bytes + def test_multiple_records_take_the_max(self): low = _record(peak_gib=12) high = _record(peak_gib=20) @@ -436,6 +486,58 @@ def test_non_server_warmup_keeps_default_frames(self): assert num_frames == 61 +class TestCalibrationFrames: + def _wan_like_args(self, *, performance_mode: str = "auto") -> SimpleNamespace: + return SimpleNamespace( + pipeline_config=SimpleNamespace( + task_type=ModelTaskType.T2V, + adjust_num_frames=lambda n: n, + ), + enable_breakable_cuda_graph=False, + performance_mode=performance_mode, + ) + + def test_capped_video_gets_a_smaller_calibration_size(self): + calibration = _resolve_calibration_num_frames( + self._wan_like_args(), + SimpleNamespace(num_frames=81), + 17, + server_based_warmup=True, + ) + assert calibration == 9 + + def test_uncapped_video_needs_no_calibration(self): + calibration = _resolve_calibration_num_frames( + self._wan_like_args(), + SimpleNamespace(num_frames=13), + 13, + server_based_warmup=True, + ) + assert calibration is None + + def test_manual_mode_skips_calibration(self): + calibration = _resolve_calibration_num_frames( + self._wan_like_args(performance_mode="manual"), + SimpleNamespace(num_frames=81), + 17, + server_based_warmup=True, + ) + assert calibration is None + + def test_frame_contract_collapse_skips_calibration(self): + # LongLive2 aligns both 17 and 9 to 29 latent-block frames: a second + # measurement at the same size adds nothing. + args = SimpleNamespace( + pipeline_config=LongLive2T2VConfig(), + enable_breakable_cuda_graph=False, + performance_mode="auto", + ) + calibration = _resolve_calibration_num_frames( + args, SimpleNamespace(num_frames=61), 29, server_based_warmup=True + ) + assert calibration is None + + class TestAutoResidencySkipReason: def _base_args(self, **overrides) -> SimpleNamespace: args = SimpleNamespace( From 2df7feecc7a4c9c869c32111a824afd25b924bb9 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 17 Aug 2026 09:28:55 +0800 Subject: [PATCH 03/90] [diffusion] narrate auto residency adjustment with equivalent server args After warmup the log now states that residency adjustment is starting, what changed (mode -> resident per component), the equivalent --component-residency flags users can pin to freeze the placement, and the SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY kill switch. Rollbacks log that the startup-configured residency is in effect again. Co-Authored-By: Claude Fable 5 --- .../runtime/managers/gpu_worker.py | 8 +++++++ .../memory_managers/auto_residency.py | 21 +++++++++++++++++++ .../multimodal_gen/runtime/server_warmup.py | 4 ++++ .../test/unit/test_auto_residency.py | 16 ++++++++++++++ 4 files changed, 49 insertions(+) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 8c072c493ff8..498846f7c257 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -65,6 +65,7 @@ collect_promotion_candidates, component_resident_size_bytes, estimate_default_workload_peak_bytes, + format_applied_changes, format_plan_summary, plan_auto_residency, plan_summary_payload, @@ -1164,6 +1165,7 @@ def apply_auto_residency(self) -> OutputBatch: ) if self.is_output_rank: logger.info("%s", summary) + logger.info("%s", format_applied_changes(plan=plan)) return OutputBatch( output=plan_summary_payload( plan=plan, workload=workload, status=PROMOTION_STATUS_PROMOTED @@ -1181,6 +1183,12 @@ def rollback_auto_residency(self) -> OutputBatch: error=f"auto residency rollback failed: {rank_errors[0]}", output={"status": PROMOTION_STATUS_ROLLBACK_FAILED}, ) + if self.is_output_rank: + logger.info( + "Auto residency: rolled back; the residency configured at " + "startup is in effect again (no equivalent server-arg " + "changes remain)." + ) return OutputBatch(output={"status": PROMOTION_STATUS_ROLLED_BACK}) def _build_auto_residency_report( diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 8621e2ef3713..abd313e33ae8 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -495,6 +495,27 @@ def format_plan_summary( ) +def format_applied_changes(*, plan: AutoResidencyPlan) -> str: + """Describe the applied residency changes as equivalent server args. + + Users can pin the printed ``--component-residency`` flags to freeze this + placement, or disable the adjustment entirely with the kill switch. + """ + changes = "; ".join( + f"{candidate.component_name}: {candidate.residency_mode} -> resident" + for candidate in plan.promotions + ) + equivalent = "--component-residency " + " ".join( + f"{candidate.component_name}=resident" for candidate in plan.promotions + ) + return ( + f"Auto residency: adjusted {changes}. " + f"Equivalent server args: {equivalent}. " + f"Pin these flags to make this placement explicit, or set " + f"SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY=1 to disable auto adjustment." + ) + + def plan_summary_payload( *, plan: AutoResidencyPlan, workload: DefaultWorkload, status: str ) -> dict: diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index 438b5ee0b0ae..7651e6042f5c 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -180,6 +180,10 @@ async def maybe_apply_auto_residency( logger.debug("Auto residency: skipped (%s)", skip_reason) return + logger.info( + "Server warmup complete; adjusting component residency for the " + "default workload (--performance-mode auto)." + ) response = await forward(AutoResidencyReq(action="apply")) status = _auto_residency_status(response) if status == PROMOTION_STATUS_ROLLBACK_FAILED: diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 3728feca401f..4d3b5692a850 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -21,6 +21,7 @@ collect_promotion_candidates, component_resident_size_bytes, estimate_default_workload_peak_bytes, + format_applied_changes, plan_auto_residency, rollback_promotions, ) @@ -454,6 +455,21 @@ def test_enable_offload_does_not_double_register_hooks(self): assert manager.register_hooks_calls == 0 +class TestAppliedChangesLog: + def test_reports_equivalent_server_args_and_kill_switch(self): + plan = _plan_for( + [ + _candidate("text_encoder", mode=LAYERWISE_OFFLOAD, weight_gib=7), + _candidate("vae", weight_gib=1), + ] + ) + message = format_applied_changes(plan=plan) + assert "text_encoder: layerwise-offload -> resident" in message + assert "vae: component-offload -> resident" in message + assert "--component-residency text_encoder=resident vae=resident" in message + assert "SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY=1" in message + + class TestWarmupFrameAdjustment: def _server_args(self, *, bcg: bool = False) -> SimpleNamespace: return SimpleNamespace( From 414b3b1db527f60ad9824423273223fcaa9783f8 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 18 Aug 2026 15:22:59 +0800 Subject: [PATCH 04/90] [diffusion] unify post-request residency hint with auto promotion logic get_can_stay_resident_components now reuses collect_promotion_candidates and the shared H2D-savings ranking, so the post-request hint names the same components in the same order auto mode would promote (still raw capacity: no reserve/margin, and explicitly offloaded components stay listed). Drops the static OFFLOAD_DISABLE_RECOMMENDATION_ORDER and points the hint at --component-residency instead of legacy flags. Co-Authored-By: Claude Fable 5 --- .../runtime/managers/gpu_worker.py | 70 ++++++++----------- .../memory_managers/auto_residency.py | 26 ++++--- .../test/unit/test_auto_residency.py | 26 +++++++ 3 files changed, 73 insertions(+), 49 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 498846f7c257..65db6e79f995 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -13,7 +13,6 @@ import numpy as np import torch -import torch.nn as nn from setproctitle import setproctitle from sglang.multimodal_gen.runtime.utils.logging_utils import ( # isort: skip @@ -63,22 +62,18 @@ WarmupMemoryRecord, apply_promotions, collect_promotion_candidates, - component_resident_size_bytes, estimate_default_workload_peak_bytes, format_applied_changes, format_plan_summary, plan_auto_residency, plan_summary_payload, + rank_candidates_by_h2d_savings, resolve_default_workload, rollback_promotions, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( peek_global_component_residency_manager, ) -from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( - COMPONENT_OFFLOAD, - LAYERWISE_OFFLOAD, -) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( configure_layerwise_offload_modules, ) @@ -122,14 +117,6 @@ logger = init_logger(__name__) -OFFLOAD_DISABLE_RECOMMENDATION_ORDER = ( - "vae", - "image_encoder", - "text_encoder", - "text_encoder_2", - "transformer", -) - @dataclass class _ExpandedOutputParts: @@ -397,8 +384,9 @@ def do_mem_analysis(self, output_batch: OutputBatch): logger.debug( "GPU memory: peak=%.2f GB, allocated=%.2f GB, pool=%.2f GB (%.1f%%), " "headroom=%.2f GB. Components that can remain on GPU: %s. " - "Adjust --cpu-offload-components or --layerwise-offload-components " - "to change residency.", + "Make it explicit with --component-residency =resident; " + "--performance-mode auto with server warmup applies safe " + "promotions automatically.", peak_reserved_gb, peak_allocated_gb, pool_overhead_gb, @@ -1080,35 +1068,35 @@ def _finalize_expanded_parts( def get_can_stay_resident_components( self, remaining_gpu_mem_gb: float ) -> List[str]: + """Which currently offloaded components would fit in the headroom. + + Reuses the auto-residency candidate collection and its H2D-savings + greedy order, so this post-request hint names the same components, + in the same order, that ``--performance-mode auto`` would promote. + Unlike the promotion plan it applies no reserve or margin: it reports + raw capacity after this request, not a placement decision. """ - Calculate which components can stay resident on GPU without being offloaded. - """ - can_stay_resident = [] if not self.pipeline: - return can_stay_resident - - # Size from the live modules, not pipeline.memory_usages: the loader - # records the GPU-load delta, which is ~0 for exactly the offloaded - # components this recommendation is about. - modules = self.pipeline.modules - ordered_names = [ - name for name in OFFLOAD_DISABLE_RECOMMENDATION_ORDER if name in modules - ] - ordered_names.extend( - name for name in modules if name not in OFFLOAD_DISABLE_RECOMMENDATION_ORDER + return [] + + workload = resolve_default_workload(self.server_args) + candidates = collect_promotion_candidates( + modules=self.pipeline.modules, + residency_mode_of=self.server_args.residency_mode, + # The hint also covers explicitly offloaded components: the user + # chose offload and should learn when the headroom no longer + # requires it (promotion itself never touches explicit ones). + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=self.pipeline.component_residency_strategies, + num_inference_steps=workload.num_inference_steps, ) - for name in ordered_names: - module = modules[name] - if not isinstance(module, nn.Module): - continue - residency_mode = self.server_args.residency_mode(name) - if residency_mode not in (COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD): - continue - usage_gb = component_resident_size_bytes(module, residency_mode) / GIB_BYTES - if 0 < usage_gb <= remaining_gpu_mem_gb: - can_stay_resident.append(name) - remaining_gpu_mem_gb -= usage_gb + can_stay_resident = [] + for candidate in rank_candidates_by_h2d_savings(candidates): + usage_gb = candidate.promoted_weight_bytes / GIB_BYTES + if usage_gb <= remaining_gpu_mem_gb: + can_stay_resident.append(candidate.component_name) + remaining_gpu_mem_gb -= usage_gb return can_stay_resident def apply_auto_residency(self) -> OutputBatch: diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index abd313e33ae8..7cdb291b98a2 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -310,6 +310,23 @@ def collect_promotion_candidates( return candidates +def rank_candidates_by_h2d_savings( + candidates: Iterable[PromotionCandidate], +) -> list[PromotionCandidate]: + """Biggest per-request H2D savings first; name breaks ties deterministically. + + Shared by the promotion plan and the post-request residency hint so the + hint always lists components in the order auto mode would promote them. + """ + return sorted( + candidates, + key=lambda candidate: ( + -candidate.h2d_bytes_per_request, + candidate.component_name, + ), + ) + + def _skip_plan(reason: str) -> AutoResidencyPlan: return AutoResidencyPlan(skip_reason=reason) @@ -369,14 +386,7 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP if not candidates: return _skip_plan("no implicitly offloaded components to promote") - # Biggest per-request H2D savings first; name breaks ties deterministically. - ordered = sorted( - candidates, - key=lambda candidate: ( - -candidate.h2d_bytes_per_request, - candidate.component_name, - ), - ) + ordered = rank_candidates_by_h2d_savings(candidates) promotions = [] promoted_bytes = 0 for candidate in ordered: diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 4d3b5692a850..eddd1a048973 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -23,6 +23,7 @@ estimate_default_workload_peak_bytes, format_applied_changes, plan_auto_residency, + rank_candidates_by_h2d_savings, rollback_promotions, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( @@ -455,6 +456,31 @@ def test_enable_offload_does_not_double_register_hooks(self): assert manager.register_hooks_calls == 0 +class TestCandidateRanking: + def test_post_request_hint_order_matches_promotion_order(self): + # The recommendation log and the promotion plan share one ranking, so + # the hint lists components in the order auto mode would promote them. + candidates = [ + _candidate("vae", weight_gib=1), + _candidate( + "transformer", mode=LAYERWISE_OFFLOAD, weight_gib=10, h2d_gib=500 + ), + _candidate("text_encoder", weight_gib=7), + ] + ranked = rank_candidates_by_h2d_savings(candidates) + assert [candidate.component_name for candidate in ranked] == [ + "transformer", + "text_encoder", + "vae", + ] + plan = plan_auto_residency( + reports=[_report(budget_gib=1000, estimated_gib=50, candidates=candidates)] + ) + assert [candidate.component_name for candidate in plan.promotions] == [ + candidate.component_name for candidate in ranked + ] + + class TestAppliedChangesLog: def test_reports_equivalent_server_args_and_kill_switch(self): plan = _plan_for( From f856eb682fde35a7f61a7710a4bdc1d11006d37d Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 18 Aug 2026 23:04:15 +0800 Subject: [PATCH 05/90] [diffusion] harden auto residency against review findings Measurement correctness: - empty_cache + reset before each measured warmup forward, so a later shape reusing the previous request's allocator pool cannot flatten the calibration slope to zero (under-estimate -> post-ready OOM) - run the calibration forward AFTER the main warmup request: one-time allocations land in the larger point and can only steepen the slope - only is_warmup reqs are measured; torch-compile prewarm reqs run a different offload layout and poisoned the fit in either direction - record failed warmup forwards even on the propagate re-raise path - skip promotion entirely under --enable-torch-compile: compile warmup runs a stripped memory layout (layerwise DiT, aux components evicted) whose peaks are far below a real request's - unknown default workload now skips instead of trusting the capped warmup peak with no margin; supported_resolutions picks the largest - warmup and target frame counts share one contract helper that also applies the real request's num_gpus latent alignment (extracted from SamplingParams._adjust_visual_fields) Distributed and rollback robustness: - fence everything before the first all-gather into a skip report (a raise there parked peer ranks in the collective until timeout); worker-side dp>1 guard; pipeline-None guard - error gathers filter on `is not None` and describe_error() keeps empty-string exceptions (str(AssertionError()) == "") visible - rollback failures inside apply raise AutoResidencyRollbackError and reach ROLLBACK_FAILED (abort) instead of masquerading as rolled_back; rollback_promotions undoes every component and aggregates errors - disable_offload re-arms hooks when load_all_layers fails, so a failed promotion cannot leave a hook-less enabled manager serving (1,) placeholders - mechanism/mode mismatches (offload_during_compile window) are excluded from candidates and guarded at apply/rollback - _required_resident_components records an owner feature; auto-residency rollback can no longer release a loader's hard requirement Startup contract: - apply-RPC failures honor the warmup fail-open contract instead of SIGTERMing startup; the post-rollback restore warmup keeps the fail-closed contract of explicit --warmup-resolutions - calibration request is gated on auto_residency_skip_reason, so the kill switch / quantized / manual configs no longer pay an extra warmup forward they never consume - warmup records are consumed exactly once and the handler is one-shot; re-warm passes no longer double-advance the warmup progress bar - debug residency hint is fenced (never fails a completed request), resolves the default workload once, and sizes layerwise components from metadata instead of materializing per-weight views Co-Authored-By: Claude Fable 5 --- .../configs/sample/sampling_params.py | 105 ++++---- .../runtime/managers/gpu_worker.py | 246 +++++++++++++----- .../memory_managers/auto_residency.py | 152 ++++++++--- .../memory_managers/layerwise_offload.py | 34 ++- .../runtime/server_args/server_args.py | 21 +- .../multimodal_gen/runtime/server_warmup.py | 54 +++- .../runtime/warmup_request_builder.py | 125 ++++++--- .../test/unit/test_auto_residency.py | 201 ++++++++++++-- 8 files changed, 725 insertions(+), 213 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index f28227fec6ab..a846de1a587a 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -78,6 +78,58 @@ def _sanitize_filename(name: str, replacement: str = "_", max_length: int = 150) return ascii_name +_SEQUENCE_SHARD_PIPELINE_FAMILIES = ("wan", "helios", "joy", "cosmos3") + + +def resolve_sequence_shard( + pipeline_config: Any, enable_sequence_shard: bool | None +) -> bool: + """Whether this pipeline shards the sequence dim instead of aligning frames. + + Shared by ``SamplingParams._adjust_visual_fields`` and the synthetic + warmup builder so warmup requests follow the same frame contract as real + requests. + """ + pipeline_name_lower = pipeline_config.__class__.__name__.lower() + return any( + family in pipeline_name_lower for family in _SEQUENCE_SHARD_PIPELINE_FAMILIES + ) and (enable_sequence_shard is None or enable_sequence_shard) + + +def align_num_frames_for_num_gpus( + num_frames: int, + *, + num_gpus: int, + vae_config: Any, + round_down: bool, +) -> int: + """Align the latent frame count to be divisible by ``num_gpus``.""" + if num_gpus <= 1: + return num_frames + use_temporal_scaling_frames = vae_config.use_temporal_scaling_frames + temporal_scale_factor = vae_config.arch_config.temporal_compression_ratio + + if use_temporal_scaling_frames: + orig_latent_num_frames = (num_frames - 1) // temporal_scale_factor + 1 + else: + orig_latent_num_frames = num_frames + + if orig_latent_num_frames % num_gpus == 0: + return num_frames + + if round_down: + # Ensure we have at least 1 batch per GPU + new_latent_num_frames = max(1, (orig_latent_num_frames // num_gpus)) * num_gpus + else: + new_latent_num_frames = math.ceil(orig_latent_num_frames / num_gpus) * num_gpus + + if use_temporal_scaling_frames: + # Convert back to frames, keeping num_frames-1 a multiple of the + # temporal scale factor + return (new_latent_num_frames - 1) * temporal_scale_factor + 1 + return new_latent_num_frames + + class DataType(Enum): IMAGE = auto() VIDEO = auto() @@ -725,14 +777,9 @@ def _adjust_visual_fields(self, server_args, pipeline_config): ) logger.warning(error_msg) - pipeline_name_lower = server_args.pipeline_config.__class__.__name__.lower() - - if ( - "wan" in pipeline_name_lower - or "helios" in pipeline_name_lower - or "joy" in pipeline_name_lower - or "cosmos3" in pipeline_name_lower - ) and (self.enable_sequence_shard is None or self.enable_sequence_shard): + if resolve_sequence_shard( + server_args.pipeline_config, self.enable_sequence_shard + ): self.enable_sequence_shard = True logger.debug("Automatically enabled enable_sequence_shard") else: @@ -765,43 +812,13 @@ def _adjust_visual_fields(self, server_args, pipeline_config): ) if self.adjust_frames: - # Adjust number of frames based on number of GPUs for video task - use_temporal_scaling_frames = ( - pipeline_config.vae_config.use_temporal_scaling_frames - ) - num_frames = self.num_frames - num_gpus = server_args.num_gpus - temporal_scale_factor = ( - pipeline_config.vae_config.arch_config.temporal_compression_ratio + new_num_frames = align_num_frames_for_num_gpus( + self.num_frames, + num_gpus=server_args.num_gpus, + vae_config=pipeline_config.vae_config, + round_down=self.num_frames_round_down, ) - - if use_temporal_scaling_frames: - orig_latent_num_frames = ( - num_frames - 1 - ) // temporal_scale_factor + 1 - else: - orig_latent_num_frames = num_frames - - if orig_latent_num_frames % server_args.num_gpus != 0: - # Adjust latent frames to be divisible by number of GPUs - if self.num_frames_round_down: - # Ensure we have at least 1 batch per GPU - new_latent_num_frames = ( - max(1, (orig_latent_num_frames // num_gpus)) * num_gpus - ) - else: - new_latent_num_frames = ( - math.ceil(orig_latent_num_frames / num_gpus) * num_gpus - ) - - if use_temporal_scaling_frames: - # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor - new_num_frames = ( - new_latent_num_frames - 1 - ) * temporal_scale_factor + 1 - else: - new_num_frames = new_latent_num_frames - + if new_num_frames != self.num_frames: logger.info( "Adjusting number of frames from %s to %s based on number of GPUs (%s)", self.num_frames, diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 65db6e79f995..40ba239e7897 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -57,11 +57,13 @@ PROMOTION_STATUS_ROLLED_BACK, PROMOTION_STATUS_SKIPPED, AppliedPromotion, + AutoResidencyRollbackError, DefaultWorkload, RankResidencyReport, WarmupMemoryRecord, apply_promotions, collect_promotion_candidates, + describe_error, estimate_default_workload_peak_bytes, format_applied_changes, format_plan_summary, @@ -188,6 +190,23 @@ def __init__( # the auto-residency promotion decision before the server turns ready self._auto_residency_warmup_records: list[WarmupMemoryRecord] = [] self._auto_residency_applied: list[AppliedPromotion] = [] + # default workload resolved once for the per-request residency hint + self._cached_default_workload: DefaultWorkload | None = None + self._cached_default_workload_failed = False + + def _default_workload_for_hint(self) -> DefaultWorkload | None: + if ( + self._cached_default_workload is None + and not self._cached_default_workload_failed + ): + try: + self._cached_default_workload = resolve_default_workload( + self.server_args + ) + except Exception: + logger.debug("Default workload unresolvable", exc_info=True) + self._cached_default_workload_failed = True + return self._cached_default_workload def release_realtime_session(self, session_id: str) -> OutputBatch: """release the session of a realtime connection""" @@ -374,25 +393,37 @@ def do_mem_analysis(self, output_batch: OutputBatch): remaining_gpu_mem_gb = ( current_platform.get_device_total_memory() / (1024**3) - peak_reserved_gb ) - can_stay_resident = self.get_can_stay_resident_components(remaining_gpu_mem_gb) + try: + can_stay_resident = self.get_can_stay_resident_components( + remaining_gpu_mem_gb + ) + except Exception: + # a debug-only hint must never fail a completed request + logger.debug("Residency hint unavailable", exc_info=True) + can_stay_resident = [] pool_overhead_gb = peak_reserved_gb - peak_allocated_gb pool_overhead_pct = ( pool_overhead_gb / peak_reserved_gb * 100 if peak_reserved_gb else 0.0 ) - logger.debug( - "GPU memory: peak=%.2f GB, allocated=%.2f GB, pool=%.2f GB (%.1f%%), " - "headroom=%.2f GB. Components that can remain on GPU: %s. " + residency_hint = ( + f" Components that can remain on GPU: {can_stay_resident}. " "Make it explicit with --component-residency =resident; " "--performance-mode auto with server warmup applies safe " - "promotions automatically.", + "promotions automatically." + if can_stay_resident + else "" + ) + logger.debug( + "GPU memory: peak=%.2f GB, allocated=%.2f GB, pool=%.2f GB (%.1f%%), " + "headroom=%.2f GB.%s", peak_reserved_gb, peak_allocated_gb, pool_overhead_gb, pool_overhead_pct, remaining_gpu_mem_gb, - can_stay_resident, + residency_hint, ) def execute_forward( @@ -513,11 +544,23 @@ def _execute_forward_common( """ output_batch = None forward_failed = False + # Prewarm reqs (is_warmup=False) run a different offload layout and + # must not contaminate the calibration records. measure_server_warmup = ( - bool(req.extra.get("server_based_warmup")) and not current_platform.is_cpu() + req.is_warmup + and bool(req.extra.get("server_based_warmup")) + and not current_platform.is_cpu() + and not current_platform.is_mps() ) warmup_baseline_allocated_bytes = 0 try: + if measure_server_warmup: + # Drop the previous request's allocator pool so max_memory + # _reserved measures THIS workload's peak. Without this, a + # later warmup shape that fits in the cached pool reports the + # same peak as the earlier one and the calibration slope + # collapses to zero. + torch.get_device_module().empty_cache() if not current_platform.is_cpu() and not current_platform.is_mps(): torch.get_device_module().reset_peak_memory_stats() if measure_server_warmup: @@ -635,12 +678,16 @@ def _execute_forward_common( # clean cache if OOM if not current_platform.is_cpu(): torch.get_device_module().empty_cache() - if measure_server_warmup and output_batch is not None: - self._record_server_warmup_memory( - req=req, - baseline_allocated_bytes=warmup_baseline_allocated_bytes, - succeeded=output_batch.error is None, - ) + finally: + # also runs on the propagate_forward_errors re-raise: a warmup + # forward that never completed must still leave a failed record, + # or the estimator would plan from the remaining partial data + if measure_server_warmup: + self._record_server_warmup_memory( + req=req, + baseline_allocated_bytes=warmup_baseline_allocated_bytes, + succeeded=output_batch is not None and output_batch.error is None, + ) return output_batch def _record_server_warmup_memory( @@ -1072,14 +1119,17 @@ def get_can_stay_resident_components( Reuses the auto-residency candidate collection and its H2D-savings greedy order, so this post-request hint names the same components, - in the same order, that ``--performance-mode auto`` would promote. - Unlike the promotion plan it applies no reserve or margin: it reports - raw capacity after this request, not a placement decision. + in the same order, that ``--performance-mode auto`` would promote + (it deliberately does not maximize the component count the way the + old fixed cheap-first order did, and it omits components driven by a + pipeline-custom residency strategy -- user flags cannot act on + those). Unlike the promotion plan it applies no reserve or margin: + it reports raw capacity after this request, not a placement decision. """ - if not self.pipeline: + if not self.pipeline or not self.pipeline.modules: return [] - workload = resolve_default_workload(self.server_args) + workload = self._default_workload_for_hint() candidates = collect_promotion_candidates( modules=self.pipeline.modules, residency_mode_of=self.server_args.residency_mode, @@ -1088,7 +1138,9 @@ def get_can_stay_resident_components( # requires it (promotion itself never touches explicit ones). explicit_residency_mode_of=lambda _name: None, custom_strategy_names=self.pipeline.component_residency_strategies, - num_inference_steps=workload.num_inference_steps, + num_inference_steps=( + workload.num_inference_steps if workload is not None else 1 + ), ) can_stay_resident = [] @@ -1106,32 +1158,78 @@ def apply_auto_residency(self) -> OutputBatch: is computed from all-gathered rank reports so each rank reaches the same decision. A failure on any rank rolls every rank back to the original strategy. + + Everything before the first all-gather is fenced into a skip report: + an uncaught raise there would leave the peer ranks parked in the + collective until the group timeout. """ - workload = resolve_default_workload(self.server_args) - local_report = self._build_auto_residency_report(workload=workload) + if (self.server_args.dp_size or 1) > 1: + # DP replicas dequeue control reqs at independent times, but the + # gather below spans the whole torch world; entering it per + # replica would deadlock. (The HTTP orchestrator also gates this + # case; this guard keeps the worker safe on its own.) + return OutputBatch(output={"status": PROMOTION_STATUS_SKIPPED}) + if self._auto_residency_applied: + # one-shot per process: records measured after a promotion + # describe the promoted layout, not the startup strategy + return OutputBatch(output={"status": PROMOTION_STATUS_SKIPPED}) + + records = list(self._auto_residency_warmup_records) + # consumed exactly once; the post-promotion re-warm appends fresh + # records that must not feed a later decision + self._auto_residency_warmup_records.clear() + + try: + workload = resolve_default_workload(self.server_args) + local_report = self._build_auto_residency_report( + workload=workload, records=records + ) + except Exception as e: + logger.warning( + "Auto residency: rank %d could not build its report: %s", + self.rank, + e, + exc_info=True, + ) + workload = DefaultWorkload( + width=None, height=None, num_frames=1, num_inference_steps=1 + ) + local_report = RankResidencyReport( + rank=self.rank, + budget_bytes=0, + estimated_peak_bytes=None, + candidates=[], + skip_reason=describe_error(e), + ) reports = self._auto_residency_all_gather(local_report) plan = plan_auto_residency(reports=reports) - summary = format_plan_summary( - plan=plan, - workload=workload, - records=self._auto_residency_warmup_records, - ) + summary = format_plan_summary(plan=plan, workload=workload, records=records) if plan.skip_reason is not None or not plan.promotions: if self.is_output_rank: logger.info("%s", summary) return OutputBatch( - output=plan_summary_payload( - plan=plan, workload=workload, status=PROMOTION_STATUS_SKIPPED - ) + output=plan_summary_payload(plan=plan, status=PROMOTION_STATUS_SKIPPED) ) apply_error: str | None = None + local_rollback_failed = False try: self._auto_residency_applied = apply_promotions( plan=plan, modules=self.pipeline.modules, server_args=self.server_args, ) + except AutoResidencyRollbackError as e: + logger.error( + "Auto residency promotion failed on rank %d and the rank " + "could not roll itself back: %s", + self.rank, + e, + exc_info=True, + ) + self._auto_residency_applied = [] + apply_error = describe_error(e) + local_rollback_failed = True except Exception as e: # this rank already rolled itself back logger.error( "Auto residency promotion failed on rank %d: %s", @@ -1140,13 +1238,15 @@ def apply_auto_residency(self) -> OutputBatch: exc_info=True, ) self._auto_residency_applied = [] - apply_error = str(e) + apply_error = describe_error(e) - rank_errors = [ - error for error in self._auto_residency_all_gather(apply_error) if error - ] + gathered = self._auto_residency_all_gather((apply_error, local_rollback_failed)) + rank_errors = [error for error, _ in gathered if error is not None] + any_rollback_failed = any(failed for _, failed in gathered) if rank_errors: - return self._rollback_auto_residency_after_failure(cause=rank_errors[0]) + return self._rollback_everywhere( + cause=rank_errors[0], already_failed=any_rollback_failed + ) self._invalidate_component_strategies( [candidate.component_name for candidate in plan.promotions] @@ -1155,22 +1255,47 @@ def apply_auto_residency(self) -> OutputBatch: logger.info("%s", summary) logger.info("%s", format_applied_changes(plan=plan)) return OutputBatch( - output=plan_summary_payload( - plan=plan, workload=workload, status=PROMOTION_STATUS_PROMOTED - ) + output=plan_summary_payload(plan=plan, status=PROMOTION_STATUS_PROMOTED) ) def rollback_auto_residency(self) -> OutputBatch: """Restore the pre-promotion residency (post-promotion re-warm failed).""" + return self._rollback_everywhere(cause=None, already_failed=False) + + def _rollback_everywhere( + self, *, cause: str | None, already_failed: bool + ) -> OutputBatch: + """Roll this rank back and gather the replica-wide outcome. + + ``cause`` is the promotion failure that triggered the rollback (None + when the rollback was requested after a failed re-warm); + ``already_failed`` marks that some rank already failed its in-apply + rollback, which is fatal regardless of what the remaining ranks do. + Collective-symmetric: every rank gathers exactly once. + """ rollback_error = self._rollback_applied_promotions() - rank_errors = [ - error for error in self._auto_residency_all_gather(rollback_error) if error - ] + gathered = self._auto_residency_all_gather(rollback_error) + rank_errors = [error for error in gathered if error is not None] + if already_failed and not rank_errors: + rank_errors = ["a rank could not undo its own promotions"] if rank_errors: + prefix = ( + f"auto residency promotion failed ({cause}) and rollback failed" + if cause is not None + else "auto residency rollback failed" + ) return OutputBatch( - error=f"auto residency rollback failed: {rank_errors[0]}", + error=f"{prefix}: {rank_errors[0]}", output={"status": PROMOTION_STATUS_ROLLBACK_FAILED}, ) + if cause is not None: + return OutputBatch( + error=( + "auto residency promotion failed; original strategy " + f"restored: {cause}" + ), + output={"status": PROMOTION_STATUS_ROLLED_BACK}, + ) if self.is_output_rank: logger.info( "Auto residency: rolled back; the residency configured at " @@ -1180,10 +1305,23 @@ def rollback_auto_residency(self) -> OutputBatch: return OutputBatch(output={"status": PROMOTION_STATUS_ROLLED_BACK}) def _build_auto_residency_report( - self, *, workload: DefaultWorkload + self, *, workload: DefaultWorkload, records: List[WarmupMemoryRecord] ) -> RankResidencyReport: - records = list(self._auto_residency_warmup_records) - skip_reason = None if records else "no server warmup measurements" + skip_reason = None + if self.pipeline is None: + skip_reason = "pipeline not initialized" + elif not records: + skip_reason = "no server warmup measurements" + elif workload.workload_units() is None: + skip_reason = "default workload resolution unknown" + if skip_reason is not None: + return RankResidencyReport( + rank=self.rank, + budget_bytes=0, + estimated_peak_bytes=None, + candidates=[], + skip_reason=skip_reason, + ) estimated_peak_bytes = estimate_default_workload_peak_bytes( records=records, target_units=workload.workload_units() ) @@ -1234,31 +1372,13 @@ def _rollback_applied_promotions(self) -> str | None: e, exc_info=True, ) - return str(e) + return describe_error(e) self._invalidate_component_strategies( [promotion.component_name for promotion in applied] ) self._auto_residency_applied = [] return None - def _rollback_auto_residency_after_failure(self, *, cause: str) -> OutputBatch: - rollback_error = self._rollback_applied_promotions() - rank_errors = [ - error for error in self._auto_residency_all_gather(rollback_error) if error - ] - if rank_errors: - return OutputBatch( - error=( - f"auto residency promotion failed ({cause}) and rollback " - f"failed: {rank_errors[0]}" - ), - output={"status": PROMOTION_STATUS_ROLLBACK_FAILED}, - ) - return OutputBatch( - error=f"auto residency promotion failed; original strategy restored: {cause}", - output={"status": PROMOTION_STATUS_ROLLED_BACK}, - ) - @staticmethod def _invalidate_component_strategies(component_names: List[str]) -> None: manager = peek_global_component_residency_manager() diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 7cdb291b98a2..d5082e204b78 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -35,6 +35,7 @@ ) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, + is_layerwise_offloaded_module, ) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( is_dit_component_name, @@ -64,6 +65,21 @@ PROMOTION_STATUS_ROLLBACK_FAILED = "rollback_failed" +class AutoResidencyRollbackError(RuntimeError): + """A promotion failed AND undoing the already-applied part also failed. + + The rank is in a mixed residency state that only a restart fixes; callers + must abort startup instead of continuing on a half-rolled-back replica. + """ + + +def describe_error(error: BaseException) -> str: + """Never-empty error text (str(AssertionError()) is "" and would be + dropped by any truthiness filter).""" + text = str(error) + return f"{type(error).__name__}: {text}" if text else type(error).__name__ + + class WarmupMemoryRecord(msgspec.Struct, frozen=True): """Per-rank memory measurement of one server warmup forward.""" @@ -135,17 +151,25 @@ class AppliedPromotion(msgspec.Struct, frozen=True): def resolve_default_workload(server_args: ServerArgs) -> DefaultWorkload: """Resolve the default request shape promotion is optimized for.""" from sglang.multimodal_gen.runtime.warmup_request_builder import ( + _apply_warmup_frame_contract, get_model_sampling_defaults, ) defaults = get_model_sampling_defaults(server_args) + # __post_init__ already applied the model's _default_width/_default_height width = defaults.width height = defaults.height if (width is None or height is None) and defaults.supported_resolutions: - width, height = defaults.supported_resolutions[0] + # worst-case the target with the largest supported shape + width, height = max( + defaults.supported_resolutions, key=lambda size: size[0] * size[1] + ) num_frames = defaults.num_frames or 1 if num_frames > 1: - num_frames = server_args.pipeline_config.adjust_num_frames(num_frames) + # same contract the warmup requests and real requests get + num_frames = _apply_warmup_frame_contract( + server_args, defaults, num_frames=num_frames + ) return DefaultWorkload( width=width, height=height, @@ -173,13 +197,16 @@ def estimate_default_workload_peak_bytes( 3. One usable size: scale everything above the pre-forward allocated baseline (conservative; may block promotion but never over-promotes). - Returns None when there is no usable measurement (no records, or any - warmup forward failed -- a failed warmup means the measured peak does not - cover the full request path). + Returns None when the estimate cannot be trusted: no records, any warmup + forward failed (its peak does not cover the full request path), or the + target workload is unknown (an unknown target would silently equate the + area/frame-capped warmup peak with the real serving peak). """ records = list(records) if not records or any(not record.succeeded for record in records): return None + if target_units is None: + return None peak_by_units: dict[int, int] = {} for record in records: @@ -188,9 +215,6 @@ def estimate_default_workload_peak_bytes( peak_by_units.get(units, 0), record.peak_reserved_bytes ) - if target_units is None: - return max(peak_by_units.values()) - covering_peaks = [ peak for units, peak in peak_by_units.items() if units >= target_units ] @@ -240,15 +264,11 @@ def _module_weight_bytes(module: nn.Module) -> int: def _layerwise_offloaded_bytes(module: LayerwiseOffloadableModuleMixin) -> int: - total = 0 - for manager in module.layerwise_offload_managers: - if not manager.enabled: - continue - total += sum( - tensor.numel() * tensor.element_size() - for _, tensor in manager.iter_cpu_weights() - ) - return total + return sum( + manager.offloaded_weight_bytes() + for manager in module.layerwise_offload_managers + if manager.enabled + ) def component_resident_size_bytes(module: nn.Module, residency_mode: str) -> int: @@ -289,6 +309,13 @@ def collect_promotion_candidates( continue if is_fsdp_managed_module(module): continue + if mode == COMPONENT_OFFLOAD and is_layerwise_offloaded_module(module): + # The module is layerwise-managed despite its configured mode + # (e.g. the offload_during_compile window); flipping the flag + # would be a silent no-op while the manager keeps streaming. + continue + if mode == LAYERWISE_OFFLOAD and not is_layerwise_offloaded_module(module): + continue promoted_bytes = component_resident_size_bytes(module, mode) if promoted_bytes <= 0: continue @@ -391,6 +418,11 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP promoted_bytes = 0 for candidate in ordered: next_bytes = promoted_bytes + candidate.promoted_weight_bytes + # Known conservatism: estimated_peak already contains whatever part + # of this component was on GPU at the measured peak (streamed + # layers, keep-ready components), so those bytes are counted twice. + # The error is in the safe direction; refining it needs per-phase + # peaks. if estimated_peak + next_bytes + reserve <= budget: promotions.append(candidate) promoted_bytes = next_bytes @@ -418,7 +450,9 @@ def apply_promotions( Raises on failure after undoing the promotions already applied, so a caller observing an exception knows this rank is back on the original - strategy. + strategy -- unless the undo itself fails, which raises + ``AutoResidencyRollbackError`` (the rank is in a mixed state and startup + must abort). """ applied: list[AppliedPromotion] = [] try: @@ -428,6 +462,20 @@ def apply_promotions( raise RuntimeError( f"promotion target {candidate.component_name!r} is missing" ) + if candidate.residency_mode == LAYERWISE_OFFLOAD and not isinstance( + module, LayerwiseOffloadableModuleMixin + ): + raise RuntimeError( + f"promotion target {candidate.component_name!r} lost its " + "layerwise offload capability between planning and apply" + ) + if candidate.residency_mode == COMPONENT_OFFLOAD and ( + is_layerwise_offloaded_module(module) + ): + raise RuntimeError( + f"promotion target {candidate.component_name!r} became " + "layerwise-managed between planning and apply" + ) # Record before acting so a mid-promotion failure rolls back the # partially promoted component as well. applied.append( @@ -440,10 +488,17 @@ def apply_promotions( candidate.component_name, feature_name=AUTO_RESIDENCY_FEATURE_NAME ) if candidate.residency_mode == LAYERWISE_OFFLOAD: - assert isinstance(module, LayerwiseOffloadableModuleMixin) module.disable_offload() - except Exception: - rollback_promotions(applied=applied, modules=modules, server_args=server_args) + except Exception as apply_error: + try: + rollback_promotions( + applied=applied, modules=modules, server_args=server_args + ) + except Exception as rollback_error: + raise AutoResidencyRollbackError( + f"promotion failed ({describe_error(apply_error)}) and rollback " + f"failed ({describe_error(rollback_error)})" + ) from apply_error raise return applied @@ -454,19 +509,39 @@ def rollback_promotions( modules: Mapping[str, object], server_args: ServerArgs, ) -> None: - """Restore the original residency for previously applied promotions.""" + """Restore the original residency for previously applied promotions. + + Every promotion is undone even when one of them fails; the collected + failures are re-raised at the end so one broken component cannot leave + the later (earlier-applied) ones promoted. + """ + errors: list[str] = [] for promotion in reversed(list(applied)): - server_args.release_required_component_residency(promotion.component_name) - module = modules.get(promotion.component_name) - if not isinstance(module, nn.Module): - continue - if promotion.residency_mode == LAYERWISE_OFFLOAD: - assert isinstance(module, LayerwiseOffloadableModuleMixin) - module.enable_offload() - else: - module.to("cpu") + try: + server_args.release_required_component_residency( + promotion.component_name, + feature_name=AUTO_RESIDENCY_FEATURE_NAME, + ) + module = modules.get(promotion.component_name) + if not isinstance(module, nn.Module): + continue + if promotion.residency_mode == LAYERWISE_OFFLOAD: + if not isinstance(module, LayerwiseOffloadableModuleMixin): + raise RuntimeError("lost layerwise offload capability") + module.enable_offload() + elif is_layerwise_offloaded_module(module): + # A layerwise manager owns this module's placement (e.g. the + # offload_during_compile window); moving it to CPU behind the + # manager would strand its bookkeeping. + pass + else: + module.to("cpu") + except Exception as e: + errors.append(f"{promotion.component_name}: {describe_error(e)}") if torch.get_device_module().is_available(): torch.get_device_module().empty_cache() + if errors: + raise RuntimeError("; ".join(errors)) def format_plan_summary( @@ -526,16 +601,13 @@ def format_applied_changes(*, plan: AutoResidencyPlan) -> str: ) -def plan_summary_payload( - *, plan: AutoResidencyPlan, workload: DefaultWorkload, status: str -) -> dict: - """Structured decision summary returned to the warmup orchestrator.""" +def plan_summary_payload(*, plan: AutoResidencyPlan, status: str) -> dict: + """Minimal decision payload for the warmup orchestrator. + + The orchestrator only branches on ``status``; the human-readable detail + lives in the logged ``format_plan_summary`` line. + """ return { "status": status, - "skip_reason": plan.skip_reason, - "target_workload": workload.describe(), - "estimated_peak_bytes": plan.estimated_peak_bytes, - "reserve_bytes": plan.reserve_bytes, - "budget_bytes": plan.budget_bytes, "promoted": [candidate.component_name for candidate in plan.promotions], } 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 3e391a3c46ca..cc6562cc47cd 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 @@ -1,4 +1,5 @@ import bisect +import math import queue import re import threading @@ -1336,6 +1337,22 @@ def update_cpu_weights( return updated_names + def offloaded_weight_bytes(self) -> int: + """Total bytes of this manager's offloadable weights, from metadata. + + Size-only companion to ``iter_cpu_weights`` -- it never materializes + the per-weight buffer views, so it is safe to call per request. + """ + total = 0 + for layer_meta in self._weight_metadata.values(): + for meta in layer_meta.values(): + if meta.get("preserve_strides", False): + numel = math.prod(meta["shape"]) + else: + numel = meta["numel"] + total += int(numel) * meta["dtype"].itemsize + return total + def iter_cpu_weights(self): """Yield (name, tensor) pairs from consolidated CPU buffers. @@ -1826,10 +1843,21 @@ def disable_offload(self) -> None: if self.layerwise_offload_managers is None: return for manager in self.layerwise_offload_managers: - if manager.enabled: - manager.remove_forward_hooks() + if not manager.enabled: + continue + manager.remove_forward_hooks() + try: manager.load_all_layers() - manager.enabled = False + except Exception: + # Loading every layer is the step most likely to OOM. Re-arm + # this manager before re-raising: leaving it hook-less with + # enabled=True would let release_all() swap weights for (1,) + # placeholders that nothing ever streams back in, and + # enable_offload()'s already-armed guard would skip repair. + manager.release_all() + manager.register_forward_hooks() + raise + manager.enabled = False def enable_offload(self) -> None: """Re-enable layerwise offload: sync weights to CPU, release layers, and restore hooks.""" 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 cb0415426f72..820e682bc341 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -368,8 +368,10 @@ class ServerArgs(DisaggServerArgsMixin): pin_cpu_memory: bool = True ltx2_two_stage_device_mode: str | None = None _explicit_arg_names: set[str] = field(default_factory=set, repr=False) - _required_resident_components: set[str] = field( - default_factory=set, init=False, repr=False + # component name -> owning feature; the first (loader-time) owner wins so + # a runtime feature's rollback can never release a loader hard requirement + _required_resident_components: dict[str, str] = field( + default_factory=dict, init=False, repr=False ) _fsdp_disabled_components: set[str] = field( default_factory=set, init=False, repr=False @@ -1541,11 +1543,18 @@ def require_component_resident( f"{feature_name} requires {component_name!r} to be resident; " f"got {configured_mode!r} from --component-residency" ) - self._required_resident_components.add(component_name) + self._required_resident_components.setdefault(component_name, feature_name) + + def release_required_component_residency( + self, component_name: str, *, feature_name: str + ) -> None: + """Undo ``require_component_resident`` for the same owning feature. - def release_required_component_residency(self, component_name: str) -> None: - """Undo ``require_component_resident`` (auto-residency rollback).""" - self._required_resident_components.discard(component_name) + A release with a different owner is a no-op: an auto-residency + rollback must not clear a loader's hard residency requirement. + """ + if self._required_resident_components.get(component_name) == feature_name: + del self._required_resident_components[component_name] def should_use_fsdp_for_component(self, component_name: str) -> bool: return bool( diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index 7651e6042f5c..26e90c107a3e 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -124,6 +124,14 @@ def auto_residency_skip_reason(server_args: ServerArgs) -> str | None: return "diffusers backend" if server_args.enable_breakable_cuda_graph: return "breakable CUDA graph captures during warmup" + if server_args.enable_torch_compile: + # Compile warmup runs a deliberately stripped memory layout: the DiT + # is temporarily layerwise-offloaded (offload_during_compile) and + # _move_resident_components_for_warmup evicts resident aux components + # for the denoise. Peaks measured there are far below a real + # request's, and the post-promotion re-warm (also a warmup layout) + # cannot catch the difference. + return "torch.compile warmup uses a stripped memory layout" if envs.SGLANG_CACHE_DIT_ENABLED: return "cache-dit enabled" if server_args.batching_max_size > 1: @@ -184,7 +192,20 @@ async def maybe_apply_auto_residency( "Server warmup complete; adjusting component residency for the " "default workload (--performance-mode auto)." ) - response = await forward(AutoResidencyReq(action="apply")) + # Same fail-open contract as the warmup itself: implicit warmups must + # never abort startup, explicit --warmup-resolutions ones must succeed. + fail_open = server_args.warmup_resolutions is None + try: + response = await forward(AutoResidencyReq(action="apply")) + except Exception as e: + if not fail_open: + raise + logger.warning( + "Auto residency apply request failed; continuing on the original " + "strategy: %s", + e, + ) + return status = _auto_residency_status(response) if status == PROMOTION_STATUS_ROLLBACK_FAILED: raise RuntimeError(f"auto residency rollback failed: {response.error}") @@ -202,7 +223,9 @@ async def maybe_apply_auto_residency( # moves promoted components, rebuilds compile caches invalidated by the # placement change, and proves the promoted layout fits before ready. try: - await run_async_client_warmup(server_args, forward, fail_open=False) + await run_async_client_warmup( + server_args, forward, fail_open=False, rewarm=True + ) except Exception as e: logger.warning( "Post-promotion re-warmup failed (%s); rolling back auto residency", e @@ -215,8 +238,11 @@ async def maybe_apply_auto_residency( raise RuntimeError( f"auto residency rollback failed: {rollback.error}" ) from e - # restore warm caches for the original strategy before turning ready - await run_async_client_warmup(server_args, forward, fail_open=True) + # restore warm caches for the original strategy before turning ready, + # keeping the caller's fail-open contract + await run_async_client_warmup( + server_args, forward, fail_open=fail_open, rewarm=True + ) def format_warmup_req(req_or_group: Any) -> str: @@ -249,6 +275,7 @@ def build_client_warmup_reqs( server_args: ServerArgs, *, warmup_input_path: str | None = None, + rewarm: bool = False, ) -> list[Req]: warmup_reqs = build_warmup_reqs( server_args, @@ -261,6 +288,10 @@ def build_client_warmup_reqs( for req in warmup_reqs: if req.is_warmup: req.extra["warmup_total"] = warmup_total + if rewarm: + # a repeat pass after an auto-residency change: keep it out of + # the scheduler's warmup progress accounting (already at N/N) + req.extra["server_warmup_rewarm"] = True return warmup_reqs @@ -269,6 +300,7 @@ async def run_async_client_warmup( forward: Callable[[Req], Awaitable[OutputBatch]], *, fail_open: bool = False, + rewarm: bool = False, ) -> None: try: warmup_input_path = None @@ -276,7 +308,7 @@ async def run_async_client_warmup( warmup_input_path = prepare_warmup_image_path(server_args) for req in build_client_warmup_reqs( - server_args, warmup_input_path=warmup_input_path + server_args, warmup_input_path=warmup_input_path, rewarm=rewarm ): response = await forward(req) if response.error is not None: @@ -393,6 +425,18 @@ def _log_warmup_result( if not is_warmup: return + req = get_first_generation_req(req_or_group) + if req is not None and req.extra.get("server_warmup_rewarm"): + # auto-residency re-warm passes repeat already-counted requests; + # advancing the bar again would log N+1/N in CI + if output_batch.error is not None: + logger.warning( + "%s processing failed: %s", + self._format_warmup_req(req_or_group), + output_batch.error, + ) + return + server_based_warmup = is_server_based_warmup(req_or_group) self._warmup_processed += 1 self._advance_warmup_progress_bar(req_or_group, output_batch) diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py index dfe0c8b35baf..726da63fab78 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -17,7 +17,11 @@ from typing import Any from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType -from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams +from sglang.multimodal_gen.configs.sample.sampling_params import ( + SamplingParams, + align_num_frames_for_num_gpus, + resolve_sequence_shard, +) from sglang.multimodal_gen.registry import get_pipeline_config_classes from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.server_args import ( @@ -260,34 +264,70 @@ def _resolve_warmup_num_frames( ) warmup_num_frames = min(num_frames, frame_budget) - return server_args.pipeline_config.adjust_num_frames(warmup_num_frames) + return _apply_warmup_frame_contract( + server_args, sampling_defaults, num_frames=warmup_num_frames + ) + + +def _apply_warmup_frame_contract( + server_args: ServerArgs, sampling_defaults: SamplingParams, *, num_frames: int +) -> int: + """Re-apply the real-request frame contract to a capped warmup frame count. + + Warmup requests skip ``SamplingParams._adjust``, so the cap must run the + model frame contract itself -- without this, e.g. LongLive2's capped 17 + frames map to 5 latent frames (not divisible by its 8-frame causal block) + and every server warmup fails silently under fail-open. Pipelines that + align frames to ``num_gpus`` (rather than sharding the sequence dim) get + the same latent alignment real requests get. + """ + num_frames = server_args.pipeline_config.adjust_num_frames(num_frames) + if sampling_defaults.adjust_frames and not resolve_sequence_shard( + server_args.pipeline_config, sampling_defaults.enable_sequence_shard + ): + num_frames = align_num_frames_for_num_gpus( + num_frames, + num_gpus=server_args.num_gpus, + vae_config=server_args.pipeline_config.vae_config, + round_down=sampling_defaults.num_frames_round_down, + ) + return num_frames def _resolve_calibration_num_frames( server_args: ServerArgs, sampling_defaults: SamplingParams, - warmup_num_frames: int, *, + warmup_num_frames: int | None, server_based_warmup: bool, ) -> int | None: """A second, smaller warmup frame count for auto-residency calibration. - Only produced when the main video warmup was frame-capped under - ``--performance-mode auto``: the auto-residency estimator then has two - workload sizes to fit the constant/linear split of the peak. Returns None - when a second measurement adds no information (same adjusted frames). + Only produced when the main video warmup was frame-capped and the + auto-residency promotion is actually eligible to consume the measurement + (``auto_residency_skip_reason``): the estimator then has two workload + sizes to fit the constant/linear split of the peak. Returns None when a + second measurement adds no information (same adjusted frames) or would + never be read. """ + # local import: server_warmup imports this module at module level + from sglang.multimodal_gen.runtime.server_warmup import ( + auto_residency_skip_reason, + ) + if not server_based_warmup or not _is_video_warmup_task(server_args): return None - if server_args.performance_mode != "auto": + if warmup_num_frames is None: return None - if getattr(server_args, "enable_breakable_cuda_graph", False) is True: + if auto_residency_skip_reason(server_args) is not None: return None - default_num_frames = getattr(sampling_defaults, "num_frames", 1) or 1 + default_num_frames = sampling_defaults.num_frames or 1 if warmup_num_frames >= default_num_frames: return None - calibration_num_frames = server_args.pipeline_config.adjust_num_frames( - SERVER_WARMUP_CALIBRATION_VIDEO_FRAMES + calibration_num_frames = _apply_warmup_frame_contract( + server_args, + sampling_defaults, + num_frames=SERVER_WARMUP_CALIBRATION_VIDEO_FRAMES, ) if calibration_num_frames == warmup_num_frames or calibration_num_frames <= 0: return None @@ -408,7 +448,7 @@ def build_warmup_reqs( _resolve_calibration_num_frames( server_args, sampling_defaults, - warmup_num_frames, + warmup_num_frames=warmup_num_frames, server_based_warmup=server_based_warmup, ) if warmup_resolutions is None @@ -448,23 +488,6 @@ def build_warmup_reqs( elif negative_prompt is not None and cfg_scale is not None and cfg_scale > 1.0: req_kwargs["do_classifier_free_guidance"] = True - if calibration_num_frames is not None and (width, height) == resolutions[0]: - # smaller-workload measurement first; it also pre-warms kernels - # so the main warmup measurement is cleaner - calibration_kwargs = req_kwargs.copy() - calibration_kwargs["sampling_params"] = copy(req_kwargs["sampling_params"]) - calibration_kwargs["num_frames"] = calibration_num_frames - calibration_req = Req(**calibration_kwargs) - calibration_req.set_as_warmup(warmup_steps) - calibration_req.sampling_params.prepare_synthetic_warmup_request_for_queue( - calibration_req, server_args - ) - if return_warmup_result: - calibration_req.extra["return_warmup_result"] = True - if server_based_warmup: - calibration_req.extra["server_based_warmup"] = True - warmup_reqs.append(calibration_req) - run_real_path_prewarm = server_based_warmup and server_args.enable_torch_compile prompts = ( (DEFAULT_PLACEHOLDER_PROMPT,) + TORCH_COMPILE_REAL_PATH_PREWARM_PROMPTS @@ -493,4 +516,46 @@ def build_warmup_reqs( req.extra["server_based_warmup"] = True warmup_reqs.append(req) + if calibration_num_frames is not None and (width, height) == resolutions[0]: + # The calibration measurement runs AFTER the main warmup request: + # one-time allocations (compile scratch, cuDNN workspaces, + # allocator pool growth) then land in the larger measurement, + # which can only steepen the fitted slope -- the safe direction. + # Measuring the small shape first flattened the slope and + # under-estimated the default workload peak. + warmup_reqs.append( + _build_calibration_warmup_req( + req_kwargs, + calibration_num_frames=calibration_num_frames, + warmup_steps=warmup_steps, + return_warmup_result=return_warmup_result, + server_based_warmup=server_based_warmup, + server_args=server_args, + ) + ) + return warmup_reqs + + +def _build_calibration_warmup_req( + req_kwargs: dict, + *, + calibration_num_frames: int, + warmup_steps: int, + return_warmup_result: bool, + server_based_warmup: bool, + server_args: ServerArgs, +) -> Req: + calibration_kwargs = req_kwargs.copy() + calibration_kwargs["sampling_params"] = copy(req_kwargs["sampling_params"]) + calibration_kwargs["num_frames"] = calibration_num_frames + calibration_req = Req(**calibration_kwargs) + calibration_req.set_as_warmup(warmup_steps) + calibration_req.sampling_params.prepare_synthetic_warmup_request_for_queue( + calibration_req, server_args + ) + if return_warmup_result: + calibration_req.extra["return_warmup_result"] = True + if server_based_warmup: + calibration_req.extra["server_based_warmup"] = True + return calibration_req diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index eddd1a048973..8e43cd4dcb8a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -13,7 +13,9 @@ ACTIVATION_EXTRAPOLATION_MARGIN, GIB_BYTES, MIN_VRAM_RESERVE_BYTES, + AppliedPromotion, AutoResidencyPlan, + AutoResidencyRollbackError, PromotionCandidate, RankResidencyReport, WarmupMemoryRecord, @@ -68,12 +70,14 @@ def test_same_shape_uses_measured_peak(self): ) assert estimate == record.peak_reserved_bytes - def test_unknown_target_uses_measured_peak(self): + def test_unknown_target_disables_estimation(self): + # An unknown target would silently equate the capped warmup peak with + # the serving peak and promote with no margin at all. record = _record() estimate = estimate_default_workload_peak_bytes( records=[record], target_units=None ) - assert estimate == record.peak_reserved_bytes + assert estimate is None def test_single_capped_record_scales_only_the_activation_part(self): # Warmup capped to 832x480x17; Wan-class default is 704x1280x121. @@ -275,6 +279,7 @@ def __init__(self, tensors: dict[str, torch.Tensor]): self.enabled = True self._configured = True self._tensors = tensors + self.fail_load = False self.load_all_layers_calls = 0 self.remove_hooks_calls = 0 self.register_hooks_calls = 0 @@ -284,8 +289,15 @@ def __init__(self, tensors: dict[str, torch.Tensor]): def iter_cpu_weights(self): yield from self._tensors.items() + def offloaded_weight_bytes(self): + return sum( + tensor.numel() * tensor.element_size() for tensor in self._tensors.values() + ) + def load_all_layers(self): self.load_all_layers_calls += 1 + if self.fail_load: + raise RuntimeError("CUDA out of memory") def remove_forward_hooks(self): self.remove_hooks_calls += 1 @@ -315,7 +327,7 @@ def __init__(self): def require_component_resident(self, component_name, *, feature_name): self.required.add(component_name) - def release_required_component_residency(self, component_name): + def release_required_component_residency(self, component_name, *, feature_name): self.required.discard(component_name) @@ -390,6 +402,21 @@ def test_filters_and_sizes(self): by_name["transformer"].promoted_weight_bytes * 50 ) + def test_mechanism_mismatch_is_not_a_candidate(self): + # A module the compile window manages layerwise while its configured + # mode says component-offload: promotion would be a silent no-op. + mismatch = _FakeLayerwiseDit( + [_FakeLayerwiseManager({"layers.0.w": torch.zeros(16)})] + ) + candidates = collect_promotion_candidates( + modules={"mismatch": mismatch}, + residency_mode_of=self._modes({"mismatch": COMPONENT_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=10, + ) + assert candidates == [] + class TestApplyAndRollback: def test_component_offload_promotion_marks_resident_without_moving(self): @@ -455,6 +482,69 @@ def test_enable_offload_does_not_double_register_hooks(self): module.enable_offload() assert manager.register_hooks_calls == 0 + def test_disable_offload_failure_rearms_the_manager(self): + # load_all_layers is the step most likely to OOM; a failed disable + # must leave the manager streaming (hooks re-registered), not in a + # hook-less enabled=True state serving (1,) placeholders. + manager = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) + manager.fail_load = True + module = _FakeLayerwiseDit([manager]) + with pytest.raises(RuntimeError, match="out of memory"): + module.disable_offload() + assert manager.enabled is True + assert manager.register_hooks_calls == 1 + assert manager.release_all_calls == 1 + + def test_failed_rollback_raises_rollback_error_with_visible_cause(self): + class _BrokenArgs(_StubResidencyArgs): + def release_required_component_residency( + self, component_name, *, feature_name + ): + # str(AssertionError()) is "" -- describe_error must keep it + # visible so no truthiness filter can drop it + raise AssertionError() + + args = _BrokenArgs() + candidates = [ + _candidate("text_encoder", weight_gib=1), + _candidate("missing_component", weight_gib=1), + ] + with pytest.raises(AutoResidencyRollbackError) as exc_info: + apply_promotions( + plan=_plan_for(candidates), + modules={"text_encoder": nn.Linear(2, 2)}, + server_args=args, + ) + assert "AssertionError" in str(exc_info.value) + assert "missing_component" in str(exc_info.value) + + def test_rollback_keeps_going_past_a_broken_component(self): + released: list[str] = [] + + class _PartialArgs(_StubResidencyArgs): + def release_required_component_residency( + self, component_name, *, feature_name + ): + if component_name == "vae": + raise RuntimeError("boom") + released.append(component_name) + + args = _PartialArgs() + applied = [ + AppliedPromotion( + component_name="text_encoder", residency_mode=COMPONENT_OFFLOAD + ), + AppliedPromotion(component_name="vae", residency_mode=COMPONENT_OFFLOAD), + ] + with pytest.raises(RuntimeError, match="vae"): + rollback_promotions( + applied=applied, + modules={}, + server_args=args, + ) + # the earlier-applied promotion was still undone + assert released == ["text_encoder"] + class TestCandidateRanking: def test_post_request_hint_order_matches_promotion_order(self): @@ -497,14 +587,20 @@ def test_reports_equivalent_server_args_and_kill_switch(self): class TestWarmupFrameAdjustment: - def _server_args(self, *, bcg: bool = False) -> SimpleNamespace: + def _server_args(self, *, bcg: bool = False, num_gpus: int = 1) -> SimpleNamespace: return SimpleNamespace( pipeline_config=LongLive2T2VConfig(), enable_breakable_cuda_graph=bcg, + num_gpus=num_gpus, ) def _defaults(self) -> SimpleNamespace: - return SimpleNamespace(num_frames=61) + return SimpleNamespace( + num_frames=61, + adjust_frames=True, + enable_sequence_shard=None, + num_frames_round_down=False, + ) def test_capped_frames_keep_the_model_frame_contract(self): # LongLive2 default 61 frames is capped to 17, whose 5 latent frames @@ -527,55 +623,112 @@ def test_non_server_warmup_keeps_default_frames(self): ) assert num_frames == 61 + def test_capped_frames_get_the_gpu_alignment_real_requests_get(self): + # a frame-aligning pipeline (no sequence shard) on multiple GPUs: + # 17 frames -> 5 latent frames -> ceil to 6 latents on 2 GPUs -> 21 + args = SimpleNamespace( + pipeline_config=SimpleNamespace( + task_type=ModelTaskType.T2V, + adjust_num_frames=lambda n: n, + vae_config=SimpleNamespace( + use_temporal_scaling_frames=True, + arch_config=SimpleNamespace(temporal_compression_ratio=4), + ), + ), + enable_breakable_cuda_graph=False, + num_gpus=2, + ) + defaults = SimpleNamespace( + num_frames=81, + adjust_frames=True, + enable_sequence_shard=None, + num_frames_round_down=False, + ) + num_frames = _resolve_warmup_num_frames( + args, defaults, server_based_warmup=True + ) + assert num_frames == 21 + class TestCalibrationFrames: - def _wan_like_args(self, *, performance_mode: str = "auto") -> SimpleNamespace: + def _patch_gate(self, monkeypatch, reason: str | None = None) -> None: + monkeypatch.setattr( + "sglang.multimodal_gen.runtime.server_warmup.auto_residency_skip_reason", + lambda _args: reason, + ) + + def _wan_like_args(self) -> SimpleNamespace: return SimpleNamespace( pipeline_config=SimpleNamespace( task_type=ModelTaskType.T2V, adjust_num_frames=lambda n: n, ), - enable_breakable_cuda_graph=False, - performance_mode=performance_mode, + num_gpus=1, ) - def test_capped_video_gets_a_smaller_calibration_size(self): + def _defaults(self, num_frames: int) -> SimpleNamespace: + return SimpleNamespace( + num_frames=num_frames, + adjust_frames=True, + enable_sequence_shard=None, + num_frames_round_down=False, + ) + + def test_capped_video_gets_a_smaller_calibration_size(self, monkeypatch): + self._patch_gate(monkeypatch) calibration = _resolve_calibration_num_frames( self._wan_like_args(), - SimpleNamespace(num_frames=81), - 17, + self._defaults(81), + warmup_num_frames=17, server_based_warmup=True, ) assert calibration == 9 - def test_uncapped_video_needs_no_calibration(self): + def test_uncapped_video_needs_no_calibration(self, monkeypatch): + self._patch_gate(monkeypatch) calibration = _resolve_calibration_num_frames( self._wan_like_args(), - SimpleNamespace(num_frames=13), - 13, + self._defaults(13), + warmup_num_frames=13, server_based_warmup=True, ) assert calibration is None - def test_manual_mode_skips_calibration(self): + def test_unknown_warmup_frames_skip_calibration(self, monkeypatch): + self._patch_gate(monkeypatch) calibration = _resolve_calibration_num_frames( - self._wan_like_args(performance_mode="manual"), - SimpleNamespace(num_frames=81), - 17, + self._wan_like_args(), + self._defaults(81), + warmup_num_frames=None, server_based_warmup=True, ) assert calibration is None - def test_frame_contract_collapse_skips_calibration(self): + def test_skip_gate_disables_calibration(self, monkeypatch): + # the calibration request costs a full extra warmup forward: it must + # share the promotion's own gate (kill switch, quantized, manual, ...) + self._patch_gate(monkeypatch, reason="performance_mode=manual") + calibration = _resolve_calibration_num_frames( + self._wan_like_args(), + self._defaults(81), + warmup_num_frames=17, + server_based_warmup=True, + ) + assert calibration is None + + def test_frame_contract_collapse_skips_calibration(self, monkeypatch): # LongLive2 aligns both 17 and 9 to 29 latent-block frames: a second # measurement at the same size adds nothing. + self._patch_gate(monkeypatch) args = SimpleNamespace( pipeline_config=LongLive2T2VConfig(), - enable_breakable_cuda_graph=False, - performance_mode="auto", + num_gpus=1, ) calibration = _resolve_calibration_num_frames( - args, SimpleNamespace(num_frames=61), 29, server_based_warmup=True + args, + self._defaults(61), + warmup_num_frames=29, + server_based_warmup=True, ) assert calibration is None @@ -588,6 +741,7 @@ def _base_args(self, **overrides) -> SimpleNamespace: warmup_resolutions=None, backend="sglang", enable_breakable_cuda_graph=False, + enable_torch_compile=False, batching_max_size=1, dp_size=1, use_fsdp_inference=False, @@ -625,6 +779,9 @@ def test_manual_performance_mode(self, monkeypatch): ({"warmup_mode": "request"}, "server warmup"), ({"backend": "diffusers"}, "diffusers"), ({"enable_breakable_cuda_graph": True}, "CUDA graph"), + # compile warmup strips the memory layout (layerwise DiT + + # resident aux components on CPU): its peaks are not serving peaks + ({"enable_torch_compile": True}, "stripped memory layout"), ({"batching_max_size": 4}, "batching"), ({"dp_size": 2}, "dp replicas"), ({"use_fsdp_inference": True}, "FSDP"), From 17b51bfadfa9c898c5da4045849cb74ebc6dbf1f Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 19 Aug 2026 15:55:52 +0800 Subject: [PATCH 06/90] [diffusion] extract control-request protocol into control_requests.py The cross-process control requests (LoRA ops, shutdown, realtime release, disagg stats, auto residency) are IPC contracts between the HTTP process and the scheduler workers, not utilities; the grab-bag entrypoints/utils.py hid that. Pure move: class bodies are unchanged, the new module stays import-light because both processes load it, and openai/utils.py keeps re-exporting the LoRA protocol names. Co-Authored-By: Claude Fable 5 --- .../runtime/entrypoints/control_requests.py | 68 +++++++++++++++++++ .../entrypoints/diffusion_generator.py | 6 +- .../runtime/entrypoints/http_server.py | 4 +- .../runtime/entrypoints/openai/common_api.py | 4 +- .../openai/realtime/realtime_video_api.py | 6 +- .../runtime/entrypoints/openai/utils.py | 4 +- .../runtime/entrypoints/utils.py | 54 --------------- .../multimodal_gen/runtime/launch_server.py | 2 +- .../runtime/managers/scheduler.py | 18 ++--- .../runtime/scheduler_client.py | 16 ++--- .../multimodal_gen/runtime/server_warmup.py | 4 +- .../test/unit/test_cfg_parallel_warmup.py | 4 +- .../unit/test_diffusion_generator_shutdown.py | 2 +- .../test/unit/test_dp_routing.py | 5 +- .../test/unit/test_launch_server_shutdown.py | 2 +- 15 files changed, 113 insertions(+), 86 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py b/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py new file mode 100644 index 000000000000..fb8adde8e5f2 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Control-request protocol between the HTTP process and scheduler workers. + +These types are cross-process IPC contracts, not utilities: the HTTP side +constructs them (scheduler_client treats them as ``_CONTROL_REQ_TYPES`` and +fans them out to every replica) and each scheduler dispatches them through +``Scheduler.request_handlers``. Keep this module import-light -- both +processes import it, and the HTTP process must not drag in torch-heavy +worker modules through it. +""" + +from dataclasses import dataclass +from typing import List, Optional, Union + +import msgspec + + +@dataclass +class SetLoraReq: + lora_nickname: Union[str, List[str]] + lora_path: Optional[Union[str, List[Optional[str]]]] = None + target: Union[str, List[str]] = "all" + strength: Union[float, List[float]] = 1.0 + merge_mode: Optional[str] = None + lora_alpha: Optional[Union[int, List[Optional[int]]]] = None + + +@dataclass +class MergeLoraWeightsReq: + target: str = "all" + strength: float = 1.0 + + +@dataclass +class UnmergeLoraWeightsReq: + target: str = "all" + + +@dataclass +class ListLorasReq: + pass + + +@dataclass +class ShutdownReq: + pass + + +@dataclass +class ReleaseRealtimeSessionReq: + session_id: str + + +@dataclass +class GetDisaggStatsReq: + """Request to get disagg pipeline metrics from the scheduler.""" + + pass + + +class AutoResidencyReq(msgspec.Struct, frozen=True): + """Apply or roll back the warmup-calibrated residency promotion. + + Sent by the server warmup orchestrator after the synthetic warmup + requests finish and before the server reports ready. + """ + + action: str = "apply" # "apply" | "rollback" diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 0523f5650625..712ba7d1bea6 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -19,13 +19,15 @@ DataType, SamplingParams, ) -from sglang.multimodal_gen.runtime.entrypoints.utils import ( - GenerationResult, +from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( ListLorasReq, MergeLoraWeightsReq, SetLoraReq, ShutdownReq, UnmergeLoraWeightsReq, +) +from sglang.multimodal_gen.runtime.entrypoints.utils import ( + GenerationResult, expand_request_outputs, format_lora_message, prepare_request, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index bd4da6eee808..42a73d2b22f4 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -268,7 +268,9 @@ async def stats_endpoint(request: Request): Returns queue depth, request counts, latency, throughput, etc. Sends a GetDisaggStatsReq to the scheduler via ZMQ and returns the result. """ - from sglang.multimodal_gen.runtime.entrypoints.utils import GetDisaggStatsReq + from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( + GetDisaggStatsReq, + ) server_args: ServerArgs = request.app.state.server_args response: dict = { diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py index 9a3e6d83116f..d3c9234a1643 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py @@ -5,11 +5,13 @@ from pydantic import BaseModel, Field from sglang.multimodal_gen.registry import get_model_info -from sglang.multimodal_gen.runtime.entrypoints.utils import ( +from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( ListLorasReq, MergeLoraWeightsReq, SetLoraReq, UnmergeLoraWeightsReq, +) +from sglang.multimodal_gen.runtime.entrypoints.utils import ( format_lora_message, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_video_api.py index 305a40818912..e8573a8ff3fa 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_video_api.py @@ -8,6 +8,9 @@ import msgspec.msgpack from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( + ReleaseRealtimeSessionReq, +) from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( RealtimeEvent, RealtimeVideoGenerationsRequest, @@ -28,9 +31,6 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( process_generation_batch, ) -from sglang.multimodal_gen.runtime.entrypoints.utils import ( - ReleaseRealtimeSessionReq, -) from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.multimodal_gen.runtime.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index 6b5ef964ddfe..7d3338a779a1 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py @@ -16,12 +16,14 @@ DataType, SamplingParams, ) -from sglang.multimodal_gen.runtime.entrypoints.utils import ( +from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( ListLorasReq, MergeLoraWeightsReq, SetLoraReq, ShutdownReq, UnmergeLoraWeightsReq, +) +from sglang.multimodal_gen.runtime.entrypoints.utils import ( format_lora_message, save_outputs, ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index 73c21225c18a..329220c6a2cc 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -23,7 +23,6 @@ from typing import Any, Callable, List, Optional, Sequence, Union import imageio -import msgspec import numpy as np import torch from PIL import Image @@ -160,59 +159,6 @@ def _close_cached_cuda_video_buffer() -> None: atexit.register(_close_cached_cuda_video_buffer) -@dataclass -class SetLoraReq: - lora_nickname: Union[str, List[str]] - lora_path: Optional[Union[str, List[Optional[str]]]] = None - target: Union[str, List[str]] = "all" - strength: Union[float, List[float]] = 1.0 - merge_mode: Optional[str] = None - lora_alpha: Optional[Union[int, List[Optional[int]]]] = None - - -@dataclass -class MergeLoraWeightsReq: - target: str = "all" - strength: float = 1.0 - - -@dataclass -class UnmergeLoraWeightsReq: - target: str = "all" - - -@dataclass -class ListLorasReq: - pass - - -@dataclass -class ShutdownReq: - pass - - -@dataclass -class ReleaseRealtimeSessionReq: - session_id: str - - -@dataclass -class GetDisaggStatsReq: - """Request to get disagg pipeline metrics from the scheduler.""" - - pass - - -class AutoResidencyReq(msgspec.Struct, frozen=True): - """Apply or roll back the warmup-calibrated residency promotion. - - Sent by the server warmup orchestrator after the synthetic warmup - requests finish and before the server reports ready. - """ - - action: str = "apply" # "apply" | "rollback" - - def format_lora_message( lora_nickname: Union[str, List[str]], target: Union[str, List[str]], diff --git a/python/sglang/multimodal_gen/runtime/launch_server.py b/python/sglang/multimodal_gen/runtime/launch_server.py index af3d09c4cfca..c1b8b38f42c4 100644 --- a/python/sglang/multimodal_gen/runtime/launch_server.py +++ b/python/sglang/multimodal_gen/runtime/launch_server.py @@ -15,8 +15,8 @@ DiffusionServer, ) from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType +from sglang.multimodal_gen.runtime.entrypoints.control_requests import ShutdownReq from sglang.multimodal_gen.runtime.entrypoints.http_server import create_app -from sglang.multimodal_gen.runtime.entrypoints.utils import ShutdownReq from sglang.multimodal_gen.runtime.managers.gpu_worker import run_scheduler_process from sglang.multimodal_gen.runtime.scheduler_client import SchedulerClient from sglang.multimodal_gen.runtime.server_args import ( diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 0d1f0e5f9d17..38158348bf4f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -16,15 +16,7 @@ from sglang.multimodal_gen.runtime.disaggregation.scheduler_mixin import ( SchedulerDisaggMixin, ) -from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( - GetWeightsChecksumReqInput, - ReleaseMemoryOccupationReqInput, - ResumeMemoryOccupationReqInput, - UpdateWeightFromDiskReqInput, - UpdateWeightFromTensorCheckerReqInput, - UpdateWeightFromTensorReqInput, -) -from sglang.multimodal_gen.runtime.entrypoints.utils import ( +from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( AutoResidencyReq, GetDisaggStatsReq, ListLorasReq, @@ -34,6 +26,14 @@ ShutdownReq, UnmergeLoraWeightsReq, ) +from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( + GetWeightsChecksumReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightFromTensorCheckerReqInput, + UpdateWeightFromTensorReqInput, +) from sglang.multimodal_gen.runtime.ipc_array import ( is_local_endpoint, spill_large_arrays_to_file_refs, diff --git a/python/sglang/multimodal_gen/runtime/scheduler_client.py b/python/sglang/multimodal_gen/runtime/scheduler_client.py index 423c68d7bf25..1aade5875ec0 100644 --- a/python/sglang/multimodal_gen/runtime/scheduler_client.py +++ b/python/sglang/multimodal_gen/runtime/scheduler_client.py @@ -7,6 +7,14 @@ import zmq import zmq.asyncio +from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( + AutoResidencyReq, + ListLorasReq, + MergeLoraWeightsReq, + SetLoraReq, + ShutdownReq, + UnmergeLoraWeightsReq, +) from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( GetWeightsChecksumReqInput, ReleaseMemoryOccupationReqInput, @@ -15,14 +23,6 @@ UpdateWeightFromTensorCheckerReqInput, UpdateWeightFromTensorReqInput, ) -from sglang.multimodal_gen.runtime.entrypoints.utils import ( - AutoResidencyReq, - ListLorasReq, - MergeLoraWeightsReq, - SetLoraReq, - ShutdownReq, - UnmergeLoraWeightsReq, -) from sglang.multimodal_gen.runtime.ipc_array import materialize_file_refs from sglang.multimodal_gen.runtime.pipelines_core import Req from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index 26e90c107a3e..7e9017114a11 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -176,7 +176,9 @@ async def maybe_apply_auto_residency( continues on the original strategy; only a failed rollback raises, which aborts startup. """ - from sglang.multimodal_gen.runtime.entrypoints.utils import AutoResidencyReq + from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( + AutoResidencyReq, + ) from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( PROMOTION_STATUS_PROMOTED, PROMOTION_STATUS_ROLLBACK_FAILED, diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index 834039c5b319..02d0c905fae2 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py @@ -29,11 +29,11 @@ ) from sglang.multimodal_gen.configs.sample.longlive2 import LongLive2SamplingParams from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams -from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator -from sglang.multimodal_gen.runtime.entrypoints.utils import ( +from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( SetLoraReq, UnmergeLoraWeightsReq, ) +from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator from sglang.multimodal_gen.runtime.managers.scheduler import Scheduler from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import ( OutputBatch, diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_generator_shutdown.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_generator_shutdown.py index e43f550dd91c..d00f9e43647a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_diffusion_generator_shutdown.py +++ b/python/sglang/multimodal_gen/test/unit/test_diffusion_generator_shutdown.py @@ -3,8 +3,8 @@ from unittest.mock import Mock, patch from sglang.multimodal_gen.runtime.entrypoints import diffusion_generator as dg +from sglang.multimodal_gen.runtime.entrypoints.control_requests import ShutdownReq from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator -from sglang.multimodal_gen.runtime.entrypoints.utils import ShutdownReq class _FakeProcess: diff --git a/python/sglang/multimodal_gen/test/unit/test_dp_routing.py b/python/sglang/multimodal_gen/test/unit/test_dp_routing.py index ad0593faa396..32bcd43991db 100644 --- a/python/sglang/multimodal_gen/test/unit/test_dp_routing.py +++ b/python/sglang/multimodal_gen/test/unit/test_dp_routing.py @@ -6,7 +6,10 @@ import itertools -from sglang.multimodal_gen.runtime.entrypoints.utils import SetLoraReq, ShutdownReq +from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( + SetLoraReq, + ShutdownReq, +) from sglang.multimodal_gen.runtime.pipelines_core import Req from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.scheduler_client import ( diff --git a/python/sglang/multimodal_gen/test/unit/test_launch_server_shutdown.py b/python/sglang/multimodal_gen/test/unit/test_launch_server_shutdown.py index 308ce5b50d72..0a3949acad82 100644 --- a/python/sglang/multimodal_gen/test/unit/test_launch_server_shutdown.py +++ b/python/sglang/multimodal_gen/test/unit/test_launch_server_shutdown.py @@ -4,7 +4,7 @@ from sglang.multimodal_gen.runtime import launch_server as ls from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType -from sglang.multimodal_gen.runtime.entrypoints.utils import ShutdownReq +from sglang.multimodal_gen.runtime.entrypoints.control_requests import ShutdownReq class _FakeProcess: From 7bd0b5cacdadaf943073018e20d9bcd065484544 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 19 Aug 2026 15:56:50 +0800 Subject: [PATCH 07/90] [diffusion] migrate control requests from dataclass to msgspec.Struct Opportunistic migration per the no-dataclasses rule while the classes moved files: same fields, same defaults, positional construction still supported. No dataclasses.asdict/fields/replace call sites exist for these types, and msgspec Structs pickle across the existing broadcast_pyobj/ZMQ transport (AutoResidencyReq already proved the path). Co-Authored-By: Claude Fable 5 --- .../runtime/entrypoints/control_requests.py | 22 ++++++------------- 1 file changed, 7 insertions(+), 15 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py b/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py index fb8adde8e5f2..9df19c7f8969 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py @@ -9,14 +9,12 @@ worker modules through it. """ -from dataclasses import dataclass from typing import List, Optional, Union import msgspec -@dataclass -class SetLoraReq: +class SetLoraReq(msgspec.Struct): lora_nickname: Union[str, List[str]] lora_path: Optional[Union[str, List[Optional[str]]]] = None target: Union[str, List[str]] = "all" @@ -25,34 +23,28 @@ class SetLoraReq: lora_alpha: Optional[Union[int, List[Optional[int]]]] = None -@dataclass -class MergeLoraWeightsReq: +class MergeLoraWeightsReq(msgspec.Struct): target: str = "all" strength: float = 1.0 -@dataclass -class UnmergeLoraWeightsReq: +class UnmergeLoraWeightsReq(msgspec.Struct): target: str = "all" -@dataclass -class ListLorasReq: +class ListLorasReq(msgspec.Struct): pass -@dataclass -class ShutdownReq: +class ShutdownReq(msgspec.Struct): pass -@dataclass -class ReleaseRealtimeSessionReq: +class ReleaseRealtimeSessionReq(msgspec.Struct): session_id: str -@dataclass -class GetDisaggStatsReq: +class GetDisaggStatsReq(msgspec.Struct): """Request to get disagg pipeline metrics from the scheduler.""" pass From bfe142bf3610284941ec54d2dcef0cac96d6d743 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 19 Aug 2026 03:07:10 +0800 Subject: [PATCH 08/90] fix: preserve single-GPU warmup frame contract --- .../multimodal_gen/runtime/warmup_request_builder.py | 8 ++++++-- .../multimodal_gen/test/unit/test_cfg_parallel_warmup.py | 7 +++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py index 726da63fab78..229b2654920c 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -282,8 +282,12 @@ def _apply_warmup_frame_contract( the same latent alignment real requests get. """ num_frames = server_args.pipeline_config.adjust_num_frames(num_frames) - if sampling_defaults.adjust_frames and not resolve_sequence_shard( - server_args.pipeline_config, sampling_defaults.enable_sequence_shard + if ( + sampling_defaults.adjust_frames + and not resolve_sequence_shard( + server_args.pipeline_config, sampling_defaults.enable_sequence_shard + ) + and server_args.num_gpus > 1 ): num_frames = align_num_frames_for_num_gpus( num_frames, diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index 02d0c905fae2..a6c10fb2cc03 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py @@ -270,6 +270,7 @@ def test_diff_generator_runs_explicit_warmup_through_scheduler_client(self): server_args.warmup_steps = 1 server_args.enable_cfg_parallel = False server_args.enable_torch_compile = False + server_args.num_gpus = 1 server_args.pipeline_config.task_type = ModelTaskType.T2V server_args.pipeline_config.adjust_num_frames.side_effect = lambda value: value @@ -362,6 +363,7 @@ def test_server_based_warmup_resolutions_keep_sampling_defaults_and_caps(self): server_args.warmup_steps = 1 server_args.enable_cfg_parallel = False server_args.enable_torch_compile = False + server_args.num_gpus = 1 server_args.pipeline_config.task_type = ModelTaskType.T2V server_args.pipeline_config.adjust_num_frames.side_effect = lambda value: value @@ -474,6 +476,7 @@ def test_server_based_warmup_keeps_video_warmup_lightweight(self): server_args.warmup_steps = 1 server_args.enable_cfg_parallel = False server_args.enable_torch_compile = False + server_args.num_gpus = 1 server_args.pipeline_config.task_type = ModelTaskType.T2V server_args.pipeline_config.adjust_num_frames.side_effect = lambda value: value @@ -503,6 +506,7 @@ def test_video_warmup_preserves_model_frame_alignment(self): pipeline_config=pipeline_config, enable_breakable_cuda_graph=False, pipeline_class_name=None, + num_gpus=1, ) num_frames = _resolve_warmup_num_frames( @@ -524,6 +528,7 @@ def test_server_based_warmup_uses_video_supported_resolution_budget(self): server_args.warmup_steps = 1 server_args.enable_cfg_parallel = False server_args.enable_torch_compile = False + server_args.num_gpus = 1 server_args.pipeline_config.task_type = ModelTaskType.T2V server_args.pipeline_config.adjust_num_frames.side_effect = lambda value: value @@ -560,6 +565,7 @@ def test_ltx2_two_stage_warmup_uses_pipeline_alignment(self): server_args.warmup_steps = 1 server_args.enable_cfg_parallel = False server_args.enable_torch_compile = False + server_args.num_gpus = 1 server_args.pipeline_class_name = "LTX2TwoStageHQPipeline" server_args.pipeline_config.task_type = ModelTaskType.T2V @@ -776,6 +782,7 @@ def test_server_based_warmup_keeps_ti2v_image_input(self): server_args.warmup_steps = 1 server_args.enable_cfg_parallel = False server_args.enable_torch_compile = False + server_args.num_gpus = 1 server_args.pipeline_config.task_type = ModelTaskType.TI2V with patch( From e9db98a09812e6a8ce5c70edfc94ddb72dc65a96 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 19 Aug 2026 07:05:28 +0800 Subject: [PATCH 09/90] fix: skip auto residency with Ulysses parallelism --- python/sglang/multimodal_gen/runtime/server_warmup.py | 5 +++++ .../sglang/multimodal_gen/test/unit/test_auto_residency.py | 2 ++ 2 files changed, 7 insertions(+) diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index 7e9017114a11..9085585ef191 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -138,6 +138,11 @@ def auto_residency_skip_reason(server_args: ServerArgs) -> str | None: return "dynamic batching enabled" if (server_args.dp_size or 1) > 1: return "dp replicas warm up independently" + if (server_args.ulysses_degree or 1) > 1: + # Sequence-parallel workers can have a different execution order + # after residency changes; keep their configured placement until that + # interaction is explicitly calibrated and validated. + return "Ulysses sequence parallelism" if server_args.use_fsdp_inference: return "FSDP inference" nunchaku = server_args.nunchaku_config diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 8e43cd4dcb8a..19ea25735ac0 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -744,6 +744,7 @@ def _base_args(self, **overrides) -> SimpleNamespace: enable_torch_compile=False, batching_max_size=1, dp_size=1, + ulysses_degree=1, use_fsdp_inference=False, quantization=None, transformer_weights_path=None, @@ -784,6 +785,7 @@ def test_manual_performance_mode(self, monkeypatch): ({"enable_torch_compile": True}, "stripped memory layout"), ({"batching_max_size": 4}, "batching"), ({"dp_size": 2}, "dp replicas"), + ({"ulysses_degree": 2}, "Ulysses"), ({"use_fsdp_inference": True}, "FSDP"), ({"quantization": "fp8"}, "quantized"), ({"transformer_weights_path": "/x.safetensors"}, "quantized"), From 274e9f697920708e812148302d690cfab55f9d31 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 20 Aug 2026 09:58:56 +0800 Subject: [PATCH 10/90] [diffusion] fix: make auto residency reachable on a small card Running this on a 12 GiB RTX 3060 it never promoted anything, for three separate reasons. A warmup probe that did not fit aborted the whole warmup. The synthetic pass sends two probes; when 832x480x17f ran out of memory the loop raised, later probes never ran, and the failed record disabled the estimate. Retry the probe about half as large instead -- frames first, since they drive video activation size and cutting them leaves the spatial kernels at their serving shape -- and give up after three attempts so a failure that is not about probe size does not walk the workload down to nothing. Only memory errors are retried; anything else fails the same way at every size and should surface. A failed probe then blocked the estimate outright. That is right when the probe is at or below the target -- the card cannot hold the default workload as it is already configured, and making weights resident only adds to it -- but a probe that failed above the target says nothing about the target. Keep the block for the first case, with a skip reason that names the size that failed instead of "no usable warmup measurement", and drop the record in the second. The 4 GiB reserve floor only binds below ~40 GiB. Above that the 10% fraction dominates (10% of 80 GiB is 8 GiB), so the floor only ever applies to the cards it was not sized for: on a 12 GiB device a flat 4 GiB fences off a third of the card. Cap the floor at a fifth of the budget, which leaves datacenter behaviour unchanged. Finally, the gate itself was invisible: `--performance-mode auto` with any other warmup mode is skipped, and only at debug level, so the user who asked for auto saw nothing at all. Log it at info when auto was requested. --- .../runtime/managers/gpu_worker.py | 17 +++++ .../memory_managers/auto_residency.py | 32 +++++++--- .../multimodal_gen/runtime/server_warmup.py | 50 ++++++++++++++- .../runtime/warmup_request_builder.py | 51 +++++++++++++++ .../test/unit/test_auto_residency.py | 39 +++++++++++- .../unit/test_warmup_probe_degradation.py | 62 +++++++++++++++++++ 6 files changed, 240 insertions(+), 11 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_warmup_probe_degradation.py diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 40ba239e7897..c6303558c177 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1314,6 +1314,23 @@ def _build_auto_residency_report( skip_reason = "no server warmup measurements" elif workload.workload_units() is None: skip_reason = "default workload resolution unknown" + else: + # A probe that failed at or below the target is a measurement, not + # missing data: the card cannot hold the default workload as it is + # already configured, and promoting weights would only add to it. + blocking = [ + record + for record in records + if not record.succeeded + and record.workload_units() <= workload.workload_units() + ] + if blocking: + smallest = min(blocking, key=lambda record: record.workload_units()) + skip_reason = ( + f"warmup failed at {smallest.width}x{smallest.height}" + f"x{smallest.num_frames}f, which the default workload " + f"{workload.describe()} meets or exceeds" + ) if skip_reason is not None: return RankResidencyReport( rank=self.rank, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index d5082e204b78..ea9de2436ad0 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -56,6 +56,11 @@ # shape variance, and CUDA graph or compile pools. VRAM_RESERVE_FRACTION = 0.10 MIN_VRAM_RESERVE_BYTES = 4 * GIB_BYTES +# The absolute floor is sized for datacenter cards, where the fraction dominates +# anyway (10% of 80 GiB is 8 GiB). Below ~40 GiB the floor takes over, and on a +# 12 GiB card a flat 4 GiB would fence off a third of the device, so cap it as a +# share of what is actually there. +MAX_VRAM_RESERVE_FRACTION = 0.20 AUTO_RESIDENCY_FEATURE_NAME = "auto residency promotion" @@ -197,16 +202,26 @@ def estimate_default_workload_peak_bytes( 3. One usable size: scale everything above the pre-forward allocated baseline (conservative; may block promotion but never over-promotes). - Returns None when the estimate cannot be trusted: no records, any warmup - forward failed (its peak does not cover the full request path), or the - target workload is unknown (an unknown target would silently equate the - area/frame-capped warmup peak with the real serving peak). + Returns None when the estimate cannot be trusted: no successful records, + the target workload is unknown (an unknown target would silently equate the + area/frame-capped warmup peak with the real serving peak), or a probe at or + below the target ran out of memory. That last case is not missing data but + a measurement: the card could not hold the target as it is already + configured, so making more weights resident can only make it worse. A probe + that failed strictly above the target says nothing about the target and is + dropped instead. """ records = list(records) - if not records or any(not record.succeeded for record in records): - return None if target_units is None: return None + failed_units = [ + record.workload_units() for record in records if not record.succeeded + ] + if any(units <= target_units for units in failed_units): + return None + records = [record for record in records if record.succeeded] + if not records: + return None peak_by_units: dict[int, int] = {} for record in records: @@ -407,7 +422,10 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP estimated_peak = max(report.estimated_peak_bytes for report in reports) budget = min(report.budget_bytes for report in reports) - reserve = max(int(budget * VRAM_RESERVE_FRACTION), MIN_VRAM_RESERVE_BYTES) + reserve = max( + int(budget * VRAM_RESERVE_FRACTION), + min(MIN_VRAM_RESERVE_BYTES, int(budget * MAX_VRAM_RESERVE_FRACTION)), + ) candidates = _consensus_candidates(reports) if not candidates: diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index 9085585ef191..82df3e9a7ed0 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -16,6 +16,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.warmup_request_builder import ( build_warmup_reqs, + lighten_warmup_req, should_include_warmup_image, supports_synthetic_warmup, ) @@ -192,7 +193,10 @@ async def maybe_apply_auto_residency( skip_reason = auto_residency_skip_reason(server_args) if skip_reason is not None: - logger.debug("Auto residency: skipped (%s)", skip_reason) + # Whoever asked for auto needs to hear why it did nothing; on any other + # server this is just a note about a mode they did not pick. + log = logger.info if server_args.performance_mode == "auto" else logger.debug + log("Auto residency: skipped (%s)", skip_reason) return logger.info( @@ -252,6 +256,34 @@ async def maybe_apply_auto_residency( ) +# Enough to clear a probe that overshot the card, few enough that a failure +# which is not about probe size gives up quickly instead of walking the +# workload down to nothing. +MAX_WARMUP_DEGRADE_ATTEMPTS = 3 + + +def _is_out_of_memory(error: Any) -> bool: + text = str(error).lower() + return "out of memory" in text or "outofmemory" in text + + +def _degrade_after_oom(server_args: ServerArgs, req: Req) -> Req | None: + """Next warmup probe to try after `req` ran the card out of memory. + + Only memory failures are worth retrying smaller; anything else fails the + same way at every size and should surface instead of being shrunk away. + """ + lighter = lighten_warmup_req(server_args, req) + if lighter is None: + return None + logger.warning( + "%s ran out of memory; retrying warmup at %s", + format_warmup_req(req), + format_warmup_req(lighter), + ) + return lighter + + def format_warmup_req(req_or_group: Any) -> str: req = get_first_generation_req(req_or_group) prefix = ( @@ -318,6 +350,14 @@ async def run_async_client_warmup( server_args, warmup_input_path=warmup_input_path, rewarm=rewarm ): response = await forward(req) + for _ in range(MAX_WARMUP_DEGRADE_ATTEMPTS): + if response.error is None or not _is_out_of_memory(response.error): + break + lighter = _degrade_after_oom(server_args, req) + if lighter is None: + break + req = lighter + response = await forward(req) if response.error is not None: raise RuntimeError(response.error) except Exception: @@ -341,6 +381,14 @@ def run_sync_client_warmup( server_args, warmup_input_path=warmup_input_path ): response = forward(req) + for _ in range(MAX_WARMUP_DEGRADE_ATTEMPTS): + if response.error is None or not _is_out_of_memory(response.error): + break + lighter = _degrade_after_oom(server_args, req) + if lighter is None: + break + req = lighter + response = forward(req) if response.error is not None: raise RuntimeError(response.error) diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py index 229b2654920c..e473840ea2f6 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -14,6 +14,7 @@ """ from copy import copy +from dataclasses import replace from typing import Any from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType @@ -229,6 +230,56 @@ def _fit_resolution_to_area( ) +def _halve_num_frames(server_args: ServerArgs, num_frames: int) -> int: + """Halve the latent frame count, keeping the model's frame arithmetic.""" + if num_frames <= 1: + return num_frames + arch_config = getattr( + getattr(server_args.pipeline_config, "vae_config", None), "arch_config", None + ) + ratio = getattr(arch_config, "temporal_compression_ratio", None) + if not ratio or ratio <= 1: + return max(1, num_frames // 2) + latent_frames = (num_frames - 1) // ratio + 1 + # round up: halving 5 latent frames should land on 3 (9 frames), not 2 (5) + return ((latent_frames + 1) // 2 - 1) * ratio + 1 + + +def lighten_warmup_req(server_args: ServerArgs, req: Req) -> Req | None: + """Roughly halve a warmup probe, or None once it cannot shrink further. + + Warmup peak memory tracks width * height * num_frames, so a card that could + not hold the full probe usually holds half of it while still walking the + same code path. Frames go first: they drive video activation size, and + cutting them leaves the spatial kernels at their serving shape. + """ + params = req.sampling_params + if params is None: + return None + + num_frames = getattr(params, "num_frames", None) or 1 + lighter_frames = _halve_num_frames(server_args, num_frames) + if lighter_frames < num_frames: + return _replace_warmup_workload(req, num_frames=lighter_frames) + + width = getattr(params, "width", None) + height = getattr(params, "height", None) + if not width or not height: + return None + alignment = _warmup_resolution_alignment(server_args) + lighter = _fit_resolution_to_area(width, height, width * height // 2, alignment) + # below the alignment floor the fit rounds back up; only take a real cut + if lighter[0] * lighter[1] >= width * height: + return None + return _replace_warmup_workload(req, width=lighter[0], height=lighter[1]) + + +def _replace_warmup_workload(req: Req, **changes: int) -> Req: + lighter = copy(req) + lighter.sampling_params = replace(req.sampling_params, **changes) + return lighter + + def _is_resolution_aligned(resolution: tuple[int, int], alignment: int) -> bool: width, height = resolution return width % alignment == 0 and height % alignment == 0 diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 19ea25735ac0..a65a7b2ebde4 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -156,10 +156,37 @@ def test_multiple_records_take_the_max(self): ) assert estimate == high.peak_reserved_bytes - def test_failed_warmup_disables_estimation(self): - records = [_record(), _record(succeeded=False)] + def test_failure_at_the_target_size_disables_estimation(self): + good = _record(num_frames=9, peak_gib=8) + failed = _record(num_frames=17, succeeded=False) assert ( - estimate_default_workload_peak_bytes(records=records, target_units=None) + estimate_default_workload_peak_bytes( + records=[good, failed], target_units=failed.workload_units() + ) + is None + ) + + def test_failure_below_the_target_size_disables_estimation(self): + good = _record(num_frames=9, peak_gib=8) + failed = _record(num_frames=17, succeeded=False) + assert ( + estimate_default_workload_peak_bytes( + records=[good, failed], target_units=failed.workload_units() * 2 + ) + is None + ) + + def test_failure_above_the_target_size_is_dropped(self): + good = _record(num_frames=9, peak_gib=8) + failed = _record(num_frames=81, succeeded=False) + estimate = estimate_default_workload_peak_bytes( + records=[good, failed], target_units=good.workload_units() + ) + assert estimate == good.peak_reserved_bytes + + def test_unknown_target_disables_estimation(self): + assert ( + estimate_default_workload_peak_bytes(records=[_record()], target_units=None) is None ) @@ -231,6 +258,12 @@ def test_reserve_has_absolute_floor(self): plan = plan_auto_residency(reports=[_report(budget_gib=20, estimated_gib=10)]) assert plan.reserve_bytes == MIN_VRAM_RESERVE_BYTES + def test_reserve_floor_is_capped_on_a_small_card(self): + # A flat 4 GiB would fence off a third of a 12 GiB card; the cap keeps + # the floor at a fifth of what is actually there. + plan = plan_auto_residency(reports=[_report(budget_gib=10, estimated_gib=4)]) + assert plan.reserve_bytes == 2 * GIB_BYTES + def test_greedy_promotion_by_h2d_savings(self): # dit saves 25 GiB x 50 steps per request; text_encoder saves 30 GiB # once. dit must be promoted first, after which the text encoder no diff --git a/python/sglang/multimodal_gen/test/unit/test_warmup_probe_degradation.py b/python/sglang/multimodal_gen/test/unit/test_warmup_probe_degradation.py new file mode 100644 index 000000000000..d23471e3d2c6 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_warmup_probe_degradation.py @@ -0,0 +1,62 @@ +"""A warmup probe that does not fit is retried smaller instead of abandoned.""" + +from types import SimpleNamespace + +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.warmup_request_builder import lighten_warmup_req + + +def _server_args(temporal_compression_ratio: int = 4) -> SimpleNamespace: + arch_config = SimpleNamespace( + temporal_compression_ratio=temporal_compression_ratio, + vae_scale_factor=8, + spatial_compression_ratio=8, + ) + return SimpleNamespace( + pipeline_class_name=None, + pipeline_config=SimpleNamespace( + vae_config=SimpleNamespace(arch_config=arch_config), + vae_scale_factor=8, + ), + ) + + +def _req(width: int, height: int, num_frames: int) -> Req: + return Req( + sampling_params=SamplingParams( + width=width, height=height, num_frames=num_frames + ) + ) + + +class TestLightenWarmupReq: + def test_video_probe_halves_latent_frames_first(self): + lighter = lighten_warmup_req(_server_args(), _req(832, 480, 17)) + assert (lighter.width, lighter.height) == (832, 480) + assert lighter.num_frames == 9 + + def test_frames_step_down_to_a_single_frame(self): + server_args = _server_args() + frames = [] + req = _req(832, 480, 17) + for _ in range(4): + req = lighten_warmup_req(server_args, req) + if req is None: + break + frames.append(req.num_frames) + assert frames[:3] == [9, 5, 1] + + def test_image_probe_halves_the_area(self): + lighter = lighten_warmup_req(_server_args(), _req(1024, 1024, 1)) + assert lighter.num_frames == 1 + assert lighter.width * lighter.height <= 1024 * 1024 // 2 + assert lighter.width % 16 == 0 and lighter.height % 16 == 0 + + def test_the_original_request_is_left_alone(self): + req = _req(832, 480, 17) + lighten_warmup_req(_server_args(), req) + assert req.num_frames == 17 + + def test_a_probe_at_the_floor_cannot_shrink(self): + assert lighten_warmup_req(_server_args(), _req(16, 16, 1)) is None From a7ac4fa43c09c0001ffcc5f45d666b28d136e183 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 20 Aug 2026 19:41:41 +0800 Subject: [PATCH 11/90] fix host gather after residency promotion --- .../managers/memory_managers/layerwise_offload.py | 11 +++++++---- .../test/unit/test_host_resident_vocab_table.py | 12 ++++++++++++ 2 files changed, 19 insertions(+), 4 deletions(-) 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 cc6562cc47cd..2d2bf4be40ee 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 @@ -178,19 +178,22 @@ def restore_host_resident_tables( def _install_host_gather_hooks( module: torch.nn.Module, device: torch.device | str ) -> None: - """Run this module's gather on the host, move only the result.""" + """Run this module's gather beside its current weight, move only the result.""" - def _inputs_to_host(_module, args, kwargs): + def _inputs_to_table(_module, args, kwargs): if not args or not torch.is_tensor(args[0]): return None - return (args[0].to("cpu"),) + args[1:], kwargs + # A runtime residency promotion can later materialize this table on + # the device. Keep the gather input colocated with its current + # weight instead of unconditionally moving it back to the host. + return (args[0].to(_module.weight.device),) + args[1:], kwargs def _output_to_device(_module, _args, output): if not torch.is_tensor(output): return output return output.to(device, non_blocking=True) - module.register_forward_pre_hook(_inputs_to_host, with_kwargs=True) + module.register_forward_pre_hook(_inputs_to_table, with_kwargs=True) module.register_forward_hook(_output_to_device) diff --git a/python/sglang/multimodal_gen/test/unit/test_host_resident_vocab_table.py b/python/sglang/multimodal_gen/test/unit/test_host_resident_vocab_table.py index 5ee6b8180170..ba268a299d54 100644 --- a/python/sglang/multimodal_gen/test/unit/test_host_resident_vocab_table.py +++ b/python/sglang/multimodal_gen/test/unit/test_host_resident_vocab_table.py @@ -2,6 +2,7 @@ from unittest.mock import patch +import pytest import torch from sglang.multimodal_gen.runtime.managers.memory_managers import layerwise_offload @@ -90,6 +91,17 @@ def test_the_gather_matches_a_plain_lookup(self): restore_host_resident_tables(detach_host_resident_tables(model), "cpu") assert torch.equal(model.embed(ids), expected) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + def test_the_gather_follows_a_table_promoted_to_cuda(self): + model = _Declared() + with patch(THRESHOLD_PATH, 1024): + restore_host_resident_tables(detach_host_resident_tables(model), "cuda") + model.embed.to("cuda") + + output = model.embed(torch.tensor([[1, 2, 3]])) + + assert output.device.type == "cuda" + def test_nothing_is_hooked_when_nothing_qualifies(self): model = _Undeclared() with patch(THRESHOLD_PATH, 1024): From f5a10fe9887393244a212c914a63c113f0027aed Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 20 Aug 2026 21:43:40 +0800 Subject: [PATCH 12/90] test(diffusion): preserve golden CI baselines --- .github/workflows/pr-test-multimodal-gen.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/pr-test-multimodal-gen.yml b/.github/workflows/pr-test-multimodal-gen.yml index c06f1eabb82c..11144947695c 100644 --- a/.github/workflows/pr-test-multimodal-gen.yml +++ b/.github/workflows/pr-test-multimodal-gen.yml @@ -192,6 +192,7 @@ jobs: timeout-minutes: 120 env: RUNAI_STREAMER_MEMORY_LIMIT: 0 + SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY: 1 CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }} SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures run: | @@ -329,6 +330,7 @@ jobs: HF_TOKEN: ${{ secrets.MINIMAX_H3_HF_TOKEN || secrets.SGLANG_DIFFUSION_CI_HF_TOKEN || secrets.HF_TOKEN }} HUGGING_FACE_HUB_TOKEN: ${{ secrets.MINIMAX_H3_HF_TOKEN || secrets.SGLANG_DIFFUSION_CI_HF_TOKEN || secrets.HF_TOKEN }} RUNAI_STREAMER_MEMORY_LIMIT: 0 + SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY: 1 CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }} PARTITION_PLAN_JSON: ${{ needs.compute-diffusion-partitions.outputs.plan-2gpu }} SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures From 4a47e4663fe2ce9e35b600e8827171bd7077be99 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 20 Aug 2026 23:46:57 +0800 Subject: [PATCH 13/90] test(diffusion): preserve 1-GPU golden CI baselines --- .github/workflows/pr-test-multimodal-gen.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/pr-test-multimodal-gen.yml b/.github/workflows/pr-test-multimodal-gen.yml index 11144947695c..46773c9be220 100644 --- a/.github/workflows/pr-test-multimodal-gen.yml +++ b/.github/workflows/pr-test-multimodal-gen.yml @@ -119,6 +119,7 @@ jobs: timeout-minutes: 240 env: RUNAI_STREAMER_MEMORY_LIMIT: 0 + SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY: 1 CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }} PARTITION_PLAN_JSON: ${{ needs.compute-diffusion-partitions.outputs.plan-1gpu }} SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures From dfb2cfe22da1096e166bc7e68867a8924dd13b37 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 21 Aug 2026 00:10:39 +0800 Subject: [PATCH 14/90] diffusion-ci: test auto residency by default --- .github/workflows/pr-test-multimodal-gen.yml | 3 -- .../test/server/perf_baselines/5090.json | 4 +-- .../test/server/perf_baselines/h100.json | 36 +++++++++---------- .../sglang/multimodal_gen/test/test_utils.py | 2 +- 4 files changed, 21 insertions(+), 24 deletions(-) diff --git a/.github/workflows/pr-test-multimodal-gen.yml b/.github/workflows/pr-test-multimodal-gen.yml index 46773c9be220..c06f1eabb82c 100644 --- a/.github/workflows/pr-test-multimodal-gen.yml +++ b/.github/workflows/pr-test-multimodal-gen.yml @@ -119,7 +119,6 @@ jobs: timeout-minutes: 240 env: RUNAI_STREAMER_MEMORY_LIMIT: 0 - SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY: 1 CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }} PARTITION_PLAN_JSON: ${{ needs.compute-diffusion-partitions.outputs.plan-1gpu }} SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures @@ -193,7 +192,6 @@ jobs: timeout-minutes: 120 env: RUNAI_STREAMER_MEMORY_LIMIT: 0 - SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY: 1 CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }} SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures run: | @@ -331,7 +329,6 @@ jobs: HF_TOKEN: ${{ secrets.MINIMAX_H3_HF_TOKEN || secrets.SGLANG_DIFFUSION_CI_HF_TOKEN || secrets.HF_TOKEN }} HUGGING_FACE_HUB_TOKEN: ${{ secrets.MINIMAX_H3_HF_TOKEN || secrets.SGLANG_DIFFUSION_CI_HF_TOKEN || secrets.HF_TOKEN }} RUNAI_STREAMER_MEMORY_LIMIT: 0 - SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY: 1 CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }} PARTITION_PLAN_JSON: ${{ needs.compute-diffusion-partitions.outputs.plan-2gpu }} SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json b/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json index bd27a971a68b..1a494a457681 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json @@ -61,7 +61,7 @@ "expected_avg_denoise_ms": 338.78, "expected_median_denoise_ms": 345.95, "load_peak_vram_mb": 1126.0, - "runtime_peak_vram_mb": 13180.0, + "runtime_peak_vram_mb": 20064.0, "estimated_full_test_time_s": 94.0 }, "wan2_1_t2v_1.3b": { @@ -109,7 +109,7 @@ "expected_avg_denoise_ms": 241.53, "expected_median_denoise_ms": 265.1, "load_peak_vram_mb": 13120.0, - "runtime_peak_vram_mb": 17776.0, + "runtime_peak_vram_mb": 24372.0, "estimated_full_test_time_s": 329.8 }, "lingbot_video_moe_t2v": { diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index 634f9c186bca..3e9d33022472 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -303,7 +303,7 @@ "expected_avg_denoise_ms": 128.85, "expected_median_denoise_ms": 130.22, "load_peak_vram_mb": 23594.0, - "runtime_peak_vram_mb": 28540.0, + "runtime_peak_vram_mb": 36820.0, "estimated_full_test_time_s": 127.4 }, "flux_2_image_t2i": { @@ -395,7 +395,7 @@ "expected_avg_denoise_ms": 30.0, "expected_median_denoise_ms": 24.14, "load_peak_vram_mb": 8504.0, - "runtime_peak_vram_mb": 13732.0, + "runtime_peak_vram_mb": 20016.0, "estimated_full_test_time_s": 120.5 }, "flux_2_klein_base_image_t2i": { @@ -464,7 +464,7 @@ "expected_avg_denoise_ms": 116.49, "expected_median_denoise_ms": 117.28, "load_peak_vram_mb": 8504.0, - "runtime_peak_vram_mb": 13754.0, + "runtime_peak_vram_mb": 20038.0, "estimated_full_test_time_s": 124.4 }, "flux_2_ti2i": { @@ -697,7 +697,7 @@ "expected_avg_denoise_ms": 82.61, "expected_median_denoise_ms": 90.0, "load_peak_vram_mb": 13120.0, - "runtime_peak_vram_mb": 17776.0, + "runtime_peak_vram_mb": 24372.0, "estimated_full_test_time_s": 116.3 }, "zimage_image_t2i_fp8": { @@ -724,7 +724,7 @@ "expected_avg_denoise_ms": 103.24, "expected_median_denoise_ms": 108.73, "load_peak_vram_mb": 11344.0, - "runtime_peak_vram_mb": 15868.0, + "runtime_peak_vram_mb": 22464.0, "estimated_full_test_time_s": 123.7 }, "zimage_image_t2i_multi_lora": { @@ -751,7 +751,7 @@ "expected_avg_denoise_ms": 83.08, "expected_median_denoise_ms": 89.93, "load_peak_vram_mb": 14196.0, - "runtime_peak_vram_mb": 17612.0, + "runtime_peak_vram_mb": 24442.0, "estimated_full_test_time_s": 162.1 }, "zimage_image_t2i_2_gpus": { @@ -1042,7 +1042,7 @@ "expected_avg_denoise_ms": 76.99, "expected_median_denoise_ms": 98.28, "load_peak_vram_mb": 7814.0, - "runtime_peak_vram_mb": 18078.0, + "runtime_peak_vram_mb": 34372.0, "estimated_full_test_time_s": 126.0 }, "wan2_1_t2v_1.3b": { @@ -1110,7 +1110,7 @@ "expected_avg_denoise_ms": 143.82, "expected_median_denoise_ms": 144.57, "load_peak_vram_mb": 7814.0, - "runtime_peak_vram_mb": 18078.0, + "runtime_peak_vram_mb": 34372.0, "estimated_full_test_time_s": 129.3 }, "wan2_1_t2v_1.3b_cfg_parallel": { @@ -1178,7 +1178,7 @@ "expected_avg_denoise_ms": 73.51, "expected_median_denoise_ms": 74.05, "load_peak_vram_mb": 5422.0, - "runtime_peak_vram_mb": 11972.0, + "runtime_peak_vram_mb": 20054.0, "estimated_full_test_time_s": 127.6 }, "turbo_wan2_1_t2v_1.3b": { @@ -1338,7 +1338,7 @@ "expected_avg_denoise_ms": 253.29, "expected_median_denoise_ms": 256.21, "load_peak_vram_mb": 17610.0, - "runtime_peak_vram_mb": 39856.0, + "runtime_peak_vram_mb": 55254.0, "estimated_full_test_time_s": 141.7 }, "qwen_image_edit_2509_ti2i": { @@ -1484,7 +1484,7 @@ "expected_avg_denoise_ms": 83.86, "expected_median_denoise_ms": 112.67, "load_peak_vram_mb": 17610.0, - "runtime_peak_vram_mb": 39054.0, + "runtime_peak_vram_mb": 54914.0, "estimated_full_test_time_s": 125.2 }, "fast_hunyuan_video": { @@ -1509,7 +1509,7 @@ "expected_avg_denoise_ms": 871.39, "expected_median_denoise_ms": 1010.9, "load_peak_vram_mb": 27320.0, - "runtime_peak_vram_mb": 35926.0, + "runtime_peak_vram_mb": 37434.0, "estimated_full_test_time_s": 77.0 }, "wan2_2_i2v_a14b_2gpu": { @@ -1834,7 +1834,7 @@ "expected_avg_denoise_ms": 461.16, "expected_median_denoise_ms": 467.69, "load_peak_vram_mb": 31246.0, - "runtime_peak_vram_mb": 37788.0, + "runtime_peak_vram_mb": 45612.0, "estimated_full_test_time_s": 173.2 }, "wan2_2_t2v_a14b_lora_2gpu": { @@ -1960,7 +1960,7 @@ "expected_avg_denoise_ms": 146.48, "expected_median_denoise_ms": 147.68, "load_peak_vram_mb": 7814.0, - "runtime_peak_vram_mb": 18350.0, + "runtime_peak_vram_mb": 34422.0, "estimated_full_test_time_s": 129.6 }, "wan2_1_i2v_14b_lora_2gpu": { @@ -2099,7 +2099,7 @@ "expected_avg_denoise_ms": 242.72, "expected_median_denoise_ms": 246.92, "load_peak_vram_mb": 33110.0, - "runtime_peak_vram_mb": 38262.0, + "runtime_peak_vram_mb": 58544.0, "estimated_full_test_time_s": 135.3 }, "qwen_image_edit_2511_ti2i": { @@ -2322,7 +2322,7 @@ "expected_avg_denoise_ms": 146.61, "expected_median_denoise_ms": 147.3, "load_peak_vram_mb": 7814.0, - "runtime_peak_vram_mb": 18080.0, + "runtime_peak_vram_mb": 34372.0, "estimated_full_test_time_s": 129.3 }, "flux_2_image_t2i_upscaling_4x": { @@ -2459,7 +2459,7 @@ "expected_avg_denoise_ms": 144.68, "expected_median_denoise_ms": 144.74, "load_peak_vram_mb": 7814.0, - "runtime_peak_vram_mb": 24380.0, + "runtime_peak_vram_mb": 40502.0, "estimated_full_test_time_s": 129.3 }, "wan2_1_t2v_1.3b_frame_interp_2x_upscaling_4x": { @@ -2527,7 +2527,7 @@ "expected_avg_denoise_ms": 143.71, "expected_median_denoise_ms": 144.67, "load_peak_vram_mb": 7814.0, - "runtime_peak_vram_mb": 24388.0, + "runtime_peak_vram_mb": 40512.0, "estimated_full_test_time_s": 129.4 }, "ltx_2.3_one_stage_ti2v": { diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index b43a29fbb50f..5d2dea467d92 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -40,7 +40,7 @@ # NPU/ascend) is read from sgl-project/ci-data-diffusion, where the GT-gen workflows # publish. SGL_TEST_FILES_CI_DATA_REPO = "sgl-project/ci-data-diffusion" -SGL_TEST_FILES_CI_DATA_REVISION = "15b30030ef980756788ab40072f9223fe21a5526" +SGL_TEST_FILES_CI_DATA_REVISION = "3d96c172db9dc5d056804bf4b45919e195fa4a2e" # The NPU pin is kept as a separate branch so ascend GT can be bumped independently # when it's regenerated on its own cadence. From 45eb7af38c7caa820abfceab6586a59fa12b0098 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 21 Aug 2026 03:42:54 +0800 Subject: [PATCH 15/90] diffusion-ci: refresh LTX2 auto residency baseline --- .../sglang/multimodal_gen/test/server/perf_baselines/h100.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index 3e9d33022472..e1860afb8ec5 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -2578,7 +2578,7 @@ "expected_avg_denoise_ms": 520.0, "expected_median_denoise_ms": 494.9, "load_peak_vram_mb": 46278.0, - "runtime_peak_vram_mb": 49390.0, + "runtime_peak_vram_mb": 60330.0, "estimated_full_test_time_s": 144.2 }, "ltx_2.3_two_stage_t2v_2gpus": { From 769597f7758d5d1b8cc6db24327b975202119239 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Aug 2026 13:54:15 +0800 Subject: [PATCH 16/90] [diffusion] jointly plan serving memory residency --- docs/docs/sglang-diffusion/api/cli.mdx | 3 +- .../configs/pipeline_configs/cosmos3.py | 3 +- .../model_deployment_config.py | 6 +- .../entrypoints/openai/realtime/registry.py | 16 +- .../runtime/managers/gpu_worker.py | 112 ++- .../memory_managers/auto_residency.py | 860 ++++++++++++++++-- .../memory_managers/component_manager.py | 135 ++- .../memory_managers/host_memory_budget.py | 44 +- .../memory_managers/layerwise_offload.py | 375 +++++++- .../memory_managers/placement_budget.py | 239 +++++ .../runtime/scheduler_client.py | 90 +- .../runtime/server_args/auto_tune.py | 66 ++ .../runtime/server_args/server_args.py | 16 + .../multimodal_gen/runtime/server_warmup.py | 71 +- .../runtime/warmup_request_builder.py | 52 +- .../test/server/test_server_common.py | 40 + .../unit/realtime/test_realtime_runtime.py | 2 + .../test/unit/test_auto_residency.py | 652 ++++++++++++- .../test/unit/test_cfg_parallel_warmup.py | 47 + .../test/unit/test_component_residency.py | 127 +++ .../test/unit/test_dp_routing.py | 37 + .../test/unit/test_host_memory_budget.py | 24 + .../test/unit/test_layerwise_offload.py | 57 ++ .../test/unit/test_performance_metrics.py | 68 ++ .../test/unit/test_placement_budget.py | 286 ++++++ .../test/unit/test_server_args.py | 29 + 26 files changed, 3253 insertions(+), 204 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_placement_budget.py diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 95b4287470f8..568dcea3999c 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -85,7 +85,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 and dispatches residency from selected-GPU headroom and workload type: image DiTs stay resident above the 45 GiB threshold, while video DiT placement remains model-specific. 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 placement, measures the representative workload, and jointly selects component/layer residency against every rank and execution-phase VRAM constraint; a single-GPU run can also rebalance its HostPin allowance. The selected placement targets the model-default (or `--warmup-sampling-params`) workload with batch size 1; use explicit residency for larger production shapes or batches. 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. - `--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. Startup impact is model-dependent, so benchmark the target model before deployment. Disabled by default because checkpoint and model weights coexist temporarily, substantially increasing peak GPU memory. 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 @@ -95,6 +95,7 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis - `--enable-cfg-parallel {true|false}`: enable or explicitly disable CFG parallelism - `--encoder-parallel {auto|fold|dp|replicate}`: how text/image encoders use the GPUs in each DiT replica. `auto` TP-folds an encoder wide enough to benefit, selects batch DP when it can engage, and otherwise keeps the existing encoder TP layout; `fold` shards across the full replica whenever dimensions allow; `dp` splits a batched encode across encoder copies and composes with encoder TP; `replicate` disables folding and batch DP. Encoder collectives never cross `--dp-size` replicas. See [Encoder Parallelism](/docs/sglang-diffusion/encoder_parallel). - `--warmup-mode {off|request|server}`: control startup warmup for `sglang serve`; `off` skips warmup, `request` primes the request path, and `server` runs a full synthetic server warmup before serving traffic +- `--warmup-sampling-params JSON`: override the representative width, height, frame count, inference steps, or guidance values used by synthetic warmup and automatic residency planning. The warmup still caps executed steps and large video probes; the uncapped values describe the workload the memory planner must fit. - `--enable-torch-compile {true|false}`: compile native diffusion hot paths. When no warmup mode is configured, this also enables server warmup so first real requests do not pay compile latency. - `--offload-during-compile {true|false}`: when compile warmup is active, temporarily layerwise-offload DiT weights and move resident non-DiT components off-device so `max-autotune` fits on tighter-memory GPUs; the configured serving residency is restored before real traffic. Skipped under existing layerwise offload, Cache-DiT, or FSDP. - `--enable-breakable-cuda-graph {true|false}`: capture supported DiT forwards as breakable CUDA graph segments to reduce launch overhead. Requires `--warmup-resolutions` for every served resolution because each resolution is captured separately. A `quality=high` request is rejected when it would mount request-scoped DiT fusions that were not present during lossless graph capture; VAE-only high-quality paths remain compatible. diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py index 6ad73fa18697..bb5780b23e38 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py @@ -170,7 +170,8 @@ def supports_action_endpoint(self) -> bool: return True def get_model_deployment_config(self) -> ModelDeploymentConfig: - # Keep the DiT and VAE resident when the GPUs have the headroom. + # Fallback for offline and no-warmup paths. Server warmup calibration + # starts from an offloaded placement and derives residency directly. return ModelDeploymentConfig( keep_resident_min_available_gb=120, keep_resident_components=("dit", "vae"), diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py index 255a49af6587..6befe80b492d 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py @@ -14,9 +14,9 @@ class ModelDeploymentConfig: dit_layerwise_offload_modes: tuple[Literal["auto", "memory"], ...] = () auto_dit_offload_prefetch_size: float | None = None keep_resident_min_available_gb: float | None = None - # Per-model resident defaults. Auto mode additionally keeps an image DiT - # resident above the image workload memory threshold; video DiT placement - # stays with the model's FSDP/layerwise policy. + # Fallback residency for offline/no-warmup deployments. Eligible server + # warmup paths start load-safe and replace these coarse thresholds with a + # measured per-phase placement plan before reporting ready. keep_resident_components: tuple[OffloadComponentName, ...] = ("vae",) fsdp_auto_min_available_memory_gb: float | None = None fsdp_auto_requires_cfg: bool = True diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py index 4c682d0d75ea..005131af4ab7 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py @@ -4,11 +4,10 @@ from typing import TYPE_CHECKING -from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter import ( - BaseRealtimeModelAdapter, -) - if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter import ( + BaseRealtimeModelAdapter, + ) from sglang.multimodal_gen.runtime.server_args import ServerArgs @@ -67,3 +66,12 @@ def get_realtime_model_adapter( "Realtime video is not supported for pipeline config " f"{type(pipeline_config).__name__}; no realtime adapter is registered." ) + + +def has_realtime_model_adapter(server_args: ServerArgs) -> bool: + """Whether the resolved pipeline config has a registered adapter.""" + _register_builtin_realtime_model_adapters() + return any( + config_cls in _REALTIME_ADAPTER_REGISTRY + for config_cls in type(server_args.pipeline_config).__mro__ + ) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index c6303558c177..0d78d3f2e70e 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -43,7 +43,6 @@ get_tp_group, get_ulysses_parallel_rank, get_ulysses_parallel_world_size, - get_world_group, ) from sglang.multimodal_gen.runtime.entrypoints.utils import ( materialize_output_sample, @@ -63,10 +62,15 @@ WarmupMemoryRecord, apply_promotions, collect_promotion_candidates, + component_runtime_weight_bytes, describe_error, estimate_default_workload_peak_bytes, + estimate_workload_phase_peaks, format_applied_changes, format_plan_summary, + layerwise_host_pin_capacity_bytes, + layerwise_pinned_host_bytes, + module_uses_quantized_weights, plan_auto_residency, plan_summary_payload, rank_candidates_by_h2d_savings, @@ -76,6 +80,10 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( peek_global_component_residency_manager, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + HOST_COPY_RESERVE_BYTES, + host_memory_available_bytes, +) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( configure_layerwise_offload_modules, ) @@ -549,8 +557,7 @@ def _execute_forward_common( measure_server_warmup = ( req.is_warmup and bool(req.extra.get("server_based_warmup")) - and not current_platform.is_cpu() - and not current_platform.is_mps() + and current_platform.is_cuda() ) warmup_baseline_allocated_bytes = 0 try: @@ -693,16 +700,35 @@ def _execute_forward_common( def _record_server_warmup_memory( self, *, req: Req, baseline_allocated_bytes: int, succeeded: bool ) -> None: + phase_peaks: dict[str, int] = {} + phase_components: dict[str, tuple[str, ...]] = {} + residency_manager = peek_global_component_residency_manager() + if residency_manager is not None: + for phase_name, ( + components, + peak_bytes, + ) in residency_manager.take_warmup_phase_peaks().items(): + phase_peaks[phase_name] = peak_bytes + phase_components[phase_name] = components + request_peak = max( + int(torch.get_device_module().max_memory_reserved()), + max(phase_peaks.values(), default=0), + ) + if request_peak > max(phase_peaks.values(), default=0): + # Work after the residency-managed stage timeline (for example, + # output materialization) must remain a placement constraint. + phase_peaks["request:untracked"] = request_peak + phase_components["request:untracked"] = () self._auto_residency_warmup_records.append( WarmupMemoryRecord( width=int(req.width or 0), height=int(req.height or 0), num_frames=int(req.num_frames or 1), baseline_allocated_bytes=int(baseline_allocated_bytes), - peak_reserved_bytes=int( - torch.get_device_module().max_memory_reserved() - ), + peak_reserved_bytes=request_peak, succeeded=succeeded, + phase_peak_reserved_bytes=phase_peaks, + phase_active_components=phase_components, ) ) @@ -1117,14 +1143,11 @@ def get_can_stay_resident_components( ) -> List[str]: """Which currently offloaded components would fit in the headroom. - Reuses the auto-residency candidate collection and its H2D-savings - greedy order, so this post-request hint names the same components, - in the same order, that ``--performance-mode auto`` would promote - (it deliberately does not maximize the component count the way the - old fixed cheap-first order did, and it omits components driven by a - pipeline-custom residency strategy -- user flags cannot act on - those). Unlike the promotion plan it applies no reserve or margin: - it reports raw capacity after this request, not a placement decision. + Reuses the auto-residency candidate frontier and benefit ranking, but + remains a raw-capacity hint rather than the phase-constrained joint + plan. It omits components driven by a pipeline-custom residency + strategy because user residency flags cannot control them, and it + applies no reserve or activation margin. """ if not self.pipeline or not self.pipeline.modules: return [] @@ -1141,10 +1164,18 @@ def get_can_stay_resident_components( num_inference_steps=( workload.num_inference_steps if workload is not None else 1 ), + allow_host_pin_reallocation=self.server_args.num_gpus == 1, ) can_stay_resident = [] + seen_components: set[str] = set() for candidate in rank_candidates_by_h2d_savings(candidates): + if ( + not candidate.permanent_residency + or candidate.component_name in seen_components + ): + continue + seen_components.add(candidate.component_name) usage_gb = candidate.promoted_weight_bytes / GIB_BYTES if usage_gb <= remaining_gpu_mem_gb: can_stay_resident.append(candidate.component_name) @@ -1163,12 +1194,6 @@ def apply_auto_residency(self) -> OutputBatch: an uncaught raise there would leave the peer ranks parked in the collective until the group timeout. """ - if (self.server_args.dp_size or 1) > 1: - # DP replicas dequeue control reqs at independent times, but the - # gather below spans the whole torch world; entering it per - # replica would deadlock. (The HTTP orchestrator also gates this - # case; this guard keeps the worker safe on its own.) - return OutputBatch(output={"status": PROMOTION_STATUS_SKIPPED}) if self._auto_residency_applied: # one-shot per process: records measured after a promotion # describe the promoted layout, not the startup strategy @@ -1315,6 +1340,18 @@ def _build_auto_residency_report( elif workload.workload_units() is None: skip_reason = "default workload resolution unknown" else: + quantized_components = [ + name + for name, module in self.pipeline.modules.items() + if isinstance(module, torch.nn.Module) + and module_uses_quantized_weights(module) + ] + if quantized_components: + skip_reason = ( + "loaded quantized components: " + f"{', '.join(sorted(quantized_components))}" + ) + if skip_reason is None: # A probe that failed at or below the target is a measurement, not # missing data: the card cannot hold the default workload as it is # already configured, and promoting weights would only add to it. @@ -1339,8 +1376,20 @@ def _build_auto_residency_report( candidates=[], skip_reason=skip_reason, ) + runtime_weights_by_component = component_runtime_weight_bytes( + self.pipeline.modules + ) estimated_peak_bytes = estimate_default_workload_peak_bytes( - records=records, target_units=workload.workload_units() + records=records, + target_units=workload.workload_units(), + constant_weight_bytes=max(runtime_weights_by_component.values(), default=0), + ) + estimated_phase_peaks, active_components_by_phase = ( + estimate_workload_phase_peaks( + records=records, + target_units=workload.workload_units(), + component_weight_bytes=runtime_weights_by_component, + ) ) # What this process may still grow into: free VRAM plus what its # allocator already reserved. Unlike the raw device total, this @@ -1355,21 +1404,34 @@ def _build_auto_residency_report( explicit_residency_mode_of=self.server_args.explicit_residency_mode, custom_strategy_names=self.pipeline.component_residency_strategies, num_inference_steps=workload.num_inference_steps, + allow_host_pin_reallocation=self.server_args.num_gpus == 1, ) return RankResidencyReport( rank=self.rank, budget_bytes=budget_bytes, estimated_peak_bytes=estimated_peak_bytes, + estimated_peak_bytes_by_phase=estimated_phase_peaks, + active_components_by_phase=active_components_by_phase, + node_rank=self.server_args.node_rank, + pinned_host_bytes=layerwise_pinned_host_bytes(self.pipeline.modules), + host_pin_capacity_bytes=layerwise_host_pin_capacity_bytes( + self.pipeline.modules + ), + host_transition_headroom_bytes=max( + 0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES + ), candidates=candidates, skip_reason=skip_reason, ) def _auto_residency_all_gather(self, obj: Any) -> list[Any]: - world_group = get_world_group() - if world_group.world_size == 1: + replica_group = get_replica_group() + if replica_group.world_size == 1: return [obj] - gathered: list[Any] = [None] * world_group.world_size - torch.distributed.all_gather_object(gathered, obj, group=world_group.cpu_group) + gathered: list[Any] = [None] * replica_group.world_size + torch.distributed.all_gather_object( + gathered, obj, group=replica_group.cpu_group + ) return gathered def _rollback_applied_promotions(self) -> str | None: diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index ea9de2436ad0..7b3c21b5a8df 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -16,16 +16,23 @@ Promotion targets the model default workload only (default resolution, default frames, batch=1). Larger shapes, batches, or multi-image inputs need explicit ``--component-residency``. + +Loading and serving are deliberately separate placement states. The existing +auto policy provides the initial state; when that state can complete loading +and calibration, this module optimizes the long-lived serving state and +validates the transition with a post-placement warmup. It does not force a +single placement to serve two different lifecycle objectives. """ from __future__ import annotations -from typing import TYPE_CHECKING, Callable, Iterable, Mapping +from typing import TYPE_CHECKING, Callable, Iterable, Mapping, Sequence import msgspec import torch import torch.nn as nn +from sglang.multimodal_gen.runtime.layers.linear import LinearBase from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD, @@ -33,13 +40,25 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import ( is_fsdp_managed_module, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + HOST_COPY_RESERVE_BYTES, + host_memory_available_bytes, +) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, + compute_streamed_layers, is_layerwise_offloaded_module, + iter_materialized_weights, + release_unused_pinned_memory, ) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( is_dit_component_name, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.placement_budget import ( + NoFeasiblePlacementError, + PlacementOption, + optimize_placement, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger if TYPE_CHECKING: @@ -94,6 +113,8 @@ class WarmupMemoryRecord(msgspec.Struct, frozen=True): baseline_allocated_bytes: int peak_reserved_bytes: int succeeded: bool + phase_peak_reserved_bytes: dict[str, int] = {} + phase_active_components: dict[str, tuple[str, ...]] = {} def workload_units(self) -> int: return max(1, self.width) * max(1, self.height) * max(1, self.num_frames) @@ -107,6 +128,32 @@ class PromotionCandidate(msgspec.Struct, frozen=True): promoted_weight_bytes: int # Estimated per-request host-to-device traffic this promotion removes. h2d_bytes_per_request: int + # Layerwise candidates jointly choose stage-scoped GPU residency and host + # pinning. None is used by ordinary component-offload promotion. + target_layerwise_resident_layers: tuple[int, ...] | None = None + target_layerwise_pinned_layers: tuple[tuple[int, ...], ...] | None = None + pinned_host_delta_bytes: int = 0 + host_unpin_scratch_bytes: int = 0 + host_pin_scratch_bytes: int = 0 + permanent_residency: bool = False + # Device-memory delta relative to the measured placement. A component + # already loaded for its own phase has a different delta from phases where + # it is absent; keeping both avoids adding the same weights twice. + active_device_delta_bytes: int = 0 + inactive_device_delta_bytes: int = 0 + + def option_key(self) -> str: + if self.target_layerwise_resident_layers is None: + return f"{self.component_name}:resident" + layer_counts = ",".join( + str(count) for count in self.target_layerwise_resident_layers + ) + pinned = "|".join( + ",".join(str(index) for index in indices) or "-" + for indices in self.target_layerwise_pinned_layers or () + ) + permanence = "permanent" if self.permanent_residency else "stage" + return f"{self.component_name}:{permanence}:layers={layer_counts}:pins={pinned}" class DefaultWorkload(msgspec.Struct, frozen=True): @@ -134,6 +181,12 @@ class RankResidencyReport(msgspec.Struct, frozen=True): rank: int budget_bytes: int estimated_peak_bytes: int | None + estimated_peak_bytes_by_phase: dict[str, int] = {} + active_components_by_phase: dict[str, tuple[str, ...]] = {} + node_rank: int = 0 + pinned_host_bytes: int = 0 + host_pin_capacity_bytes: int = 0 + host_transition_headroom_bytes: int = 0 candidates: list[PromotionCandidate] = [] skip_reason: str | None = None @@ -144,6 +197,7 @@ class AutoResidencyPlan(msgspec.Struct, frozen=True): estimated_peak_bytes: int = 0 reserve_bytes: int = 0 budget_bytes: int = 0 + resource_budget_bytes: dict[str, int] = {} promotions: list[PromotionCandidate] = [] skip_reason: str | None = None @@ -151,6 +205,10 @@ class AutoResidencyPlan(msgspec.Struct, frozen=True): class AppliedPromotion(msgspec.Struct, frozen=True): component_name: str residency_mode: str + previous_layerwise_resident_layers: tuple[int, ...] | None = None + previous_layerwise_pinned_layers: tuple[tuple[int, ...], ...] | None = None + layerwise_offload_disabled: bool = False + pinned_host_changed: bool = False def resolve_default_workload(server_args: ServerArgs) -> DefaultWorkload: @@ -187,6 +245,7 @@ def estimate_default_workload_peak_bytes( *, records: Iterable[WarmupMemoryRecord], target_units: int | None, + constant_weight_bytes: int = 0, ) -> int | None: """Extrapolate warmup peaks to the default workload. @@ -243,7 +302,10 @@ def estimate_default_workload_peak_bytes( )[:2] slope = (large_peak - small_peak) / (large_units - small_units) if slope >= 0: - constant = large_peak - slope * large_units + constant = max( + large_peak - slope * large_units, + min(constant_weight_bytes, large_peak), + ) return int( constant + slope * target_units * ACTIVATION_EXTRAPOLATION_MARGIN ) @@ -253,7 +315,9 @@ def estimate_default_workload_peak_bytes( estimates = [] for record in records: peak = record.peak_reserved_bytes - baseline = min(record.baseline_allocated_bytes, peak) + baseline = min( + max(record.baseline_allocated_bytes, constant_weight_bytes), peak + ) activation = peak - baseline ratio = target_units / record.workload_units() estimates.append( @@ -262,12 +326,73 @@ def estimate_default_workload_peak_bytes( return max(estimates) -def _module_weight_bytes(module: nn.Module) -> int: - """Full weight+buffer footprint, reading through layerwise CPU buffers.""" - from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( - iter_materialized_weights, +def estimate_workload_phase_peaks( + *, + records: Iterable[WarmupMemoryRecord], + target_units: int | None, + component_weight_bytes: Mapping[str, int], +) -> tuple[dict[str, int], dict[str, tuple[str, ...]]]: + """Estimate each measured execution phase at the target workload. + + A component already active in a phase is part of that phase's measured + peak. Keeping a component resident therefore adds no weight bytes to its + own component-offload phase, while it adds the full footprint to phases + where the component was absent. Returning both the estimated peaks and the + conservative intersection of active components preserves that distinction. + """ + records = list(records) + successful = [record for record in records if record.succeeded] + phase_names = sorted( + { + phase_name + for record in successful + for phase_name in record.phase_peak_reserved_bytes + } ) + estimated_peaks: dict[str, int] = {} + active_components: dict[str, tuple[str, ...]] = {} + for phase_name in phase_names: + phase_records = [ + record + for record in successful + if phase_name in record.phase_peak_reserved_bytes + ] + if not phase_records: + continue + active_sets = [ + set(record.phase_active_components.get(phase_name, ())) + for record in phase_records + ] + active = set.intersection(*active_sets) if active_sets else set() + weight_floor = sum(component_weight_bytes.get(name, 0) for name in active) + phase_measurements = [ + WarmupMemoryRecord( + width=record.width, + height=record.height, + num_frames=record.num_frames, + baseline_allocated_bytes=min( + record.baseline_allocated_bytes, + record.phase_peak_reserved_bytes[phase_name], + ), + peak_reserved_bytes=record.phase_peak_reserved_bytes[phase_name], + succeeded=True, + ) + for record in phase_records + ] + estimate = estimate_default_workload_peak_bytes( + records=phase_measurements, + target_units=target_units, + constant_weight_bytes=weight_floor, + ) + if estimate is None: + continue + estimated_peaks[phase_name] = estimate + active_components[phase_name] = tuple(sorted(active)) + return estimated_peaks, active_components + +def _module_weight_bytes(module: nn.Module) -> int: + """Full weight+buffer footprint, reading through layerwise CPU buffers.""" param_bytes = sum( tensor.numel() * tensor.element_size() for _, tensor in iter_materialized_weights(module) @@ -278,9 +403,78 @@ def _module_weight_bytes(module: nn.Module) -> int: return param_bytes + buffer_bytes +def module_uses_quantized_weights(module: nn.Module) -> bool: + """Whether a loaded native component contains quantized weight storage.""" + for submodule in module.modules(): + if isinstance(submodule, LinearBase) and submodule.quant_config is not None: + return True + quantized_float_dtypes = {torch.float8_e4m3fn, torch.float8_e5m2} + return any( + not tensor.is_floating_point() or tensor.dtype in quantized_float_dtypes + for _, tensor in iter_materialized_weights(module) + ) + + +def component_runtime_weight_bytes(modules: Mapping[str, object]) -> dict[str, int]: + """Persistent plus streamed-window weight bytes during a component use.""" + runtime_bytes: dict[str, int] = {} + for name, module in modules.items(): + if not isinstance(module, nn.Module): + continue + full_bytes = _module_weight_bytes(module) + if not isinstance(module, LayerwiseOffloadableModuleMixin): + runtime_bytes[name] = full_bytes + continue + managers = [ + manager for manager in module.layerwise_offload_managers if manager.enabled + ] + if not managers: + runtime_bytes[name] = full_bytes + continue + managed_bytes = sum(manager.offloaded_weight_bytes() for manager in managers) + managed_peak = sum( + manager.peak_managed_device_weight_bytes() for manager in managers + ) + runtime_bytes[name] = max(0, full_bytes - managed_bytes) + managed_peak + return runtime_bytes + + +def layerwise_pinned_host_bytes(modules: Mapping[str, object]) -> int: + """Pinned layer-store bytes in this process, without alias double counts.""" + seen_managers: set[int] = set() + total = 0 + for module in modules.values(): + if not isinstance(module, LayerwiseOffloadableModuleMixin): + continue + for manager in module.layerwise_offload_managers: + manager_id = id(manager) + if manager_id in seen_managers: + continue + seen_managers.add(manager_id) + total += manager.pinned_host_weight_bytes() + return total + + +def layerwise_host_pin_capacity_bytes(modules: Mapping[str, object]) -> int: + """This process's non-overlapping HostPin allowance.""" + seen_budgets: set[int] = set() + total = 0 + for module in modules.values(): + if not isinstance(module, LayerwiseOffloadableModuleMixin): + continue + for manager in module.layerwise_offload_managers: + budget = manager.pin_budget + budget_id = id(budget) + if budget_id in seen_budgets: + continue + seen_budgets.add(budget_id) + total += max(0, budget.available_bytes - budget.reserve_bytes) + return total + + def _layerwise_offloaded_bytes(module: LayerwiseOffloadableModuleMixin) -> int: return sum( - manager.offloaded_weight_bytes() + sum(manager.layer_weight_bytes().values()) for manager in module.layerwise_offload_managers if manager.enabled ) @@ -295,6 +489,98 @@ def component_resident_size_bytes(module: nn.Module, residency_mode: str) -> int return _module_weight_bytes(module) +def _layerwise_transfer_work_bytes( + *, + managers: Sequence, + resident_layers: tuple[int, ...], + pinned_layers: tuple[tuple[int, ...], ...], + uses_per_streamed_layer: int, +) -> int: + """Relative transfer work for one request under a layerwise placement. + + Pageable copies are weighted 2x: the CUDA driver stages them through an + internal pinned buffer and the copy cannot run ahead of compute. This is a + conservative ordering metric rather than a latency prediction. + """ + total = 0 + for manager, resident_count, pinned_indices in zip( + managers, resident_layers, pinned_layers + ): + streamed = set( + compute_streamed_layers( + num_layers=manager.num_layers, + resident_layers=resident_count, + policy=manager.residency_policy, + ) + ) + pinned = set(pinned_indices) + for layer_idx, weight_bytes in manager.layer_weight_bytes().items(): + uses = uses_per_streamed_layer if layer_idx in streamed else 1 + transfer_multiplier = 1 if layer_idx in pinned else 2 + total += uses * transfer_multiplier * weight_bytes + return total + + +def _layerwise_pin_targets( + *, + managers: Sequence, + resident_layers: tuple[int, ...], + current_pinned_layers: tuple[tuple[int, ...], ...], + uses_per_streamed_layer: int, +) -> list[tuple[tuple[int, ...], ...]]: + """Useful HostPin frontiers for one resident-layer placement. + + Streamed layers have the highest pin value because they transfer every + denoise step. Prefixes of that deterministic value order cover every + capacity breakpoint without enumerating all 2^N layer subsets. + """ + ordered: list[tuple[int, int, int]] = [] + for manager_index, (manager, resident_count) in enumerate( + zip(managers, resident_layers) + ): + streamed = set( + compute_streamed_layers( + num_layers=manager.num_layers, + resident_layers=resident_count, + policy=manager.residency_policy, + ) + ) + if not manager.pin_cpu_memory: + continue + for layer_idx in manager.pinnable_layer_indices(): + uses = uses_per_streamed_layer if layer_idx in streamed else 1 + ordered.append((-uses, manager_index, layer_idx)) + ordered.sort() + + targets: list[tuple[tuple[int, ...], ...]] = [current_pinned_layers] + selected = [set() for _ in managers] + targets.append(tuple(() for _ in managers)) + for _, manager_index, layer_idx in ordered: + selected[manager_index].add(layer_idx) + targets.append(tuple(tuple(sorted(indices)) for indices in selected)) + return list(dict.fromkeys(targets)) + + +def _layerwise_host_transition_bytes( + *, + managers: Sequence, + current_pinned_layers: tuple[tuple[int, ...], ...], + target_pinned_layers: tuple[tuple[int, ...], ...], +) -> tuple[int, int]: + """Temporary host allocations needed by the release/acquire phases.""" + unpin_bytes = 0 + pin_bytes = 0 + for manager, current_indices, target_indices in zip( + managers, current_pinned_layers, target_pinned_layers + ): + layer_bytes = manager.layer_host_store_bytes() + current = set(current_indices) + target = set(target_indices) + unpin_bytes += sum(layer_bytes.get(index, 0) for index in current - target) + pin_bytes += sum(layer_bytes.get(index, 0) for index in target - current) + return unpin_bytes, pin_bytes + + def collect_promotion_candidates( *, modules: Mapping[str, object], @@ -302,6 +588,7 @@ def collect_promotion_candidates( explicit_residency_mode_of: Callable[[str], str | None], custom_strategy_names: Iterable[str], num_inference_steps: int, + allow_host_pin_reallocation: bool = True, ) -> list[PromotionCandidate]: """List implicitly offloaded native components eligible for promotion. @@ -341,14 +628,169 @@ def collect_promotion_candidates( if mode == LAYERWISE_OFFLOAD and is_dit_component_name(name) else 1 ) - candidates.append( - PromotionCandidate( - component_name=name, - residency_mode=mode, - promoted_weight_bytes=promoted_bytes, - h2d_bytes_per_request=promoted_bytes * uses_per_request, + if mode == COMPONENT_OFFLOAD: + candidates.append( + PromotionCandidate( + component_name=name, + residency_mode=mode, + promoted_weight_bytes=promoted_bytes, + h2d_bytes_per_request=promoted_bytes, + permanent_residency=True, + active_device_delta_bytes=0, + inactive_device_delta_bytes=promoted_bytes, + ) ) + continue + + if not isinstance(module, LayerwiseOffloadableModuleMixin): + continue + managers = [ + manager for manager in module.layerwise_offload_managers if manager.enabled + ] + if not managers: + continue + current_resident_layers = tuple(manager.resident_layers for manager in managers) + current_pinned_layers = tuple( + manager.pinned_layer_indices() for manager in managers + ) + current_resident_bytes = sum( + manager.resident_weight_bytes() for manager in managers ) + current_peak_device_bytes = sum( + manager.peak_managed_device_weight_bytes() for manager in managers + ) + current_pinned_bytes = sum( + manager.pinned_host_weight_bytes() for manager in managers + ) + current_transfer_work = _layerwise_transfer_work_bytes( + managers=managers, + resident_layers=current_resident_layers, + pinned_layers=current_pinned_layers, + uses_per_streamed_layer=uses_per_request, + ) + + resident_targets = {current_resident_layers} + if is_dit_component_name(name) and num_inference_steps > 1: + for target in range( + max(current_resident_layers, default=0), + max(manager.num_layers for manager in managers) + 1, + ): + counts = tuple(min(target, manager.num_layers) for manager in managers) + if all( + count >= current + for count, current in zip(counts, current_resident_layers) + ): + resident_targets.add(counts) + + for target_resident_layers in sorted(resident_targets): + target_resident_bytes = sum( + manager.resident_weight_bytes(count) + for manager, count in zip(managers, target_resident_layers) + ) + target_peak_device_bytes = sum( + manager.peak_managed_device_weight_bytes(count) + for manager, count in zip(managers, target_resident_layers) + ) + incremental_bytes = target_resident_bytes - current_resident_bytes + pin_targets = ( + _layerwise_pin_targets( + managers=managers, + resident_layers=target_resident_layers, + current_pinned_layers=current_pinned_layers, + uses_per_streamed_layer=uses_per_request, + ) + if allow_host_pin_reallocation + else [current_pinned_layers] + ) + for target_pinned_layers in pin_targets: + if ( + target_resident_layers == current_resident_layers + and target_pinned_layers == current_pinned_layers + ): + continue + target_pinned_bytes = sum( + sum( + manager.layer_host_store_bytes().get(layer_idx, 0) + for layer_idx in pinned_indices + ) + for manager, pinned_indices in zip(managers, target_pinned_layers) + ) + target_transfer_work = _layerwise_transfer_work_bytes( + managers=managers, + resident_layers=target_resident_layers, + pinned_layers=target_pinned_layers, + uses_per_streamed_layer=uses_per_request, + ) + unpin_scratch, pin_scratch = _layerwise_host_transition_bytes( + managers=managers, + current_pinned_layers=current_pinned_layers, + target_pinned_layers=target_pinned_layers, + ) + candidates.append( + PromotionCandidate( + component_name=name, + residency_mode=mode, + promoted_weight_bytes=incremental_bytes, + h2d_bytes_per_request=( + current_transfer_work - target_transfer_work + ), + target_layerwise_resident_layers=target_resident_layers, + target_layerwise_pinned_layers=target_pinned_layers, + pinned_host_delta_bytes=( + target_pinned_bytes - current_pinned_bytes + ), + host_unpin_scratch_bytes=unpin_scratch, + host_pin_scratch_bytes=pin_scratch, + active_device_delta_bytes=max( + 0, target_peak_device_bytes - current_peak_device_bytes + ), + inactive_device_delta_bytes=0, + ) + ) + + full_resident_layers = tuple(manager.num_layers for manager in managers) + # Fully resident layers never read their host stores. On one worker we + # can release those pins and reuse the allowance. Multi-worker pin + # migration needs a node-coordinated host-RAM scratch budget because + # repacking briefly holds both the old and new layer buffers. + permanent_pin_targets = [current_pinned_layers] + if allow_host_pin_reallocation: + permanent_pin_targets.append(tuple(() for _ in managers)) + for permanent_pins in dict.fromkeys(permanent_pin_targets): + permanent_pinned_bytes = sum( + sum( + manager.layer_host_store_bytes().get(layer_idx, 0) + for layer_idx in pinned_indices + ) + for manager, pinned_indices in zip(managers, permanent_pins) + ) + unpin_scratch, pin_scratch = _layerwise_host_transition_bytes( + managers=managers, + current_pinned_layers=current_pinned_layers, + target_pinned_layers=permanent_pins, + ) + candidates.append( + PromotionCandidate( + component_name=name, + residency_mode=mode, + promoted_weight_bytes=promoted_bytes - current_resident_bytes, + h2d_bytes_per_request=current_transfer_work, + target_layerwise_resident_layers=full_resident_layers, + target_layerwise_pinned_layers=permanent_pins, + pinned_host_delta_bytes=( + permanent_pinned_bytes - current_pinned_bytes + ), + host_unpin_scratch_bytes=unpin_scratch, + host_pin_scratch_bytes=pin_scratch, + permanent_residency=True, + active_device_delta_bytes=max( + 0, promoted_bytes - current_peak_device_bytes + ), + # Stage-scoped resident layers are released after the use, + # so the complete managed footprint is new elsewhere. + inactive_device_delta_bytes=promoted_bytes, + ) + ) return candidates @@ -365,6 +807,7 @@ def rank_candidates_by_h2d_savings( key=lambda candidate: ( -candidate.h2d_bytes_per_request, candidate.component_name, + candidate.option_key(), ), ) @@ -373,6 +816,45 @@ def _skip_plan(reason: str) -> AutoResidencyPlan: return AutoResidencyPlan(skip_reason=reason) +def _vram_reserve_bytes(budget_bytes: int) -> int: + return max( + int(budget_bytes * VRAM_RESERVE_FRACTION), + min( + MIN_VRAM_RESERVE_BYTES, + int(budget_bytes * MAX_VRAM_RESERVE_FRACTION), + ), + ) + + +def _binding_phase_constraints( + report: RankResidencyReport, +) -> list[tuple[str, int, tuple[str, ...]]]: + """Keep only the highest peak for each active-component placement. + + Every candidate has the same device delta in phases with the same active + component set. The lower peaks are therefore provably redundant; removing + them reduces the exact optimizer's resource dimension without changing + its feasible placements. + """ + phase_peaks = report.estimated_peak_bytes_by_phase + if not phase_peaks: + if report.estimated_peak_bytes is None: + return [] + phase_peaks = {"request": report.estimated_peak_bytes} + binding: dict[tuple[str, ...], tuple[str, int]] = {} + for phase_name, phase_peak in phase_peaks.items(): + active = tuple(sorted(report.active_components_by_phase.get(phase_name, ()))) + current = binding.get(active) + if current is None or (phase_peak, phase_name) > (current[1], current[0]): + binding[active] = (phase_name, phase_peak) + return [ + (f"gpu:rank{report.rank}:{phase_name}", phase_peak, active) + for active, (phase_name, phase_peak) in sorted( + binding.items(), key=lambda item: item[1][0] + ) + ] + + def _consensus_candidates( reports: list[RankResidencyReport], ) -> list[PromotionCandidate]: @@ -382,29 +864,63 @@ def _consensus_candidates( fits on one rank but overflows another. """ per_rank_maps = [ - {candidate.component_name: candidate for candidate in report.candidates} + {candidate.option_key(): candidate for candidate in report.candidates} for report in reports ] - common_names = set(per_rank_maps[0]) + common_keys = set(per_rank_maps[0]) for candidate_map in per_rank_maps[1:]: - common_names &= set(candidate_map) + common_keys &= set(candidate_map) merged = [] - for name in sorted(common_names): - rank_candidates = [candidate_map[name] for candidate_map in per_rank_maps] + for option_key in sorted(common_keys): + rank_candidates = [candidate_map[option_key] for candidate_map in per_rank_maps] modes = {candidate.residency_mode for candidate in rank_candidates} if len(modes) != 1: continue + component_names = {candidate.component_name for candidate in rank_candidates} + targets = { + candidate.target_layerwise_resident_layers for candidate in rank_candidates + } + pinned_targets = { + candidate.target_layerwise_pinned_layers for candidate in rank_candidates + } + permanent = {candidate.permanent_residency for candidate in rank_candidates} + if ( + len(component_names) != 1 + or len(targets) != 1 + or len(pinned_targets) != 1 + or len(permanent) != 1 + ): + continue merged.append( PromotionCandidate( - component_name=name, + component_name=component_names.pop(), residency_mode=modes.pop(), promoted_weight_bytes=max( candidate.promoted_weight_bytes for candidate in rank_candidates ), - h2d_bytes_per_request=max( + h2d_bytes_per_request=min( candidate.h2d_bytes_per_request for candidate in rank_candidates ), + target_layerwise_resident_layers=targets.pop(), + target_layerwise_pinned_layers=pinned_targets.pop(), + pinned_host_delta_bytes=max( + candidate.pinned_host_delta_bytes for candidate in rank_candidates + ), + host_unpin_scratch_bytes=max( + candidate.host_unpin_scratch_bytes for candidate in rank_candidates + ), + host_pin_scratch_bytes=max( + candidate.host_pin_scratch_bytes for candidate in rank_candidates + ), + permanent_residency=permanent.pop(), + active_device_delta_bytes=max( + candidate.active_device_delta_bytes for candidate in rank_candidates + ), + inactive_device_delta_bytes=max( + candidate.inactive_device_delta_bytes + for candidate in rank_candidates + ), ) ) return merged @@ -422,33 +938,100 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP estimated_peak = max(report.estimated_peak_bytes for report in reports) budget = min(report.budget_bytes for report in reports) - reserve = max( - int(budget * VRAM_RESERVE_FRACTION), - min(MIN_VRAM_RESERVE_BYTES, int(budget * MAX_VRAM_RESERVE_FRACTION)), - ) + reserves_by_rank = { + report.rank: _vram_reserve_bytes(report.budget_bytes) for report in reports + } + reserve = max(reserves_by_rank.values()) candidates = _consensus_candidates(reports) if not candidates: return _skip_plan("no implicitly offloaded components to promote") - ordered = rank_candidates_by_h2d_savings(candidates) - promotions = [] - promoted_bytes = 0 - for candidate in ordered: - next_bytes = promoted_bytes + candidate.promoted_weight_bytes - # Known conservatism: estimated_peak already contains whatever part - # of this component was on GPU at the measured peak (streamed - # layers, keep-ready components), so those bytes are counted twice. - # The error is in the safe direction; refining it needs per-phase - # peaks. - if estimated_peak + next_bytes + reserve <= budget: - promotions.append(candidate) - promoted_bytes = next_bytes + phase_constraints = { + report.rank: _binding_phase_constraints(report) for report in reports + } + resource_budgets: dict[str, int] = {} + for report in reports: + for resource_name, phase_peak, _ in phase_constraints[report.rank]: + resource_budgets[resource_name] = ( + report.budget_bytes - phase_peak - reserves_by_rank[report.rank] + ) + has_hostpin_options = any( + candidate.pinned_host_delta_bytes != 0 + or candidate.host_unpin_scratch_bytes != 0 + or candidate.host_pin_scratch_bytes != 0 + for candidate in candidates + ) + if has_hostpin_options: + if len(reports) != 1: + return _skip_plan("dynamic HostPin placement requires a single-worker node") + report = reports[0] + resource_budgets[f"hostpin:node{report.node_rank}"] = ( + report.host_pin_capacity_bytes - report.pinned_host_bytes + ) + resource_budgets[f"hostram:node{report.node_rank}:unpin"] = ( + report.host_transition_headroom_bytes + ) + resource_budgets[f"hostram:node{report.node_rank}:pin"] = ( + report.host_transition_headroom_bytes + ) + + candidate_by_key = {candidate.option_key(): candidate for candidate in candidates} + report_candidates = [ + {candidate.option_key(): candidate for candidate in report.candidates} + for report in reports + ] + + options = [] + for candidate in candidates: + resource_deltas: dict[str, int] = {} + for report, rank_candidates in zip(reports, report_candidates): + rank_candidate = rank_candidates[candidate.option_key()] + for resource_name, _, active_components in phase_constraints[report.rank]: + component_is_active = candidate.component_name in active_components + phase_cost = ( + rank_candidate.active_device_delta_bytes + if component_is_active + else rank_candidate.inactive_device_delta_bytes + ) + resource_deltas[resource_name] = phase_cost + if has_hostpin_options: + host_resource = f"hostpin:node{report.node_rank}" + resource_deltas[host_resource] = ( + resource_deltas.get(host_resource, 0) + + rank_candidate.pinned_host_delta_bytes + ) + resource_deltas[f"hostram:node{report.node_rank}:unpin"] = ( + rank_candidate.host_unpin_scratch_bytes + ) + resource_deltas[f"hostram:node{report.node_rank}:pin"] = ( + rank_candidate.host_pin_scratch_bytes + ) + options.append( + PlacementOption( + group_key=candidate.component_name, + option_key=candidate.option_key(), + resource_delta_bytes=resource_deltas, + estimated_latency_savings=candidate.h2d_bytes_per_request, + ) + ) + + try: + placement = optimize_placement( + options, + resource_budget_bytes=resource_budgets, + ) + except NoFeasiblePlacementError as error: + return _skip_plan(str(error)) + promotions = rank_candidates_by_h2d_savings( + candidate_by_key[selection.option_key] for selection in placement.selections + ) return AutoResidencyPlan( estimated_peak_bytes=estimated_peak, reserve_bytes=reserve, budget_bytes=budget, + resource_budget_bytes=resource_budgets, promotions=promotions, ) @@ -464,7 +1047,8 @@ def apply_promotions( Component-offload modules are only re-marked resident; the physical move happens on their next use (the post-promotion re-warmup) through ``ResidentStrategy``, which also applies any per-use target dtype. - Layerwise modules load all layers now and drop their hooks. + Layerwise options either retain a planned number of layers during the + denoise stage or load the whole component and drop its hooks. Raises on failure after undoing the promotions already applied, so a caller observing an exception knows this rank is back on the original @@ -474,7 +1058,19 @@ def apply_promotions( """ applied: list[AppliedPromotion] = [] try: - for candidate in plan.promotions: + ordered_promotions = sorted( + plan.promotions, + key=lambda candidate: candidate.pinned_host_delta_bytes, + ) + promotion_modules: dict[str, nn.Module] = {} + snapshots: dict[str, AppliedPromotion] = {} + previous_pins_by_component: dict[str, tuple[tuple[int, ...], ...]] = {} + for candidate in ordered_promotions: + if candidate.component_name in promotion_modules: + raise RuntimeError( + f"multiple placement options selected for " + f"{candidate.component_name!r}" + ) module = modules.get(candidate.component_name) if not isinstance(module, nn.Module): raise RuntimeError( @@ -494,18 +1090,105 @@ def apply_promotions( f"promotion target {candidate.component_name!r} became " "layerwise-managed between planning and apply" ) - # Record before acting so a mid-promotion failure rolls back the - # partially promoted component as well. - applied.append( - AppliedPromotion( + promotion_modules[candidate.component_name] = module + if candidate.target_layerwise_resident_layers is not None: + if not isinstance(module, LayerwiseOffloadableModuleMixin): + raise RuntimeError( + f"partial promotion target {candidate.component_name!r} " + "lost its layerwise offload capability" + ) + previous = tuple( + manager.resident_layers + for manager in module.layerwise_offload_managers + ) + previous_pinned = module.layerwise_pinned_layers() + previous_pins_by_component[candidate.component_name] = previous_pinned + snapshots[candidate.component_name] = AppliedPromotion( component_name=candidate.component_name, residency_mode=candidate.residency_mode, + previous_layerwise_resident_layers=previous, + previous_layerwise_pinned_layers=previous_pinned, + layerwise_offload_disabled=candidate.permanent_residency, + pinned_host_changed=( + candidate.target_layerwise_pinned_layers != previous_pinned + ), ) + if candidate.target_layerwise_pinned_layers is None: + raise RuntimeError( + f"layerwise placement {candidate.option_key()!r} has no " + "host pin target" + ) + continue + + snapshots[candidate.component_name] = AppliedPromotion( + component_name=candidate.component_name, + residency_mode=candidate.residency_mode, + ) + pinning_changes = [ + candidate + for candidate in ordered_promotions + if candidate.target_layerwise_pinned_layers is not None + and candidate.target_layerwise_pinned_layers + != previous_pins_by_component[candidate.component_name] + ] + if pinning_changes: + host_headroom = max( + 0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES + ) + required_unpin = sum( + candidate.host_unpin_scratch_bytes for candidate in pinning_changes ) + required_pin = sum( + candidate.host_pin_scratch_bytes for candidate in pinning_changes + ) + if max(required_unpin, required_pin) > host_headroom: + raise MemoryError( + "host memory changed after planning: HostPin repack needs " + f"{max(required_unpin, required_pin) / GIB_BYTES:.2f} GiB " + f"but {host_headroom / GIB_BYTES:.2f} GiB is available" + ) + applied_names: set[str] = set() + # 1. release every component's old allowance + # 2. return cached pinned blocks to the driver + # 3. acquire the final allowances + # This avoids a transient HostPin oversubscription when one component + # hands its budget to another. + for candidate in pinning_changes: + applied.append(snapshots[candidate.component_name]) + applied_names.add(candidate.component_name) + module = promotion_modules[candidate.component_name] + assert isinstance(module, LayerwiseOffloadableModuleMixin) + module.release_layerwise_pins_outside( + candidate.target_layerwise_pinned_layers or () + ) + if pinning_changes: + release_unused_pinned_memory() + for candidate in ordered_promotions: + if candidate.component_name not in applied_names: + applied.append(snapshots[candidate.component_name]) + applied_names.add(candidate.component_name) + module = promotion_modules[candidate.component_name] + if candidate.target_layerwise_resident_layers is not None: + assert isinstance(module, LayerwiseOffloadableModuleMixin) + module.set_layerwise_pinned_layers( + candidate.target_layerwise_pinned_layers or () + ) + if candidate.permanent_residency: + server_args.require_component_resident( + candidate.component_name, + feature_name=AUTO_RESIDENCY_FEATURE_NAME, + ) + module.disable_offload() + else: + module.set_layerwise_resident_layer_counts( + candidate.target_layerwise_resident_layers + ) + continue server_args.require_component_resident( candidate.component_name, feature_name=AUTO_RESIDENCY_FEATURE_NAME ) if candidate.residency_mode == LAYERWISE_OFFLOAD: + assert isinstance(module, LayerwiseOffloadableModuleMixin) module.disable_offload() except Exception as apply_error: try: @@ -533,14 +1216,51 @@ def rollback_promotions( failures are re-raised at the end so one broken component cannot leave the later (earlier-applied) ones promoted. """ + applied = list(applied) errors: list[str] = [] - for promotion in reversed(list(applied)): + pinning_changes = [ + promotion + for promotion in applied + if promotion.pinned_host_changed + and promotion.previous_layerwise_pinned_layers is not None + ] + for promotion in pinning_changes: + try: + module = modules.get(promotion.component_name) + if not isinstance(module, LayerwiseOffloadableModuleMixin): + raise RuntimeError("lost layerwise offload capability") + module.release_layerwise_pins_outside( + promotion.previous_layerwise_pinned_layers or () + ) + except Exception as e: + errors.append(f"{promotion.component_name}: {describe_error(e)}") + if pinning_changes: + release_unused_pinned_memory() + for promotion in reversed(applied): try: + module = modules.get(promotion.component_name) + if promotion.previous_layerwise_resident_layers is not None: + if not isinstance(module, LayerwiseOffloadableModuleMixin): + raise RuntimeError("lost layerwise offload capability") + if promotion.layerwise_offload_disabled: + module.enable_offload() + server_args.release_required_component_residency( + promotion.component_name, + feature_name=AUTO_RESIDENCY_FEATURE_NAME, + ) + module.restore_layerwise_resident_layers( + promotion.previous_layerwise_resident_layers + ) + if promotion.previous_layerwise_pinned_layers is None: + raise RuntimeError("lost previous layerwise host placement") + module.restore_layerwise_pinned_layers( + promotion.previous_layerwise_pinned_layers + ) + continue server_args.release_required_component_residency( promotion.component_name, feature_name=AUTO_RESIDENCY_FEATURE_NAME, ) - module = modules.get(promotion.component_name) if not isinstance(module, nn.Module): continue if promotion.residency_mode == LAYERWISE_OFFLOAD: @@ -572,12 +1292,7 @@ def format_plan_summary( if plan.skip_reason is not None: return f"Auto residency: skipped ({plan.skip_reason})" promoted = ( - ", ".join( - f"{candidate.component_name}" - f"({candidate.residency_mode}, " - f"{candidate.promoted_weight_bytes / GIB_BYTES:.1f} GiB)" - for candidate in plan.promotions - ) + ", ".join(_format_candidate_summary(candidate) for candidate in plan.promotions) or "none" ) measured = ", ".join( @@ -605,11 +1320,17 @@ def format_applied_changes(*, plan: AutoResidencyPlan) -> str: placement, or disable the adjustment entirely with the kill switch. """ changes = "; ".join( - f"{candidate.component_name}: {candidate.residency_mode} -> resident" - for candidate in plan.promotions + _format_promotion_change(candidate) for candidate in plan.promotions ) - equivalent = "--component-residency " + " ".join( - f"{candidate.component_name}=resident" for candidate in plan.promotions + component_args = [ + f"{candidate.component_name}=resident" + for candidate in plan.promotions + if candidate.permanent_residency + ] + equivalent = ( + "--component-residency " + " ".join(component_args) + if component_args + else "none (the selected partial layer/HostPin placement is auto-only)" ) return ( f"Auto residency: adjusted {changes}. " @@ -619,6 +1340,33 @@ def format_applied_changes(*, plan: AutoResidencyPlan) -> str: ) +def _format_promotion_change(candidate: PromotionCandidate) -> str: + if candidate.target_layerwise_resident_layers is not None: + pin_counts = tuple( + len(indices) for indices in candidate.target_layerwise_pinned_layers or () + ) + scope = "permanent" if candidate.permanent_residency else "stage-scoped" + return ( + f"{candidate.component_name}: {scope} resident layers=" + f"{candidate.target_layerwise_resident_layers}, pinned layers=" + f"{pin_counts}" + ) + return f"{candidate.component_name}: {candidate.residency_mode} -> resident" + + +def _format_candidate_summary(candidate: PromotionCandidate) -> str: + target = "" + if candidate.target_layerwise_resident_layers is not None: + target = ( + f", layers={candidate.target_layerwise_resident_layers}, " + f"pins={tuple(len(indices) for indices in candidate.target_layerwise_pinned_layers or ())}" + ) + return ( + f"{candidate.component_name}({candidate.residency_mode}{target}, " + f"{candidate.promoted_weight_bytes / GIB_BYTES:.1f} GiB)" + ) + + def plan_summary_payload(*, plan: AutoResidencyPlan, status: str) -> dict: """Minimal decision payload for the warmup orchestrator. diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py index 8bdcf4b9d458..e49006a09498 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py @@ -129,6 +129,11 @@ def __init__( ] = {} self._uses_seen: dict[str, ComponentUse] = {} self._modules_seen: dict[str, nn.Module] = {} + self._track_warmup_memory = False + self._warmup_phase_key: str | None = None + self._warmup_phase_components: tuple[str, ...] = () + self._warmup_phase_peaks: dict[str, tuple[tuple[str, ...], int]] = {} + self._completed_warmup_phase_peaks: dict[str, tuple[tuple[str, ...], int]] = {} def refresh_pipeline(self, pipeline: ComponentResidencyPipeline) -> None: custom_strategies = dict(pipeline.component_residency_strategies) @@ -188,6 +193,21 @@ def begin_request( self._ordered_uses = tuple( use for uses in self._stage_uses_by_index for use in uses ) + self._track_warmup_memory = ( + self.state.batch_is_warmup + and current_platform.is_cuda() + and torch.get_device_module().is_available() + ) + self._warmup_phase_key = None + self._warmup_phase_components = () + self._warmup_phase_peaks = {} + self._completed_warmup_phase_peaks = {} + if self._track_warmup_memory: + # GPUWorker reset the request peak before entering the pipeline. + # Start the first interval without another reset so preprocessing + # before stage 0 remains part of the placement constraints. + self._warmup_phase_key = "request:before-stage" + self._warmup_phase_components = self._warmup_active_components() @staticmethod def _is_warmup_batch(batch: ResidencyBatch | list[ResidencyBatch]) -> bool: @@ -205,6 +225,11 @@ def before_stage( self.state.stage_index = stage_index self.state.stage_name = self.stage_name(stage) self.state.next_stage_name = self._next_stage_name(stage_index) + if self._track_warmup_memory: + self._begin_warmup_phase( + key=f"{stage_index}:{self.state.stage_name}:setup", + components=self._warmup_active_components(), + ) def begin_stage(self) -> None: """Prepare a stage that declares one uninterrupted component use.""" @@ -214,6 +239,7 @@ def begin_stage(self) -> None: def end_stage(self) -> None: """Close the component interval owned by the current stage.""" + self._record_warmup_phase_peak() if self._active_use is None: return if self._active_use.stage_name != self.state.stage_name: @@ -244,7 +270,13 @@ def begin_use(self, use: ComponentUse, module: nn.Module | None = None) -> None: self._active_use_module is not None and active_module is not self._active_use_module ) + requires_prepare = active_module is not None and ( + self._active_use_module is None + or module_changed + or use.target_dtype != previous_use.target_dtype + ) if module_changed: + self._begin_warmup_transition(previous_use, None) self._disable_active_nvtx() self._finish_use( previous_use, @@ -252,18 +284,19 @@ def begin_use(self, use: ComponentUse, module: nn.Module | None = None) -> None: keep_on_warmup=False, force=True, ) - if active_module is not None and ( - self._active_use_module is None - or module_changed - or use.target_dtype != previous_use.target_dtype - ): + self._begin_warmup_transition(None, use) + elif requires_prepare: + self._begin_warmup_transition(previous_use, use) + if requires_prepare: active_module = self._prepare_forward_use(use, module=active_module) + self._begin_warmup_use(use) self._active_use = use self._active_use_module = active_module self.state.current_use = use self._enable_nvtx_for_use(use, active_module) return if self._active_use is not None: + self._begin_warmup_transition(self._active_use, None) self._disable_active_nvtx() self._finish_use( self._active_use, @@ -273,8 +306,10 @@ def begin_use(self, use: ComponentUse, module: nn.Module | None = None) -> None: self._active_use = None self._active_use_module = None self.state.current_use = None + self._begin_warmup_transition(None, use) self._mark_current_use(use) module = self._prepare_forward_use(use, module=module) + self._begin_warmup_use(use) self._active_use = use self._active_use_module = module self._enable_nvtx_for_use(use, module) @@ -284,6 +319,7 @@ def end_use(self, use: ComponentUse, module: nn.Module | None = None) -> None: """End one sequential component use interval.""" if self._active_use is None or not self._same_use(self._active_use, use): return + self._begin_warmup_transition(self._active_use, None) self._disable_active_nvtx() self._finish_use( self._active_use, @@ -297,6 +333,7 @@ def end_use(self, use: ComponentUse, module: nn.Module | None = None) -> None: self._active_use = None self._active_use_module = None self.state.current_use = None + self._begin_warmup_between_uses() self._prefetch_next_memory_intensive_use() @contextmanager @@ -359,6 +396,7 @@ def finish_active_use(self, *, prefetch_next: bool = True) -> None: if self._active_use is None: return active_use = self._active_use + self._begin_warmup_transition(active_use, None) self._disable_active_nvtx() self._finish_use( active_use, @@ -368,6 +406,7 @@ def finish_active_use(self, *, prefetch_next: bool = True) -> None: self._active_use = None self._active_use_module = None self.state.current_use = None + self._begin_warmup_between_uses() if prefetch_next: self._prefetch_next_memory_intensive_use() @@ -515,11 +554,97 @@ def finish_request(self) -> None: not self._is_single_dit_component(component_name) or keep_single_dit ) strategy = self.strategy_for(component_name, module) + if self._track_warmup_memory: + will_prepare = self.state.batch_is_warmup and preferred + self._begin_warmup_phase( + key=f"request:cleanup:{component_name}", + components=self._warmup_active_components( + (use,) if will_prepare else () + ), + ) was_on_supported_device = self._module_on_supported_device(module) strategy.finish_request(module, use, self.state, preferred=preferred) self._empty_cache_after_large_release( use, strategy, module, was_on_supported_device ) + if self._track_warmup_memory: + self._record_warmup_phase_peak() + self._warmup_phase_peaks["idle"] = ( + self._warmup_active_components(), + int(torch.get_device_module().memory_reserved()), + ) + self._completed_warmup_phase_peaks = dict(self._warmup_phase_peaks) + self._track_warmup_memory = False + + def _begin_warmup_phase(self, *, key: str, components: tuple[str, ...]) -> None: + if not self._track_warmup_memory: + return + self._record_warmup_phase_peak() + self._warmup_phase_key = key + self._warmup_phase_components = tuple(sorted(set(components))) + torch.get_device_module().reset_peak_memory_stats() + + def _begin_warmup_transition( + self, previous: ComponentUse | None, upcoming: ComponentUse | None + ) -> None: + if not self._track_warmup_memory: + return + previous_name = previous.component_name if previous is not None else "idle" + upcoming_name = upcoming.component_name if upcoming is not None else "idle" + self._begin_warmup_phase( + key=( + f"{self.state.stage_index}:{self.state.stage_name}:transition:" + f"{previous_name}->{upcoming_name}" + ), + components=self._warmup_active_components( + tuple(use for use in (previous, upcoming) if use is not None) + ), + ) + + def _begin_warmup_use(self, use: ComponentUse) -> None: + phase = use.phase or use.component_name + self._begin_warmup_phase( + key=f"{self.state.stage_index}:{self.state.stage_name}:use:{phase}", + components=self._warmup_active_components((use,)), + ) + + def _begin_warmup_between_uses(self) -> None: + self._begin_warmup_phase( + key=f"{self.state.stage_index}:{self.state.stage_name}:between", + components=self._warmup_active_components(), + ) + + def _warmup_active_components( + self, active_uses: Sequence[ComponentUse] = () + ) -> tuple[str, ...]: + components = {use.component_name for use in active_uses} + for component_name, module in self.pipeline.modules.items(): + if not isinstance(module, nn.Module): + continue + if is_layerwise_offloaded_module(module): + continue + if self._module_on_supported_device(module): + components.add(component_name) + return tuple(sorted(components)) + + def _record_warmup_phase_peak(self) -> None: + if not self._track_warmup_memory or self._warmup_phase_key is None: + return + peak = int(torch.get_device_module().max_memory_reserved()) + previous = self._warmup_phase_peaks.get(self._warmup_phase_key) + if previous is None or peak > previous[1]: + self._warmup_phase_peaks[self._warmup_phase_key] = ( + self._warmup_phase_components, + peak, + ) + + def take_warmup_phase_peaks( + self, + ) -> dict[str, tuple[tuple[str, ...], int]]: + """Return and clear the most recently completed warmup phase peaks.""" + peaks = self._completed_warmup_phase_peaks + self._completed_warmup_phase_peaks = {} + return peaks def stage_name(self, stage: ComponentResidencyStage) -> str: return self._stage_names_by_id.get(id(stage), stage.__class__.__name__) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py index 3c6293f1ecfb..37a6fa9d8807 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py @@ -168,9 +168,19 @@ def host_copies_would_not_fit(weight_bytes: int) -> bool: return weight_bytes >= host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES +def host_pin_reserve_bytes(available_bytes: int) -> int: + return max( + int(available_bytes * HOST_RESERVE_FRACTION), + MIN_HOST_RESERVE_BYTES, + ) + + class HostPinBudget: """Hands out pinned-host-memory allowances until the headroom runs out. + This enforces a capacity selected elsewhere; it does not choose component + placement independently from the VRAM planner. + Pinning is not all-or-nothing per process: a component whose weights stream once per request gains far less from pinning than one re-streamed on every denoise step. So the hot components are offered the budget first, and a cold @@ -184,15 +194,33 @@ class HostPinBudget: it only fires when the bytes genuinely do not fit. """ - def __init__(self, available_bytes: int | None = None) -> None: + def __init__( + self, + available_bytes: int | None = None, + *, + reserve_bytes: int | None = None, + ) -> None: if available_bytes is None: available_bytes = host_memory_available_bytes() self.available_bytes = available_bytes - self.reserve_bytes = max( - int(available_bytes * HOST_RESERVE_FRACTION), MIN_HOST_RESERVE_BYTES + self.reserve_bytes = ( + host_pin_reserve_bytes(available_bytes) + if reserve_bytes is None + else reserve_bytes ) self.committed_bytes = 0 + @classmethod + def for_local_worker(cls, local_worker_count: int) -> "HostPinBudget": + """Give one worker a non-overlapping share of the node allowance.""" + worker_count = max(1, local_worker_count) + node_available = host_memory_available_bytes() + node_spendable = max(0, node_available - host_pin_reserve_bytes(node_available)) + return cls( + available_bytes=node_spendable // worker_count, + reserve_bytes=0, + ) + @property def spendable_bytes(self) -> int: return max(0, self.available_bytes - self.reserve_bytes - self.committed_bytes) @@ -223,6 +251,16 @@ def request(self, *, component_name: str, weight_bytes: int) -> bool: ) return False + def release(self, weight_bytes: int) -> None: + """Return a previously committed allowance after buffers are unpinned.""" + if weight_bytes <= 0: + return + if weight_bytes > self.committed_bytes: + raise ValueError( + "cannot release more pinned-host memory than was committed" + ) + self.committed_bytes -= weight_bytes + def pin_benefit_bytes(*, weight_bytes: int, uses_per_request: int) -> int: """Host-to-device bytes a pin would cover for one request. 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 2d2bf4be40ee..c07088c65985 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 @@ -1,9 +1,10 @@ import bisect +import gc import math import queue import re import threading -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from contextlib import nullcontext from time import perf_counter from typing import Any, Dict, List, Optional, Set, Tuple @@ -44,6 +45,20 @@ logger = init_logger(__name__) +def release_unused_pinned_memory() -> None: + """Return unreferenced cached HostPin blocks to the CUDA driver. + + ``torch.cuda.empty_cache()`` only covers device allocations. PyTorch's + host caching allocator currently exposes its corresponding operation via + this internal binding; keeping the call isolated makes the compatibility + boundary explicit until a public API is available. + """ + if not current_platform.is_cuda(): + return + gc.collect() + torch._C._host_emptyCache() + + def compute_streamed_layers( *, num_layers: int, resident_layers: int, policy: str ) -> tuple[int, ...]: @@ -382,6 +397,7 @@ def __init__( # was never reached. A private budget reads the same host limit. self._pin_budget = pin_budget if pin_budget is not None else HostPinBudget() self._pin_component_name = pin_component_name + self._layer_hosting: Dict[int, str] = {} # an explicit MPS zero avoids staging the next layer alongside the # active one; MPS has no transfer overlap to recover from that cost self.prefetch_size = ( @@ -561,19 +577,43 @@ def _initialize(self) -> None: def _layer_byte_totals( self, layer_groups: Dict ) -> Tuple[Dict[int, int], Dict[int, int]]: - """Per layer: (all weight bytes, the subset that are checkpoint views).""" + """Per layer: (planned host-store bytes, checkpoint-view bytes).""" totals: Dict[int, int] = {} mapped: Dict[int, int] = {} for layer_idx, dtype_to_params in layer_groups.items(): total = 0 from_mapping = 0 - for weights in dtype_to_params.values(): + for dtype, weights in dtype_to_params.items(): + contiguous_numel = 0 + anonymous_numel = 0 for _, weight in weights: tensor = self._to_local_tensor(weight) - nbytes = tensor.untyped_storage().nbytes() - total += nbytes + if tensor.is_contiguous(): + contiguous_numel = self._align_numel_offset( + contiguous_numel, dtype + ) + contiguous_numel += tensor.numel() + if not self._mapped_regions.holds(tensor): + anonymous_numel = self._align_numel_offset( + anonymous_numel, dtype + ) + anonymous_numel += tensor.numel() + continue + storage_numel = ( + 0 + if tensor.numel() == 0 + else 1 + + sum( + (size - 1) * abs(stride) + for size, stride in zip(tensor.shape, tensor.stride()) + ) + ) + storage_bytes = storage_numel * tensor.element_size() + total += storage_bytes if self._mapped_regions.holds(tensor): - from_mapping += nbytes + from_mapping += storage_bytes + total += contiguous_numel * dtype.itemsize + from_mapping += (contiguous_numel - anonymous_numel) * dtype.itemsize totals[layer_idx] = total mapped[layer_idx] = from_mapping return totals, mapped @@ -607,7 +647,7 @@ def _plan_layer_hosting(self, layer_groups: Dict) -> Dict[int, str]: pinned_bytes = 0 hosting: Dict[int, str] = {} pin_order: List[int] = [] - spendable = self._pin_budget.spendable_bytes if self._pin_budget else 0 + spendable = self._pin_budget.spendable_bytes streamed = [idx for idx in self._streamed_order if idx in totals] resident = [idx for idx in sorted(totals) if idx not in set(streamed)] for layer_idx in streamed + resident: @@ -655,7 +695,7 @@ def anonymous_new_bytes() -> int: anonymous_new_bytes() / 1024**3, host_memory_available_bytes() / 1024**3, ) - if pinned_bytes and self._pin_budget is not None: + if pinned_bytes: self._pin_budget.request( component_name=self._pin_component_name, weight_bytes=pinned_bytes ) @@ -700,6 +740,7 @@ def _initialize_layer_weights(self) -> None: ).append((name, tensor)) layer_hosting = self._plan_layer_hosting(layer_groups) + self._layer_hosting = dict(layer_hosting) # 2. concat and offload (in pinned memory) for layer_idx, dtype_to_params in layer_groups.items(): @@ -1356,6 +1397,218 @@ def offloaded_weight_bytes(self) -> int: total += int(numel) * meta["dtype"].itemsize return total + def layer_weight_bytes(self) -> dict[int, int]: + """Physical bytes copied to the device for each managed layer.""" + totals = {layer_idx: 0 for layer_idx in self._weight_metadata} + for layer_idx, buffers in self._consolidated_cpu_weights.items(): + totals[layer_idx] += sum( + buffer.untyped_storage().nbytes() for buffer in buffers.values() + ) + for layer_idx, tensors in self._strided_cpu_weights.items(): + totals[layer_idx] += sum( + tensor.untyped_storage().nbytes() for tensor in tensors.values() + ) + for layer_idx, tensors in self._mapped_cpu_weights.items(): + totals[layer_idx] += sum( + tensor.numel() * tensor.element_size() for tensor in tensors.values() + ) + return totals + + def layer_host_store_bytes(self) -> dict[int, int]: + """Physical anonymous host-store bytes booked for each managed layer.""" + totals = {layer_idx: 0 for layer_idx in self._weight_metadata} + for layer_idx, buffers in self._consolidated_cpu_weights.items(): + totals[layer_idx] += sum( + buffer.untyped_storage().nbytes() for buffer in buffers.values() + ) + for layer_idx, tensors in self._strided_cpu_weights.items(): + totals[layer_idx] += sum( + tensor.untyped_storage().nbytes() for tensor in tensors.values() + ) + return totals + + def pinned_host_weight_bytes(self) -> int: + """Logical managed-weight bytes currently backed by pinned memory.""" + layer_bytes = self.layer_host_store_bytes() + return sum( + layer_bytes.get(layer_idx, 0) + for layer_idx, placement in self._layer_hosting.items() + if placement == "pinned" + ) + + def pinnable_layer_indices(self) -> tuple[int, ...]: + """Layers whose anonymous host store can switch pinning in place.""" + return tuple( + layer_idx + for layer_idx, placement in sorted(self._layer_hosting.items()) + if placement != "mapped" + ) + + def pinned_layer_indices(self) -> tuple[int, ...]: + return tuple( + layer_idx + for layer_idx, placement in sorted(self._layer_hosting.items()) + if placement == "pinned" + ) + + def _repack_layer_host_store(self, layer_idx: int, *, pinned: bool) -> None: + if self._mapped_cpu_weights.get(layer_idx): + raise ValueError( + f"layer {layer_idx} retains checkpoint mappings and cannot be repacked" + ) + + consolidated: Dict[torch.dtype, torch.Tensor] = {} + for dtype, source in self._consolidated_cpu_weights.get(layer_idx, {}).items(): + target = torch.empty( + source.shape, + dtype=dtype, + pin_memory=pinned, + ) + target.copy_(source) + consolidated[dtype] = target + + strided: Dict[str, torch.Tensor] = {} + for name, source in self._strided_cpu_weights.get(layer_idx, {}).items(): + target = torch.empty_strided( + size=source.shape, + stride=source.stride(), + dtype=source.dtype, + pin_memory=pinned, + ) + target.copy_(source) + strided[name] = target + + self._consolidated_cpu_weights[layer_idx] = consolidated + self._strided_cpu_weights[layer_idx] = strided + + def _set_layer_pinned(self, layer_idx: int, *, pinned: bool) -> None: + current = self._layer_hosting[layer_idx] + if (current == "pinned") == pinned: + return + layer_bytes = self.layer_host_store_bytes().get(layer_idx, 0) + if pinned and not self._pin_budget.request( + component_name=self._pin_component_name, + weight_bytes=layer_bytes, + ): + raise MemoryError( + f"host pin budget cannot pin {self._pin_component_name} " + f"layer {layer_idx}" + ) + try: + self._repack_layer_host_store(layer_idx, pinned=pinned) + except Exception: + if pinned: + self._pin_budget.release(layer_bytes) + raise + if not pinned: + self._pin_budget.release(layer_bytes) + self._layer_hosting[layer_idx] = "pinned" if pinned else "pageable" + + def set_pinned_layers(self, pinned_layers: Iterable[int]) -> tuple[int, ...]: + """Reassign the manager's anonymous host stores between pin tiers. + + Mapped layers stay mapped: materializing them would consume host RAM + outside the pin budget and can turn a bounded configuration into a host + OOM. Unpins run before new pins, allowing one component to hand its + allowance to another without a transient HostPin oversubscription. + """ + target = set(pinned_layers) + pinnable = set(self.pinnable_layer_indices()) + if not target <= pinnable: + raise ValueError( + f"pinned layer selection contains non-pinnable layers: " + f"{sorted(target - pinnable)}" + ) + previous = self.pinned_layer_indices() + current = set(previous) + to_unpin = sorted(current - target) + to_pin = sorted(target - current) + changed_unpinned: list[int] = [] + changed_pinned: list[int] = [] + try: + for layer_idx in to_unpin: + self._set_layer_pinned(layer_idx, pinned=False) + changed_unpinned.append(layer_idx) + for layer_idx in to_pin: + self._set_layer_pinned(layer_idx, pinned=True) + changed_pinned.append(layer_idx) + except Exception: + for layer_idx in reversed(changed_pinned): + self._set_layer_pinned(layer_idx, pinned=False) + for layer_idx in reversed(changed_unpinned): + self._set_layer_pinned(layer_idx, pinned=True) + raise + return previous + + @property + def pin_budget(self) -> HostPinBudget: + return self._pin_budget + + def resident_weight_bytes(self, resident_layers: int | None = None) -> int: + """Managed bytes retained across denoise steps for a layer count.""" + count = self.resident_layers if resident_layers is None else resident_layers + streamed = compute_streamed_layers( + num_layers=self.num_layers, + resident_layers=count, + policy=self.residency_policy, + ) + resident = set(range(self.num_layers)) - set(streamed) + layer_bytes = self.layer_weight_bytes() + return sum(layer_bytes.get(layer_idx, 0) for layer_idx in resident) + + def peak_managed_device_weight_bytes( + self, resident_layers: int | None = None + ) -> int: + """Conservative managed-weight working set during this component's use. + + Resident layers coexist with the copy window. The exact layers in that + window change throughout the forward, so use the largest streamed + layers as a placement-independent upper bound. This is tighter than + charging every newly resident byte on top of the measured phase peak: + the measured peak already contains the current copy window. + """ + count = self.resident_layers if resident_layers is None else resident_layers + streamed = compute_streamed_layers( + num_layers=self.num_layers, + resident_layers=count, + policy=self.residency_policy, + ) + layer_bytes = self.layer_weight_bytes() + resident = set(range(self.num_layers)) - set(streamed) + resident_bytes = sum(layer_bytes.get(index, 0) for index in resident) + copy_window = min(max(1, self.prefetch_size), len(streamed)) + streamed_window_bytes = sum( + sorted((layer_bytes.get(index, 0) for index in streamed), reverse=True)[ + :copy_window + ] + ) + return resident_bytes + streamed_window_bytes + + @torch.compiler.disable + def set_resident_layers(self, resident_layers: int) -> int: + """Change the stage-scoped resident-layer count between requests. + + CPU stores and hooks stay intact. The next component use loads the new + resident set once and keeps it across denoise steps; request teardown + still releases every layer, so other pipeline phases recover the VRAM. + """ + target = min(max(0, int(resident_layers)), self.num_layers) + previous = self.resident_layers + if target == previous: + return previous + self.release_all() + self.resident_layers = target + self._streamed_order = compute_streamed_layers( + num_layers=self.num_layers, + resident_layers=target, + policy=self.residency_policy, + ) + self._resident_set = frozenset(range(self.num_layers)) - set( + self._streamed_order + ) + self._residency_active = False + return previous + def iter_cpu_weights(self): """Yield (name, tensor) pairs from consolidated CPU buffers. @@ -1832,6 +2085,93 @@ def prepare_for_next_req(self): for manager in self.layerwise_offload_managers: manager.prepare_for_next_req(non_blocking=True) + def set_layerwise_resident_layers(self, resident_layers: int) -> tuple[int, ...]: + """Apply one CLI-compatible resident-layer value to every layer group.""" + targets = tuple( + min(max(0, int(resident_layers)), manager.num_layers) + for manager in self.layerwise_offload_managers + ) + return self.set_layerwise_resident_layer_counts(targets) + + def set_layerwise_resident_layer_counts( + self, resident_layers: Sequence[int] + ) -> tuple[int, ...]: + """Apply an exact resident-layer count to every managed layer group.""" + if len(resident_layers) != len(self.layerwise_offload_managers): + raise ValueError("layerwise resident-layer group count changed") + previous = tuple( + manager.resident_layers for manager in self.layerwise_offload_managers + ) + updated: list[LayerwiseOffloadManager] = [] + try: + for manager, count in zip(self.layerwise_offload_managers, resident_layers): + manager.set_resident_layers(count) + updated.append(manager) + except Exception: + for manager, count in zip(updated, previous): + manager.set_resident_layers(count) + raise + return previous + + def layerwise_pinned_host_bytes(self) -> int: + return sum( + manager.pinned_host_weight_bytes() + for manager in self.layerwise_offload_managers + if manager.enabled + ) + + def layerwise_pinned_layers(self) -> tuple[tuple[int, ...], ...]: + return tuple( + manager.pinned_layer_indices() + for manager in self.layerwise_offload_managers + ) + + def _apply_layerwise_pinned_layers( + self, pinned_layers: Sequence[Sequence[int]] + ) -> None: + if len(pinned_layers) != len(self.layerwise_offload_managers): + raise ValueError("layerwise host-placement group count changed") + targets = [set(indices) for indices in pinned_layers] + # Release every allowance first. A later manager can then consume the + # bytes without the application order creating a transient rejection. + for manager, target in zip(self.layerwise_offload_managers, targets): + retained = set(manager.pinned_layer_indices()) & target + manager.set_pinned_layers(retained) + for manager, target in zip(self.layerwise_offload_managers, targets): + manager.set_pinned_layers(target) + + def set_layerwise_pinned_layers( + self, pinned_layers: Sequence[Sequence[int]] + ) -> tuple[tuple[int, ...], ...]: + previous = self.layerwise_pinned_layers() + try: + self._apply_layerwise_pinned_layers(pinned_layers) + except Exception: + self._apply_layerwise_pinned_layers(previous) + raise + return previous + + def release_layerwise_pins_outside( + self, pinned_layers: Sequence[Sequence[int]] + ) -> None: + """Release pins not used by a target before any new pins are acquired.""" + if len(pinned_layers) != len(self.layerwise_offload_managers): + raise ValueError("layerwise host-placement group count changed") + for manager, target in zip(self.layerwise_offload_managers, pinned_layers): + retained = set(manager.pinned_layer_indices()) & set(target) + manager.set_pinned_layers(retained) + + def restore_layerwise_pinned_layers( + self, pinned_layers: Sequence[Sequence[int]] + ) -> None: + self._apply_layerwise_pinned_layers(pinned_layers) + + def restore_layerwise_resident_layers(self, resident_layers: Sequence[int]) -> None: + """Restore exact per-group counts captured before an automatic change.""" + if len(resident_layers) != len(self.layerwise_offload_managers): + raise ValueError("layerwise resident-layer group count changed") + self.set_layerwise_resident_layer_counts(resident_layers) + def disable_offload(self) -> None: """Disable layerwise offload: load all layers to GPU and remove hooks. @@ -2108,18 +2448,18 @@ def _default_num_inference_steps() -> int: # back to 1 here turns the benefit ranking into a bare-bytes # ranking, and a once-per-request encoder can then outrank the # stepped DiT for the pin budget. - model_path = getattr(server_args, "model_path", None) + model_path = server_args.model_path if model_path: model_info = get_model_info( model_path, - backend=getattr(server_args, "backend", None), - model_id=getattr(server_args, "model_id", None), + backend=server_args.backend, + model_id=server_args.model_id, ) if model_info is not None: sampling_cls = model_info.sampling_param_cls if sampling_cls is None: return 1 - steps = getattr(sampling_cls(), "num_inference_steps", None) + steps = sampling_cls().num_inference_steps if not steps: return 1 return max(1, int(steps)) @@ -2153,8 +2493,15 @@ def _h2d_bytes_a_pin_would_save(name: str) -> int: key=_h2d_bytes_a_pin_would_save, reverse=True, ) - pin_budget = HostPinBudget() - logger.info("Layerwise offload host memory: %s", describe_host_memory()) + local_worker_count = max(1, server_args.num_gpus // server_args.nnodes) + pin_budget = HostPinBudget.for_local_worker(local_worker_count) + logger.info( + "Layerwise offload host memory: %s; this worker may pin %.2f GiB " + "after splitting the node allowance across %d local workers", + describe_host_memory(), + pin_budget.spendable_bytes / 1024**3, + local_worker_count, + ) for component_name in selected_pipeline_component_names: module = modules[component_name] diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py new file mode 100644 index 000000000000..1fc8be4509f4 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Joint placement optimization across every constrained memory resource.""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping + +import msgspec + + +class NoFeasiblePlacementError(ValueError): + pass + + +class PlacementOption(msgspec.Struct, frozen=True): + """One alternative placement for a component or layer group. + + Resource costs are deltas from the currently active placement. This is the + normalized form of absolute placement costs: the current placement is the + implicit zero option, positive values consume headroom, and negative values + release a resource. Keeping one resource map lets a single decision satisfy + rank/phase VRAM and node-scoped pinned-host constraints for one placement + state. A lifecycle with different load and serving placements must solve + those states separately and validate the transition between them. + + ``estimated_latency_savings`` is an integer supplied by the caller. It can + be transfer nanoseconds when measured bandwidth is available, or avoided + transfer bytes when only a stable ordering is known. + """ + + group_key: str + option_key: str + resource_delta_bytes: dict[str, int] + estimated_latency_savings: int + + +class PlacementPlan(msgspec.Struct, frozen=True): + resource_delta_bytes: dict[str, int] = {} + estimated_latency_savings: int = 0 + selections: list[PlacementOption] = [] + + +# Resource deltas, utility, selected options. +_State = tuple[tuple[int, ...], int, tuple[PlacementOption, ...]] + + +def _selection_key(state: _State) -> tuple[str, ...]: + return tuple(option.option_key for option in state[2]) + + +def _better_same_cost(candidate: _State, incumbent: _State) -> bool: + if candidate[1] != incumbent[1]: + return candidate[1] > incumbent[1] + return _selection_key(candidate) < _selection_key(incumbent) + + +def _dominates(candidate: _State, other: _State) -> bool: + """Whether candidate is at least as useful and no more expensive.""" + return candidate[1] >= other[1] and all( + left <= right for left, right in zip(candidate[0], other[0]) + ) + + +def _pareto_prune(states: Iterable[_State]) -> list[_State]: + """Drop states dominated across every constrained resource.""" + best_at_cost: dict[tuple[int, ...], _State] = {} + for state in states: + incumbent = best_at_cost.get(state[0]) + if incumbent is None or _better_same_cost(state, incumbent): + best_at_cost[state[0]] = state + + ordered = sorted( + best_at_cost.values(), + key=lambda state: ( + sum(state[0]), + state[0], + -state[1], + _selection_key(state), + ), + ) + frontier: list[_State] = [] + for state in ordered: + if any(_dominates(existing, state) for existing in frontier): + continue + frontier = [ + existing for existing in frontier if not _dominates(state, existing) + ] + frontier.append(state) + return frontier + + +def optimize_placement( + options: Iterable[PlacementOption], + *, + resource_budget_bytes: Mapping[str, int], +) -> PlacementPlan: + """Choose at most one alternative per group under every resource budget. + + Not selecting an alternative keeps that group at its current placement. + All alternatives are evaluated with the same selection vector, so every + constraint describing this placement state sees the same decision. A + choice may trade one resource for another by using a negative delta. The + result is exact for the supplied option frontier, resource model and + utility values; it does not imply that distinct load and serving states + should share one selection vector, or that the supplied frontier contains + every physically possible placement. + """ + grouped: dict[str, list[PlacementOption]] = {} + option_keys: set[str] = set() + all_resource_names = tuple(sorted(resource_budget_bytes)) + for option in options: + if option.option_key in option_keys: + raise ValueError(f"duplicate placement option key {option.option_key!r}") + option_keys.add(option.option_key) + unknown_resources = option.resource_delta_bytes.keys() - all_resource_names + if unknown_resources: + raise ValueError( + f"placement option {option.option_key!r} uses unknown resources: " + f"{sorted(unknown_resources)}" + ) + grouped.setdefault(option.group_key, []).append(option) + + # A resource whose worst possible positive use still fits cannot constrain + # the choice. Removing it before Pareto construction is exact and matters + # for the common case where a large-host server can pin every candidate: + # all pin-prefix alternatives then collapse to the highest-utility one for + # each GPU placement instead of multiplying the global frontier. + resource_names = tuple( + resource_name + for resource_name in all_resource_names + if sum( + max( + 0, + *( + option.resource_delta_bytes.get(resource_name, 0) + for option in group_options + ), + ) + for group_options in grouped.values() + ) + > resource_budget_bytes[resource_name] + ) + resource_budgets = tuple( + int(resource_budget_bytes[name]) for name in resource_names + ) + + ordered_groups = [] + for group_key in sorted(grouped): + group_options = sorted(grouped[group_key], key=lambda option: option.option_key) + local_states = [((0,) * len(resource_names), 0, ())] + local_states.extend( + ( + tuple( + option.resource_delta_bytes.get(resource_name, 0) + for resource_name in resource_names + ), + option.estimated_latency_savings, + (option,), + ) + for option in group_options + ) + ordered_groups.append( + [state[2][0] for state in _pareto_prune(local_states) if state[2]] + ) + minimum_remaining: list[tuple[int, ...]] = [ + (0,) * len(resource_names) for _ in range(len(ordered_groups) + 1) + ] + for group_index in range(len(ordered_groups) - 1, -1, -1): + group_options = ordered_groups[group_index] + next_minimum = minimum_remaining[group_index + 1] + minimum_remaining[group_index] = tuple( + next_delta + + min( + 0, + *( + option.resource_delta_bytes.get(resource_name, 0) + for option in group_options + ), + ) + for next_delta, resource_name in zip(next_minimum, resource_names) + ) + + frontier: list[_State] = [((0,) * len(resource_names), 0, ())] + for group_index, group_options in enumerate(ordered_groups): + candidates = list(frontier) + for state in frontier: + for option in group_options: + resource_deltas = tuple( + current + option.resource_delta_bytes.get(resource_name, 0) + for current, resource_name in zip(state[0], resource_names) + ) + candidates.append( + ( + resource_deltas, + state[1] + option.estimated_latency_savings, + state[2] + (option,), + ) + ) + remaining = minimum_remaining[group_index + 1] + reachable = [ + state + for state in candidates + if all( + used + releasable <= budget + for used, releasable, budget in zip( + state[0], remaining, resource_budgets + ) + ) + ] + frontier = _pareto_prune(reachable) + + frontier = [ + state + for state in frontier + if all(used <= budget for used, budget in zip(state[0], resource_budgets)) + ] + if not frontier: + raise NoFeasiblePlacementError("no placement satisfies all resource budgets") + + best = min( + frontier, + key=lambda state: ( + -state[1], + sum(state[0]), + state[0], + _selection_key(state), + ), + ) + full_resource_deltas = { + resource_name: sum( + option.resource_delta_bytes.get(resource_name, 0) for option in best[2] + ) + for resource_name in all_resource_names + } + return PlacementPlan( + resource_delta_bytes=full_resource_deltas, + estimated_latency_savings=best[1], + selections=list(best[2]), + ) diff --git a/python/sglang/multimodal_gen/runtime/scheduler_client.py b/python/sglang/multimodal_gen/runtime/scheduler_client.py index 1aade5875ec0..5d2c95617e9c 100644 --- a/python/sglang/multimodal_gen/runtime/scheduler_client.py +++ b/python/sglang/multimodal_gen/runtime/scheduler_client.py @@ -1,3 +1,4 @@ +import asyncio import itertools import pickle import time @@ -135,6 +136,11 @@ def _select_replica(batch: Any, dp_size: int, counter: "itertools.count") -> int return next(counter) % dp_size +def _is_warmup_batch(batch: Any) -> bool: + reqs = batch if isinstance(batch, list) else [batch] + return bool(reqs) and all(isinstance(req, Req) and req.is_warmup for req in reqs) + + def _merge_fanout_results(results: list[Any]) -> Any: """One reply for a control op sent to every replica: the first error wins, because "succeeded" must mean succeeded everywhere.""" @@ -144,6 +150,62 @@ def _merge_fanout_results(results: list[Any]) -> Any: return results[0] +def _auto_residency_status(result: Any) -> str | None: + if not isinstance(result, OutputBatch) or not isinstance(result.output, dict): + return None + return result.output.get("status") + + +def _merge_auto_residency_results(results: list[Any]) -> Any: + """Preserve the strongest replica-wide auto-residency outcome. + + Replicas plan against their own measured headroom, so one may keep the + original placement while another promotes. If any replica promoted, the + orchestrator must still run the validation warmup everywhere. A failed + local rollback remains fatal and takes precedence over that success. + """ + rollback_failed = next( + ( + result + for result in results + if _auto_residency_status(result) == "rollback_failed" + ), + None, + ) + if rollback_failed is not None: + return rollback_failed + + promoted = next( + (result for result in results if _auto_residency_status(result) == "promoted"), + None, + ) + if promoted is not None: + replica_errors = [ + result.error + for result in results + if isinstance(result, OutputBatch) and result.error is not None + ] + if replica_errors: + logger.warning( + "Auto residency kept the original placement on %d replica(s): %s", + len(replica_errors), + replica_errors[0], + ) + return promoted + + rolled_back = next( + ( + result + for result in results + if _auto_residency_status(result) == "rolled_back" + ), + None, + ) + if rolled_back is not None: + return rolled_back + return _merge_fanout_results(results) + + class SchedulerClient: """ A synchronous, singleton client for communicating with the Scheduler service. @@ -189,9 +251,13 @@ def _forward_one(self, endpoint: str, batch: Any, timeout_ms: int | None) -> Any def _forward_routed(self, batch: Any, timeout_ms: int | None) -> Any: self.request_logger.log_received_request(batch) endpoints = self.server_args.scheduler_endpoints - if isinstance(batch, _CONTROL_REQ_TYPES): + if isinstance(batch, _CONTROL_REQ_TYPES) or _is_warmup_batch(batch): results = [self._forward_one(ep, batch, timeout_ms) for ep in endpoints] - output_batch = _merge_fanout_results(results) + output_batch = ( + _merge_auto_residency_results(results) + if isinstance(batch, AutoResidencyReq) + else _merge_fanout_results(results) + ) else: replica = _select_replica(batch, len(endpoints), self._replica_counter) output_batch = self._forward_one(endpoints[replica], batch, timeout_ms) @@ -263,12 +329,20 @@ async def forward(self, batch: Any, timeout_ms: int | None = None) -> Any: ) endpoints = self.server_args.scheduler_endpoints - if isinstance(batch, _CONTROL_REQ_TYPES): - # replica state (weights, LoRA, memory) must change everywhere - results = [ - await self._forward_one(ep, batch, timeout_ms) for ep in endpoints - ] - output_batch = _merge_fanout_results(results) + if isinstance(batch, _CONTROL_REQ_TYPES) or _is_warmup_batch(batch): + # Every replica must receive the state change or calibration + # request; the measured values and selected placements may differ. + results = await asyncio.gather( + *( + self._forward_one(endpoint, batch, timeout_ms) + for endpoint in endpoints + ) + ) + output_batch = ( + _merge_auto_residency_results(results) + if isinstance(batch, AutoResidencyReq) + else _merge_fanout_results(results) + ) else: replica = _select_replica(batch, len(endpoints), self._replica_counter) output_batch = await self._forward_one( diff --git a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py index 350c27b4548e..ee9fde16621a 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -10,6 +10,13 @@ from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import ( ModelDeploymentConfig, ) +from sglang.multimodal_gen.configs.quantization.nunchaku import ( + NunchakuSVDQuantArgs, +) +from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType +from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.registry import ( + has_realtime_model_adapter, +) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( LAYERWISE_OFFLOAD_ALL_COMPONENTS, LAYERWISE_OFFLOAD_DIT_GROUP, @@ -41,6 +48,54 @@ DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB = 120.0 +def auto_residency_args_skip_reason(server_args: ServerArgs) -> str | None: + """Return why args cannot use warmup-calibrated residency.""" + if envs.SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY: + return "disabled via SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY" + if server_args.performance_mode != "auto": + return f"performance_mode={server_args.performance_mode}" + if ( + server_args.warmup_mode != "server" + or server_args.disagg_role != RoleType.MONOLITHIC + ): + return "no synthetic server warmup to calibrate from" + task_type = server_args.pipeline_config.task_type + if not (task_type.is_visual_gen() or task_type.is_mesh_gen()): + return "no synthetic server warmup to calibrate from" + if has_realtime_model_adapter(server_args): + return "realtime serving has no representative synthetic warmup" + if server_args.backend == "diffusers": + return "diffusers backend" + if server_args.enable_breakable_cuda_graph: + return "breakable CUDA graph captures during warmup" + if server_args.enable_torch_compile: + # Compile warmup temporarily evicts resident auxiliaries and may + # layerwise-offload the DiT, so its peak is not a serving peak. + return "torch.compile warmup uses a stripped memory layout" + if envs.SGLANG_CACHE_DIT_ENABLED: + return "cache-dit enabled" + if server_args.batching_max_size > 1: + return "dynamic batching enabled" + nunchaku = server_args.nunchaku_config + svdquant_enabled = nunchaku is not None and ( + not isinstance(nunchaku, NunchakuSVDQuantArgs) or nunchaku.enable_svdquant + ) + if ( + server_args.quantization is not None + or server_args.component_quantizations + or server_args.transformer_weights_path + or svdquant_enabled + ): + # Residency changes can perturb quantized numerical paths through a + # different allocation and kernel layout. + return "quantized checkpoint" + if server_args.direct_gpu_weight_loading: + return "direct GPU weight loading requires a fixed resident placement" + if not current_platform.is_cuda(): + return "requires CUDA" + return None + + class ServerArgsAutoTuner: """Auto-tunes the server-arg for the given performance-mode, based on practical deployment experience with different model architectures""" @@ -51,6 +106,10 @@ def __init__(self, server_args: ServerArgs): def _deployment_config(self) -> ModelDeploymentConfig: return self.server_args.pipeline_config.get_model_deployment_config() + def _uses_warmup_calibrated_residency(self) -> bool: + """Whether warmup will replace coarse pre-load residency defaults.""" + return auto_residency_args_skip_reason(self.server_args) is None + def _resolve_keep_resident_min_available_gb( self, deployment_config: ModelDeploymentConfig ) -> float | None: @@ -123,6 +182,11 @@ def maybe_adjust_auto_component_residency_after_offload(self) -> None: args = self.server_args if args.performance_mode != "auto" or current_platform.is_cpu(): return + if self._uses_warmup_calibrated_residency(): + # Keep the load-safe placement until warmup has measured the real + # workload. The post-warmup planner replaces model/card thresholds + # with component sizes and per-phase headroom. + return # Explicit placement is component-scoped; unmatched components still # receive automatic defaults. @@ -517,6 +581,8 @@ def _filter_high_memory_resident_components( args = self.server_args if args.performance_mode != "auto" or current_platform.is_cpu(): return components + if self._uses_warmup_calibrated_residency(): + return components deployment_config = self._deployment_config() threshold_gb = self._resolve_keep_resident_min_available_gb(deployment_config) 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 820e682bc341..4ac80edffcef 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -417,6 +417,10 @@ class ServerArgs(DisaggServerArgsMixin): warmup_resolutions: list[str] = None warmup_steps: int = 1 + # JSON overrides for the representative request shape used by synthetic + # warmup and automatic residency planning. Execution remains bounded by + # warmup_steps and the server warmup frame/area caps. + warmup_sampling_params: dict[str, Any] | str | None = None disable_autocast: bool | None = None @@ -2253,6 +2257,18 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: default=ServerArgs.warmup_steps, help="The number of warmup steps to perform for each resolution.", ) + parser.add_argument( + "--warmup-sampling-params", + type=str, + default=ServerArgs.warmup_sampling_params, + help=( + "JSON object overriding model sampling defaults for synthetic " + "warmup and auto residency planning, for example " + '\'{"width":832,"height":480,"num_frames":9,' + '"num_inference_steps":4}\'. Warmup still applies its ' + "bounded execution caps." + ), + ) # component residency and legacy offload controls parser.add_argument( "--component-residency", diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index 82df3e9a7ed0..2dfc67c07711 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -7,11 +7,17 @@ from tqdm.auto import tqdm +from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.registry import ( + has_realtime_model_adapter, +) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import ( OutputBatch, Req, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.server_args.auto_tune import ( + auto_residency_args_skip_reason, +) from sglang.multimodal_gen.runtime.utils.image_io import save_base64_image_to_path from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.warmup_request_builder import ( @@ -75,15 +81,7 @@ def should_run_server_warmup(server_args: ServerArgs) -> bool: def is_realtime_serving(server_args: ServerArgs) -> bool: """Synthetic warmup has no realtime session state.""" - try: - from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.registry import ( - get_realtime_model_adapter, - ) - - get_realtime_model_adapter(server_args) - return True - except Exception: - return False + return has_realtime_model_adapter(server_args) def should_run_synthetic_server_warmup(server_args: ServerArgs) -> bool: @@ -103,64 +101,17 @@ def should_run_explicit_client_warmup(server_args: ServerArgs) -> bool: def auto_residency_skip_reason(server_args: ServerArgs) -> str | None: - """Server-args-only gate for the warmup-calibrated residency promotion. + """Final gate for warmup-calibrated residency promotion. Only rules out paths the promotion was not designed for; the workers re-check per component (explicit placement, FSDP modules, custom strategies, missing sizes) and per measurement. """ - from sglang.multimodal_gen import envs - from sglang.multimodal_gen.configs.quantization.nunchaku import ( - NunchakuSVDQuantArgs, - ) - from sglang.multimodal_gen.runtime.platforms import current_platform - - if envs.SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY: - return "disabled via SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY" - if server_args.performance_mode != "auto": - return f"performance_mode={server_args.performance_mode}" + args_reason = auto_residency_args_skip_reason(server_args) + if args_reason is not None: + return args_reason if not should_run_synthetic_server_warmup(server_args): return "no synthetic server warmup to calibrate from" - if server_args.backend == "diffusers": - return "diffusers backend" - if server_args.enable_breakable_cuda_graph: - return "breakable CUDA graph captures during warmup" - if server_args.enable_torch_compile: - # Compile warmup runs a deliberately stripped memory layout: the DiT - # is temporarily layerwise-offloaded (offload_during_compile) and - # _move_resident_components_for_warmup evicts resident aux components - # for the denoise. Peaks measured there are far below a real - # request's, and the post-promotion re-warm (also a warmup layout) - # cannot catch the difference. - return "torch.compile warmup uses a stripped memory layout" - if envs.SGLANG_CACHE_DIT_ENABLED: - return "cache-dit enabled" - if server_args.batching_max_size > 1: - return "dynamic batching enabled" - if (server_args.dp_size or 1) > 1: - return "dp replicas warm up independently" - if (server_args.ulysses_degree or 1) > 1: - # Sequence-parallel workers can have a different execution order - # after residency changes; keep their configured placement until that - # interaction is explicitly calibrated and validated. - return "Ulysses sequence parallelism" - if server_args.use_fsdp_inference: - return "FSDP inference" - nunchaku = server_args.nunchaku_config - svdquant_enabled = nunchaku is not None and ( - not isinstance(nunchaku, NunchakuSVDQuantArgs) or nunchaku.enable_svdquant - ) - if ( - server_args.quantization is not None - or server_args.transformer_weights_path - or svdquant_enabled - ): - # Residency changes shift the memory layout, which measurably moved - # fp8-quantized DiT outputs in the past (see PR #29649); keep - # quantized checkpoints on their configured strategy. - return "quantized checkpoint" - if not current_platform.is_cuda(): - return "requires CUDA" return None diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py index e473840ea2f6..ee5558eb9582 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -13,8 +13,9 @@ request transport path as real generation. """ +import json from copy import copy -from dataclasses import replace +from dataclasses import fields, replace from typing import Any from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType @@ -29,6 +30,9 @@ ServerArgs, is_ltx2_two_stage_pipeline_name, ) +from sglang.multimodal_gen.runtime.server_args.auto_tune import ( + auto_residency_args_skip_reason, +) from sglang.multimodal_gen.runtime.utils.common import parse_size from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -59,13 +63,44 @@ def get_model_sampling_defaults(server_args: ServerArgs) -> SamplingParams: config_classes = get_pipeline_config_classes(pipeline_class_name) if config_classes is not None: _, sampling_params_cls = config_classes - return sampling_params_cls() + defaults = sampling_params_cls() + return _apply_warmup_sampling_overrides(server_args, defaults) - return SamplingParams.from_pretrained( + defaults = SamplingParams.from_pretrained( server_args.model_path, backend=server_args.backend, model_id=server_args.model_id, ) + return _apply_warmup_sampling_overrides(server_args, defaults) + + +def _apply_warmup_sampling_overrides( + server_args: ServerArgs, defaults: SamplingParams +) -> SamplingParams: + value = server_args.warmup_sampling_params + if value is None: + return defaults + if isinstance(value, str): + try: + value = json.loads(value) + except json.JSONDecodeError as error: + raise ValueError( + "--warmup-sampling-params must be a JSON object" + ) from error + if not isinstance(value, dict): + raise ValueError("--warmup-sampling-params must be a JSON object") + field_names = {item.name for item in fields(defaults)} + unknown = value.keys() - field_names + if unknown: + raise ValueError( + "invalid --warmup-sampling-params fields: " f"{', '.join(sorted(unknown))}" + ) + updated = copy(defaults) + for name, field_value in value.items(): + # Some model contracts intentionally expose fixed dataclass fields + # with init=False; a warmup workload still needs to mirror the request. + object.__setattr__(updated, name, field_value) + return updated def _resolve_default_warmup_resolution( @@ -359,22 +394,17 @@ def _resolve_calibration_num_frames( """A second, smaller warmup frame count for auto-residency calibration. Only produced when the main video warmup was frame-capped and the - auto-residency promotion is actually eligible to consume the measurement - (``auto_residency_skip_reason``): the estimator then has two workload + auto-residency promotion is actually eligible to consume the measurement: + the estimator then has two workload sizes to fit the constant/linear split of the peak. Returns None when a second measurement adds no information (same adjusted frames) or would never be read. """ - # local import: server_warmup imports this module at module level - from sglang.multimodal_gen.runtime.server_warmup import ( - auto_residency_skip_reason, - ) - if not server_based_warmup or not _is_video_warmup_task(server_args): return None if warmup_num_frames is None: return None - if auto_residency_skip_reason(server_args) is not None: + if auto_residency_args_skip_reason(server_args) is not None: return None default_num_frames = sampling_defaults.num_frames or 1 if warmup_num_frames >= default_num_frames: diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index d1f524dafad4..829ea67dfe45 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -10,6 +10,7 @@ import json import os import queue +import shlex import threading import time from pathlib import Path @@ -111,6 +112,36 @@ def _print_case_log_separator(case_id: str, state: str) -> None: ) +def _case_warmup_sampling_params(case: DiffusionTestCase) -> dict[str, Any]: + """Representative workload passed to startup residency calibration.""" + sampling = case.sampling_params + if sampling is None: + return {} + overrides: dict[str, Any] = {} + if sampling.output_size: + width, height = sampling.output_size.lower().split("x", 1) + overrides.update(width=int(width), height=int(height)) + if sampling.num_frames is not None: + overrides["num_frames"] = sampling.num_frames + elif sampling.fps is not None: + # /v1/videos resolves an omitted num_frames from seconds * fps. + overrides["num_frames"] = sampling.seconds * sampling.fps + if sampling.fps is not None: + overrides["fps"] = sampling.fps + for name in ("num_inference_steps", "guidance_scale", "true_cfg_scale"): + value = sampling.extras.get(name) + if value is not None: + overrides[name] = value + return overrides + + +def _has_server_option(extra_args: str, case_extras: list[str], option: str) -> bool: + tokens = shlex.split(extra_args) + for extra in case_extras: + tokens.extend(shlex.split(extra)) + return option in tokens or any(token.startswith(f"{option}=") for token in tokens) + + @pytest.fixture def diffusion_server(case: DiffusionTestCase) -> ServerContext: """Start a diffusion server for a single case and tear it down afterwards.""" @@ -135,6 +166,15 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: extra_args = os.environ.get("SGLANG_TEST_SERVE_ARGS", "") extra_args = f"--model-type diffusion {extra_args}".strip() + has_explicit_warmup_workload = _has_server_option( + extra_args, server_args.extras, "--warmup-sampling-params" + ) + if not has_explicit_warmup_workload: + warmup_sampling_params = _case_warmup_sampling_params(case) + if warmup_sampling_params: + encoded = json.dumps(warmup_sampling_params, separators=(",", ":")) + extra_args += f" --warmup-sampling-params {shlex.quote(encoded)}" + extra_args += f" --num-gpus {server_args.num_gpus}" if server_args.tp_size is not None: diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py index 4cecc97554f4..6d8d5e3d161a 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py @@ -35,6 +35,7 @@ ) from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.registry import ( get_realtime_model_adapter, + has_realtime_model_adapter, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world.lingbot_world_causal_denoising import ( @@ -976,6 +977,7 @@ def test_realtime_input_validation_reuses_generator_across_chunks(): def test_realtime_registry_resolves_lingbot_adapter(): server_args = SimpleNamespace(pipeline_config=LingBotWorldCausalDMDConfig()) + assert has_realtime_model_adapter(server_args) adapter = get_realtime_model_adapter(server_args) assert isinstance(adapter, lingbot_realtime.LingBotWorldRealtimeAdapter) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index a65a7b2ebde4..b43ccfb57811 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -9,6 +9,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType from sglang.multimodal_gen.configs.pipeline_configs.longlive2 import LongLive2T2VConfig +from sglang.multimodal_gen.runtime.layers.linear import LinearBase from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( ACTIVATION_EXTRAPOLATION_MARGIN, GIB_BYTES, @@ -22,8 +23,11 @@ apply_promotions, collect_promotion_candidates, component_resident_size_bytes, + component_runtime_weight_bytes, estimate_default_workload_peak_bytes, + estimate_workload_phase_peaks, format_applied_changes, + module_uses_quantized_weights, plan_auto_residency, rank_candidates_by_h2d_savings, rollback_promotions, @@ -140,6 +144,37 @@ def test_negative_slope_falls_back_to_single_point_formula(self): ) assert estimate == fallback + def test_weight_floor_prevents_scaling_constant_component_memory(self): + small = WarmupMemoryRecord( + width=832, + height=480, + num_frames=9, + baseline_allocated_bytes=GIB_BYTES, + peak_reserved_bytes=int(29.5 * GIB_BYTES), + succeeded=True, + ) + large = WarmupMemoryRecord( + width=832, + height=480, + num_frames=17, + baseline_allocated_bytes=GIB_BYTES, + peak_reserved_bytes=int(29.4 * GIB_BYTES), + succeeded=True, + ) + target_units = 1024 * 1024 * 81 + + estimate = estimate_default_workload_peak_bytes( + records=[small, large], + target_units=target_units, + constant_weight_bytes=int(29.4 * GIB_BYTES), + ) + fallback_without_weights = estimate_default_workload_peak_bytes( + records=[small, large], target_units=target_units + ) + + assert estimate < 50 * GIB_BYTES + assert estimate < fallback_without_weights + def test_covering_measurement_bounds_the_target(self): capped = _record(num_frames=17, peak_gib=12) full = _record(width=1280, height=720, num_frames=81, peak_gib=30) @@ -195,6 +230,51 @@ def test_no_records_disables_estimation(self): estimate_default_workload_peak_bytes(records=[], target_units=None) is None ) + def test_phase_estimation_preserves_active_component_membership(self): + small = _record(num_frames=9, peak_gib=30) + large = _record(num_frames=17, peak_gib=32) + small = WarmupMemoryRecord( + width=small.width, + height=small.height, + num_frames=small.num_frames, + baseline_allocated_bytes=small.baseline_allocated_bytes, + peak_reserved_bytes=small.peak_reserved_bytes, + succeeded=small.succeeded, + phase_peak_reserved_bytes={"denoise": 30 * GIB_BYTES}, + phase_active_components={"denoise": ("transformer",)}, + ) + large = WarmupMemoryRecord( + width=large.width, + height=large.height, + num_frames=large.num_frames, + baseline_allocated_bytes=large.baseline_allocated_bytes, + peak_reserved_bytes=large.peak_reserved_bytes, + succeeded=large.succeeded, + phase_peak_reserved_bytes={"denoise": 32 * GIB_BYTES}, + phase_active_components={"denoise": ("transformer",)}, + ) + + peaks, active = estimate_workload_phase_peaks( + records=[small, large], + target_units=832 * 480 * 81, + component_weight_bytes={"transformer": 28 * GIB_BYTES}, + ) + + assert peaks["denoise"] >= 32 * GIB_BYTES + assert active == {"denoise": ("transformer",)} + + +def test_loaded_native_quantization_config_is_detected_with_bf16_weights(): + linear = LinearBase.__new__(LinearBase) + nn.Module.__init__(linear) + linear.quant_config = object() + linear.register_parameter( + "weight", + nn.Parameter(torch.ones(1, dtype=torch.bfloat16), requires_grad=False), + ) + + assert module_uses_quantized_weights(linear) + def _candidate( name: str, @@ -203,12 +283,18 @@ def _candidate( weight_gib: int = 10, h2d_gib: int | None = None, ) -> PromotionCandidate: + layerwise = mode == LAYERWISE_OFFLOAD return PromotionCandidate( component_name=name, residency_mode=mode, promoted_weight_bytes=weight_gib * GIB_BYTES, h2d_bytes_per_request=(h2d_gib if h2d_gib is not None else weight_gib) * GIB_BYTES, + target_layerwise_resident_layers=(1,) if layerwise else None, + target_layerwise_pinned_layers=((),) if layerwise else None, + permanent_residency=True, + active_device_delta_bytes=0, + inactive_device_delta_bytes=weight_gib * GIB_BYTES, ) @@ -219,6 +305,12 @@ def _report( estimated_gib: int | None = 50, candidates: list[PromotionCandidate] | None = None, skip_reason: str | None = None, + phase_peaks_gib: dict[str, int] | None = None, + phase_components: dict[str, tuple[str, ...]] | None = None, + node_rank: int = 0, + pinned_host_gib: int = 0, + host_pin_capacity_gib: int = 0, + host_transition_headroom_gib: int = 0, ) -> RankResidencyReport: return RankResidencyReport( rank=rank, @@ -226,6 +318,14 @@ def _report( estimated_peak_bytes=( None if estimated_gib is None else estimated_gib * GIB_BYTES ), + estimated_peak_bytes_by_phase={ + name: value * GIB_BYTES for name, value in (phase_peaks_gib or {}).items() + }, + active_components_by_phase=phase_components or {}, + node_rank=node_rank, + pinned_host_bytes=pinned_host_gib * GIB_BYTES, + host_pin_capacity_bytes=host_pin_capacity_gib * GIB_BYTES, + host_transition_headroom_bytes=(host_transition_headroom_gib * GIB_BYTES), candidates=candidates if candidates is not None else [_candidate("vae")], skip_reason=skip_reason, ) @@ -306,12 +406,250 @@ def test_no_candidates_skips(self): plan = plan_auto_residency(reports=[_report(candidates=[])]) assert plan.skip_reason is not None + def test_component_active_phase_is_not_double_counted(self): + # Cosmos3-Super TP2: the 66 GiB denoise peak already includes the + # roughly 60 GiB/rank DiT. Keeping it resident adds those bytes only + # to the 4 GiB idle phase, matching the measured 73 GiB resident peak + # instead of inventing a 126 GiB denoise peak. + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=139, + estimated_gib=66, + candidates=[_candidate("transformer", weight_gib=60)], + phase_peaks_gib={"denoise": 66, "idle": 4}, + phase_components={"denoise": ("transformer",), "idle": ()}, + ) + ] + ) + + assert [candidate.component_name for candidate in plan.promotions] == [ + "transformer" + ] + + def test_other_phase_can_block_component_residency(self): + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=139, + estimated_gib=90, + candidates=[_candidate("transformer", weight_gib=60)], + phase_peaks_gib={"denoise": 66, "decode": 90}, + phase_components={"denoise": ("transformer",), "decode": ()}, + ) + ] + ) + + assert not plan.promotions + + def test_redundant_phase_constraints_collapse_to_the_highest_peak(self): + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=100, + estimated_gib=70, + candidates=[_candidate("transformer", weight_gib=10)], + phase_peaks_gib={ + "denoise:first": 60, + "denoise:later": 70, + "decode": 50, + }, + phase_components={ + "denoise:first": ("transformer",), + "denoise:later": ("transformer",), + "decode": ("vae",), + }, + ) + ] + ) + + assert "gpu:rank0:denoise:first" not in plan.resource_budget_bytes + assert plan.resource_budget_bytes["gpu:rank0:denoise:later"] == (20 * GIB_BYTES) + assert "gpu:rank0:decode" in plan.resource_budget_bytes + + def test_negative_gpu_headroom_is_not_treated_as_zero_capacity(self): + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=80, + estimated_gib=76, + candidates=[_candidate("transformer", weight_gib=10)], + phase_peaks_gib={"denoise": 76}, + phase_components={"denoise": ("transformer",)}, + ) + ] + ) + + assert plan.promotions == [] + assert plan.skip_reason == "no placement satisfies all resource budgets" + + def test_each_rank_keeps_a_reserve_scaled_to_its_own_gpu(self): + candidate = _candidate("vae", weight_gib=1) + plan = plan_auto_residency( + reports=[ + _report( + rank=0, + budget_gib=40, + estimated_gib=20, + candidates=[candidate], + phase_peaks_gib={"decode": 20}, + phase_components={"decode": ("vae",)}, + ), + _report( + rank=1, + budget_gib=140, + estimated_gib=100, + candidates=[candidate], + phase_peaks_gib={"decode": 100}, + phase_components={"decode": ("vae",)}, + ), + ] + ) + + assert plan.resource_budget_bytes["gpu:rank0:decode"] == 16 * GIB_BYTES + assert plan.resource_budget_bytes["gpu:rank1:decode"] == 26 * GIB_BYTES + assert plan.reserve_bytes == 14 * GIB_BYTES + + def test_partial_layer_residency_is_selected_when_full_dit_does_not_fit(self): + full = _candidate( + "transformer", + mode=LAYERWISE_OFFLOAD, + weight_gib=30, + h2d_gib=300, + ) + partial = PromotionCandidate( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + promoted_weight_bytes=10 * GIB_BYTES, + h2d_bytes_per_request=90 * GIB_BYTES, + target_layerwise_resident_layers=(12,), + target_layerwise_pinned_layers=((),), + active_device_delta_bytes=10 * GIB_BYTES, + ) + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=100, + estimated_gib=70, + candidates=[full, partial], + phase_peaks_gib={"denoise": 70}, + phase_components={"denoise": ("transformer",)}, + ) + ] + ) + + assert len(plan.promotions) == 1 + assert plan.promotions[0].target_layerwise_resident_layers == (12,) + + def test_node_hostpin_is_optimized_with_the_same_selection_vector(self): + cold_pageable = PromotionCandidate( + component_name="cold_encoder", + residency_mode=LAYERWISE_OFFLOAD, + promoted_weight_bytes=0, + h2d_bytes_per_request=-GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + pinned_host_delta_bytes=-10 * GIB_BYTES, + ) + hot_pinned = PromotionCandidate( + component_name="hot_dit", + residency_mode=LAYERWISE_OFFLOAD, + promoted_weight_bytes=0, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=10 * GIB_BYTES, + ) + plan = plan_auto_residency( + reports=[ + _report( + pinned_host_gib=10, + host_pin_capacity_gib=10, + candidates=[cold_pageable, hot_pinned], + ) + ] + ) + + assert [candidate.component_name for candidate in plan.promotions] == [ + "cold_encoder", + "hot_dit", + ] + assert plan.resource_budget_bytes["hostpin:node0"] == 0 + + def test_dynamic_hostpin_is_not_planned_across_workers(self): + pin_more = PromotionCandidate( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + promoted_weight_bytes=0, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=10 * GIB_BYTES, + ) + plan = plan_auto_residency( + reports=[ + _report( + rank=0, + host_pin_capacity_gib=20, + candidates=[pin_more], + ), + _report( + rank=1, + host_pin_capacity_gib=20, + candidates=[pin_more], + ), + ] + ) + + assert plan.promotions == [] + assert plan.skip_reason == ( + "dynamic HostPin placement requires a single-worker node" + ) + + def test_hostpin_repack_must_fit_transition_headroom(self): + pin_more = PromotionCandidate( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + promoted_weight_bytes=0, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=4 * GIB_BYTES, + host_pin_scratch_bytes=6 * GIB_BYTES, + ) + plan = plan_auto_residency( + reports=[ + _report( + host_pin_capacity_gib=10, + host_transition_headroom_gib=5, + candidates=[pin_more], + ) + ] + ) + + assert plan.promotions == [] + assert plan.resource_budget_bytes["hostram:node0:pin"] == 5 * GIB_BYTES + class _FakeLayerwiseManager: - def __init__(self, tensors: dict[str, torch.Tensor]): + def __init__( + self, + tensors: dict[str, torch.Tensor], + *, + resident_layers: int = 0, + event_log: list[str] | None = None, + event_name: str = "manager", + ): self.enabled = True self._configured = True self._tensors = tensors + self.num_layers = max(1, len(tensors)) + self.resident_layers = resident_layers + self.residency_policy = "leading" + self.pin_cpu_memory = True + self._pinned_layers: tuple[int, ...] = () + self._event_log = event_log + self._event_name = event_name self.fail_load = False self.load_all_layers_calls = 0 self.remove_hooks_calls = 0 @@ -327,6 +665,57 @@ def offloaded_weight_bytes(self): tensor.numel() * tensor.element_size() for tensor in self._tensors.values() ) + def resident_weight_bytes(self, resident_layers=None): + count = self.resident_layers if resident_layers is None else resident_layers + return ( + self.offloaded_weight_bytes() + * min(count, self.num_layers) + // self.num_layers + ) + + def peak_managed_device_weight_bytes(self, resident_layers=None): + count = self.resident_layers if resident_layers is None else resident_layers + return ( + self.offloaded_weight_bytes() + * min(count + 1, self.num_layers) + // self.num_layers + ) + + def layer_weight_bytes(self): + total = self.offloaded_weight_bytes() + base, remainder = divmod(total, self.num_layers) + return { + layer_idx: base + (1 if layer_idx < remainder else 0) + for layer_idx in range(self.num_layers) + } + + def layer_host_store_bytes(self): + return self.layer_weight_bytes() + + def pinned_host_weight_bytes(self): + layer_bytes = self.layer_weight_bytes() + return sum(layer_bytes[layer_idx] for layer_idx in self._pinned_layers) + + def pinnable_layer_indices(self): + return tuple(range(self.num_layers)) + + def pinned_layer_indices(self): + return self._pinned_layers + + def set_pinned_layers(self, pinned_layers): + previous = self._pinned_layers + target = tuple(sorted(pinned_layers)) + if target != previous and self._event_log is not None: + action = "pin" if target else "unpin" + self._event_log.append(f"{self._event_name}:{action}") + self._pinned_layers = target + return previous + + def set_resident_layers(self, resident_layers): + previous = self.resident_layers + self.resident_layers = min(max(0, resident_layers), self.num_layers) + return previous + def load_all_layers(self): self.load_all_layers_calls += 1 if self.fail_load: @@ -388,6 +777,20 @@ def test_layerwise_counts_manager_cpu_buffers_not_placeholders(self): expected = sum(t.numel() * t.element_size() for t in cpu_weights.values()) assert component_resident_size_bytes(module, LAYERWISE_OFFLOAD) == expected + def test_layerwise_runtime_floor_uses_only_the_streaming_window(self): + manager = _FakeLayerwiseManager( + { + "layers.0.w": torch.zeros(16), + "layers.1.w": torch.zeros(16), + "layers.2.w": torch.zeros(16), + } + ) + module = _FakeLayerwiseDit([manager]) + + runtime = component_runtime_weight_bytes({"transformer": module}) + + assert runtime["transformer"] == manager.peak_managed_device_weight_bytes() + class TestCollectPromotionCandidates: def _modes(self, mapping): @@ -424,15 +827,24 @@ def test_filters_and_sizes(self): custom_strategy_names={"custom"}, num_inference_steps=50, ) - by_name = {c.component_name: c for c in candidates} + permanent_candidates = { + candidate.component_name: candidate + for candidate in candidates + if candidate.permanent_residency + } # resident vae, explicit placement, custom strategy, non-module: out - assert set(by_name) == {"text_encoder", "transformer"} - assert by_name["text_encoder"].h2d_bytes_per_request == ( - by_name["text_encoder"].promoted_weight_bytes + assert set(permanent_candidates) == {"text_encoder", "transformer"} + assert permanent_candidates["text_encoder"].h2d_bytes_per_request == ( + permanent_candidates["text_encoder"].promoted_weight_bytes ) # layerwise DiT re-streams its layers once per denoise step - assert by_name["transformer"].h2d_bytes_per_request == ( - by_name["transformer"].promoted_weight_bytes * 50 + assert permanent_candidates["transformer"].h2d_bytes_per_request == ( + permanent_candidates["transformer"].promoted_weight_bytes * 100 + ) + assert any( + candidate.component_name == "transformer" + and candidate.target_layerwise_resident_layers == (1,) + for candidate in candidates ) def test_mechanism_mismatch_is_not_a_candidate(self): @@ -450,6 +862,56 @@ def test_mechanism_mismatch_is_not_a_candidate(self): ) assert candidates == [] + def test_fully_pinned_component_can_release_hostpin_for_a_hotter_component(self): + manager = _FakeLayerwiseManager( + { + "layers.0.w": torch.zeros(16), + "layers.1.w": torch.zeros(16), + } + ) + manager._pinned_layers = (0, 1) + module = _FakeLayerwiseDit([manager]) + + candidates = collect_promotion_candidates( + modules={"text_encoder": module}, + residency_mode_of=self._modes({"text_encoder": LAYERWISE_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=10, + ) + + assert any( + candidate.target_layerwise_pinned_layers == ((),) + and candidate.pinned_host_delta_bytes < 0 + and candidate.host_unpin_scratch_bytes > 0 + for candidate in candidates + ) + + def test_multi_worker_candidate_keeps_existing_hostpin_placement(self): + manager = _FakeLayerwiseManager( + { + "layers.0.w": torch.zeros(16), + "layers.1.w": torch.zeros(16), + } + ) + manager._pinned_layers = (0, 1) + module = _FakeLayerwiseDit([manager]) + + candidates = collect_promotion_candidates( + modules={"transformer": module}, + residency_mode_of=self._modes({"transformer": LAYERWISE_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=10, + allow_host_pin_reallocation=False, + ) + permanent = next( + candidate for candidate in candidates if candidate.permanent_residency + ) + + assert permanent.target_layerwise_pinned_layers == ((0, 1),) + assert permanent.pinned_host_delta_bytes == 0 + class TestApplyAndRollback: def test_component_offload_promotion_marks_resident_without_moving(self): @@ -490,6 +952,49 @@ def test_layerwise_promotion_loads_all_layers_and_rollback_rearms(self): assert manager.register_hooks_calls == 1 assert args.required == set() + def test_partial_layerwise_promotion_restores_exact_group_counts(self): + managers = [ + _FakeLayerwiseManager( + { + "layers.0.w": torch.zeros(16), + "layers.1.w": torch.zeros(16), + "layers.2.w": torch.zeros(16), + }, + resident_layers=1, + ), + _FakeLayerwiseManager( + { + "layers.0.w": torch.zeros(16), + "layers.1.w": torch.zeros(16), + } + ), + ] + module = _FakeLayerwiseDit(managers) + args = _StubResidencyArgs() + candidate = PromotionCandidate( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + promoted_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=10 * GIB_BYTES, + target_layerwise_resident_layers=(2, 2), + target_layerwise_pinned_layers=((), ()), + ) + + applied = apply_promotions( + plan=_plan_for([candidate]), + modules={"transformer": module}, + server_args=args, + ) + + assert [manager.resident_layers for manager in managers] == [2, 2] + assert args.required == set() + rollback_promotions( + applied=applied, + modules={"transformer": module}, + server_args=args, + ) + assert [manager.resident_layers for manager in managers] == [1, 0] + def test_mid_failure_rolls_back_already_applied_promotions(self): manager = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) module = _FakeLayerwiseDit([manager]) @@ -506,7 +1011,111 @@ def test_mid_failure_rolls_back_already_applied_promotions(self): ) assert args.required == set() assert manager.enabled is True - assert manager.register_hooks_calls == 1 + # Validation finishes before any placement is touched. + assert manager.register_hooks_calls == 0 + + def test_hostpin_handoff_releases_all_components_before_acquiring( + self, monkeypatch + ): + events: list[str] = [] + source_manager = _FakeLayerwiseManager( + {"layers.0.w": torch.zeros(16)}, + event_log=events, + event_name="source", + ) + source_manager._pinned_layers = (0,) + target_manager = _FakeLayerwiseManager( + {"layers.0.w": torch.zeros(16)}, + event_log=events, + event_name="target", + ) + modules = { + "source": _FakeLayerwiseDit([source_manager]), + "target": _FakeLayerwiseDit([target_manager]), + } + candidates = [ + PromotionCandidate( + component_name="source", + residency_mode=LAYERWISE_OFFLOAD, + promoted_weight_bytes=0, + h2d_bytes_per_request=-1, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + pinned_host_delta_bytes=-source_manager.offloaded_weight_bytes(), + ), + PromotionCandidate( + component_name="target", + residency_mode=LAYERWISE_OFFLOAD, + promoted_weight_bytes=0, + h2d_bytes_per_request=10, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=target_manager.offloaded_weight_bytes(), + ), + ] + monkeypatch.setattr( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "auto_residency.release_unused_pinned_memory", + lambda: events.append("flush"), + ) + + applied = apply_promotions( + plan=_plan_for(candidates), + modules=modules, + server_args=_StubResidencyArgs(), + ) + + assert events == ["source:unpin", "flush", "target:pin"] + assert all(promotion.pinned_host_changed for promotion in applied) + + def test_hostpin_rollback_releases_new_owner_before_restoring_old_owner( + self, monkeypatch + ): + events: list[str] = [] + source_manager = _FakeLayerwiseManager( + {"layers.0.w": torch.zeros(16)}, + event_log=events, + event_name="source", + ) + target_manager = _FakeLayerwiseManager( + {"layers.0.w": torch.zeros(16)}, + event_log=events, + event_name="target", + ) + target_manager._pinned_layers = (0,) + modules = { + "source": _FakeLayerwiseDit([source_manager]), + "target": _FakeLayerwiseDit([target_manager]), + } + applied = [ + AppliedPromotion( + component_name="source", + residency_mode=LAYERWISE_OFFLOAD, + previous_layerwise_resident_layers=(0,), + previous_layerwise_pinned_layers=((0,),), + pinned_host_changed=True, + ), + AppliedPromotion( + component_name="target", + residency_mode=LAYERWISE_OFFLOAD, + previous_layerwise_resident_layers=(0,), + previous_layerwise_pinned_layers=((),), + pinned_host_changed=True, + ), + ] + monkeypatch.setattr( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "auto_residency.release_unused_pinned_memory", + lambda: events.append("flush"), + ) + + rollback_promotions( + applied=applied, + modules=modules, + server_args=_StubResidencyArgs(), + ) + + assert events == ["target:unpin", "flush", "source:pin"] def test_enable_offload_does_not_double_register_hooks(self): manager = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) @@ -613,7 +1222,7 @@ def test_reports_equivalent_server_args_and_kill_switch(self): ] ) message = format_applied_changes(plan=plan) - assert "text_encoder: layerwise-offload -> resident" in message + assert "text_encoder: permanent resident layers=(1,)" in message assert "vae: component-offload -> resident" in message assert "--component-residency text_encoder=resident vae=resident" in message assert "SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY=1" in message @@ -686,7 +1295,7 @@ def test_capped_frames_get_the_gpu_alignment_real_requests_get(self): class TestCalibrationFrames: def _patch_gate(self, monkeypatch, reason: str | None = None) -> None: monkeypatch.setattr( - "sglang.multimodal_gen.runtime.server_warmup.auto_residency_skip_reason", + "sglang.multimodal_gen.runtime.warmup_request_builder.auto_residency_args_skip_reason", lambda _args: reason, ) @@ -772,6 +1381,7 @@ def _base_args(self, **overrides) -> SimpleNamespace: performance_mode="auto", warmup_mode="server", warmup_resolutions=None, + disagg_role="monolithic", backend="sglang", enable_breakable_cuda_graph=False, enable_torch_compile=False, @@ -780,8 +1390,10 @@ def _base_args(self, **overrides) -> SimpleNamespace: ulysses_degree=1, use_fsdp_inference=False, quantization=None, + component_quantizations={}, transformer_weights_path=None, nunchaku_config=None, + direct_gpu_weight_loading=False, pipeline_config=SimpleNamespace(task_type=ModelTaskType.T2V), ) for key, value in overrides.items(): @@ -811,17 +1423,17 @@ def test_manual_performance_mode(self, monkeypatch): "overrides, expected_fragment", [ ({"warmup_mode": "request"}, "server warmup"), + ({"disagg_role": "denoiser"}, "server warmup"), ({"backend": "diffusers"}, "diffusers"), ({"enable_breakable_cuda_graph": True}, "CUDA graph"), # compile warmup strips the memory layout (layerwise DiT + # resident aux components on CPU): its peaks are not serving peaks ({"enable_torch_compile": True}, "stripped memory layout"), ({"batching_max_size": 4}, "batching"), - ({"dp_size": 2}, "dp replicas"), - ({"ulysses_degree": 2}, "Ulysses"), - ({"use_fsdp_inference": True}, "FSDP"), ({"quantization": "fp8"}, "quantized"), + ({"component_quantizations": {"transformer": "fp8"}}, "quantized"), ({"transformer_weights_path": "/x.safetensors"}, "quantized"), + ({"direct_gpu_weight_loading": True}, "direct GPU"), ], ) def test_excluded_paths(self, monkeypatch, overrides, expected_fragment): @@ -836,6 +1448,20 @@ def test_cache_dit_excluded(self, monkeypatch): reason = self._skip_reason(self._base_args()) assert reason is not None and "cache-dit" in reason + @pytest.mark.parametrize( + "overrides", + [ + {"dp_size": 2}, + {"ulysses_degree": 2}, + {"use_fsdp_inference": True}, + ], + ) + def test_parallel_paths_reach_platform_gate(self, monkeypatch, overrides): + monkeypatch.delenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", raising=False) + monkeypatch.delenv("SGLANG_CACHE_DIT_ENABLED", raising=False) + reason = self._skip_reason(self._base_args(**overrides)) + assert reason is None or reason == "requires CUDA" + def test_eligible_path_reaches_platform_gate(self, monkeypatch): monkeypatch.delenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", raising=False) monkeypatch.delenv("SGLANG_CACHE_DIT_ENABLED", raising=False) diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index a6c10fb2cc03..59a63c776c4c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py @@ -28,6 +28,7 @@ LongLive2T2VConfig, ) from sglang.multimodal_gen.configs.sample.longlive2 import LongLive2SamplingParams +from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( SetLoraReq, @@ -53,6 +54,7 @@ from sglang.multimodal_gen.runtime.warmup_request_builder import ( DEFAULT_PLACEHOLDER_PROMPT, SERVER_WARMUP_IMAGE_FALLBACK_RESOLUTION, + _apply_warmup_sampling_overrides, _resolve_warmup_num_frames, build_warmup_reqs, should_include_warmup_image, @@ -108,6 +110,51 @@ def _make_validation_server_args(enable_cfg_parallel: bool) -> MagicMock: class TestWarmupReqCfgParallel(unittest.TestCase): """Warmup request construction and req-based warmup guards.""" + def test_sampling_workload_override_accepts_json(self): + defaults = SamplingParams( + width=1024, + height=1024, + num_frames=81, + num_inference_steps=35, + ) + server_args = SimpleNamespace( + warmup_sampling_params=( + '{"width":832,"height":480,"num_frames":9,' '"num_inference_steps":4}' + ) + ) + + overridden = _apply_warmup_sampling_overrides(server_args, defaults) + + self.assertEqual( + ( + overridden.width, + overridden.height, + overridden.num_frames, + overridden.num_inference_steps, + ), + (832, 480, 9, 4), + ) + self.assertEqual((defaults.width, defaults.height), (1024, 1024)) + + def test_sampling_workload_override_rejects_unknown_field(self): + server_args = SimpleNamespace( + warmup_sampling_params={"not_a_sampling_field": 1} + ) + + with self.assertRaisesRegex(ValueError, "invalid --warmup-sampling-params"): + _apply_warmup_sampling_overrides(server_args, SamplingParams()) + + def test_sampling_workload_override_supports_fixed_model_fields(self): + defaults = MiniMaxH3SamplingParams() + server_args = SimpleNamespace( + warmup_sampling_params={"num_frames": 49, "fps": 12} + ) + + overridden = _apply_warmup_sampling_overrides(server_args, defaults) + + self.assertEqual((overridden.num_frames, overridden.fps), (49, 12)) + self.assertEqual((defaults.num_frames, defaults.fps), (1, 24)) + def test_warmup_req_cfg_parallel_sets_do_cfg(self): server_args = _make_bare_scheduler(enable_cfg_parallel=True).server_args sampling_defaults = SamplingParams() diff --git a/python/sglang/multimodal_gen/test/unit/test_component_residency.py b/python/sglang/multimodal_gen/test/unit/test_component_residency.py index e4e52fd66075..da06bf7d1d95 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_residency.py @@ -15,6 +15,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.text_encoding import ( RealtimeTextEncodingStage, ) +from sglang.multimodal_gen.runtime.platforms import current_platform def test_component_offload_releases_preferred_component_after_request(): @@ -158,6 +159,132 @@ def component_uses(self, server_args, stage_name=None): return self.uses +def test_warmup_records_use_and_transition_peaks(monkeypatch): + device_module = SimpleNamespace( + is_available=lambda: True, + reset_peak_memory_stats=Mock(), + max_memory_reserved=lambda: 11, + memory_reserved=lambda: 3, + ) + monkeypatch.setattr(torch, "get_device_module", lambda: device_module) + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + + use = ComponentUse("denoise", "transformer") + stage = _Stage(use) + module = torch.nn.Linear(2, 2) + pipeline = SimpleNamespace( + modules={"transformer": module}, + _stage_name_mapping={"denoise": stage}, + component_residency_strategies={}, + ) + server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + manager = ComponentResidencyManager(pipeline, server_args) + manager.strategy_for = Mock(return_value=Mock()) + manager.refresh_pipeline(pipeline) + manager.begin_request([stage], SimpleNamespace(is_warmup=True), server_args) + + manager.before_stage(stage, 0, SimpleNamespace(is_warmup=True), server_args) + manager.begin_stage() + manager.end_stage() + manager.finish_request() + + peaks = manager.take_warmup_phase_peaks() + assert peaks["request:before-stage"] == ((), 11) + assert peaks["0:denoise:setup"] == ((), 11) + assert peaks["0:denoise:transition:idle->transformer"] == ( + ("transformer",), + 11, + ) + assert peaks["0:denoise:use:transformer"] == (("transformer",), 11) + assert peaks["0:denoise:transition:transformer->idle"] == ( + ("transformer",), + 11, + ) + assert peaks["0:denoise:between"] == ((), 11) + assert peaks["request:cleanup:transformer"] == (("transformer",), 11) + assert peaks["idle"] == ((), 3) + + +def test_warmup_records_same_component_dtype_prepare_as_transition(monkeypatch): + device_module = SimpleNamespace( + is_available=lambda: True, + reset_peak_memory_stats=Mock(), + max_memory_reserved=lambda: 11, + memory_reserved=lambda: 3, + ) + monkeypatch.setattr(torch, "get_device_module", lambda: device_module) + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + + first = ComponentUse("stage", "transformer", target_dtype=torch.float16) + second = ComponentUse("stage", "transformer", target_dtype=torch.bfloat16) + stage = _Stage(first, second) + module = torch.nn.Linear(2, 2) + pipeline = SimpleNamespace( + modules={"transformer": module}, + _stage_name_mapping={"stage": stage}, + component_residency_strategies={}, + ) + server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + manager = ComponentResidencyManager(pipeline, server_args) + strategy = Mock() + manager.strategy_for = Mock(return_value=strategy) + manager.refresh_pipeline(pipeline) + manager.begin_request([stage], SimpleNamespace(is_warmup=True), server_args) + manager.before_stage(stage, 0, SimpleNamespace(is_warmup=True), server_args) + + manager.begin_use(first, module=module) + manager.begin_use(second, module=module) + + manager._record_warmup_phase_peak() + assert manager._warmup_phase_peaks[ + "0:stage:transition:transformer->transformer" + ] == (("transformer",), 11) + assert strategy.prepare_for_use.call_count == 2 + + +def test_warmup_splits_sequential_component_transition(monkeypatch): + device_module = SimpleNamespace( + is_available=lambda: True, + reset_peak_memory_stats=Mock(), + max_memory_reserved=lambda: 11, + memory_reserved=lambda: 3, + ) + monkeypatch.setattr(torch, "get_device_module", lambda: device_module) + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + + first = ComponentUse("stage", "text_encoder") + second = ComponentUse("stage", "transformer") + stage = _Stage(first, second) + modules = { + "text_encoder": torch.nn.Linear(2, 2), + "transformer": torch.nn.Linear(2, 2), + } + pipeline = SimpleNamespace( + modules=modules, + _stage_name_mapping={"stage": stage}, + component_residency_strategies={}, + ) + server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + manager = ComponentResidencyManager(pipeline, server_args) + manager.strategy_for = Mock(return_value=Mock()) + manager.refresh_pipeline(pipeline) + manager.begin_request([stage], SimpleNamespace(is_warmup=True), server_args) + manager.before_stage(stage, 0, SimpleNamespace(is_warmup=True), server_args) + + manager.begin_use(first) + manager.begin_use(second) + manager._record_warmup_phase_peak() + + assert manager._warmup_phase_peaks["0:stage:transition:text_encoder->idle"] == ( + ("text_encoder",), + 11, + ) + assert manager._warmup_phase_peaks["0:stage:transition:idle->transformer"] == ( + ("transformer",), + 11, + ) + + def _manager_for_stage(stage, modules): pipeline = SimpleNamespace( modules=modules, diff --git a/python/sglang/multimodal_gen/test/unit/test_dp_routing.py b/python/sglang/multimodal_gen/test/unit/test_dp_routing.py index 32bcd43991db..bc6572590e8f 100644 --- a/python/sglang/multimodal_gen/test/unit/test_dp_routing.py +++ b/python/sglang/multimodal_gen/test/unit/test_dp_routing.py @@ -7,6 +7,7 @@ import itertools from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( + AutoResidencyReq, SetLoraReq, ShutdownReq, ) @@ -14,6 +15,8 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.scheduler_client import ( _CONTROL_REQ_TYPES, + _is_warmup_batch, + _merge_auto_residency_results, _merge_fanout_results, _select_replica, ) @@ -45,12 +48,20 @@ def test_session_requests_stick_to_one_replica(): def test_control_reqs_are_recognized(): + assert isinstance(AutoResidencyReq(), _CONTROL_REQ_TYPES) assert isinstance(SetLoraReq(lora_nickname="x", lora_path="y"), _CONTROL_REQ_TYPES) assert isinstance(ShutdownReq(), _CONTROL_REQ_TYPES) assert not isinstance([_req()], _CONTROL_REQ_TYPES) assert not isinstance(_req(), _CONTROL_REQ_TYPES) +def test_warmup_reqs_are_fanned_out_to_every_replica(): + assert _is_warmup_batch(_req(is_warmup=True)) + assert _is_warmup_batch([_req(is_warmup=True), _req(is_warmup=True)]) + assert not _is_warmup_batch(_req()) + assert not _is_warmup_batch([]) + + def test_fanout_merge_surfaces_the_failing_replica(): ok = OutputBatch(output=None) bad = OutputBatch(error="replica 1 exploded") @@ -58,6 +69,32 @@ def test_fanout_merge_surfaces_the_failing_replica(): assert _merge_fanout_results([ok, OutputBatch(output=None)]) is ok +def test_auto_residency_fanout_rewarms_when_any_replica_promotes(): + skipped = OutputBatch(output={"status": "skipped"}) + promoted = OutputBatch(output={"status": "promoted"}) + assert _merge_auto_residency_results([skipped, promoted]) is promoted + + +def test_auto_residency_fanout_tolerates_a_locally_rolled_back_replica(): + restored = OutputBatch( + output={"status": "rolled_back"}, error="original placement restored" + ) + promoted = OutputBatch(output={"status": "promoted"}) + assert _merge_auto_residency_results([restored, promoted]) is promoted + + +def test_auto_residency_fanout_never_hides_failed_rollback(): + failed = OutputBatch(output={"status": "rollback_failed"}, error="restart required") + promoted = OutputBatch(output={"status": "promoted"}) + assert _merge_auto_residency_results([promoted, failed]) is failed + + +def test_auto_residency_fanout_reports_successful_rollback(): + skipped = OutputBatch(output={"status": "skipped"}) + restored = OutputBatch(output={"status": "rolled_back"}) + assert _merge_auto_residency_results([skipped, restored]) is restored + + def test_scheduler_endpoints_one_per_replica(): args = ServerArgs.__new__(ServerArgs) args.host = "localhost" diff --git a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py b/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py index 5caeb91e3962..1395628b8a6b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py @@ -1,5 +1,6 @@ """Pinned host memory is planned against the cgroup cap, not the whole machine.""" +import pytest import torch import torch.nn as nn @@ -195,6 +196,18 @@ def test_a_suffix_of_the_reported_path_is_found_under_the_mount( class TestHostPinBudget: + def test_local_workers_split_one_node_allowance(self, monkeypatch): + monkeypatch.setattr( + host_memory_budget, + "host_memory_available_bytes", + lambda: 40 * GIB_BYTES, + ) + + budget = HostPinBudget.for_local_worker(4) + + # The node keeps one 2 GiB reserve, then workers split the remainder. + assert budget.spendable_bytes == int(9.5 * GIB_BYTES) + def test_a_component_that_fits_is_granted(self): budget = HostPinBudget(available_bytes=40 * GIB_BYTES) assert budget.request(component_name="dit", weight_bytes=20 * GIB_BYTES) @@ -217,6 +230,17 @@ def test_a_later_component_is_denied_once_the_budget_is_spent(self): component_name="text_encoder", weight_bytes=20 * GIB_BYTES ) + def test_released_allowance_can_be_reassigned(self): + budget = HostPinBudget(available_bytes=40 * GIB_BYTES) + assert budget.request(component_name="encoder", weight_bytes=30 * GIB_BYTES) + budget.release(20 * GIB_BYTES) + assert budget.request(component_name="dit", weight_bytes=20 * GIB_BYTES) + + def test_cannot_release_uncommitted_allowance(self): + budget = HostPinBudget(available_bytes=40 * GIB_BYTES) + with pytest.raises(ValueError, match="more pinned-host memory"): + budget.release(GIB_BYTES) + def test_the_cap_binds_even_for_the_first_component(self): # priority comes from asking first, not from being allowed to overrun budget = HostPinBudget(available_bytes=8 * GIB_BYTES) 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 1fbc65cbcf35..68a79059d337 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -206,6 +206,8 @@ def _server_args(**kwargs): component_residency=None, disagg_role=RoleType.MONOLITHIC, performance_mode="auto", + num_gpus=1, + nnodes=1, _required_resident_components=set(), _component_layerwise_capabilities={}, _explicit_arg_names=set(), @@ -899,6 +901,29 @@ def test_resident_layers_off_by_default_streams_everything(monkeypatch): assert 2 not in manager._gpu_layers +def test_resident_layer_size_and_runtime_reconfiguration(monkeypatch): + _patch_fake_device(monkeypatch) + manager = _resident_manager( + _MultiBlockModel(4), num_layers=4, prefetch_size=1, resident_layers=2 + ) + full_bytes = sum(manager.layer_weight_bytes().values()) + + assert manager.resident_weight_bytes(0) == 0 + assert manager.resident_weight_bytes(4) == full_bytes + assert manager.resident_weight_bytes(2) * 2 == full_bytes + assert manager.peak_managed_device_weight_bytes(0) == full_bytes // 4 + assert manager.peak_managed_device_weight_bytes(2) == 3 * full_bytes // 4 + assert manager.peak_managed_device_weight_bytes(4) == full_bytes + + _arm_residency(manager) + previous = manager.set_resident_layers(1) + assert previous == 2 + assert manager.resident_layers == 1 + assert manager._resident_set == {0} + assert manager._residency_active is False + assert not manager._gpu_layers + + def test_prepare_for_next_req_repins_residents(monkeypatch): _patch_fake_device(monkeypatch) manager = _resident_manager( @@ -1511,6 +1536,38 @@ def test_weights_are_copied_when_they_fit(tmp_path, monkeypatch): assert manager._consolidated_cpu_weights[0] +def test_anonymous_layer_store_can_reassign_its_pin_budget(tmp_path, monkeypatch): + if not pathlib.Path("/proc/self/maps").exists(): + pytest.skip("needs /proc to tell a mapping from anonymous memory") + manager = _mapped_manager( + tmp_path, + monkeypatch, + available_gib=64, + pin_budget_bytes=4 * 1024**3, + num_blocks=2, + ) + original = manager.pinned_layer_indices() + assert original == (0, 1) + committed = manager._pin_budget.committed_bytes + pinned_store_bytes = sum( + manager.layer_host_store_bytes()[layer_idx] for layer_idx in original + ) + + previous = manager.set_pinned_layers(()) + assert previous == original + assert manager.pinned_layer_indices() == () + assert manager._pin_budget.committed_bytes == committed - pinned_store_bytes + assert all( + not buffer.is_pinned() + for stores in manager._consolidated_cpu_weights.values() + for buffer in stores.values() + ) + + manager.set_pinned_layers(previous) + assert manager.pinned_layer_indices() == original + assert manager._pin_budget.committed_bytes == committed + + def test_a_mapped_weight_is_not_written_back(tmp_path, monkeypatch): if not pathlib.Path("/proc/self/maps").exists(): pytest.skip("needs /proc to tell a mapping from anonymous memory") diff --git a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py index d141ae8f095e..67b183a1cd2b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -1,4 +1,5 @@ import json +from types import SimpleNamespace from unittest.mock import Mock, patch import pytest @@ -6,6 +7,10 @@ import sglang.multimodal_gen.runtime.managers.gpu_worker as gpu_worker_module from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker +from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( + DefaultWorkload, + WarmupMemoryRecord, +) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.perf_logger import ( @@ -89,6 +94,69 @@ def test_worker_records_replica_load_and_runtime_peaks(): assert metrics.memory_snapshots["runtime_peak"].peak_reserved_mb == 3584.0 +def test_server_warmup_preserves_peak_after_managed_stage_timeline(): + worker = GPUWorker.__new__(GPUWorker) + worker._auto_residency_warmup_records = [] + residency_manager = Mock() + residency_manager.take_warmup_phase_peaks.return_value = { + "0:denoise:use:transformer": (("transformer",), 10) + } + device_module = Mock() + device_module.max_memory_reserved.return_value = 12 + req = SimpleNamespace(width=64, height=64, num_frames=1) + + with ( + patch.object(torch, "get_device_module", return_value=device_module), + patch.object( + gpu_worker_module, + "peek_global_component_residency_manager", + return_value=residency_manager, + ), + ): + worker._record_server_warmup_memory( + req=req, + baseline_allocated_bytes=3, + succeeded=True, + ) + + record = worker._auto_residency_warmup_records[0] + assert record.peak_reserved_bytes == 12 + assert record.phase_peak_reserved_bytes["request:untracked"] == 12 + assert record.phase_active_components["request:untracked"] == () + + +def test_loaded_quantized_checkpoint_disables_auto_residency(): + worker = GPUWorker.__new__(GPUWorker) + worker.rank = 0 + quantized_module = torch.nn.Module() + quantized_module.register_parameter( + "weight", + torch.nn.Parameter(torch.ones(1, dtype=torch.int8), requires_grad=False), + ) + worker.pipeline = SimpleNamespace(modules={"transformer": quantized_module}) + + report = worker._build_auto_residency_report( + workload=DefaultWorkload( + width=64, + height=64, + num_frames=1, + num_inference_steps=1, + ), + records=[ + WarmupMemoryRecord( + width=64, + height=64, + num_frames=1, + baseline_allocated_bytes=1, + peak_reserved_bytes=2, + succeeded=True, + ) + ], + ) + + assert report.skip_reason == "loaded quantized components: transformer" + + def test_baseline_config_loads_per_scenario_peak_vram(tmp_path): path = tmp_path / "baseline.json" path.write_text( diff --git a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py new file mode 100644 index 000000000000..bea10327cb0f --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from sglang.multimodal_gen.runtime.managers.memory_managers.placement_budget import ( + PlacementOption, + optimize_placement, +) + + +def _option( + group: str, + option: str, + *, + device: int = 0, + pinned: int = 0, + savings: int = 0, +) -> PlacementOption: + return PlacementOption( + group_key=group, + option_key=option, + resource_delta_bytes={ + "gpu:rank0:runtime": device, + "hostpin:node0": pinned, + }, + estimated_latency_savings=savings, + ) + + +def test_joint_plan_beats_independent_greedy_allocations(): + plan = optimize_placement( + [ + _option("a", "a-resident", device=10, savings=10), + _option("a", "a-pinned", pinned=10, savings=9), + _option("b", "b-resident", device=10, savings=100), + _option("b", "b-pinned", pinned=10, savings=1), + ], + resource_budget_bytes={ + "gpu:rank0:runtime": 10, + "hostpin:node0": 10, + }, + ) + + assert [option.option_key for option in plan.selections] == [ + "a-pinned", + "b-resident", + ] + assert plan.estimated_latency_savings == 109 + + +def test_selects_at_most_one_option_per_group(): + plan = optimize_placement( + [ + _option("dit", "half", device=5, savings=8), + _option("dit", "full", device=10, savings=20), + ], + resource_budget_bytes={ + "gpu:rank0:runtime": 20, + "hostpin:node0": 0, + }, + ) + + assert [option.option_key for option in plan.selections] == ["full"] + + +def test_exact_optimizer_can_choose_multiple_smaller_components(): + plan = optimize_placement( + [ + _option("large", "large", device=10, savings=100), + _option("small-a", "small-a", device=6, savings=70), + _option("small-b", "small-b", device=4, savings=40), + ], + resource_budget_bytes={ + "gpu:rank0:runtime": 10, + "hostpin:node0": 0, + }, + ) + + assert [option.option_key for option in plan.selections] == [ + "small-a", + "small-b", + ] + assert plan.estimated_latency_savings == 110 + + +def test_ties_prefer_lower_resource_cost_then_stable_key(): + plan = optimize_placement( + [ + _option("a", "higher-cost", device=2, savings=10), + _option("a", "z", device=1, savings=10), + _option("a", "a", device=1, savings=10), + ], + resource_budget_bytes={ + "gpu:rank0:runtime": 2, + "hostpin:node0": 0, + }, + ) + + assert [option.option_key for option in plan.selections] == ["a"] + assert plan.resource_delta_bytes == { + "gpu:rank0:runtime": 1, + "hostpin:node0": 0, + } + + +def test_one_static_placement_satisfies_all_observed_phase_constraints(): + plan = optimize_placement( + [ + PlacementOption( + group_key="a", + option_key="a-resident", + resource_delta_bytes={ + "gpu:rank0:encode": 8, + "gpu:rank0:denoise": 2, + "hostpin:node0": 0, + }, + estimated_latency_savings=9, + ), + PlacementOption( + group_key="a", + option_key="a-pinned", + resource_delta_bytes={ + "gpu:rank0:encode": 1, + "gpu:rank0:denoise": 1, + "hostpin:node0": 8, + }, + estimated_latency_savings=5, + ), + PlacementOption( + group_key="b", + option_key="b-resident", + resource_delta_bytes={ + "gpu:rank0:encode": 2, + "gpu:rank0:denoise": 8, + "hostpin:node0": 0, + }, + estimated_latency_savings=10, + ), + PlacementOption( + group_key="b", + option_key="b-pinned", + resource_delta_bytes={ + "gpu:rank0:encode": 1, + "gpu:rank0:denoise": 1, + "hostpin:node0": 8, + }, + estimated_latency_savings=4, + ), + ], + resource_budget_bytes={ + "gpu:rank0:encode": 10, + "gpu:rank0:denoise": 10, + "hostpin:node0": 8, + }, + ) + + assert [option.option_key for option in plan.selections] == [ + "a-resident", + "b-resident", + ] + assert plan.resource_delta_bytes == { + "gpu:rank0:encode": 10, + "gpu:rank0:denoise": 10, + "hostpin:node0": 0, + } + + +def test_negative_headroom_forces_the_plan_to_release_a_resource(): + plan = optimize_placement( + [ + PlacementOption( + group_key="z-cold-encoder", + option_key="cold-encoder-pageable", + resource_delta_bytes={"hostpin:node0": -8}, + estimated_latency_savings=-1, + ) + ], + resource_budget_bytes={"hostpin:node0": -4}, + ) + + assert [option.option_key for option in plan.selections] == [ + "cold-encoder-pageable" + ] + + +def test_negative_headroom_without_a_release_is_infeasible(): + with pytest.raises(ValueError, match="no placement satisfies"): + optimize_placement( + [], + resource_budget_bytes={"hostpin:node0": -1}, + ) + + +def test_unknown_resource_is_rejected(): + with pytest.raises(ValueError, match="unknown resources"): + optimize_placement( + [ + PlacementOption( + group_key="dit", + option_key="resident", + resource_delta_bytes={"gpu:rank0:denoise": 1}, + estimated_latency_savings=1, + ) + ], + resource_budget_bytes={"gpu:rank0:load": 1}, + ) + + +def test_signed_deltas_trade_pinned_host_for_device_memory(): + plan = optimize_placement( + [ + PlacementOption( + group_key="dit", + option_key="dit-resident", + resource_delta_bytes={ + "gpu:rank0:runtime": 8, + "hostpin:node0": -8, + }, + estimated_latency_savings=20, + ), + PlacementOption( + group_key="encoder", + option_key="encoder-pinned", + resource_delta_bytes={ + "gpu:rank0:runtime": 0, + "hostpin:node0": 8, + }, + estimated_latency_savings=10, + ), + ], + resource_budget_bytes={ + "gpu:rank0:runtime": 8, + "hostpin:node0": 0, + }, + ) + + assert [option.option_key for option in plan.selections] == [ + "dit-resident", + "encoder-pinned", + ] + assert plan.resource_delta_bytes["hostpin:node0"] == 0 + + +def test_joint_plan_can_accept_a_local_regression_to_reassign_hostpin(): + plan = optimize_placement( + [ + PlacementOption( + group_key="cold-encoder", + option_key="cold-encoder-pageable", + resource_delta_bytes={"hostpin:node0": -10}, + estimated_latency_savings=-1, + ), + PlacementOption( + group_key="a-hot-dit", + option_key="hot-dit-pinned", + resource_delta_bytes={"hostpin:node0": 10}, + estimated_latency_savings=100, + ), + ], + resource_budget_bytes={"hostpin:node0": 0}, + ) + + assert [option.option_key for option in plan.selections] == [ + "hot-dit-pinned", + "cold-encoder-pageable", + ] + assert plan.estimated_latency_savings == 99 + + +def test_nonbinding_resource_is_removed_without_losing_full_accounting(): + plan = optimize_placement( + [ + _option("dit", "dit-low", device=5, pinned=20, savings=10), + _option("dit", "dit-high", device=5, pinned=80, savings=20), + ], + resource_budget_bytes={ + "gpu:rank0:runtime": 5, + "hostpin:node0": 100, + }, + ) + + assert [option.option_key for option in plan.selections] == ["dit-high"] + assert plan.resource_delta_bytes == { + "gpu:rank0:runtime": 5, + "hostpin:node0": 80, + } diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 6f954eb308bc..081212933e17 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -1423,6 +1423,35 @@ def test_explicit_dit_layerwise_component_wins_over_auto_residency(self): self.assertEqual(args.layerwise_offload_components, ["dit"]) self.assertEqual(args.residency_mode("transformer"), LAYERWISE_OFFLOAD) + def test_server_warmup_defers_high_memory_residency_to_calibration(self): + args = self._from_dict_with_pipeline_config( + Cosmos3Config(), + memory_gb=140, + kwargs={"performance_mode": "auto", "warmup_mode": "server"}, + ) + + self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD) + self.assertEqual(args.residency_mode("vae"), LAYERWISE_OFFLOAD) + + def test_no_warmup_keeps_model_residency_fallback(self): + args = self._from_dict_with_pipeline_config( + Cosmos3Config(), + memory_gb=140, + kwargs={"performance_mode": "auto", "warmup_mode": "off"}, + ) + + self.assertEqual(args.residency_mode("transformer"), RESIDENT) + self.assertEqual(args.residency_mode("vae"), RESIDENT) + + def test_realtime_server_keeps_model_residency_fallback(self): + args = self._from_dict_with_pipeline_config( + SanaWMRealtimeConfig(), + memory_gb=140, + kwargs={"performance_mode": "auto", "warmup_mode": "server"}, + ) + + self.assertEqual(args.residency_mode("transformer"), RESIDENT) + def test_pipeline_configs_declare_auto_tune_hints(self): qwen_deployment = QwenImagePipelineConfig().get_model_deployment_config() cosmos3_deployment = Cosmos3Config( From d4316d3ec77f992b8efbb04a66930a3dcaa47245 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Aug 2026 14:18:09 +0800 Subject: [PATCH 17/90] [diffusion] fix auto residency unit fixtures --- .../test/unit/test_auto_residency.py | 32 +++++++++++++------ .../test/unit/test_cfg_parallel_warmup.py | 4 +++ .../test/unit/test_component_residency.py | 4 ++- .../test/unit/test_layerwise_offload.py | 3 ++ 4 files changed, 33 insertions(+), 10 deletions(-) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index b43ccfb57811..a2d4cccb087b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -511,11 +511,16 @@ def test_each_rank_keeps_a_reserve_scaled_to_its_own_gpu(self): assert plan.reserve_bytes == 14 * GIB_BYTES def test_partial_layer_residency_is_selected_when_full_dit_does_not_fit(self): - full = _candidate( - "transformer", - mode=LAYERWISE_OFFLOAD, - weight_gib=30, - h2d_gib=300, + full = PromotionCandidate( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + promoted_weight_bytes=30 * GIB_BYTES, + h2d_bytes_per_request=300 * GIB_BYTES, + target_layerwise_resident_layers=(30,), + target_layerwise_pinned_layers=((),), + permanent_residency=True, + active_device_delta_bytes=30 * GIB_BYTES, + inactive_device_delta_bytes=30 * GIB_BYTES, ) partial = PromotionCandidate( component_name="transformer", @@ -570,9 +575,11 @@ def test_node_hostpin_is_optimized_with_the_same_selection_vector(self): ] ) + # Plans are reported by benefit. apply_promotions independently sorts + # pin changes so the release still happens before the acquisition. assert [candidate.component_name for candidate in plan.promotions] == [ - "cold_encoder", "hot_dit", + "cold_encoder", ] assert plan.resource_budget_bytes["hostpin:node0"] == 0 @@ -1147,18 +1154,23 @@ def release_required_component_residency( raise AssertionError() args = _BrokenArgs() + failing_manager = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) + failing_manager.fail_load = True candidates = [ _candidate("text_encoder", weight_gib=1), - _candidate("missing_component", weight_gib=1), + _candidate("transformer", mode=LAYERWISE_OFFLOAD, weight_gib=1), ] with pytest.raises(AutoResidencyRollbackError) as exc_info: apply_promotions( plan=_plan_for(candidates), - modules={"text_encoder": nn.Linear(2, 2)}, + modules={ + "text_encoder": nn.Linear(2, 2), + "transformer": _FakeLayerwiseDit([failing_manager]), + }, server_args=args, ) assert "AssertionError" in str(exc_info.value) - assert "missing_component" in str(exc_info.value) + assert "CUDA out of memory" in str(exc_info.value) def test_rollback_keeps_going_past_a_broken_component(self): released: list[str] = [] @@ -1232,6 +1244,7 @@ class TestWarmupFrameAdjustment: def _server_args(self, *, bcg: bool = False, num_gpus: int = 1) -> SimpleNamespace: return SimpleNamespace( pipeline_config=LongLive2T2VConfig(), + pipeline_class_name=None, enable_breakable_cuda_graph=bcg, num_gpus=num_gpus, ) @@ -1277,6 +1290,7 @@ def test_capped_frames_get_the_gpu_alignment_real_requests_get(self): arch_config=SimpleNamespace(temporal_compression_ratio=4), ), ), + pipeline_class_name=None, enable_breakable_cuda_graph=False, num_gpus=2, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index 59a63c776c4c..285d81424889 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py @@ -617,6 +617,10 @@ def test_ltx2_two_stage_warmup_uses_pipeline_alignment(self): server_args.pipeline_config.task_type = ModelTaskType.T2V server_args.pipeline_config.vae_scale_factor = 32 + server_args.pipeline_config.vae_config = SimpleNamespace( + use_temporal_scaling_frames=True, + arch_config=SimpleNamespace(temporal_compression_ratio=8), + ) server_args.pipeline_config.adjust_num_frames.return_value = 25 server_args.num_gpus = 2 diff --git a/python/sglang/multimodal_gen/test/unit/test_component_residency.py b/python/sglang/multimodal_gen/test/unit/test_component_residency.py index da06bf7d1d95..25cc782e9556 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_residency.py @@ -201,7 +201,9 @@ def test_warmup_records_use_and_transition_peaks(monkeypatch): 11, ) assert peaks["0:denoise:between"] == ((), 11) - assert peaks["request:cleanup:transformer"] == (("transformer",), 11) + # A non-preferred component is being released during cleanup, so it is no + # longer part of the placement that follows this transition. + assert peaks["request:cleanup:transformer"] == ((), 11) assert peaks["idle"] == ((), 3) 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 68a79059d337..87424ba5ebd9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -229,6 +229,9 @@ def _server_args(**kwargs): # the pin budget ranks candidates by bytes x steps, and reads the step # count off the pipeline's sampling defaults pipeline_class_name=None, + model_path=None, + backend="sglang", + model_id=None, ) defaults.update(kwargs) return _TestServerArgs(**defaults) From d83fc8771a463e9be0206dcdd2faa29c459aa4b8 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Aug 2026 15:20:52 +0800 Subject: [PATCH 18/90] [diffusion] calibrate residency at serving shape --- .../memory_managers/auto_residency.py | 30 ++--- .../multimodal_gen/runtime/server_warmup.py | 11 +- .../runtime/warmup_request_builder.py | 113 ++++++++++-------- .../test/server/test_server_utils.py | 2 +- .../test/unit/test_auto_residency.py | 77 ++++++------ .../test/unit/test_cfg_parallel_warmup.py | 53 ++++++++ .../test/unit/test_health_warmup_gate.py | 14 ++- 7 files changed, 186 insertions(+), 114 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 7b3c21b5a8df..129248485680 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -2,16 +2,15 @@ """Warmup-calibrated automatic component residency promotion. Under ``--performance-mode auto`` with server warmup, each rank measures the -peak GPU memory of the synthetic warmup requests, extrapolates it to the -model's default workload, and promotes implicitly offloaded components +peak GPU memory of bounded synthetic warmup requests and a low-step probe at +the complete default serving shape, then promotes implicitly offloaded components (component offload -> resident, layerwise offload -> fully loaded) when the estimate plus the promoted weights still fits under a safety reserve. -The estimate splits the measured peak into a persistent part (weights and -buffers alive before the forward) and an activation part (everything above -it); only the activation part scales with the workload ratio. Scaling the -whole peak would multiply resident weights by the video frame/area cap ratio -(~16x for Wan-class defaults) and promotion would never trigger. +When no full-shape measurement is available, the fallback estimate splits the +measured peak into persistent weights and workload-scaled activations. Scaling +the whole peak would multiply resident weights by the video frame/area cap +ratio (~16x for Wan-class defaults) and promotion would never trigger. Promotion targets the model default workload only (default resolution, default frames, batch=1). Larger shapes, batches, or multi-image inputs need @@ -214,25 +213,12 @@ class AppliedPromotion(msgspec.Struct, frozen=True): def resolve_default_workload(server_args: ServerArgs) -> DefaultWorkload: """Resolve the default request shape promotion is optimized for.""" from sglang.multimodal_gen.runtime.warmup_request_builder import ( - _apply_warmup_frame_contract, get_model_sampling_defaults, + resolve_default_workload_shape, ) defaults = get_model_sampling_defaults(server_args) - # __post_init__ already applied the model's _default_width/_default_height - width = defaults.width - height = defaults.height - if (width is None or height is None) and defaults.supported_resolutions: - # worst-case the target with the largest supported shape - width, height = max( - defaults.supported_resolutions, key=lambda size: size[0] * size[1] - ) - num_frames = defaults.num_frames or 1 - if num_frames > 1: - # same contract the warmup requests and real requests get - num_frames = _apply_warmup_frame_contract( - server_args, defaults, num_frames=num_frames - ) + width, height, num_frames = resolve_default_workload_shape(server_args, defaults) return DefaultWorkload( width=width, height=height, diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index 2dfc67c07711..f26206045ad5 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -237,9 +237,14 @@ def _degrade_after_oom(server_args: ServerArgs, req: Req) -> Req | None: def format_warmup_req(req_or_group: Any) -> str: req = get_first_generation_req(req_or_group) - prefix = ( - "server warmup req" if is_server_based_warmup(req_or_group) else "warmup req" - ) + if req is not None and req.extra.get("auto_residency_full_shape_probe"): + prefix = "auto residency probe" + else: + prefix = ( + "server warmup req" + if is_server_based_warmup(req_or_group) + else "warmup req" + ) if req is None: return prefix diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py index ee5558eb9582..b2b5d37bf2d8 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -6,6 +6,11 @@ sampling defaults, then keeps startup bounded by choosing common low-cost resolution/frame buckets and trimming the denoising step count. +When warmup-calibrated auto residency is active, a second probe restores the +full default serving shape but keeps the trimmed step count. Memory depends on +the activation shape rather than the number of repeated denoising steps, so +this directly measures placement headroom without running a full generation. + Image models may run a tiny second step because first/last step paths often initialize different kernels or scheduler state. Video models cap frames and steps to keep startup bounded. Explicit warmup resolutions share this builder; @@ -49,10 +54,6 @@ SERVER_WARMUP_VIDEO_MAX_AREA = 832 * 480 SERVER_WARMUP_MAX_VIDEO_FRAMES = 17 SERVER_WARMUP_LTX2_TWO_STAGE_MAX_VIDEO_FRAMES = 25 -# Second, smaller frame count used to calibrate auto residency: two workload -# sizes let the estimator fit peak = constant + slope * units instead of -# scaling one measured peak (whose constant part dominates under offload). -SERVER_WARMUP_CALIBRATION_VIDEO_FRAMES = 9 SERVER_WARMUP_IMAGE_STEPS = 2 SERVER_WARMUP_VIDEO_STEPS = 2 @@ -384,39 +385,53 @@ def _apply_warmup_frame_contract( return num_frames -def _resolve_calibration_num_frames( +def resolve_default_workload_shape( + server_args: ServerArgs, + sampling_defaults: SamplingParams, +) -> tuple[int | None, int | None, int]: + """Resolve the serving shape used by warmup-based memory planning.""" + width = sampling_defaults.width + height = sampling_defaults.height + if (width is None or height is None) and sampling_defaults.supported_resolutions: + width, height = max( + sampling_defaults.supported_resolutions, + key=lambda size: size[0] * size[1], + ) + num_frames = sampling_defaults.num_frames or 1 + if num_frames > 1: + num_frames = _apply_warmup_frame_contract( + server_args, sampling_defaults, num_frames=num_frames + ) + return width, height, num_frames + + +def _resolve_auto_residency_probe_shape( server_args: ServerArgs, sampling_defaults: SamplingParams, *, - warmup_num_frames: int | None, + warmup_shape: tuple[int, int, int | None], server_based_warmup: bool, -) -> int | None: - """A second, smaller warmup frame count for auto-residency calibration. - - Only produced when the main video warmup was frame-capped and the - auto-residency promotion is actually eligible to consume the measurement: - the estimator then has two workload - sizes to fit the constant/linear split of the peak. Returns None when a - second measurement adds no information (same adjusted frames) or would - never be read. +) -> tuple[int, int, int] | None: + """Return a full serving-shape memory probe when bounded warmup is smaller. + + The probe still runs only the bounded warmup step count. Denoising steps + repeat the same activation shape, so one full-shape forward measures the + placement constraints directly without paying for a full generation or + extrapolating a small-shape allocator peak. """ - if not server_based_warmup or not _is_video_warmup_task(server_args): - return None - if warmup_num_frames is None: + if not server_based_warmup: return None if auto_residency_args_skip_reason(server_args) is not None: return None - default_num_frames = sampling_defaults.num_frames or 1 - if warmup_num_frames >= default_num_frames: - return None - calibration_num_frames = _apply_warmup_frame_contract( - server_args, - sampling_defaults, - num_frames=SERVER_WARMUP_CALIBRATION_VIDEO_FRAMES, + width, height, num_frames = resolve_default_workload_shape( + server_args, sampling_defaults ) - if calibration_num_frames == warmup_num_frames or calibration_num_frames <= 0: + if width is None or height is None: + return None + target = (width, height, num_frames) + if target == warmup_shape: return None - return calibration_num_frames + return target def _effective_cfg_scale(sampling_defaults: SamplingParams) -> float | None: @@ -529,11 +544,11 @@ def build_warmup_reqs( sampling_defaults, server_based_warmup=server_based_warmup, ) - calibration_num_frames = ( - _resolve_calibration_num_frames( + auto_residency_probe_shape = ( + _resolve_auto_residency_probe_shape( server_args, sampling_defaults, - warmup_num_frames=warmup_num_frames, + warmup_shape=(width, height, warmup_num_frames), server_based_warmup=server_based_warmup, ) if warmup_resolutions is None @@ -601,17 +616,11 @@ def build_warmup_reqs( req.extra["server_based_warmup"] = True warmup_reqs.append(req) - if calibration_num_frames is not None and (width, height) == resolutions[0]: - # The calibration measurement runs AFTER the main warmup request: - # one-time allocations (compile scratch, cuDNN workspaces, - # allocator pool growth) then land in the larger measurement, - # which can only steepen the fitted slope -- the safe direction. - # Measuring the small shape first flattened the slope and - # under-estimated the default workload peak. + if auto_residency_probe_shape is not None and (width, height) == resolutions[0]: warmup_reqs.append( - _build_calibration_warmup_req( + _build_auto_residency_probe_req( req_kwargs, - calibration_num_frames=calibration_num_frames, + probe_shape=auto_residency_probe_shape, warmup_steps=warmup_steps, return_warmup_result=return_warmup_result, server_based_warmup=server_based_warmup, @@ -622,25 +631,27 @@ def build_warmup_reqs( return warmup_reqs -def _build_calibration_warmup_req( +def _build_auto_residency_probe_req( req_kwargs: dict, *, - calibration_num_frames: int, + probe_shape: tuple[int, int, int], warmup_steps: int, return_warmup_result: bool, server_based_warmup: bool, server_args: ServerArgs, ) -> Req: - calibration_kwargs = req_kwargs.copy() - calibration_kwargs["sampling_params"] = copy(req_kwargs["sampling_params"]) - calibration_kwargs["num_frames"] = calibration_num_frames - calibration_req = Req(**calibration_kwargs) - calibration_req.set_as_warmup(warmup_steps) - calibration_req.sampling_params.prepare_synthetic_warmup_request_for_queue( - calibration_req, server_args + width, height, num_frames = probe_shape + probe_kwargs = req_kwargs.copy() + probe_kwargs["sampling_params"] = copy(req_kwargs["sampling_params"]) + probe_kwargs.update(width=width, height=height, num_frames=num_frames) + probe_req = Req(**probe_kwargs) + probe_req.set_as_warmup(warmup_steps) + probe_req.sampling_params.prepare_synthetic_warmup_request_for_queue( + probe_req, server_args ) if return_warmup_result: - calibration_req.extra["return_warmup_result"] = True + probe_req.extra["return_warmup_result"] = True if server_based_warmup: - calibration_req.extra["server_based_warmup"] = True - return calibration_req + probe_req.extra["server_based_warmup"] = True + probe_req.extra["auto_residency_full_shape_probe"] = True + return probe_req diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py index 0a4b3034620e..3fb8cac07898 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -405,7 +405,7 @@ def start(self) -> ServerContext: "--log-level=debug", ] if self.extra_args.strip(): - command.extend(self.extra_args.strip().split()) + command.extend(shlex.split(self.extra_args)) access_log_exclude_flag = "--uvicorn-access-log-exclude-prefixes" if not any(arg.startswith(access_log_exclude_flag) for arg in command): command.extend(["--uvicorn-access-log-exclude-prefixes", "/health"]) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index a2d4cccb087b..205e74ee7293 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -42,7 +42,7 @@ ) from sglang.multimodal_gen.runtime.warmup_request_builder import ( SERVER_WARMUP_MAX_VIDEO_FRAMES, - _resolve_calibration_num_frames, + _resolve_auto_residency_probe_shape, _resolve_warmup_num_frames, ) @@ -1306,7 +1306,7 @@ def test_capped_frames_get_the_gpu_alignment_real_requests_get(self): assert num_frames == 21 -class TestCalibrationFrames: +class TestAutoResidencyProbeShape: def _patch_gate(self, monkeypatch, reason: str | None = None) -> None: monkeypatch.setattr( "sglang.multimodal_gen.runtime.warmup_request_builder.auto_residency_args_skip_reason", @@ -1322,71 +1322,80 @@ def _wan_like_args(self) -> SimpleNamespace: num_gpus=1, ) - def _defaults(self, num_frames: int) -> SimpleNamespace: + def _defaults( + self, + num_frames: int, + *, + width: int | None = 1280, + height: int | None = 720, + supported_resolutions=None, + ) -> SimpleNamespace: return SimpleNamespace( + width=width, + height=height, num_frames=num_frames, + supported_resolutions=supported_resolutions, adjust_frames=True, enable_sequence_shard=None, num_frames_round_down=False, ) - def test_capped_video_gets_a_smaller_calibration_size(self, monkeypatch): + def test_capped_video_gets_a_full_shape_probe(self, monkeypatch): self._patch_gate(monkeypatch) - calibration = _resolve_calibration_num_frames( + probe = _resolve_auto_residency_probe_shape( self._wan_like_args(), self._defaults(81), - warmup_num_frames=17, + warmup_shape=(832, 480, 17), server_based_warmup=True, ) - assert calibration == 9 + assert probe == (1280, 720, 81) - def test_uncapped_video_needs_no_calibration(self, monkeypatch): + def test_matching_warmup_needs_no_probe(self, monkeypatch): self._patch_gate(monkeypatch) - calibration = _resolve_calibration_num_frames( + probe = _resolve_auto_residency_probe_shape( self._wan_like_args(), - self._defaults(13), - warmup_num_frames=13, + self._defaults(17, width=832, height=480), + warmup_shape=(832, 480, 17), server_based_warmup=True, ) - assert calibration is None + assert probe is None - def test_unknown_warmup_frames_skip_calibration(self, monkeypatch): + def test_unknown_target_resolution_skips_probe(self, monkeypatch): self._patch_gate(monkeypatch) - calibration = _resolve_calibration_num_frames( + probe = _resolve_auto_residency_probe_shape( self._wan_like_args(), - self._defaults(81), - warmup_num_frames=None, + self._defaults(81, width=None, height=None), + warmup_shape=(832, 480, 17), server_based_warmup=True, ) - assert calibration is None + assert probe is None - def test_skip_gate_disables_calibration(self, monkeypatch): - # the calibration request costs a full extra warmup forward: it must + def test_skip_gate_disables_probe(self, monkeypatch): + # The probe costs a full extra warmup forward, so it must # share the promotion's own gate (kill switch, quantized, manual, ...) self._patch_gate(monkeypatch, reason="performance_mode=manual") - calibration = _resolve_calibration_num_frames( + probe = _resolve_auto_residency_probe_shape( self._wan_like_args(), self._defaults(81), - warmup_num_frames=17, + warmup_shape=(832, 480, 17), server_based_warmup=True, ) - assert calibration is None + assert probe is None - def test_frame_contract_collapse_skips_calibration(self, monkeypatch): - # LongLive2 aligns both 17 and 9 to 29 latent-block frames: a second - # measurement at the same size adds nothing. + def test_supported_resolution_fills_missing_target_size(self, monkeypatch): self._patch_gate(monkeypatch) - args = SimpleNamespace( - pipeline_config=LongLive2T2VConfig(), - num_gpus=1, - ) - calibration = _resolve_calibration_num_frames( - args, - self._defaults(61), - warmup_num_frames=29, + probe = _resolve_auto_residency_probe_shape( + self._wan_like_args(), + self._defaults( + 81, + width=None, + height=None, + supported_resolutions=[(832, 480), (1024, 1024)], + ), + warmup_shape=(832, 480, 17), server_based_warmup=True, ) - assert calibration is None + assert probe == (1024, 1024, 81) class TestAutoResidencySkipReason: diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index 285d81424889..897ec55c15a1 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py @@ -382,6 +382,59 @@ def test_server_based_warmup_uses_model_default_negative_prompt(self): self.assertTrue(req.extra["return_warmup_result"]) self.assertTrue(req.extra["server_based_warmup"]) + def test_auto_residency_adds_a_full_serving_shape_probe(self): + server_args = SimpleNamespace( + warmup_steps=1, + enable_cfg_parallel=False, + enable_torch_compile=False, + enable_breakable_cuda_graph=False, + pipeline_class_name=None, + num_gpus=1, + pipeline_config=SimpleNamespace( + task_type=ModelTaskType.T2V, + adjust_num_frames=lambda value: value, + vae_stride=None, + vae_scale_factor=None, + vae_config=SimpleNamespace(arch_config=None), + ), + is_arg_explicitly_set=lambda _name: False, + ) + sampling_defaults = SamplingParams( + width=1280, + height=720, + num_frames=81, + num_inference_steps=35, + adjust_frames=False, + supported_resolutions=[(1280, 720), (832, 480)], + ) + with ( + patch( + "sglang.multimodal_gen.runtime.warmup_request_builder.get_model_sampling_defaults", + return_value=sampling_defaults, + ), + patch( + "sglang.multimodal_gen.runtime.warmup_request_builder.auto_residency_args_skip_reason", + return_value=None, + ), + ): + reqs = build_warmup_reqs( + server_args, + warmup_resolutions=None, + server_based_warmup=True, + ) + + self.assertEqual( + [(req.width, req.height, req.num_frames) for req in reqs], + [(832, 480, 17), (1280, 720, 81)], + ) + self.assertNotIn("auto_residency_full_shape_probe", reqs[0].extra) + self.assertTrue(reqs[1].extra["auto_residency_full_shape_probe"]) + self.assertEqual(reqs[1].num_inference_steps, 2) + self.assertIn( + "auto residency probe (1280x720x81f, 2/35 steps)", + format_warmup_req(reqs[1]), + ) + def test_server_based_warmup_uses_model_default_resolution(self): server_args = MagicMock() server_args.warmup_steps = 1 diff --git a/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py b/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py index a3898bad831a..2d88774a5f32 100644 --- a/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py +++ b/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py @@ -134,8 +134,14 @@ def test_waits_for_health_after_http_startup(self): ["http://127.0.0.1:11000/health"] * 2, ) - def test_start_cleans_up_process_when_readiness_fails(self): - manager = ServerManager("test-model", port=11000, wait_deadline=1) + def test_start_preserves_quoted_args_and_cleans_up_on_readiness_failure(self): + warmup = '{"height":720,"width":1280}' + manager = ServerManager( + "test-model", + port=11000, + wait_deadline=1, + extra_args=f"--warmup-sampling-params '{warmup}'", + ) process = SimpleNamespace(pid=123, stdout=None) with tempfile.TemporaryDirectory() as temp_dir: @@ -148,7 +154,7 @@ def test_start_cleans_up_process_when_readiness_fails(self): mock.patch( "sglang.multimodal_gen.test.server.test_server_utils.subprocess.Popen", return_value=process, - ), + ) as subprocess_popen, mock.patch.object( manager, "_wait_for_ready", @@ -165,6 +171,8 @@ def test_start_cleans_up_process_when_readiness_fails(self): manager.start() kill_process.assert_called_once_with(123) + command = subprocess_popen.call_args.args[0] + self.assertEqual(command[command.index("--warmup-sampling-params") + 1], warmup) if __name__ == "__main__": From 481e404d3628b6fa83a679c0bb19d379cda95c66 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Aug 2026 16:13:47 +0800 Subject: [PATCH 19/90] [diffusion] calibrate prequantized residency --- .../runtime/managers/gpu_worker.py | 13 ----- .../memory_managers/auto_residency.py | 13 ----- .../runtime/server_args/auto_tune.py | 5 +- .../test/server/test_server_common.py | 6 ++- .../test/unit/test_auto_residency.py | 14 ----- .../test/unit/test_health_warmup_gate.py | 25 +++++++++ .../test/unit/test_performance_metrics.py | 51 ++++++++++++------- 7 files changed, 65 insertions(+), 62 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 0d78d3f2e70e..e5347b0fd9a6 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -70,7 +70,6 @@ format_plan_summary, layerwise_host_pin_capacity_bytes, layerwise_pinned_host_bytes, - module_uses_quantized_weights, plan_auto_residency, plan_summary_payload, rank_candidates_by_h2d_savings, @@ -1339,18 +1338,6 @@ def _build_auto_residency_report( skip_reason = "no server warmup measurements" elif workload.workload_units() is None: skip_reason = "default workload resolution unknown" - else: - quantized_components = [ - name - for name, module in self.pipeline.modules.items() - if isinstance(module, torch.nn.Module) - and module_uses_quantized_weights(module) - ] - if quantized_components: - skip_reason = ( - "loaded quantized components: " - f"{', '.join(sorted(quantized_components))}" - ) if skip_reason is None: # A probe that failed at or below the target is a measurement, not # missing data: the card cannot hold the default workload as it is diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 129248485680..c328bfe7cf72 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -31,7 +31,6 @@ import torch import torch.nn as nn -from sglang.multimodal_gen.runtime.layers.linear import LinearBase from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD, @@ -389,18 +388,6 @@ def _module_weight_bytes(module: nn.Module) -> int: return param_bytes + buffer_bytes -def module_uses_quantized_weights(module: nn.Module) -> bool: - """Whether a loaded native component contains quantized weight storage.""" - for submodule in module.modules(): - if isinstance(submodule, LinearBase) and submodule.quant_config is not None: - return True - quantized_float_dtypes = {torch.float8_e4m3fn, torch.float8_e5m2} - return any( - not tensor.is_floating_point() or tensor.dtype in quantized_float_dtypes - for _, tensor in iter_materialized_weights(module) - ) - - def component_runtime_weight_bytes(modules: Mapping[str, object]) -> dict[str, int]: """Persistent plus streamed-window weight bytes during a component use.""" runtime_bytes: dict[str, int] = {} diff --git a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py index ee9fde16621a..80045e77c13c 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -86,8 +86,9 @@ def auto_residency_args_skip_reason(server_args: ServerArgs) -> str | None: or server_args.transformer_weights_path or svdquant_enabled ): - # Residency changes can perturb quantized numerical paths through a - # different allocation and kernel layout. + # Explicit and online quantization require a fixed loading path. + # Self-describing pre-quantized checkpoints are safe to move after + # loading: residency changes storage location, not quantized values. return "quantized checkpoint" if server_args.direct_gpu_weight_loading: return "direct GPU weight loading requires a fixed resident placement" diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 829ea67dfe45..894d65dbf202 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -121,7 +121,11 @@ def _case_warmup_sampling_params(case: DiffusionTestCase) -> dict[str, Any]: if sampling.output_size: width, height = sampling.output_size.lower().split("x", 1) overrides.update(width=int(width), height=int(height)) - if sampling.num_frames is not None: + if case.server_args.modality == "image": + # Omni pipelines can expose a video-oriented model default even when + # this case calls /v1/images, whose implicit temporal extent is one. + overrides["num_frames"] = 1 + elif sampling.num_frames is not None: overrides["num_frames"] = sampling.num_frames elif sampling.fps is not None: # /v1/videos resolves an omitted num_frames from seconds * fps. diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 205e74ee7293..fd13d439212f 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -9,7 +9,6 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType from sglang.multimodal_gen.configs.pipeline_configs.longlive2 import LongLive2T2VConfig -from sglang.multimodal_gen.runtime.layers.linear import LinearBase from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( ACTIVATION_EXTRAPOLATION_MARGIN, GIB_BYTES, @@ -27,7 +26,6 @@ estimate_default_workload_peak_bytes, estimate_workload_phase_peaks, format_applied_changes, - module_uses_quantized_weights, plan_auto_residency, rank_candidates_by_h2d_savings, rollback_promotions, @@ -264,18 +262,6 @@ def test_phase_estimation_preserves_active_component_membership(self): assert active == {"denoise": ("transformer",)} -def test_loaded_native_quantization_config_is_detected_with_bf16_weights(): - linear = LinearBase.__new__(LinearBase) - nn.Module.__init__(linear) - linear.quant_config = object() - linear.register_parameter( - "weight", - nn.Parameter(torch.ones(1, dtype=torch.bfloat16), requires_grad=False), - ) - - assert module_uses_quantized_weights(linear) - - def _candidate( name: str, *, diff --git a/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py b/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py index 2d88774a5f32..f72efe1d1634 100644 --- a/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py +++ b/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py @@ -18,6 +18,9 @@ health_generate, liveness, ) +from sglang.multimodal_gen.test.server.test_server_common import ( + _case_warmup_sampling_params, +) from sglang.multimodal_gen.test.server.test_server_utils import ServerManager @@ -112,6 +115,28 @@ def poll(self): class TestServerManagerReadiness(unittest.TestCase): + def test_image_case_warmup_uses_one_frame_for_omni_model(self): + case = SimpleNamespace( + server_args=SimpleNamespace(modality="image"), + sampling_params=SimpleNamespace( + output_size="832x480", + num_frames=None, + fps=None, + seconds=1, + extras={"num_inference_steps": 35}, + ), + ) + + self.assertEqual( + _case_warmup_sampling_params(case), + { + "width": 832, + "height": 480, + "num_frames": 1, + "num_inference_steps": 35, + }, + ) + def test_waits_for_health_after_http_startup(self): manager = ServerManager("test-model", port=11000, wait_deadline=1) with tempfile.TemporaryDirectory() as temp_dir: diff --git a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py index 67b183a1cd2b..bbc11d6e2c57 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -125,7 +125,7 @@ def test_server_warmup_preserves_peak_after_managed_stage_timeline(): assert record.phase_active_components["request:untracked"] == () -def test_loaded_quantized_checkpoint_disables_auto_residency(): +def test_loaded_prequantized_checkpoint_can_use_auto_residency(): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 quantized_module = torch.nn.Module() @@ -133,28 +133,41 @@ def test_loaded_quantized_checkpoint_disables_auto_residency(): "weight", torch.nn.Parameter(torch.ones(1, dtype=torch.int8), requires_grad=False), ) - worker.pipeline = SimpleNamespace(modules={"transformer": quantized_module}) - - report = worker._build_auto_residency_report( - workload=DefaultWorkload( - width=64, - height=64, - num_frames=1, - num_inference_steps=1, - ), - records=[ - WarmupMemoryRecord( + worker.pipeline = SimpleNamespace( + modules={"transformer": quantized_module}, + component_residency_strategies={}, + ) + worker.server_args = SimpleNamespace( + num_gpus=1, + node_rank=0, + residency_mode=lambda _name: "resident", + explicit_residency_mode=lambda _name: None, + ) + device_module = Mock() + device_module.mem_get_info.return_value = (100, 100) + device_module.memory_reserved.return_value = 0 + + with patch.object(torch, "get_device_module", return_value=device_module): + report = worker._build_auto_residency_report( + workload=DefaultWorkload( width=64, height=64, num_frames=1, - baseline_allocated_bytes=1, - peak_reserved_bytes=2, - succeeded=True, - ) - ], - ) + num_inference_steps=1, + ), + records=[ + WarmupMemoryRecord( + width=64, + height=64, + num_frames=1, + baseline_allocated_bytes=1, + peak_reserved_bytes=2, + succeeded=True, + ) + ], + ) - assert report.skip_reason == "loaded quantized components: transformer" + assert report.skip_reason is None def test_baseline_config_loads_per_scenario_peak_vram(tmp_path): From ae3dc54070d94852e1357ee4c294ab881543f56a Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Aug 2026 17:43:40 +0800 Subject: [PATCH 20/90] [diffusion] harden residency workload calibration --- .../runtime/managers/gpu_worker.py | 2 + .../memory_managers/auto_residency.py | 55 ++++++++- .../runtime/warmup_request_builder.py | 60 +++------- .../test/server/test_server_common.py | 5 +- .../test/unit/test_auto_residency.py | 105 ++++++++++++++++-- .../test/unit/test_cfg_parallel_warmup.py | 14 +-- .../test/unit/test_health_warmup_gate.py | 17 +++ 7 files changed, 191 insertions(+), 67 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index e5347b0fd9a6..191bb3fefeca 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -74,6 +74,7 @@ plan_summary_payload, rank_candidates_by_h2d_savings, resolve_default_workload, + resolve_measured_default_workload, rollback_promotions, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( @@ -1205,6 +1206,7 @@ def apply_auto_residency(self) -> OutputBatch: try: workload = resolve_default_workload(self.server_args) + workload = resolve_measured_default_workload(workload, records) local_report = self._build_auto_residency_report( workload=workload, records=records ) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index c328bfe7cf72..a9f3fcf74a18 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -226,6 +226,35 @@ def resolve_default_workload(server_args: ServerArgs) -> DefaultWorkload: ) +def resolve_measured_default_workload( + workload: DefaultWorkload, records: Iterable[WarmupMemoryRecord] +) -> DefaultWorkload: + """Fill an implicit default resolution from the executed warmup. + + Image-edit pipelines can derive their output size from the input image, so + the sampling defaults legitimately omit width and height. The warmup record + is captured after input validation and therefore contains the effective + serving shape. Keep the model-default frame count because video warmup may + intentionally cap frames before measurement. + """ + if workload.workload_units() is not None: + return workload + measured = [ + record + for record in records + if record.succeeded and record.width > 0 and record.height > 0 + ] + if not measured: + return workload + representative = max(measured, key=lambda record: record.width * record.height) + return DefaultWorkload( + width=representative.width, + height=representative.height, + num_frames=workload.num_frames, + num_inference_steps=workload.num_inference_steps, + ) + + def estimate_default_workload_peak_bytes( *, records: Iterable[WarmupMemoryRecord], @@ -801,25 +830,43 @@ def _vram_reserve_bytes(budget_bytes: int) -> int: def _binding_phase_constraints( report: RankResidencyReport, + candidate_component_names: set[str], ) -> list[tuple[str, int, tuple[str, ...]]]: - """Keep only the highest peak for each active-component placement. + """Keep the binding measured and conservative unobserved phases. Every candidate has the same device delta in phases with the same active component set. The lower peaks are therefore provably redundant; removing them reduces the exact optimizer's resource dimension without changing - its feasible placements. + its feasible placements. A short calibration may not reach a later + timestep-routed component (for example the second DiT in a two-stage + denoiser), so give each unobserved candidate a synthetic phase at the + measured request peak. This prevents permanent residents from being + placed in memory that the later component may need. """ phase_peaks = report.estimated_peak_bytes_by_phase if not phase_peaks: if report.estimated_peak_bytes is None: return [] phase_peaks = {"request": report.estimated_peak_bytes} + request_peak = ( + report.estimated_peak_bytes + if report.estimated_peak_bytes is not None + else max(phase_peaks.values()) + ) binding: dict[tuple[str, ...], tuple[str, int]] = {} for phase_name, phase_peak in phase_peaks.items(): active = tuple(sorted(report.active_components_by_phase.get(phase_name, ()))) current = binding.get(active) if current is None or (phase_peak, phase_name) > (current[1], current[0]): binding[active] = (phase_name, phase_peak) + observed_components = { + component_name for active in binding for component_name in active + } + for component_name in sorted(candidate_component_names - observed_components): + binding[(component_name,)] = ( + f"unobserved:{component_name}", + request_peak, + ) return [ (f"gpu:rank{report.rank}:{phase_name}", phase_peak, active) for active, (phase_name, phase_peak) in sorted( @@ -920,8 +967,10 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP if not candidates: return _skip_plan("no implicitly offloaded components to promote") + candidate_component_names = {candidate.component_name for candidate in candidates} phase_constraints = { - report.rank: _binding_phase_constraints(report) for report in reports + report.rank: _binding_phase_constraints(report, candidate_component_names) + for report in reports } resource_budgets: dict[str, int] = {} for report in reports: diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py index b2b5d37bf2d8..0fc445e4a0f5 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -6,10 +6,11 @@ sampling defaults, then keeps startup bounded by choosing common low-cost resolution/frame buckets and trimming the denoising step count. -When warmup-calibrated auto residency is active, a second probe restores the -full default serving shape but keeps the trimmed step count. Memory depends on -the activation shape rather than the number of repeated denoising steps, so -this directly measures placement headroom without running a full generation. +When warmup-calibrated auto residency is active, warmup uses the full default +serving shape but keeps the trimmed step count. Memory depends on the +activation shape rather than the number of repeated denoising steps, so this +directly measures placement headroom without running a full generation or an +extra stateful pipeline request. Image models may run a tiny second step because first/last step paths often initialize different kernels or scheduler state. Video models cap frames and @@ -405,14 +406,14 @@ def resolve_default_workload_shape( return width, height, num_frames -def _resolve_auto_residency_probe_shape( +def _resolve_auto_residency_warmup_shape( server_args: ServerArgs, sampling_defaults: SamplingParams, *, warmup_shape: tuple[int, int, int | None], server_based_warmup: bool, ) -> tuple[int, int, int] | None: - """Return a full serving-shape memory probe when bounded warmup is smaller. + """Return the full serving shape when bounded warmup is smaller. The probe still runs only the bounded warmup step count. Denoising steps repeat the same activation shape, so one full-shape forward measures the @@ -544,8 +545,8 @@ def build_warmup_reqs( sampling_defaults, server_based_warmup=server_based_warmup, ) - auto_residency_probe_shape = ( - _resolve_auto_residency_probe_shape( + auto_residency_warmup_shape = ( + _resolve_auto_residency_warmup_shape( server_args, sampling_defaults, warmup_shape=(width, height, warmup_num_frames), @@ -554,6 +555,9 @@ def build_warmup_reqs( if warmup_resolutions is None else None ) + if auto_residency_warmup_shape is not None: + width, height, warmup_num_frames = auto_residency_warmup_shape + resolutions[0] = (width, height) # build warmup reqs warmup_reqs = [] @@ -614,44 +618,8 @@ def build_warmup_reqs( req.extra["return_warmup_result"] = True if server_based_warmup: req.extra["server_based_warmup"] = True + if auto_residency_warmup_shape is not None: + req.extra["auto_residency_full_shape_probe"] = True warmup_reqs.append(req) - if auto_residency_probe_shape is not None and (width, height) == resolutions[0]: - warmup_reqs.append( - _build_auto_residency_probe_req( - req_kwargs, - probe_shape=auto_residency_probe_shape, - warmup_steps=warmup_steps, - return_warmup_result=return_warmup_result, - server_based_warmup=server_based_warmup, - server_args=server_args, - ) - ) - return warmup_reqs - - -def _build_auto_residency_probe_req( - req_kwargs: dict, - *, - probe_shape: tuple[int, int, int], - warmup_steps: int, - return_warmup_result: bool, - server_based_warmup: bool, - server_args: ServerArgs, -) -> Req: - width, height, num_frames = probe_shape - probe_kwargs = req_kwargs.copy() - probe_kwargs["sampling_params"] = copy(req_kwargs["sampling_params"]) - probe_kwargs.update(width=width, height=height, num_frames=num_frames) - probe_req = Req(**probe_kwargs) - probe_req.set_as_warmup(warmup_steps) - probe_req.sampling_params.prepare_synthetic_warmup_request_for_queue( - probe_req, server_args - ) - if return_warmup_result: - probe_req.extra["return_warmup_result"] = True - if server_based_warmup: - probe_req.extra["server_based_warmup"] = True - probe_req.extra["auto_residency_full_shape_probe"] = True - return probe_req diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 894d65dbf202..7981daf519f6 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -22,6 +22,7 @@ import requests from openai import OpenAI +from sglang.multimodal_gen.runtime.entrypoints.openai.utils import DEFAULT_FPS from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord @@ -127,9 +128,9 @@ def _case_warmup_sampling_params(case: DiffusionTestCase) -> dict[str, Any]: overrides["num_frames"] = 1 elif sampling.num_frames is not None: overrides["num_frames"] = sampling.num_frames - elif sampling.fps is not None: + else: # /v1/videos resolves an omitted num_frames from seconds * fps. - overrides["num_frames"] = sampling.seconds * sampling.fps + overrides["num_frames"] = sampling.seconds * (sampling.fps or DEFAULT_FPS) if sampling.fps is not None: overrides["fps"] = sampling.fps for name in ("num_inference_steps", "guidance_scale", "true_cfg_scale"): diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index fd13d439212f..5413c1510c24 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -16,6 +16,7 @@ AppliedPromotion, AutoResidencyPlan, AutoResidencyRollbackError, + DefaultWorkload, PromotionCandidate, RankResidencyReport, WarmupMemoryRecord, @@ -28,6 +29,7 @@ format_applied_changes, plan_auto_residency, rank_candidates_by_h2d_savings, + resolve_measured_default_workload, rollback_promotions, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( @@ -40,7 +42,7 @@ ) from sglang.multimodal_gen.runtime.warmup_request_builder import ( SERVER_WARMUP_MAX_VIDEO_FRAMES, - _resolve_auto_residency_probe_shape, + _resolve_auto_residency_warmup_shape, _resolve_warmup_num_frames, ) @@ -317,6 +319,60 @@ def _report( ) +class TestResolveMeasuredDefaultWorkload: + def test_uses_effective_warmup_resolution_for_implicit_image_size(self): + workload = DefaultWorkload( + width=None, + height=None, + num_frames=1, + num_inference_steps=40, + ) + + resolved = resolve_measured_default_workload( + workload, + [ + _record(width=512, height=512), + _record(width=1024, height=1024), + ], + ) + + assert resolved == DefaultWorkload( + width=1024, + height=1024, + num_frames=1, + num_inference_steps=40, + ) + + def test_keeps_default_frames_when_warmup_caps_video(self): + workload = DefaultWorkload( + width=None, + height=None, + num_frames=81, + num_inference_steps=30, + ) + + resolved = resolve_measured_default_workload( + workload, [_record(width=832, height=480, num_frames=17)] + ) + + assert resolved.num_frames == 81 + + def test_does_not_replace_explicit_default_shape(self): + workload = DefaultWorkload( + width=1280, + height=720, + num_frames=81, + num_inference_steps=30, + ) + + assert ( + resolve_measured_default_workload( + workload, [_record(width=512, height=512)] + ) + is workload + ) + + class TestPlanAutoResidency: def test_rank_skip_reason_propagates(self): plan = plan_auto_residency( @@ -428,6 +484,37 @@ def test_other_phase_can_block_component_residency(self): assert not plan.promotions + def test_unobserved_later_component_gets_a_conservative_phase(self): + first_transformer = _candidate("first_transformer", weight_gib=26, h2d_gib=100) + second_transformer = PromotionCandidate( + component_name="second_transformer", + residency_mode=LAYERWISE_OFFLOAD, + promoted_weight_bytes=26 * GIB_BYTES, + h2d_bytes_per_request=90 * GIB_BYTES, + target_layerwise_resident_layers=(40,), + target_layerwise_pinned_layers=((),), + active_device_delta_bytes=26 * GIB_BYTES, + inactive_device_delta_bytes=0, + ) + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=80, + estimated_gib=30, + candidates=[first_transformer, second_transformer], + phase_peaks_gib={"early_denoise": 30}, + phase_components={ + "early_denoise": ("first_transformer",), + }, + ) + ] + ) + + assert [candidate.component_name for candidate in plan.promotions] == [ + "first_transformer" + ] + assert "gpu:rank0:unobserved:second_transformer" in plan.resource_budget_bytes + def test_redundant_phase_constraints_collapse_to_the_highest_peak(self): plan = plan_auto_residency( reports=[ @@ -1292,7 +1379,7 @@ def test_capped_frames_get_the_gpu_alignment_real_requests_get(self): assert num_frames == 21 -class TestAutoResidencyProbeShape: +class TestAutoResidencyWarmupShape: def _patch_gate(self, monkeypatch, reason: str | None = None) -> None: monkeypatch.setattr( "sglang.multimodal_gen.runtime.warmup_request_builder.auto_residency_args_skip_reason", @@ -1328,7 +1415,7 @@ def _defaults( def test_capped_video_gets_a_full_shape_probe(self, monkeypatch): self._patch_gate(monkeypatch) - probe = _resolve_auto_residency_probe_shape( + probe = _resolve_auto_residency_warmup_shape( self._wan_like_args(), self._defaults(81), warmup_shape=(832, 480, 17), @@ -1338,7 +1425,7 @@ def test_capped_video_gets_a_full_shape_probe(self, monkeypatch): def test_matching_warmup_needs_no_probe(self, monkeypatch): self._patch_gate(monkeypatch) - probe = _resolve_auto_residency_probe_shape( + probe = _resolve_auto_residency_warmup_shape( self._wan_like_args(), self._defaults(17, width=832, height=480), warmup_shape=(832, 480, 17), @@ -1348,7 +1435,7 @@ def test_matching_warmup_needs_no_probe(self, monkeypatch): def test_unknown_target_resolution_skips_probe(self, monkeypatch): self._patch_gate(monkeypatch) - probe = _resolve_auto_residency_probe_shape( + probe = _resolve_auto_residency_warmup_shape( self._wan_like_args(), self._defaults(81, width=None, height=None), warmup_shape=(832, 480, 17), @@ -1357,10 +1444,10 @@ def test_unknown_target_resolution_skips_probe(self, monkeypatch): assert probe is None def test_skip_gate_disables_probe(self, monkeypatch): - # The probe costs a full extra warmup forward, so it must - # share the promotion's own gate (kill switch, quantized, manual, ...) + # Full-shape calibration must share the promotion's own gate (kill + # switch, quantized, manual, ...). self._patch_gate(monkeypatch, reason="performance_mode=manual") - probe = _resolve_auto_residency_probe_shape( + probe = _resolve_auto_residency_warmup_shape( self._wan_like_args(), self._defaults(81), warmup_shape=(832, 480, 17), @@ -1370,7 +1457,7 @@ def test_skip_gate_disables_probe(self, monkeypatch): def test_supported_resolution_fills_missing_target_size(self, monkeypatch): self._patch_gate(monkeypatch) - probe = _resolve_auto_residency_probe_shape( + probe = _resolve_auto_residency_warmup_shape( self._wan_like_args(), self._defaults( 81, diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index 897ec55c15a1..2d2b7957e41a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py @@ -382,7 +382,7 @@ def test_server_based_warmup_uses_model_default_negative_prompt(self): self.assertTrue(req.extra["return_warmup_result"]) self.assertTrue(req.extra["server_based_warmup"]) - def test_auto_residency_adds_a_full_serving_shape_probe(self): + def test_auto_residency_uses_one_full_serving_shape_probe(self): server_args = SimpleNamespace( warmup_steps=1, enable_cfg_parallel=False, @@ -423,16 +423,16 @@ def test_auto_residency_adds_a_full_serving_shape_probe(self): server_based_warmup=True, ) + self.assertEqual(len(reqs), 1) self.assertEqual( - [(req.width, req.height, req.num_frames) for req in reqs], - [(832, 480, 17), (1280, 720, 81)], + (reqs[0].width, reqs[0].height, reqs[0].num_frames), + (1280, 720, 81), ) - self.assertNotIn("auto_residency_full_shape_probe", reqs[0].extra) - self.assertTrue(reqs[1].extra["auto_residency_full_shape_probe"]) - self.assertEqual(reqs[1].num_inference_steps, 2) + self.assertTrue(reqs[0].extra["auto_residency_full_shape_probe"]) + self.assertEqual(reqs[0].num_inference_steps, 2) self.assertIn( "auto residency probe (1280x720x81f, 2/35 steps)", - format_warmup_req(reqs[1]), + format_warmup_req(reqs[0]), ) def test_server_based_warmup_uses_model_default_resolution(self): diff --git a/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py b/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py index f72efe1d1634..52f68952be03 100644 --- a/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py +++ b/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py @@ -137,6 +137,23 @@ def test_image_case_warmup_uses_one_frame_for_omni_model(self): }, ) + def test_video_case_warmup_matches_endpoint_default_fps(self): + case = SimpleNamespace( + server_args=SimpleNamespace(modality="video"), + sampling_params=SimpleNamespace( + output_size="", + num_frames=None, + fps=None, + seconds=2, + extras={}, + ), + ) + + self.assertEqual( + _case_warmup_sampling_params(case), + {"num_frames": 48}, + ) + def test_waits_for_health_after_http_startup(self): manager = ServerManager("test-model", port=11000, wait_deadline=1) with tempfile.TemporaryDirectory() as temp_dir: From de89cedc4b508c959f7ff78c4ed83319d8300b5c Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Aug 2026 17:43:56 +0800 Subject: [PATCH 21/90] [diffusion] gate unsupported quantized offload --- .../runtime/loader/transformer_load_utils.py | 44 +++++++++++++++++++ .../runtime/models/dits/base.py | 4 ++ .../runtime/models/dits/ideogram.py | 4 ++ .../test/unit/test_transformer_quant.py | 33 ++++++++++++++ 4 files changed, 85 insertions(+) diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 3216ce3bc980..51895c7a0bb3 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -252,6 +252,43 @@ def _reject_explicit_component_selector( ) +def _enforce_quantized_component_offload_capability( + *, + server_args: ServerArgs, + quant_config: Optional[QuantizationConfig], + component_name: str | None, + model_cls: type[nn.Module], +) -> None: + if quant_config is None or model_cls.supports_quantized_component_offload: + return + if not _uses_component_offload( + server_args, + component_name, + legacy_enabled=bool(server_args.dit_cpu_offload), + ): + return + + feature_name = f"quantized {model_cls.__name__} checkpoints" + _reject_explicit_component_selector( + server_args, + component_name, + feature_name=feature_name, + ) + if component_name is None: + server_args.dit_cpu_offload = False + else: + server_args.require_component_resident( + component_name, + feature_name=feature_name, + ) + logger.info( + "%s keeps %s resident because its quantized runtime layout does not " + "support component offload", + model_cls.__name__, + component_name or "the transformer", + ) + + class _NunchakuQuantAdapter(_TransformerQuantAdapter): """Adapter for Nunchaku checkpoints""" @@ -813,6 +850,13 @@ def resolve_transformer_quant_load_spec( vars(model_cls).get("param_names_mapping", {}) ) + _enforce_quantized_component_offload_capability( + server_args=server_args, + quant_config=quant_config, + component_name=component_name, + model_cls=model_cls, + ) + nunchaku_config = server_args.nunchaku_config # resolve target param dtype diff --git a/python/sglang/multimodal_gen/runtime/models/dits/base.py b/python/sglang/multimodal_gen/runtime/models/dits/base.py index e5d545ca95eb..aa1042de9f83 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/base.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/base.py @@ -32,6 +32,10 @@ class BaseDiT(nn.Module, ABC): # shard conditions left in the root group stays sharded unless the entry # point is registered; loaders read this and register each name. _fsdp_forward_methods: tuple[str, ...] = () + # Quantized modules normally support the same coarse CPU residency as + # floating-point modules. Implementations whose packed runtime layout + # cannot survive that round trip override this after quant detection. + supports_quantized_component_offload: bool = True param_names_mapping: dict reverse_param_names_mapping: dict hidden_size: int diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py b/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py index 30cecf1cbc21..b86275cd1872 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py @@ -640,6 +640,10 @@ class Ideogram4Transformer2DModel(BaseDiT, LayerwiseOffloadableModuleMixin): AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA, } + # The coarse CPU component-offload load changes the packed Ideogram + # numerical path even if the module is promoted before the real request. + # Layerwise offload uses its own managed storage layout. + supports_quantized_component_offload = False param_names_mapping = Ideogram4DiTConfig().arch_config.param_names_mapping reverse_param_names_mapping = {} diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index 4a847ec7a4ca..b4185e5bb0e9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -96,6 +96,7 @@ def _loads(input_str, _flags=None): ) from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( TransformerQuantLoadSpec, + _enforce_quantized_component_offload_capability, _filter_duplicate_precision_variant_safetensors, _Flux2Nvfp4FallbackAdapter, _needs_device_weight_postprocess, @@ -136,6 +137,14 @@ def get_name(cls): return "modelopt_fp4" +class _QuantizedOffloadSafeModel: + supports_quantized_component_offload = True + + +class _QuantizedOffloadUnsafeModel: + supports_quantized_component_offload = False + + def _make_quant_config(name: str, **attrs): cls = type( f"_Fake{name.title().replace('_', '')}QuantConfig", @@ -149,6 +158,30 @@ def _make_quant_config(name: str, **attrs): class TestTransformerQuantHelpers(unittest.TestCase): + def test_model_capability_can_keep_quantized_transformer_resident(self): + server_args = SimpleNamespace(dit_cpu_offload=True) + + _enforce_quantized_component_offload_capability( + server_args=server_args, + quant_config=_FakeQuantConfig(), + component_name=None, + model_cls=_QuantizedOffloadUnsafeModel, + ) + + self.assertFalse(server_args.dit_cpu_offload) + + def test_model_capability_preserves_supported_quantized_offload(self): + server_args = SimpleNamespace(dit_cpu_offload=True) + + _enforce_quantized_component_offload_capability( + server_args=server_args, + quant_config=_FakeQuantConfig(), + component_name=None, + model_cls=_QuantizedOffloadSafeModel, + ) + + self.assertTrue(server_args.dit_cpu_offload) + def test_autoround_config_is_inferred_and_remapped_to_native_prefixes(self): layer_config = { "bits": 4, From 2f22951a07cd2bf849cdb6f764781bc48e4695ea Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Aug 2026 17:44:09 +0800 Subject: [PATCH 22/90] [diffusion] keep runtime LoRA weights inference-only --- .../runtime/layers/lora/linear.py | 10 +++++---- .../test/unit/test_lora_inference_mode.py | 21 +++++++++++++++++++ 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py index 2bba150510eb..28af3ca14165 100644 --- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py @@ -220,10 +220,12 @@ def set_lora_weights( If False, append to existing list (for multi-LoRA support). output_offset: Optional constant output term paired with this adapter """ - lora_A_param = torch.nn.Parameter( - A - ) # share storage with weights in the pipeline - lora_B_param = torch.nn.Parameter(B) + # This runtime never trains adapters. In particular, weights loaded + # during a warmup may be inference tensors; leaving Parameter's + # requires_grad default enabled makes a later dynamic-LoRA request try + # to read their unavailable version counters while sharding views. + lora_A_param = torch.nn.Parameter(A, requires_grad=False) + lora_B_param = torch.nn.Parameter(B, requires_grad=False) output_offset_param = ( torch.nn.Parameter(output_offset, requires_grad=False) if output_offset is not None diff --git a/python/sglang/multimodal_gen/test/unit/test_lora_inference_mode.py b/python/sglang/multimodal_gen/test/unit/test_lora_inference_mode.py index ae1b7db42242..4a4bd83bacb8 100644 --- a/python/sglang/multimodal_gen/test/unit/test_lora_inference_mode.py +++ b/python/sglang/multimodal_gen/test/unit/test_lora_inference_mode.py @@ -49,3 +49,24 @@ def test_lora_merge_unmerge_handles_inference_base_weight(): assert not layer.merged assert not layer.base_layer.weight.is_inference() assert torch.allclose(layer.base_layer.weight, base_weight) + + +def test_dynamic_lora_reuses_inference_weights_without_autograd_tracking(): + base_layer = nn.Linear(4, 3, bias=False) + layer = LinearWithLoRA(base_layer, lora_rank=2, lora_alpha=2) + + with torch.inference_mode(): + layer.set_lora_weights( + torch.ones(2, 4), + torch.ones(3, 2), + clear_existing=True, + merge_weights=False, + ) + + assert layer.lora_A.is_inference() + assert layer.lora_B.is_inference() + assert not layer.lora_A.requires_grad + assert not layer.lora_B.requires_grad + with torch.no_grad(): + sharded_view = layer.lora_B[:2, :] + assert sharded_view.shape == (2, 2) From 992f4c1a033161b0f219808530f08b3cb0cce105 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Aug 2026 17:44:24 +0800 Subject: [PATCH 23/90] [diffusion] avoid CUDA sync in timestep logs --- .../runtime/pipelines_core/stages/causal_denoising.py | 3 +-- .../pipelines_core/stages/timestep_preparation.py | 11 ++++++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py index 5a225e731959..0a9db62d7914 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py @@ -218,9 +218,8 @@ def _prepare_causal_dmd_timesteps( (scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)) ) timesteps = scheduler_timesteps[1000 - timesteps] - timesteps = timesteps.to(device) logger.info("Using timesteps: %s", timesteps) - return timesteps + return timesteps.to(device) def _prepare_causal_dmd_image_kwargs( self, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py index 21430f51c825..5abbb2ba4e5e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py @@ -162,7 +162,12 @@ def forward( batch.timesteps = timesteps batch.scheduler = scheduler if not batch.is_warmup: - self.log_debug("timesteps: %s", timesteps) + self.log_debug( + "timesteps prepared: shape=%s, dtype=%s, device=%s", + tuple(timesteps.shape), + timesteps.dtype, + timesteps.device, + ) return batch def build_dedup_fingerprint( @@ -244,10 +249,10 @@ def forward( ) timesteps = scheduler_timesteps[num_train_timesteps - timesteps] + if not batch.is_warmup: + self.log_debug("DMD timesteps: %s", timesteps) batch.timesteps = timesteps.to(get_local_torch_device()) batch.scheduler = scheduler - if not batch.is_warmup: - self.log_debug("DMD timesteps: %s", batch.timesteps) return batch def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: From c044a0bd4174a7ff4a32bbf79168bc292b3ba28c Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Aug 2026 18:08:20 +0800 Subject: [PATCH 24/90] [diffusion] preserve undeclared quantized offload capability --- .../runtime/loader/transformer_load_utils.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 51895c7a0bb3..44d912378cb5 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -259,7 +259,13 @@ def _enforce_quantized_component_offload_capability( component_name: str | None, model_cls: type[nn.Module], ) -> None: - if quant_config is None or model_cls.supports_quantized_component_offload: + supports_offload = True + for cls in model_cls.__mro__: + capability = vars(cls).get("supports_quantized_component_offload") + if capability is not None: + supports_offload = bool(capability) + break + if quant_config is None or supports_offload: return if not _uses_component_offload( server_args, From c497b0f081c1ce682ad2c536fc9b71e75eac74fb Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Aug 2026 18:46:36 +0800 Subject: [PATCH 25/90] [diffusion] plan residency from live VRAM --- .../runtime/managers/gpu_worker.py | 35 ++++++------- .../memory_managers/auto_residency.py | 30 ++++++----- .../memory_managers/component_manager.py | 37 ++++++++++---- .../test/unit/test_auto_residency.py | 50 +++++++++++++------ .../test/unit/test_component_residency.py | 48 +++++++++--------- .../test/unit/test_performance_metrics.py | 42 ++++++++++++++-- 6 files changed, 155 insertions(+), 87 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 191bb3fefeca..9d592c79b305 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -562,11 +562,9 @@ def _execute_forward_common( warmup_baseline_allocated_bytes = 0 try: if measure_server_warmup: - # Drop the previous request's allocator pool so max_memory - # _reserved measures THIS workload's peak. Without this, a - # later warmup shape that fits in the cached pool reports the - # same peak as the earlier one and the calibration slope - # collapses to zero. + # Drop the previous request's allocator pool so each probe + # starts from the same placement and can return released + # component storage before its allocated peak is measured. torch.get_device_module().empty_cache() if not current_platform.is_cpu() and not current_platform.is_mps(): torch.get_device_module().reset_peak_memory_stats() @@ -700,24 +698,23 @@ def _execute_forward_common( def _record_server_warmup_memory( self, *, req: Req, baseline_allocated_bytes: int, succeeded: bool ) -> None: - phase_peaks: dict[str, int] = {} + phase_allocated_peaks: dict[str, int] = {} phase_components: dict[str, tuple[str, ...]] = {} residency_manager = peek_global_component_residency_manager() if residency_manager is not None: - for phase_name, ( - components, - peak_bytes, - ) in residency_manager.take_warmup_phase_peaks().items(): - phase_peaks[phase_name] = peak_bytes - phase_components[phase_name] = components - request_peak = max( - int(torch.get_device_module().max_memory_reserved()), - max(phase_peaks.values(), default=0), + for phase_name, peak in residency_manager.take_warmup_phase_peaks().items(): + phase_allocated_peaks[phase_name] = peak.allocated_bytes + phase_components[phase_name] = peak.active_components + request_allocated_peak = max( + int(torch.get_device_module().max_memory_allocated()), + max(phase_allocated_peaks.values(), default=0), ) - if request_peak > max(phase_peaks.values(), default=0): + if request_allocated_peak > max(phase_allocated_peaks.values(), default=0): # Work after the residency-managed stage timeline (for example, # output materialization) must remain a placement constraint. - phase_peaks["request:untracked"] = request_peak + # A reserved-only increase is reclaimable allocator cache, not a + # second live placement, and is covered by the explicit reserve. + phase_allocated_peaks["request:untracked"] = request_allocated_peak phase_components["request:untracked"] = () self._auto_residency_warmup_records.append( WarmupMemoryRecord( @@ -725,9 +722,9 @@ def _record_server_warmup_memory( height=int(req.height or 0), num_frames=int(req.num_frames or 1), baseline_allocated_bytes=int(baseline_allocated_bytes), - peak_reserved_bytes=request_peak, + peak_allocated_bytes=request_allocated_peak, succeeded=succeeded, - phase_peak_reserved_bytes=phase_peaks, + phase_peak_allocated_bytes=phase_allocated_peaks, phase_active_components=phase_components, ) ) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index a9f3fcf74a18..0444638cfc48 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -109,9 +109,9 @@ class WarmupMemoryRecord(msgspec.Struct, frozen=True): height: int num_frames: int baseline_allocated_bytes: int - peak_reserved_bytes: int + peak_allocated_bytes: int succeeded: bool - phase_peak_reserved_bytes: dict[str, int] = {} + phase_peak_allocated_bytes: dict[str, int] = {} phase_active_components: dict[str, tuple[str, ...]] = {} def workload_units(self) -> int: @@ -261,7 +261,12 @@ def estimate_default_workload_peak_bytes( target_units: int | None, constant_weight_bytes: int = 0, ) -> int | None: - """Extrapolate warmup peaks to the default workload. + """Extrapolate live warmup memory to the default workload. + + Real measurements use allocated bytes. Cached allocator blocks are + reclaimable and are covered by the explicit VRAM reserve; treating them as + live memory would charge the same storage again when adding resident + weights. Preference order: 1. A measurement at or above the target workload bounds the peak directly. @@ -299,9 +304,8 @@ def estimate_default_workload_peak_bytes( peak_by_units: dict[int, int] = {} for record in records: units = record.workload_units() - peak_by_units[units] = max( - peak_by_units.get(units, 0), record.peak_reserved_bytes - ) + peak = record.peak_allocated_bytes + peak_by_units[units] = max(peak_by_units.get(units, 0), peak) covering_peaks = [ peak for units, peak in peak_by_units.items() if units >= target_units @@ -328,7 +332,7 @@ def estimate_default_workload_peak_bytes( estimates = [] for record in records: - peak = record.peak_reserved_bytes + peak = record.peak_allocated_bytes baseline = min( max(record.baseline_allocated_bytes, constant_weight_bytes), peak ) @@ -360,7 +364,7 @@ def estimate_workload_phase_peaks( { phase_name for record in successful - for phase_name in record.phase_peak_reserved_bytes + for phase_name in record.phase_peak_allocated_bytes } ) estimated_peaks: dict[str, int] = {} @@ -369,7 +373,7 @@ def estimate_workload_phase_peaks( phase_records = [ record for record in successful - if phase_name in record.phase_peak_reserved_bytes + if phase_name in record.phase_peak_allocated_bytes ] if not phase_records: continue @@ -386,9 +390,9 @@ def estimate_workload_phase_peaks( num_frames=record.num_frames, baseline_allocated_bytes=min( record.baseline_allocated_bytes, - record.phase_peak_reserved_bytes[phase_name], + record.phase_peak_allocated_bytes[phase_name], ), - peak_reserved_bytes=record.phase_peak_reserved_bytes[phase_name], + peak_allocated_bytes=record.phase_peak_allocated_bytes[phase_name], succeeded=True, ) for record in phase_records @@ -1319,11 +1323,11 @@ def format_plan_summary( ) measured = ", ".join( f"{record.width}x{record.height}x{record.num_frames}f=" - f"{record.peak_reserved_bytes / GIB_BYTES:.1f}GiB" + f"{record.peak_allocated_bytes / GIB_BYTES:.1f}GiB" for record in records if record.succeeded ) - measured_part = f"measured=[{measured}], " if measured else "" + measured_part = f"measured_allocated=[{measured}], " if measured else "" return ( f"Auto residency: target={workload.describe()} " f"steps={workload.num_inference_steps}, " diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py index e49006a09498..7c85abf564e7 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py @@ -58,6 +58,12 @@ class ResidencyState: batch_is_warmup: bool = False +@dataclass(frozen=True, slots=True) +class WarmupPhasePeak: + active_components: tuple[str, ...] + allocated_bytes: int + + class ResidencyBatch(Protocol): is_warmup: bool @@ -132,8 +138,8 @@ def __init__( self._track_warmup_memory = False self._warmup_phase_key: str | None = None self._warmup_phase_components: tuple[str, ...] = () - self._warmup_phase_peaks: dict[str, tuple[tuple[str, ...], int]] = {} - self._completed_warmup_phase_peaks: dict[str, tuple[tuple[str, ...], int]] = {} + self._warmup_phase_peaks: dict[str, WarmupPhasePeak] = {} + self._completed_warmup_phase_peaks: dict[str, WarmupPhasePeak] = {} def refresh_pipeline(self, pipeline: ComponentResidencyPipeline) -> None: custom_strategies = dict(pipeline.component_residency_strategies) @@ -569,9 +575,9 @@ def finish_request(self) -> None: ) if self._track_warmup_memory: self._record_warmup_phase_peak() - self._warmup_phase_peaks["idle"] = ( - self._warmup_active_components(), - int(torch.get_device_module().memory_reserved()), + self._warmup_phase_peaks["idle"] = WarmupPhasePeak( + active_components=self._warmup_active_components(), + allocated_bytes=int(torch.get_device_module().memory_allocated()), ) self._completed_warmup_phase_peaks = dict(self._warmup_phase_peaks) self._track_warmup_memory = False @@ -630,17 +636,26 @@ def _warmup_active_components( def _record_warmup_phase_peak(self) -> None: if not self._track_warmup_memory or self._warmup_phase_key is None: return - peak = int(torch.get_device_module().max_memory_reserved()) + peak = WarmupPhasePeak( + active_components=self._warmup_phase_components, + allocated_bytes=int(torch.get_device_module().max_memory_allocated()), + ) previous = self._warmup_phase_peaks.get(self._warmup_phase_key) - if previous is None or peak > previous[1]: - self._warmup_phase_peaks[self._warmup_phase_key] = ( - self._warmup_phase_components, - peak, + if previous is None: + self._warmup_phase_peaks[self._warmup_phase_key] = peak + else: + self._warmup_phase_peaks[self._warmup_phase_key] = WarmupPhasePeak( + active_components=tuple( + sorted( + set(previous.active_components) & set(peak.active_components) + ) + ), + allocated_bytes=max(previous.allocated_bytes, peak.allocated_bytes), ) def take_warmup_phase_peaks( self, - ) -> dict[str, tuple[tuple[str, ...], int]]: + ) -> dict[str, WarmupPhasePeak]: """Return and clear the most recently completed warmup phase peaks.""" peaks = self._completed_warmup_phase_peaks self._completed_warmup_phase_peaks = {} diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 5413c1510c24..274083b290c6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -61,7 +61,7 @@ def _record( height=height, num_frames=num_frames, baseline_allocated_bytes=baseline_gib * GIB_BYTES, - peak_reserved_bytes=peak_gib * GIB_BYTES, + peak_allocated_bytes=peak_gib * GIB_BYTES, succeeded=succeeded, ) @@ -72,7 +72,7 @@ def test_same_shape_uses_measured_peak(self): estimate = estimate_default_workload_peak_bytes( records=[record], target_units=record.workload_units() ) - assert estimate == record.peak_reserved_bytes + assert estimate == record.peak_allocated_bytes def test_unknown_target_disables_estimation(self): # An unknown target would silently equate the capped warmup peak with @@ -93,14 +93,14 @@ def test_single_capped_record_scales_only_the_activation_part(self): estimate = estimate_default_workload_peak_bytes( records=[record], target_units=target_units ) - activation = record.peak_reserved_bytes - record.baseline_allocated_bytes + activation = record.peak_allocated_bytes - record.baseline_allocated_bytes expected = record.baseline_allocated_bytes + int( activation * ratio * ACTIVATION_EXTRAPOLATION_MARGIN ) assert estimate == expected # Scaling the whole peak would inflate the estimate by the resident # weights times the cap ratio and promotion would never trigger. - naive = int(record.peak_reserved_bytes * ratio) + naive = int(record.peak_allocated_bytes * ratio) assert estimate < naive def test_two_point_fit_separates_constant_from_linear(self): @@ -116,10 +116,10 @@ def test_two_point_fit_separates_constant_from_linear(self): estimate = estimate_default_workload_peak_bytes( records=[small, large], target_units=target_units ) - slope = (large.peak_reserved_bytes - small.peak_reserved_bytes) / ( + slope = (large.peak_allocated_bytes - small.peak_allocated_bytes) / ( large.workload_units() - small.workload_units() ) - constant = large.peak_reserved_bytes - slope * large.workload_units() + constant = large.peak_allocated_bytes - slope * large.workload_units() expected = int( constant + slope * target_units * ACTIVATION_EXTRAPOLATION_MARGIN ) @@ -150,7 +150,7 @@ def test_weight_floor_prevents_scaling_constant_component_memory(self): height=480, num_frames=9, baseline_allocated_bytes=GIB_BYTES, - peak_reserved_bytes=int(29.5 * GIB_BYTES), + peak_allocated_bytes=int(29.5 * GIB_BYTES), succeeded=True, ) large = WarmupMemoryRecord( @@ -158,7 +158,7 @@ def test_weight_floor_prevents_scaling_constant_component_memory(self): height=480, num_frames=17, baseline_allocated_bytes=GIB_BYTES, - peak_reserved_bytes=int(29.4 * GIB_BYTES), + peak_allocated_bytes=int(29.4 * GIB_BYTES), succeeded=True, ) target_units = 1024 * 1024 * 81 @@ -181,7 +181,7 @@ def test_covering_measurement_bounds_the_target(self): estimate = estimate_default_workload_peak_bytes( records=[capped, full], target_units=1280 * 720 * 81 ) - assert estimate == full.peak_reserved_bytes + assert estimate == full.peak_allocated_bytes def test_multiple_records_take_the_max(self): low = _record(peak_gib=12) @@ -189,7 +189,7 @@ def test_multiple_records_take_the_max(self): estimate = estimate_default_workload_peak_bytes( records=[low, high], target_units=low.workload_units() ) - assert estimate == high.peak_reserved_bytes + assert estimate == high.peak_allocated_bytes def test_failure_at_the_target_size_disables_estimation(self): good = _record(num_frames=9, peak_gib=8) @@ -217,7 +217,7 @@ def test_failure_above_the_target_size_is_dropped(self): estimate = estimate_default_workload_peak_bytes( records=[good, failed], target_units=good.workload_units() ) - assert estimate == good.peak_reserved_bytes + assert estimate == good.peak_allocated_bytes def test_unknown_target_disables_estimation(self): assert ( @@ -238,9 +238,9 @@ def test_phase_estimation_preserves_active_component_membership(self): height=small.height, num_frames=small.num_frames, baseline_allocated_bytes=small.baseline_allocated_bytes, - peak_reserved_bytes=small.peak_reserved_bytes, + peak_allocated_bytes=small.peak_allocated_bytes, succeeded=small.succeeded, - phase_peak_reserved_bytes={"denoise": 30 * GIB_BYTES}, + phase_peak_allocated_bytes={"denoise": 30 * GIB_BYTES}, phase_active_components={"denoise": ("transformer",)}, ) large = WarmupMemoryRecord( @@ -248,9 +248,9 @@ def test_phase_estimation_preserves_active_component_membership(self): height=large.height, num_frames=large.num_frames, baseline_allocated_bytes=large.baseline_allocated_bytes, - peak_reserved_bytes=large.peak_reserved_bytes, + peak_allocated_bytes=large.peak_allocated_bytes, succeeded=large.succeeded, - phase_peak_reserved_bytes={"denoise": 32 * GIB_BYTES}, + phase_peak_allocated_bytes={"denoise": 32 * GIB_BYTES}, phase_active_components={"denoise": ("transformer",)}, ) @@ -263,6 +263,26 @@ def test_phase_estimation_preserves_active_component_membership(self): assert peaks["denoise"] >= 32 * GIB_BYTES assert active == {"denoise": ("transformer",)} + def test_phase_estimation_uses_allocated_peak(self): + record = WarmupMemoryRecord( + width=1024, + height=1024, + num_frames=1, + baseline_allocated_bytes=5 * GIB_BYTES, + peak_allocated_bytes=12 * GIB_BYTES, + succeeded=True, + phase_peak_allocated_bytes={"denoise": 11 * GIB_BYTES}, + phase_active_components={"denoise": ("transformer",)}, + ) + + peaks, _ = estimate_workload_phase_peaks( + records=[record], + target_units=record.workload_units(), + component_weight_bytes={"transformer": 10 * GIB_BYTES}, + ) + + assert peaks["denoise"] == 11 * GIB_BYTES + def _candidate( name: str, diff --git a/python/sglang/multimodal_gen/test/unit/test_component_residency.py b/python/sglang/multimodal_gen/test/unit/test_component_residency.py index 25cc782e9556..a9c911a1d828 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_residency.py @@ -7,6 +7,7 @@ ComponentResidencyManager, ComponentUse, ResidencyState, + WarmupPhasePeak, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import ( ComponentOffloadStrategy, @@ -163,8 +164,8 @@ def test_warmup_records_use_and_transition_peaks(monkeypatch): device_module = SimpleNamespace( is_available=lambda: True, reset_peak_memory_stats=Mock(), - max_memory_reserved=lambda: 11, - memory_reserved=lambda: 3, + max_memory_allocated=lambda: 7, + memory_allocated=lambda: 2, ) monkeypatch.setattr(torch, "get_device_module", lambda: device_module) monkeypatch.setattr(current_platform, "is_cuda", lambda: True) @@ -189,30 +190,29 @@ def test_warmup_records_use_and_transition_peaks(monkeypatch): manager.finish_request() peaks = manager.take_warmup_phase_peaks() - assert peaks["request:before-stage"] == ((), 11) - assert peaks["0:denoise:setup"] == ((), 11) - assert peaks["0:denoise:transition:idle->transformer"] == ( - ("transformer",), - 11, - ) - assert peaks["0:denoise:use:transformer"] == (("transformer",), 11) - assert peaks["0:denoise:transition:transformer->idle"] == ( - ("transformer",), - 11, - ) - assert peaks["0:denoise:between"] == ((), 11) + inactive_peak = WarmupPhasePeak((), 7) + transformer_peak = WarmupPhasePeak(("transformer",), 7) + assert peaks["request:before-stage"] == inactive_peak + assert peaks["0:denoise:setup"] == inactive_peak + assert peaks["0:denoise:transition:idle->transformer"] == transformer_peak + assert peaks["0:denoise:use:transformer"] == transformer_peak + assert peaks["0:denoise:transition:transformer->idle"] == transformer_peak + assert peaks["0:denoise:between"] == inactive_peak # A non-preferred component is being released during cleanup, so it is no # longer part of the placement that follows this transition. - assert peaks["request:cleanup:transformer"] == ((), 11) - assert peaks["idle"] == ((), 3) + assert peaks["request:cleanup:transformer"] == inactive_peak + assert peaks["idle"] == WarmupPhasePeak( + active_components=(), + allocated_bytes=2, + ) def test_warmup_records_same_component_dtype_prepare_as_transition(monkeypatch): device_module = SimpleNamespace( is_available=lambda: True, reset_peak_memory_stats=Mock(), - max_memory_reserved=lambda: 11, - memory_reserved=lambda: 3, + max_memory_allocated=lambda: 7, + memory_allocated=lambda: 2, ) monkeypatch.setattr(torch, "get_device_module", lambda: device_module) monkeypatch.setattr(current_platform, "is_cuda", lambda: True) @@ -240,7 +240,7 @@ def test_warmup_records_same_component_dtype_prepare_as_transition(monkeypatch): manager._record_warmup_phase_peak() assert manager._warmup_phase_peaks[ "0:stage:transition:transformer->transformer" - ] == (("transformer",), 11) + ] == WarmupPhasePeak(("transformer",), 7) assert strategy.prepare_for_use.call_count == 2 @@ -248,8 +248,8 @@ def test_warmup_splits_sequential_component_transition(monkeypatch): device_module = SimpleNamespace( is_available=lambda: True, reset_peak_memory_stats=Mock(), - max_memory_reserved=lambda: 11, - memory_reserved=lambda: 3, + max_memory_allocated=lambda: 7, + memory_allocated=lambda: 2, ) monkeypatch.setattr(torch, "get_device_module", lambda: device_module) monkeypatch.setattr(current_platform, "is_cuda", lambda: True) @@ -278,12 +278,10 @@ def test_warmup_splits_sequential_component_transition(monkeypatch): manager._record_warmup_phase_peak() assert manager._warmup_phase_peaks["0:stage:transition:text_encoder->idle"] == ( - ("text_encoder",), - 11, + WarmupPhasePeak(("text_encoder",), 7) ) assert manager._warmup_phase_peaks["0:stage:transition:idle->transformer"] == ( - ("transformer",), - 11, + WarmupPhasePeak(("transformer",), 7) ) diff --git a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py index bbc11d6e2c57..315a147dfd00 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -11,6 +11,9 @@ DefaultWorkload, WarmupMemoryRecord, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( + WarmupPhasePeak, +) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.perf_logger import ( @@ -99,9 +102,10 @@ def test_server_warmup_preserves_peak_after_managed_stage_timeline(): worker._auto_residency_warmup_records = [] residency_manager = Mock() residency_manager.take_warmup_phase_peaks.return_value = { - "0:denoise:use:transformer": (("transformer",), 10) + "0:denoise:use:transformer": WarmupPhasePeak(("transformer",), 8) } device_module = Mock() + device_module.max_memory_allocated.return_value = 9 device_module.max_memory_reserved.return_value = 12 req = SimpleNamespace(width=64, height=64, num_frames=1) @@ -120,11 +124,41 @@ def test_server_warmup_preserves_peak_after_managed_stage_timeline(): ) record = worker._auto_residency_warmup_records[0] - assert record.peak_reserved_bytes == 12 - assert record.phase_peak_reserved_bytes["request:untracked"] == 12 + assert record.peak_allocated_bytes == 9 + assert record.phase_peak_allocated_bytes["request:untracked"] == 9 assert record.phase_active_components["request:untracked"] == () +def test_server_warmup_does_not_treat_allocator_cache_as_untracked_live_memory(): + worker = GPUWorker.__new__(GPUWorker) + worker._auto_residency_warmup_records = [] + residency_manager = Mock() + residency_manager.take_warmup_phase_peaks.return_value = { + "0:denoise:use:transformer": WarmupPhasePeak(("transformer",), 8) + } + device_module = Mock() + device_module.max_memory_allocated.return_value = 8 + device_module.max_memory_reserved.return_value = 12 + req = SimpleNamespace(width=64, height=64, num_frames=1) + + with ( + patch.object(torch, "get_device_module", return_value=device_module), + patch.object( + gpu_worker_module, + "peek_global_component_residency_manager", + return_value=residency_manager, + ), + ): + worker._record_server_warmup_memory( + req=req, + baseline_allocated_bytes=3, + succeeded=True, + ) + + record = worker._auto_residency_warmup_records[0] + assert "request:untracked" not in record.phase_peak_allocated_bytes + + def test_loaded_prequantized_checkpoint_can_use_auto_residency(): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 @@ -161,7 +195,7 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): height=64, num_frames=1, baseline_allocated_bytes=1, - peak_reserved_bytes=2, + peak_allocated_bytes=2, succeeded=True, ) ], From ebff2f42ffb0b2e0522f072922451eec840f73ec Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Aug 2026 21:14:16 +0800 Subject: [PATCH 26/90] [diffusion] converge auto residency after warmup --- .../runtime/managers/gpu_worker.py | 119 ++++++++--- .../memory_managers/auto_residency.py | 37 ++-- .../memory_managers/component_manager.py | 19 ++ .../multimodal_gen/runtime/server_warmup.py | 107 ++++++---- .../test/unit/test_auto_residency.py | 15 +- .../test/unit/test_auto_residency_warmup.py | 191 ++++++++++++++++++ .../test/unit/test_component_residency.py | 49 +++++ .../test/unit/test_performance_metrics.py | 52 ++++- 8 files changed, 490 insertions(+), 99 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 9d592c79b305..a26dd6860a24 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -198,6 +198,7 @@ def __init__( # the auto-residency promotion decision before the server turns ready self._auto_residency_warmup_records: list[WarmupMemoryRecord] = [] self._auto_residency_applied: list[AppliedPromotion] = [] + self._auto_residency_round_sizes: list[int] = [] # default workload resolved once for the per-request residency hint self._cached_default_workload: DefaultWorkload | None = None self._cached_default_workload_failed = False @@ -700,11 +701,13 @@ def _record_server_warmup_memory( ) -> None: phase_allocated_peaks: dict[str, int] = {} phase_components: dict[str, tuple[str, ...]] = {} + untracked_active_components: tuple[str, ...] = () residency_manager = peek_global_component_residency_manager() if residency_manager is not None: for phase_name, peak in residency_manager.take_warmup_phase_peaks().items(): phase_allocated_peaks[phase_name] = peak.allocated_bytes phase_components[phase_name] = peak.active_components + untracked_active_components = residency_manager.current_device_components() request_allocated_peak = max( int(torch.get_device_module().max_memory_allocated()), max(phase_allocated_peaks.values(), default=0), @@ -715,7 +718,7 @@ def _record_server_warmup_memory( # A reserved-only increase is reclaimable allocator cache, not a # second live placement, and is covered by the explicit reserve. phase_allocated_peaks["request:untracked"] = request_allocated_peak - phase_components["request:untracked"] = () + phase_components["request:untracked"] = untracked_active_components self._auto_residency_warmup_records.append( WarmupMemoryRecord( width=int(req.width or 0), @@ -1180,25 +1183,23 @@ def get_can_stay_resident_components( return can_stay_resident def apply_auto_residency(self) -> OutputBatch: - """Apply the warmup-calibrated residency promotion on this replica. + """Apply one warmup-calibrated residency promotion round. Every rank executes this handler at the same queue position; the plan is computed from all-gathered rank reports so each rank reaches the - same decision. A failure on any rank rolls every rank back to the - original strategy. + same decision. Post-promotion warmup can expose more safe headroom, so + the server may call this again with measurements from the new layout. + A failure on any rank rolls every rank back to the previously + calibrated placement. Everything before the first all-gather is fenced into a skip report: an uncaught raise there would leave the peer ranks parked in the collective until the group timeout. """ - if self._auto_residency_applied: - # one-shot per process: records measured after a promotion - # describe the promoted layout, not the startup strategy - return OutputBatch(output={"status": PROMOTION_STATUS_SKIPPED}) - records = list(self._auto_residency_warmup_records) - # consumed exactly once; the post-promotion re-warm appends fresh - # records that must not feed a later decision + # Each round must describe one placement. Intersecting phase ownership + # across old and newly promoted layouts would double-count weights that + # are already resident in the new layout. self._auto_residency_warmup_records.clear() try: @@ -1236,12 +1237,20 @@ def apply_auto_residency(self) -> OutputBatch: apply_error: str | None = None local_rollback_failed = False + # Keep round history aligned across ranks. A zero entry means this + # rank's transactional apply failed before committing the round, so a + # peer-triggered rollback must leave earlier calibrated rounds alone. + self._auto_residency_round_sizes.append(0) try: - self._auto_residency_applied = apply_promotions( + newly_applied = apply_promotions( plan=plan, modules=self.pipeline.modules, server_args=self.server_args, ) + if not newly_applied: + raise RuntimeError("promotion plan applied no residency changes") + self._auto_residency_applied.extend(newly_applied) + self._auto_residency_round_sizes[-1] = len(newly_applied) except AutoResidencyRollbackError as e: logger.error( "Auto residency promotion failed on rank %d and the rank " @@ -1250,7 +1259,6 @@ def apply_auto_residency(self) -> OutputBatch: e, exc_info=True, ) - self._auto_residency_applied = [] apply_error = describe_error(e) local_rollback_failed = True except Exception as e: # this rank already rolled itself back @@ -1260,7 +1268,6 @@ def apply_auto_residency(self) -> OutputBatch: e, exc_info=True, ) - self._auto_residency_applied = [] apply_error = describe_error(e) gathered = self._auto_residency_all_gather((apply_error, local_rollback_failed)) @@ -1268,7 +1275,9 @@ def apply_auto_residency(self) -> OutputBatch: any_rollback_failed = any(failed for _, failed in gathered) if rank_errors: return self._rollback_everywhere( - cause=rank_errors[0], already_failed=any_rollback_failed + cause=rank_errors[0], + already_failed=any_rollback_failed, + latest_round_only=True, ) self._invalidate_component_strategies( @@ -1282,11 +1291,17 @@ def apply_auto_residency(self) -> OutputBatch: ) def rollback_auto_residency(self) -> OutputBatch: - """Restore the pre-promotion residency (post-promotion re-warm failed).""" - return self._rollback_everywhere(cause=None, already_failed=False) + """Undo the latest promotion round after its calibration failed.""" + return self._rollback_everywhere( + cause=None, already_failed=False, latest_round_only=True + ) def _rollback_everywhere( - self, *, cause: str | None, already_failed: bool + self, + *, + cause: str | None, + already_failed: bool, + latest_round_only: bool = False, ) -> OutputBatch: """Roll this rank back and gather the replica-wide outcome. @@ -1296,7 +1311,9 @@ def _rollback_everywhere( rollback, which is fatal regardless of what the remaining ranks do. Collective-symmetric: every rank gathers exactly once. """ - rollback_error = self._rollback_applied_promotions() + rollback_error = self._rollback_applied_promotions( + latest_round_only=latest_round_only + ) gathered = self._auto_residency_all_gather(rollback_error) rank_errors = [error for error in gathered if error is not None] if already_failed and not rank_errors: @@ -1312,19 +1329,29 @@ def _rollback_everywhere( output={"status": PROMOTION_STATUS_ROLLBACK_FAILED}, ) if cause is not None: + restored = ( + "previously calibrated placement" + if self._auto_residency_applied + else "original strategy" + ) return OutputBatch( error=( - "auto residency promotion failed; original strategy " - f"restored: {cause}" + f"auto residency promotion failed; {restored} restored: {cause}" ), output={"status": PROMOTION_STATUS_ROLLED_BACK}, ) if self.is_output_rank: - logger.info( - "Auto residency: rolled back; the residency configured at " - "startup is in effect again (no equivalent server-arg " - "changes remain)." - ) + if self._auto_residency_applied: + logger.info( + "Auto residency: rolled back the latest promotion round; " + "the previously calibrated placement remains active." + ) + else: + logger.info( + "Auto residency: rolled back; the residency configured at " + "startup is in effect again (no equivalent server-arg " + "changes remain)." + ) return OutputBatch(output={"status": PROMOTION_STATUS_ROLLED_BACK}) def _build_auto_residency_report( @@ -1384,13 +1411,19 @@ def _build_auto_residency_report( budget_bytes = int(free_bytes) + int( torch.get_device_module().memory_reserved() ) + test_cap_gib = envs.SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB + if test_cap_gib is not None: + budget_bytes = min(budget_bytes, int(test_cap_gib * GIB_BYTES)) candidates = collect_promotion_candidates( modules=self.pipeline.modules, residency_mode_of=self.server_args.residency_mode, explicit_residency_mode_of=self.server_args.explicit_residency_mode, custom_strategy_names=self.pipeline.component_residency_strategies, num_inference_steps=workload.num_inference_steps, - allow_host_pin_reallocation=self.server_args.num_gpus == 1, + allow_host_pin_reallocation=True, + ) + local_worker_count = max( + 1, self.server_args.num_gpus // self.server_args.nnodes ) return RankResidencyReport( rank=self.rank, @@ -1405,7 +1438,8 @@ def _build_auto_residency_report( ), host_transition_headroom_bytes=max( 0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES - ), + ) + // local_worker_count, candidates=candidates, skip_reason=skip_reason, ) @@ -1420,10 +1454,24 @@ def _auto_residency_all_gather(self, obj: Any) -> list[Any]: ) return gathered - def _rollback_applied_promotions(self) -> str | None: - if not self._auto_residency_applied: - return None - applied = self._auto_residency_applied + def _rollback_applied_promotions( + self, *, latest_round_only: bool = False + ) -> str | None: + if latest_round_only: + if not self._auto_residency_round_sizes: + return None + round_size = self._auto_residency_round_sizes[-1] + if round_size == 0: + self._auto_residency_round_sizes.pop() + return None + if round_size > len(self._auto_residency_applied): + return "auto residency round history is inconsistent" + applied = self._auto_residency_applied[-round_size:] + else: + if not self._auto_residency_applied: + self._auto_residency_round_sizes = [] + return None + applied = self._auto_residency_applied try: rollback_promotions( applied=applied, @@ -1441,7 +1489,12 @@ def _rollback_applied_promotions(self) -> str | None: self._invalidate_component_strategies( [promotion.component_name for promotion in applied] ) - self._auto_residency_applied = [] + if latest_round_only and self._auto_residency_round_sizes: + del self._auto_residency_applied[-len(applied) :] + self._auto_residency_round_sizes.pop() + else: + self._auto_residency_applied = [] + self._auto_residency_round_sizes = [] return None @staticmethod diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 0444638cfc48..2b35ec804d19 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -755,10 +755,9 @@ def collect_promotion_candidates( ) full_resident_layers = tuple(manager.num_layers for manager in managers) - # Fully resident layers never read their host stores. On one worker we - # can release those pins and reuse the allowance. Multi-worker pin - # migration needs a node-coordinated host-RAM scratch budget because - # repacking briefly holds both the old and new layer buffers. + # Fully resident layers never read their host stores, so their pins can + # be released and reused. Every worker plans against its non-overlapping + # HostPin and transition-headroom share before repacking concurrently. permanent_pin_targets = [current_pinned_layers] if allow_host_pin_reallocation: permanent_pin_targets.append(tuple(() for _ in managers)) @@ -989,18 +988,17 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP for candidate in candidates ) if has_hostpin_options: - if len(reports) != 1: - return _skip_plan("dynamic HostPin placement requires a single-worker node") - report = reports[0] - resource_budgets[f"hostpin:node{report.node_rank}"] = ( - report.host_pin_capacity_bytes - report.pinned_host_bytes - ) - resource_budgets[f"hostram:node{report.node_rank}:unpin"] = ( - report.host_transition_headroom_bytes - ) - resource_budgets[f"hostram:node{report.node_rank}:pin"] = ( - report.host_transition_headroom_bytes - ) + for report in reports: + prefix = f"node{report.node_rank}:rank{report.rank}" + resource_budgets[f"hostpin:{prefix}"] = ( + report.host_pin_capacity_bytes - report.pinned_host_bytes + ) + resource_budgets[f"hostram:{prefix}:unpin"] = ( + report.host_transition_headroom_bytes + ) + resource_budgets[f"hostram:{prefix}:pin"] = ( + report.host_transition_headroom_bytes + ) candidate_by_key = {candidate.option_key(): candidate for candidate in candidates} report_candidates = [ @@ -1022,15 +1020,16 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP ) resource_deltas[resource_name] = phase_cost if has_hostpin_options: - host_resource = f"hostpin:node{report.node_rank}" + prefix = f"node{report.node_rank}:rank{report.rank}" + host_resource = f"hostpin:{prefix}" resource_deltas[host_resource] = ( resource_deltas.get(host_resource, 0) + rank_candidate.pinned_host_delta_bytes ) - resource_deltas[f"hostram:node{report.node_rank}:unpin"] = ( + resource_deltas[f"hostram:{prefix}:unpin"] = ( rank_candidate.host_unpin_scratch_bytes ) - resource_deltas[f"hostram:node{report.node_rank}:pin"] = ( + resource_deltas[f"hostram:{prefix}:pin"] = ( rank_candidate.host_pin_scratch_bytes ) options.append( diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py index 7c85abf564e7..6a2a1f25ba7c 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py @@ -509,8 +509,11 @@ def _prefetch_use(self, use: ComponentUse) -> None: self._uses_seen[use.component_name] = use self._modules_seen[use.component_name] = module + self._begin_warmup_prefetch(use) if strategy.prefetch_for_use(module, use, self.state): self._prefetched_use_keys.add(self._use_key(use)) + else: + self._begin_warmup_between_uses() def _finish_use( self, @@ -620,6 +623,13 @@ def _begin_warmup_between_uses(self) -> None: components=self._warmup_active_components(), ) + def _begin_warmup_prefetch(self, use: ComponentUse) -> None: + phase = use.phase or use.component_name + self._begin_warmup_phase( + key=f"{self.state.stage_index}:{self.state.stage_name}:prefetch:{phase}", + components=self._warmup_active_components((use,)), + ) + def _warmup_active_components( self, active_uses: Sequence[ComponentUse] = () ) -> tuple[str, ...]: @@ -661,6 +671,15 @@ def take_warmup_phase_peaks( self._completed_warmup_phase_peaks = {} return peaks + def current_device_components(self) -> tuple[str, ...]: + """Components whose complete module is currently on the device. + + Dormant layerwise-managed modules are excluded because only their + resident window is present. Active layerwise uses are attributed by + the managed phase timeline instead. + """ + return self._warmup_active_components() + def stage_name(self, stage: ComponentResidencyStage) -> str: return self._stage_names_by_id.get(id(stage), stage.__class__.__name__) diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index f26206045ad5..e64556dfc322 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -128,10 +128,9 @@ async def maybe_apply_auto_residency( """Promote implicitly offloaded components after warmup, then re-warm. Runs between the synthetic warmup and the ready signal, so the residency - is frozen before ``/health`` turns 200. If a promotion (or the - post-promotion re-warmup) fails, the workers roll back and startup - continues on the original strategy; only a failed rollback raises, which - aborts startup. + is frozen before ``/health`` turns 200. If a promotion (or its calibration) + fails, the workers roll back that round and retain the last calibrated + placement; only a failed rollback raises and aborts startup. """ from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( AutoResidencyReq, @@ -157,40 +156,12 @@ async def maybe_apply_auto_residency( # Same fail-open contract as the warmup itself: implicit warmups must # never abort startup, explicit --warmup-resolutions ones must succeed. fail_open = server_args.warmup_resolutions is None - try: - response = await forward(AutoResidencyReq(action="apply")) - except Exception as e: - if not fail_open: - raise - logger.warning( - "Auto residency apply request failed; continuing on the original " - "strategy: %s", - e, - ) - return - status = _auto_residency_status(response) - if status == PROMOTION_STATUS_ROLLBACK_FAILED: - raise RuntimeError(f"auto residency rollback failed: {response.error}") - if response.error is not None: - logger.warning( - "Auto residency promotion not applied; continuing on the original " - "strategy: %s", - response.error, - ) - return - if status != PROMOTION_STATUS_PROMOTED: - return + promoted_any = False - # Re-run the synthetic warmup under the final residency: it physically - # moves promoted components, rebuilds compile caches invalidated by the - # placement change, and proves the promoted layout fits before ready. - try: - await run_async_client_warmup( - server_args, forward, fail_open=False, rewarm=True - ) - except Exception as e: + async def rollback_and_rewarm(error: Exception) -> None: logger.warning( - "Post-promotion re-warmup failed (%s); rolling back auto residency", e + "Post-promotion calibration failed (%s); rolling back auto residency", + error, ) rollback = await forward(AutoResidencyReq(action="rollback")) if ( @@ -199,18 +170,76 @@ async def maybe_apply_auto_residency( ): raise RuntimeError( f"auto residency rollback failed: {rollback.error}" - ) from e - # restore warm caches for the original strategy before turning ready, - # keeping the caller's fail-open contract + ) from error + # Restore warm caches for the previous calibrated placement before + # turning ready, keeping the caller's fail-open contract. await run_async_client_warmup( server_args, forward, fail_open=fail_open, rewarm=True ) + for _ in range(MAX_AUTO_RESIDENCY_ROUNDS): + try: + response = await forward(AutoResidencyReq(action="apply")) + except Exception as e: + if promoted_any: + # The worker-side apply is transactional, but an RPC failure + # does not tell us whether its response was lost after commit. + # Do not mistake the last validated round for an unvalidated + # one and roll it back; abort startup rather than serve a + # placement whose calibration state is unknown. + raise RuntimeError( + "auto residency apply failed after a calibrated promotion" + ) from e + if not fail_open: + raise + logger.warning( + "Auto residency apply request failed; continuing on the original " + "strategy: %s", + e, + ) + return + status = _auto_residency_status(response) + if status == PROMOTION_STATUS_ROLLBACK_FAILED: + raise RuntimeError(f"auto residency rollback failed: {response.error}") + if response.error is not None: + retained = ( + "the last calibrated placement" + if promoted_any + else "the original strategy" + ) + logger.warning( + "Auto residency promotion not applied; continuing on %s: %s", + retained, + response.error, + ) + return + if status != PROMOTION_STATUS_PROMOTED: + return + + promoted_any = True + # This pass physically realizes the selected placement and measures + # phases that overlap under it. The next round can then use that new + # evidence instead of permanently charging components twice. + try: + await run_async_client_warmup( + server_args, forward, fail_open=False, rewarm=True + ) + except Exception as e: + await rollback_and_rewarm(e) + return + + logger.warning( + "Auto residency stopped after %d promotion rounds; keeping the last " + "successfully calibrated placement.", + MAX_AUTO_RESIDENCY_ROUNDS, + ) + # Enough to clear a probe that overshot the card, few enough that a failure # which is not about probe size gives up quickly instead of walking the # workload down to nothing. MAX_WARMUP_DEGRADE_ATTEMPTS = 3 +MAX_AUTO_RESIDENCY_ROUNDS = 8 def _is_out_of_memory(error: Any) -> bool: diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 274083b290c6..3e8f7b7da3c8 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -674,9 +674,9 @@ def test_node_hostpin_is_optimized_with_the_same_selection_vector(self): "hot_dit", "cold_encoder", ] - assert plan.resource_budget_bytes["hostpin:node0"] == 0 + assert plan.resource_budget_bytes["hostpin:node0:rank0"] == 0 - def test_dynamic_hostpin_is_not_planned_across_workers(self): + def test_dynamic_hostpin_is_constrained_for_every_worker(self): pin_more = PromotionCandidate( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, @@ -701,10 +701,11 @@ def test_dynamic_hostpin_is_not_planned_across_workers(self): ] ) - assert plan.promotions == [] - assert plan.skip_reason == ( - "dynamic HostPin placement requires a single-worker node" - ) + assert [candidate.component_name for candidate in plan.promotions] == [ + "transformer" + ] + assert plan.resource_budget_bytes["hostpin:node0:rank0"] == 20 * GIB_BYTES + assert plan.resource_budget_bytes["hostpin:node0:rank1"] == 20 * GIB_BYTES def test_hostpin_repack_must_fit_transition_headroom(self): pin_more = PromotionCandidate( @@ -728,7 +729,7 @@ def test_hostpin_repack_must_fit_transition_headroom(self): ) assert plan.promotions == [] - assert plan.resource_budget_bytes["hostram:node0:pin"] == 5 * GIB_BYTES + assert plan.resource_budget_bytes["hostram:node0:rank0:pin"] == 5 * GIB_BYTES class _FakeLayerwiseManager: diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py new file mode 100644 index 000000000000..03090e99df9f --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -0,0 +1,191 @@ +import asyncio +import unittest +from types import SimpleNamespace +from unittest import mock + +from sglang.multimodal_gen.runtime import server_warmup +from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( + AutoResidencyReq, +) +from sglang.multimodal_gen.runtime.managers import gpu_worker as gpu_worker_module +from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker +from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( + PROMOTION_STATUS_PROMOTED, + PROMOTION_STATUS_ROLLBACK_FAILED, + PROMOTION_STATUS_ROLLED_BACK, + PROMOTION_STATUS_SKIPPED, + AppliedPromotion, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( + COMPONENT_OFFLOAD, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch + + +class TestAutoResidencyWarmup(unittest.TestCase): + def test_replans_until_the_calibrated_layout_reaches_a_fixed_point(self): + responses = iter( + [ + OutputBatch(output={"status": PROMOTION_STATUS_PROMOTED}), + OutputBatch(output={"status": PROMOTION_STATUS_PROMOTED}), + OutputBatch(output={"status": PROMOTION_STATUS_SKIPPED}), + ] + ) + actions = [] + + async def forward(req): + self.assertIsInstance(req, AutoResidencyReq) + actions.append(req.action) + return next(responses) + + rewarm = mock.AsyncMock() + server_args = SimpleNamespace( + performance_mode="auto", + warmup_resolutions=None, + ) + with ( + mock.patch.object( + server_warmup, "auto_residency_skip_reason", return_value=None + ), + mock.patch.object(server_warmup, "run_async_client_warmup", rewarm), + ): + asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) + + self.assertEqual(actions, ["apply", "apply", "apply"]) + self.assertEqual(rewarm.await_count, 2) + + def test_failed_first_calibration_rolls_back_and_rewarms(self): + actions = [] + + async def forward(req): + self.assertIsInstance(req, AutoResidencyReq) + actions.append(req.action) + status = ( + PROMOTION_STATUS_PROMOTED + if req.action == "apply" + else PROMOTION_STATUS_ROLLED_BACK + ) + return OutputBatch(output={"status": status}) + + rewarm = mock.AsyncMock(side_effect=[RuntimeError("oom"), None]) + server_args = SimpleNamespace( + performance_mode="auto", + warmup_resolutions=None, + ) + with ( + mock.patch.object( + server_warmup, "auto_residency_skip_reason", return_value=None + ), + mock.patch.object(server_warmup, "run_async_client_warmup", rewarm), + ): + asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) + + self.assertEqual(actions, ["apply", "rollback"]) + self.assertEqual(rewarm.await_count, 2) + + def test_failed_calibration_reports_rollback_failure(self): + async def forward(req): + self.assertIsInstance(req, AutoResidencyReq) + if req.action == "apply": + return OutputBatch(output={"status": PROMOTION_STATUS_PROMOTED}) + return OutputBatch( + error="rollback failed", + output={"status": PROMOTION_STATUS_ROLLBACK_FAILED}, + ) + + rewarm = mock.AsyncMock(side_effect=RuntimeError("oom")) + server_args = SimpleNamespace( + performance_mode="auto", + warmup_resolutions=None, + ) + with ( + mock.patch.object( + server_warmup, "auto_residency_skip_reason", return_value=None + ), + mock.patch.object(server_warmup, "run_async_client_warmup", rewarm), + self.assertRaisesRegex(RuntimeError, "auto residency rollback failed"), + ): + asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) + + def test_apply_rpc_failure_after_calibration_aborts_without_rollback(self): + actions = [] + responses = iter( + [ + OutputBatch(output={"status": PROMOTION_STATUS_PROMOTED}), + RuntimeError("rpc failed"), + ] + ) + + async def forward(req): + self.assertIsInstance(req, AutoResidencyReq) + actions.append(req.action) + response = next(responses) + if isinstance(response, Exception): + raise response + return response + + rewarm = mock.AsyncMock() + server_args = SimpleNamespace( + performance_mode="auto", + warmup_resolutions=None, + ) + with ( + mock.patch.object( + server_warmup, "auto_residency_skip_reason", return_value=None + ), + mock.patch.object(server_warmup, "run_async_client_warmup", rewarm), + self.assertRaisesRegex( + RuntimeError, + "auto residency apply failed after a calibrated promotion", + ), + ): + asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) + + self.assertEqual(actions, ["apply", "apply"]) + self.assertEqual(rewarm.await_count, 1) + + def test_failed_later_round_keeps_earlier_calibrated_promotions(self): + first_round = AppliedPromotion("transformer", COMPONENT_OFFLOAD) + latest_round = [ + AppliedPromotion("text_encoder", COMPONENT_OFFLOAD), + AppliedPromotion("vae", COMPONENT_OFFLOAD), + ] + worker = GPUWorker.__new__(GPUWorker) + worker.rank = 0 + worker.is_output_rank = False + worker.pipeline = SimpleNamespace(modules={}) + worker.server_args = SimpleNamespace() + worker._auto_residency_applied = [first_round, *latest_round] + worker._auto_residency_round_sizes = [1, 2] + worker._auto_residency_all_gather = lambda value: [value] + worker._invalidate_component_strategies = mock.Mock() + + with mock.patch.object(gpu_worker_module, "rollback_promotions") as rollback: + response = worker.rollback_auto_residency() + + rollback.assert_called_once_with( + applied=latest_round, + modules={}, + server_args=worker.server_args, + ) + self.assertEqual(response.output["status"], PROMOTION_STATUS_ROLLED_BACK) + self.assertEqual(worker._auto_residency_applied, [first_round]) + self.assertEqual(worker._auto_residency_round_sizes, [1]) + + def test_failed_apply_round_does_not_rollback_prior_round(self): + first_round = AppliedPromotion("transformer", COMPONENT_OFFLOAD) + worker = GPUWorker.__new__(GPUWorker) + worker.rank = 0 + worker.pipeline = SimpleNamespace(modules={}) + worker.server_args = SimpleNamespace() + worker._auto_residency_applied = [first_round] + worker._auto_residency_round_sizes = [1, 0] + worker._invalidate_component_strategies = mock.Mock() + + with mock.patch.object(gpu_worker_module, "rollback_promotions") as rollback: + error = worker._rollback_applied_promotions(latest_round_only=True) + + self.assertIsNone(error) + rollback.assert_not_called() + self.assertEqual(worker._auto_residency_applied, [first_round]) + self.assertEqual(worker._auto_residency_round_sizes, [1]) diff --git a/python/sglang/multimodal_gen/test/unit/test_component_residency.py b/python/sglang/multimodal_gen/test/unit/test_component_residency.py index a9c911a1d828..21cf29e17dc4 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_residency.py @@ -244,6 +244,55 @@ def test_warmup_records_same_component_dtype_prepare_as_transition(monkeypatch): assert strategy.prepare_for_use.call_count == 2 +def test_warmup_attributes_prefetch_peak_to_prefetched_component(monkeypatch): + device_module = SimpleNamespace( + is_available=lambda: True, + reset_peak_memory_stats=Mock(), + max_memory_allocated=lambda: 7, + memory_allocated=lambda: 2, + ) + monkeypatch.setattr(torch, "get_device_module", lambda: device_module) + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + + encoder_use = ComponentUse("encode", "text_encoder") + transformer_use = ComponentUse("denoise", "transformer", memory_intensive=True) + encode_stage = _Stage(encoder_use) + denoise_stage = _Stage(transformer_use) + modules = { + "text_encoder": torch.nn.Linear(2, 2), + "transformer": torch.nn.Linear(2, 2), + } + pipeline = SimpleNamespace( + modules=modules, + _stage_name_mapping={ + "encode": encode_stage, + "denoise": denoise_stage, + }, + component_residency_strategies={}, + ) + server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + manager = ComponentResidencyManager(pipeline, server_args) + strategy = Mock() + strategy.prefetch_for_use.return_value = True + manager.strategy_for = Mock(return_value=strategy) + manager.refresh_pipeline(pipeline) + manager.begin_request( + [encode_stage, denoise_stage], + SimpleNamespace(is_warmup=True), + server_args, + ) + + manager.before_stage(encode_stage, 0, SimpleNamespace(is_warmup=True), server_args) + manager.begin_stage() + manager.end_stage() + manager.before_stage(denoise_stage, 1, SimpleNamespace(is_warmup=True), server_args) + + assert manager._warmup_phase_peaks[ + "0:encode:prefetch:transformer" + ] == WarmupPhasePeak(("transformer",), 7) + assert manager._warmup_phase_peaks["0:encode:between"] == WarmupPhasePeak((), 7) + + def test_warmup_splits_sequential_component_transition(monkeypatch): device_module = SimpleNamespace( is_available=lambda: True, diff --git a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py index 315a147dfd00..57914657f68c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -8,6 +8,7 @@ import sglang.multimodal_gen.runtime.managers.gpu_worker as gpu_worker_module from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( + GIB_BYTES, DefaultWorkload, WarmupMemoryRecord, ) @@ -104,6 +105,7 @@ def test_server_warmup_preserves_peak_after_managed_stage_timeline(): residency_manager.take_warmup_phase_peaks.return_value = { "0:denoise:use:transformer": WarmupPhasePeak(("transformer",), 8) } + residency_manager.current_device_components.return_value = ("transformer",) device_module = Mock() device_module.max_memory_allocated.return_value = 9 device_module.max_memory_reserved.return_value = 12 @@ -126,7 +128,7 @@ def test_server_warmup_preserves_peak_after_managed_stage_timeline(): record = worker._auto_residency_warmup_records[0] assert record.peak_allocated_bytes == 9 assert record.phase_peak_allocated_bytes["request:untracked"] == 9 - assert record.phase_active_components["request:untracked"] == () + assert record.phase_active_components["request:untracked"] == ("transformer",) def test_server_warmup_does_not_treat_allocator_cache_as_untracked_live_memory(): @@ -136,6 +138,7 @@ def test_server_warmup_does_not_treat_allocator_cache_as_untracked_live_memory() residency_manager.take_warmup_phase_peaks.return_value = { "0:denoise:use:transformer": WarmupPhasePeak(("transformer",), 8) } + residency_manager.current_device_components.return_value = ("transformer",) device_module = Mock() device_module.max_memory_allocated.return_value = 8 device_module.max_memory_reserved.return_value = 12 @@ -173,6 +176,7 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): ) worker.server_args = SimpleNamespace( num_gpus=1, + nnodes=1, node_rank=0, residency_mode=lambda _name: "resident", explicit_residency_mode=lambda _name: None, @@ -204,6 +208,52 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): assert report.skip_reason is None +def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): + worker = GPUWorker.__new__(GPUWorker) + worker.rank = 0 + worker.pipeline = SimpleNamespace( + modules={"transformer": torch.nn.Linear(1, 1)}, + component_residency_strategies={}, + ) + worker.server_args = SimpleNamespace( + num_gpus=1, + nnodes=1, + node_rank=0, + residency_mode=lambda _name: "resident", + explicit_residency_mode=lambda _name: None, + ) + device_module = Mock() + device_module.mem_get_info.return_value = (100 * GIB_BYTES, 140 * GIB_BYTES) + device_module.memory_reserved.return_value = 2 * GIB_BYTES + monkeypatch.setattr( + gpu_worker_module.envs, + "SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB", + 80, + ) + + with patch.object(torch, "get_device_module", return_value=device_module): + report = worker._build_auto_residency_report( + workload=DefaultWorkload( + width=64, + height=64, + num_frames=1, + num_inference_steps=1, + ), + records=[ + WarmupMemoryRecord( + width=64, + height=64, + num_frames=1, + baseline_allocated_bytes=1, + peak_allocated_bytes=2, + succeeded=True, + ) + ], + ) + + assert report.budget_bytes == 80 * GIB_BYTES + + def test_baseline_config_loads_per_scenario_peak_vram(tmp_path): path = tmp_path / "baseline.json" path.write_text( From 9c8626ee95ffc54978cddf0bc18ab2ead9fba219 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Aug 2026 21:14:29 +0800 Subject: [PATCH 27/90] [diffusion] keep LTX LoRA adapters versioned --- .../model_specific_stages/ltx_2/upsampling.py | 5 +++- .../test/unit/test_lora_inference_mode.py | 28 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py index 9afa06359bee..75d1a26a0b1c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py @@ -50,7 +50,10 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: if self.pipeline.should_skip_ltx2_lora_switch_stage(): batch.extra["ltx2_phase"] = self.phase return batch - self.pipeline.switch_lora_phase(self.phase, batch=batch) + # The installed adapter survives this stage. Give its tensors version + # counters for later offload stages, which run outside inference mode. + with torch.inference_mode(False), torch.no_grad(): + self.pipeline.switch_lora_phase(self.phase, batch=batch) batch.extra["ltx2_phase"] = self.phase return batch diff --git a/python/sglang/multimodal_gen/test/unit/test_lora_inference_mode.py b/python/sglang/multimodal_gen/test/unit/test_lora_inference_mode.py index 4a4bd83bacb8..7ffea359aabc 100644 --- a/python/sglang/multimodal_gen/test/unit/test_lora_inference_mode.py +++ b/python/sglang/multimodal_gen/test/unit/test_lora_inference_mode.py @@ -1,3 +1,5 @@ +from types import SimpleNamespace + import torch from torch import nn @@ -5,6 +7,9 @@ LinearWithLoRA, _compute_lora_delta, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.upsampling import ( + LTX2LoRASwitchStage, +) def test_stacked_lora_delta_preserves_projection_order(): @@ -70,3 +75,26 @@ def test_dynamic_lora_reuses_inference_weights_without_autograd_tracking(): with torch.no_grad(): sharded_view = layer.lora_B[:2, :] assert sharded_view.shape == (2, 2) + + +def test_ltx2_lora_switch_creates_versioned_adapter_tensors(): + class _Pipeline: + adapter = None + + @staticmethod + def should_skip_ltx2_lora_switch_stage(): + return False + + def switch_lora_phase(self, phase, *, batch): + self.adapter = torch.ones(1) + + pipeline = _Pipeline() + stage = LTX2LoRASwitchStage(pipeline, "stage2") + batch = SimpleNamespace(extra={}) + + with torch.inference_mode(): + stage.forward(batch, SimpleNamespace()) + + assert pipeline.adapter is not None + assert not pipeline.adapter.is_inference() + assert batch.extra["ltx2_phase"] == "stage2" From 9c4875aaa8a07f11f4d1deef7afd0e9fdbcbd1c2 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 01:08:30 +0800 Subject: [PATCH 28/90] [diffusion] make auto residency benefit-aware --- .../runtime/managers/gpu_worker.py | 41 +++- .../memory_managers/auto_residency.py | 175 +++++++++++++++++- .../memory_managers/component_manager.py | 19 ++ .../memory_managers/placement_budget.py | 14 +- .../multimodal_gen/runtime/server_warmup.py | 4 + .../runtime/warmup_request_builder.py | 10 + .../test/unit/test_auto_residency.py | 156 +++++++++++++++- .../test/unit/test_auto_residency_warmup.py | 32 ++++ .../test/unit/test_cfg_parallel_warmup.py | 1 + .../test/unit/test_component_residency.py | 25 +++ .../test/unit/test_performance_metrics.py | 16 +- .../test/unit/test_placement_budget.py | 70 +++++++ 12 files changed, 550 insertions(+), 13 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index a26dd6860a24..6123e8332c53 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -64,7 +64,9 @@ collect_promotion_candidates, component_runtime_weight_bytes, describe_error, + estimate_candidate_latency_savings_ns, estimate_default_workload_peak_bytes, + estimate_default_workload_timing, estimate_workload_phase_peaks, format_applied_changes, format_plan_summary, @@ -697,7 +699,11 @@ def _execute_forward_common( return output_batch def _record_server_warmup_memory( - self, *, req: Req, baseline_allocated_bytes: int, succeeded: bool + self, + *, + req: Req, + baseline_allocated_bytes: int, + succeeded: bool, ) -> None: phase_allocated_peaks: dict[str, int] = {} phase_components: dict[str, tuple[str, ...]] = {} @@ -719,6 +725,7 @@ def _record_server_warmup_memory( # second live placement, and is covered by the explicit reserve. phase_allocated_peaks["request:untracked"] = request_allocated_peak phase_components["request:untracked"] = untracked_active_components + metrics = req.metrics self._auto_residency_warmup_records.append( WarmupMemoryRecord( width=int(req.width or 0), @@ -729,6 +736,12 @@ def _record_server_warmup_memory( succeeded=succeeded, phase_peak_allocated_bytes=phase_allocated_peaks, phase_active_components=phase_components, + num_inference_steps=max(1, int(req.num_inference_steps or 1)), + total_duration_ms=( + float(metrics.total_duration_ms) if metrics is not None else 0.0 + ), + stage_duration_ms=(dict(metrics.stages) if metrics is not None else {}), + step_duration_ms=(tuple(metrics.steps) if metrics is not None else ()), ) ) @@ -1165,6 +1178,7 @@ def get_can_stay_resident_components( workload.num_inference_steps if workload is not None else 1 ), allow_host_pin_reallocation=self.server_args.num_gpus == 1, + mixed_dtype_components=self._mixed_dtype_residency_components(), ) can_stay_resident = [] @@ -1421,6 +1435,22 @@ def _build_auto_residency_report( custom_strategy_names=self.pipeline.component_residency_strategies, num_inference_steps=workload.num_inference_steps, allow_host_pin_reallocation=True, + mixed_dtype_components=self._mixed_dtype_residency_components(), + ) + ( + estimated_request_duration_ns, + stage_duration_ns, + component_stages, + ) = estimate_default_workload_timing( + records=records, + target_units=workload.workload_units(), + target_num_inference_steps=workload.num_inference_steps, + ) + candidate_latency_savings_ns = estimate_candidate_latency_savings_ns( + candidates=candidates, + request_duration_ns=estimated_request_duration_ns, + stage_duration_ns=stage_duration_ns, + component_stages=component_stages, ) local_worker_count = max( 1, self.server_args.num_gpus // self.server_args.nnodes @@ -1440,6 +1470,8 @@ def _build_auto_residency_report( 0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES ) // local_worker_count, + estimated_request_duration_ns=estimated_request_duration_ns, + candidate_latency_savings_ns=candidate_latency_savings_ns, candidates=candidates, skip_reason=skip_reason, ) @@ -1454,6 +1486,13 @@ def _auto_residency_all_gather(self, obj: Any) -> list[Any]: ) return gathered + @staticmethod + def _mixed_dtype_residency_components() -> set[str]: + manager = peek_global_component_residency_manager() + if manager is None: + return set() + return manager.components_with_mixed_use_dtypes() + def _rollback_applied_promotions( self, *, latest_round_only: bool = False ) -> str | None: diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 2b35ec804d19..609626e4a967 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -25,6 +25,7 @@ from __future__ import annotations +import statistics from typing import TYPE_CHECKING, Callable, Iterable, Mapping, Sequence import msgspec @@ -79,6 +80,14 @@ # share of what is actually there. MAX_VRAM_RESERVE_FRACTION = 0.20 +# A feasible placement is not automatically useful. Charge every placement a +# latency-risk cost before joint optimization so small, noisy predicted gains +# do not spend memory. The raw estimate is already an upper bound: transfer +# time is capped by the measured stage that can possibly wait for it. +ESTIMATED_PINNED_H2D_BYTES_PER_SECOND = 24 * GIB_BYTES +MIN_PLACEMENT_LATENCY_RISK_NS = 50_000_000 +PLACEMENT_LATENCY_RISK_FRACTION = 0.005 + AUTO_RESIDENCY_FEATURE_NAME = "auto residency promotion" PROMOTION_STATUS_SKIPPED = "skipped" @@ -113,6 +122,10 @@ class WarmupMemoryRecord(msgspec.Struct, frozen=True): succeeded: bool phase_peak_allocated_bytes: dict[str, int] = {} phase_active_components: dict[str, tuple[str, ...]] = {} + num_inference_steps: int = 1 + total_duration_ms: float = 0.0 + stage_duration_ms: dict[str, float] = {} + step_duration_ms: tuple[float, ...] = () def workload_units(self) -> int: return max(1, self.width) * max(1, self.height) * max(1, self.num_frames) @@ -185,6 +198,8 @@ class RankResidencyReport(msgspec.Struct, frozen=True): pinned_host_bytes: int = 0 host_pin_capacity_bytes: int = 0 host_transition_headroom_bytes: int = 0 + estimated_request_duration_ns: int = 0 + candidate_latency_savings_ns: dict[str, int] = {} candidates: list[PromotionCandidate] = [] skip_reason: str | None = None @@ -255,6 +270,129 @@ def resolve_measured_default_workload( ) +def estimate_default_workload_timing( + *, + records: Iterable[WarmupMemoryRecord], + target_units: int | None, + target_num_inference_steps: int, +) -> tuple[int, dict[str, int], dict[str, tuple[str, ...]]]: + """Estimate full-request and stage durations from the warmup workload. + + Non-denoise stages run once. Denoising time scales with the requested step + count; the full-shape probe intentionally executes only a few steps, so + using its raw total would make every one-shot encoder transfer look + important relative to a long video request. + """ + successful = [record for record in records if record.succeeded] + if not successful: + return 0, {}, {} + if target_units is not None: + at_target = [ + record for record in successful if record.workload_units() >= target_units + ] + if at_target: + successful = at_target + representative = max( + successful, + key=lambda record: ( + record.workload_units(), + record.total_duration_ms, + ), + ) + if representative.total_duration_ms <= 0 or not representative.stage_duration_ms: + return 0, {}, {} + + component_stages: dict[str, set[str]] = {} + for phase_name, components in representative.phase_active_components.items(): + fields = phase_name.split(":", 2) + if len(fields) < 2 or not fields[0].isdigit(): + continue + stage_name = fields[1] + if stage_name not in representative.stage_duration_ms: + continue + for component_name in components: + component_stages.setdefault(component_name, set()).add(stage_name) + + denoising_stages = { + stage_name + for component_name, stage_names in component_stages.items() + if is_dit_component_name(component_name) + for stage_name in stage_names + } + if not denoising_stages: + denoising_stages = { + stage_name + for stage_name in representative.stage_duration_ms + if stage_name.endswith("DenoisingStage") + and not stage_name.endswith("BeforeDenoisingStage") + } + + measured_steps = max(1, representative.num_inference_steps) + target_steps = max(1, target_num_inference_steps) + denoising_scale = target_steps / measured_steps + measured_denoising_ms = sum( + representative.stage_duration_ms[stage_name] for stage_name in denoising_stages + ) + if measured_denoising_ms > 0 and representative.step_duration_ms: + steady_steps = ( + representative.step_duration_ms[1:] + if len(representative.step_duration_ms) > 1 + else representative.step_duration_ms + ) + steady_step_ms = statistics.median(steady_steps) + non_step_denoising_ms = max( + 0.0, + measured_denoising_ms - sum(representative.step_duration_ms), + ) + target_denoising_ms = non_step_denoising_ms + steady_step_ms * target_steps + denoising_scale = target_denoising_ms / measured_denoising_ms + + stage_duration_ns: dict[str, int] = {} + for stage_name, duration_ms in representative.stage_duration_ms.items(): + scale = denoising_scale if stage_name in denoising_stages else 1.0 + stage_duration_ns[stage_name] = max(0, int(duration_ms * scale * 1_000_000)) + + measured_stage_ms = sum(representative.stage_duration_ms.values()) + untracked_ms = max(0.0, representative.total_duration_ms - measured_stage_ms) + request_duration_ns = sum(stage_duration_ns.values()) + int( + untracked_ms * 1_000_000 + ) + + return ( + request_duration_ns, + stage_duration_ns, + { + component_name: tuple(sorted(stage_names)) + for component_name, stage_names in component_stages.items() + }, + ) + + +def estimate_candidate_latency_savings_ns( + *, + candidates: Iterable[PromotionCandidate], + request_duration_ns: int, + stage_duration_ns: Mapping[str, int], + component_stages: Mapping[str, tuple[str, ...]], +) -> dict[str, int]: + """Upper-bound blocking H2D time removed by each placement option.""" + estimates: dict[str, int] = {} + for candidate in candidates: + transfer_ns = int( + candidate.h2d_bytes_per_request + / ESTIMATED_PINNED_H2D_BYTES_PER_SECOND + * 1_000_000_000 + ) + stages = component_stages.get(candidate.component_name, ()) + stage_cap_ns = sum(stage_duration_ns.get(stage, 0) for stage in stages) + if stage_cap_ns <= 0: + stage_cap_ns = request_duration_ns + estimates[candidate.option_key()] = ( + min(transfer_ns, stage_cap_ns) if transfer_ns > 0 else transfer_ns + ) + return estimates + + def estimate_default_workload_peak_bytes( *, records: Iterable[WarmupMemoryRecord], @@ -595,6 +733,7 @@ def collect_promotion_candidates( custom_strategy_names: Iterable[str], num_inference_steps: int, allow_host_pin_reallocation: bool = True, + mixed_dtype_components: Iterable[str] = (), ) -> list[PromotionCandidate]: """List implicitly offloaded native components eligible for promotion. @@ -603,6 +742,7 @@ def collect_promotion_candidates( pipeline-custom residency strategy are never touched. """ custom_names = set(custom_strategy_names) + mixed_dtype_names = set(mixed_dtype_components) candidates = [] for name in sorted(modules): module = modules[name] @@ -754,6 +894,14 @@ def collect_promotion_candidates( ) ) + # Layerwise stores have one fixed dtype. ResidentStrategy instead casts + # at every declared use, so a component with mixed use dtypes cannot be + # switched permanently without changing its numerical path. HostPin + # repacking and stage-scoped layer residency keep the layerwise strategy + # active and remain valid candidates. + if name in mixed_dtype_names: + continue + full_resident_layers = tuple(manager.num_layers for manager in managers) # Fully resident layers never read their host stores, so their pins can # be released and reused. Every worker plans against its non-overlapping @@ -1006,6 +1154,21 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP for report in reports ] + use_latency_utility = all( + report.estimated_request_duration_ns > 0 and report.candidate_latency_savings_ns + for report in reports + ) + placement_latency_risk_ns = ( + max( + MIN_PLACEMENT_LATENCY_RISK_NS, + int( + max(report.estimated_request_duration_ns for report in reports) + * PLACEMENT_LATENCY_RISK_FRACTION + ), + ) + if use_latency_utility + else 0 + ) options = [] for candidate in candidates: resource_deltas: dict[str, int] = {} @@ -1037,7 +1200,17 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP group_key=candidate.component_name, option_key=candidate.option_key(), resource_delta_bytes=resource_deltas, - estimated_latency_savings=candidate.h2d_bytes_per_request, + estimated_latency_savings=( + min( + report.candidate_latency_savings_ns.get( + candidate.option_key(), 0 + ) + for report in reports + ) + - placement_latency_risk_ns + if use_latency_utility + else candidate.h2d_bytes_per_request + ), ) ) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py index 6a2a1f25ba7c..83bb8239c183 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py @@ -174,6 +174,25 @@ def invalidate_component_strategies(self, component_names: Iterable[str]) -> Non for component_name in component_names: self._strategy_cache.pop(component_name, None) + def components_with_mixed_use_dtypes(self) -> set[str]: + """Components whose stages request more than one runtime dtype. + + Layerwise offload keeps its managed weight stores in one fixed dtype, + while resident placement honors each ``ComponentUse.target_dtype``. + Switching such a component between those strategies would therefore + change its numerical path, even when both placements fit. + """ + dtypes_by_component: dict[str, set[torch.dtype | None]] = {} + for use in self._ordered_uses: + dtypes_by_component.setdefault(use.component_name, set()).add( + use.target_dtype + ) + return { + component_name + for component_name, dtypes in dtypes_by_component.items() + if len(dtypes) > 1 + } + def begin_request( self, stages: Sequence[ComponentResidencyStage], diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py index 1fc8be4509f4..91ec6ee9ce73 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py @@ -19,9 +19,11 @@ class PlacementOption(msgspec.Struct, frozen=True): normalized form of absolute placement costs: the current placement is the implicit zero option, positive values consume headroom, and negative values release a resource. Keeping one resource map lets a single decision satisfy - rank/phase VRAM and node-scoped pinned-host constraints for one placement - state. A lifecycle with different load and serving placements must solve - those states separately and validate the transition between them. + rank/phase VRAM and node-scoped pinned-host constraints. A lifecycle option + may describe different load and serving placements by including load, + transition, and runtime resource dimensions in this same vector. The caller + must construct only compatible lifecycle combinations; independently solved + load and serving choices cannot be composed safely after the fact. ``estimated_latency_savings`` is an integer supplied by the caller. It can be transfer nanoseconds when measured bandwidth is available, or avoided @@ -101,9 +103,9 @@ def optimize_placement( constraint describing this placement state sees the same decision. A choice may trade one resource for another by using a negative delta. The result is exact for the supplied option frontier, resource model and - utility values; it does not imply that distinct load and serving states - should share one selection vector, or that the supplied frontier contains - every physically possible placement. + utility values. Distinct load and serving states may be encoded in one + option, but the result does not imply that the supplied frontier contains + every physically possible lifecycle plan. """ grouped: dict[str, list[PlacementOption]] = {} option_keys: set[str] = set() diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index e64556dfc322..15ddcd6b5fa5 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -202,6 +202,10 @@ async def rollback_and_rewarm(error: Exception) -> None: if status == PROMOTION_STATUS_ROLLBACK_FAILED: raise RuntimeError(f"auto residency rollback failed: {response.error}") if response.error is not None: + if status == PROMOTION_STATUS_ROLLED_BACK: + await run_async_client_warmup( + server_args, forward, fail_open=fail_open, rewarm=True + ) retained = ( "the last calibrated placement" if promoted_any diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py index 0fc445e4a0f5..02a6c9dbb782 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -555,6 +555,11 @@ def build_warmup_reqs( if warmup_resolutions is None else None ) + collect_auto_residency_metrics = ( + warmup_resolutions is None + and server_based_warmup + and auto_residency_args_skip_reason(server_args) is None + ) if auto_residency_warmup_shape is not None: width, height, warmup_num_frames = auto_residency_warmup_shape resolutions[0] = (width, height) @@ -618,6 +623,11 @@ def build_warmup_reqs( req.extra["return_warmup_result"] = True if server_based_warmup: req.extra["server_based_warmup"] = True + if collect_auto_residency_metrics: + # Stage timers already synchronize around warmup stages. Keep + # their values for residency planning instead of discarding + # the measurements after paying that cost. + req.metrics.suppress_stage_breakdown = False if auto_residency_warmup_shape is not None: req.extra["auto_residency_full_shape_probe"] = True warmup_reqs.append(req) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 3e8f7b7da3c8..6841c720307d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -24,7 +24,9 @@ collect_promotion_candidates, component_resident_size_bytes, component_runtime_weight_bytes, + estimate_candidate_latency_savings_ns, estimate_default_workload_peak_bytes, + estimate_default_workload_timing, estimate_workload_phase_peaks, format_applied_changes, plan_auto_residency, @@ -55,6 +57,11 @@ def _record( baseline_gib=10, peak_gib=12, succeeded=True, + num_inference_steps=1, + total_duration_ms=0.0, + stage_duration_ms=None, + step_duration_ms=(), + phase_active_components=None, ) -> WarmupMemoryRecord: return WarmupMemoryRecord( width=width, @@ -63,9 +70,94 @@ def _record( baseline_allocated_bytes=baseline_gib * GIB_BYTES, peak_allocated_bytes=peak_gib * GIB_BYTES, succeeded=succeeded, + num_inference_steps=num_inference_steps, + total_duration_ms=total_duration_ms, + stage_duration_ms=stage_duration_ms or {}, + step_duration_ms=step_duration_ms, + phase_active_components=phase_active_components or {}, ) +class TestEstimateDefaultWorkloadTiming: + def test_scales_only_denoising_with_target_steps(self): + record = _record( + num_inference_steps=2, + total_duration_ms=1_200, + stage_duration_ms={ + "TextEncodingStage": 100, + "DenoisingStage": 1_000, + }, + phase_active_components={ + "0:TextEncodingStage:use:text_encoder": ("text_encoder",), + "1:DenoisingStage:use:transformer": ("transformer",), + }, + ) + + request_ns, stage_ns, component_stages = estimate_default_workload_timing( + records=[record], + target_units=record.workload_units(), + target_num_inference_steps=40, + ) + + assert stage_ns == { + "TextEncodingStage": 100_000_000, + "DenoisingStage": 20_000_000_000, + } + assert request_ns == 20_200_000_000 + assert component_stages == { + "text_encoder": ("TextEncodingStage",), + "transformer": ("DenoisingStage",), + } + + def test_uses_steady_step_instead_of_scaling_first_step_setup(self): + record = _record( + num_inference_steps=2, + total_duration_ms=720, + stage_duration_ms={ + "TextEncodingStage": 100, + "DenoisingStage": 620, + }, + step_duration_ms=(500, 100), + phase_active_components={ + "0:TextEncodingStage:use:text_encoder": ("text_encoder",), + "1:DenoisingStage:use:transformer": ("transformer",), + }, + ) + + request_ns, stage_ns, _ = estimate_default_workload_timing( + records=[record], + target_units=record.workload_units(), + target_num_inference_steps=10, + ) + + assert stage_ns["DenoisingStage"] == 1_020_000_000 + assert request_ns == 1_120_000_000 + + def test_transfer_savings_is_capped_by_own_stage(self): + candidates = [ + _candidate("text_encoder", weight_gib=10, h2d_gib=10), + _candidate("transformer", weight_gib=10, h2d_gib=10), + _candidate("cold_encoder", weight_gib=10, h2d_gib=-1), + ] + + savings = estimate_candidate_latency_savings_ns( + candidates=candidates, + request_duration_ns=10_000_000_000, + stage_duration_ns={ + "TextEncodingStage": 80_000_000, + "DenoisingStage": 5_000_000_000, + }, + component_stages={ + "text_encoder": ("TextEncodingStage",), + "transformer": ("DenoisingStage",), + }, + ) + + assert savings[candidates[0].option_key()] == 80_000_000 + assert savings[candidates[1].option_key()] > 400_000_000 + assert savings[candidates[2].option_key()] < 0 + + class TestEstimateDefaultWorkloadPeak: def test_same_shape_uses_measured_peak(self): record = _record() @@ -319,6 +411,8 @@ def _report( pinned_host_gib: int = 0, host_pin_capacity_gib: int = 0, host_transition_headroom_gib: int = 0, + estimated_request_duration_ns: int = 0, + candidate_latency_savings_ns: dict[str, int] | None = None, ) -> RankResidencyReport: return RankResidencyReport( rank=rank, @@ -334,6 +428,8 @@ def _report( pinned_host_bytes=pinned_host_gib * GIB_BYTES, host_pin_capacity_bytes=host_pin_capacity_gib * GIB_BYTES, host_transition_headroom_bytes=(host_transition_headroom_gib * GIB_BYTES), + estimated_request_duration_ns=estimated_request_duration_ns, + candidate_latency_savings_ns=candidate_latency_savings_ns or {}, candidates=candidates if candidates is not None else [_candidate("vae")], skip_reason=skip_reason, ) @@ -426,7 +522,7 @@ def test_reserve_floor_is_capped_on_a_small_card(self): plan = plan_auto_residency(reports=[_report(budget_gib=10, estimated_gib=4)]) assert plan.reserve_bytes == 2 * GIB_BYTES - def test_greedy_promotion_by_h2d_savings(self): + def test_optimizes_promotion_by_h2d_savings(self): # dit saves 25 GiB x 50 steps per request; text_encoder saves 30 GiB # once. dit must be promoted first, after which the text encoder no # longer fits (50 + 25 + 30 + 10 > 100). @@ -439,6 +535,34 @@ def test_greedy_promotion_by_h2d_savings(self): ) assert [c.component_name for c in plan.promotions] == ["dit"] + def test_long_request_does_not_spend_vram_on_tiny_encoder_gain(self): + transformer = _candidate( + "transformer", + mode=LAYERWISE_OFFLOAD, + weight_gib=40, + h2d_gib=40 * 40, + ) + text_encoder = _candidate("text_encoder", weight_gib=8, h2d_gib=8) + candidates = [transformer, text_encoder] + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=80, + estimated_gib=20, + candidates=candidates, + estimated_request_duration_ns=66_000_000_000, + candidate_latency_savings_ns={ + transformer.option_key(): 10_000_000_000, + text_encoder.option_key(): 83_000_000, + }, + ) + ] + ) + + assert [candidate.component_name for candidate in plan.promotions] == [ + "transformer" + ] + def test_smaller_component_still_fits_after_skipping_a_big_one(self): candidates = [ _candidate("text_encoder", weight_gib=45, h2d_gib=100), @@ -664,12 +788,18 @@ def test_node_hostpin_is_optimized_with_the_same_selection_vector(self): pinned_host_gib=10, host_pin_capacity_gib=10, candidates=[cold_pageable, hot_pinned], + estimated_request_duration_ns=1_000_000_000, + candidate_latency_savings_ns={ + cold_pageable.option_key(): 10_000_000, + hot_pinned.option_key(): 500_000_000, + }, ) ] ) - # Plans are reported by benefit. apply_promotions independently sorts - # pin changes so the release still happens before the acquisition. + # The encoder's gain is below the risk floor, but its hostpin release + # makes the valuable DiT placement feasible. A joint solver must keep + # this cross-component trade instead of filtering it locally. assert [candidate.component_name for candidate in plan.promotions] == [ "hot_dit", "cold_encoder", @@ -1013,6 +1143,26 @@ def test_multi_worker_candidate_keeps_existing_hostpin_placement(self): assert permanent.target_layerwise_pinned_layers == ((0, 1),) assert permanent.pinned_host_delta_bytes == 0 + def test_mixed_dtype_layerwise_component_is_not_permanently_promoted(self): + manager = _FakeLayerwiseManager( + { + "layers.0.w": torch.zeros(16), + "layers.1.w": torch.zeros(16), + } + ) + module = _FakeLayerwiseDit([manager]) + + candidates = collect_promotion_candidates( + modules={"vae": module}, + residency_mode_of=self._modes({"vae": LAYERWISE_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=10, + mixed_dtype_components={"vae"}, + ) + + assert not any(candidate.permanent_residency for candidate in candidates) + class TestApplyAndRollback: def test_component_offload_promotion_marks_resident_without_moving(self): diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py index 03090e99df9f..e85e5600e3b3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -23,6 +23,38 @@ class TestAutoResidencyWarmup(unittest.TestCase): + def test_worker_rollback_response_rewarms_the_restored_layout(self): + actions = [] + + async def forward(req): + self.assertIsInstance(req, AutoResidencyReq) + actions.append(req.action) + return OutputBatch( + error="post-promotion warmup output changed", + output={"status": PROMOTION_STATUS_ROLLED_BACK}, + ) + + rewarm = mock.AsyncMock() + server_args = SimpleNamespace( + performance_mode="auto", + warmup_resolutions=None, + ) + with ( + mock.patch.object( + server_warmup, "auto_residency_skip_reason", return_value=None + ), + mock.patch.object(server_warmup, "run_async_client_warmup", rewarm), + ): + asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) + + self.assertEqual(actions, ["apply"]) + rewarm.assert_awaited_once_with( + server_args, + forward, + fail_open=True, + rewarm=True, + ) + def test_replans_until_the_calibrated_layout_reaches_a_fixed_point(self): responses = iter( [ diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index 2d2b7957e41a..f722e3b00018 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py @@ -429,6 +429,7 @@ def test_auto_residency_uses_one_full_serving_shape_probe(self): (1280, 720, 81), ) self.assertTrue(reqs[0].extra["auto_residency_full_shape_probe"]) + self.assertFalse(reqs[0].metrics.suppress_stage_breakdown) self.assertEqual(reqs[0].num_inference_steps, 2) self.assertIn( "auto residency probe (1280x720x81f, 2/35 steps)", diff --git a/python/sglang/multimodal_gen/test/unit/test_component_residency.py b/python/sglang/multimodal_gen/test/unit/test_component_residency.py index 21cf29e17dc4..39ece58b4a70 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_residency.py @@ -160,6 +160,31 @@ def component_uses(self, server_args, stage_name=None): return self.uses +def test_reports_components_with_mixed_use_dtypes(): + encode = _Stage(ComponentUse("encode", "vae", target_dtype=torch.float32)) + decode = _Stage(ComponentUse("decode", "vae", target_dtype=torch.bfloat16)) + denoise = _Stage( + ComponentUse("denoise", "transformer", target_dtype=torch.bfloat16) + ) + pipeline = SimpleNamespace( + modules={}, + _stage_name_mapping={ + "encode": encode, + "denoise": denoise, + "decode": decode, + }, + component_residency_strategies={}, + ) + server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + manager = ComponentResidencyManager(pipeline, server_args) + + manager.begin_request( + [encode, denoise, decode], SimpleNamespace(is_warmup=True), server_args + ) + + assert manager.components_with_mixed_use_dtypes() == {"vae"} + + def test_warmup_records_use_and_transition_peaks(monkeypatch): device_module = SimpleNamespace( is_available=lambda: True, diff --git a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py index 57914657f68c..d644e7e0143b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -109,7 +109,13 @@ def test_server_warmup_preserves_peak_after_managed_stage_timeline(): device_module = Mock() device_module.max_memory_allocated.return_value = 9 device_module.max_memory_reserved.return_value = 12 - req = SimpleNamespace(width=64, height=64, num_frames=1) + req = SimpleNamespace( + width=64, + height=64, + num_frames=1, + num_inference_steps=1, + metrics=None, + ) with ( patch.object(torch, "get_device_module", return_value=device_module), @@ -142,7 +148,13 @@ def test_server_warmup_does_not_treat_allocator_cache_as_untracked_live_memory() device_module = Mock() device_module.max_memory_allocated.return_value = 8 device_module.max_memory_reserved.return_value = 12 - req = SimpleNamespace(width=64, height=64, num_frames=1) + req = SimpleNamespace( + width=64, + height=64, + num_frames=1, + num_inference_steps=1, + metrics=None, + ) with ( patch.object(torch, "get_device_module", return_value=device_module), diff --git a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py index bea10327cb0f..e41b1d468305 100644 --- a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py @@ -165,6 +165,76 @@ def test_one_static_placement_satisfies_all_observed_phase_constraints(): } +def test_lifecycle_plan_jointly_constrains_load_transition_runtime_and_hostpin(): + plan = optimize_placement( + [ + PlacementOption( + group_key="dit-lifecycle", + option_key="dit-gpu-load-to-resident", + resource_delta_bytes={ + "gpu:rank0:load": 8, + "gpu:rank0:transition": 8, + "gpu:rank0:runtime": 6, + "hostpin:node0": 0, + }, + estimated_latency_savings=100, + ), + PlacementOption( + group_key="dit-lifecycle", + option_key="dit-sharded-load-to-layerwise", + resource_delta_bytes={ + "gpu:rank0:load": 4, + "gpu:rank0:transition": 3, + "gpu:rank0:runtime": 3, + "hostpin:node0": 8, + }, + estimated_latency_savings=80, + ), + PlacementOption( + group_key="encoder-lifecycle", + option_key="encoder-gpu-load-to-resident", + resource_delta_bytes={ + "gpu:rank0:load": 2, + "gpu:rank0:transition": 4, + "gpu:rank0:runtime": 4, + "hostpin:node0": 0, + }, + estimated_latency_savings=50, + ), + PlacementOption( + group_key="encoder-lifecycle", + option_key="encoder-cpu-load-to-pinned", + resource_delta_bytes={ + "gpu:rank0:load": 0, + "gpu:rank0:transition": 1, + "gpu:rank0:runtime": 1, + "hostpin:node0": 4, + }, + estimated_latency_savings=20, + ), + ], + resource_budget_bytes={ + "gpu:rank0:load": 10, + "gpu:rank0:transition": 10, + "gpu:rank0:runtime": 10, + "hostpin:node0": 8, + }, + ) + + # Independently optimizing load and runtime would pick both resident + # choices, but their transition uses 12 bytes and is infeasible. + assert [option.option_key for option in plan.selections] == [ + "dit-sharded-load-to-layerwise", + "encoder-gpu-load-to-resident", + ] + assert plan.resource_delta_bytes == { + "gpu:rank0:load": 6, + "gpu:rank0:transition": 7, + "gpu:rank0:runtime": 7, + "hostpin:node0": 8, + } + + def test_negative_headroom_forces_the_plan_to_release_a_resource(): plan = optimize_placement( [ From 1261612f3667092543da0e9214148c8093846977 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 01:29:04 +0800 Subject: [PATCH 29/90] [diffusion] let auto residency rebalance layerwise memory --- .../memory_managers/auto_residency.py | 30 ++++---- .../test/unit/test_auto_residency.py | 71 +++++++++++++++++++ 2 files changed, 88 insertions(+), 13 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 609626e4a967..49a9dfba3d36 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -672,11 +672,12 @@ def _layerwise_pin_targets( current_pinned_layers: tuple[tuple[int, ...], ...], uses_per_streamed_layer: int, ) -> list[tuple[tuple[int, ...], ...]]: - """Useful HostPin frontiers for one resident-layer placement. + """Bounded useful HostPin frontier for one resident-layer placement. - Streamed layers have the highest pin value because they transfer every - denoise step. Prefixes of that deterministic value order cover every - capacity breakpoint without enumerating all 2^N layer subsets. + Streamed layers have the highest value per byte because they transfer every + denoise step. Deterministic prefixes keep candidate generation linear in + the layer count; the joint planner still evaluates those targets against + rank VRAM, HostPin, and transition headroom together. """ ordered: list[tuple[int, int, int]] = [] for manager_index, (manager, resident_count) in enumerate( @@ -818,15 +819,11 @@ def collect_promotion_candidates( resident_targets = {current_resident_layers} if is_dit_component_name(name) and num_inference_steps > 1: for target in range( - max(current_resident_layers, default=0), + 0, max(manager.num_layers for manager in managers) + 1, ): counts = tuple(min(target, manager.num_layers) for manager in managers) - if all( - count >= current - for count, current in zip(counts, current_resident_layers) - ): - resident_targets.add(counts) + resident_targets.add(counts) for target_resident_layers in sorted(resident_targets): target_resident_bytes = sum( @@ -887,8 +884,8 @@ def collect_promotion_candidates( ), host_unpin_scratch_bytes=unpin_scratch, host_pin_scratch_bytes=pin_scratch, - active_device_delta_bytes=max( - 0, target_peak_device_bytes - current_peak_device_bytes + active_device_delta_bytes=( + target_peak_device_bytes - current_peak_device_bytes ), inactive_device_delta_bytes=0, ) @@ -1258,7 +1255,14 @@ def apply_promotions( try: ordered_promotions = sorted( plan.promotions, - key=lambda candidate: candidate.pinned_host_delta_bytes, + key=lambda candidate: ( + max( + candidate.active_device_delta_bytes, + candidate.inactive_device_delta_bytes, + ), + candidate.pinned_host_delta_bytes, + candidate.option_key(), + ), ) promotion_modules: dict[str, nn.Module] = {} snapshots: dict[str, AppliedPromotion] = {} diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 6841c720307d..5f12221b508d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -763,6 +763,46 @@ def test_partial_layer_residency_is_selected_when_full_dit_does_not_fit(self): assert len(plan.promotions) == 1 assert plan.promotions[0].target_layerwise_resident_layers == (12,) + def test_layerwise_release_can_unlock_a_more_valuable_component(self): + release_cold_layers = PromotionCandidate( + component_name="transformer_2", + residency_mode=LAYERWISE_OFFLOAD, + promoted_weight_bytes=-20 * GIB_BYTES, + h2d_bytes_per_request=-20 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + active_device_delta_bytes=-20 * GIB_BYTES, + ) + keep_hot_encoder = PromotionCandidate( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + promoted_weight_bytes=30 * GIB_BYTES, + h2d_bytes_per_request=100 * GIB_BYTES, + permanent_residency=True, + active_device_delta_bytes=0, + inactive_device_delta_bytes=30 * GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=100, + estimated_gib=80, + candidates=[release_cold_layers, keep_hot_encoder], + phase_peaks_gib={"denoise": 80, "encode": 80}, + phase_components={ + "denoise": ("transformer_2",), + "encode": ("text_encoder",), + }, + ) + ] + ) + + assert {candidate.component_name for candidate in plan.promotions} == { + "transformer_2", + "text_encoder", + } + def test_node_hostpin_is_optimized_with_the_same_selection_vector(self): cold_pageable = PromotionCandidate( component_name="cold_encoder", @@ -1093,6 +1133,37 @@ def test_mechanism_mismatch_is_not_a_candidate(self): ) assert candidates == [] + def test_layerwise_dit_can_release_existing_resident_layers(self): + manager = _FakeLayerwiseManager( + { + "layers.0.w": torch.zeros(16), + "layers.1.w": torch.zeros(16), + "layers.2.w": torch.zeros(16), + "layers.3.w": torch.zeros(16), + }, + resident_layers=3, + ) + module = _FakeLayerwiseDit([manager]) + + candidates = collect_promotion_candidates( + modules={"transformer": module}, + residency_mode_of=self._modes({"transformer": LAYERWISE_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=10, + ) + + target = next( + candidate + for candidate in candidates + if not candidate.permanent_residency + and candidate.target_layerwise_resident_layers == (1,) + and candidate.target_layerwise_pinned_layers == ((),) + ) + assert target.promoted_weight_bytes < 0 + assert target.active_device_delta_bytes < 0 + assert target.h2d_bytes_per_request < 0 + def test_fully_pinned_component_can_release_hostpin_for_a_hotter_component(self): manager = _FakeLayerwiseManager( { From 219186865d86ae36e775145f1e49749a0d0b14c6 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 01:46:01 +0800 Subject: [PATCH 30/90] [diffusion] preserve timing utility across ranks --- .../memory_managers/auto_residency.py | 17 ++++++--- .../test/unit/test_auto_residency.py | 37 +++++++++++++++++++ 2 files changed, 49 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 49a9dfba3d36..31cd2f8b44bc 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -1151,15 +1151,22 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP for report in reports ] - use_latency_utility = all( - report.estimated_request_duration_ns > 0 and report.candidate_latency_savings_ns + # Request timing is produced by the output rank. Other SPMD ranks still + # contribute their VRAM and HostPin constraints, but their empty metrics + # must not discard the replica's measured latency utility and fall back to + # the much more aggressive transfer-byte ordering. + timed_reports = [ + report for report in reports - ) + if report.estimated_request_duration_ns > 0 + and report.candidate_latency_savings_ns + ] + use_latency_utility = bool(timed_reports) placement_latency_risk_ns = ( max( MIN_PLACEMENT_LATENCY_RISK_NS, int( - max(report.estimated_request_duration_ns for report in reports) + max(report.estimated_request_duration_ns for report in timed_reports) * PLACEMENT_LATENCY_RISK_FRACTION ), ) @@ -1202,7 +1209,7 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP report.candidate_latency_savings_ns.get( candidate.option_key(), 0 ) - for report in reports + for report in timed_reports ) - placement_latency_risk_ns if use_latency_utility diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 5f12221b508d..c22d83e1cc47 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -563,6 +563,43 @@ def test_long_request_does_not_spend_vram_on_tiny_encoder_gain(self): "transformer" ] + def test_output_rank_timing_applies_to_untimed_peer_ranks(self): + transformer = _candidate( + "transformer", + mode=LAYERWISE_OFFLOAD, + weight_gib=40, + h2d_gib=40 * 40, + ) + text_encoder = _candidate("text_encoder", weight_gib=8, h2d_gib=8) + candidates = [transformer, text_encoder] + timing = { + transformer.option_key(): 10_000_000_000, + text_encoder.option_key(): 83_000_000, + } + + plan = plan_auto_residency( + reports=[ + _report( + rank=0, + budget_gib=80, + estimated_gib=20, + candidates=candidates, + estimated_request_duration_ns=66_000_000_000, + candidate_latency_savings_ns=timing, + ), + _report( + rank=1, + budget_gib=80, + estimated_gib=20, + candidates=candidates, + ), + ] + ) + + assert [candidate.component_name for candidate in plan.promotions] == [ + "transformer" + ] + def test_smaller_component_still_fits_after_skipping_a_big_one(self): candidates = [ _candidate("text_encoder", weight_gib=45, h2d_gib=100), From ba9612da8ea3e57e1341d4d094126226c5360904 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 03:52:13 +0800 Subject: [PATCH 31/90] [diffusion] solve complete residency target states --- .../runtime/managers/gpu_worker.py | 147 ++-- .../memory_managers/auto_residency.py | 827 ++++++++++++------ .../memory_managers/component_manager.py | 25 +- .../memory_managers/placement_budget.py | 115 ++- .../runtime/scheduler_client.py | 10 +- .../runtime/server_args/server_args.py | 17 + .../multimodal_gen/runtime/server_warmup.py | 38 +- .../test/unit/test_auto_residency.py | 545 +++++++++--- .../test/unit/test_auto_residency_warmup.py | 52 +- .../test/unit/test_component_residency.py | 12 +- .../test/unit/test_dp_routing.py | 14 +- .../test/unit/test_performance_metrics.py | 8 + .../test/unit/test_placement_budget.py | 94 ++ 13 files changed, 1369 insertions(+), 535 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 6123e8332c53..180c2a526204 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -50,18 +50,19 @@ save_outputs, ) from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( + AUTO_RESIDENCY_FEATURE_NAME, GIB_BYTES, - PROMOTION_STATUS_PROMOTED, - PROMOTION_STATUS_ROLLBACK_FAILED, - PROMOTION_STATUS_ROLLED_BACK, - PROMOTION_STATUS_SKIPPED, - AppliedPromotion, + PLACEMENT_STATUS_ADJUSTED, + PLACEMENT_STATUS_ROLLBACK_FAILED, + PLACEMENT_STATUS_ROLLED_BACK, + PLACEMENT_STATUS_SKIPPED, + AppliedResidencyChange, AutoResidencyRollbackError, DefaultWorkload, RankResidencyReport, WarmupMemoryRecord, - apply_promotions, - collect_promotion_candidates, + apply_residency_changes, + collect_residency_targets, component_runtime_weight_bytes, describe_error, estimate_candidate_latency_savings_ns, @@ -77,7 +78,7 @@ rank_candidates_by_h2d_savings, resolve_default_workload, resolve_measured_default_workload, - rollback_promotions, + rollback_residency_changes, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( peek_global_component_residency_manager, @@ -197,9 +198,9 @@ def __init__( self._realtime_sessions = RealtimeSessionCache(max_sessions=1) self.memory_occupation: MemoryOccupationController | None = None # per-rank memory measurements of server warmup forwards; consumed by - # the auto-residency promotion decision before the server turns ready + # the auto-residency placement decision before the server turns ready self._auto_residency_warmup_records: list[WarmupMemoryRecord] = [] - self._auto_residency_applied: list[AppliedPromotion] = [] + self._auto_residency_applied: list[AppliedResidencyChange] = [] self._auto_residency_round_sizes: list[int] = [] # default workload resolved once for the per-request residency hint self._cached_default_workload: DefaultWorkload | None = None @@ -422,7 +423,7 @@ def do_mem_analysis(self, output_batch: OutputBatch): f" Components that can remain on GPU: {can_stay_resident}. " "Make it explicit with --component-residency =resident; " "--performance-mode auto with server warmup applies safe " - "promotions automatically." + "adjustments automatically." if can_stay_resident else "" ) @@ -562,6 +563,16 @@ def _execute_forward_common( and bool(req.extra.get("server_based_warmup")) and current_platform.is_cuda() ) + warmup_workload = ( + ( + int(req.width or 0), + int(req.height or 0), + int(req.num_frames or 1), + max(1, int(req.num_inference_steps or 1)), + ) + if measure_server_warmup + else None + ) warmup_baseline_allocated_bytes = 0 try: if measure_server_warmup: @@ -691,8 +702,10 @@ def _execute_forward_common( # forward that never completed must still leave a failed record, # or the estimator would plan from the remaining partial data if measure_server_warmup: + assert warmup_workload is not None self._record_server_warmup_memory( req=req, + workload=warmup_workload, baseline_allocated_bytes=warmup_baseline_allocated_bytes, succeeded=output_batch is not None and output_batch.error is None, ) @@ -702,17 +715,20 @@ def _record_server_warmup_memory( self, *, req: Req, + workload: tuple[int, int, int, int], baseline_allocated_bytes: int, succeeded: bool, ) -> None: phase_allocated_peaks: dict[str, int] = {} phase_components: dict[str, tuple[str, ...]] = {} + phase_used_components: dict[str, tuple[str, ...]] = {} untracked_active_components: tuple[str, ...] = () residency_manager = peek_global_component_residency_manager() if residency_manager is not None: for phase_name, peak in residency_manager.take_warmup_phase_peaks().items(): phase_allocated_peaks[phase_name] = peak.allocated_bytes phase_components[phase_name] = peak.active_components + phase_used_components[phase_name] = peak.used_components untracked_active_components = residency_manager.current_device_components() request_allocated_peak = max( int(torch.get_device_module().max_memory_allocated()), @@ -725,18 +741,21 @@ def _record_server_warmup_memory( # second live placement, and is covered by the explicit reserve. phase_allocated_peaks["request:untracked"] = request_allocated_peak phase_components["request:untracked"] = untracked_active_components + phase_used_components["request:untracked"] = () metrics = req.metrics + width, height, num_frames, num_inference_steps = workload self._auto_residency_warmup_records.append( WarmupMemoryRecord( - width=int(req.width or 0), - height=int(req.height or 0), - num_frames=int(req.num_frames or 1), + width=width, + height=height, + num_frames=num_frames, baseline_allocated_bytes=int(baseline_allocated_bytes), peak_allocated_bytes=request_allocated_peak, succeeded=succeeded, phase_peak_allocated_bytes=phase_allocated_peaks, phase_active_components=phase_components, - num_inference_steps=max(1, int(req.num_inference_steps or 1)), + phase_used_components=phase_used_components, + num_inference_steps=num_inference_steps, total_duration_ms=( float(metrics.total_duration_ms) if metrics is not None else 0.0 ), @@ -1166,12 +1185,12 @@ def get_can_stay_resident_components( return [] workload = self._default_workload_for_hint() - candidates = collect_promotion_candidates( + candidates = collect_residency_targets( modules=self.pipeline.modules, residency_mode_of=self.server_args.residency_mode, # The hint also covers explicitly offloaded components: the user # chose offload and should learn when the headroom no longer - # requires it (promotion itself never touches explicit ones). + # requires it (automatic adjustment never touches explicit ones). explicit_residency_mode_of=lambda _name: None, custom_strategy_names=self.pipeline.component_residency_strategies, num_inference_steps=( @@ -1190,18 +1209,18 @@ def get_can_stay_resident_components( ): continue seen_components.add(candidate.component_name) - usage_gb = candidate.promoted_weight_bytes / GIB_BYTES + usage_gb = candidate.target_resident_weight_bytes / GIB_BYTES if usage_gb <= remaining_gpu_mem_gb: can_stay_resident.append(candidate.component_name) remaining_gpu_mem_gb -= usage_gb return can_stay_resident def apply_auto_residency(self) -> OutputBatch: - """Apply one warmup-calibrated residency promotion round. + """Apply one warmup-calibrated residency adjustment round. Every rank executes this handler at the same queue position; the plan is computed from all-gathered rank reports so each rank reaches the - same decision. Post-promotion warmup can expose more safe headroom, so + same decision. Post-adjustment warmup can expose more safe headroom, so the server may call this again with measurements from the new layout. A failure on any rank rolls every rank back to the previously calibrated placement. @@ -1212,7 +1231,7 @@ def apply_auto_residency(self) -> OutputBatch: """ records = list(self._auto_residency_warmup_records) # Each round must describe one placement. Intersecting phase ownership - # across old and newly promoted layouts would double-count weights that + # across old and newly adjusted layouts would double-count weights that # are already resident in the new layout. self._auto_residency_warmup_records.clear() @@ -1242,11 +1261,11 @@ def apply_auto_residency(self) -> OutputBatch: reports = self._auto_residency_all_gather(local_report) plan = plan_auto_residency(reports=reports) summary = format_plan_summary(plan=plan, workload=workload, records=records) - if plan.skip_reason is not None or not plan.promotions: + if plan.skip_reason is not None or not plan.changes: if self.is_output_rank: logger.info("%s", summary) return OutputBatch( - output=plan_summary_payload(plan=plan, status=PROMOTION_STATUS_SKIPPED) + output=plan_summary_payload(plan=plan, status=PLACEMENT_STATUS_SKIPPED) ) apply_error: str | None = None @@ -1256,18 +1275,18 @@ def apply_auto_residency(self) -> OutputBatch: # peer-triggered rollback must leave earlier calibrated rounds alone. self._auto_residency_round_sizes.append(0) try: - newly_applied = apply_promotions( + newly_applied = apply_residency_changes( plan=plan, modules=self.pipeline.modules, server_args=self.server_args, ) if not newly_applied: - raise RuntimeError("promotion plan applied no residency changes") + raise RuntimeError("placement plan applied no residency changes") self._auto_residency_applied.extend(newly_applied) self._auto_residency_round_sizes[-1] = len(newly_applied) except AutoResidencyRollbackError as e: logger.error( - "Auto residency promotion failed on rank %d and the rank " + "Auto residency adjustment failed on rank %d and the rank " "could not roll itself back: %s", self.rank, e, @@ -1277,7 +1296,7 @@ def apply_auto_residency(self) -> OutputBatch: local_rollback_failed = True except Exception as e: # this rank already rolled itself back logger.error( - "Auto residency promotion failed on rank %d: %s", + "Auto residency adjustment failed on rank %d: %s", self.rank, e, exc_info=True, @@ -1295,17 +1314,17 @@ def apply_auto_residency(self) -> OutputBatch: ) self._invalidate_component_strategies( - [candidate.component_name for candidate in plan.promotions] + [candidate.component_name for candidate in plan.changes] ) if self.is_output_rank: logger.info("%s", summary) logger.info("%s", format_applied_changes(plan=plan)) return OutputBatch( - output=plan_summary_payload(plan=plan, status=PROMOTION_STATUS_PROMOTED) + output=plan_summary_payload(plan=plan, status=PLACEMENT_STATUS_ADJUSTED) ) def rollback_auto_residency(self) -> OutputBatch: - """Undo the latest promotion round after its calibration failed.""" + """Undo the latest adjustment round after its calibration failed.""" return self._rollback_everywhere( cause=None, already_failed=False, latest_round_only=True ) @@ -1319,28 +1338,28 @@ def _rollback_everywhere( ) -> OutputBatch: """Roll this rank back and gather the replica-wide outcome. - ``cause`` is the promotion failure that triggered the rollback (None + ``cause`` is the adjustment failure that triggered the rollback (None when the rollback was requested after a failed re-warm); ``already_failed`` marks that some rank already failed its in-apply rollback, which is fatal regardless of what the remaining ranks do. Collective-symmetric: every rank gathers exactly once. """ - rollback_error = self._rollback_applied_promotions( + rollback_error = self._rollback_applied_residency_changes( latest_round_only=latest_round_only ) gathered = self._auto_residency_all_gather(rollback_error) rank_errors = [error for error in gathered if error is not None] if already_failed and not rank_errors: - rank_errors = ["a rank could not undo its own promotions"] + rank_errors = ["a rank could not undo its own residency changes"] if rank_errors: prefix = ( - f"auto residency promotion failed ({cause}) and rollback failed" + f"auto residency adjustment failed ({cause}) and rollback failed" if cause is not None else "auto residency rollback failed" ) return OutputBatch( error=f"{prefix}: {rank_errors[0]}", - output={"status": PROMOTION_STATUS_ROLLBACK_FAILED}, + output={"status": PLACEMENT_STATUS_ROLLBACK_FAILED}, ) if cause is not None: restored = ( @@ -1350,14 +1369,14 @@ def _rollback_everywhere( ) return OutputBatch( error=( - f"auto residency promotion failed; {restored} restored: {cause}" + f"auto residency adjustment failed; {restored} restored: {cause}" ), - output={"status": PROMOTION_STATUS_ROLLED_BACK}, + output={"status": PLACEMENT_STATUS_ROLLED_BACK}, ) if self.is_output_rank: if self._auto_residency_applied: logger.info( - "Auto residency: rolled back the latest promotion round; " + "Auto residency: rolled back the latest adjustment round; " "the previously calibrated placement remains active." ) else: @@ -1366,7 +1385,7 @@ def _rollback_everywhere( "startup is in effect again (no equivalent server-arg " "changes remain)." ) - return OutputBatch(output={"status": PROMOTION_STATUS_ROLLED_BACK}) + return OutputBatch(output={"status": PLACEMENT_STATUS_ROLLED_BACK}) def _build_auto_residency_report( self, *, workload: DefaultWorkload, records: List[WarmupMemoryRecord] @@ -1381,7 +1400,7 @@ def _build_auto_residency_report( if skip_reason is None: # A probe that failed at or below the target is a measurement, not # missing data: the card cannot hold the default workload as it is - # already configured, and promoting weights would only add to it. + # already configured, and increasing residency would only add to it. blocking = [ record for record in records @@ -1403,20 +1422,24 @@ def _build_auto_residency_report( candidates=[], skip_reason=skip_reason, ) + target_units = workload.workload_units() + assert target_units is not None runtime_weights_by_component = component_runtime_weight_bytes( self.pipeline.modules ) estimated_peak_bytes = estimate_default_workload_peak_bytes( records=records, - target_units=workload.workload_units(), + target_units=target_units, constant_weight_bytes=max(runtime_weights_by_component.values(), default=0), ) - estimated_phase_peaks, active_components_by_phase = ( - estimate_workload_phase_peaks( - records=records, - target_units=workload.workload_units(), - component_weight_bytes=runtime_weights_by_component, - ) + ( + estimated_phase_peaks, + active_components_by_phase, + used_components_by_phase, + ) = estimate_workload_phase_peaks( + records=records, + target_units=target_units, + component_weight_bytes=runtime_weights_by_component, ) # What this process may still grow into: free VRAM plus what its # allocator already reserved. Unlike the raw device total, this @@ -1428,29 +1451,30 @@ def _build_auto_residency_report( test_cap_gib = envs.SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB if test_cap_gib is not None: budget_bytes = min(budget_bytes, int(test_cap_gib * GIB_BYTES)) - candidates = collect_promotion_candidates( + candidates = collect_residency_targets( modules=self.pipeline.modules, residency_mode_of=self.server_args.residency_mode, + baseline_residency_mode_of=(self.server_args.configured_residency_mode), explicit_residency_mode_of=self.server_args.explicit_residency_mode, custom_strategy_names=self.pipeline.component_residency_strategies, num_inference_steps=workload.num_inference_steps, allow_host_pin_reallocation=True, mixed_dtype_components=self._mixed_dtype_residency_components(), + auto_resident_components={ + name + for name in self.pipeline.modules + if self.server_args.component_residency_requirement(name) + == AUTO_RESIDENCY_FEATURE_NAME + }, ) - ( - estimated_request_duration_ns, - stage_duration_ns, - component_stages, - ) = estimate_default_workload_timing( + estimated_request_duration_ns, _, _ = estimate_default_workload_timing( records=records, - target_units=workload.workload_units(), + target_units=target_units, target_num_inference_steps=workload.num_inference_steps, ) candidate_latency_savings_ns = estimate_candidate_latency_savings_ns( candidates=candidates, request_duration_ns=estimated_request_duration_ns, - stage_duration_ns=stage_duration_ns, - component_stages=component_stages, ) local_worker_count = max( 1, self.server_args.num_gpus // self.server_args.nnodes @@ -1459,8 +1483,13 @@ def _build_auto_residency_report( rank=self.rank, budget_bytes=budget_bytes, estimated_peak_bytes=estimated_peak_bytes, + target_workload_measured=any( + record.succeeded and record.workload_units() >= target_units + for record in records + ), estimated_peak_bytes_by_phase=estimated_phase_peaks, active_components_by_phase=active_components_by_phase, + used_components_by_phase=used_components_by_phase, node_rank=self.server_args.node_rank, pinned_host_bytes=layerwise_pinned_host_bytes(self.pipeline.modules), host_pin_capacity_bytes=layerwise_host_pin_capacity_bytes( @@ -1493,7 +1522,7 @@ def _mixed_dtype_residency_components() -> set[str]: return set() return manager.components_with_mixed_use_dtypes() - def _rollback_applied_promotions( + def _rollback_applied_residency_changes( self, *, latest_round_only: bool = False ) -> str | None: if latest_round_only: @@ -1512,7 +1541,7 @@ def _rollback_applied_promotions( return None applied = self._auto_residency_applied try: - rollback_promotions( + rollback_residency_changes( applied=applied, modules=self.pipeline.modules, server_args=self.server_args, @@ -1526,7 +1555,7 @@ def _rollback_applied_promotions( ) return describe_error(e) self._invalidate_component_strategies( - [promotion.component_name for promotion in applied] + [adjustment.component_name for adjustment in applied] ) if latest_round_only and self._auto_residency_round_sizes: del self._auto_residency_applied[-len(applied) :] diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 31cd2f8b44bc..58935e50ec02 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -1,18 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 -"""Warmup-calibrated automatic component residency promotion. +"""Warmup-calibrated automatic component residency placement. Under ``--performance-mode auto`` with server warmup, each rank measures the peak GPU memory of bounded synthetic warmup requests and a low-step probe at -the complete default serving shape, then promotes implicitly offloaded components -(component offload -> resident, layerwise offload -> fully loaded) when the -estimate plus the promoted weights still fits under a safety reserve. +the complete default serving shape, then selects a complete serving placement +for every eligible component under the measured memory constraints. When no full-shape measurement is available, the fallback estimate splits the measured peak into persistent weights and workload-scaled activations. Scaling the whole peak would multiply resident weights by the video frame/area cap -ratio (~16x for Wan-class defaults) and promotion would never trigger. +ratio (~16x for Wan-class defaults) and residency adjustment would never trigger. -Promotion targets the model default workload only (default resolution, +The planner targets the model default workload only (default resolution, default frames, batch=1). Larger shapes, batches, or multi-image inputs need explicit ``--component-residency``. @@ -35,6 +34,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD, + RESIDENT, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import ( is_fsdp_managed_module, @@ -70,34 +70,36 @@ # Activation memory rarely scales perfectly linearly with workload units; # pad the extrapolated activation part before checking the budget. ACTIVATION_EXTRAPOLATION_MARGIN = 1.2 -# Always keep this much VRAM free after promotion for allocator slack, -# shape variance, and CUDA graph or compile pools. -VRAM_RESERVE_FRACTION = 0.10 +# A target-shape measurement plus the mandatory post-placement warmup justifies +# a tighter reserve than an extrapolated estimate. Both retain an absolute +# floor for allocator slack, shape variance, and CUDA graph or compile pools. +MEASURED_VRAM_RESERVE_FRACTION = 0.05 +EXTRAPOLATED_VRAM_RESERVE_FRACTION = 0.10 MIN_VRAM_RESERVE_BYTES = 4 * GIB_BYTES -# The absolute floor is sized for datacenter cards, where the fraction dominates -# anyway (10% of 80 GiB is 8 GiB). Below ~40 GiB the floor takes over, and on a -# 12 GiB card a flat 4 GiB would fence off a third of the device, so cap it as a -# share of what is actually there. +# The absolute floor is sized for datacenter cards, where either fraction can +# dominate. On a 12 GiB card a flat 4 GiB would fence off a third of the device, +# so cap the floor as a share of what is actually there. MAX_VRAM_RESERVE_FRACTION = 0.20 -# A feasible placement is not automatically useful. Charge every placement a -# latency-risk cost before joint optimization so small, noisy predicted gains -# do not spend memory. The raw estimate is already an upper bound: transfer -# time is capped by the measured stage that can possibly wait for it. +# A feasible placement is not automatically useful. Predictions inside this +# interval are treated as latency-equivalent, then the joint optimizer chooses +# the option with the lowest additional VRAM and HostPin use. The raw estimate +# is already an upper bound: transfer time is capped by the measured request. ESTIMATED_PINNED_H2D_BYTES_PER_SECOND = 24 * GIB_BYTES -MIN_PLACEMENT_LATENCY_RISK_NS = 50_000_000 -PLACEMENT_LATENCY_RISK_FRACTION = 0.005 +MIN_LATENCY_EQUIVALENCE_NS = 50_000_000 +MAX_LATENCY_EQUIVALENCE_NS = 100_000_000 +LATENCY_EQUIVALENCE_FRACTION = 0.01 -AUTO_RESIDENCY_FEATURE_NAME = "auto residency promotion" +AUTO_RESIDENCY_FEATURE_NAME = "auto residency" -PROMOTION_STATUS_SKIPPED = "skipped" -PROMOTION_STATUS_PROMOTED = "promoted" -PROMOTION_STATUS_ROLLED_BACK = "rolled_back" -PROMOTION_STATUS_ROLLBACK_FAILED = "rollback_failed" +PLACEMENT_STATUS_SKIPPED = "skipped" +PLACEMENT_STATUS_ADJUSTED = "adjusted" +PLACEMENT_STATUS_ROLLED_BACK = "rolled_back" +PLACEMENT_STATUS_ROLLBACK_FAILED = "rollback_failed" class AutoResidencyRollbackError(RuntimeError): - """A promotion failed AND undoing the already-applied part also failed. + """An adjustment failed and undoing its applied part also failed. The rank is in a mixed residency state that only a restart fixes; callers must abort startup instead of continuing on a half-rolled-back replica. @@ -122,6 +124,7 @@ class WarmupMemoryRecord(msgspec.Struct, frozen=True): succeeded: bool phase_peak_allocated_bytes: dict[str, int] = {} phase_active_components: dict[str, tuple[str, ...]] = {} + phase_used_components: dict[str, tuple[str, ...]] = {} num_inference_steps: int = 1 total_duration_ms: float = 0.0 stage_duration_ms: dict[str, float] = {} @@ -131,16 +134,16 @@ def workload_units(self) -> int: return max(1, self.width) * max(1, self.height) * max(1, self.num_frames) -class PromotionCandidate(msgspec.Struct, frozen=True): - """An implicitly offloaded component that could become resident.""" +class ResidencyTarget(msgspec.Struct, frozen=True): + """One complete target state for an auto-managed component.""" component_name: str residency_mode: str - promoted_weight_bytes: int - # Estimated per-request host-to-device traffic this promotion removes. + target_resident_weight_bytes: int + # Estimated per-request host-to-device traffic this target removes. h2d_bytes_per_request: int # Layerwise candidates jointly choose stage-scoped GPU residency and host - # pinning. None is used by ordinary component-offload promotion. + # pinning. None is used by ordinary component placement. target_layerwise_resident_layers: tuple[int, ...] | None = None target_layerwise_pinned_layers: tuple[tuple[int, ...], ...] | None = None pinned_host_delta_bytes: int = 0 @@ -151,9 +154,32 @@ class PromotionCandidate(msgspec.Struct, frozen=True): # already loaded for its own phase has a different delta from phases where # it is absent; keeping both avoids adding the same weights twice. active_device_delta_bytes: int = 0 + # Delta when the component is already present because of async prefetch, + # but is not the semantic owner of this phase. + present_device_delta_bytes: int = 0 inactive_device_delta_bytes: int = 0 + # None preserves the historical derived target for hand-built callers: + # partial layerwise targets remain layerwise, every other option is + # resident. Generated complete-state frontiers set this explicitly. + target_residency_mode: str | None = None + current_placement: bool = False + target_device_weight_bytes: int = 0 + target_pinned_host_bytes: int = 0 + + def target_mode(self) -> str: + if self.target_residency_mode is not None: + return self.target_residency_mode + if ( + self.target_layerwise_resident_layers is not None + and not self.permanent_residency + ): + return LAYERWISE_OFFLOAD + return RESIDENT def option_key(self) -> str: + target_mode = self.target_mode() + if target_mode == COMPONENT_OFFLOAD: + return f"{self.component_name}:component-offload" if self.target_layerwise_resident_layers is None: return f"{self.component_name}:resident" layer_counts = ",".join( @@ -168,7 +194,7 @@ def option_key(self) -> str: class DefaultWorkload(msgspec.Struct, frozen=True): - """The model-default request shape the promotion is calibrated for.""" + """The model-default request shape the planner is calibrated for.""" width: int | None height: int | None @@ -187,45 +213,48 @@ def describe(self) -> str: class RankResidencyReport(msgspec.Struct, frozen=True): - """One rank's inputs to the replica-wide promotion decision.""" + """One rank's inputs to the replica-wide placement decision.""" rank: int budget_bytes: int estimated_peak_bytes: int | None + target_workload_measured: bool = False estimated_peak_bytes_by_phase: dict[str, int] = {} active_components_by_phase: dict[str, tuple[str, ...]] = {} + used_components_by_phase: dict[str, tuple[str, ...]] = {} node_rank: int = 0 pinned_host_bytes: int = 0 host_pin_capacity_bytes: int = 0 host_transition_headroom_bytes: int = 0 estimated_request_duration_ns: int = 0 candidate_latency_savings_ns: dict[str, int] = {} - candidates: list[PromotionCandidate] = [] + candidates: list[ResidencyTarget] = [] skip_reason: str | None = None class AutoResidencyPlan(msgspec.Struct, frozen=True): - """Deterministic promotion decision shared by every rank.""" + """Deterministic placement changes shared by every rank.""" estimated_peak_bytes: int = 0 reserve_bytes: int = 0 budget_bytes: int = 0 resource_budget_bytes: dict[str, int] = {} - promotions: list[PromotionCandidate] = [] + changes: list[ResidencyTarget] = [] skip_reason: str | None = None -class AppliedPromotion(msgspec.Struct, frozen=True): +class AppliedResidencyChange(msgspec.Struct, frozen=True): component_name: str residency_mode: str previous_layerwise_resident_layers: tuple[int, ...] | None = None previous_layerwise_pinned_layers: tuple[tuple[int, ...], ...] | None = None - layerwise_offload_disabled: bool = False + previous_layerwise_offload_enabled: bool = True + previous_required_resident: bool = False pinned_host_changed: bool = False def resolve_default_workload(server_args: ServerArgs) -> DefaultWorkload: - """Resolve the default request shape promotion is optimized for.""" + """Resolve the default request shape the planner is optimized for.""" from sglang.multimodal_gen.runtime.warmup_request_builder import ( get_model_sampling_defaults, resolve_default_workload_shape, @@ -303,7 +332,10 @@ def estimate_default_workload_timing( return 0, {}, {} component_stages: dict[str, set[str]] = {} - for phase_name, components in representative.phase_active_components.items(): + phase_components = ( + representative.phase_used_components or representative.phase_active_components + ) + for phase_name, components in phase_components.items(): fields = phase_name.split(":", 2) if len(fields) < 2 or not fields[0].isdigit(): continue @@ -370,12 +402,15 @@ def estimate_default_workload_timing( def estimate_candidate_latency_savings_ns( *, - candidates: Iterable[PromotionCandidate], + candidates: Iterable[ResidencyTarget], request_duration_ns: int, - stage_duration_ns: Mapping[str, int], - component_stages: Mapping[str, tuple[str, ...]], ) -> dict[str, int]: - """Upper-bound blocking H2D time removed by each placement option.""" + """Upper-bound H2D time removed by each complete placement option. + + An async prefetch can consume copy-engine and memory bandwidth outside the + component's own stage, so that stage is not a sound cap. The complete + request duration is the only generally valid upper bound. + """ estimates: dict[str, int] = {} for candidate in candidates: transfer_ns = int( @@ -383,12 +418,10 @@ def estimate_candidate_latency_savings_ns( / ESTIMATED_PINNED_H2D_BYTES_PER_SECOND * 1_000_000_000 ) - stages = component_stages.get(candidate.component_name, ()) - stage_cap_ns = sum(stage_duration_ns.get(stage, 0) for stage in stages) - if stage_cap_ns <= 0: - stage_cap_ns = request_duration_ns estimates[candidate.option_key()] = ( - min(transfer_ns, stage_cap_ns) if transfer_ns > 0 else transfer_ns + min(transfer_ns, request_duration_ns) + if transfer_ns > 0 and request_duration_ns > 0 + else transfer_ns ) return estimates @@ -416,7 +449,7 @@ def estimate_default_workload_peak_bytes( activations -- measured on Wan2.1-14B, single-point scaling overestimated a ~30 GiB peak as ~183 GiB. 3. One usable size: scale everything above the pre-forward allocated - baseline (conservative; may block promotion but never over-promotes). + baseline (conservative; may block adjustment but never over-allocates). Returns None when the estimate cannot be trusted: no successful records, the target workload is unknown (an unknown target would silently equate the @@ -487,39 +520,52 @@ def estimate_workload_phase_peaks( records: Iterable[WarmupMemoryRecord], target_units: int | None, component_weight_bytes: Mapping[str, int], -) -> tuple[dict[str, int], dict[str, tuple[str, ...]]]: +) -> tuple[ + dict[str, int], + dict[str, tuple[str, ...]], + dict[str, tuple[str, ...]], +]: """Estimate each measured execution phase at the target workload. A component already active in a phase is part of that phase's measured peak. Keeping a component resident therefore adds no weight bytes to its own component-offload phase, while it adds the full footprint to phases - where the component was absent. Returning both the estimated peaks and the - conservative intersection of active components preserves that distinction. + where the component was absent. Measurements with different active layouts + remain separate constraints; combining one layout's peak with another + layout's component set would describe a state that never occurred. """ - records = list(records) successful = [record for record in records if record.succeeded] - phase_names = sorted( - { - phase_name - for record in successful - for phase_name in record.phase_peak_allocated_bytes - } - ) + if target_units is not None: + covering = [ + record for record in successful if record.workload_units() >= target_units + ] + if covering: + successful = covering + + grouped: dict[ + tuple[str, tuple[str, ...], tuple[str, ...]], list[WarmupMemoryRecord] + ] = {} + for record in successful: + used_by_phase = record.phase_used_components or record.phase_active_components + for phase_name in record.phase_peak_allocated_bytes: + active = tuple(sorted(record.phase_active_components.get(phase_name, ()))) + used = tuple(sorted(used_by_phase.get(phase_name, ()))) + grouped.setdefault((phase_name, active, used), []).append(record) + + layouts_per_phase: dict[str, int] = {} + for phase_name, _, _ in grouped: + layouts_per_phase[phase_name] = layouts_per_phase.get(phase_name, 0) + 1 + estimated_peaks: dict[str, int] = {} active_components: dict[str, tuple[str, ...]] = {} - for phase_name in phase_names: - phase_records = [ - record - for record in successful - if phase_name in record.phase_peak_allocated_bytes - ] - if not phase_records: - continue - active_sets = [ - set(record.phase_active_components.get(phase_name, ())) - for record in phase_records - ] - active = set.intersection(*active_sets) if active_sets else set() + used_components: dict[str, tuple[str, ...]] = {} + layout_indices: dict[str, int] = {} + for (phase_name, active, used), phase_records in sorted(grouped.items()): + output_name = phase_name + if layouts_per_phase[phase_name] > 1: + index = layout_indices.get(phase_name, 0) + layout_indices[phase_name] = index + 1 + output_name = f"{phase_name}:layout:{index}" weight_floor = sum(component_weight_bytes.get(name, 0) for name in active) phase_measurements = [ WarmupMemoryRecord( @@ -542,9 +588,10 @@ def estimate_workload_phase_peaks( ) if estimate is None: continue - estimated_peaks[phase_name] = estimate - active_components[phase_name] = tuple(sorted(active)) - return estimated_peaks, active_components + estimated_peaks[output_name] = estimate + active_components[output_name] = active + used_components[output_name] = used + return estimated_peaks, active_components, used_components def _module_weight_bytes(module: nn.Module) -> int: @@ -625,7 +672,7 @@ def _layerwise_offloaded_bytes(module: LayerwiseOffloadableModuleMixin) -> int: def component_resident_size_bytes(module: nn.Module, residency_mode: str) -> int: - """Extra GPU bytes a promotion of this component would pin resident.""" + """GPU weight bytes used by a resident version of this component.""" if residency_mode == LAYERWISE_OFFLOAD: if not isinstance(module, LayerwiseOffloadableModuleMixin): return 0 @@ -726,24 +773,30 @@ def _layerwise_host_transition_bytes( return unpin_bytes, pin_bytes -def collect_promotion_candidates( +def collect_residency_targets( *, modules: Mapping[str, object], residency_mode_of: Callable[[str], str], + baseline_residency_mode_of: Callable[[str], str] | None = None, explicit_residency_mode_of: Callable[[str], str | None], custom_strategy_names: Iterable[str], num_inference_steps: int, allow_host_pin_reallocation: bool = True, mixed_dtype_components: Iterable[str] = (), -) -> list[PromotionCandidate]: - """List implicitly offloaded native components eligible for promotion. - - Only components whose residency was chosen by the auto policy qualify: - explicit placements, FSDP-managed modules, and components driven by a - pipeline-custom residency strategy are never touched. + auto_resident_components: Iterable[str] = (), +) -> list[ResidencyTarget]: + """Build complete target-state frontiers for auto-managed components. + + Every option is expressed relative to the currently measured placement, + but its transfer utility is absolute within the component's frontier. + Including the current and lower-memory states lets later calibration rounds + replace an earlier choice instead of being limited to monotonic upgrades. """ custom_names = set(custom_strategy_names) mixed_dtype_names = set(mixed_dtype_components) + auto_resident_names = set(auto_resident_components) + if baseline_residency_mode_of is None: + baseline_residency_mode_of = residency_mode_of candidates = [] for name in sorted(modules): module = modules[name] @@ -751,79 +804,122 @@ def collect_promotion_candidates( continue if name in custom_names: continue - mode = residency_mode_of(name) - if mode not in (COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD): + baseline_mode = baseline_residency_mode_of(name) + if baseline_mode not in (COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD): continue if explicit_residency_mode_of(name) is not None: continue if is_fsdp_managed_module(module): continue - if mode == COMPONENT_OFFLOAD and is_layerwise_offloaded_module(module): - # The module is layerwise-managed despite its configured mode - # (e.g. the offload_during_compile window); flipping the flag - # would be a silent no-op while the manager keeps streaming. - continue - if mode == LAYERWISE_OFFLOAD and not is_layerwise_offloaded_module(module): + current_mode = residency_mode_of(name) + if current_mode == RESIDENT and name not in auto_resident_names: + # A loader or another runtime feature owns this hard requirement. continue - promoted_bytes = component_resident_size_bytes(module, mode) - if promoted_bytes <= 0: - continue - # A layerwise DiT re-streams its layers once per denoise forward; - # every other offloaded component transfers once per request. - uses_per_request = ( - max(1, num_inference_steps) - if mode == LAYERWISE_OFFLOAD and is_dit_component_name(name) - else 1 - ) - if mode == COMPONENT_OFFLOAD: - candidates.append( - PromotionCandidate( - component_name=name, - residency_mode=mode, - promoted_weight_bytes=promoted_bytes, - h2d_bytes_per_request=promoted_bytes, - permanent_residency=True, - active_device_delta_bytes=0, - inactive_device_delta_bytes=promoted_bytes, - ) + + if baseline_mode == COMPONENT_OFFLOAD: + if current_mode not in (COMPONENT_OFFLOAD, RESIDENT): + continue + if is_layerwise_offloaded_module(module): + # The module is layerwise-managed despite its configured mode + # (e.g. the offload_during_compile window). Moving it behind + # the manager would strand its bookkeeping. + continue + weight_bytes = _module_weight_bytes(module) + if weight_bytes <= 0: + continue + current_resident = current_mode == RESIDENT + candidates.extend( + [ + ResidencyTarget( + component_name=name, + residency_mode=baseline_mode, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + active_device_delta_bytes=0, + inactive_device_delta_bytes=( + -weight_bytes if current_resident else 0 + ), + present_device_delta_bytes=( + -weight_bytes if current_resident else 0 + ), + current_placement=not current_resident, + ), + ResidencyTarget( + component_name=name, + residency_mode=baseline_mode, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=weight_bytes, + h2d_bytes_per_request=weight_bytes, + permanent_residency=True, + active_device_delta_bytes=0, + inactive_device_delta_bytes=( + 0 if current_resident else weight_bytes + ), + present_device_delta_bytes=0, + current_placement=current_resident, + target_device_weight_bytes=weight_bytes, + ), + ] ) continue if not isinstance(module, LayerwiseOffloadableModuleMixin): continue - managers = [ - manager for manager in module.layerwise_offload_managers if manager.enabled - ] + managers = list(module.layerwise_offload_managers) if not managers: continue + enabled = {manager.enabled for manager in managers} + if len(enabled) != 1: + logger.warning( + "Skipping auto residency for %s: layerwise groups disagree on " + "whether offload is enabled", + name, + ) + continue + current_permanent = not enabled.pop() + if current_permanent != (current_mode == RESIDENT): + # The module is layerwise-managed despite its configured mode + # or is hard-required resident by a different owner. + continue + managed_weight_bytes = sum( + sum(manager.layer_weight_bytes().values()) for manager in managers + ) + if managed_weight_bytes <= 0: + continue + # A layerwise DiT re-streams its layers once per denoise forward; + # every other offloaded component transfers once per request. + uses_per_request = ( + max(1, num_inference_steps) if is_dit_component_name(name) else 1 + ) current_resident_layers = tuple(manager.resident_layers for manager in managers) current_pinned_layers = tuple( manager.pinned_layer_indices() for manager in managers ) - current_resident_bytes = sum( - manager.resident_weight_bytes() for manager in managers - ) - current_peak_device_bytes = sum( - manager.peak_managed_device_weight_bytes() for manager in managers + current_peak_device_bytes = ( + managed_weight_bytes + if current_permanent + else sum(manager.peak_managed_device_weight_bytes() for manager in managers) ) + current_inactive_device_bytes = managed_weight_bytes if current_permanent else 0 current_pinned_bytes = sum( manager.pinned_host_weight_bytes() for manager in managers ) - current_transfer_work = _layerwise_transfer_work_bytes( + maximum_transfer_work = _layerwise_transfer_work_bytes( managers=managers, - resident_layers=current_resident_layers, - pinned_layers=current_pinned_layers, + resident_layers=tuple(0 for _ in managers), + pinned_layers=tuple(() for _ in managers), uses_per_streamed_layer=uses_per_request, ) - resident_targets = {current_resident_layers} - if is_dit_component_name(name) and num_inference_steps > 1: - for target in range( - 0, - max(manager.num_layers for manager in managers) + 1, - ): - counts = tuple(min(target, manager.num_layers) for manager in managers) - resident_targets.add(counts) + resident_targets = set() + for target in range( + 0, + max(manager.num_layers for manager in managers) + 1, + ): + resident_targets.add( + tuple(min(target, manager.num_layers) for manager in managers) + ) for target_resident_layers in sorted(resident_targets): target_resident_bytes = sum( @@ -834,7 +930,6 @@ def collect_promotion_candidates( manager.peak_managed_device_weight_bytes(count) for manager, count in zip(managers, target_resident_layers) ) - incremental_bytes = target_resident_bytes - current_resident_bytes pin_targets = ( _layerwise_pin_targets( managers=managers, @@ -846,11 +941,6 @@ def collect_promotion_candidates( else [current_pinned_layers] ) for target_pinned_layers in pin_targets: - if ( - target_resident_layers == current_resident_layers - and target_pinned_layers == current_pinned_layers - ): - continue target_pinned_bytes = sum( sum( manager.layer_host_store_bytes().get(layer_idx, 0) @@ -870,12 +960,13 @@ def collect_promotion_candidates( target_pinned_layers=target_pinned_layers, ) candidates.append( - PromotionCandidate( + ResidencyTarget( component_name=name, - residency_mode=mode, - promoted_weight_bytes=incremental_bytes, + residency_mode=baseline_mode, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=target_resident_bytes, h2d_bytes_per_request=( - current_transfer_work - target_transfer_work + maximum_transfer_work - target_transfer_work ), target_layerwise_resident_layers=target_resident_layers, target_layerwise_pinned_layers=target_pinned_layers, @@ -887,7 +978,15 @@ def collect_promotion_candidates( active_device_delta_bytes=( target_peak_device_bytes - current_peak_device_bytes ), - inactive_device_delta_bytes=0, + inactive_device_delta_bytes=-current_inactive_device_bytes, + present_device_delta_bytes=-current_inactive_device_bytes, + current_placement=( + not current_permanent + and target_resident_layers == current_resident_layers + and target_pinned_layers == current_pinned_layers + ), + target_device_weight_bytes=target_resident_bytes, + target_pinned_host_bytes=target_pinned_bytes, ) ) @@ -920,11 +1019,12 @@ def collect_promotion_candidates( target_pinned_layers=permanent_pins, ) candidates.append( - PromotionCandidate( + ResidencyTarget( component_name=name, - residency_mode=mode, - promoted_weight_bytes=promoted_bytes - current_resident_bytes, - h2d_bytes_per_request=current_transfer_work, + residency_mode=baseline_mode, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=managed_weight_bytes, + h2d_bytes_per_request=maximum_transfer_work, target_layerwise_resident_layers=full_resident_layers, target_layerwise_pinned_layers=permanent_pins, pinned_host_delta_bytes=( @@ -933,24 +1033,34 @@ def collect_promotion_candidates( host_unpin_scratch_bytes=unpin_scratch, host_pin_scratch_bytes=pin_scratch, permanent_residency=True, - active_device_delta_bytes=max( - 0, promoted_bytes - current_peak_device_bytes + active_device_delta_bytes=( + managed_weight_bytes - current_peak_device_bytes ), # Stage-scoped resident layers are released after the use, # so the complete managed footprint is new elsewhere. - inactive_device_delta_bytes=promoted_bytes, + inactive_device_delta_bytes=( + managed_weight_bytes - current_inactive_device_bytes + ), + present_device_delta_bytes=( + managed_weight_bytes - current_inactive_device_bytes + ), + current_placement=( + current_permanent and permanent_pins == current_pinned_layers + ), + target_device_weight_bytes=managed_weight_bytes, + target_pinned_host_bytes=permanent_pinned_bytes, ) ) return candidates def rank_candidates_by_h2d_savings( - candidates: Iterable[PromotionCandidate], -) -> list[PromotionCandidate]: + candidates: Iterable[ResidencyTarget], +) -> list[ResidencyTarget]: """Biggest per-request H2D savings first; name breaks ties deterministically. - Shared by the promotion plan and the post-request residency hint so the - hint always lists components in the order auto mode would promote them. + Shared by the placement plan and the post-request residency hint so both + use the same benefit ordering. """ return sorted( candidates, @@ -966,9 +1076,14 @@ def _skip_plan(reason: str) -> AutoResidencyPlan: return AutoResidencyPlan(skip_reason=reason) -def _vram_reserve_bytes(budget_bytes: int) -> int: +def _vram_reserve_bytes(budget_bytes: int, *, target_workload_measured: bool) -> int: + reserve_fraction = ( + MEASURED_VRAM_RESERVE_FRACTION + if target_workload_measured + else EXTRAPOLATED_VRAM_RESERVE_FRACTION + ) return max( - int(budget_bytes * VRAM_RESERVE_FRACTION), + int(budget_bytes * reserve_fraction), min( MIN_VRAM_RESERVE_BYTES, int(budget_bytes * MAX_VRAM_RESERVE_FRACTION), @@ -979,7 +1094,7 @@ def _vram_reserve_bytes(budget_bytes: int) -> int: def _binding_phase_constraints( report: RankResidencyReport, candidate_component_names: set[str], -) -> list[tuple[str, int, tuple[str, ...]]]: +) -> list[tuple[str, int, tuple[str, ...], tuple[str, ...]]]: """Keep the binding measured and conservative unobserved phases. Every candidate has the same device delta in phases with the same active @@ -1001,23 +1116,30 @@ def _binding_phase_constraints( if report.estimated_peak_bytes is not None else max(phase_peaks.values()) ) - binding: dict[tuple[str, ...], tuple[str, int]] = {} + binding: dict[tuple[tuple[str, ...], tuple[str, ...]], tuple[str, int]] = {} for phase_name, phase_peak in phase_peaks.items(): - active = tuple(sorted(report.active_components_by_phase.get(phase_name, ()))) - current = binding.get(active) + present = tuple( + sorted( + set(report.active_components_by_phase.get(phase_name, ())) + & candidate_component_names + ) + ) + used = tuple(sorted(report.used_components_by_phase.get(phase_name, ()))) + layout = (present, used) + current = binding.get(layout) if current is None or (phase_peak, phase_name) > (current[1], current[0]): - binding[active] = (phase_name, phase_peak) + binding[layout] = (phase_name, phase_peak) observed_components = { - component_name for active in binding for component_name in active + component_name for (_, used) in binding for component_name in used } for component_name in sorted(candidate_component_names - observed_components): - binding[(component_name,)] = ( + binding[((component_name,), (component_name,))] = ( f"unobserved:{component_name}", request_peak, ) return [ - (f"gpu:rank{report.rank}:{phase_name}", phase_peak, active) - for active, (phase_name, phase_peak) in sorted( + (f"gpu:rank{report.rank}:{phase_name}", phase_peak, present, used) + for (present, used), (phase_name, phase_peak) in sorted( binding.items(), key=lambda item: item[1][0] ) ] @@ -1025,10 +1147,10 @@ def _binding_phase_constraints( def _consensus_candidates( reports: list[RankResidencyReport], -) -> list[PromotionCandidate]: +) -> list[ResidencyTarget]: """Merge per-rank candidates: keep components every rank agrees on. - Sizes are worst-cased with the per-rank maximum so a promotion never + Sizes are worst-cased with the per-rank maximum so an adjustment never fits on one rank but overflows another. """ per_rank_maps = [ @@ -1053,19 +1175,26 @@ def _consensus_candidates( candidate.target_layerwise_pinned_layers for candidate in rank_candidates } permanent = {candidate.permanent_residency for candidate in rank_candidates} + target_modes = { + candidate.target_residency_mode for candidate in rank_candidates + } + current = {candidate.current_placement for candidate in rank_candidates} if ( len(component_names) != 1 or len(targets) != 1 or len(pinned_targets) != 1 or len(permanent) != 1 + or len(target_modes) != 1 + or len(current) != 1 ): continue merged.append( - PromotionCandidate( + ResidencyTarget( component_name=component_names.pop(), residency_mode=modes.pop(), - promoted_weight_bytes=max( - candidate.promoted_weight_bytes for candidate in rank_candidates + target_resident_weight_bytes=max( + candidate.target_resident_weight_bytes + for candidate in rank_candidates ), h2d_bytes_per_request=min( candidate.h2d_bytes_per_request for candidate in rank_candidates @@ -1085,17 +1214,30 @@ def _consensus_candidates( active_device_delta_bytes=max( candidate.active_device_delta_bytes for candidate in rank_candidates ), + present_device_delta_bytes=max( + candidate.present_device_delta_bytes + for candidate in rank_candidates + ), inactive_device_delta_bytes=max( candidate.inactive_device_delta_bytes for candidate in rank_candidates ), + target_residency_mode=target_modes.pop(), + current_placement=current.pop(), + target_device_weight_bytes=max( + candidate.target_device_weight_bytes + for candidate in rank_candidates + ), + target_pinned_host_bytes=max( + candidate.target_pinned_host_bytes for candidate in rank_candidates + ), ) ) return merged def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyPlan: - """Turn the gathered rank reports into one deterministic promotion plan.""" + """Turn gathered rank reports into one deterministic placement plan.""" if not reports: return _skip_plan("no rank reports") for report in reports: @@ -1106,14 +1248,24 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP estimated_peak = max(report.estimated_peak_bytes for report in reports) budget = min(report.budget_bytes for report in reports) + # The warmup request is executed collectively. Some non-output ranks may + # not retain the effective request shape in their local Req, but one rank's + # covering measurement proves that the replica executed the target shape. + target_workload_measured = any( + report.target_workload_measured for report in reports + ) reserves_by_rank = { - report.rank: _vram_reserve_bytes(report.budget_bytes) for report in reports + report.rank: _vram_reserve_bytes( + report.budget_bytes, + target_workload_measured=target_workload_measured, + ) + for report in reports } reserve = max(reserves_by_rank.values()) candidates = _consensus_candidates(reports) if not candidates: - return _skip_plan("no implicitly offloaded components to promote") + return _skip_plan("no eligible residency alternatives") candidate_component_names = {candidate.component_name for candidate in candidates} phase_constraints = { @@ -1122,9 +1274,13 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP } resource_budgets: dict[str, int] = {} for report in reports: - for resource_name, phase_peak, _ in phase_constraints[report.rank]: - resource_budgets[resource_name] = ( - report.budget_bytes - phase_peak - reserves_by_rank[report.rank] + for resource_name, phase_peak, _, _ in phase_constraints[report.rank]: + # The measured placement has already completed warmup. A negative + # reserve headroom therefore means "do not grow this phase", not + # that the current zero-delta placement is infeasible. + resource_budgets[resource_name] = max( + 0, + report.budget_bytes - phase_peak - reserves_by_rank[report.rank], ) has_hostpin_options = any( candidate.pinned_host_delta_bytes != 0 @@ -1135,8 +1291,9 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP if has_hostpin_options: for report in reports: prefix = f"node{report.node_rank}:rank{report.rank}" - resource_budgets[f"hostpin:{prefix}"] = ( - report.host_pin_capacity_bytes - report.pinned_host_bytes + resource_budgets[f"hostpin:{prefix}"] = max( + 0, + report.host_pin_capacity_bytes - report.pinned_host_bytes, ) resource_budgets[f"hostram:{prefix}:unpin"] = ( report.host_transition_headroom_bytes @@ -1162,12 +1319,17 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP and report.candidate_latency_savings_ns ] use_latency_utility = bool(timed_reports) - placement_latency_risk_ns = ( + latency_equivalence_ns = ( max( - MIN_PLACEMENT_LATENCY_RISK_NS, - int( - max(report.estimated_request_duration_ns for report in timed_reports) - * PLACEMENT_LATENCY_RISK_FRACTION + MIN_LATENCY_EQUIVALENCE_NS, + min( + MAX_LATENCY_EQUIVALENCE_NS, + int( + max( + report.estimated_request_duration_ns for report in timed_reports + ) + * LATENCY_EQUIVALENCE_FRACTION + ), ), ) if use_latency_utility @@ -1178,13 +1340,18 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP resource_deltas: dict[str, int] = {} for report, rank_candidates in zip(reports, report_candidates): rank_candidate = rank_candidates[candidate.option_key()] - for resource_name, _, active_components in phase_constraints[report.rank]: - component_is_active = candidate.component_name in active_components - phase_cost = ( - rank_candidate.active_device_delta_bytes - if component_is_active - else rank_candidate.inactive_device_delta_bytes - ) + for ( + resource_name, + _, + present_components, + used_components, + ) in phase_constraints[report.rank]: + if candidate.component_name in used_components: + phase_cost = rank_candidate.active_device_delta_bytes + elif candidate.component_name in present_components: + phase_cost = rank_candidate.present_device_delta_bytes + else: + phase_cost = rank_candidate.inactive_device_delta_bytes resource_deltas[resource_name] = phase_cost if has_hostpin_options: prefix = f"node{report.node_rank}:rank{report.rank}" @@ -1211,57 +1378,78 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP ) for report in timed_reports ) - - placement_latency_risk_ns if use_latency_utility else candidate.h2d_bytes_per_request ), + placement_cost_bytes=( + candidate.target_device_weight_bytes + or max(0, candidate.target_resident_weight_bytes), + candidate.target_pinned_host_bytes + or max(0, candidate.pinned_host_delta_bytes), + ), ) ) + candidates_by_component: dict[str, list[ResidencyTarget]] = {} + for candidate in candidates: + candidates_by_component.setdefault(candidate.component_name, []).append( + candidate + ) + complete_state_frontier = all( + any(candidate.current_placement for candidate in component_candidates) + for component_candidates in candidates_by_component.values() + ) try: placement = optimize_placement( options, resource_budget_bytes=resource_budgets, + estimated_latency_tolerance=( + latency_equivalence_ns if use_latency_utility else 0 + ), + require_selection_from_every_group=complete_state_frontier, ) except NoFeasiblePlacementError as error: return _skip_plan(str(error)) - promotions = rank_candidates_by_h2d_savings( - candidate_by_key[selection.option_key] for selection in placement.selections - ) + changed_candidates = [] + for selection in placement.selections: + candidate = candidate_by_key[selection.option_key] + if not candidate.current_placement: + changed_candidates.append(candidate) + changes = rank_candidates_by_h2d_savings(changed_candidates) return AutoResidencyPlan( estimated_peak_bytes=estimated_peak, reserve_bytes=reserve, budget_bytes=budget, resource_budget_bytes=resource_budgets, - promotions=promotions, + changes=changes, ) -def apply_promotions( +def apply_residency_changes( *, plan: AutoResidencyPlan, modules: Mapping[str, object], server_args: ServerArgs, -) -> list[AppliedPromotion]: - """Flip the planned components to resident on this rank. +) -> list[AppliedResidencyChange]: + """Apply complete target states transactionally on this rank. - Component-offload modules are only re-marked resident; the physical move - happens on their next use (the post-promotion re-warmup) through - ``ResidentStrategy``, which also applies any per-use target dtype. - Layerwise options either retain a planned number of layers during the - denoise stage or load the whole component and drop its hooks. + Component-offload targets move back to CPU immediately; resident targets + are materialized on their next use by ``ResidentStrategy``. Layerwise + targets may change resident counts and HostPin placement or disable the + streaming hooks entirely. Because targets can move in either direction, + a later calibration round can replace an earlier placement. - Raises on failure after undoing the promotions already applied, so a + Raises on failure after undoing the changes already applied, so a caller observing an exception knows this rank is back on the original strategy -- unless the undo itself fails, which raises ``AutoResidencyRollbackError`` (the rank is in a mixed state and startup must abort). """ - applied: list[AppliedPromotion] = [] + applied: list[AppliedResidencyChange] = [] try: - ordered_promotions = sorted( - plan.promotions, + ordered_changes = sorted( + plan.changes, key=lambda candidate: ( max( candidate.active_device_delta_bytes, @@ -1271,11 +1459,11 @@ def apply_promotions( candidate.option_key(), ), ) - promotion_modules: dict[str, nn.Module] = {} - snapshots: dict[str, AppliedPromotion] = {} + target_modules: dict[str, nn.Module] = {} + snapshots: dict[str, AppliedResidencyChange] = {} previous_pins_by_component: dict[str, tuple[tuple[int, ...], ...]] = {} - for candidate in ordered_promotions: - if candidate.component_name in promotion_modules: + for candidate in ordered_changes: + if candidate.component_name in target_modules: raise RuntimeError( f"multiple placement options selected for " f"{candidate.component_name!r}" @@ -1283,27 +1471,31 @@ def apply_promotions( module = modules.get(candidate.component_name) if not isinstance(module, nn.Module): raise RuntimeError( - f"promotion target {candidate.component_name!r} is missing" + f"residency target {candidate.component_name!r} is missing" ) if candidate.residency_mode == LAYERWISE_OFFLOAD and not isinstance( module, LayerwiseOffloadableModuleMixin ): raise RuntimeError( - f"promotion target {candidate.component_name!r} lost its " + f"residency target {candidate.component_name!r} lost its " "layerwise offload capability between planning and apply" ) if candidate.residency_mode == COMPONENT_OFFLOAD and ( is_layerwise_offloaded_module(module) ): raise RuntimeError( - f"promotion target {candidate.component_name!r} became " + f"residency target {candidate.component_name!r} became " "layerwise-managed between planning and apply" ) - promotion_modules[candidate.component_name] = module - if candidate.target_layerwise_resident_layers is not None: + target_modules[candidate.component_name] = module + previous_required_resident = ( + server_args.component_residency_requirement(candidate.component_name) + == AUTO_RESIDENCY_FEATURE_NAME + ) + if candidate.residency_mode == LAYERWISE_OFFLOAD: if not isinstance(module, LayerwiseOffloadableModuleMixin): raise RuntimeError( - f"partial promotion target {candidate.component_name!r} " + f"partial residency target {candidate.component_name!r} " "lost its layerwise offload capability" ) previous = tuple( @@ -1311,13 +1503,22 @@ def apply_promotions( for manager in module.layerwise_offload_managers ) previous_pinned = module.layerwise_pinned_layers() + enabled = { + manager.enabled for manager in module.layerwise_offload_managers + } + if len(enabled) != 1: + raise RuntimeError( + f"placement target {candidate.component_name!r} has mixed " + "layerwise manager states" + ) previous_pins_by_component[candidate.component_name] = previous_pinned - snapshots[candidate.component_name] = AppliedPromotion( + snapshots[candidate.component_name] = AppliedResidencyChange( component_name=candidate.component_name, residency_mode=candidate.residency_mode, previous_layerwise_resident_layers=previous, previous_layerwise_pinned_layers=previous_pinned, - layerwise_offload_disabled=candidate.permanent_residency, + previous_layerwise_offload_enabled=enabled.pop(), + previous_required_resident=previous_required_resident, pinned_host_changed=( candidate.target_layerwise_pinned_layers != previous_pinned ), @@ -1329,13 +1530,14 @@ def apply_promotions( ) continue - snapshots[candidate.component_name] = AppliedPromotion( + snapshots[candidate.component_name] = AppliedResidencyChange( component_name=candidate.component_name, residency_mode=candidate.residency_mode, + previous_required_resident=previous_required_resident, ) pinning_changes = [ candidate - for candidate in ordered_promotions + for candidate in ordered_changes if candidate.target_layerwise_pinned_layers is not None and candidate.target_layerwise_pinned_layers != previous_pins_by_component[candidate.component_name] @@ -1365,126 +1567,161 @@ def apply_promotions( for candidate in pinning_changes: applied.append(snapshots[candidate.component_name]) applied_names.add(candidate.component_name) - module = promotion_modules[candidate.component_name] + module = target_modules[candidate.component_name] assert isinstance(module, LayerwiseOffloadableModuleMixin) module.release_layerwise_pins_outside( candidate.target_layerwise_pinned_layers or () ) if pinning_changes: release_unused_pinned_memory() - for candidate in ordered_promotions: + for candidate in ordered_changes: if candidate.component_name not in applied_names: applied.append(snapshots[candidate.component_name]) applied_names.add(candidate.component_name) - module = promotion_modules[candidate.component_name] - if candidate.target_layerwise_resident_layers is not None: + module = target_modules[candidate.component_name] + target_mode = candidate.target_mode() + if candidate.residency_mode == LAYERWISE_OFFLOAD: assert isinstance(module, LayerwiseOffloadableModuleMixin) module.set_layerwise_pinned_layers( candidate.target_layerwise_pinned_layers or () ) - if candidate.permanent_residency: + if target_mode == RESIDENT: server_args.require_component_resident( candidate.component_name, feature_name=AUTO_RESIDENCY_FEATURE_NAME, ) - module.disable_offload() + if is_layerwise_offloaded_module(module): + module.disable_offload() else: + if not is_layerwise_offloaded_module(module): + module.enable_offload() + server_args.release_required_component_residency( + candidate.component_name, + feature_name=AUTO_RESIDENCY_FEATURE_NAME, + ) module.set_layerwise_resident_layer_counts( candidate.target_layerwise_resident_layers ) continue - server_args.require_component_resident( - candidate.component_name, feature_name=AUTO_RESIDENCY_FEATURE_NAME - ) - if candidate.residency_mode == LAYERWISE_OFFLOAD: - assert isinstance(module, LayerwiseOffloadableModuleMixin) - module.disable_offload() + if target_mode == RESIDENT: + server_args.require_component_resident( + candidate.component_name, + feature_name=AUTO_RESIDENCY_FEATURE_NAME, + ) + elif target_mode == COMPONENT_OFFLOAD: + server_args.release_required_component_residency( + candidate.component_name, + feature_name=AUTO_RESIDENCY_FEATURE_NAME, + ) + module.to("cpu") + else: + raise RuntimeError( + f"unsupported target mode {target_mode!r} for " + f"{candidate.component_name!r}" + ) except Exception as apply_error: try: - rollback_promotions( + rollback_residency_changes( applied=applied, modules=modules, server_args=server_args ) except Exception as rollback_error: raise AutoResidencyRollbackError( - f"promotion failed ({describe_error(apply_error)}) and rollback " + f"residency adjustment failed ({describe_error(apply_error)}) and rollback " f"failed ({describe_error(rollback_error)})" ) from apply_error raise return applied -def rollback_promotions( +def rollback_residency_changes( *, - applied: Iterable[AppliedPromotion], + applied: Iterable[AppliedResidencyChange], modules: Mapping[str, object], server_args: ServerArgs, ) -> None: - """Restore the original residency for previously applied promotions. + """Restore the exact placement that preceded a target-state change. - Every promotion is undone even when one of them fails; the collected + Every change is undone even when one of them fails; the collected failures are re-raised at the end so one broken component cannot leave - the later (earlier-applied) ones promoted. + the later (earlier-applied) ones adjusted. """ applied = list(applied) errors: list[str] = [] pinning_changes = [ - promotion - for promotion in applied - if promotion.pinned_host_changed - and promotion.previous_layerwise_pinned_layers is not None + adjustment + for adjustment in applied + if adjustment.pinned_host_changed + and adjustment.previous_layerwise_pinned_layers is not None ] - for promotion in pinning_changes: + for adjustment in pinning_changes: try: - module = modules.get(promotion.component_name) + module = modules.get(adjustment.component_name) if not isinstance(module, LayerwiseOffloadableModuleMixin): raise RuntimeError("lost layerwise offload capability") module.release_layerwise_pins_outside( - promotion.previous_layerwise_pinned_layers or () + adjustment.previous_layerwise_pinned_layers or () ) except Exception as e: - errors.append(f"{promotion.component_name}: {describe_error(e)}") + errors.append(f"{adjustment.component_name}: {describe_error(e)}") if pinning_changes: release_unused_pinned_memory() - for promotion in reversed(applied): + for adjustment in reversed(applied): try: - module = modules.get(promotion.component_name) - if promotion.previous_layerwise_resident_layers is not None: + module = modules.get(adjustment.component_name) + if adjustment.previous_layerwise_resident_layers is not None: if not isinstance(module, LayerwiseOffloadableModuleMixin): raise RuntimeError("lost layerwise offload capability") - if promotion.layerwise_offload_disabled: + currently_enabled = { + manager.enabled for manager in module.layerwise_offload_managers + } + if len(currently_enabled) != 1: + raise RuntimeError("layerwise managers have mixed enabled states") + current_enabled = currently_enabled.pop() + previous_enabled = adjustment.previous_layerwise_offload_enabled + if previous_enabled and not current_enabled: module.enable_offload() + elif not previous_enabled and current_enabled: + module.disable_offload() + if adjustment.previous_required_resident: + server_args.require_component_resident( + adjustment.component_name, + feature_name=AUTO_RESIDENCY_FEATURE_NAME, + ) + else: server_args.release_required_component_residency( - promotion.component_name, + adjustment.component_name, feature_name=AUTO_RESIDENCY_FEATURE_NAME, ) module.restore_layerwise_resident_layers( - promotion.previous_layerwise_resident_layers + adjustment.previous_layerwise_resident_layers ) - if promotion.previous_layerwise_pinned_layers is None: + if adjustment.previous_layerwise_pinned_layers is None: raise RuntimeError("lost previous layerwise host placement") module.restore_layerwise_pinned_layers( - promotion.previous_layerwise_pinned_layers + adjustment.previous_layerwise_pinned_layers ) continue - server_args.release_required_component_residency( - promotion.component_name, - feature_name=AUTO_RESIDENCY_FEATURE_NAME, - ) + if adjustment.previous_required_resident: + server_args.require_component_resident( + adjustment.component_name, + feature_name=AUTO_RESIDENCY_FEATURE_NAME, + ) + else: + server_args.release_required_component_residency( + adjustment.component_name, + feature_name=AUTO_RESIDENCY_FEATURE_NAME, + ) if not isinstance(module, nn.Module): continue - if promotion.residency_mode == LAYERWISE_OFFLOAD: - if not isinstance(module, LayerwiseOffloadableModuleMixin): - raise RuntimeError("lost layerwise offload capability") - module.enable_offload() - elif is_layerwise_offloaded_module(module): + if is_layerwise_offloaded_module(module): # A layerwise manager owns this module's placement (e.g. the # offload_during_compile window); moving it to CPU behind the # manager would strand its bookkeeping. pass - else: + elif not adjustment.previous_required_resident: module.to("cpu") except Exception as e: - errors.append(f"{promotion.component_name}: {describe_error(e)}") + errors.append(f"{adjustment.component_name}: {describe_error(e)}") if torch.get_device_module().is_available(): torch.get_device_module().empty_cache() if errors: @@ -1500,8 +1737,8 @@ def format_plan_summary( """One-line decision summary for the startup log.""" if plan.skip_reason is not None: return f"Auto residency: skipped ({plan.skip_reason})" - promoted = ( - ", ".join(_format_candidate_summary(candidate) for candidate in plan.promotions) + changes = ( + ", ".join(_format_candidate_summary(candidate) for candidate in plan.changes) or "none" ) measured = ", ".join( @@ -1518,7 +1755,7 @@ def format_plan_summary( f"estimated_peak={plan.estimated_peak_bytes / GIB_BYTES:.1f} GiB, " f"reserve={plan.reserve_bytes / GIB_BYTES:.1f} GiB, " f"budget={plan.budget_bytes / GIB_BYTES:.1f} GiB, " - f"promoted=[{promoted}]" + f"changes=[{changes}]" ) @@ -1529,18 +1766,23 @@ def format_applied_changes(*, plan: AutoResidencyPlan) -> str: placement, or disable the adjustment entirely with the kill switch. """ changes = "; ".join( - _format_promotion_change(candidate) for candidate in plan.promotions + _format_residency_change(candidate) for candidate in plan.changes ) - component_args = [ - f"{candidate.component_name}=resident" - for candidate in plan.promotions - if candidate.permanent_residency - ] + component_args = [] + has_auto_only_change = False + for candidate in plan.changes: + target_mode = candidate.target_mode() + if target_mode in (COMPONENT_OFFLOAD, RESIDENT): + component_args.append(f"{candidate.component_name}={target_mode}") + else: + has_auto_only_change = True equivalent = ( "--component-residency " + " ".join(component_args) if component_args - else "none (the selected partial layer/HostPin placement is auto-only)" + else "none" ) + if has_auto_only_change: + equivalent += " (partial layer/HostPin placement remains auto-only)" return ( f"Auto residency: adjusted {changes}. " f"Equivalent server args: {equivalent}. " @@ -1549,31 +1791,34 @@ def format_applied_changes(*, plan: AutoResidencyPlan) -> str: ) -def _format_promotion_change(candidate: PromotionCandidate) -> str: - if candidate.target_layerwise_resident_layers is not None: +def _format_residency_change(candidate: ResidencyTarget) -> str: + target_mode = candidate.target_mode() + if ( + target_mode == LAYERWISE_OFFLOAD + and candidate.target_layerwise_resident_layers is not None + ): pin_counts = tuple( len(indices) for indices in candidate.target_layerwise_pinned_layers or () ) - scope = "permanent" if candidate.permanent_residency else "stage-scoped" return ( - f"{candidate.component_name}: {scope} resident layers=" + f"{candidate.component_name}: layerwise resident layers=" f"{candidate.target_layerwise_resident_layers}, pinned layers=" f"{pin_counts}" ) - return f"{candidate.component_name}: {candidate.residency_mode} -> resident" + return ( + f"{candidate.component_name}: {candidate.residency_mode} -> " f"{target_mode}" + ) -def _format_candidate_summary(candidate: PromotionCandidate) -> str: - target = "" - if candidate.target_layerwise_resident_layers is not None: - target = ( +def _format_candidate_summary(candidate: ResidencyTarget) -> str: + target_mode = candidate.target_mode() + details = "" + if target_mode == LAYERWISE_OFFLOAD: + details = ( f", layers={candidate.target_layerwise_resident_layers}, " f"pins={tuple(len(indices) for indices in candidate.target_layerwise_pinned_layers or ())}" ) - return ( - f"{candidate.component_name}({candidate.residency_mode}{target}, " - f"{candidate.promoted_weight_bytes / GIB_BYTES:.1f} GiB)" - ) + return f"{candidate.component_name}({target_mode}{details})" def plan_summary_payload(*, plan: AutoResidencyPlan, status: str) -> dict: @@ -1584,5 +1829,5 @@ def plan_summary_payload(*, plan: AutoResidencyPlan, status: str) -> dict: """ return { "status": status, - "promoted": [candidate.component_name for candidate in plan.promotions], + "changed": [candidate.component_name for candidate in plan.changes], } diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py index 83bb8239c183..2ee48a55436b 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py @@ -62,6 +62,7 @@ class ResidencyState: class WarmupPhasePeak: active_components: tuple[str, ...] allocated_bytes: int + used_components: tuple[str, ...] = () class ResidencyBatch(Protocol): @@ -138,6 +139,7 @@ def __init__( self._track_warmup_memory = False self._warmup_phase_key: str | None = None self._warmup_phase_components: tuple[str, ...] = () + self._warmup_phase_used_components: tuple[str, ...] = () self._warmup_phase_peaks: dict[str, WarmupPhasePeak] = {} self._completed_warmup_phase_peaks: dict[str, WarmupPhasePeak] = {} @@ -225,6 +227,7 @@ def begin_request( ) self._warmup_phase_key = None self._warmup_phase_components = () + self._warmup_phase_used_components = () self._warmup_phase_peaks = {} self._completed_warmup_phase_peaks = {} if self._track_warmup_memory: @@ -254,6 +257,7 @@ def before_stage( self._begin_warmup_phase( key=f"{stage_index}:{self.state.stage_name}:setup", components=self._warmup_active_components(), + used_components=(), ) def begin_stage(self) -> None: @@ -589,6 +593,7 @@ def finish_request(self) -> None: components=self._warmup_active_components( (use,) if will_prepare else () ), + used_components=(component_name,) if will_prepare else (), ) was_on_supported_device = self._module_on_supported_device(module) strategy.finish_request(module, use, self.state, preferred=preferred) @@ -600,16 +605,24 @@ def finish_request(self) -> None: self._warmup_phase_peaks["idle"] = WarmupPhasePeak( active_components=self._warmup_active_components(), allocated_bytes=int(torch.get_device_module().memory_allocated()), + used_components=(), ) self._completed_warmup_phase_peaks = dict(self._warmup_phase_peaks) self._track_warmup_memory = False - def _begin_warmup_phase(self, *, key: str, components: tuple[str, ...]) -> None: + def _begin_warmup_phase( + self, + *, + key: str, + components: tuple[str, ...], + used_components: tuple[str, ...], + ) -> None: if not self._track_warmup_memory: return self._record_warmup_phase_peak() self._warmup_phase_key = key self._warmup_phase_components = tuple(sorted(set(components))) + self._warmup_phase_used_components = tuple(sorted(set(used_components))) torch.get_device_module().reset_peak_memory_stats() def _begin_warmup_transition( @@ -627,6 +640,9 @@ def _begin_warmup_transition( components=self._warmup_active_components( tuple(use for use in (previous, upcoming) if use is not None) ), + used_components=tuple( + use.component_name for use in (previous, upcoming) if use is not None + ), ) def _begin_warmup_use(self, use: ComponentUse) -> None: @@ -634,12 +650,14 @@ def _begin_warmup_use(self, use: ComponentUse) -> None: self._begin_warmup_phase( key=f"{self.state.stage_index}:{self.state.stage_name}:use:{phase}", components=self._warmup_active_components((use,)), + used_components=(use.component_name,), ) def _begin_warmup_between_uses(self) -> None: self._begin_warmup_phase( key=f"{self.state.stage_index}:{self.state.stage_name}:between", components=self._warmup_active_components(), + used_components=(), ) def _begin_warmup_prefetch(self, use: ComponentUse) -> None: @@ -647,6 +665,7 @@ def _begin_warmup_prefetch(self, use: ComponentUse) -> None: self._begin_warmup_phase( key=f"{self.state.stage_index}:{self.state.stage_name}:prefetch:{phase}", components=self._warmup_active_components((use,)), + used_components=(use.component_name,), ) def _warmup_active_components( @@ -668,6 +687,7 @@ def _record_warmup_phase_peak(self) -> None: peak = WarmupPhasePeak( active_components=self._warmup_phase_components, allocated_bytes=int(torch.get_device_module().max_memory_allocated()), + used_components=self._warmup_phase_used_components, ) previous = self._warmup_phase_peaks.get(self._warmup_phase_key) if previous is None: @@ -680,6 +700,9 @@ def _record_warmup_phase_peak(self) -> None: ) ), allocated_bytes=max(previous.allocated_bytes, peak.allocated_bytes), + used_components=tuple( + sorted(set(previous.used_components) & set(peak.used_components)) + ), ) def take_warmup_phase_peaks( diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py index 91ec6ee9ce73..f7bc77c0da0c 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py @@ -34,20 +34,29 @@ class PlacementOption(msgspec.Struct, frozen=True): option_key: str resource_delta_bytes: dict[str, int] estimated_latency_savings: int + # Lexicographic soft cost used only after latency-equivalent placements + # have been identified. Resource budgets remain the hard safety limits. + placement_cost_bytes: tuple[int, ...] = () class PlacementPlan(msgspec.Struct, frozen=True): resource_delta_bytes: dict[str, int] = {} estimated_latency_savings: int = 0 + placement_cost_bytes: tuple[int, ...] = () selections: list[PlacementOption] = [] -# Resource deltas, utility, selected options. -_State = tuple[tuple[int, ...], int, tuple[PlacementOption, ...]] +# Resource deltas, utility, soft placement cost, selected options. +_State = tuple[ + tuple[int, ...], + int, + tuple[int, ...], + tuple[PlacementOption, ...], +] def _selection_key(state: _State) -> tuple[str, ...]: - return tuple(option.option_key for option in state[2]) + return tuple(option.option_key for option in state[3]) def _better_same_cost(candidate: _State, incumbent: _State) -> bool: @@ -58,24 +67,28 @@ def _better_same_cost(candidate: _State, incumbent: _State) -> bool: def _dominates(candidate: _State, other: _State) -> bool: """Whether candidate is at least as useful and no more expensive.""" - return candidate[1] >= other[1] and all( - left <= right for left, right in zip(candidate[0], other[0]) + return ( + candidate[1] >= other[1] + and all(left <= right for left, right in zip(candidate[0], other[0])) + and all(left <= right for left, right in zip(candidate[2], other[2])) ) def _pareto_prune(states: Iterable[_State]) -> list[_State]: """Drop states dominated across every constrained resource.""" - best_at_cost: dict[tuple[int, ...], _State] = {} + best_at_cost: dict[tuple[tuple[int, ...], tuple[int, ...]], _State] = {} for state in states: - incumbent = best_at_cost.get(state[0]) + cost_key = (state[0], state[2]) + incumbent = best_at_cost.get(cost_key) if incumbent is None or _better_same_cost(state, incumbent): - best_at_cost[state[0]] = state + best_at_cost[cost_key] = state ordered = sorted( best_at_cost.values(), key=lambda state: ( sum(state[0]), state[0], + state[2], -state[1], _selection_key(state), ), @@ -95,10 +108,15 @@ def optimize_placement( options: Iterable[PlacementOption], *, resource_budget_bytes: Mapping[str, int], + estimated_latency_tolerance: int = 0, + require_selection_from_every_group: bool = False, ) -> PlacementPlan: """Choose at most one alternative per group under every resource budget. Not selecting an alternative keeps that group at its current placement. + Callers that provide complete target states may instead require exactly + one option per group, allowing the solver to replace or demote a current + placement rather than only add to it. All alternatives are evaluated with the same selection vector, so every constraint describing this placement state sees the same decision. A choice may trade one resource for another by using a negative delta. The @@ -107,6 +125,21 @@ def optimize_placement( option, but the result does not imply that the supplied frontier contains every physically possible lifecycle plan. """ + if estimated_latency_tolerance < 0: + raise ValueError("estimated_latency_tolerance must be non-negative") + + options = list(options) + placement_cost_dimensions = max( + (len(option.placement_cost_bytes) for option in options), default=0 + ) + for option in options: + if option.placement_cost_bytes and ( + len(option.placement_cost_bytes) != placement_cost_dimensions + ): + raise ValueError( + "all non-empty placement costs must have the same dimensions" + ) + grouped: dict[str, list[PlacementOption]] = {} option_keys: set[str] = set() all_resource_names = tuple(sorted(resource_budget_bytes)) @@ -149,7 +182,16 @@ def optimize_placement( ordered_groups = [] for group_key in sorted(grouped): group_options = sorted(grouped[group_key], key=lambda option: option.option_key) - local_states = [((0,) * len(resource_names), 0, ())] + local_states = [] + if not require_selection_from_every_group: + local_states.append( + ( + (0,) * len(resource_names), + 0, + (0,) * placement_cost_dimensions, + (), + ) + ) local_states.extend( ( tuple( @@ -157,12 +199,13 @@ def optimize_placement( for resource_name in resource_names ), option.estimated_latency_savings, + option.placement_cost_bytes or (0,) * placement_cost_dimensions, (option,), ) for option in group_options ) ordered_groups.append( - [state[2][0] for state in _pareto_prune(local_states) if state[2]] + [state[3][0] for state in _pareto_prune(local_states) if state[3]] ) minimum_remaining: list[tuple[int, ...]] = [ (0,) * len(resource_names) for _ in range(len(ordered_groups) + 1) @@ -172,19 +215,33 @@ def optimize_placement( next_minimum = minimum_remaining[group_index + 1] minimum_remaining[group_index] = tuple( next_delta - + min( - 0, - *( + + ( + min( option.resource_delta_bytes.get(resource_name, 0) for option in group_options - ), + ) + if require_selection_from_every_group + else min( + 0, + *( + option.resource_delta_bytes.get(resource_name, 0) + for option in group_options + ), + ) ) for next_delta, resource_name in zip(next_minimum, resource_names) ) - frontier: list[_State] = [((0,) * len(resource_names), 0, ())] + frontier: list[_State] = [ + ( + (0,) * len(resource_names), + 0, + (0,) * placement_cost_dimensions, + (), + ) + ] for group_index, group_options in enumerate(ordered_groups): - candidates = list(frontier) + candidates = [] if require_selection_from_every_group else list(frontier) for state in frontier: for option in group_options: resource_deltas = tuple( @@ -195,7 +252,15 @@ def optimize_placement( ( resource_deltas, state[1] + option.estimated_latency_savings, - state[2] + (option,), + tuple( + current + delta + for current, delta in zip( + state[2], + option.placement_cost_bytes + or (0,) * placement_cost_dimensions, + ) + ), + state[3] + (option,), ) ) remaining = minimum_remaining[group_index + 1] @@ -219,23 +284,31 @@ def optimize_placement( if not frontier: raise NoFeasiblePlacementError("no placement satisfies all resource budgets") + best_utility = max(state[1] for state in frontier) + latency_equivalent = [ + state + for state in frontier + if state[1] >= best_utility - estimated_latency_tolerance + ] best = min( - frontier, + latency_equivalent, key=lambda state: ( - -state[1], + state[2], sum(state[0]), state[0], + -state[1], _selection_key(state), ), ) full_resource_deltas = { resource_name: sum( - option.resource_delta_bytes.get(resource_name, 0) for option in best[2] + option.resource_delta_bytes.get(resource_name, 0) for option in best[3] ) for resource_name in all_resource_names } return PlacementPlan( resource_delta_bytes=full_resource_deltas, estimated_latency_savings=best[1], - selections=list(best[2]), + placement_cost_bytes=best[2], + selections=list(best[3]), ) diff --git a/python/sglang/multimodal_gen/runtime/scheduler_client.py b/python/sglang/multimodal_gen/runtime/scheduler_client.py index 5d2c95617e9c..052b35ef2e51 100644 --- a/python/sglang/multimodal_gen/runtime/scheduler_client.py +++ b/python/sglang/multimodal_gen/runtime/scheduler_client.py @@ -160,7 +160,7 @@ def _merge_auto_residency_results(results: list[Any]) -> Any: """Preserve the strongest replica-wide auto-residency outcome. Replicas plan against their own measured headroom, so one may keep the - original placement while another promotes. If any replica promoted, the + original placement while another adjusts. If any replica changed, the orchestrator must still run the validation warmup everywhere. A failed local rollback remains fatal and takes precedence over that success. """ @@ -175,11 +175,11 @@ def _merge_auto_residency_results(results: list[Any]) -> Any: if rollback_failed is not None: return rollback_failed - promoted = next( - (result for result in results if _auto_residency_status(result) == "promoted"), + adjusted = next( + (result for result in results if _auto_residency_status(result) == "adjusted"), None, ) - if promoted is not None: + if adjusted is not None: replica_errors = [ result.error for result in results @@ -191,7 +191,7 @@ def _merge_auto_residency_results(results: list[Any]) -> Any: len(replica_errors), replica_errors[0], ) - return promoted + return adjusted rolled_back = next( ( 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 4ac80edffcef..6ce88280f84d 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -1494,6 +1494,19 @@ def residency_mode(self, component_name: str) -> str: if component_name in self._required_resident_components: return RESIDENT + return self.configured_residency_mode(component_name) + + def configured_residency_mode(self, component_name: str) -> str: + """Return placement from user controls and startup auto tuning. + + Unlike :meth:`residency_mode`, this excludes runtime requirements such + as warmup-calibrated auto residency. A runtime planner can therefore + recover and reconsider the load-safe baseline without overriding an + explicit user choice. + """ + if current_platform.is_cpu(): + return RESIDENT + explicit_mode = self.explicit_residency_mode(component_name) if explicit_mode is not None: return explicit_mode @@ -1529,6 +1542,10 @@ def residency_mode(self, component_name: str) -> str: return COMPONENT_OFFLOAD if self.vae_cpu_offload else RESIDENT return RESIDENT + def component_residency_requirement(self, component_name: str) -> str | None: + """Return the feature that currently hard-requires GPU residency.""" + return self._required_resident_components.get(component_name) + def should_cpu_offload_component(self, component_name: str) -> bool: return self.residency_mode(component_name) == COMPONENT_OFFLOAD diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index 15ddcd6b5fa5..d15fa1141184 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -101,9 +101,9 @@ def should_run_explicit_client_warmup(server_args: ServerArgs) -> bool: def auto_residency_skip_reason(server_args: ServerArgs) -> str | None: - """Final gate for warmup-calibrated residency promotion. + """Final gate for warmup-calibrated residency placement. - Only rules out paths the promotion was not designed for; the workers + Only rules out paths the planner was not designed for; the workers re-check per component (explicit placement, FSDP modules, custom strategies, missing sizes) and per measurement. """ @@ -125,10 +125,10 @@ async def maybe_apply_auto_residency( server_args: ServerArgs, forward: Callable[[Req], Awaitable[OutputBatch]], ) -> None: - """Promote implicitly offloaded components after warmup, then re-warm. + """Adjust implicit component residency after warmup, then re-warm. Runs between the synthetic warmup and the ready signal, so the residency - is frozen before ``/health`` turns 200. If a promotion (or its calibration) + is frozen before ``/health`` turns 200. If an adjustment or its calibration fails, the workers roll back that round and retain the last calibrated placement; only a failed rollback raises and aborts startup. """ @@ -136,9 +136,9 @@ async def maybe_apply_auto_residency( AutoResidencyReq, ) from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( - PROMOTION_STATUS_PROMOTED, - PROMOTION_STATUS_ROLLBACK_FAILED, - PROMOTION_STATUS_ROLLED_BACK, + PLACEMENT_STATUS_ADJUSTED, + PLACEMENT_STATUS_ROLLBACK_FAILED, + PLACEMENT_STATUS_ROLLED_BACK, ) skip_reason = auto_residency_skip_reason(server_args) @@ -156,17 +156,17 @@ async def maybe_apply_auto_residency( # Same fail-open contract as the warmup itself: implicit warmups must # never abort startup, explicit --warmup-resolutions ones must succeed. fail_open = server_args.warmup_resolutions is None - promoted_any = False + adjusted_any = False async def rollback_and_rewarm(error: Exception) -> None: logger.warning( - "Post-promotion calibration failed (%s); rolling back auto residency", + "Post-adjustment calibration failed (%s); rolling back auto residency", error, ) rollback = await forward(AutoResidencyReq(action="rollback")) if ( rollback.error is not None - or _auto_residency_status(rollback) != PROMOTION_STATUS_ROLLED_BACK + or _auto_residency_status(rollback) != PLACEMENT_STATUS_ROLLED_BACK ): raise RuntimeError( f"auto residency rollback failed: {rollback.error}" @@ -181,14 +181,14 @@ async def rollback_and_rewarm(error: Exception) -> None: try: response = await forward(AutoResidencyReq(action="apply")) except Exception as e: - if promoted_any: + if adjusted_any: # The worker-side apply is transactional, but an RPC failure # does not tell us whether its response was lost after commit. # Do not mistake the last validated round for an unvalidated # one and roll it back; abort startup rather than serve a # placement whose calibration state is unknown. raise RuntimeError( - "auto residency apply failed after a calibrated promotion" + "auto residency apply failed after a calibrated adjustment" ) from e if not fail_open: raise @@ -199,28 +199,28 @@ async def rollback_and_rewarm(error: Exception) -> None: ) return status = _auto_residency_status(response) - if status == PROMOTION_STATUS_ROLLBACK_FAILED: + if status == PLACEMENT_STATUS_ROLLBACK_FAILED: raise RuntimeError(f"auto residency rollback failed: {response.error}") if response.error is not None: - if status == PROMOTION_STATUS_ROLLED_BACK: + if status == PLACEMENT_STATUS_ROLLED_BACK: await run_async_client_warmup( server_args, forward, fail_open=fail_open, rewarm=True ) retained = ( "the last calibrated placement" - if promoted_any + if adjusted_any else "the original strategy" ) logger.warning( - "Auto residency promotion not applied; continuing on %s: %s", + "Auto residency adjustment not applied; continuing on %s: %s", retained, response.error, ) return - if status != PROMOTION_STATUS_PROMOTED: + if status != PLACEMENT_STATUS_ADJUSTED: return - promoted_any = True + adjusted_any = True # This pass physically realizes the selected placement and measures # phases that overlap under it. The next round can then use that new # evidence instead of permanently charging components twice. @@ -233,7 +233,7 @@ async def rollback_and_rewarm(error: Exception) -> None: return logger.warning( - "Auto residency stopped after %d promotion rounds; keeping the last " + "Auto residency stopped after %d adjustment rounds; keeping the last " "successfully calibrated placement.", MAX_AUTO_RESIDENCY_ROUNDS, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index c22d83e1cc47..531e0130715c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for warmup-calibrated auto residency promotion.""" +"""Unit tests for warmup-calibrated auto residency adjustment.""" from types import SimpleNamespace @@ -11,17 +11,18 @@ from sglang.multimodal_gen.configs.pipeline_configs.longlive2 import LongLive2T2VConfig from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( ACTIVATION_EXTRAPOLATION_MARGIN, + AUTO_RESIDENCY_FEATURE_NAME, GIB_BYTES, MIN_VRAM_RESERVE_BYTES, - AppliedPromotion, + AppliedResidencyChange, AutoResidencyPlan, AutoResidencyRollbackError, DefaultWorkload, - PromotionCandidate, RankResidencyReport, + ResidencyTarget, WarmupMemoryRecord, - apply_promotions, - collect_promotion_candidates, + apply_residency_changes, + collect_residency_targets, component_resident_size_bytes, component_runtime_weight_bytes, estimate_candidate_latency_savings_ns, @@ -32,7 +33,7 @@ plan_auto_residency, rank_candidates_by_h2d_savings, resolve_measured_default_workload, - rollback_promotions, + rollback_residency_changes, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( COMPONENT_OFFLOAD, @@ -133,7 +134,7 @@ def test_uses_steady_step_instead_of_scaling_first_step_setup(self): assert stage_ns["DenoisingStage"] == 1_020_000_000 assert request_ns == 1_120_000_000 - def test_transfer_savings_is_capped_by_own_stage(self): + def test_transfer_savings_is_capped_by_request_not_component_stage(self): candidates = [ _candidate("text_encoder", weight_gib=10, h2d_gib=10), _candidate("transformer", weight_gib=10, h2d_gib=10), @@ -143,17 +144,9 @@ def test_transfer_savings_is_capped_by_own_stage(self): savings = estimate_candidate_latency_savings_ns( candidates=candidates, request_duration_ns=10_000_000_000, - stage_duration_ns={ - "TextEncodingStage": 80_000_000, - "DenoisingStage": 5_000_000_000, - }, - component_stages={ - "text_encoder": ("TextEncodingStage",), - "transformer": ("DenoisingStage",), - }, ) - assert savings[candidates[0].option_key()] == 80_000_000 + assert savings[candidates[0].option_key()] > 400_000_000 assert savings[candidates[1].option_key()] > 400_000_000 assert savings[candidates[2].option_key()] < 0 @@ -191,7 +184,7 @@ def test_single_capped_record_scales_only_the_activation_part(self): ) assert estimate == expected # Scaling the whole peak would inflate the estimate by the resident - # weights times the cap ratio and promotion would never trigger. + # weights times the cap ratio and adjustment would never trigger. naive = int(record.peak_allocated_bytes * ratio) assert estimate < naive @@ -346,7 +339,7 @@ def test_phase_estimation_preserves_active_component_membership(self): phase_active_components={"denoise": ("transformer",)}, ) - peaks, active = estimate_workload_phase_peaks( + peaks, active, used = estimate_workload_phase_peaks( records=[small, large], target_units=832 * 480 * 81, component_weight_bytes={"transformer": 28 * GIB_BYTES}, @@ -354,6 +347,7 @@ def test_phase_estimation_preserves_active_component_membership(self): assert peaks["denoise"] >= 32 * GIB_BYTES assert active == {"denoise": ("transformer",)} + assert used == active def test_phase_estimation_uses_allocated_peak(self): record = WarmupMemoryRecord( @@ -367,7 +361,7 @@ def test_phase_estimation_uses_allocated_peak(self): phase_active_components={"denoise": ("transformer",)}, ) - peaks, _ = estimate_workload_phase_peaks( + peaks, _, _ = estimate_workload_phase_peaks( records=[record], target_units=record.workload_units(), component_weight_bytes={"transformer": 10 * GIB_BYTES}, @@ -375,6 +369,79 @@ def test_phase_estimation_uses_allocated_peak(self): assert peaks["denoise"] == 11 * GIB_BYTES + def test_phase_estimation_prefers_target_layout_over_smaller_warmup(self): + small = WarmupMemoryRecord( + width=256, + height=256, + num_frames=9, + baseline_allocated_bytes=2 * GIB_BYTES, + peak_allocated_bytes=30 * GIB_BYTES, + succeeded=True, + phase_peak_allocated_bytes={"denoise": 30 * GIB_BYTES}, + phase_active_components={"denoise": ()}, + ) + target = WarmupMemoryRecord( + width=768, + height=512, + num_frames=25, + baseline_allocated_bytes=2 * GIB_BYTES, + peak_allocated_bytes=48 * GIB_BYTES, + succeeded=True, + phase_peak_allocated_bytes={"denoise": 48 * GIB_BYTES}, + phase_active_components={"denoise": ("transformer",)}, + ) + + peaks, active, used = estimate_workload_phase_peaks( + records=[small, target], + target_units=target.workload_units(), + component_weight_bytes={"transformer": 40 * GIB_BYTES}, + ) + + assert peaks == {"denoise": 48 * GIB_BYTES} + assert active == {"denoise": ("transformer",)} + assert used == active + + def test_phase_estimation_keeps_distinct_active_layouts_separate(self): + transformer_phase = WarmupMemoryRecord( + width=768, + height=512, + num_frames=25, + baseline_allocated_bytes=2 * GIB_BYTES, + peak_allocated_bytes=48 * GIB_BYTES, + succeeded=True, + phase_peak_allocated_bytes={"denoise": 48 * GIB_BYTES}, + phase_active_components={"denoise": ("transformer",)}, + ) + encoder_phase = WarmupMemoryRecord( + width=768, + height=512, + num_frames=25, + baseline_allocated_bytes=2 * GIB_BYTES, + peak_allocated_bytes=20 * GIB_BYTES, + succeeded=True, + phase_peak_allocated_bytes={"denoise": 20 * GIB_BYTES}, + phase_active_components={"denoise": ("text_encoder",)}, + ) + + peaks, active, used = estimate_workload_phase_peaks( + records=[transformer_phase, encoder_phase], + target_units=transformer_phase.workload_units(), + component_weight_bytes={ + "transformer": 40 * GIB_BYTES, + "text_encoder": 10 * GIB_BYTES, + }, + ) + + assert peaks == { + "denoise:layout:0": 20 * GIB_BYTES, + "denoise:layout:1": 48 * GIB_BYTES, + } + assert active == { + "denoise:layout:0": ("text_encoder",), + "denoise:layout:1": ("transformer",), + } + assert used == active + def _candidate( name: str, @@ -382,12 +449,12 @@ def _candidate( mode: str = COMPONENT_OFFLOAD, weight_gib: int = 10, h2d_gib: int | None = None, -) -> PromotionCandidate: +) -> ResidencyTarget: layerwise = mode == LAYERWISE_OFFLOAD - return PromotionCandidate( + return ResidencyTarget( component_name=name, residency_mode=mode, - promoted_weight_bytes=weight_gib * GIB_BYTES, + target_resident_weight_bytes=weight_gib * GIB_BYTES, h2d_bytes_per_request=(h2d_gib if h2d_gib is not None else weight_gib) * GIB_BYTES, target_layerwise_resident_layers=(1,) if layerwise else None, @@ -403,14 +470,16 @@ def _report( rank: int = 0, budget_gib: int = 100, estimated_gib: int | None = 50, - candidates: list[PromotionCandidate] | None = None, + candidates: list[ResidencyTarget] | None = None, skip_reason: str | None = None, phase_peaks_gib: dict[str, int] | None = None, phase_components: dict[str, tuple[str, ...]] | None = None, + phase_present_components: dict[str, tuple[str, ...]] | None = None, node_rank: int = 0, pinned_host_gib: int = 0, host_pin_capacity_gib: int = 0, host_transition_headroom_gib: int = 0, + target_workload_measured: bool = False, estimated_request_duration_ns: int = 0, candidate_latency_savings_ns: dict[str, int] | None = None, ) -> RankResidencyReport: @@ -420,10 +489,16 @@ def _report( estimated_peak_bytes=( None if estimated_gib is None else estimated_gib * GIB_BYTES ), + target_workload_measured=target_workload_measured, estimated_peak_bytes_by_phase={ name: value * GIB_BYTES for name, value in (phase_peaks_gib or {}).items() }, - active_components_by_phase=phase_components or {}, + active_components_by_phase=( + phase_present_components + if phase_present_components is not None + else phase_components or {} + ), + used_components_by_phase=phase_components or {}, node_rank=node_rank, pinned_host_bytes=pinned_host_gib * GIB_BYTES, host_pin_capacity_bytes=host_pin_capacity_gib * GIB_BYTES, @@ -495,7 +570,7 @@ def test_rank_skip_reason_propagates(self): reports=[_report(), _report(rank=1, skip_reason="no measurements")] ) assert plan.skip_reason == "rank 1: no measurements" - assert not plan.promotions + assert not plan.changes def test_missing_estimate_skips(self): plan = plan_auto_residency(reports=[_report(estimated_gib=None)]) @@ -522,6 +597,42 @@ def test_reserve_floor_is_capped_on_a_small_card(self): plan = plan_auto_residency(reports=[_report(budget_gib=10, estimated_gib=4)]) assert plan.reserve_bytes == 2 * GIB_BYTES + def test_target_workload_measurement_uses_tighter_reserve(self): + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=100, + estimated_gib=50, + target_workload_measured=True, + ) + ] + ) + assert plan.reserve_bytes == 5 * GIB_BYTES + + def test_target_workload_measurement_is_replica_wide(self): + candidate = _candidate("vae", weight_gib=1) + plan = plan_auto_residency( + reports=[ + _report( + rank=0, + budget_gib=100, + estimated_gib=50, + candidates=[candidate], + target_workload_measured=True, + ), + _report( + rank=1, + budget_gib=140, + estimated_gib=70, + candidates=[candidate], + ), + ] + ) + + assert plan.reserve_bytes == 7 * GIB_BYTES + assert plan.resource_budget_bytes["gpu:rank0:request"] == 45 * GIB_BYTES + assert plan.resource_budget_bytes["gpu:rank1:request"] == 63 * GIB_BYTES + def test_optimizes_promotion_by_h2d_savings(self): # dit saves 25 GiB x 50 steps per request; text_encoder saves 30 GiB # once. dit must be promoted first, after which the text encoder no @@ -533,7 +644,7 @@ def test_optimizes_promotion_by_h2d_savings(self): plan = plan_auto_residency( reports=[_report(budget_gib=100, estimated_gib=50, candidates=candidates)] ) - assert [c.component_name for c in plan.promotions] == ["dit"] + assert [c.component_name for c in plan.changes] == ["dit"] def test_long_request_does_not_spend_vram_on_tiny_encoder_gain(self): transformer = _candidate( @@ -559,7 +670,7 @@ def test_long_request_does_not_spend_vram_on_tiny_encoder_gain(self): ] ) - assert [candidate.component_name for candidate in plan.promotions] == [ + assert [candidate.component_name for candidate in plan.changes] == [ "transformer" ] @@ -596,7 +707,7 @@ def test_output_rank_timing_applies_to_untimed_peer_ranks(self): ] ) - assert [candidate.component_name for candidate in plan.promotions] == [ + assert [candidate.component_name for candidate in plan.changes] == [ "transformer" ] @@ -608,7 +719,7 @@ def test_smaller_component_still_fits_after_skipping_a_big_one(self): plan = plan_auto_residency( reports=[_report(budget_gib=100, estimated_gib=50, candidates=candidates)] ) - assert [c.component_name for c in plan.promotions] == ["vae"] + assert [c.component_name for c in plan.changes] == ["vae"] def test_consensus_requires_component_on_every_rank(self): plan = plan_auto_residency( @@ -620,10 +731,10 @@ def test_consensus_requires_component_on_every_rank(self): _report(rank=1, candidates=[_candidate("vae", weight_gib=12)]), ] ) - names = [c.component_name for c in plan.promotions] + names = [c.component_name for c in plan.changes] assert names == ["vae"] # worst-case size across ranks - assert plan.promotions[0].promoted_weight_bytes == 12 * GIB_BYTES + assert plan.changes[0].target_resident_weight_bytes == 12 * GIB_BYTES def test_no_candidates_skips(self): plan = plan_auto_residency(reports=[_report(candidates=[])]) @@ -646,10 +757,56 @@ def test_component_active_phase_is_not_double_counted(self): ] ) - assert [candidate.component_name for candidate in plan.promotions] == [ + assert [candidate.component_name for candidate in plan.changes] == [ "transformer" ] + def test_async_prefetch_presence_is_not_charged_twice(self): + offload = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + current_placement=True, + ) + resident = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=42 * GIB_BYTES, + h2d_bytes_per_request=42 * GIB_BYTES, + inactive_device_delta_bytes=42 * GIB_BYTES, + present_device_delta_bytes=0, + target_device_weight_bytes=42 * GIB_BYTES, + ) + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=95, + estimated_gib=50, + target_workload_measured=True, + phase_peaks_gib={ + "text_and_prefetch": 50, + "denoise": 50, + }, + phase_components={ + "text_and_prefetch": ("text_encoder",), + "denoise": ("transformer",), + }, + phase_present_components={ + "text_and_prefetch": ("text_encoder", "transformer"), + "denoise": ("transformer",), + }, + candidates=[offload, resident], + ) + ] + ) + + assert [candidate.option_key() for candidate in plan.changes] == [ + "transformer:resident" + ] + def test_other_phase_can_block_component_residency(self): plan = plan_auto_residency( reports=[ @@ -663,14 +820,14 @@ def test_other_phase_can_block_component_residency(self): ] ) - assert not plan.promotions + assert not plan.changes def test_unobserved_later_component_gets_a_conservative_phase(self): first_transformer = _candidate("first_transformer", weight_gib=26, h2d_gib=100) - second_transformer = PromotionCandidate( + second_transformer = ResidencyTarget( component_name="second_transformer", residency_mode=LAYERWISE_OFFLOAD, - promoted_weight_bytes=26 * GIB_BYTES, + target_resident_weight_bytes=26 * GIB_BYTES, h2d_bytes_per_request=90 * GIB_BYTES, target_layerwise_resident_layers=(40,), target_layerwise_pinned_layers=((),), @@ -691,7 +848,7 @@ def test_unobserved_later_component_gets_a_conservative_phase(self): ] ) - assert [candidate.component_name for candidate in plan.promotions] == [ + assert [candidate.component_name for candidate in plan.changes] == [ "first_transformer" ] assert "gpu:rank0:unobserved:second_transformer" in plan.resource_budget_bytes @@ -721,7 +878,7 @@ def test_redundant_phase_constraints_collapse_to_the_highest_peak(self): assert plan.resource_budget_bytes["gpu:rank0:denoise:later"] == (20 * GIB_BYTES) assert "gpu:rank0:decode" in plan.resource_budget_bytes - def test_negative_gpu_headroom_is_not_treated_as_zero_capacity(self): + def test_validated_baseline_with_no_reserve_headroom_accepts_zero_delta(self): plan = plan_auto_residency( reports=[ _report( @@ -734,8 +891,10 @@ def test_negative_gpu_headroom_is_not_treated_as_zero_capacity(self): ] ) - assert plan.promotions == [] - assert plan.skip_reason == "no placement satisfies all resource budgets" + assert [candidate.component_name for candidate in plan.changes] == [ + "transformer" + ] + assert plan.resource_budget_bytes["gpu:rank0:denoise"] == 0 def test_each_rank_keeps_a_reserve_scaled_to_its_own_gpu(self): candidate = _candidate("vae", weight_gib=1) @@ -765,10 +924,10 @@ def test_each_rank_keeps_a_reserve_scaled_to_its_own_gpu(self): assert plan.reserve_bytes == 14 * GIB_BYTES def test_partial_layer_residency_is_selected_when_full_dit_does_not_fit(self): - full = PromotionCandidate( + full = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, - promoted_weight_bytes=30 * GIB_BYTES, + target_resident_weight_bytes=30 * GIB_BYTES, h2d_bytes_per_request=300 * GIB_BYTES, target_layerwise_resident_layers=(30,), target_layerwise_pinned_layers=((),), @@ -776,10 +935,10 @@ def test_partial_layer_residency_is_selected_when_full_dit_does_not_fit(self): active_device_delta_bytes=30 * GIB_BYTES, inactive_device_delta_bytes=30 * GIB_BYTES, ) - partial = PromotionCandidate( + partial = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, - promoted_weight_bytes=10 * GIB_BYTES, + target_resident_weight_bytes=10 * GIB_BYTES, h2d_bytes_per_request=90 * GIB_BYTES, target_layerwise_resident_layers=(12,), target_layerwise_pinned_layers=((),), @@ -797,23 +956,23 @@ def test_partial_layer_residency_is_selected_when_full_dit_does_not_fit(self): ] ) - assert len(plan.promotions) == 1 - assert plan.promotions[0].target_layerwise_resident_layers == (12,) + assert len(plan.changes) == 1 + assert plan.changes[0].target_layerwise_resident_layers == (12,) def test_layerwise_release_can_unlock_a_more_valuable_component(self): - release_cold_layers = PromotionCandidate( + release_cold_layers = ResidencyTarget( component_name="transformer_2", residency_mode=LAYERWISE_OFFLOAD, - promoted_weight_bytes=-20 * GIB_BYTES, + target_resident_weight_bytes=-20 * GIB_BYTES, h2d_bytes_per_request=-20 * GIB_BYTES, target_layerwise_resident_layers=(0,), target_layerwise_pinned_layers=((),), active_device_delta_bytes=-20 * GIB_BYTES, ) - keep_hot_encoder = PromotionCandidate( + keep_hot_encoder = ResidencyTarget( component_name="text_encoder", residency_mode=COMPONENT_OFFLOAD, - promoted_weight_bytes=30 * GIB_BYTES, + target_resident_weight_bytes=30 * GIB_BYTES, h2d_bytes_per_request=100 * GIB_BYTES, permanent_residency=True, active_device_delta_bytes=0, @@ -835,25 +994,25 @@ def test_layerwise_release_can_unlock_a_more_valuable_component(self): ] ) - assert {candidate.component_name for candidate in plan.promotions} == { + assert {candidate.component_name for candidate in plan.changes} == { "transformer_2", "text_encoder", } def test_node_hostpin_is_optimized_with_the_same_selection_vector(self): - cold_pageable = PromotionCandidate( + cold_pageable = ResidencyTarget( component_name="cold_encoder", residency_mode=LAYERWISE_OFFLOAD, - promoted_weight_bytes=0, + target_resident_weight_bytes=0, h2d_bytes_per_request=-GIB_BYTES, target_layerwise_resident_layers=(0,), target_layerwise_pinned_layers=((),), pinned_host_delta_bytes=-10 * GIB_BYTES, ) - hot_pinned = PromotionCandidate( + hot_pinned = ResidencyTarget( component_name="hot_dit", residency_mode=LAYERWISE_OFFLOAD, - promoted_weight_bytes=0, + target_resident_weight_bytes=0, h2d_bytes_per_request=100 * GIB_BYTES, target_layerwise_resident_layers=(0,), target_layerwise_pinned_layers=((0,),), @@ -877,17 +1036,17 @@ def test_node_hostpin_is_optimized_with_the_same_selection_vector(self): # The encoder's gain is below the risk floor, but its hostpin release # makes the valuable DiT placement feasible. A joint solver must keep # this cross-component trade instead of filtering it locally. - assert [candidate.component_name for candidate in plan.promotions] == [ + assert [candidate.component_name for candidate in plan.changes] == [ "hot_dit", "cold_encoder", ] assert plan.resource_budget_bytes["hostpin:node0:rank0"] == 0 def test_dynamic_hostpin_is_constrained_for_every_worker(self): - pin_more = PromotionCandidate( + pin_more = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, - promoted_weight_bytes=0, + target_resident_weight_bytes=0, h2d_bytes_per_request=100 * GIB_BYTES, target_layerwise_resident_layers=(0,), target_layerwise_pinned_layers=((0,),), @@ -908,17 +1067,17 @@ def test_dynamic_hostpin_is_constrained_for_every_worker(self): ] ) - assert [candidate.component_name for candidate in plan.promotions] == [ + assert [candidate.component_name for candidate in plan.changes] == [ "transformer" ] assert plan.resource_budget_bytes["hostpin:node0:rank0"] == 20 * GIB_BYTES assert plan.resource_budget_bytes["hostpin:node0:rank1"] == 20 * GIB_BYTES def test_hostpin_repack_must_fit_transition_headroom(self): - pin_more = PromotionCandidate( + pin_more = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, - promoted_weight_bytes=0, + target_resident_weight_bytes=0, h2d_bytes_per_request=100 * GIB_BYTES, target_layerwise_resident_layers=(0,), target_layerwise_pinned_layers=((0,),), @@ -935,9 +1094,79 @@ def test_hostpin_repack_must_fit_transition_headroom(self): ] ) - assert plan.promotions == [] + assert plan.changes == [] assert plan.resource_budget_bytes["hostram:node0:rank0:pin"] == 5 * GIB_BYTES + def test_complete_state_frontier_can_replace_an_earlier_resident(self): + transformer_offload = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=10 * GIB_BYTES, + h2d_bytes_per_request=0, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + active_device_delta_bytes=-20 * GIB_BYTES, + target_device_weight_bytes=10 * GIB_BYTES, + ) + transformer_resident = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=30 * GIB_BYTES, + h2d_bytes_per_request=40 * GIB_BYTES, + target_layerwise_resident_layers=(20,), + target_layerwise_pinned_layers=((),), + current_placement=True, + target_device_weight_bytes=30 * GIB_BYTES, + ) + encoder_offload = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + current_placement=True, + ) + encoder_resident = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=30 * GIB_BYTES, + h2d_bytes_per_request=80 * GIB_BYTES, + inactive_device_delta_bytes=30 * GIB_BYTES, + target_device_weight_bytes=30 * GIB_BYTES, + ) + candidates = [ + transformer_offload, + transformer_resident, + encoder_offload, + encoder_resident, + ] + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=100, + estimated_gib=90, + target_workload_measured=True, + phase_peaks_gib={ + "transformer_use": 70, + "encoder_use": 90, + }, + phase_components={ + "transformer_use": ("transformer",), + "encoder_use": ("text_encoder",), + }, + candidates=candidates, + ) + ] + ) + + assert {candidate.option_key() for candidate in plan.changes} == { + "transformer:stage:layers=0:pins=-", + "text_encoder:resident", + } + class _FakeLayerwiseManager: def __init__( @@ -1049,7 +1278,7 @@ def __init__(self, managers: list[_FakeLayerwiseManager]): class _StubResidencyArgs: - """Duck-typed stand-in for the two ServerArgs hooks promotions use.""" + """Duck-typed stand-in for the two ServerArgs hooks adjustments use.""" def __init__(self): self.required: set[str] = set() @@ -1060,13 +1289,16 @@ def require_component_resident(self, component_name, *, feature_name): def release_required_component_residency(self, component_name, *, feature_name): self.required.discard(component_name) + def component_residency_requirement(self, component_name): + return AUTO_RESIDENCY_FEATURE_NAME if component_name in self.required else None -def _plan_for(candidates: list[PromotionCandidate]) -> AutoResidencyPlan: + +def _plan_for(candidates: list[ResidencyTarget]) -> AutoResidencyPlan: return AutoResidencyPlan( estimated_peak_bytes=GIB_BYTES, reserve_bytes=MIN_VRAM_RESERVE_BYTES, budget_bytes=100 * GIB_BYTES, - promotions=candidates, + changes=candidates, ) @@ -1100,7 +1332,7 @@ def test_layerwise_runtime_floor_uses_only_the_streaming_window(self): assert runtime["transformer"] == manager.peak_managed_device_weight_bytes() -class TestCollectPromotionCandidates: +class TestCollectResidencyTargets: def _modes(self, mapping): return lambda name: mapping.get(name, RESIDENT) @@ -1126,7 +1358,7 @@ def test_filters_and_sizes(self): "custom": COMPONENT_OFFLOAD, } ) - candidates = collect_promotion_candidates( + candidates = collect_residency_targets( modules=modules, residency_mode_of=modes, explicit_residency_mode_of=lambda name: ( @@ -1143,11 +1375,11 @@ def test_filters_and_sizes(self): # resident vae, explicit placement, custom strategy, non-module: out assert set(permanent_candidates) == {"text_encoder", "transformer"} assert permanent_candidates["text_encoder"].h2d_bytes_per_request == ( - permanent_candidates["text_encoder"].promoted_weight_bytes + permanent_candidates["text_encoder"].target_resident_weight_bytes ) # layerwise DiT re-streams its layers once per denoise step assert permanent_candidates["transformer"].h2d_bytes_per_request == ( - permanent_candidates["transformer"].promoted_weight_bytes * 100 + permanent_candidates["transformer"].target_resident_weight_bytes * 100 ) assert any( candidate.component_name == "transformer" @@ -1157,11 +1389,11 @@ def test_filters_and_sizes(self): def test_mechanism_mismatch_is_not_a_candidate(self): # A module the compile window manages layerwise while its configured - # mode says component-offload: promotion would be a silent no-op. + # mode says component-offload: adjustment would be a silent no-op. mismatch = _FakeLayerwiseDit( [_FakeLayerwiseManager({"layers.0.w": torch.zeros(16)})] ) - candidates = collect_promotion_candidates( + candidates = collect_residency_targets( modules={"mismatch": mismatch}, residency_mode_of=self._modes({"mismatch": COMPONENT_OFFLOAD}), explicit_residency_mode_of=lambda _name: None, @@ -1182,7 +1414,7 @@ def test_layerwise_dit_can_release_existing_resident_layers(self): ) module = _FakeLayerwiseDit([manager]) - candidates = collect_promotion_candidates( + candidates = collect_residency_targets( modules={"transformer": module}, residency_mode_of=self._modes({"transformer": LAYERWISE_OFFLOAD}), explicit_residency_mode_of=lambda _name: None, @@ -1197,9 +1429,9 @@ def test_layerwise_dit_can_release_existing_resident_layers(self): and candidate.target_layerwise_resident_layers == (1,) and candidate.target_layerwise_pinned_layers == ((),) ) - assert target.promoted_weight_bytes < 0 + assert target.target_resident_weight_bytes > 0 assert target.active_device_delta_bytes < 0 - assert target.h2d_bytes_per_request < 0 + assert target.h2d_bytes_per_request > 0 def test_fully_pinned_component_can_release_hostpin_for_a_hotter_component(self): manager = _FakeLayerwiseManager( @@ -1211,7 +1443,7 @@ def test_fully_pinned_component_can_release_hostpin_for_a_hotter_component(self) manager._pinned_layers = (0, 1) module = _FakeLayerwiseDit([manager]) - candidates = collect_promotion_candidates( + candidates = collect_residency_targets( modules={"text_encoder": module}, residency_mode_of=self._modes({"text_encoder": LAYERWISE_OFFLOAD}), explicit_residency_mode_of=lambda _name: None, @@ -1236,7 +1468,7 @@ def test_multi_worker_candidate_keeps_existing_hostpin_placement(self): manager._pinned_layers = (0, 1) module = _FakeLayerwiseDit([manager]) - candidates = collect_promotion_candidates( + candidates = collect_residency_targets( modules={"transformer": module}, residency_mode_of=self._modes({"transformer": LAYERWISE_OFFLOAD}), explicit_residency_mode_of=lambda _name: None, @@ -1260,7 +1492,7 @@ def test_mixed_dtype_layerwise_component_is_not_permanently_promoted(self): ) module = _FakeLayerwiseDit([manager]) - candidates = collect_promotion_candidates( + candidates = collect_residency_targets( modules={"vae": module}, residency_mode_of=self._modes({"vae": LAYERWISE_OFFLOAD}), explicit_residency_mode_of=lambda _name: None, @@ -1271,12 +1503,49 @@ def test_mixed_dtype_layerwise_component_is_not_permanently_promoted(self): assert not any(candidate.permanent_residency for candidate in candidates) + def test_non_dit_layerwise_frontier_includes_partial_residency(self): + manager = _FakeLayerwiseManager( + {f"layers.{index}.w": torch.zeros(16) for index in range(4)} + ) + module = _FakeLayerwiseDit([manager]) + + candidates = collect_residency_targets( + modules={"text_encoder": module}, + residency_mode_of=self._modes({"text_encoder": LAYERWISE_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=1, + ) + + assert any( + candidate.target_layerwise_resident_layers == (2,) + and candidate.target_mode() == LAYERWISE_OFFLOAD + for candidate in candidates + ) + + def test_auto_resident_component_keeps_offload_as_a_target(self): + module = nn.Linear(4, 4) + candidates = collect_residency_targets( + modules={"text_encoder": module}, + residency_mode_of=self._modes({"text_encoder": RESIDENT}), + baseline_residency_mode_of=self._modes({"text_encoder": COMPONENT_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=1, + auto_resident_components={"text_encoder"}, + ) + + by_mode = {candidate.target_mode(): candidate for candidate in candidates} + assert set(by_mode) == {COMPONENT_OFFLOAD, RESIDENT} + assert by_mode[RESIDENT].current_placement + assert by_mode[COMPONENT_OFFLOAD].inactive_device_delta_bytes < 0 + class TestApplyAndRollback: def test_component_offload_promotion_marks_resident_without_moving(self): module = nn.Linear(4, 4) args = _StubResidencyArgs() - applied = apply_promotions( + applied = apply_residency_changes( plan=_plan_for([_candidate("text_encoder", weight_gib=1)]), modules={"text_encoder": module}, server_args=args, @@ -1284,18 +1553,45 @@ def test_component_offload_promotion_marks_resident_without_moving(self): assert args.required == {"text_encoder"} assert [p.component_name for p in applied] == ["text_encoder"] - rollback_promotions( + rollback_residency_changes( applied=applied, modules={"text_encoder": module}, server_args=args ) assert args.required == set() + def test_component_resident_can_be_demoted_and_rollback_restores_it(self): + module = nn.Linear(4, 4) + args = _StubResidencyArgs() + args.required.add("text_encoder") + candidate = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + inactive_device_delta_bytes=-GIB_BYTES, + ) + + applied = apply_residency_changes( + plan=_plan_for([candidate]), + modules={"text_encoder": module}, + server_args=args, + ) + assert args.required == set() + + rollback_residency_changes( + applied=applied, + modules={"text_encoder": module}, + server_args=args, + ) + assert args.required == {"text_encoder"} + def test_layerwise_promotion_loads_all_layers_and_rollback_rearms(self): manager = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) module = _FakeLayerwiseDit([manager]) args = _StubResidencyArgs() candidate = _candidate("transformer", mode=LAYERWISE_OFFLOAD, weight_gib=1) - applied = apply_promotions( + applied = apply_residency_changes( plan=_plan_for([candidate]), modules={"transformer": module}, server_args=args, @@ -1304,7 +1600,7 @@ def test_layerwise_promotion_loads_all_layers_and_rollback_rearms(self): assert manager.remove_hooks_calls == 1 assert manager.enabled is False - rollback_promotions( + rollback_residency_changes( applied=applied, modules={"transformer": module}, server_args=args ) assert manager.enabled is True @@ -1330,16 +1626,16 @@ def test_partial_layerwise_promotion_restores_exact_group_counts(self): ] module = _FakeLayerwiseDit(managers) args = _StubResidencyArgs() - candidate = PromotionCandidate( + candidate = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, - promoted_weight_bytes=GIB_BYTES, + target_resident_weight_bytes=GIB_BYTES, h2d_bytes_per_request=10 * GIB_BYTES, target_layerwise_resident_layers=(2, 2), target_layerwise_pinned_layers=((), ()), ) - applied = apply_promotions( + applied = apply_residency_changes( plan=_plan_for([candidate]), modules={"transformer": module}, server_args=args, @@ -1347,13 +1643,54 @@ def test_partial_layerwise_promotion_restores_exact_group_counts(self): assert [manager.resident_layers for manager in managers] == [2, 2] assert args.required == set() - rollback_promotions( + rollback_residency_changes( applied=applied, modules={"transformer": module}, server_args=args, ) assert [manager.resident_layers for manager in managers] == [1, 0] + def test_permanent_layerwise_can_be_demoted_and_rollback_restores_it(self): + manager = _FakeLayerwiseManager( + { + "layers.0.w": torch.zeros(16), + "layers.1.w": torch.zeros(16), + }, + resident_layers=2, + ) + module = _FakeLayerwiseDit([manager]) + module.disable_offload() + args = _StubResidencyArgs() + args.required.add("transformer") + candidate = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + target_layerwise_resident_layers=(1,), + target_layerwise_pinned_layers=((),), + ) + + applied = apply_residency_changes( + plan=_plan_for([candidate]), + modules={"transformer": module}, + server_args=args, + ) + + assert manager.enabled is True + assert manager.resident_layers == 1 + assert args.required == set() + + rollback_residency_changes( + applied=applied, + modules={"transformer": module}, + server_args=args, + ) + assert manager.enabled is False + assert manager.resident_layers == 2 + assert args.required == {"transformer"} + def test_mid_failure_rolls_back_already_applied_promotions(self): manager = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) module = _FakeLayerwiseDit([manager]) @@ -1363,7 +1700,7 @@ def test_mid_failure_rolls_back_already_applied_promotions(self): _candidate("missing_component", weight_gib=1), ] with pytest.raises(RuntimeError, match="missing_component"): - apply_promotions( + apply_residency_changes( plan=_plan_for(candidates), modules={"transformer": module}, server_args=args, @@ -1393,19 +1730,19 @@ def test_hostpin_handoff_releases_all_components_before_acquiring( "target": _FakeLayerwiseDit([target_manager]), } candidates = [ - PromotionCandidate( + ResidencyTarget( component_name="source", residency_mode=LAYERWISE_OFFLOAD, - promoted_weight_bytes=0, + target_resident_weight_bytes=0, h2d_bytes_per_request=-1, target_layerwise_resident_layers=(0,), target_layerwise_pinned_layers=((),), pinned_host_delta_bytes=-source_manager.offloaded_weight_bytes(), ), - PromotionCandidate( + ResidencyTarget( component_name="target", residency_mode=LAYERWISE_OFFLOAD, - promoted_weight_bytes=0, + target_resident_weight_bytes=0, h2d_bytes_per_request=10, target_layerwise_resident_layers=(0,), target_layerwise_pinned_layers=((0,),), @@ -1418,14 +1755,14 @@ def test_hostpin_handoff_releases_all_components_before_acquiring( lambda: events.append("flush"), ) - applied = apply_promotions( + applied = apply_residency_changes( plan=_plan_for(candidates), modules=modules, server_args=_StubResidencyArgs(), ) assert events == ["source:unpin", "flush", "target:pin"] - assert all(promotion.pinned_host_changed for promotion in applied) + assert all(adjustment.pinned_host_changed for adjustment in applied) def test_hostpin_rollback_releases_new_owner_before_restoring_old_owner( self, monkeypatch @@ -1447,14 +1784,14 @@ def test_hostpin_rollback_releases_new_owner_before_restoring_old_owner( "target": _FakeLayerwiseDit([target_manager]), } applied = [ - AppliedPromotion( + AppliedResidencyChange( component_name="source", residency_mode=LAYERWISE_OFFLOAD, previous_layerwise_resident_layers=(0,), previous_layerwise_pinned_layers=((0,),), pinned_host_changed=True, ), - AppliedPromotion( + AppliedResidencyChange( component_name="target", residency_mode=LAYERWISE_OFFLOAD, previous_layerwise_resident_layers=(0,), @@ -1468,7 +1805,7 @@ def test_hostpin_rollback_releases_new_owner_before_restoring_old_owner( lambda: events.append("flush"), ) - rollback_promotions( + rollback_residency_changes( applied=applied, modules=modules, server_args=_StubResidencyArgs(), @@ -1513,7 +1850,7 @@ def release_required_component_residency( _candidate("transformer", mode=LAYERWISE_OFFLOAD, weight_gib=1), ] with pytest.raises(AutoResidencyRollbackError) as exc_info: - apply_promotions( + apply_residency_changes( plan=_plan_for(candidates), modules={ "text_encoder": nn.Linear(2, 2), @@ -1537,24 +1874,26 @@ def release_required_component_residency( args = _PartialArgs() applied = [ - AppliedPromotion( + AppliedResidencyChange( component_name="text_encoder", residency_mode=COMPONENT_OFFLOAD ), - AppliedPromotion(component_name="vae", residency_mode=COMPONENT_OFFLOAD), + AppliedResidencyChange( + component_name="vae", residency_mode=COMPONENT_OFFLOAD + ), ] with pytest.raises(RuntimeError, match="vae"): - rollback_promotions( + rollback_residency_changes( applied=applied, modules={}, server_args=args, ) - # the earlier-applied promotion was still undone + # the earlier-applied adjustment was still undone assert released == ["text_encoder"] class TestCandidateRanking: def test_post_request_hint_order_matches_promotion_order(self): - # The recommendation log and the promotion plan share one ranking, so + # The recommendation log and the adjustment plan share one ranking, so # the hint lists components in the order auto mode would promote them. candidates = [ _candidate("vae", weight_gib=1), @@ -1572,7 +1911,7 @@ def test_post_request_hint_order_matches_promotion_order(self): plan = plan_auto_residency( reports=[_report(budget_gib=1000, estimated_gib=50, candidates=candidates)] ) - assert [candidate.component_name for candidate in plan.promotions] == [ + assert [candidate.component_name for candidate in plan.changes] == [ candidate.component_name for candidate in ranked ] @@ -1586,7 +1925,7 @@ def test_reports_equivalent_server_args_and_kill_switch(self): ] ) message = format_applied_changes(plan=plan) - assert "text_encoder: permanent resident layers=(1,)" in message + assert "text_encoder: layerwise-offload -> resident" in message assert "vae: component-offload -> resident" in message assert "--component-residency text_encoder=resident vae=resident" in message assert "SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY=1" in message @@ -1723,7 +2062,7 @@ def test_unknown_target_resolution_skips_probe(self, monkeypatch): assert probe is None def test_skip_gate_disables_probe(self, monkeypatch): - # Full-shape calibration must share the promotion's own gate (kill + # Full-shape calibration must share the adjustment's own gate (kill # switch, quantized, manual, ...). self._patch_gate(monkeypatch, reason="performance_mode=manual") probe = _resolve_auto_residency_warmup_shape( diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py index e85e5600e3b3..4b48dbb3d657 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -10,11 +10,11 @@ from sglang.multimodal_gen.runtime.managers import gpu_worker as gpu_worker_module from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( - PROMOTION_STATUS_PROMOTED, - PROMOTION_STATUS_ROLLBACK_FAILED, - PROMOTION_STATUS_ROLLED_BACK, - PROMOTION_STATUS_SKIPPED, - AppliedPromotion, + PLACEMENT_STATUS_ADJUSTED, + PLACEMENT_STATUS_ROLLBACK_FAILED, + PLACEMENT_STATUS_ROLLED_BACK, + PLACEMENT_STATUS_SKIPPED, + AppliedResidencyChange, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( COMPONENT_OFFLOAD, @@ -30,8 +30,8 @@ async def forward(req): self.assertIsInstance(req, AutoResidencyReq) actions.append(req.action) return OutputBatch( - error="post-promotion warmup output changed", - output={"status": PROMOTION_STATUS_ROLLED_BACK}, + error="post-adjustment warmup output changed", + output={"status": PLACEMENT_STATUS_ROLLED_BACK}, ) rewarm = mock.AsyncMock() @@ -58,9 +58,9 @@ async def forward(req): def test_replans_until_the_calibrated_layout_reaches_a_fixed_point(self): responses = iter( [ - OutputBatch(output={"status": PROMOTION_STATUS_PROMOTED}), - OutputBatch(output={"status": PROMOTION_STATUS_PROMOTED}), - OutputBatch(output={"status": PROMOTION_STATUS_SKIPPED}), + OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), + OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), + OutputBatch(output={"status": PLACEMENT_STATUS_SKIPPED}), ] ) actions = [] @@ -93,9 +93,9 @@ async def forward(req): self.assertIsInstance(req, AutoResidencyReq) actions.append(req.action) status = ( - PROMOTION_STATUS_PROMOTED + PLACEMENT_STATUS_ADJUSTED if req.action == "apply" - else PROMOTION_STATUS_ROLLED_BACK + else PLACEMENT_STATUS_ROLLED_BACK ) return OutputBatch(output={"status": status}) @@ -119,10 +119,10 @@ def test_failed_calibration_reports_rollback_failure(self): async def forward(req): self.assertIsInstance(req, AutoResidencyReq) if req.action == "apply": - return OutputBatch(output={"status": PROMOTION_STATUS_PROMOTED}) + return OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}) return OutputBatch( error="rollback failed", - output={"status": PROMOTION_STATUS_ROLLBACK_FAILED}, + output={"status": PLACEMENT_STATUS_ROLLBACK_FAILED}, ) rewarm = mock.AsyncMock(side_effect=RuntimeError("oom")) @@ -143,7 +143,7 @@ def test_apply_rpc_failure_after_calibration_aborts_without_rollback(self): actions = [] responses = iter( [ - OutputBatch(output={"status": PROMOTION_STATUS_PROMOTED}), + OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), RuntimeError("rpc failed"), ] ) @@ -168,7 +168,7 @@ async def forward(req): mock.patch.object(server_warmup, "run_async_client_warmup", rewarm), self.assertRaisesRegex( RuntimeError, - "auto residency apply failed after a calibrated promotion", + "auto residency apply failed after a calibrated adjustment", ), ): asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) @@ -177,10 +177,10 @@ async def forward(req): self.assertEqual(rewarm.await_count, 1) def test_failed_later_round_keeps_earlier_calibrated_promotions(self): - first_round = AppliedPromotion("transformer", COMPONENT_OFFLOAD) + first_round = AppliedResidencyChange("transformer", COMPONENT_OFFLOAD) latest_round = [ - AppliedPromotion("text_encoder", COMPONENT_OFFLOAD), - AppliedPromotion("vae", COMPONENT_OFFLOAD), + AppliedResidencyChange("text_encoder", COMPONENT_OFFLOAD), + AppliedResidencyChange("vae", COMPONENT_OFFLOAD), ] worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 @@ -192,7 +192,9 @@ def test_failed_later_round_keeps_earlier_calibrated_promotions(self): worker._auto_residency_all_gather = lambda value: [value] worker._invalidate_component_strategies = mock.Mock() - with mock.patch.object(gpu_worker_module, "rollback_promotions") as rollback: + with mock.patch.object( + gpu_worker_module, "rollback_residency_changes" + ) as rollback: response = worker.rollback_auto_residency() rollback.assert_called_once_with( @@ -200,12 +202,12 @@ def test_failed_later_round_keeps_earlier_calibrated_promotions(self): modules={}, server_args=worker.server_args, ) - self.assertEqual(response.output["status"], PROMOTION_STATUS_ROLLED_BACK) + self.assertEqual(response.output["status"], PLACEMENT_STATUS_ROLLED_BACK) self.assertEqual(worker._auto_residency_applied, [first_round]) self.assertEqual(worker._auto_residency_round_sizes, [1]) def test_failed_apply_round_does_not_rollback_prior_round(self): - first_round = AppliedPromotion("transformer", COMPONENT_OFFLOAD) + first_round = AppliedResidencyChange("transformer", COMPONENT_OFFLOAD) worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 worker.pipeline = SimpleNamespace(modules={}) @@ -214,8 +216,10 @@ def test_failed_apply_round_does_not_rollback_prior_round(self): worker._auto_residency_round_sizes = [1, 0] worker._invalidate_component_strategies = mock.Mock() - with mock.patch.object(gpu_worker_module, "rollback_promotions") as rollback: - error = worker._rollback_applied_promotions(latest_round_only=True) + with mock.patch.object( + gpu_worker_module, "rollback_residency_changes" + ) as rollback: + error = worker._rollback_applied_residency_changes(latest_round_only=True) self.assertIsNone(error) rollback.assert_not_called() diff --git a/python/sglang/multimodal_gen/test/unit/test_component_residency.py b/python/sglang/multimodal_gen/test/unit/test_component_residency.py index 39ece58b4a70..2da9bd5f1eeb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_residency.py @@ -216,7 +216,9 @@ def test_warmup_records_use_and_transition_peaks(monkeypatch): peaks = manager.take_warmup_phase_peaks() inactive_peak = WarmupPhasePeak((), 7) - transformer_peak = WarmupPhasePeak(("transformer",), 7) + transformer_peak = WarmupPhasePeak( + ("transformer",), 7, used_components=("transformer",) + ) assert peaks["request:before-stage"] == inactive_peak assert peaks["0:denoise:setup"] == inactive_peak assert peaks["0:denoise:transition:idle->transformer"] == transformer_peak @@ -265,7 +267,7 @@ def test_warmup_records_same_component_dtype_prepare_as_transition(monkeypatch): manager._record_warmup_phase_peak() assert manager._warmup_phase_peaks[ "0:stage:transition:transformer->transformer" - ] == WarmupPhasePeak(("transformer",), 7) + ] == WarmupPhasePeak(("transformer",), 7, used_components=("transformer",)) assert strategy.prepare_for_use.call_count == 2 @@ -314,7 +316,7 @@ def test_warmup_attributes_prefetch_peak_to_prefetched_component(monkeypatch): assert manager._warmup_phase_peaks[ "0:encode:prefetch:transformer" - ] == WarmupPhasePeak(("transformer",), 7) + ] == WarmupPhasePeak(("transformer",), 7, used_components=("transformer",)) assert manager._warmup_phase_peaks["0:encode:between"] == WarmupPhasePeak((), 7) @@ -352,10 +354,10 @@ def test_warmup_splits_sequential_component_transition(monkeypatch): manager._record_warmup_phase_peak() assert manager._warmup_phase_peaks["0:stage:transition:text_encoder->idle"] == ( - WarmupPhasePeak(("text_encoder",), 7) + WarmupPhasePeak(("text_encoder",), 7, used_components=("text_encoder",)) ) assert manager._warmup_phase_peaks["0:stage:transition:idle->transformer"] == ( - WarmupPhasePeak(("transformer",), 7) + WarmupPhasePeak(("transformer",), 7, used_components=("transformer",)) ) diff --git a/python/sglang/multimodal_gen/test/unit/test_dp_routing.py b/python/sglang/multimodal_gen/test/unit/test_dp_routing.py index bc6572590e8f..fc3438c54dfb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_dp_routing.py +++ b/python/sglang/multimodal_gen/test/unit/test_dp_routing.py @@ -69,24 +69,24 @@ def test_fanout_merge_surfaces_the_failing_replica(): assert _merge_fanout_results([ok, OutputBatch(output=None)]) is ok -def test_auto_residency_fanout_rewarms_when_any_replica_promotes(): +def test_auto_residency_fanout_rewarms_when_any_replica_adjusts(): skipped = OutputBatch(output={"status": "skipped"}) - promoted = OutputBatch(output={"status": "promoted"}) - assert _merge_auto_residency_results([skipped, promoted]) is promoted + adjusted = OutputBatch(output={"status": "adjusted"}) + assert _merge_auto_residency_results([skipped, adjusted]) is adjusted def test_auto_residency_fanout_tolerates_a_locally_rolled_back_replica(): restored = OutputBatch( output={"status": "rolled_back"}, error="original placement restored" ) - promoted = OutputBatch(output={"status": "promoted"}) - assert _merge_auto_residency_results([restored, promoted]) is promoted + adjusted = OutputBatch(output={"status": "adjusted"}) + assert _merge_auto_residency_results([restored, adjusted]) is adjusted def test_auto_residency_fanout_never_hides_failed_rollback(): failed = OutputBatch(output={"status": "rollback_failed"}, error="restart required") - promoted = OutputBatch(output={"status": "promoted"}) - assert _merge_auto_residency_results([promoted, failed]) is failed + adjusted = OutputBatch(output={"status": "adjusted"}) + assert _merge_auto_residency_results([adjusted, failed]) is failed def test_auto_residency_fanout_reports_successful_rollback(): diff --git a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py index d644e7e0143b..384f964af449 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -127,12 +127,15 @@ def test_server_warmup_preserves_peak_after_managed_stage_timeline(): ): worker._record_server_warmup_memory( req=req, + workload=(128, 96, 9, 2), baseline_allocated_bytes=3, succeeded=True, ) record = worker._auto_residency_warmup_records[0] assert record.peak_allocated_bytes == 9 + assert (record.width, record.height, record.num_frames) == (128, 96, 9) + assert record.num_inference_steps == 2 assert record.phase_peak_allocated_bytes["request:untracked"] == 9 assert record.phase_active_components["request:untracked"] == ("transformer",) @@ -166,6 +169,7 @@ def test_server_warmup_does_not_treat_allocator_cache_as_untracked_live_memory() ): worker._record_server_warmup_memory( req=req, + workload=(64, 64, 1, 1), baseline_allocated_bytes=3, succeeded=True, ) @@ -191,7 +195,9 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): nnodes=1, node_rank=0, residency_mode=lambda _name: "resident", + configured_residency_mode=lambda _name: "resident", explicit_residency_mode=lambda _name: None, + component_residency_requirement=lambda _name: None, ) device_module = Mock() device_module.mem_get_info.return_value = (100, 100) @@ -232,7 +238,9 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): nnodes=1, node_rank=0, residency_mode=lambda _name: "resident", + configured_residency_mode=lambda _name: "resident", explicit_residency_mode=lambda _name: None, + component_residency_requirement=lambda _name: None, ) device_module = Mock() device_module.mem_get_info.return_value = (100 * GIB_BYTES, 140 * GIB_BYTES) diff --git a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py index e41b1d468305..26df22ce3ff2 100644 --- a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py @@ -103,6 +103,100 @@ def test_ties_prefer_lower_resource_cost_then_stable_key(): } +def test_complete_state_selection_can_replace_current_placement(): + plan = optimize_placement( + [ + PlacementOption( + group_key="transformer", + option_key="transformer:resident", + resource_delta_bytes={"vram": 0}, + estimated_latency_savings=100, + ), + PlacementOption( + group_key="transformer", + option_key="transformer:offload", + resource_delta_bytes={"vram": -40}, + estimated_latency_savings=0, + ), + PlacementOption( + group_key="text_encoder", + option_key="text_encoder:resident", + resource_delta_bytes={"vram": 30}, + estimated_latency_savings=200, + ), + PlacementOption( + group_key="text_encoder", + option_key="text_encoder:offload", + resource_delta_bytes={"vram": 0}, + estimated_latency_savings=0, + ), + ], + resource_budget_bytes={"vram": 0}, + require_selection_from_every_group=True, + ) + + assert {option.option_key for option in plan.selections} == { + "transformer:offload", + "text_encoder:resident", + } + + +def test_latency_equivalent_plan_prefers_lower_soft_memory_cost(): + plan = optimize_placement( + [ + PlacementOption( + group_key="transformer", + option_key="transformer-resident", + resource_delta_bytes={"gpu:rank0:runtime": 40}, + estimated_latency_savings=1_000, + placement_cost_bytes=(40, 0), + ), + PlacementOption( + group_key="text_encoder", + option_key="text-encoder-resident", + resource_delta_bytes={"gpu:rank0:runtime": 10}, + estimated_latency_savings=60, + placement_cost_bytes=(10, 0), + ), + ], + resource_budget_bytes={"gpu:rank0:runtime": 100}, + estimated_latency_tolerance=100, + ) + + assert [option.option_key for option in plan.selections] == ["transformer-resident"] + assert plan.estimated_latency_savings == 1_000 + assert plan.placement_cost_bytes == (40, 0) + + +def test_latency_tolerance_is_global_not_per_option(): + plan = optimize_placement( + [ + PlacementOption( + group_key="a", + option_key="a-resident", + resource_delta_bytes={}, + estimated_latency_savings=60, + placement_cost_bytes=(1,), + ), + PlacementOption( + group_key="b", + option_key="b-resident", + resource_delta_bytes={}, + estimated_latency_savings=60, + placement_cost_bytes=(1,), + ), + ], + resource_budget_bytes={}, + estimated_latency_tolerance=100, + ) + + # Neither 60-unit option would survive a per-option 100-unit cutoff. Their + # combined benefit first establishes a non-zero global utility floor; the + # least costly placement within that floor then keeps one of them. + assert [option.option_key for option in plan.selections] == ["a-resident"] + assert plan.estimated_latency_savings == 60 + + def test_one_static_placement_satisfies_all_observed_phase_constraints(): plan = optimize_placement( [ From a268535efd4c54415f0851e6ea1fc5e469bb90f4 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 03:57:40 +0800 Subject: [PATCH 32/90] [diffusion] enumerate layerwise residency ratios --- .../memory_managers/auto_residency.py | 45 ++++++++++++++----- .../test/unit/test_auto_residency.py | 16 +++++++ 2 files changed, 51 insertions(+), 10 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 58935e50ec02..d66649625881 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -753,6 +753,40 @@ def _layerwise_pin_targets( return list(dict.fromkeys(targets)) +def _layerwise_resident_targets(managers: Sequence) -> list[tuple[int, ...]]: + """All resident-count tuples expressible by the component-level CLI knob. + + ``--layerwise-resident-layers`` accepts either one absolute count or one + ratio for a component. Those are different paths when a component owns + layer stacks of unequal length: for stacks of 4 and 8 layers, an absolute + value of 2 gives ``(2, 2)``, while a ratio of 0.25 gives ``(1, 2)``. + Enumerate every interval where the rounded ratio result is constant, plus + every absolute count, so auto placement does not silently omit a state the + explicit interface can represent. + """ + layer_counts = tuple(manager.num_layers for manager in managers) + targets = {tuple(0 for _ in managers)} + + for count in range(1, max(layer_counts) + 1): + targets.add(tuple(min(count, num_layers) for num_layers in layer_counts)) + + boundaries = {0.0, 1.0} + for num_layers in layer_counts: + boundaries.update( + (count + 0.5) / num_layers + for count in range(num_layers) + if count + 0.5 < num_layers + ) + ordered_boundaries = sorted(boundaries) + for lower, upper in zip(ordered_boundaries, ordered_boundaries[1:]): + ratio = (lower + upper) / 2 + targets.add( + tuple(max(1, int(round(ratio * num_layers))) for num_layers in layer_counts) + ) + + return sorted(targets) + + def _layerwise_host_transition_bytes( *, managers: Sequence, @@ -912,16 +946,7 @@ def collect_residency_targets( uses_per_streamed_layer=uses_per_request, ) - resident_targets = set() - for target in range( - 0, - max(manager.num_layers for manager in managers) + 1, - ): - resident_targets.add( - tuple(min(target, manager.num_layers) for manager in managers) - ) - - for target_resident_layers in sorted(resident_targets): + for target_resident_layers in _layerwise_resident_targets(managers): target_resident_bytes = sum( manager.resident_weight_bytes(count) for manager, count in zip(managers, target_resident_layers) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 531e0130715c..67403d274f0d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -21,6 +21,7 @@ RankResidencyReport, ResidencyTarget, WarmupMemoryRecord, + _layerwise_resident_targets, apply_residency_changes, collect_residency_targets, component_resident_size_bytes, @@ -1336,6 +1337,21 @@ class TestCollectResidencyTargets: def _modes(self, mapping): return lambda name: mapping.get(name, RESIDENT) + def test_resident_frontier_includes_absolute_and_ratio_cli_states(self): + managers = [ + _FakeLayerwiseManager( + {f"layers.{index}.w": torch.zeros(16) for index in range(4)} + ), + _FakeLayerwiseManager( + {f"layers.{index}.w": torch.zeros(16) for index in range(8)} + ), + ] + + targets = _layerwise_resident_targets(managers) + + assert (2, 2) in targets + assert (1, 2) in targets + def test_filters_and_sizes(self): te = nn.Linear(8, 8) vae = nn.Linear(4, 4) From 03b8922ebbf62d1ae11fba007e28a7c4ed900465 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 04:03:17 +0800 Subject: [PATCH 33/90] [diffusion] document calibrated residency scope --- docs/docs/sglang-diffusion/api/cli.mdx | 2 +- docs/docs/sglang-diffusion/deployment_cookbook.mdx | 10 +++++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 568dcea3999c..c21c08b7e8c0 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -85,7 +85,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; a single-GPU run can also rebalance its HostPin allowance. The selected placement targets the model-default (or `--warmup-sampling-params`) workload with batch size 1; use explicit residency for larger production shapes or batches. 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 placement, measures the representative workload, and jointly selects component/layer residency against every rank and execution-phase VRAM constraint plus HostPin limits. Multi-process HostPin planning stays inside non-overlapping per-worker shares. The selected placement targets the model-default (or `--warmup-sampling-params`) workload with batch size 1; use explicit residency for larger production shapes or batches. 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. - `--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. Startup impact is model-dependent, so benchmark the target model before deployment. Disabled by default because checkpoint and model weights coexist temporarily, substantially increasing peak GPU memory. 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 diff --git a/docs/docs/sglang-diffusion/deployment_cookbook.mdx b/docs/docs/sglang-diffusion/deployment_cookbook.mdx index c917bfb5c8e7..e3384126126c 100644 --- a/docs/docs/sglang-diffusion/deployment_cookbook.mdx +++ b/docs/docs/sglang-diffusion/deployment_cookbook.mdx @@ -115,7 +115,7 @@ See [OpenAI API: Served model name](/docs/sglang-diffusion/api/openai_api#served `auto` - Default. Keeps legacy safe offload defaults and uses FSDP/CFG only on validated multi-GPU deployments where FSDP can replace DiT offload. + Default. Server warmup starts from a load-safe layout, measures the representative workload, then adjusts eligible component and layer residency before the ready signal. Offline and unsupported paths keep the static safe defaults. `speed` @@ -128,7 +128,11 @@ 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. Model-specific components such as VAEs become resident only when their configured memory threshold is met. `memory` instead keeps the VAE in its default layerwise set to maximize memory headroom; use `--component-residency vae=resident` for a measured recipe with sufficient capacity. 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. +With an eligible `sglang serve` warmup, `auto` records allocated-memory peaks for the representative request and their active component layouts. One placement decision then satisfies every rank and execution-phase VRAM constraint together with HostPin and HostPin-transition limits. The target frontier covers component offload, resident components, layerwise resident counts, and pinned layer-store reassignment. A post-adjustment warmup realizes the selected state and validates its memory peak before the server becomes ready. Later rounds may replace an earlier automatic choice rather than only increasing residency. + +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. + +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. 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. @@ -147,7 +151,7 @@ When `torch.compile` is enabled, `--offload-during-compile` stays on by default. Breakable CUDA graph is a separate manual opt-in for supported image pipelines. If you enable `--enable-breakable-cuda-graph`, declare every served resolution in `--warmup-resolutions` so warmup captures matching graph signatures. -The preset is intentionally coarse. A future continuous value such as `0.0` to `1.0` could express the speed-memory tradeoff more precisely, but it would need model-specific memory models and clearer user expectations. Until then, use the preset plus explicit flags for overrides. +The calibrated planner is exact for the candidate states and memory constraints it receives, but transfer-based latency estimates cannot model every prefetch overlap or bandwidth interaction. Use explicit placement for a benchmarked production recipe when stable latency matters more than automatic adaptation. Examples: From 6d03a14ed70c432b7b0a8fa423c58093d62aa657 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 04:22:07 +0800 Subject: [PATCH 34/90] [diffusion] constrain residency placement transitions --- .../runtime/managers/gpu_worker.py | 3 + .../memory_managers/auto_residency.py | 85 +++++++- .../test/unit/test_auto_residency.py | 191 ++++++++++++++++++ .../test/unit/test_performance_metrics.py | 3 + 4 files changed, 273 insertions(+), 9 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 180c2a526204..5bde6e0a7e9d 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1499,6 +1499,9 @@ def _build_auto_residency_report( 0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES ) // local_worker_count, + device_transition_allocated_bytes=int( + torch.get_device_module().memory_allocated() + ), estimated_request_duration_ns=estimated_request_duration_ns, candidate_latency_savings_ns=candidate_latency_savings_ns, candidates=candidates, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index d66649625881..15fb1c6a020d 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -149,6 +149,12 @@ class ResidencyTarget(msgspec.Struct, frozen=True): pinned_host_delta_bytes: int = 0 host_unpin_scratch_bytes: int = 0 host_pin_scratch_bytes: int = 0 + host_materialize_scratch_bytes: int = 0 + # Signed device-memory delta while applying the placement before the + # validation warmup. Layerwise -> resident materializes every managed + # layer immediately; a demotion can release those bytes first and fund a + # later materialization in the same transaction. + device_transition_delta_bytes: int = 0 permanent_residency: bool = False # Device-memory delta relative to the measured placement. A component # already loaded for its own phase has a different delta from phases where @@ -226,6 +232,7 @@ class RankResidencyReport(msgspec.Struct, frozen=True): pinned_host_bytes: int = 0 host_pin_capacity_bytes: int = 0 host_transition_headroom_bytes: int = 0 + device_transition_allocated_bytes: int = 0 estimated_request_duration_ns: int = 0 candidate_latency_savings_ns: dict[str, int] = {} candidates: list[ResidencyTarget] = [] @@ -870,6 +877,12 @@ def collect_residency_targets( target_residency_mode=COMPONENT_OFFLOAD, target_resident_weight_bytes=0, h2d_bytes_per_request=0, + host_materialize_scratch_bytes=( + weight_bytes if current_resident else 0 + ), + device_transition_delta_bytes=( + -weight_bytes if current_resident else 0 + ), active_device_delta_bytes=0, inactive_device_delta_bytes=( -weight_bytes if current_resident else 0 @@ -886,6 +899,7 @@ def collect_residency_targets( target_resident_weight_bytes=weight_bytes, h2d_bytes_per_request=weight_bytes, permanent_residency=True, + device_transition_delta_bytes=0, active_device_delta_bytes=0, inactive_device_delta_bytes=( 0 if current_resident else weight_bytes @@ -1000,6 +1014,9 @@ def collect_residency_targets( ), host_unpin_scratch_bytes=unpin_scratch, host_pin_scratch_bytes=pin_scratch, + device_transition_delta_bytes=( + -managed_weight_bytes if current_permanent else 0 + ), active_device_delta_bytes=( target_peak_device_bytes - current_peak_device_bytes ), @@ -1057,6 +1074,9 @@ def collect_residency_targets( ), host_unpin_scratch_bytes=unpin_scratch, host_pin_scratch_bytes=pin_scratch, + device_transition_delta_bytes=( + 0 if current_permanent else managed_weight_bytes + ), permanent_residency=True, active_device_delta_bytes=( managed_weight_bytes - current_peak_device_bytes @@ -1235,6 +1255,14 @@ def _consensus_candidates( host_pin_scratch_bytes=max( candidate.host_pin_scratch_bytes for candidate in rank_candidates ), + host_materialize_scratch_bytes=max( + candidate.host_materialize_scratch_bytes + for candidate in rank_candidates + ), + device_transition_delta_bytes=max( + candidate.device_transition_delta_bytes + for candidate in rank_candidates + ), permanent_residency=permanent.pop(), active_device_delta_bytes=max( candidate.active_device_delta_bytes for candidate in rank_candidates @@ -1307,13 +1335,25 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP 0, report.budget_bytes - phase_peak - reserves_by_rank[report.rank], ) - has_hostpin_options = any( + has_device_transition_options = any( + candidate.device_transition_delta_bytes != 0 for candidate in candidates + ) + if has_device_transition_options: + for report in reports: + resource_budgets[f"gpu:rank{report.rank}:placement-transition"] = max( + 0, + report.budget_bytes + - report.device_transition_allocated_bytes + - reserves_by_rank[report.rank], + ) + has_host_resources = any( candidate.pinned_host_delta_bytes != 0 or candidate.host_unpin_scratch_bytes != 0 or candidate.host_pin_scratch_bytes != 0 + or candidate.host_materialize_scratch_bytes != 0 for candidate in candidates ) - if has_hostpin_options: + if has_host_resources: for report in reports: prefix = f"node{report.node_rank}:rank{report.rank}" resource_budgets[f"hostpin:{prefix}"] = max( @@ -1326,6 +1366,9 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP resource_budgets[f"hostram:{prefix}:pin"] = ( report.host_transition_headroom_bytes ) + resource_budgets[f"hostram:{prefix}:materialize"] = ( + report.host_transition_headroom_bytes + ) candidate_by_key = {candidate.option_key(): candidate for candidate in candidates} report_candidates = [ @@ -1378,7 +1421,11 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP else: phase_cost = rank_candidate.inactive_device_delta_bytes resource_deltas[resource_name] = phase_cost - if has_hostpin_options: + if has_device_transition_options: + resource_deltas[f"gpu:rank{report.rank}:placement-transition"] = ( + rank_candidate.device_transition_delta_bytes + ) + if has_host_resources: prefix = f"node{report.node_rank}:rank{report.rank}" host_resource = f"hostpin:{prefix}" resource_deltas[host_resource] = ( @@ -1391,6 +1438,9 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP resource_deltas[f"hostram:{prefix}:pin"] = ( rank_candidate.host_pin_scratch_bytes ) + resource_deltas[f"hostram:{prefix}:materialize"] = ( + rank_candidate.host_materialize_scratch_bytes + ) options.append( PlacementOption( group_key=candidate.component_name, @@ -1476,6 +1526,7 @@ def apply_residency_changes( ordered_changes = sorted( plan.changes, key=lambda candidate: ( + candidate.device_transition_delta_bytes, max( candidate.active_device_delta_bytes, candidate.inactive_device_delta_bytes, @@ -1567,20 +1618,36 @@ def apply_residency_changes( and candidate.target_layerwise_pinned_layers != previous_pins_by_component[candidate.component_name] ] - if pinning_changes: + host_transition_changes = [ + candidate + for candidate in ordered_changes + if candidate.host_unpin_scratch_bytes + or candidate.host_pin_scratch_bytes + or candidate.host_materialize_scratch_bytes + ] + if host_transition_changes: host_headroom = max( 0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES ) required_unpin = sum( - candidate.host_unpin_scratch_bytes for candidate in pinning_changes + candidate.host_unpin_scratch_bytes + for candidate in host_transition_changes ) required_pin = sum( - candidate.host_pin_scratch_bytes for candidate in pinning_changes + candidate.host_pin_scratch_bytes + for candidate in host_transition_changes + ) + required_materialize = sum( + candidate.host_materialize_scratch_bytes + for candidate in host_transition_changes + ) + required_host_headroom = max( + required_unpin, required_pin, required_materialize ) - if max(required_unpin, required_pin) > host_headroom: + if required_host_headroom > host_headroom: raise MemoryError( - "host memory changed after planning: HostPin repack needs " - f"{max(required_unpin, required_pin) / GIB_BYTES:.2f} GiB " + "host memory changed after planning: placement transition needs " + f"{required_host_headroom / GIB_BYTES:.2f} GiB " f"but {host_headroom / GIB_BYTES:.2f} GiB is available" ) applied_names: set[str] = set() diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 67403d274f0d..1f2c303bc7f2 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -480,6 +480,7 @@ def _report( pinned_host_gib: int = 0, host_pin_capacity_gib: int = 0, host_transition_headroom_gib: int = 0, + device_transition_allocated_gib: int = 0, target_workload_measured: bool = False, estimated_request_duration_ns: int = 0, candidate_latency_savings_ns: dict[str, int] | None = None, @@ -504,6 +505,7 @@ def _report( pinned_host_bytes=pinned_host_gib * GIB_BYTES, host_pin_capacity_bytes=host_pin_capacity_gib * GIB_BYTES, host_transition_headroom_bytes=(host_transition_headroom_gib * GIB_BYTES), + device_transition_allocated_bytes=(device_transition_allocated_gib * GIB_BYTES), estimated_request_duration_ns=estimated_request_duration_ns, candidate_latency_savings_ns=candidate_latency_savings_ns or {}, candidates=candidates if candidates is not None else [_candidate("vae")], @@ -1098,6 +1100,98 @@ def test_hostpin_repack_must_fit_transition_headroom(self): assert plan.changes == [] assert plan.resource_budget_bytes["hostram:node0:rank0:pin"] == 5 * GIB_BYTES + def test_component_demotion_must_fit_host_materialization_headroom(self): + demote = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=100 * GIB_BYTES, + host_materialize_scratch_bytes=6 * GIB_BYTES, + device_transition_delta_bytes=-6 * GIB_BYTES, + ) + plan = plan_auto_residency( + reports=[ + _report( + host_transition_headroom_gib=5, + candidates=[demote], + ) + ] + ) + + assert plan.changes == [] + assert ( + plan.resource_budget_bytes["hostram:node0:rank0:materialize"] + == 5 * GIB_BYTES + ) + + def test_layerwise_materialization_must_fit_transition_vram(self): + resident = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=10 * GIB_BYTES, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(10,), + target_layerwise_pinned_layers=((),), + device_transition_delta_bytes=10 * GIB_BYTES, + ) + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=100, + estimated_gib=50, + device_transition_allocated_gib=90, + target_workload_measured=True, + candidates=[resident], + ) + ] + ) + + assert plan.changes == [] + assert ( + plan.resource_budget_bytes["gpu:rank0:placement-transition"] + == 5 * GIB_BYTES + ) + + def test_device_release_can_fund_later_layerwise_materialization(self): + release = ResidencyTarget( + component_name="cold_transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=-10 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + device_transition_delta_bytes=-10 * GIB_BYTES, + ) + materialize = ResidencyTarget( + component_name="hot_transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=10 * GIB_BYTES, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(10,), + target_layerwise_pinned_layers=((),), + device_transition_delta_bytes=10 * GIB_BYTES, + ) + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=100, + estimated_gib=50, + device_transition_allocated_gib=95, + target_workload_measured=True, + candidates=[release, materialize], + ) + ] + ) + + assert {candidate.component_name for candidate in plan.changes} == { + "cold_transformer", + "hot_transformer", + } + def test_complete_state_frontier_can_replace_an_earlier_resident(self): transformer_offload = ResidencyTarget( component_name="transformer", @@ -1256,6 +1350,8 @@ def set_resident_layers(self, resident_layers): def load_all_layers(self): self.load_all_layers_calls += 1 + if self._event_log is not None: + self._event_log.append(f"{self._event_name}:load") if self.fail_load: raise RuntimeError("CUDA out of memory") @@ -1267,6 +1363,8 @@ def register_forward_hooks(self): def sync_all_layers_to_cpu(self): self.sync_to_cpu_calls += 1 + if self._event_log is not None: + self._event_log.append(f"{self._event_name}:sync") def release_all(self): self.release_all_calls += 1 @@ -1555,6 +1653,50 @@ def test_auto_resident_component_keeps_offload_as_a_target(self): assert set(by_mode) == {COMPONENT_OFFLOAD, RESIDENT} assert by_mode[RESIDENT].current_placement assert by_mode[COMPONENT_OFFLOAD].inactive_device_delta_bytes < 0 + assert by_mode[COMPONENT_OFFLOAD].device_transition_delta_bytes < 0 + assert by_mode[COMPONENT_OFFLOAD].host_materialize_scratch_bytes > 0 + + def test_layerwise_frontier_accounts_for_immediate_materialization(self): + manager = _FakeLayerwiseManager( + { + "layers.0.w": torch.zeros(16), + "layers.1.w": torch.zeros(16), + } + ) + module = _FakeLayerwiseDit([manager]) + + candidates = collect_residency_targets( + modules={"transformer": module}, + residency_mode_of=self._modes({"transformer": LAYERWISE_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=10, + ) + resident = next( + candidate for candidate in candidates if candidate.permanent_residency + ) + assert ( + resident.device_transition_delta_bytes == manager.offloaded_weight_bytes() + ) + + module.disable_offload() + candidates = collect_residency_targets( + modules={"transformer": module}, + residency_mode_of=self._modes({"transformer": RESIDENT}), + baseline_residency_mode_of=self._modes({"transformer": LAYERWISE_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=10, + auto_resident_components={"transformer"}, + ) + streamed = next( + candidate + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + ) + assert ( + streamed.device_transition_delta_bytes == -manager.offloaded_weight_bytes() + ) class TestApplyAndRollback: @@ -1780,6 +1922,55 @@ def test_hostpin_handoff_releases_all_components_before_acquiring( assert events == ["source:unpin", "flush", "target:pin"] assert all(adjustment.pinned_host_changed for adjustment in applied) + def test_vram_handoff_releases_before_materializing(self): + events: list[str] = [] + source_manager = _FakeLayerwiseManager( + {"layers.0.w": torch.zeros(16)}, + event_log=events, + event_name="source", + ) + target_manager = _FakeLayerwiseManager( + {"layers.0.w": torch.zeros(16)}, + event_log=events, + event_name="target", + ) + source = _FakeLayerwiseDit([source_manager]) + target = _FakeLayerwiseDit([target_manager]) + source.disable_offload() + events.clear() + candidates = [ + ResidencyTarget( + component_name="source", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=-1, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + device_transition_delta_bytes=-source_manager.offloaded_weight_bytes(), + active_device_delta_bytes=100, + ), + ResidencyTarget( + component_name="target", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=target_manager.offloaded_weight_bytes(), + h2d_bytes_per_request=10, + target_layerwise_resident_layers=(1,), + target_layerwise_pinned_layers=((),), + device_transition_delta_bytes=target_manager.offloaded_weight_bytes(), + active_device_delta_bytes=-100, + ), + ] + + apply_residency_changes( + plan=_plan_for(candidates), + modules={"source": source, "target": target}, + server_args=_StubResidencyArgs(), + ) + + assert events == ["source:sync", "target:load"] + def test_hostpin_rollback_releases_new_owner_before_restoring_old_owner( self, monkeypatch ): diff --git a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py index 384f964af449..e79440262bed 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -202,6 +202,7 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): device_module = Mock() device_module.mem_get_info.return_value = (100, 100) device_module.memory_reserved.return_value = 0 + device_module.memory_allocated.return_value = 1 with patch.object(torch, "get_device_module", return_value=device_module): report = worker._build_auto_residency_report( @@ -245,6 +246,7 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): device_module = Mock() device_module.mem_get_info.return_value = (100 * GIB_BYTES, 140 * GIB_BYTES) device_module.memory_reserved.return_value = 2 * GIB_BYTES + device_module.memory_allocated.return_value = 60 * GIB_BYTES monkeypatch.setattr( gpu_worker_module.envs, "SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB", @@ -272,6 +274,7 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): ) assert report.budget_bytes == 80 * GIB_BYTES + assert report.device_transition_allocated_bytes == 60 * GIB_BYTES def test_baseline_config_loads_per_scenario_peak_vram(tmp_path): From f8afe141f37768119a0da0f9606f38dd6b197066 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 07:02:37 +0800 Subject: [PATCH 35/90] [diffusion] scale exact placement optimization --- .../memory_managers/placement_budget.py | 273 ++++++++++++------ .../test/unit/test_placement_budget.py | 162 +++++++++++ 2 files changed, 347 insertions(+), 88 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py index f7bc77c0da0c..4f1e3fe08560 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py @@ -104,6 +104,43 @@ def _pareto_prune(states: Iterable[_State]) -> list[_State]: return frontier +def _drop_locally_latency_ineligible_options( + options: Iterable[PlacementOption], + *, + resource_names: tuple[str, ...], + estimated_latency_tolerance: int, +) -> list[PlacementOption]: + """Drop an option that alone spends more than the global latency slack. + + When two options have the same constrained-resource vector, replacing the + slower one cannot make any complete plan infeasible. If that replacement + gains more than the final plan's entire latency tolerance, the slower + option cannot belong to a latency-equivalent optimum regardless of the + choices made by other groups. This matters when an unconstrained HostPin + dimension leaves many otherwise identical pin prefixes in the frontier. + """ + options = list(options) + best_utility_by_resources: dict[tuple[int, ...], int] = {} + resource_vector_by_key: dict[str, tuple[int, ...]] = {} + for option in options: + resources = tuple( + option.resource_delta_bytes.get(resource_name, 0) + for resource_name in resource_names + ) + resource_vector_by_key[option.option_key] = resources + best_utility_by_resources[resources] = max( + best_utility_by_resources.get(resources, option.estimated_latency_savings), + option.estimated_latency_savings, + ) + return [ + option + for option in options + if option.estimated_latency_savings + >= best_utility_by_resources[resource_vector_by_key[option.option_key]] + - estimated_latency_tolerance + ] + + def optimize_placement( options: Iterable[PlacementOption], *, @@ -179,9 +216,20 @@ def optimize_placement( int(resource_budget_bytes[name]) for name in resource_names ) + # Each local frontier includes the implicit "keep current placement" choice + # when selection is optional. Keep the branch representation compact so the + # global solve does not materialize and repeatedly Pareto-scan a large + # Cartesian product of layer-residency and HostPin prefixes. ordered_groups = [] for group_key in sorted(grouped): - group_options = sorted(grouped[group_key], key=lambda option: option.option_key) + group_options = sorted( + _drop_locally_latency_ineligible_options( + grouped[group_key], + resource_names=resource_names, + estimated_latency_tolerance=estimated_latency_tolerance, + ), + key=lambda option: option.option_key, + ) local_states = [] if not require_selection_from_every_group: local_states.append( @@ -204,102 +252,151 @@ def optimize_placement( ) for option in group_options ) - ordered_groups.append( - [state[3][0] for state in _pareto_prune(local_states) if state[3]] - ) - minimum_remaining: list[tuple[int, ...]] = [ - (0,) * len(resource_names) for _ in range(len(ordered_groups) + 1) - ] - for group_index in range(len(ordered_groups) - 1, -1, -1): - group_options = ordered_groups[group_index] - next_minimum = minimum_remaining[group_index + 1] - minimum_remaining[group_index] = tuple( - next_delta - + ( - min( - option.resource_delta_bytes.get(resource_name, 0) - for option in group_options - ) - if require_selection_from_every_group - else min( - 0, - *( - option.resource_delta_bytes.get(resource_name, 0) - for option in group_options - ), + ordered_groups.append(_pareto_prune(local_states)) + + def _suffix_bounds(groups): + minimum_resources = [(0,) * len(resource_names) for _ in range(len(groups) + 1)] + maximum_utility = [0] * (len(groups) + 1) + minimum_cost = [ + (0,) * placement_cost_dimensions for _ in range(len(groups) + 1) + ] + for group_index in range(len(groups) - 1, -1, -1): + group = groups[group_index] + minimum_resources[group_index] = tuple( + following + min(choice[0][resource_index] for choice in group) + for resource_index, following in enumerate( + minimum_resources[group_index + 1] ) ) - for next_delta, resource_name in zip(next_minimum, resource_names) + maximum_utility[group_index] = maximum_utility[group_index + 1] + max( + choice[1] for choice in group + ) + minimum_cost[group_index] = tuple( + following + min(choice[2][cost_index] for choice in group) + for cost_index, following in enumerate(minimum_cost[group_index + 1]) + ) + return minimum_resources, maximum_utility, minimum_cost + + # First maximize utility. Choices with the same constrained-resource vector + # collapse to the fastest one; HostPin-only soft-cost alternatives therefore + # do not multiply this pass when HostPin is non-binding. + utility_groups = [] + for group in ordered_groups: + best_by_resources: dict[tuple[int, ...], _State] = {} + for choice in group: + incumbent = best_by_resources.get(choice[0]) + if incumbent is None or (-choice[1], choice[2], _selection_key(choice)) < ( + -incumbent[1], + incumbent[2], + _selection_key(incumbent), + ): + best_by_resources[choice[0]] = choice + utility_groups.append( + sorted( + best_by_resources.values(), + key=lambda choice: (-choice[1], choice[0], _selection_key(choice)), + ) ) - frontier: list[_State] = [ - ( - (0,) * len(resource_names), - 0, - (0,) * placement_cost_dimensions, - (), - ) - ] - for group_index, group_options in enumerate(ordered_groups): - candidates = [] if require_selection_from_every_group else list(frontier) - for state in frontier: - for option in group_options: - resource_deltas = tuple( - current + option.resource_delta_bytes.get(resource_name, 0) - for current, resource_name in zip(state[0], resource_names) - ) - candidates.append( - ( - resource_deltas, - state[1] + option.estimated_latency_savings, - tuple( - current + delta - for current, delta in zip( - state[2], - option.placement_cost_bytes - or (0,) * placement_cost_dimensions, - ) - ), - state[3] + (option,), - ) - ) - remaining = minimum_remaining[group_index + 1] - reachable = [ - state - for state in candidates - if all( - used + releasable <= budget - for used, releasable, budget in zip( - state[0], remaining, resource_budgets - ) + minimum_resources, maximum_utility, _ = _suffix_bounds(utility_groups) + best_utility: int | None = None + + def _maximize_utility( + group_index: int, resources: tuple[int, ...], utility: int + ) -> None: + nonlocal best_utility + if ( + best_utility is not None + and utility + maximum_utility[group_index] <= best_utility + ): + return + if any( + used + releasable > budget + for used, releasable, budget in zip( + resources, minimum_resources[group_index], resource_budgets + ) + ): + return + if group_index == len(utility_groups): + best_utility = utility + return + for choice in utility_groups[group_index]: + _maximize_utility( + group_index + 1, + tuple(left + right for left, right in zip(resources, choice[0])), + utility + choice[1], ) - ] - frontier = _pareto_prune(reachable) - frontier = [ - state - for state in frontier - if all(used <= budget for used, budget in zip(state[0], resource_budgets)) - ] - if not frontier: + _maximize_utility(0, (0,) * len(resource_names), 0) + if best_utility is None: raise NoFeasiblePlacementError("no placement satisfies all resource budgets") - best_utility = max(state[1] for state in frontier) - latency_equivalent = [ - state - for state in frontier - if state[1] >= best_utility - estimated_latency_tolerance - ] - best = min( - latency_equivalent, - key=lambda state: ( - state[2], - sum(state[0]), - state[0], - -state[1], - _selection_key(state), - ), + # Then minimize the lexicographic soft placement cost among every plan in + # the single global latency-equivalence window. Depth-first bounds avoid the + # quadratic global Pareto scans that large layerwise frontiers otherwise + # trigger while preserving the exact final ordering. + minimum_resources, maximum_utility, minimum_cost = _suffix_bounds(ordered_groups) + utility_floor = best_utility - estimated_latency_tolerance + best: _State | None = None + + def _minimize_cost( + group_index: int, + resources: tuple[int, ...], + utility: int, + cost: tuple[int, ...], + selections: tuple[PlacementOption, ...], + ) -> None: + nonlocal best + if utility + maximum_utility[group_index] < utility_floor: + return + if any( + used + releasable > budget + for used, releasable, budget in zip( + resources, minimum_resources[group_index], resource_budgets + ) + ): + return + if best is not None: + lower_cost = tuple( + current + remaining + for current, remaining in zip(cost, minimum_cost[group_index]) + ) + if lower_cost > best[2]: + return + if group_index == len(ordered_groups): + candidate = (resources, utility, cost, selections) + if best is None or ( + candidate[2], + sum(candidate[0]), + candidate[0], + -candidate[1], + _selection_key(candidate), + ) < ( + best[2], + sum(best[0]), + best[0], + -best[1], + _selection_key(best), + ): + best = candidate + return + for choice in ordered_groups[group_index]: + _minimize_cost( + group_index + 1, + tuple(left + right for left, right in zip(resources, choice[0])), + utility + choice[1], + tuple(left + right for left, right in zip(cost, choice[2])), + selections + choice[3], + ) + + _minimize_cost( + 0, + (0,) * len(resource_names), + 0, + (0,) * placement_cost_dimensions, + (), ) + assert best is not None full_resource_deltas = { resource_name: sum( option.resource_delta_bytes.get(resource_name, 0) for option in best[3] diff --git a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py index 26df22ce3ff2..d94f0a99f09d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py @@ -1,8 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 +import itertools +import random + import pytest from sglang.multimodal_gen.runtime.managers.memory_managers.placement_budget import ( + NoFeasiblePlacementError, PlacementOption, optimize_placement, ) @@ -27,6 +31,75 @@ def _option( ) +def _exhaustive_plan_key( + options: list[PlacementOption], + *, + resource_budget_bytes: dict[str, int], + estimated_latency_tolerance: int, + require_selection_from_every_group: bool, +): + resource_names = tuple(sorted(resource_budget_bytes)) + groups: dict[str, list[PlacementOption | None]] = {} + for option in options: + groups.setdefault(option.group_key, []).append(option) + choices = [] + for group_key in sorted(groups): + group = sorted(groups[group_key], key=lambda option: option.option_key) + if not require_selection_from_every_group: + group = [None, *group] + choices.append(group) + + feasible = [] + for selections in itertools.product(*choices): + selected = tuple(option for option in selections if option is not None) + resources = tuple( + sum(option.resource_delta_bytes.get(name, 0) for option in selected) + for name in resource_names + ) + if any( + used > resource_budget_bytes[name] + for name, used in zip(resource_names, resources) + ): + continue + utility = sum(option.estimated_latency_savings for option in selected) + cost_dimensions = max( + (len(option.placement_cost_bytes) for option in options), default=0 + ) + cost = tuple( + sum( + ( + option.placement_cost_bytes[index] + if option.placement_cost_bytes + else 0 + ) + for option in selected + ) + for index in range(cost_dimensions) + ) + feasible.append((selected, resources, utility, cost)) + + if not feasible: + return None + best_utility = max(state[2] for state in feasible) + utility_floor = best_utility - estimated_latency_tolerance + selected, resources, utility, cost = min( + (state for state in feasible if state[2] >= utility_floor), + key=lambda state: ( + state[3], + sum(state[1]), + state[1], + -state[2], + tuple(option.option_key for option in state[0]), + ), + ) + return ( + tuple(option.option_key for option in selected), + resources, + utility, + cost, + ) + + def test_joint_plan_beats_independent_greedy_allocations(): plan = optimize_placement( [ @@ -448,3 +521,92 @@ def test_nonbinding_resource_is_removed_without_losing_full_accounting(): "gpu:rank0:runtime": 5, "hostpin:node0": 80, } + + +def test_nonbinding_pin_prefixes_keep_only_globally_latency_eligible_states(): + plan = optimize_placement( + [ + PlacementOption( + group_key="dit", + option_key="dit-pin-none", + resource_delta_bytes={"gpu": 5, "hostpin": 0}, + estimated_latency_savings=800, + placement_cost_bytes=(5, 0), + ), + PlacementOption( + group_key="dit", + option_key="dit-pin-most", + resource_delta_bytes={"gpu": 5, "hostpin": 80}, + estimated_latency_savings=950, + placement_cost_bytes=(5, 80), + ), + PlacementOption( + group_key="dit", + option_key="dit-pin-all", + resource_delta_bytes={"gpu": 5, "hostpin": 100}, + estimated_latency_savings=1_000, + placement_cost_bytes=(5, 100), + ), + ], + resource_budget_bytes={"gpu": 5, "hostpin": 100}, + estimated_latency_tolerance=100, + ) + + # No-pin is over the global tolerance and cannot be selected. Pin-most is + # still latency-equivalent, so the soft memory tie-break chooses it. + assert [option.option_key for option in plan.selections] == ["dit-pin-most"] + + +def test_branch_and_bound_matches_exhaustive_multiresource_search(): + rng = random.Random(20260827) + resource_names = ("load-vram", "runtime-vram", "hostpin") + + for _ in range(200): + options = [] + for group_index in range(rng.randint(1, 5)): + for option_index in range(rng.randint(1, 4)): + options.append( + PlacementOption( + group_key=f"group-{group_index}", + option_key=f"group-{group_index}:option-{option_index}", + resource_delta_bytes={ + name: rng.randint(-5, 8) for name in resource_names + }, + estimated_latency_savings=rng.randint(-5, 20), + placement_cost_bytes=(rng.randint(0, 10), rng.randint(0, 10)), + ) + ) + budgets = {name: rng.randint(-2, 12) for name in resource_names} + tolerance = rng.randint(0, 5) + require_every_group = rng.choice((False, True)) + expected = _exhaustive_plan_key( + options, + resource_budget_bytes=budgets, + estimated_latency_tolerance=tolerance, + require_selection_from_every_group=require_every_group, + ) + + if expected is None: + with pytest.raises(NoFeasiblePlacementError): + optimize_placement( + options, + resource_budget_bytes=budgets, + estimated_latency_tolerance=tolerance, + require_selection_from_every_group=require_every_group, + ) + continue + + plan = optimize_placement( + options, + resource_budget_bytes=budgets, + estimated_latency_tolerance=tolerance, + require_selection_from_every_group=require_every_group, + ) + expected_keys, expected_resources, expected_utility, expected_cost = expected + assert tuple(option.option_key for option in plan.selections) == expected_keys + assert ( + tuple(plan.resource_delta_bytes[name] for name in sorted(resource_names)) + == expected_resources + ) + assert plan.estimated_latency_savings == expected_utility + assert plan.placement_cost_bytes == expected_cost From 2a4125180a54531d5ab1b49de066c39788ea5236 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 07:02:56 +0800 Subject: [PATCH 36/90] [diffusion] add lazy layerwise placement targets --- .../runtime/managers/gpu_worker.py | 28 +- .../memory_managers/auto_residency.py | 463 +++++++++++++++--- .../memory_managers/host_memory_budget.py | 25 +- .../memory_managers/layerwise_offload.py | 103 +++- .../runtime/server_args/server_args.py | 61 ++- .../test/unit/test_auto_residency.py | 221 ++++++++- .../test/unit/test_layerwise_offload.py | 4 + .../test/unit/test_performance_metrics.py | 6 + 8 files changed, 784 insertions(+), 127 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 5bde6e0a7e9d..0013ee70ddc4 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -50,7 +50,6 @@ save_outputs, ) from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( - AUTO_RESIDENCY_FEATURE_NAME, GIB_BYTES, PLACEMENT_STATUS_ADJUSTED, PLACEMENT_STATUS_ROLLBACK_FAILED, @@ -83,6 +82,9 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( peek_global_component_residency_manager, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( + RESIDENT, +) from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( HOST_COPY_RESERVE_BYTES, host_memory_available_bytes, @@ -202,6 +204,10 @@ def __init__( self._auto_residency_warmup_records: list[WarmupMemoryRecord] = [] self._auto_residency_applied: list[AppliedResidencyChange] = [] self._auto_residency_round_sizes: list[int] = [] + # Keep every round on one serving objective. Re-warm measurements + # validate and refine memory constraints, but placement-dependent + # latency must not rewrite the utility function and cause oscillation. + self._auto_residency_reference_request_duration_ns: int | None = None # default workload resolved once for the per-request residency hint self._cached_default_workload: DefaultWorkload | None = None self._cached_default_workload_failed = False @@ -1463,15 +1469,29 @@ def _build_auto_residency_report( auto_resident_components={ name for name in self.pipeline.modules - if self.server_args.component_residency_requirement(name) - == AUTO_RESIDENCY_FEATURE_NAME + if self.server_args.auto_residency_mode(name) == RESIDENT }, + layerwise_tuning_of=lambda name, dit_group: ( + self.server_args.layerwise_tuning_for(name, dit_group=dit_group) + ), + pin_cpu_memory=self.server_args.pin_cpu_memory, ) - estimated_request_duration_ns, _, _ = estimate_default_workload_timing( + measured_request_duration_ns, _, _ = estimate_default_workload_timing( records=records, target_units=target_units, target_num_inference_steps=workload.num_inference_steps, ) + reference_request_duration_ns = self.__dict__.get( + "_auto_residency_reference_request_duration_ns" + ) + if reference_request_duration_ns is None and measured_request_duration_ns > 0: + reference_request_duration_ns = measured_request_duration_ns + self._auto_residency_reference_request_duration_ns = ( + reference_request_duration_ns + ) + estimated_request_duration_ns = ( + reference_request_duration_ns or measured_request_duration_ns + ) candidate_latency_savings_ns = estimate_candidate_latency_savings_ns( candidates=candidates, request_duration_ns=estimated_request_duration_ns, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 15fb1c6a020d..1713731c31e4 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -25,6 +25,7 @@ from __future__ import annotations import statistics +from itertools import chain from typing import TYPE_CHECKING, Callable, Iterable, Mapping, Sequence import msgspec @@ -42,10 +43,12 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( HOST_COPY_RESERVE_BYTES, host_memory_available_bytes, + tensor_storage_bytes, ) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, compute_streamed_layers, + estimate_layer_weight_bytes, is_layerwise_offloaded_module, iter_materialized_weights, release_unused_pinned_memory, @@ -90,8 +93,6 @@ MAX_LATENCY_EQUIVALENCE_NS = 100_000_000 LATENCY_EQUIVALENCE_FRACTION = 0.01 -AUTO_RESIDENCY_FEATURE_NAME = "auto residency" - PLACEMENT_STATUS_SKIPPED = "skipped" PLACEMENT_STATUS_ADJUSTED = "adjusted" PLACEMENT_STATUS_ROLLED_BACK = "rolled_back" @@ -256,7 +257,8 @@ class AppliedResidencyChange(msgspec.Struct, frozen=True): previous_layerwise_resident_layers: tuple[int, ...] | None = None previous_layerwise_pinned_layers: tuple[tuple[int, ...], ...] | None = None previous_layerwise_offload_enabled: bool = True - previous_required_resident: bool = False + previous_layerwise_configured: bool = True + previous_auto_residency_mode: str | None = None pinned_host_changed: bool = False @@ -603,14 +605,12 @@ def estimate_workload_phase_peaks( def _module_weight_bytes(module: nn.Module) -> int: """Full weight+buffer footprint, reading through layerwise CPU buffers.""" - param_bytes = sum( - tensor.numel() * tensor.element_size() - for _, tensor in iter_materialized_weights(module) - ) - buffer_bytes = sum( - tensor.numel() * tensor.element_size() for tensor in module.buffers() + return tensor_storage_bytes( + chain( + (tensor for _, tensor in iter_materialized_weights(module)), + module.buffers(), + ) ) - return param_bytes + buffer_bytes def component_runtime_weight_bytes(modules: Mapping[str, object]) -> dict[str, int]: @@ -794,6 +794,206 @@ def _layerwise_resident_targets(managers: Sequence) -> list[tuple[int, ...]]: return sorted(targets) +class _EstimatedLayerwiseManager: + """Size-only layer group used before a real offload manager exists.""" + + def __init__( + self, + *, + layer_bytes: dict[int, int], + prefetch_size: int, + residency_policy: str, + pin_cpu_memory: bool, + ) -> None: + self._layer_bytes = layer_bytes + self.num_layers = len(layer_bytes) + self.prefetch_size = min(max(1, prefetch_size), self.num_layers) + self.residency_policy = residency_policy + self.pin_cpu_memory = pin_cpu_memory + self.resident_layers = 0 + + def layer_weight_bytes(self) -> dict[int, int]: + return self._layer_bytes + + def layer_host_store_bytes(self) -> dict[int, int]: + return self._layer_bytes + + def pinned_host_weight_bytes(self) -> int: + return 0 + + def pinned_layer_indices(self) -> tuple[int, ...]: + return () + + def pinnable_layer_indices(self) -> tuple[int, ...]: + if not self.pin_cpu_memory: + return () + return tuple(self._layer_bytes) + + def resident_weight_bytes(self, resident_layers: int | None = None) -> int: + count = self.resident_layers if resident_layers is None else resident_layers + streamed = set( + compute_streamed_layers( + num_layers=self.num_layers, + resident_layers=count, + policy=self.residency_policy, + ) + ) + return sum( + weight_bytes + for layer_index, weight_bytes in self._layer_bytes.items() + if layer_index not in streamed + ) + + def peak_managed_device_weight_bytes( + self, resident_layers: int | None = None + ) -> int: + count = self.resident_layers if resident_layers is None else resident_layers + streamed = compute_streamed_layers( + num_layers=self.num_layers, + resident_layers=count, + policy=self.residency_policy, + ) + resident_bytes = self.resident_weight_bytes(count) + copy_window = min(self.prefetch_size, len(streamed)) + streamed_window_bytes = sum( + sorted( + (self._layer_bytes.get(index, 0) for index in streamed), + reverse=True, + )[:copy_window] + ) + return resident_bytes + streamed_window_bytes + + +def _estimate_unconfigured_layerwise_managers( + *, + module: LayerwiseOffloadableModuleMixin, + prefetch_value: float, + residency_policy: str, + pin_cpu_memory: bool, +) -> list[_EstimatedLayerwiseManager]: + """Read an unloaded module's block sizes without allocating CPU stores.""" + named_modules = dict(module.named_modules()) + estimates = [] + for layer_name in module.layer_names: + layers = named_modules.get(layer_name) + if not isinstance(layers, (nn.ModuleList, nn.Sequential)) or not layers: + continue + num_layers = len(layers) + if prefetch_value < 1.0: + prefetch_size = 1 + int(round(prefetch_value * (num_layers - 1))) + else: + prefetch_size = int(prefetch_value) + estimates.append( + _EstimatedLayerwiseManager( + layer_bytes=estimate_layer_weight_bytes(layers), + prefetch_size=prefetch_size, + residency_policy=residency_policy, + pin_cpu_memory=pin_cpu_memory, + ) + ) + return estimates + + +def _unconfigured_layerwise_targets( + *, + component_name: str, + module: LayerwiseOffloadableModuleMixin, + current_mode: str, + full_weight_bytes: int, + num_inference_steps: int, + prefetch_value: float, + residency_policy: str, + pin_cpu_memory: bool, +) -> tuple[list[ResidencyTarget], int]: + """Virtual layerwise frontier for a component still using coarse offload. + + The first selected virtual state configures real managers lazily. It starts + pageable and is remeasured immediately; the next fixed-point round then + exposes the exact HostPin frontier from the realized stores. + """ + managers = _estimate_unconfigured_layerwise_managers( + module=module, + prefetch_value=prefetch_value, + residency_policy=residency_policy, + pin_cpu_memory=pin_cpu_memory, + ) + if not managers: + return [], 0 + managed_weight_bytes = sum( + sum(manager.layer_weight_bytes().values()) for manager in managers + ) + if managed_weight_bytes <= 0: + return [], 0 + unmanaged_weight_bytes = max(0, full_weight_bytes - managed_weight_bytes) + uses_per_request = ( + max(1, num_inference_steps) if is_dit_component_name(component_name) else 1 + ) + empty_pins = tuple(() for _ in managers) + maximum_transfer_work = _layerwise_transfer_work_bytes( + managers=managers, + resident_layers=tuple(0 for _ in managers), + pinned_layers=empty_pins, + uses_per_streamed_layer=uses_per_request, + ) + current_resident = current_mode == RESIDENT + current_inactive_bytes = full_weight_bytes if current_resident else 0 + initial_managed_bytes = sum( + manager.peak_managed_device_weight_bytes(0) for manager in managers + ) + host_materialize_scratch = max( + ( + weight_bytes + for manager in managers + for weight_bytes in manager.layer_host_store_bytes().values() + ), + default=0, + ) + targets = [] + for resident_layers in _layerwise_resident_targets(managers): + resident_bytes = sum( + manager.resident_weight_bytes(count) + for manager, count in zip(managers, resident_layers) + ) + active_managed_bytes = sum( + manager.peak_managed_device_weight_bytes(count) + for manager, count in zip(managers, resident_layers) + ) + transfer_work = _layerwise_transfer_work_bytes( + managers=managers, + resident_layers=resident_layers, + pinned_layers=empty_pins, + uses_per_streamed_layer=uses_per_request, + ) + targets.append( + ResidencyTarget( + component_name=component_name, + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=resident_bytes, + h2d_bytes_per_request=maximum_transfer_work - transfer_work, + target_layerwise_resident_layers=resident_layers, + target_layerwise_pinned_layers=empty_pins, + host_materialize_scratch_bytes=host_materialize_scratch, + device_transition_delta_bytes=( + unmanaged_weight_bytes + + initial_managed_bytes + - current_inactive_bytes + ), + active_device_delta_bytes=( + unmanaged_weight_bytes + active_managed_bytes - full_weight_bytes + ), + present_device_delta_bytes=( + unmanaged_weight_bytes + active_managed_bytes - full_weight_bytes + ), + inactive_device_delta_bytes=( + unmanaged_weight_bytes - current_inactive_bytes + ), + target_device_weight_bytes=unmanaged_weight_bytes + resident_bytes, + ) + ) + return targets, maximum_transfer_work + + def _layerwise_host_transition_bytes( *, managers: Sequence, @@ -825,6 +1025,8 @@ def collect_residency_targets( allow_host_pin_reallocation: bool = True, mixed_dtype_components: Iterable[str] = (), auto_resident_components: Iterable[str] = (), + layerwise_tuning_of: Callable[[str, bool], tuple[float, float, str]] | None = None, + pin_cpu_memory: bool = True, ) -> list[ResidencyTarget]: """Build complete target-state frontiers for auto-managed components. @@ -857,7 +1059,17 @@ def collect_residency_targets( # A loader or another runtime feature owns this hard requirement. continue - if baseline_mode == COMPONENT_OFFLOAD: + has_layerwise_managers = isinstance( + module, LayerwiseOffloadableModuleMixin + ) and bool(module.layerwise_offload_managers) + if has_layerwise_managers and current_mode == COMPONENT_OFFLOAD: + # Temporary managers used by offload-during-compile do not own the + # serving placement. Only an effective layerwise mode proves that + # these managers form the current serving frontier. + continue + frontier_mode = LAYERWISE_OFFLOAD if has_layerwise_managers else baseline_mode + + if frontier_mode == COMPONENT_OFFLOAD: if current_mode not in (COMPONENT_OFFLOAD, RESIDENT): continue if is_layerwise_offloaded_module(module): @@ -869,6 +1081,35 @@ def collect_residency_targets( if weight_bytes <= 0: continue current_resident = current_mode == RESIDENT + virtual_targets: list[ResidencyTarget] = [] + maximum_transfer_work = 0 + if ( + isinstance(module, LayerwiseOffloadableModuleMixin) + and layerwise_tuning_of is not None + ): + prefetch_value, _, residency_policy = layerwise_tuning_of( + name, module.layerwise_offload_dit_group_enabled + ) + virtual_targets, maximum_transfer_work = ( + _unconfigured_layerwise_targets( + component_name=name, + module=module, + current_mode=current_mode, + full_weight_bytes=weight_bytes, + num_inference_steps=num_inference_steps, + prefetch_value=prefetch_value, + residency_policy=residency_policy, + pin_cpu_memory=pin_cpu_memory, + ) + ) + uses_per_request = ( + max(1, num_inference_steps) if is_dit_component_name(name) else 1 + ) + component_transfer_work = ( + maximum_transfer_work // uses_per_request + if maximum_transfer_work + else 0 + ) candidates.extend( [ ResidencyTarget( @@ -876,7 +1117,9 @@ def collect_residency_targets( residency_mode=baseline_mode, target_residency_mode=COMPONENT_OFFLOAD, target_resident_weight_bytes=0, - h2d_bytes_per_request=0, + h2d_bytes_per_request=max( + 0, maximum_transfer_work - component_transfer_work + ), host_materialize_scratch_bytes=( weight_bytes if current_resident else 0 ), @@ -897,7 +1140,7 @@ def collect_residency_targets( residency_mode=baseline_mode, target_residency_mode=RESIDENT, target_resident_weight_bytes=weight_bytes, - h2d_bytes_per_request=weight_bytes, + h2d_bytes_per_request=(maximum_transfer_work or weight_bytes), permanent_residency=True, device_transition_delta_bytes=0, active_device_delta_bytes=0, @@ -910,6 +1153,7 @@ def collect_residency_targets( ), ] ) + candidates.extend(virtual_targets) continue if not isinstance(module, LayerwiseOffloadableModuleMixin): @@ -953,12 +1197,53 @@ def collect_residency_targets( current_pinned_bytes = sum( manager.pinned_host_weight_bytes() for manager in managers ) - maximum_transfer_work = _layerwise_transfer_work_bytes( + layerwise_maximum_transfer_work = _layerwise_transfer_work_bytes( managers=managers, resident_layers=tuple(0 for _ in managers), pinned_layers=tuple(() for _ in managers), uses_per_streamed_layer=uses_per_request, ) + full_weight_bytes = _module_weight_bytes(module) + unmanaged_weight_bytes = max(0, full_weight_bytes - managed_weight_bytes) + component_transfer_work = 2 * full_weight_bytes + maximum_transfer_work = max( + layerwise_maximum_transfer_work, component_transfer_work + ) + + empty_resident_layers = tuple(0 for _ in managers) + empty_pinned_layers = tuple(() for _ in managers) + unpin_scratch, _ = _layerwise_host_transition_bytes( + managers=managers, + current_pinned_layers=current_pinned_layers, + target_pinned_layers=empty_pinned_layers, + ) + candidates.append( + ResidencyTarget( + component_name=name, + residency_mode=frontier_mode, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=max( + 0, maximum_transfer_work - component_transfer_work + ), + target_layerwise_resident_layers=empty_resident_layers, + target_layerwise_pinned_layers=empty_pinned_layers, + pinned_host_delta_bytes=-current_pinned_bytes, + host_unpin_scratch_bytes=unpin_scratch, + device_transition_delta_bytes=( + 0 if current_permanent else managed_weight_bytes + ), + active_device_delta_bytes=( + full_weight_bytes - current_peak_device_bytes + ), + present_device_delta_bytes=( + full_weight_bytes - current_peak_device_bytes + ), + inactive_device_delta_bytes=-current_inactive_device_bytes, + target_device_weight_bytes=0, + target_pinned_host_bytes=0, + ) + ) for target_resident_layers in _layerwise_resident_targets(managers): target_resident_bytes = sum( @@ -1001,7 +1286,7 @@ def collect_residency_targets( candidates.append( ResidencyTarget( component_name=name, - residency_mode=baseline_mode, + residency_mode=frontier_mode, target_residency_mode=LAYERWISE_OFFLOAD, target_resident_weight_bytes=target_resident_bytes, h2d_bytes_per_request=( @@ -1027,7 +1312,9 @@ def collect_residency_targets( and target_resident_layers == current_resident_layers and target_pinned_layers == current_pinned_layers ), - target_device_weight_bytes=target_resident_bytes, + target_device_weight_bytes=( + unmanaged_weight_bytes + target_resident_bytes + ), target_pinned_host_bytes=target_pinned_bytes, ) ) @@ -1063,7 +1350,7 @@ def collect_residency_targets( candidates.append( ResidencyTarget( component_name=name, - residency_mode=baseline_mode, + residency_mode=frontier_mode, target_residency_mode=RESIDENT, target_resident_weight_bytes=managed_weight_bytes, h2d_bytes_per_request=maximum_transfer_work, @@ -1092,7 +1379,7 @@ def collect_residency_targets( current_placement=( current_permanent and permanent_pins == current_pinned_layers ), - target_device_weight_bytes=managed_weight_bytes, + target_device_weight_bytes=full_weight_bytes, target_pinned_host_bytes=permanent_pinned_bytes, ) ) @@ -1549,7 +1836,11 @@ def apply_residency_changes( raise RuntimeError( f"residency target {candidate.component_name!r} is missing" ) - if candidate.residency_mode == LAYERWISE_OFFLOAD and not isinstance( + layerwise_transition = ( + candidate.residency_mode == LAYERWISE_OFFLOAD + or candidate.target_mode() == LAYERWISE_OFFLOAD + ) + if layerwise_transition and not isinstance( module, LayerwiseOffloadableModuleMixin ): raise RuntimeError( @@ -1564,24 +1855,32 @@ def apply_residency_changes( "layerwise-managed between planning and apply" ) target_modules[candidate.component_name] = module - previous_required_resident = ( - server_args.component_residency_requirement(candidate.component_name) - == AUTO_RESIDENCY_FEATURE_NAME + previous_auto_mode = server_args.auto_residency_mode( + candidate.component_name ) - if candidate.residency_mode == LAYERWISE_OFFLOAD: + if layerwise_transition: if not isinstance(module, LayerwiseOffloadableModuleMixin): raise RuntimeError( f"partial residency target {candidate.component_name!r} " "lost its layerwise offload capability" ) - previous = tuple( - manager.resident_layers - for manager in module.layerwise_offload_managers + previous_configured = bool(module.layerwise_offload_managers) + previous = ( + tuple( + manager.resident_layers + for manager in module.layerwise_offload_managers + ) + if previous_configured + else () + ) + previous_pinned = ( + module.layerwise_pinned_layers() if previous_configured else () + ) + enabled = ( + {manager.enabled for manager in module.layerwise_offload_managers} + if previous_configured + else {False} ) - previous_pinned = module.layerwise_pinned_layers() - enabled = { - manager.enabled for manager in module.layerwise_offload_managers - } if len(enabled) != 1: raise RuntimeError( f"placement target {candidate.component_name!r} has mixed " @@ -1594,9 +1893,11 @@ def apply_residency_changes( previous_layerwise_resident_layers=previous, previous_layerwise_pinned_layers=previous_pinned, previous_layerwise_offload_enabled=enabled.pop(), - previous_required_resident=previous_required_resident, + previous_layerwise_configured=previous_configured, + previous_auto_residency_mode=previous_auto_mode, pinned_host_changed=( - candidate.target_layerwise_pinned_layers != previous_pinned + previous_configured + and candidate.target_layerwise_pinned_layers != previous_pinned ), ) if candidate.target_layerwise_pinned_layers is None: @@ -1609,11 +1910,12 @@ def apply_residency_changes( snapshots[candidate.component_name] = AppliedResidencyChange( component_name=candidate.component_name, residency_mode=candidate.residency_mode, - previous_required_resident=previous_required_resident, + previous_auto_residency_mode=previous_auto_mode, ) pinning_changes = [ candidate for candidate in ordered_changes + if snapshots[candidate.component_name].previous_layerwise_configured if candidate.target_layerwise_pinned_layers is not None and candidate.target_layerwise_pinned_layers != previous_pins_by_component[candidate.component_name] @@ -1672,38 +1974,54 @@ def apply_residency_changes( applied_names.add(candidate.component_name) module = target_modules[candidate.component_name] target_mode = candidate.target_mode() - if candidate.residency_mode == LAYERWISE_OFFLOAD: + if ( + candidate.residency_mode == LAYERWISE_OFFLOAD + or target_mode == LAYERWISE_OFFLOAD + ): assert isinstance(module, LayerwiseOffloadableModuleMixin) + snapshot = snapshots[candidate.component_name] + if not snapshot.previous_layerwise_configured: + module.configure_layerwise_offload( + server_args, + pin_budget=server_args.host_pin_budget(), + component_name=candidate.component_name, + pin_during_initialization=False, + ) + if not module.layerwise_offload_managers: + raise RuntimeError( + f"residency target {candidate.component_name!r} " + "did not configure any layerwise groups" + ) module.set_layerwise_pinned_layers( candidate.target_layerwise_pinned_layers or () ) - if target_mode == RESIDENT: - server_args.require_component_resident( - candidate.component_name, - feature_name=AUTO_RESIDENCY_FEATURE_NAME, + if target_mode == COMPONENT_OFFLOAD: + module.remove_layerwise_offload() + server_args.set_auto_residency_mode( + candidate.component_name, COMPONENT_OFFLOAD + ) + module.to("cpu") + elif target_mode == RESIDENT: + server_args.set_auto_residency_mode( + candidate.component_name, RESIDENT ) if is_layerwise_offloaded_module(module): module.disable_offload() else: if not is_layerwise_offloaded_module(module): module.enable_offload() - server_args.release_required_component_residency( - candidate.component_name, - feature_name=AUTO_RESIDENCY_FEATURE_NAME, + server_args.set_auto_residency_mode( + candidate.component_name, LAYERWISE_OFFLOAD ) module.set_layerwise_resident_layer_counts( candidate.target_layerwise_resident_layers ) continue if target_mode == RESIDENT: - server_args.require_component_resident( - candidate.component_name, - feature_name=AUTO_RESIDENCY_FEATURE_NAME, - ) + server_args.set_auto_residency_mode(candidate.component_name, RESIDENT) elif target_mode == COMPONENT_OFFLOAD: - server_args.release_required_component_residency( - candidate.component_name, - feature_name=AUTO_RESIDENCY_FEATURE_NAME, + server_args.set_auto_residency_mode( + candidate.component_name, COMPONENT_OFFLOAD ) module.to("cpu") else: @@ -1763,9 +2081,34 @@ def rollback_residency_changes( if adjustment.previous_layerwise_resident_layers is not None: if not isinstance(module, LayerwiseOffloadableModuleMixin): raise RuntimeError("lost layerwise offload capability") + if not adjustment.previous_layerwise_configured: + module.remove_layerwise_offload() + if adjustment.previous_auto_residency_mode is None: + server_args.clear_auto_residency_mode(adjustment.component_name) + else: + server_args.set_auto_residency_mode( + adjustment.component_name, + adjustment.previous_auto_residency_mode, + ) + if ( + server_args.residency_mode(adjustment.component_name) + != RESIDENT + ): + module.to("cpu") + continue currently_enabled = { manager.enabled for manager in module.layerwise_offload_managers } + if not currently_enabled: + module.configure_layerwise_offload( + server_args, + pin_budget=server_args.host_pin_budget(), + component_name=adjustment.component_name, + pin_during_initialization=False, + ) + currently_enabled = { + manager.enabled for manager in module.layerwise_offload_managers + } if len(currently_enabled) != 1: raise RuntimeError("layerwise managers have mixed enabled states") current_enabled = currently_enabled.pop() @@ -1774,15 +2117,12 @@ def rollback_residency_changes( module.enable_offload() elif not previous_enabled and current_enabled: module.disable_offload() - if adjustment.previous_required_resident: - server_args.require_component_resident( - adjustment.component_name, - feature_name=AUTO_RESIDENCY_FEATURE_NAME, - ) + if adjustment.previous_auto_residency_mode is None: + server_args.clear_auto_residency_mode(adjustment.component_name) else: - server_args.release_required_component_residency( + server_args.set_auto_residency_mode( adjustment.component_name, - feature_name=AUTO_RESIDENCY_FEATURE_NAME, + adjustment.previous_auto_residency_mode, ) module.restore_layerwise_resident_layers( adjustment.previous_layerwise_resident_layers @@ -1793,15 +2133,12 @@ def rollback_residency_changes( adjustment.previous_layerwise_pinned_layers ) continue - if adjustment.previous_required_resident: - server_args.require_component_resident( - adjustment.component_name, - feature_name=AUTO_RESIDENCY_FEATURE_NAME, - ) + if adjustment.previous_auto_residency_mode is None: + server_args.clear_auto_residency_mode(adjustment.component_name) else: - server_args.release_required_component_residency( + server_args.set_auto_residency_mode( adjustment.component_name, - feature_name=AUTO_RESIDENCY_FEATURE_NAME, + adjustment.previous_auto_residency_mode, ) if not isinstance(module, nn.Module): continue @@ -1810,7 +2147,7 @@ def rollback_residency_changes( # offload_during_compile window); moving it to CPU behind the # manager would strand its bookkeeping. pass - elif not adjustment.previous_required_resident: + elif server_args.residency_mode(adjustment.component_name) != RESIDENT: module.to("cpu") except Exception as e: errors.append(f"{adjustment.component_name}: {describe_error(e)}") diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py index 37a6fa9d8807..7edac38d1cd9 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py @@ -14,8 +14,11 @@ """ import os +from collections.abc import Iterable import psutil +import torch +from torch.distributed.tensor import DTensor from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -272,24 +275,32 @@ def pin_benefit_bytes(*, weight_bytes: int, uses_per_request: int) -> int: return max(0, weight_bytes) * max(1, uses_per_request) -def module_weight_bytes(module) -> int: - """Bytes of parameters and buffers a module would hand to the host.""" - seen: set[int] = set() +def tensor_storage_bytes(tensors: Iterable[torch.Tensor]) -> int: + """Physical bytes backing tensors, deduplicated across aliases and views.""" + seen: set[tuple[torch.device, int]] = set() total = 0 - for tensor in list(module.parameters()) + list(module.buffers()): + for tensor in tensors: + if isinstance(tensor, DTensor): + tensor = tensor.to_local() try: storage = tensor.untyped_storage() pointer = storage.data_ptr() storage_bytes = storage.nbytes() - except RuntimeError: + except (AttributeError, RuntimeError): continue - if pointer == 0 or pointer in seen: + storage_key = (tensor.device, pointer) + if pointer == 0 or storage_key in seen: continue - seen.add(pointer) + seen.add(storage_key) total += storage_bytes return total +def module_weight_bytes(module) -> int: + """Physical bytes of parameters and buffers a module would hand to the host.""" + return tensor_storage_bytes((*module.parameters(), *module.buffers())) + + def describe_host_memory() -> str: """One line for startup logs: what the cap is and where it came from.""" capped = cgroup_memory_limit_bytes() 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 c07088c65985..a161a9eeccd9 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 @@ -45,6 +45,55 @@ logger = init_logger(__name__) +def _align_numel_offset( + offset: int, dtype: torch.dtype, alignment_bytes: int = 32 +) -> int: + element_size = torch.empty((), dtype=dtype).element_size() + alignment_numel = max(1, alignment_bytes // element_size) + remainder = offset % alignment_numel + return offset if remainder == 0 else offset + alignment_numel - remainder + + +def estimate_layer_weight_bytes( + layers: torch.nn.ModuleList | torch.nn.Sequential, +) -> dict[int, int]: + """Physical per-layer store sizes using the real manager's packing rules.""" + grouped: dict[int, dict[torch.dtype, list[torch.Tensor]]] = {} + for name, tensor in layers.named_parameters(): + layer_index_text, separator, _ = name.partition(".") + if not separator or not layer_index_text.isdigit(): + continue + layer_index = int(layer_index_text) + local_tensor = tensor.to_local() if isinstance(tensor, DTensor) else tensor + grouped.setdefault(layer_index, {}).setdefault(local_tensor.dtype, []).append( + local_tensor + ) + + totals = {layer_index: 0 for layer_index in range(len(layers))} + for layer_index, dtype_to_tensors in grouped.items(): + for dtype, tensors in dtype_to_tensors.items(): + contiguous_numel = 0 + for tensor in tensors: + if tensor.is_contiguous(): + contiguous_numel = _align_numel_offset(contiguous_numel, dtype) + contiguous_numel += tensor.numel() + continue + storage_numel = ( + 0 + if tensor.numel() == 0 + else 1 + + sum( + (size - 1) * abs(stride) + for size, stride in zip(tensor.shape, tensor.stride()) + ) + ) + totals[layer_index] += storage_numel * tensor.element_size() + totals[layer_index] += ( + contiguous_numel * torch.empty((), dtype=dtype).element_size() + ) + return totals + + def release_unused_pinned_memory() -> None: """Return unreferenced cached HostPin blocks to the CUDA driver. @@ -382,6 +431,7 @@ def __init__( residency_policy: str = RESIDENCY_POLICY_LEADING, pin_budget: HostPinBudget | None = None, pin_component_name: str = "layerwise offload", + pin_during_initialization: bool = True, ) -> None: self.model = model self.layers_attr_str = layers_attr_str @@ -397,6 +447,7 @@ def __init__( # was never reached. A private budget reads the same host limit. self._pin_budget = pin_budget if pin_budget is not None else HostPinBudget() self._pin_component_name = pin_component_name + self._pin_during_initialization = pin_during_initialization self._layer_hosting: Dict[int, str] = {} # an explicit MPS zero avoids staging the next layer alongside the # active one; MPS has no transfer overlap to recover from that cost @@ -536,21 +587,6 @@ def _get_shared_empty_tensor_for_target( ) -> torch.Tensor: return self._wrap_for_target(target, self._get_shared_empty_tensor(dtype)) - @staticmethod - def _get_alignment_numel(dtype: torch.dtype, alignment_bytes: int = 32) -> int: - element_size = torch.empty((), dtype=dtype).element_size() - return max(1, alignment_bytes // element_size) - - @classmethod - def _align_numel_offset( - cls, offset: int, dtype: torch.dtype, alignment_bytes: int = 32 - ) -> int: - alignment_numel = cls._get_alignment_numel(dtype, alignment_bytes) - remainder = offset % alignment_numel - if remainder == 0: - return offset - return offset + alignment_numel - remainder - @torch.compiler.disable def _initialize(self) -> None: if not self.enabled: @@ -589,12 +625,10 @@ def _layer_byte_totals( for _, weight in weights: tensor = self._to_local_tensor(weight) if tensor.is_contiguous(): - contiguous_numel = self._align_numel_offset( - contiguous_numel, dtype - ) + contiguous_numel = _align_numel_offset(contiguous_numel, dtype) contiguous_numel += tensor.numel() if not self._mapped_regions.holds(tensor): - anonymous_numel = self._align_numel_offset( + anonymous_numel = _align_numel_offset( anonymous_numel, dtype ) anonymous_numel += tensor.numel() @@ -652,7 +686,11 @@ def _plan_layer_hosting(self, layer_groups: Dict) -> Dict[int, str]: resident = [idx for idx in sorted(totals) if idx not in set(streamed)] for layer_idx in streamed + resident: layer_bytes = totals[layer_idx] - if self.pin_cpu_memory and pinned_bytes + layer_bytes <= spendable: + if ( + self.pin_cpu_memory + and self._pin_during_initialization + and pinned_bytes + layer_bytes <= spendable + ): hosting[layer_idx] = "pinned" pinned_bytes += layer_bytes pin_order.append(layer_idx) @@ -821,7 +859,7 @@ def _initialize_layer_weights(self) -> None: # satisfy a 32-byte alignment contract. Reusing one flat buffer # is still fine, but each logical tensor slice must start on an # aligned offset inside that buffer. - current_offset = self._align_numel_offset(current_offset, dtype) + current_offset = _align_numel_offset(current_offset, dtype) aligned_offsets[name] = current_offset current_offset += local_weight.numel() @@ -1938,6 +1976,7 @@ def configure_layerwise_offload( *, pin_budget: HostPinBudget | None = None, component_name: str | None = None, + pin_during_initialization: bool = True, ): self.park_non_layer_weights_between_uses = ( server_args.performance_mode == "memory" @@ -2016,6 +2055,7 @@ def configure_layerwise_offload( pin_cpu_memory=server_args.pin_cpu_memory, pin_budget=pin_budget, pin_component_name=pin_component_name, + pin_during_initialization=pin_during_initialization, prefetch_size=prefetch_size, resident_layers=resident_layers, initialize=False, @@ -2216,6 +2256,25 @@ def enable_offload(self) -> None: manager.release_all() manager.register_forward_hooks() + def remove_layerwise_offload(self) -> None: + """Restore an ordinary component-managed module. + + Every real layer is materialized before the manager stores are dropped, + and HostPin allowances are returned first. Auto residency uses this for + both a selected component-offload target and transactional rollback of + a lazily configured layerwise target. + """ + managers = list(self.layerwise_offload_managers or ()) + if not managers: + return + for manager in managers: + manager.set_pinned_layers(()) + self.restore_non_layer_weights() + self.disable_offload() + self.layerwise_offload_managers = [] + self._parked_non_layer_weights.clear() + self._park_placeholders.clear() + def iter_materialized_weights(module: torch.nn.Module): """Yield (name, tensor) pairs with materialized weights, even under offload. @@ -2494,7 +2553,7 @@ def _h2d_bytes_a_pin_would_save(name: str) -> int: reverse=True, ) local_worker_count = max(1, server_args.num_gpus // server_args.nnodes) - pin_budget = HostPinBudget.for_local_worker(local_worker_count) + pin_budget = server_args.host_pin_budget() logger.info( "Layerwise offload host memory: %s; this worker may pin %.2f GiB " "after splitting the node allowance across %d local workers", 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 6ce88280f84d..257e7f9c869e 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -14,7 +14,7 @@ import tempfile from dataclasses import field from enum import Enum -from typing import Any, List, Literal, Optional +from typing import TYPE_CHECKING, Any, List, Literal, Optional import addict import yaml @@ -88,6 +88,11 @@ logger = init_logger(__name__) +if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + HostPinBudget, + ) + LTX2_TWO_STAGE_DEVICE_MODES = ("original", "resident") LTX2_TWO_STAGE_DEVICE_MODE_CHOICES = LTX2_TWO_STAGE_DEVICE_MODES LTX2_TWO_STAGE_PIPELINE_NAMES = ("LTX2TwoStagePipeline", "LTX2TwoStageHQPipeline") @@ -373,6 +378,13 @@ class ServerArgs(DisaggServerArgsMixin): _required_resident_components: dict[str, str] = field( default_factory=dict, init=False, repr=False ) + # Warmup-calibrated serving placement. This is deliberately separate from + # configured CLI placement and loader hard requirements: calibration can + # move among all three runtime modes, while a loader-owned resident + # requirement must remain authoritative. + _auto_residency_modes: dict[str, str] = field( + default_factory=dict, init=False, repr=False + ) _fsdp_disabled_components: set[str] = field( default_factory=set, init=False, repr=False ) @@ -891,10 +903,7 @@ def _resolve_default_ltx2_two_stage_device_mode(self) -> str: device_total_memory_gb = ( current_platform.get_device_total_memory() / BYTES_PER_GB ) - if ( - "H200" in device_name - or device_total_memory_gb >= LTX2_RESIDENT_AUTO_ENABLE_MEM_GB - ): + if device_total_memory_gb >= LTX2_RESIDENT_AUTO_ENABLE_MEM_GB: logger.info( "Automatically set ltx2_two_stage_device_mode=resident for high-memory CUDA GPU (%s, %.2f GiB total)", device_name, @@ -1493,6 +1502,9 @@ def residency_mode(self, component_name: str) -> str: return RESIDENT if component_name in self._required_resident_components: return RESIDENT + auto_mode = self.auto_residency_mode(component_name) + if auto_mode is not None: + return auto_mode return self.configured_residency_mode(component_name) @@ -1546,6 +1558,43 @@ def component_residency_requirement(self, component_name: str) -> str | None: """Return the feature that currently hard-requires GPU residency.""" return self._required_resident_components.get(component_name) + def auto_residency_mode(self, component_name: str) -> str | None: + """Return the warmup-calibrated serving override, if one is active.""" + return self.__dict__.get("_auto_residency_modes", {}).get(component_name) + + def host_pin_budget(self) -> "HostPinBudget": + """Return this worker's shared HostPin ledger. + + Initial and lazily configured layerwise components must spend from the + same allowance. Keeping this process-local object out of dataclass + fields also keeps ServerArgs serialization and logging unchanged. + """ + from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + HostPinBudget, + ) + + budget = self.__dict__.get("_host_pin_budget") + if budget is None: + local_worker_count = max(1, self.num_gpus // self.nnodes) + budget = HostPinBudget.for_local_worker(local_worker_count) + self.__dict__["_host_pin_budget"] = budget + return budget + + def set_auto_residency_mode(self, component_name: str, mode: str) -> None: + """Set one runtime placement without rewriting its configured policy.""" + if mode not in (RESIDENT, COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD): + raise ValueError(f"invalid auto residency mode {mode!r}") + owner = self._required_resident_components.get(component_name) + if owner is not None and mode != RESIDENT: + raise ValueError( + f"cannot set {component_name!r} to {mode!r}; " + f"{owner} requires it resident" + ) + self.__dict__.setdefault("_auto_residency_modes", {})[component_name] = mode + + def clear_auto_residency_mode(self, component_name: str) -> None: + self.__dict__.get("_auto_residency_modes", {}).pop(component_name, None) + def should_cpu_offload_component(self, component_name: str) -> bool: return self.residency_mode(component_name) == COMPONENT_OFFLOAD @@ -2452,7 +2501,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "LTX-2.3 two-stage device residency mode: " "'original' keeps official two-stage semantics without premerged stage2, " "'resident' keeps both transformers resident on GPU. " - "Default is auto: resident on H200/high-memory CUDA GPUs, otherwise original." + "Default is auto: resident on CUDA GPUs with at least 130 GiB, otherwise original." ), ) parser.add_argument( diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 1f2c303bc7f2..21d2d9fd0bf4 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -11,7 +11,6 @@ from sglang.multimodal_gen.configs.pipeline_configs.longlive2 import LongLive2T2VConfig from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( ACTIVATION_EXTRAPOLATION_MARGIN, - AUTO_RESIDENCY_FEATURE_NAME, GIB_BYTES, MIN_VRAM_RESERVE_BYTES, AppliedResidencyChange, @@ -1376,20 +1375,50 @@ def __init__(self, managers: list[_FakeLayerwiseManager]): self.layerwise_offload_managers = managers +class _FakeLazyLayerwiseDit(LayerwiseOffloadableModuleMixin, nn.Module): + layer_names = ["layers"] + + def __init__(self, num_layers: int = 4): + nn.Module.__init__(self) + self.layers = nn.ModuleList(nn.Linear(4, 4) for _ in range(num_layers)) + self.layerwise_offload_managers = [] + + def configure_layerwise_offload( + self, + server_args, + *, + pin_budget=None, + component_name=None, + pin_during_initialization=True, + ): + del server_args, pin_budget, component_name, pin_during_initialization + tensors = { + f"layers.{index}.weight": layer.weight + for index, layer in enumerate(self.layers) + } + self.layerwise_offload_managers = [_FakeLayerwiseManager(tensors)] + + class _StubResidencyArgs: """Duck-typed stand-in for the two ServerArgs hooks adjustments use.""" def __init__(self): - self.required: set[str] = set() + self.auto_modes: dict[str, str] = {} + + def auto_residency_mode(self, component_name): + return self.auto_modes.get(component_name) - def require_component_resident(self, component_name, *, feature_name): - self.required.add(component_name) + def residency_mode(self, component_name): + return self.auto_modes.get(component_name, COMPONENT_OFFLOAD) - def release_required_component_residency(self, component_name, *, feature_name): - self.required.discard(component_name) + def set_auto_residency_mode(self, component_name, mode): + self.auto_modes[component_name] = mode - def component_residency_requirement(self, component_name): - return AUTO_RESIDENCY_FEATURE_NAME if component_name in self.required else None + def clear_auto_residency_mode(self, component_name): + self.auto_modes.pop(component_name, None) + + def host_pin_budget(self): + return None def _plan_for(candidates: list[ResidencyTarget]) -> AutoResidencyPlan: @@ -1636,6 +1665,81 @@ def test_non_dit_layerwise_frontier_includes_partial_residency(self): and candidate.target_mode() == LAYERWISE_OFFLOAD for candidate in candidates ) + component = next( + candidate + for candidate in candidates + if candidate.target_mode() == COMPONENT_OFFLOAD + ) + assert component.target_layerwise_pinned_layers == ((),) + assert component.pinned_host_delta_bytes == 0 + + def test_component_offload_frontier_can_expand_layerwise_lazily(self): + module = _FakeLazyLayerwiseDit(num_layers=4) + + candidates = collect_residency_targets( + modules={"transformer": module}, + residency_mode_of=self._modes({"transformer": COMPONENT_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=10, + layerwise_tuning_of=lambda _name, _dit_group: (0.0, 0.0, "leading"), + ) + + component = next( + candidate + for candidate in candidates + if candidate.target_mode() == COMPONENT_OFFLOAD + ) + full_stage = next( + candidate + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + and candidate.target_layerwise_resident_layers == (4,) + ) + assert component.current_placement + assert full_stage.h2d_bytes_per_request == component.h2d_bytes_per_request + assert any( + candidate.target_layerwise_resident_layers == (2,) + for candidate in candidates + ) + assert any(candidate.target_mode() == RESIDENT for candidate in candidates) + + def test_lazy_layerwise_frontier_keeps_buffers_out_of_managed_weights(self): + module = _FakeLazyLayerwiseDit(num_layers=2) + module.layers[0].register_buffer("cache", torch.zeros(4096)) + + candidates = collect_residency_targets( + modules={"transformer": module}, + residency_mode_of=self._modes({"transformer": COMPONENT_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=10, + layerwise_tuning_of=lambda _name, _dit_group: (0.0, 0.0, "leading"), + ) + + full_stage = next( + candidate + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + and candidate.target_layerwise_resident_layers == (2,) + ) + # Each Linear packs a 16-element weight and 4-element bias. The buffer + # remains resident and is not copied into the layerwise store. + assert full_stage.target_resident_weight_bytes == 2 * 20 * 4 + assert full_stage.target_device_weight_bytes == ( + full_stage.target_resident_weight_bytes + + module.layers[0].cache.untyped_storage().nbytes() + ) + + def test_component_weight_footprint_deduplicates_tied_storage(self): + module = nn.Module() + backing = torch.empty(1024) + module.register_parameter("left", nn.Parameter(backing[:512])) + module.register_parameter("right", nn.Parameter(backing[512:])) + + assert component_runtime_weight_bytes({"transformer": module}) == { + "transformer": backing.untyped_storage().nbytes() + } def test_auto_resident_component_keeps_offload_as_a_target(self): module = nn.Linear(4, 4) @@ -1700,6 +1804,77 @@ def test_layerwise_frontier_accounts_for_immediate_materialization(self): class TestApplyAndRollback: + def test_lazy_layerwise_configuration_rolls_back_to_component_offload(self): + module = _FakeLazyLayerwiseDit(num_layers=4) + args = _StubResidencyArgs() + candidates = collect_residency_targets( + modules={"transformer": module}, + residency_mode_of=lambda _name: COMPONENT_OFFLOAD, + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=10, + layerwise_tuning_of=lambda _name, _dit_group: (0.0, 0.0, "leading"), + ) + target = next( + candidate + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + and candidate.target_layerwise_resident_layers == (2,) + ) + + applied = apply_residency_changes( + plan=_plan_for([target]), + modules={"transformer": module}, + server_args=args, + ) + + assert len(module.layerwise_offload_managers) == 1 + assert module.layerwise_offload_managers[0].resident_layers == 2 + assert args.auto_modes == {"transformer": LAYERWISE_OFFLOAD} + + rollback_residency_changes( + applied=applied, + modules={"transformer": module}, + server_args=args, + ) + assert module.layerwise_offload_managers == [] + assert args.auto_modes == {} + + def test_layerwise_can_become_component_offload_and_rollback(self): + module = _FakeLazyLayerwiseDit(num_layers=4) + args = _StubResidencyArgs() + module.configure_layerwise_offload(args) + args.auto_modes["text_encoder"] = LAYERWISE_OFFLOAD + candidates = collect_residency_targets( + modules={"text_encoder": module}, + residency_mode_of=lambda _name: LAYERWISE_OFFLOAD, + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=1, + ) + target = next( + candidate + for candidate in candidates + if candidate.target_mode() == COMPONENT_OFFLOAD + ) + + applied = apply_residency_changes( + plan=_plan_for([target]), + modules={"text_encoder": module}, + server_args=args, + ) + + assert module.layerwise_offload_managers == [] + assert args.auto_modes == {"text_encoder": COMPONENT_OFFLOAD} + + rollback_residency_changes( + applied=applied, + modules={"text_encoder": module}, + server_args=args, + ) + assert len(module.layerwise_offload_managers) == 1 + assert args.auto_modes == {"text_encoder": LAYERWISE_OFFLOAD} + def test_component_offload_promotion_marks_resident_without_moving(self): module = nn.Linear(4, 4) args = _StubResidencyArgs() @@ -1708,18 +1883,18 @@ def test_component_offload_promotion_marks_resident_without_moving(self): modules={"text_encoder": module}, server_args=args, ) - assert args.required == {"text_encoder"} + assert args.auto_modes == {"text_encoder": RESIDENT} assert [p.component_name for p in applied] == ["text_encoder"] rollback_residency_changes( applied=applied, modules={"text_encoder": module}, server_args=args ) - assert args.required == set() + assert args.auto_modes == {} def test_component_resident_can_be_demoted_and_rollback_restores_it(self): module = nn.Linear(4, 4) args = _StubResidencyArgs() - args.required.add("text_encoder") + args.auto_modes["text_encoder"] = RESIDENT candidate = ResidencyTarget( component_name="text_encoder", residency_mode=COMPONENT_OFFLOAD, @@ -1734,14 +1909,14 @@ def test_component_resident_can_be_demoted_and_rollback_restores_it(self): modules={"text_encoder": module}, server_args=args, ) - assert args.required == set() + assert args.auto_modes == {"text_encoder": COMPONENT_OFFLOAD} rollback_residency_changes( applied=applied, modules={"text_encoder": module}, server_args=args, ) - assert args.required == {"text_encoder"} + assert args.auto_modes == {"text_encoder": RESIDENT} def test_layerwise_promotion_loads_all_layers_and_rollback_rearms(self): manager = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) @@ -1763,7 +1938,7 @@ def test_layerwise_promotion_loads_all_layers_and_rollback_rearms(self): ) assert manager.enabled is True assert manager.register_hooks_calls == 1 - assert args.required == set() + assert args.auto_modes == {} def test_partial_layerwise_promotion_restores_exact_group_counts(self): managers = [ @@ -1800,7 +1975,7 @@ def test_partial_layerwise_promotion_restores_exact_group_counts(self): ) assert [manager.resident_layers for manager in managers] == [2, 2] - assert args.required == set() + assert args.auto_modes == {"transformer": LAYERWISE_OFFLOAD} rollback_residency_changes( applied=applied, modules={"transformer": module}, @@ -1819,7 +1994,7 @@ def test_permanent_layerwise_can_be_demoted_and_rollback_restores_it(self): module = _FakeLayerwiseDit([manager]) module.disable_offload() args = _StubResidencyArgs() - args.required.add("transformer") + args.auto_modes["transformer"] = RESIDENT candidate = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, @@ -1838,7 +2013,7 @@ def test_permanent_layerwise_can_be_demoted_and_rollback_restores_it(self): assert manager.enabled is True assert manager.resident_layers == 1 - assert args.required == set() + assert args.auto_modes == {"transformer": LAYERWISE_OFFLOAD} rollback_residency_changes( applied=applied, @@ -1847,7 +2022,7 @@ def test_permanent_layerwise_can_be_demoted_and_rollback_restores_it(self): ) assert manager.enabled is False assert manager.resident_layers == 2 - assert args.required == {"transformer"} + assert args.auto_modes == {"transformer": RESIDENT} def test_mid_failure_rolls_back_already_applied_promotions(self): manager = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) @@ -1863,7 +2038,7 @@ def test_mid_failure_rolls_back_already_applied_promotions(self): modules={"transformer": module}, server_args=args, ) - assert args.required == set() + assert args.auto_modes == {} assert manager.enabled is True # Validation finishes before any placement is touched. assert manager.register_hooks_calls == 0 @@ -2042,9 +2217,7 @@ def test_disable_offload_failure_rearms_the_manager(self): def test_failed_rollback_raises_rollback_error_with_visible_cause(self): class _BrokenArgs(_StubResidencyArgs): - def release_required_component_residency( - self, component_name, *, feature_name - ): + def clear_auto_residency_mode(self, component_name): # str(AssertionError()) is "" -- describe_error must keep it # visible so no truthiness filter can drop it raise AssertionError() @@ -2072,9 +2245,7 @@ def test_rollback_keeps_going_past_a_broken_component(self): released: list[str] = [] class _PartialArgs(_StubResidencyArgs): - def release_required_component_residency( - self, component_name, *, feature_name - ): + def clear_auto_residency_mode(self, component_name): if component_name == "vae": raise RuntimeError("boom") released.append(component_name) 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 87424ba5ebd9..6496d31f759a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -184,11 +184,14 @@ def __init__(self, enabled: bool) -> None: class _TestServerArgs(SimpleNamespace): canonical_residency_mode = ServerArgs.canonical_residency_mode + configured_residency_mode = ServerArgs.configured_residency_mode explicit_residency_mode = ServerArgs.explicit_residency_mode _legacy_component_offload_flag = staticmethod( ServerArgs._legacy_component_offload_flag ) residency_mode = ServerArgs.residency_mode + auto_residency_mode = ServerArgs.auto_residency_mode + host_pin_budget = ServerArgs.host_pin_budget is_arg_explicitly_set = ServerArgs.is_arg_explicitly_set is_explicit_layerwise_offload_component = ( ServerArgs.is_explicit_layerwise_offload_component @@ -209,6 +212,7 @@ def _server_args(**kwargs): num_gpus=1, nnodes=1, _required_resident_components=set(), + _auto_residency_modes={}, _component_layerwise_capabilities={}, _explicit_arg_names=set(), cpu_offload_components=None, diff --git a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py index e79440262bed..2e3cb79dcf9b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -198,6 +198,9 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): configured_residency_mode=lambda _name: "resident", explicit_residency_mode=lambda _name: None, component_residency_requirement=lambda _name: None, + auto_residency_mode=lambda _name: None, + layerwise_tuning_for=lambda _name, *, dit_group: (0.0, 0.0, "leading"), + pin_cpu_memory=True, ) device_module = Mock() device_module.mem_get_info.return_value = (100, 100) @@ -242,6 +245,9 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): configured_residency_mode=lambda _name: "resident", explicit_residency_mode=lambda _name: None, component_residency_requirement=lambda _name: None, + auto_residency_mode=lambda _name: None, + layerwise_tuning_for=lambda _name, *, dit_group: (0.0, 0.0, "leading"), + pin_cpu_memory=True, ) device_module = Mock() device_module.mem_get_info.return_value = (100 * GIB_BYTES, 140 * GIB_BYTES) From 618f2ce384a6ce38f182788779f0641973af0f3f Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 08:56:20 +0800 Subject: [PATCH 37/90] [diffusion] guard auto residency calibration regressions --- .../runtime/managers/gpu_worker.py | 28 +++++++++++ .../memory_managers/auto_residency.py | 21 +++++---- .../memory_managers/layerwise_offload.py | 47 ++++++++++++++++--- .../test/server/perf_baselines/5090.json | 2 +- .../test/server/perf_baselines/h100.json | 6 +-- .../test/unit/test_auto_residency.py | 9 ++++ .../test/unit/test_auto_residency_warmup.py | 41 ++++++++++++++++ 7 files changed, 135 insertions(+), 19 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 0013ee70ddc4..43357ac2a45f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -51,10 +51,12 @@ ) from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( GIB_BYTES, + MAX_LATENCY_EQUIVALENCE_NS, PLACEMENT_STATUS_ADJUSTED, PLACEMENT_STATUS_ROLLBACK_FAILED, PLACEMENT_STATUS_ROLLED_BACK, PLACEMENT_STATUS_SKIPPED, + POST_ADJUSTMENT_REGRESSION_FRACTION, AppliedResidencyChange, AutoResidencyRollbackError, DefaultWorkload, @@ -1265,6 +1267,31 @@ def apply_auto_residency(self) -> OutputBatch: skip_reason=describe_error(e), ) reports = self._auto_residency_all_gather(local_report) + if self._auto_residency_round_sizes: + regressions = [] + for report in reports: + reference_ns = report.estimated_request_duration_ns + measured_ns = report.measured_request_duration_ns + tolerance_ns = max( + MAX_LATENCY_EQUIVALENCE_NS, + int(reference_ns * POST_ADJUSTMENT_REGRESSION_FRACTION), + ) + if reference_ns > 0 and measured_ns > reference_ns + tolerance_ns: + regressions.append((measured_ns - reference_ns, report)) + if regressions: + _, regressed = max(regressions, key=lambda item: item[0]) + cause = ( + "post-adjustment calibration regressed request duration " + f"from {regressed.estimated_request_duration_ns / 1e9:.2f}s " + f"to {regressed.measured_request_duration_ns / 1e9:.2f}s" + ) + if self.is_output_rank: + logger.warning("Auto residency: %s; rolling back", cause) + return self._rollback_everywhere( + cause=cause, + already_failed=False, + latest_round_only=True, + ) plan = plan_auto_residency(reports=reports) summary = format_plan_summary(plan=plan, workload=workload, records=records) if plan.skip_reason is not None or not plan.changes: @@ -1523,6 +1550,7 @@ def _build_auto_residency_report( torch.get_device_module().memory_allocated() ), estimated_request_duration_ns=estimated_request_duration_ns, + measured_request_duration_ns=measured_request_duration_ns, candidate_latency_savings_ns=candidate_latency_savings_ns, candidates=candidates, skip_reason=skip_reason, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 1713731c31e4..ac7f28f185f5 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -92,6 +92,11 @@ MIN_LATENCY_EQUIVALENCE_NS = 50_000_000 MAX_LATENCY_EQUIVALENCE_NS = 100_000_000 LATENCY_EQUIVALENCE_FRACTION = 0.01 +# The transfer model ranks feasible placements; the mandatory warmup is the +# authority on whether a selected placement actually helped. Allow normal +# measurement noise, but undo a round whose calibrated request is materially +# slower than the original layout. +POST_ADJUSTMENT_REGRESSION_FRACTION = 0.05 PLACEMENT_STATUS_SKIPPED = "skipped" PLACEMENT_STATUS_ADJUSTED = "adjusted" @@ -235,6 +240,7 @@ class RankResidencyReport(msgspec.Struct, frozen=True): host_transition_headroom_bytes: int = 0 device_transition_allocated_bytes: int = 0 estimated_request_duration_ns: int = 0 + measured_request_duration_ns: int = 0 candidate_latency_savings_ns: dict[str, int] = {} candidates: list[ResidencyTarget] = [] skip_reason: str | None = None @@ -1996,11 +2002,10 @@ def apply_residency_changes( candidate.target_layerwise_pinned_layers or () ) if target_mode == COMPONENT_OFFLOAD: - module.remove_layerwise_offload() + module.remove_layerwise_offload(to_cpu=True) server_args.set_auto_residency_mode( candidate.component_name, COMPONENT_OFFLOAD ) - module.to("cpu") elif target_mode == RESIDENT: server_args.set_auto_residency_mode( candidate.component_name, RESIDENT @@ -2082,7 +2087,6 @@ def rollback_residency_changes( if not isinstance(module, LayerwiseOffloadableModuleMixin): raise RuntimeError("lost layerwise offload capability") if not adjustment.previous_layerwise_configured: - module.remove_layerwise_offload() if adjustment.previous_auto_residency_mode is None: server_args.clear_auto_residency_mode(adjustment.component_name) else: @@ -2090,11 +2094,12 @@ def rollback_residency_changes( adjustment.component_name, adjustment.previous_auto_residency_mode, ) - if ( - server_args.residency_mode(adjustment.component_name) - != RESIDENT - ): - module.to("cpu") + module.remove_layerwise_offload( + to_cpu=( + server_args.residency_mode(adjustment.component_name) + != RESIDENT + ) + ) continue currently_enabled = { manager.enabled for manager in module.layerwise_offload_managers 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 a161a9eeccd9..7e6b7ebeaeb6 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 @@ -2256,21 +2256,54 @@ def enable_offload(self) -> None: manager.release_all() manager.register_forward_hooks() - def remove_layerwise_offload(self) -> None: + def remove_layerwise_offload(self, *, to_cpu: bool = False) -> None: """Restore an ordinary component-managed module. - Every real layer is materialized before the manager stores are dropped, - and HostPin allowances are returned first. Auto residency uses this for - both a selected component-offload target and transactional rollback of - a lazily configured layerwise target. + HostPin allowances are returned first. Resident targets materialize + every real layer on the device. Component-offload targets instead bind + the manager's existing host tensors directly, avoiding a transient full + device copy that can OOM while undoing a lazily configured placement. """ managers = list(self.layerwise_offload_managers or ()) if not managers: return for manager in managers: manager.set_pinned_layers(()) - self.restore_non_layer_weights() - self.disable_offload() + if to_cpu: + parameters = dict(self.named_parameters()) + with torch.inference_mode(False), torch.no_grad(): + for manager in managers: + if manager.enabled: + # Inference weights are immutable, but preserve any + # materialized update before releasing device tensors. + manager.sync_all_layers_to_cpu() + manager.release_all() + manager.remove_forward_hooks() + manager.enabled = False + for name, host_tensor in manager.iter_cpu_weights(): + parameter = parameters.get(name) + if parameter is None: + raise RuntimeError( + f"layerwise parameter {name!r} disappeared" + ) + local_tensor = host_tensor.detach() + parameter.data = ( + manager._wrap_for_target(parameter, local_tensor) + if isinstance(parameter, DTensor) + else local_tensor + ) + for name, host_tensor in self._parked_non_layer_weights.items(): + parameter = parameters.get(name) + if parameter is not None: + parameter.data = host_tensor + if current_platform.is_mps() and hasattr( + self, "_mps_cpu_non_layer_parameters" + ): + self.restore_mps_cpu_non_layer_weights() + self.to("cpu") + else: + self.restore_non_layer_weights() + self.disable_offload() self.layerwise_offload_managers = [] self._parked_non_layer_weights.clear() self._park_placeholders.clear() diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json b/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json index 1a494a457681..2012eaee4829 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json @@ -85,7 +85,7 @@ "expected_avg_denoise_ms": 426.59, "expected_median_denoise_ms": 427.94, "load_peak_vram_mb": 4934.0, - "runtime_peak_vram_mb": 18246.0, + "runtime_peak_vram_mb": 30428.0, "estimated_full_test_time_s": 160.9 }, "zimage_image_t2i": { diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index e1860afb8ec5..3f6f8b7f329a 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -1397,7 +1397,7 @@ "expected_avg_denoise_ms": 938.38, "expected_median_denoise_ms": 945.14, "load_peak_vram_mb": 43244.0, - "runtime_peak_vram_mb": 47874.0, + "runtime_peak_vram_mb": 59772.0, "estimated_full_test_time_s": 160.2 }, "qwen_image_layered_i2i": { @@ -1509,7 +1509,7 @@ "expected_avg_denoise_ms": 871.39, "expected_median_denoise_ms": 1010.9, "load_peak_vram_mb": 27320.0, - "runtime_peak_vram_mb": 37434.0, + "runtime_peak_vram_mb": 48504.0, "estimated_full_test_time_s": 77.0 }, "wan2_2_i2v_a14b_2gpu": { @@ -2158,7 +2158,7 @@ "expected_avg_denoise_ms": 551.91, "expected_median_denoise_ms": 555.79, "load_peak_vram_mb": 43244.0, - "runtime_peak_vram_mb": 47874.0, + "runtime_peak_vram_mb": 59770.0, "estimated_full_test_time_s": 143.7 }, "fsdp-inference": { diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 21d2d9fd0bf4..070fb5708617 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -482,6 +482,7 @@ def _report( device_transition_allocated_gib: int = 0, target_workload_measured: bool = False, estimated_request_duration_ns: int = 0, + measured_request_duration_ns: int = 0, candidate_latency_savings_ns: dict[str, int] | None = None, ) -> RankResidencyReport: return RankResidencyReport( @@ -506,6 +507,7 @@ def _report( host_transition_headroom_bytes=(host_transition_headroom_gib * GIB_BYTES), device_transition_allocated_bytes=(device_transition_allocated_gib * GIB_BYTES), estimated_request_duration_ns=estimated_request_duration_ns, + measured_request_duration_ns=measured_request_duration_ns, candidate_latency_savings_ns=candidate_latency_savings_ns or {}, candidates=candidates if candidates is not None else [_candidate("vae")], skip_reason=skip_reason, @@ -1831,6 +1833,7 @@ def test_lazy_layerwise_configuration_rolls_back_to_component_offload(self): assert len(module.layerwise_offload_managers) == 1 assert module.layerwise_offload_managers[0].resident_layers == 2 assert args.auto_modes == {"transformer": LAYERWISE_OFFLOAD} + manager = module.layerwise_offload_managers[0] rollback_residency_changes( applied=applied, @@ -1839,11 +1842,14 @@ def test_lazy_layerwise_configuration_rolls_back_to_component_offload(self): ) assert module.layerwise_offload_managers == [] assert args.auto_modes == {} + assert module.layers[0].weight.device.type == "cpu" + assert manager.load_all_layers_calls == 0 def test_layerwise_can_become_component_offload_and_rollback(self): module = _FakeLazyLayerwiseDit(num_layers=4) args = _StubResidencyArgs() module.configure_layerwise_offload(args) + manager = module.layerwise_offload_managers[0] args.auto_modes["text_encoder"] = LAYERWISE_OFFLOAD candidates = collect_residency_targets( modules={"text_encoder": module}, @@ -1866,6 +1872,9 @@ def test_layerwise_can_become_component_offload_and_rollback(self): assert module.layerwise_offload_managers == [] assert args.auto_modes == {"text_encoder": COMPONENT_OFFLOAD} + assert module.layers[0].weight.device.type == "cpu" + assert target.target_resident_weight_bytes == 0 + assert manager.load_all_layers_calls == 0 rollback_residency_changes( applied=applied, diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py index 4b48dbb3d657..c74acbb04f29 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -15,6 +15,7 @@ PLACEMENT_STATUS_ROLLED_BACK, PLACEMENT_STATUS_SKIPPED, AppliedResidencyChange, + RankResidencyReport, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( COMPONENT_OFFLOAD, @@ -23,6 +24,46 @@ class TestAutoResidencyWarmup(unittest.TestCase): + def test_worker_rolls_back_a_materially_slower_calibrated_round(self): + worker = GPUWorker.__new__(GPUWorker) + worker.rank = 0 + worker.is_output_rank = False + worker.server_args = SimpleNamespace() + worker._auto_residency_warmup_records = [object()] + worker._auto_residency_round_sizes = [1] + report = RankResidencyReport( + rank=0, + budget_bytes=1, + estimated_peak_bytes=1, + estimated_request_duration_ns=10_000_000_000, + measured_request_duration_ns=20_000_000_000, + ) + worker._build_auto_residency_report = mock.Mock(return_value=report) + worker._auto_residency_all_gather = mock.Mock(return_value=[report]) + rolled_back = OutputBatch( + error="request duration regressed", + output={"status": PLACEMENT_STATUS_ROLLED_BACK}, + ) + worker._rollback_everywhere = mock.Mock(return_value=rolled_back) + + with mock.patch.object( + gpu_worker_module, + "resolve_default_workload", + return_value=SimpleNamespace(), + ), mock.patch.object( + gpu_worker_module, + "resolve_measured_default_workload", + return_value=SimpleNamespace(), + ): + response = worker.apply_auto_residency() + + self.assertIs(response, rolled_back) + worker._rollback_everywhere.assert_called_once() + kwargs = worker._rollback_everywhere.call_args.kwargs + self.assertIn("10.00s to 20.00s", kwargs["cause"]) + self.assertFalse(kwargs["already_failed"]) + self.assertTrue(kwargs["latest_round_only"]) + def test_worker_rollback_response_rewarms_the_restored_layout(self): actions = [] From 06442527876d6f4914321f670811dd63dafee77b Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 09:37:22 +0800 Subject: [PATCH 38/90] [diffusion] restore layerwise pins after rollback --- .../runtime/managers/memory_managers/auto_residency.py | 2 ++ python/sglang/multimodal_gen/test/unit/test_auto_residency.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index ac7f28f185f5..ed09f7dfef64 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -2073,6 +2073,8 @@ def rollback_residency_changes( module = modules.get(adjustment.component_name) if not isinstance(module, LayerwiseOffloadableModuleMixin): raise RuntimeError("lost layerwise offload capability") + if not module.layerwise_offload_managers: + continue module.release_layerwise_pins_outside( adjustment.previous_layerwise_pinned_layers or () ) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 070fb5708617..b073322d98f5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -1850,6 +1850,7 @@ def test_layerwise_can_become_component_offload_and_rollback(self): args = _StubResidencyArgs() module.configure_layerwise_offload(args) manager = module.layerwise_offload_managers[0] + manager.set_pinned_layers((0,)) args.auto_modes["text_encoder"] = LAYERWISE_OFFLOAD candidates = collect_residency_targets( modules={"text_encoder": module}, @@ -1883,6 +1884,7 @@ def test_layerwise_can_become_component_offload_and_rollback(self): ) assert len(module.layerwise_offload_managers) == 1 assert args.auto_modes == {"text_encoder": LAYERWISE_OFFLOAD} + assert module.layerwise_offload_managers[0].pinned_layer_indices() == (0,) def test_component_offload_promotion_marks_resident_without_moving(self): module = nn.Linear(4, 4) From f73d112776627a86cc3807c6667dd449aa1f5e92 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 10:26:03 +0800 Subject: [PATCH 39/90] [diffusion] constrain auto residency lifecycle phases --- .../runtime/managers/gpu_worker.py | 32 +++ .../memory_managers/auto_residency.py | 246 ++++++++++++++++-- .../memory_managers/component_manager.py | 53 ++++ .../runtime/pipelines_core/lora/pipeline.py | 44 +++- .../test/unit/test_auto_residency.py | 235 ++++++++++++++++- .../test/unit/test_auto_residency_warmup.py | 48 ++++ .../test/unit/test_component_residency.py | 38 +++ .../test/unit/test_lora_pipeline.py | 26 +- 8 files changed, 678 insertions(+), 44 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 43357ac2a45f..3d7953fbf2de 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -64,6 +64,7 @@ WarmupMemoryRecord, apply_residency_changes, collect_residency_targets, + component_current_device_weight_bytes, component_runtime_weight_bytes, describe_error, estimate_candidate_latency_savings_ns, @@ -730,6 +731,7 @@ def _record_server_warmup_memory( phase_allocated_peaks: dict[str, int] = {} phase_components: dict[str, tuple[str, ...]] = {} phase_used_components: dict[str, tuple[str, ...]] = {} + phase_full_weight_transition_components: dict[str, tuple[str, ...]] = {} untracked_active_components: tuple[str, ...] = () residency_manager = peek_global_component_residency_manager() if residency_manager is not None: @@ -737,6 +739,9 @@ def _record_server_warmup_memory( phase_allocated_peaks[phase_name] = peak.allocated_bytes phase_components[phase_name] = peak.active_components phase_used_components[phase_name] = peak.used_components + phase_full_weight_transition_components[phase_name] = ( + peak.full_weight_transition_components + ) untracked_active_components = residency_manager.current_device_components() request_allocated_peak = max( int(torch.get_device_module().max_memory_allocated()), @@ -750,6 +755,7 @@ def _record_server_warmup_memory( phase_allocated_peaks["request:untracked"] = request_allocated_peak phase_components["request:untracked"] = untracked_active_components phase_used_components["request:untracked"] = () + phase_full_weight_transition_components["request:untracked"] = () metrics = req.metrics width, height, num_frames, num_inference_steps = workload self._auto_residency_warmup_records.append( @@ -763,6 +769,9 @@ def _record_server_warmup_memory( phase_peak_allocated_bytes=phase_allocated_peaks, phase_active_components=phase_components, phase_used_components=phase_used_components, + phase_full_weight_transition_components=( + phase_full_weight_transition_components + ), num_inference_steps=num_inference_steps, total_duration_ms=( float(metrics.total_duration_ms) if metrics is not None else 0.0 @@ -1294,6 +1303,22 @@ def apply_auto_residency(self) -> OutputBatch: ) plan = plan_auto_residency(reports=reports) summary = format_plan_summary(plan=plan, workload=workload, records=records) + if ( + self._auto_residency_round_sizes + and plan.current_placement_reserve_shortfall_bytes > 0 + ): + shortfall_gib = plan.current_placement_reserve_shortfall_bytes / GIB_BYTES + if self.is_output_rank: + logger.warning( + "Auto residency calibration exceeded the VRAM reserve by " + "%.1f GiB; rolling back the latest adjustment round.", + shortfall_gib, + ) + return self._rollback_everywhere( + cause=f"VRAM reserve exceeded by {shortfall_gib:.1f} GiB", + already_failed=False, + latest_round_only=True, + ) if plan.skip_reason is not None or not plan.changes: if self.is_output_rank: logger.info("%s", summary) @@ -1469,6 +1494,7 @@ def _build_auto_residency_report( estimated_phase_peaks, active_components_by_phase, used_components_by_phase, + full_weight_transition_components_by_phase, ) = estimate_workload_phase_peaks( records=records, target_units=target_units, @@ -1537,6 +1563,12 @@ def _build_auto_residency_report( estimated_peak_bytes_by_phase=estimated_phase_peaks, active_components_by_phase=active_components_by_phase, used_components_by_phase=used_components_by_phase, + full_weight_transition_components_by_phase=( + full_weight_transition_components_by_phase + ), + current_device_weight_bytes_by_component=( + component_current_device_weight_bytes(self.pipeline.modules) + ), node_rank=self.server_args.node_rank, pinned_host_bytes=layerwise_pinned_host_bytes(self.pipeline.modules), host_pin_capacity_bytes=layerwise_host_pin_capacity_bytes( diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index ed09f7dfef64..3c3209aeb971 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -61,6 +61,7 @@ PlacementOption, optimize_placement, ) +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger if TYPE_CHECKING: @@ -131,6 +132,7 @@ class WarmupMemoryRecord(msgspec.Struct, frozen=True): phase_peak_allocated_bytes: dict[str, int] = {} phase_active_components: dict[str, tuple[str, ...]] = {} phase_used_components: dict[str, tuple[str, ...]] = {} + phase_full_weight_transition_components: dict[str, tuple[str, ...]] = {} num_inference_steps: int = 1 total_duration_ms: float = 0.0 stage_duration_ms: dict[str, float] = {} @@ -234,6 +236,8 @@ class RankResidencyReport(msgspec.Struct, frozen=True): estimated_peak_bytes_by_phase: dict[str, int] = {} active_components_by_phase: dict[str, tuple[str, ...]] = {} used_components_by_phase: dict[str, tuple[str, ...]] = {} + full_weight_transition_components_by_phase: dict[str, tuple[str, ...]] = {} + current_device_weight_bytes_by_component: dict[str, int] = {} node_rank: int = 0 pinned_host_bytes: int = 0 host_pin_capacity_bytes: int = 0 @@ -255,6 +259,7 @@ class AutoResidencyPlan(msgspec.Struct, frozen=True): resource_budget_bytes: dict[str, int] = {} changes: list[ResidencyTarget] = [] skip_reason: str | None = None + current_placement_reserve_shortfall_bytes: int = 0 class AppliedResidencyChange(msgspec.Struct, frozen=True): @@ -266,6 +271,17 @@ class AppliedResidencyChange(msgspec.Struct, frozen=True): previous_layerwise_configured: bool = True previous_auto_residency_mode: str | None = None pinned_host_changed: bool = False + applied_device_delta_bytes: int = 0 + + +def residency_device_growth_bytes(candidate: ResidencyTarget) -> int: + """Largest modeled accelerator increase while realizing one target.""" + return max( + candidate.device_transition_delta_bytes, + candidate.active_device_delta_bytes, + candidate.present_device_delta_bytes, + candidate.inactive_device_delta_bytes, + ) def resolve_default_workload(server_args: ServerArgs) -> DefaultWorkload: @@ -539,6 +555,7 @@ def estimate_workload_phase_peaks( dict[str, int], dict[str, tuple[str, ...]], dict[str, tuple[str, ...]], + dict[str, tuple[str, ...]], ]: """Estimate each measured execution phase at the target workload. @@ -558,24 +575,38 @@ def estimate_workload_phase_peaks( successful = covering grouped: dict[ - tuple[str, tuple[str, ...], tuple[str, ...]], list[WarmupMemoryRecord] + tuple[str, tuple[str, ...], tuple[str, ...], tuple[str, ...]], + list[WarmupMemoryRecord], ] = {} for record in successful: used_by_phase = record.phase_used_components or record.phase_active_components for phase_name in record.phase_peak_allocated_bytes: active = tuple(sorted(record.phase_active_components.get(phase_name, ()))) used = tuple(sorted(used_by_phase.get(phase_name, ()))) - grouped.setdefault((phase_name, active, used), []).append(record) + full_weight_transitions = tuple( + sorted( + record.phase_full_weight_transition_components.get(phase_name, ()) + ) + ) + grouped.setdefault( + (phase_name, active, used, full_weight_transitions), [] + ).append(record) layouts_per_phase: dict[str, int] = {} - for phase_name, _, _ in grouped: + for phase_name, _, _, _ in grouped: layouts_per_phase[phase_name] = layouts_per_phase.get(phase_name, 0) + 1 estimated_peaks: dict[str, int] = {} active_components: dict[str, tuple[str, ...]] = {} used_components: dict[str, tuple[str, ...]] = {} + full_weight_transition_components: dict[str, tuple[str, ...]] = {} layout_indices: dict[str, int] = {} - for (phase_name, active, used), phase_records in sorted(grouped.items()): + for ( + phase_name, + active, + used, + full_weight_transitions, + ), phase_records in sorted(grouped.items()): output_name = phase_name if layouts_per_phase[phase_name] > 1: index = layout_indices.get(phase_name, 0) @@ -606,7 +637,13 @@ def estimate_workload_phase_peaks( estimated_peaks[output_name] = estimate active_components[output_name] = active used_components[output_name] = used - return estimated_peaks, active_components, used_components + full_weight_transition_components[output_name] = full_weight_transitions + return ( + estimated_peaks, + active_components, + used_components, + full_weight_transition_components, + ) def _module_weight_bytes(module: nn.Module) -> int: @@ -643,6 +680,22 @@ def component_runtime_weight_bytes(modules: Mapping[str, object]) -> dict[str, i return runtime_bytes +def component_current_device_weight_bytes( + modules: Mapping[str, object], +) -> dict[str, int]: + """Physical component storage currently resident on the local accelerator.""" + result: dict[str, int] = {} + for name, module in modules.items(): + if not isinstance(module, nn.Module): + continue + result[name] = tensor_storage_bytes( + tensor + for tensor in chain(module.parameters(), module.buffers()) + if current_platform.is_device_type(tensor.device.type) + ) + return result + + def layerwise_pinned_host_bytes(modules: Mapping[str, object]) -> int: """Pinned layer-store bytes in this process, without alias double counts.""" seen_managers: set[int] = set() @@ -941,6 +994,8 @@ def _unconfigured_layerwise_targets( pinned_layers=empty_pins, uses_per_streamed_layer=uses_per_request, ) + coarse_transfer_work = maximum_transfer_work // uses_per_request + coarse_savings = max(0, maximum_transfer_work - coarse_transfer_work) current_resident = current_mode == RESIDENT current_inactive_bytes = full_weight_bytes if current_resident else 0 initial_managed_bytes = sum( @@ -976,7 +1031,14 @@ def _unconfigured_layerwise_targets( residency_mode=COMPONENT_OFFLOAD, target_residency_mode=LAYERWISE_OFFLOAD, target_resident_weight_bytes=resident_bytes, - h2d_bytes_per_request=maximum_transfer_work - transfer_work, + # This strategy has not run yet, so transfer bytes cannot + # establish that it beats the calibrated coarse path. It may + # tie that path as a memory alternative; only a subsequent + # calibration may justify further layerwise tuning. + h2d_bytes_per_request=min( + coarse_savings, + max(0, maximum_transfer_work - transfer_work), + ), target_layerwise_resident_layers=resident_layers, target_layerwise_pinned_layers=empty_pins, host_materialize_scratch_bytes=host_materialize_scratch, @@ -1410,8 +1472,15 @@ def rank_candidates_by_h2d_savings( ) -def _skip_plan(reason: str) -> AutoResidencyPlan: - return AutoResidencyPlan(skip_reason=reason) +def _skip_plan( + reason: str, *, current_placement_reserve_shortfall_bytes: int = 0 +) -> AutoResidencyPlan: + return AutoResidencyPlan( + skip_reason=reason, + current_placement_reserve_shortfall_bytes=( + current_placement_reserve_shortfall_bytes + ), + ) def _vram_reserve_bytes(budget_bytes: int, *, target_workload_measured: bool) -> int: @@ -1432,7 +1501,7 @@ def _vram_reserve_bytes(budget_bytes: int, *, target_workload_measured: bool) -> def _binding_phase_constraints( report: RankResidencyReport, candidate_component_names: set[str], -) -> list[tuple[str, int, tuple[str, ...], tuple[str, ...]]]: +) -> list[tuple[str, int, tuple[str, ...], tuple[str, ...], tuple[str, ...]]]: """Keep the binding measured and conservative unobserved phases. Every candidate has the same device delta in phases with the same active @@ -1442,7 +1511,9 @@ def _binding_phase_constraints( timestep-routed component (for example the second DiT in a two-stage denoiser), so give each unobserved candidate a synthetic phase at the measured request peak. This prevents permanent residents from being - placed in memory that the later component may need. + placed in memory that the later component may need. Components left on the + device at request end also carry into the next request, so add their weight + to cold-start phases that did not observe them yet. """ phase_peaks = report.estimated_peak_bytes_by_phase if not phase_peaks: @@ -1454,32 +1525,84 @@ def _binding_phase_constraints( if report.estimated_peak_bytes is not None else max(phase_peaks.values()) ) - binding: dict[tuple[tuple[str, ...], tuple[str, ...]], tuple[str, int]] = {} + current_device_bytes = dict(report.current_device_weight_bytes_by_component) + for candidate in report.candidates: + if ( + candidate.current_placement + and candidate.component_name not in current_device_bytes + ): + current_device_bytes[candidate.component_name] = ( + candidate.target_device_weight_bytes + ) + steady_state_components = { + component_name + for phase_name, components in report.active_components_by_phase.items() + if phase_name == "idle" or phase_name.startswith("idle:layout:") + for component_name in components + } + + binding: dict[ + tuple[tuple[str, ...], tuple[str, ...], tuple[str, ...]], tuple[str, int] + ] = {} for phase_name, phase_peak in phase_peaks.items(): + measured_components = set(report.active_components_by_phase.get(phase_name, ())) + carried_components = steady_state_components - measured_components + phase_peak += sum( + current_device_bytes.get(component_name, 0) + for component_name in carried_components + ) present = tuple( sorted( - set(report.active_components_by_phase.get(phase_name, ())) + (measured_components | steady_state_components) & candidate_component_names ) ) used = tuple(sorted(report.used_components_by_phase.get(phase_name, ()))) - layout = (present, used) + full_weight_transitions = tuple( + sorted( + set( + report.full_weight_transition_components_by_phase.get( + phase_name, () + ) + ) + & candidate_component_names + ) + ) + layout = (present, used, full_weight_transitions) current = binding.get(layout) if current is None or (phase_peak, phase_name) > (current[1], current[0]): binding[layout] = (phase_name, phase_peak) observed_components = { - component_name for (_, used) in binding for component_name in used + component_name + for (_, used, full_weight_transitions) in binding + for component_name in set(used) | set(full_weight_transitions) } + steady_request_peak = max( + request_peak, + *(phase_peak for _, phase_peak in binding.values()), + ) for component_name in sorted(candidate_component_names - observed_components): - binding[((component_name,), (component_name,))] = ( + present = tuple( + sorted( + (steady_state_components | {component_name}) & candidate_component_names + ) + ) + binding[(present, (component_name,), ())] = ( f"unobserved:{component_name}", - request_peak, + steady_request_peak, ) return [ - (f"gpu:rank{report.rank}:{phase_name}", phase_peak, present, used) - for (present, used), (phase_name, phase_peak) in sorted( - binding.items(), key=lambda item: item[1][0] + ( + f"gpu:rank{report.rank}:{phase_name}", + phase_peak, + present, + used, + full_weight_transitions, ) + for (present, used, full_weight_transitions), ( + phase_name, + phase_peak, + ) in sorted(binding.items(), key=lambda item: item[1][0]) ] @@ -1608,19 +1731,54 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP for report in reports } reserve = max(reserves_by_rank.values()) + current_placement_reserve_shortfall = max( + ( + max( + [ + report.estimated_peak_bytes or 0, + *report.estimated_peak_bytes_by_phase.values(), + ] + ) + + reserves_by_rank[report.rank] + - report.budget_bytes + ) + for report in reports + ) + current_placement_reserve_shortfall = max(0, current_placement_reserve_shortfall) candidates = _consensus_candidates(reports) if not candidates: - return _skip_plan("no eligible residency alternatives") + return _skip_plan( + "no eligible residency alternatives", + current_placement_reserve_shortfall_bytes=( + current_placement_reserve_shortfall + ), + ) candidate_component_names = {candidate.component_name for candidate in candidates} phase_constraints = { report.rank: _binding_phase_constraints(report, candidate_component_names) for report in reports } + current_placement_reserve_shortfall = max( + [ + current_placement_reserve_shortfall, + *( + max( + ( + phase_peak + reserves_by_rank[report.rank] - report.budget_bytes + for _, phase_peak, _, _, _ in phase_constraints[report.rank] + ), + default=0, + ) + for report in reports + ), + ] + ) + current_placement_reserve_shortfall = max(0, current_placement_reserve_shortfall) resource_budgets: dict[str, int] = {} for report in reports: - for resource_name, phase_peak, _, _ in phase_constraints[report.rank]: + for resource_name, phase_peak, _, _, _ in phase_constraints[report.rank]: # The measured placement has already completed warmup. A negative # reserve headroom therefore means "do not grow this phase", not # that the current zero-delta placement is infeasible. @@ -1706,11 +1864,25 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP _, present_components, used_components, + full_weight_transition_components, ) in phase_constraints[report.rank]: - if candidate.component_name in used_components: + if candidate.component_name in full_weight_transition_components: + # The measured phase already contains the component's full + # temporary materialization under every target mode. + phase_cost = 0 + elif candidate.component_name in used_components: phase_cost = rank_candidate.active_device_delta_bytes elif candidate.component_name in present_components: phase_cost = rank_candidate.present_device_delta_bytes + elif rank_candidate.permanent_residency: + # This phase observed none of the component. A permanent + # target introduces its complete footprint, including the + # unmanaged weights of a formerly layerwise component. + phase_cost = max( + rank_candidate.target_device_weight_bytes, + rank_candidate.target_resident_weight_bytes, + 0, + ) else: phase_cost = rank_candidate.inactive_device_delta_bytes resource_deltas[resource_name] = phase_cost @@ -1777,7 +1949,12 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP require_selection_from_every_group=complete_state_frontier, ) except NoFeasiblePlacementError as error: - return _skip_plan(str(error)) + return _skip_plan( + str(error), + current_placement_reserve_shortfall_bytes=( + current_placement_reserve_shortfall + ), + ) changed_candidates = [] for selection in placement.selections: candidate = candidate_by_key[selection.option_key] @@ -1791,6 +1968,7 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP budget_bytes=budget, resource_budget_bytes=resource_budgets, changes=changes, + current_placement_reserve_shortfall_bytes=(current_placement_reserve_shortfall), ) @@ -1864,6 +2042,7 @@ def apply_residency_changes( previous_auto_mode = server_args.auto_residency_mode( candidate.component_name ) + applied_device_delta_bytes = residency_device_growth_bytes(candidate) if layerwise_transition: if not isinstance(module, LayerwiseOffloadableModuleMixin): raise RuntimeError( @@ -1905,6 +2084,7 @@ def apply_residency_changes( previous_configured and candidate.target_layerwise_pinned_layers != previous_pinned ), + applied_device_delta_bytes=applied_device_delta_bytes, ) if candidate.target_layerwise_pinned_layers is None: raise RuntimeError( @@ -1917,6 +2097,7 @@ def apply_residency_changes( component_name=candidate.component_name, residency_mode=candidate.residency_mode, previous_auto_residency_mode=previous_auto_mode, + applied_device_delta_bytes=applied_device_delta_bytes, ) pinning_changes = [ candidate @@ -2082,7 +2263,24 @@ def rollback_residency_changes( errors.append(f"{adjustment.component_name}: {describe_error(e)}") if pinning_changes: release_unused_pinned_memory() - for adjustment in reversed(applied): + reverse_apply_order = list(reversed(applied)) + release_first = [ + adjustment + for adjustment in reverse_apply_order + if adjustment.applied_device_delta_bytes > 0 + ] + remaining = [ + adjustment + for adjustment in reverse_apply_order + if adjustment.applied_device_delta_bytes <= 0 + ] + for index, adjustment in enumerate(release_first + remaining): + if ( + index == len(release_first) + and release_first + and torch.get_device_module().is_available() + ): + torch.get_device_module().empty_cache() try: module = modules.get(adjustment.component_name) if adjustment.previous_layerwise_resident_layers is not None: diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py index 2ee48a55436b..3b8851eaa23c 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py @@ -63,6 +63,7 @@ class WarmupPhasePeak: active_components: tuple[str, ...] allocated_bytes: int used_components: tuple[str, ...] = () + full_weight_transition_components: tuple[str, ...] = () class ResidencyBatch(Protocol): @@ -140,6 +141,7 @@ def __init__( self._warmup_phase_key: str | None = None self._warmup_phase_components: tuple[str, ...] = () self._warmup_phase_used_components: tuple[str, ...] = () + self._warmup_phase_full_weight_transition_components: tuple[str, ...] = () self._warmup_phase_peaks: dict[str, WarmupPhasePeak] = {} self._completed_warmup_phase_peaks: dict[str, WarmupPhasePeak] = {} @@ -228,6 +230,7 @@ def begin_request( self._warmup_phase_key = None self._warmup_phase_components = () self._warmup_phase_used_components = () + self._warmup_phase_full_weight_transition_components = () self._warmup_phase_peaks = {} self._completed_warmup_phase_peaks = {} if self._track_warmup_memory: @@ -260,6 +263,43 @@ def before_stage( used_components=(), ) + @contextmanager + def full_weight_transition(self, component_names: Iterable[str]) -> Iterator[None]: + """Measure request logic that temporarily materializes complete weights.""" + names = tuple(sorted(set(component_names))) + if not self._track_warmup_memory or not names: + yield + return + previous_phase = ( + self._warmup_phase_key, + self._warmup_phase_components, + self._warmup_phase_used_components, + self._warmup_phase_full_weight_transition_components, + ) + self._begin_warmup_phase( + key=( + f"{self.state.stage_index}:{self.state.stage_name}:" + f"full-weight-transition:{','.join(names)}" + ), + components=self._warmup_active_components(), + used_components=(), + full_weight_transition_components=names, + ) + try: + yield + finally: + previous_key, components, used_components, transitions = previous_phase + if previous_key is None: + self._record_warmup_phase_peak() + self._warmup_phase_key = None + else: + self._begin_warmup_phase( + key=previous_key, + components=components, + used_components=used_components, + full_weight_transition_components=transitions, + ) + def begin_stage(self) -> None: """Prepare a stage that declares one uninterrupted component use.""" stage_uses = self._stage_uses_by_index[self.state.stage_index] @@ -616,6 +656,7 @@ def _begin_warmup_phase( key: str, components: tuple[str, ...], used_components: tuple[str, ...], + full_weight_transition_components: tuple[str, ...] = (), ) -> None: if not self._track_warmup_memory: return @@ -623,6 +664,9 @@ def _begin_warmup_phase( self._warmup_phase_key = key self._warmup_phase_components = tuple(sorted(set(components))) self._warmup_phase_used_components = tuple(sorted(set(used_components))) + self._warmup_phase_full_weight_transition_components = tuple( + sorted(set(full_weight_transition_components)) + ) torch.get_device_module().reset_peak_memory_stats() def _begin_warmup_transition( @@ -688,6 +732,9 @@ def _record_warmup_phase_peak(self) -> None: active_components=self._warmup_phase_components, allocated_bytes=int(torch.get_device_module().max_memory_allocated()), used_components=self._warmup_phase_used_components, + full_weight_transition_components=( + self._warmup_phase_full_weight_transition_components + ), ) previous = self._warmup_phase_peaks.get(self._warmup_phase_key) if previous is None: @@ -703,6 +750,12 @@ def _record_warmup_phase_peak(self) -> None: used_components=tuple( sorted(set(previous.used_components) & set(peak.used_components)) ), + full_weight_transition_components=tuple( + sorted( + set(previous.full_weight_transition_components) + & set(peak.full_weight_transition_components) + ) + ), ) def take_warmup_phase_peaks( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py index f4412c0f1e99..dd48efe962df 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py @@ -289,24 +289,40 @@ def _temporarily_disable_offload( yield [] return - # clear device cache to free up unused memory - if torch.get_device_module().is_available(): - torch.get_device_module().synchronize() - torch.get_device_module().empty_cache() + offloaded_module_names = [ + module_name + for module_name in module_names + if ( + (module := self.modules.get(module_name)) is not None + and is_layerwise_offloaded_module(module) + ) + ] + residency_transition = ( + self.component_residency_manager.full_weight_transition( + offloaded_module_names + ) + if self.component_residency_manager is not None + else nullcontext() + ) - offload_disabled_modules = [] - for module_name in module_names: - module = self.modules.get(module_name) - if module is not None and is_layerwise_offloaded_module(module): + with residency_transition: + # clear device cache to free up unused memory + if torch.get_device_module().is_available(): + torch.get_device_module().synchronize() + torch.get_device_module().empty_cache() + + offload_disabled_modules = [] + for module_name in offloaded_module_names: + module = self.modules[module_name] module.disable_offload() offload_disabled_modules.append(module) - try: - yield offload_disabled_modules - finally: - # Re-enable layerwise offload: sync weights to CPU and restore hooks - for module in offload_disabled_modules: - module.enable_offload() + try: + yield offload_disabled_modules + finally: + # Re-enable layerwise offload: sync weights to CPU and restore hooks + for module in offload_disabled_modules: + module.enable_offload() def _needs_lora_weight_update_context( self, diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index b073322d98f5..dddff0f830fd 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -63,6 +63,7 @@ def _record( stage_duration_ms=None, step_duration_ms=(), phase_active_components=None, + phase_full_weight_transition_components=None, ) -> WarmupMemoryRecord: return WarmupMemoryRecord( width=width, @@ -76,6 +77,9 @@ def _record( stage_duration_ms=stage_duration_ms or {}, step_duration_ms=step_duration_ms, phase_active_components=phase_active_components or {}, + phase_full_weight_transition_components=( + phase_full_weight_transition_components or {} + ), ) @@ -339,7 +343,7 @@ def test_phase_estimation_preserves_active_component_membership(self): phase_active_components={"denoise": ("transformer",)}, ) - peaks, active, used = estimate_workload_phase_peaks( + peaks, active, used, _ = estimate_workload_phase_peaks( records=[small, large], target_units=832 * 480 * 81, component_weight_bytes={"transformer": 28 * GIB_BYTES}, @@ -361,7 +365,7 @@ def test_phase_estimation_uses_allocated_peak(self): phase_active_components={"denoise": ("transformer",)}, ) - peaks, _, _ = estimate_workload_phase_peaks( + peaks, _, _, _ = estimate_workload_phase_peaks( records=[record], target_units=record.workload_units(), component_weight_bytes={"transformer": 10 * GIB_BYTES}, @@ -369,6 +373,27 @@ def test_phase_estimation_uses_allocated_peak(self): assert peaks["denoise"] == 11 * GIB_BYTES + def test_phase_estimation_preserves_full_weight_transition_components(self): + record = WarmupMemoryRecord( + width=1024, + height=1024, + num_frames=1, + baseline_allocated_bytes=2 * GIB_BYTES, + peak_allocated_bytes=4 * GIB_BYTES, + succeeded=True, + phase_peak_allocated_bytes={"lora_switch": 4 * GIB_BYTES}, + phase_full_weight_transition_components={"lora_switch": ("transformer",)}, + ) + + peaks, _, _, transitions = estimate_workload_phase_peaks( + records=[record], + target_units=record.workload_units(), + component_weight_bytes={"transformer": 2 * GIB_BYTES}, + ) + + assert peaks == {"lora_switch": 4 * GIB_BYTES} + assert transitions == {"lora_switch": ("transformer",)} + def test_phase_estimation_prefers_target_layout_over_smaller_warmup(self): small = WarmupMemoryRecord( width=256, @@ -391,7 +416,7 @@ def test_phase_estimation_prefers_target_layout_over_smaller_warmup(self): phase_active_components={"denoise": ("transformer",)}, ) - peaks, active, used = estimate_workload_phase_peaks( + peaks, active, used, _ = estimate_workload_phase_peaks( records=[small, target], target_units=target.workload_units(), component_weight_bytes={"transformer": 40 * GIB_BYTES}, @@ -423,7 +448,7 @@ def test_phase_estimation_keeps_distinct_active_layouts_separate(self): phase_active_components={"denoise": ("text_encoder",)}, ) - peaks, active, used = estimate_workload_phase_peaks( + peaks, active, used, _ = estimate_workload_phase_peaks( records=[transformer_phase, encoder_phase], target_units=transformer_phase.workload_units(), component_weight_bytes={ @@ -475,6 +500,8 @@ def _report( phase_peaks_gib: dict[str, int] | None = None, phase_components: dict[str, tuple[str, ...]] | None = None, phase_present_components: dict[str, tuple[str, ...]] | None = None, + phase_full_weight_transition_components: dict[str, tuple[str, ...]] | None = None, + component_weight_gib: dict[str, int] | None = None, node_rank: int = 0, pinned_host_gib: int = 0, host_pin_capacity_gib: int = 0, @@ -501,6 +528,13 @@ def _report( else phase_components or {} ), used_components_by_phase=phase_components or {}, + full_weight_transition_components_by_phase=( + phase_full_weight_transition_components or {} + ), + current_device_weight_bytes_by_component={ + name: value * GIB_BYTES + for name, value in (component_weight_gib or {}).items() + }, node_rank=node_rank, pinned_host_bytes=pinned_host_gib * GIB_BYTES, host_pin_capacity_bytes=host_pin_capacity_gib * GIB_BYTES, @@ -811,6 +845,56 @@ def test_async_prefetch_presence_is_not_charged_twice(self): "transformer:resident" ] + def test_request_end_residency_is_carried_into_the_next_request(self): + text_offload = ResidencyTarget( + component_name="text_encoder", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + current_placement=True, + ) + text_resident = ResidencyTarget( + component_name="text_encoder", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=20 * GIB_BYTES, + target_layerwise_resident_layers=(48,), + target_layerwise_pinned_layers=((),), + active_device_delta_bytes=20 * GIB_BYTES, + inactive_device_delta_bytes=20 * GIB_BYTES, + target_device_weight_bytes=20 * GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=80, + estimated_gib=40, + target_workload_measured=True, + phase_peaks_gib={"text_encode": 30, "idle": 40}, + phase_components={ + "text_encode": ("text_encoder",), + "idle": (), + }, + phase_present_components={ + "text_encode": ("text_encoder",), + "idle": ("transformer",), + }, + component_weight_gib={"transformer": 40}, + candidates=[text_offload, text_resident], + ) + ] + ) + + # The first warmup encoded text before the retained DiT existed. The + # next request starts with that 40 GiB DiT, so a resident encoder would + # need 90 GiB before the explicit reserve and is not feasible. + assert plan.changes == [] + def test_other_phase_can_block_component_residency(self): plan = plan_auto_residency( reports=[ @@ -826,6 +910,102 @@ def test_other_phase_can_block_component_residency(self): assert not plan.changes + def test_inactive_phase_charges_full_permanent_target(self): + layerwise = ResidencyTarget( + component_name="text_encoder", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + current_placement=True, + ) + resident = ResidencyTarget( + component_name="text_encoder", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=20 * GIB_BYTES, + target_layerwise_resident_layers=(48,), + target_layerwise_pinned_layers=((),), + permanent_residency=True, + inactive_device_delta_bytes=20 * GIB_BYTES, + target_device_weight_bytes=25 * GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=80, + estimated_gib=55, + target_workload_measured=True, + phase_peaks_gib={"denoise": 55}, + candidates=[layerwise, resident], + ) + ] + ) + + # The managed-layer delta would fit in the 21 GiB post-reserve + # headroom, but the complete permanent 25 GiB footprint does not. + assert plan.changes == [] + + def test_full_materialization_phase_does_not_double_count_transitioned_weights( + self, + ): + component_offload = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + current_placement=True, + ) + layerwise = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(1,), + target_layerwise_pinned_layers=((),), + active_device_delta_bytes=GIB_BYTES, + inactive_device_delta_bytes=GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=90, + estimated_gib=50, + target_workload_measured=True, + phase_peaks_gib={"lora_switch": 50}, + phase_full_weight_transition_components={ + "lora_switch": ("transformer",) + }, + candidates=[component_offload, layerwise], + ) + ] + ) + + assert [candidate.target_mode() for candidate in plan.changes] == [ + LAYERWISE_OFFLOAD + ] + + def test_reports_reserve_shortfall_for_a_newly_calibrated_placement(self): + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=30, + estimated_gib=28, + target_workload_measured=True, + phase_peaks_gib={"decode": 28}, + ) + ] + ) + + assert plan.current_placement_reserve_shortfall_bytes == 2 * GIB_BYTES + def test_unobserved_later_component_gets_a_conservative_phase(self): first_transformer = _candidate("first_transformer", weight_gib=26, h2d_gib=100) second_transformer = ResidencyTarget( @@ -1700,6 +1880,11 @@ def test_component_offload_frontier_can_expand_layerwise_lazily(self): ) assert component.current_placement assert full_stage.h2d_bytes_per_request == component.h2d_bytes_per_request + assert all( + candidate.h2d_bytes_per_request <= component.h2d_bytes_per_request + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + ) assert any( candidate.target_layerwise_resident_layers == (2,) for candidate in candidates @@ -1784,7 +1969,6 @@ def test_layerwise_frontier_accounts_for_immediate_materialization(self): assert ( resident.device_transition_delta_bytes == manager.offloaded_weight_bytes() ) - module.disable_offload() candidates = collect_residency_targets( modules={"transformer": module}, @@ -2157,6 +2341,47 @@ def test_vram_handoff_releases_before_materializing(self): assert events == ["source:sync", "target:load"] + def test_vram_rollback_releases_before_restoring_resident_layers(self): + events: list[str] = [] + + class _TrackedModule(nn.Module): + def to(self, *args, **kwargs): + events.append("text_encoder:release") + return self + + manager = _FakeLayerwiseManager( + {"layers.0.w": torch.zeros(16)}, + event_log=events, + event_name="transformer", + ) + transformer = _FakeLayerwiseDit([manager]) + applied = [ + AppliedResidencyChange( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + applied_device_delta_bytes=10, + ), + AppliedResidencyChange( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + previous_layerwise_resident_layers=(1,), + previous_layerwise_pinned_layers=((),), + previous_layerwise_offload_enabled=False, + applied_device_delta_bytes=-10, + ), + ] + + rollback_residency_changes( + applied=applied, + modules={ + "text_encoder": _TrackedModule(), + "transformer": transformer, + }, + server_args=_StubResidencyArgs(), + ) + + assert events == ["text_encoder:release", "transformer:load"] + def test_hostpin_rollback_releases_new_owner_before_restoring_old_owner( self, monkeypatch ): diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py index c74acbb04f29..934362204e50 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -10,11 +10,14 @@ from sglang.multimodal_gen.runtime.managers import gpu_worker as gpu_worker_module from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( + GIB_BYTES, PLACEMENT_STATUS_ADJUSTED, PLACEMENT_STATUS_ROLLBACK_FAILED, PLACEMENT_STATUS_ROLLED_BACK, PLACEMENT_STATUS_SKIPPED, AppliedResidencyChange, + AutoResidencyPlan, + DefaultWorkload, RankResidencyReport, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( @@ -64,6 +67,51 @@ def test_worker_rolls_back_a_materially_slower_calibrated_round(self): self.assertFalse(kwargs["already_failed"]) self.assertTrue(kwargs["latest_round_only"]) + def test_new_placement_that_spends_the_reserve_rolls_back(self): + worker = GPUWorker.__new__(GPUWorker) + worker.rank = 0 + worker.is_output_rank = False + worker.pipeline = SimpleNamespace(modules={}) + worker.server_args = SimpleNamespace() + worker._auto_residency_warmup_records = [] + worker._auto_residency_round_sizes = [1] + worker._build_auto_residency_report = mock.Mock( + return_value=RankResidencyReport( + rank=0, + budget_bytes=30 * GIB_BYTES, + estimated_peak_bytes=28 * GIB_BYTES, + ) + ) + worker._auto_residency_all_gather = lambda value: [value] + expected = OutputBatch(output={"status": PLACEMENT_STATUS_ROLLED_BACK}) + worker._rollback_everywhere = mock.Mock(return_value=expected) + workload = DefaultWorkload(832, 480, 24, 50) + plan = AutoResidencyPlan( + current_placement_reserve_shortfall_bytes=2 * GIB_BYTES + ) + + with ( + mock.patch.object( + gpu_worker_module, "resolve_default_workload", return_value=workload + ), + mock.patch.object( + gpu_worker_module, + "resolve_measured_default_workload", + return_value=workload, + ), + mock.patch.object( + gpu_worker_module, "plan_auto_residency", return_value=plan + ), + ): + response = worker.apply_auto_residency() + + self.assertIs(response, expected) + worker._rollback_everywhere.assert_called_once_with( + cause="VRAM reserve exceeded by 2.0 GiB", + already_failed=False, + latest_round_only=True, + ) + def test_worker_rollback_response_rewarms_the_restored_layout(self): actions = [] diff --git a/python/sglang/multimodal_gen/test/unit/test_component_residency.py b/python/sglang/multimodal_gen/test/unit/test_component_residency.py index 2da9bd5f1eeb..0cbbb8fb74eb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_residency.py @@ -234,6 +234,44 @@ def test_warmup_records_use_and_transition_peaks(monkeypatch): ) +def test_warmup_records_full_weight_transition_without_preparing(monkeypatch): + device_module = SimpleNamespace( + is_available=lambda: True, + reset_peak_memory_stats=Mock(), + max_memory_allocated=lambda: 7, + memory_allocated=lambda: 2, + ) + monkeypatch.setattr(torch, "get_device_module", lambda: device_module) + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + + stage = _Stage() + module = torch.nn.Linear(2, 2) + pipeline = SimpleNamespace( + modules={"transformer": module}, + _stage_name_mapping={"lora_switch": stage}, + component_residency_strategies={}, + ) + server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + manager = ComponentResidencyManager(pipeline, server_args) + manager.strategy_for = Mock() + manager.refresh_pipeline(pipeline) + manager.begin_request([stage], SimpleNamespace(is_warmup=True), server_args) + + manager.before_stage(stage, 0, SimpleNamespace(is_warmup=True), server_args) + with manager.full_weight_transition(("transformer",)): + pass + + assert manager._warmup_phase_peaks[ + "0:lora_switch:full-weight-transition:transformer" + ] == WarmupPhasePeak( + (), + 7, + full_weight_transition_components=("transformer",), + ) + assert manager._warmup_phase_key == "0:lora_switch:setup" + manager.strategy_for.assert_not_called() + + def test_warmup_records_same_component_dtype_prepare_as_transition(monkeypatch): device_module = SimpleNamespace( is_available=lambda: True, diff --git a/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py b/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py index 28b2a2340201..816f8a9c2371 100644 --- a/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py +++ b/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py @@ -1,7 +1,7 @@ from collections import defaultdict from contextlib import contextmanager, nullcontext from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import torch @@ -45,12 +45,36 @@ def _make_pipeline(layer: BaseLayerWithLoRA) -> _TestLoRAPipeline: pipeline.lora_layers_transformer_2 = {} pipeline.lora_layers_critic = {} pipeline.is_lora_merged = {} + pipeline.component_residency_manager = None pipeline.lora_adapters["adapter"]["linear.lora_A"] = torch.ones(1, 2) pipeline.lora_adapters["adapter"]["linear.lora_B"] = torch.ones(2, 1) return pipeline +def test_temporary_offload_materialization_is_reported_to_residency_manager(): + pipeline = _make_pipeline(_make_layer()) + module = Mock() + pipeline.modules = {"transformer": module} + residency_manager = Mock() + residency_manager.full_weight_transition.return_value = nullcontext() + pipeline.component_residency_manager = residency_manager + + with patch( + "sglang.multimodal_gen.runtime.pipelines_core.lora.pipeline." + "is_layerwise_offloaded_module", + return_value=True, + ): + with pipeline._temporarily_disable_offload( + target="transformer", use_module_names_only=True + ): + pass + + residency_manager.full_weight_transition.assert_called_once_with(["transformer"]) + module.disable_offload.assert_called_once_with() + module.enable_offload.assert_called_once_with() + + def test_merge_cache_only_accepts_cpu_backed_weights(): pipeline = _make_pipeline(_make_layer()) cpu_cache = pipeline._merge_cache_for( From cfa139609bf03d07b26725e12fdf2dd3edb404b2 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 12:37:25 +0800 Subject: [PATCH 40/90] [diffusion] refresh auto residency CI baselines --- .../test/server/perf_baselines/h100.json | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index 3f6f8b7f329a..d0cdfbcd40d1 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -104,7 +104,7 @@ "expected_avg_denoise_ms": 247.95, "expected_median_denoise_ms": 249.01, "load_peak_vram_mb": 41840.0, - "runtime_peak_vram_mb": 46938.0, + "runtime_peak_vram_mb": 58554.0, "estimated_full_test_time_s": 133.1 }, "qwen_image_t2i_2_gpus": { @@ -317,7 +317,7 @@ "ImageVAEEncodingStage": 0.01 }, "denoise_step_ms": { - "0": 65.28, + "0": 1082.71, "1": 129.11, "2": 440.87, "3": 457.02, @@ -778,7 +778,7 @@ "expected_avg_denoise_ms": 51.13, "expected_median_denoise_ms": 54.3, "load_peak_vram_mb": 13120.0, - "runtime_peak_vram_mb": 17602.0, + "runtime_peak_vram_mb": 24262.0, "estimated_full_test_time_s": 75.0 }, "qwen_image_edit_ti2i": { @@ -847,7 +847,7 @@ "expected_avg_denoise_ms": 633.05, "expected_median_denoise_ms": 636.16, "load_peak_vram_mb": 43244.0, - "runtime_peak_vram_mb": 48250.0, + "runtime_peak_vram_mb": 59796.0, "estimated_full_test_time_s": 153.6 }, "joyai_image_edit_ti2i": { @@ -1499,7 +1499,7 @@ }, "denoise_step_ms": { "0": 286.89, - "1": 167.02, + "1": 828.85, "2": 1010.47, "3": 1009.41, "4": 1011.33, @@ -1708,7 +1708,7 @@ "expected_avg_denoise_ms": 1923.51, "expected_median_denoise_ms": 1933.7, "load_peak_vram_mb": 35350.0, - "runtime_peak_vram_mb": 51230.0, + "runtime_peak_vram_mb": 58158.0, "estimated_full_test_time_s": 248.9 }, "wan2_2_t2v_a14b_2gpu": { @@ -2044,7 +2044,7 @@ "DecodingStage": 7.37 }, "denoise_step_ms": { - "0": 66.76, + "0": 1195.42, "1": 66.48, "2": 189.28, "3": 226.37, @@ -2336,7 +2336,7 @@ "ImageVAEEncodingStage": 0.01 }, "denoise_step_ms": { - "0": 66.32, + "0": 1110.98, "1": 341.96, "2": 419.84, "3": 460.86, @@ -2796,7 +2796,7 @@ "expected_avg_denoise_ms": 802.06, "expected_median_denoise_ms": 747.62, "load_peak_vram_mb": 56048.0, - "runtime_peak_vram_mb": 59804.0, + "runtime_peak_vram_mb": 66586.0, "estimated_full_test_time_s": 363.2 }, "qwen_image_t2i_cache_dit_scm_config_diffusers_1gpu": { From 95a4ee953d3c4ef5d98df015ef46695767f2bab1 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 14:40:55 +0800 Subject: [PATCH 41/90] [diffusion] account for layerwise LoRA transition memory --- .../memory_managers/auto_residency.py | 33 +++++++++++++-- .../test/server/perf_baselines/h100.json | 2 +- .../test/unit/test_auto_residency.py | 41 +++++++++++++++++++ 3 files changed, 72 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 3c3209aeb971..23f9fc32fa21 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -1826,6 +1826,14 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP {candidate.option_key(): candidate for candidate in report.candidates} for report in reports ] + current_report_candidates = [ + { + candidate.component_name: candidate + for candidate in report.candidates + if candidate.current_placement + } + for report in reports + ] # Request timing is produced by the output rank. Other SPMD ranks still # contribute their VRAM and HostPin constraints, but their empty metrics @@ -1857,8 +1865,13 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP options = [] for candidate in candidates: resource_deltas: dict[str, int] = {} - for report, rank_candidates in zip(reports, report_candidates): + for report, rank_candidates, current_rank_candidates in zip( + reports, report_candidates, current_report_candidates + ): rank_candidate = rank_candidates[candidate.option_key()] + current_rank_candidate = current_rank_candidates.get( + candidate.component_name + ) for ( resource_name, _, @@ -1868,8 +1881,22 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP ) in phase_constraints[report.rank]: if candidate.component_name in full_weight_transition_components: # The measured phase already contains the component's full - # temporary materialization under every target mode. - phase_cost = 0 + # temporary materialization. Partial layerwise residency + # remains live alongside that temporary copy (notably + # while LoRA switches own full-weight snapshots), so only + # its resident-layer delta is additional. Coarse offload + # and permanent residency replace the measured temporary + # materialization and therefore add no phase cost. + phase_cost = ( + rank_candidate.target_resident_weight_bytes + - ( + current_rank_candidate.target_resident_weight_bytes + if current_rank_candidate is not None + else 0 + ) + if rank_candidate.target_mode() == LAYERWISE_OFFLOAD + else 0 + ) elif candidate.component_name in used_components: phase_cost = rank_candidate.active_device_delta_bytes elif candidate.component_name in present_components: diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index d0cdfbcd40d1..5fde9f15c888 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -1647,7 +1647,7 @@ "TextEncodingStage": 327.33, "DecodingStage": 1357.99, "TimestepPreparationStage": 1.32, - "ImageEncodingStage": 22.1, + "ImageEncodingStage": 451.16, "LatentPreparationStage": 0.1, "DenoisingStage": 96179.65, "ImageVAEEncodingStage": 910.6 diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index dddff0f830fd..4c8d40013a79 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -992,6 +992,47 @@ def test_full_materialization_phase_does_not_double_count_transitioned_weights( LAYERWISE_OFFLOAD ] + def test_full_materialization_phase_accounts_for_partial_resident_layers(self): + component_offload = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + current_placement=True, + ) + layerwise = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(1,), + target_layerwise_pinned_layers=((),), + active_device_delta_bytes=GIB_BYTES, + inactive_device_delta_bytes=GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=54, + estimated_gib=50, + target_workload_measured=True, + phase_peaks_gib={"lora_switch": 50}, + phase_full_weight_transition_components={ + "lora_switch": ("transformer",) + }, + candidates=[component_offload, layerwise], + ) + ] + ) + + # The measured full-weight transition plus the 4 GiB reserve already + # consumes the budget. A partial resident layer remains live alongside + # that transition and must not be selected into the missing headroom. + assert plan.changes == [] + def test_reports_reserve_shortfall_for_a_newly_calibrated_placement(self): plan = plan_auto_residency( reports=[ From 7a4db981b178081f252a6440774e2774f5eec6f9 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 16:34:56 +0800 Subject: [PATCH 42/90] [diffusion] record coarse LoRA weight transitions --- .../memory_managers/auto_residency.py | 34 ++------------- .../runtime/pipelines_core/lora/pipeline.py | 7 ++-- .../test/unit/test_auto_residency.py | 41 ------------------- .../test/unit/test_lora_pipeline.py | 23 +++++++++++ 4 files changed, 30 insertions(+), 75 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 23f9fc32fa21..e756afbcbdfd 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -1826,15 +1826,6 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP {candidate.option_key(): candidate for candidate in report.candidates} for report in reports ] - current_report_candidates = [ - { - candidate.component_name: candidate - for candidate in report.candidates - if candidate.current_placement - } - for report in reports - ] - # Request timing is produced by the output rank. Other SPMD ranks still # contribute their VRAM and HostPin constraints, but their empty metrics # must not discard the replica's measured latency utility and fall back to @@ -1865,13 +1856,8 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP options = [] for candidate in candidates: resource_deltas: dict[str, int] = {} - for report, rank_candidates, current_rank_candidates in zip( - reports, report_candidates, current_report_candidates - ): + for report, rank_candidates in zip(reports, report_candidates): rank_candidate = rank_candidates[candidate.option_key()] - current_rank_candidate = current_rank_candidates.get( - candidate.component_name - ) for ( resource_name, _, @@ -1881,22 +1867,8 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP ) in phase_constraints[report.rank]: if candidate.component_name in full_weight_transition_components: # The measured phase already contains the component's full - # temporary materialization. Partial layerwise residency - # remains live alongside that temporary copy (notably - # while LoRA switches own full-weight snapshots), so only - # its resident-layer delta is additional. Coarse offload - # and permanent residency replace the measured temporary - # materialization and therefore add no phase cost. - phase_cost = ( - rank_candidate.target_resident_weight_bytes - - ( - current_rank_candidate.target_resident_weight_bytes - if current_rank_candidate is not None - else 0 - ) - if rank_candidate.target_mode() == LAYERWISE_OFFLOAD - else 0 - ) + # temporary materialization under every target mode. + phase_cost = 0 elif candidate.component_name in used_components: phase_cost = rank_candidate.active_device_delta_bytes elif candidate.component_name in present_components: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py index dd48efe962df..5b0aa40227a6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py @@ -297,10 +297,11 @@ def _temporarily_disable_offload( and is_layerwise_offloaded_module(module) ) ] + # Record every target, including coarse component offload. The current + # coarse path can already materialize the full component here, and the + # planner must preserve that phase if it later chooses layerwise mode. residency_transition = ( - self.component_residency_manager.full_weight_transition( - offloaded_module_names - ) + self.component_residency_manager.full_weight_transition(module_names) if self.component_residency_manager is not None else nullcontext() ) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 4c8d40013a79..dddff0f830fd 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -992,47 +992,6 @@ def test_full_materialization_phase_does_not_double_count_transitioned_weights( LAYERWISE_OFFLOAD ] - def test_full_materialization_phase_accounts_for_partial_resident_layers(self): - component_offload = ResidencyTarget( - component_name="transformer", - residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=COMPONENT_OFFLOAD, - target_resident_weight_bytes=0, - h2d_bytes_per_request=0, - current_placement=True, - ) - layerwise = ResidencyTarget( - component_name="transformer", - residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=GIB_BYTES, - h2d_bytes_per_request=100 * GIB_BYTES, - target_layerwise_resident_layers=(1,), - target_layerwise_pinned_layers=((),), - active_device_delta_bytes=GIB_BYTES, - inactive_device_delta_bytes=GIB_BYTES, - ) - - plan = plan_auto_residency( - reports=[ - _report( - budget_gib=54, - estimated_gib=50, - target_workload_measured=True, - phase_peaks_gib={"lora_switch": 50}, - phase_full_weight_transition_components={ - "lora_switch": ("transformer",) - }, - candidates=[component_offload, layerwise], - ) - ] - ) - - # The measured full-weight transition plus the 4 GiB reserve already - # consumes the budget. A partial resident layer remains live alongside - # that transition and must not be selected into the missing headroom. - assert plan.changes == [] - def test_reports_reserve_shortfall_for_a_newly_calibrated_placement(self): plan = plan_auto_residency( reports=[ diff --git a/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py b/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py index 816f8a9c2371..3e0d8395d852 100644 --- a/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py +++ b/python/sglang/multimodal_gen/test/unit/test_lora_pipeline.py @@ -75,6 +75,29 @@ def test_temporary_offload_materialization_is_reported_to_residency_manager(): module.enable_offload.assert_called_once_with() +def test_component_offload_weight_update_is_reported_without_disabling_offload(): + pipeline = _make_pipeline(_make_layer()) + module = Mock() + pipeline.modules = {"transformer": module} + residency_manager = Mock() + residency_manager.full_weight_transition.return_value = nullcontext() + pipeline.component_residency_manager = residency_manager + + with patch( + "sglang.multimodal_gen.runtime.pipelines_core.lora.pipeline." + "is_layerwise_offloaded_module", + return_value=False, + ): + with pipeline._temporarily_disable_offload( + target="transformer", use_module_names_only=True + ): + pass + + residency_manager.full_weight_transition.assert_called_once_with(["transformer"]) + module.disable_offload.assert_not_called() + module.enable_offload.assert_not_called() + + def test_merge_cache_only_accepts_cpu_backed_weights(): pipeline = _make_pipeline(_make_layer()) cpu_cache = pipeline._merge_cache_for( From c4ce470f058fdbc5777c29d7b2862068b73692c7 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 18:40:51 +0800 Subject: [PATCH 43/90] [diffusion] account for unmaterialized LoRA transitions --- .../memory_managers/auto_residency.py | 60 ++++++++++++++-- .../test/unit/test_auto_residency.py | 68 +++++++++++++++++-- 2 files changed, 120 insertions(+), 8 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index e756afbcbdfd..cd2825be93ab 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -1826,6 +1826,27 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP {candidate.option_key(): candidate for candidate in report.candidates} for report in reports ] + current_report_candidates = [ + { + candidate.component_name: candidate + for candidate in report.candidates + if candidate.current_placement + } + for report in reports + ] + full_device_weight_bytes_by_component = [ + { + component_name: max( + candidate.target_device_weight_bytes + for candidate in report.candidates + if candidate.component_name == component_name + ) + for component_name in { + candidate.component_name for candidate in report.candidates + } + } + for report in reports + ] # Request timing is produced by the output rank. Other SPMD ranks still # contribute their VRAM and HostPin constraints, but their empty metrics # must not discard the replica's measured latency utility and fall back to @@ -1856,8 +1877,21 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP options = [] for candidate in candidates: resource_deltas: dict[str, int] = {} - for report, rank_candidates in zip(reports, report_candidates): + for ( + report, + rank_candidates, + current_rank_candidates, + rank_full_device_weight_bytes, + ) in zip( + reports, + report_candidates, + current_report_candidates, + full_device_weight_bytes_by_component, + ): rank_candidate = rank_candidates[candidate.option_key()] + current_rank_candidate = current_rank_candidates.get( + candidate.component_name + ) for ( resource_name, _, @@ -1865,10 +1899,28 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP used_components, full_weight_transition_components, ) in phase_constraints[report.rank]: - if candidate.component_name in full_weight_transition_components: - # The measured phase already contains the component's full - # temporary materialization under every target mode. + transition_measured_full_weights = ( + candidate.component_name in full_weight_transition_components + and ( + candidate.component_name in present_components + or ( + current_rank_candidate is not None + and current_rank_candidate.target_mode() + in (LAYERWISE_OFFLOAD, RESIDENT) + ) + ) + ) + if transition_measured_full_weights: + # The measured phase already contains the complete weights. phase_cost = 0 + elif ( + candidate.component_name in full_weight_transition_components + and rank_candidate.target_mode() == LAYERWISE_OFFLOAD + ): + # Coarse offload can update CPU weights without putting the + # component on the device. Layerwise offload must materialize + # every layer for the same update, including non-resident ones. + phase_cost = rank_full_device_weight_bytes[candidate.component_name] elif candidate.component_name in used_components: phase_cost = rank_candidate.active_device_delta_bytes elif candidate.component_name in present_components: diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index dddff0f830fd..336473e1d80e 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -950,9 +950,7 @@ def test_inactive_phase_charges_full_permanent_target(self): # headroom, but the complete permanent 25 GiB footprint does not. assert plan.changes == [] - def test_full_materialization_phase_does_not_double_count_transitioned_weights( - self, - ): + def test_unmaterialized_weight_update_charges_full_layerwise_target(self): component_offload = ResidencyTarget( component_name="transformer", residency_mode=COMPONENT_OFFLOAD, @@ -971,6 +969,68 @@ def test_full_materialization_phase_does_not_double_count_transitioned_weights( target_layerwise_pinned_layers=((),), active_device_delta_bytes=GIB_BYTES, inactive_device_delta_bytes=GIB_BYTES, + target_device_weight_bytes=GIB_BYTES, + ) + resident = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=40 * GIB_BYTES, + h2d_bytes_per_request=0, + permanent_residency=True, + inactive_device_delta_bytes=40 * GIB_BYTES, + target_device_weight_bytes=40 * GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=90, + estimated_gib=50, + target_workload_measured=True, + phase_peaks_gib={"lora_switch": 50}, + phase_full_weight_transition_components={ + "lora_switch": ("transformer",) + }, + candidates=[component_offload, layerwise, resident], + ) + ] + ) + + assert plan.changes == [] + + def test_materialized_layerwise_transition_does_not_double_count_weights(self): + current_layerwise = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + current_placement=True, + ) + layerwise = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(1,), + target_layerwise_pinned_layers=((),), + active_device_delta_bytes=GIB_BYTES, + inactive_device_delta_bytes=GIB_BYTES, + target_device_weight_bytes=GIB_BYTES, + ) + resident = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=40 * GIB_BYTES, + h2d_bytes_per_request=0, + permanent_residency=True, + inactive_device_delta_bytes=40 * GIB_BYTES, + target_device_weight_bytes=40 * GIB_BYTES, ) plan = plan_auto_residency( @@ -983,7 +1043,7 @@ def test_full_materialization_phase_does_not_double_count_transitioned_weights( phase_full_weight_transition_components={ "lora_switch": ("transformer",) }, - candidates=[component_offload, layerwise], + candidates=[current_layerwise, layerwise, resident], ) ] ) From ade6ab6db08f6de5c613ce78f5857dd8739cebaf Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 20:41:41 +0800 Subject: [PATCH 44/90] [diffusion] distinguish measured residency transitions --- .../memory_managers/auto_residency.py | 39 ++++++++++--- .../test/unit/test_auto_residency.py | 58 +++++++++++++++++++ 2 files changed, 88 insertions(+), 9 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index cd2825be93ab..f9b64481ad48 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -1501,7 +1501,16 @@ def _vram_reserve_bytes(budget_bytes: int, *, target_workload_measured: bool) -> def _binding_phase_constraints( report: RankResidencyReport, candidate_component_names: set[str], -) -> list[tuple[str, int, tuple[str, ...], tuple[str, ...], tuple[str, ...]]]: +) -> list[ + tuple[ + str, + int, + tuple[str, ...], + tuple[str, ...], + tuple[str, ...], + tuple[str, ...], + ] +]: """Keep the binding measured and conservative unobserved phases. Every candidate has the same device delta in phases with the same active @@ -1542,7 +1551,13 @@ def _binding_phase_constraints( } binding: dict[ - tuple[tuple[str, ...], tuple[str, ...], tuple[str, ...]], tuple[str, int] + tuple[ + tuple[str, ...], + tuple[str, ...], + tuple[str, ...], + tuple[str, ...], + ], + tuple[str, int], ] = {} for phase_name, phase_peak in phase_peaks.items(): measured_components = set(report.active_components_by_phase.get(phase_name, ())) @@ -1557,6 +1572,7 @@ def _binding_phase_constraints( & candidate_component_names ) ) + measured = tuple(sorted(measured_components & candidate_component_names)) used = tuple(sorted(report.used_components_by_phase.get(phase_name, ()))) full_weight_transitions = tuple( sorted( @@ -1568,13 +1584,13 @@ def _binding_phase_constraints( & candidate_component_names ) ) - layout = (present, used, full_weight_transitions) + layout = (present, measured, used, full_weight_transitions) current = binding.get(layout) if current is None or (phase_peak, phase_name) > (current[1], current[0]): binding[layout] = (phase_name, phase_peak) observed_components = { component_name - for (_, used, full_weight_transitions) in binding + for (_, _, used, full_weight_transitions) in binding for component_name in set(used) | set(full_weight_transitions) } steady_request_peak = max( @@ -1587,7 +1603,7 @@ def _binding_phase_constraints( (steady_state_components | {component_name}) & candidate_component_names ) ) - binding[(present, (component_name,), ())] = ( + binding[(present, (), (component_name,), ())] = ( f"unobserved:{component_name}", steady_request_peak, ) @@ -1596,10 +1612,11 @@ def _binding_phase_constraints( f"gpu:rank{report.rank}:{phase_name}", phase_peak, present, + measured, used, full_weight_transitions, ) - for (present, used, full_weight_transitions), ( + for (present, measured, used, full_weight_transitions), ( phase_name, phase_peak, ) in sorted(binding.items(), key=lambda item: item[1][0]) @@ -1767,7 +1784,7 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP max( ( phase_peak + reserves_by_rank[report.rank] - report.budget_bytes - for _, phase_peak, _, _, _ in phase_constraints[report.rank] + for _, phase_peak, _, _, _, _ in phase_constraints[report.rank] ), default=0, ) @@ -1778,7 +1795,7 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP current_placement_reserve_shortfall = max(0, current_placement_reserve_shortfall) resource_budgets: dict[str, int] = {} for report in reports: - for resource_name, phase_peak, _, _, _ in phase_constraints[report.rank]: + for resource_name, phase_peak, _, _, _, _ in phase_constraints[report.rank]: # The measured placement has already completed warmup. A negative # reserve headroom therefore means "do not grow this phase", not # that the current zero-delta placement is infeasible. @@ -1896,13 +1913,17 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP resource_name, _, present_components, + measured_components, used_components, full_weight_transition_components, ) in phase_constraints[report.rank]: + # ``present_components`` also includes request-end placement + # carried into this phase. Only the measured set proves that + # the phase peak already captured a complete materialization. transition_measured_full_weights = ( candidate.component_name in full_weight_transition_components and ( - candidate.component_name in present_components + candidate.component_name in measured_components or ( current_rank_candidate is not None and current_rank_candidate.target_mode() diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 336473e1d80e..59779a12fec2 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -989,6 +989,10 @@ def test_unmaterialized_weight_update_charges_full_layerwise_target(self): estimated_gib=50, target_workload_measured=True, phase_peaks_gib={"lora_switch": 50}, + phase_present_components={ + "lora_switch": (), + "idle": ("transformer",), + }, phase_full_weight_transition_components={ "lora_switch": ("transformer",) }, @@ -999,6 +1003,60 @@ def test_unmaterialized_weight_update_charges_full_layerwise_target(self): assert plan.changes == [] + def test_measured_component_transition_does_not_double_count_weights(self): + component_offload = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + current_placement=True, + ) + layerwise = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(1,), + target_layerwise_pinned_layers=((),), + active_device_delta_bytes=GIB_BYTES, + inactive_device_delta_bytes=GIB_BYTES, + target_device_weight_bytes=GIB_BYTES, + ) + resident = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=40 * GIB_BYTES, + h2d_bytes_per_request=0, + permanent_residency=True, + inactive_device_delta_bytes=40 * GIB_BYTES, + target_device_weight_bytes=40 * GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=90, + estimated_gib=50, + target_workload_measured=True, + phase_peaks_gib={"lora_switch": 50}, + phase_present_components={ + "lora_switch": ("transformer",), + }, + phase_full_weight_transition_components={ + "lora_switch": ("transformer",) + }, + candidates=[component_offload, layerwise, resident], + ) + ] + ) + + assert [candidate.target_mode() for candidate in plan.changes] == [ + LAYERWISE_OFFLOAD + ] + def test_materialized_layerwise_transition_does_not_double_count_weights(self): current_layerwise = ResidencyTarget( component_name="transformer", From 137af033b2400db1bc6938c7e4de5085ce3b935b Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 02:57:33 +0800 Subject: [PATCH 45/90] [diffusion] score residency relative to calibrated placement --- .../memory_managers/auto_residency.py | 40 ++++++++++++++++++- .../test/unit/test_auto_residency.py | 36 +++++++++++++++++ 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index f9b64481ad48..3b7fe88004fd 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -436,14 +436,50 @@ def estimate_candidate_latency_savings_ns( candidates: Iterable[ResidencyTarget], request_duration_ns: int, ) -> dict[str, int]: - """Upper-bound H2D time removed by each complete placement option. + """Estimate each complete placement relative to the measured one. An async prefetch can consume copy-engine and memory bandwidth outside the component's own stage, so that stage is not a sound cap. The complete - request duration is the only generally valid upper bound. + request duration bounds removable latency. Added transfer time is not + capped: an unmeasured mechanism may be slower than the whole request. """ + candidates = list(candidates) + candidates_by_component: dict[str, list[ResidencyTarget]] = {} + for candidate in candidates: + candidates_by_component.setdefault(candidate.component_name, []).append( + candidate + ) + estimates: dict[str, int] = {} + for component_candidates in candidates_by_component.values(): + current = next( + ( + candidate + for candidate in component_candidates + if candidate.current_placement + ), + None, + ) + if current is None: + continue + + for candidate in component_candidates: + relative_savings_ns = int( + (candidate.h2d_bytes_per_request - current.h2d_bytes_per_request) + / ESTIMATED_PINNED_H2D_BYTES_PER_SECOND + * 1_000_000_000 + ) + # The request duration bounds removable latency, but not latency an + # unmeasured placement can add. Capping both directions can make a + # repeatedly streamed DiT appear neutral and let an unrelated small + # gain select a several-times-slower combined placement. + if relative_savings_ns > 0 and request_duration_ns > 0: + relative_savings_ns = min(relative_savings_ns, request_duration_ns) + estimates[candidate.option_key()] = relative_savings_ns + for candidate in candidates: + if candidate.option_key() in estimates: + continue transfer_ns = int( candidate.h2d_bytes_per_request / ESTIMATED_PINNED_H2D_BYTES_PER_SECOND diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 59779a12fec2..ad7916724056 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -154,6 +154,42 @@ def test_transfer_savings_is_capped_by_request_not_component_stage(self): assert savings[candidates[1].option_key()] > 400_000_000 assert savings[candidates[2].option_key()] < 0 + def test_measured_placement_keeps_full_cost_of_slower_dit_target(self): + coarse = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=140 * GIB_BYTES, + current_placement=True, + ) + virtual_layerwise = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=10 * GIB_BYTES, + h2d_bytes_per_request=20 * GIB_BYTES, + target_layerwise_resident_layers=(10,), + target_layerwise_pinned_layers=((),), + ) + resident = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=60 * GIB_BYTES, + h2d_bytes_per_request=150 * GIB_BYTES, + permanent_residency=True, + ) + + savings = estimate_candidate_latency_savings_ns( + candidates=[coarse, virtual_layerwise, resident], + request_duration_ns=1_000_000_000, + ) + + assert savings[coarse.option_key()] == 0 + assert savings[resident.option_key()] == 416_666_666 + assert savings[virtual_layerwise.option_key()] == -5_000_000_000 + class TestEstimateDefaultWorkloadPeak: def test_same_shape_uses_measured_peak(self): From 0aed0aba7401c42f6fafcf453658d2050cdcf659 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 13:22:02 +0800 Subject: [PATCH 46/90] [diffusion] calibrate residency with the CI workload --- .../runtime/warmup_request_builder.py | 8 ++- .../test/server/test_server_common.py | 24 ++++--- .../test/server/testcase_configs.py | 25 +++++-- .../test/unit/test_cfg_parallel_warmup.py | 27 ++++++++ .../test/unit/test_health_warmup_gate.py | 69 +++++++++++++++---- 5 files changed, 126 insertions(+), 27 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py index 038420365f54..86293a38699e 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -601,7 +601,13 @@ def build_warmup_reqs( "Warmup image path is required for image-input model" ) req_kwargs["prompt"] = DEFAULT_PLACEHOLDER_PROMPT - req_kwargs["image_path"] = [warmup_input_path] + default_image_path = sampling_defaults.image_path + image_count = ( + len(default_image_path) + if isinstance(default_image_path, (list, tuple)) + else 1 + ) + req_kwargs["image_path"] = [warmup_input_path] * max(1, image_count) if server_args.enable_cfg_parallel: if not req_kwargs.get("negative_prompt"): req_kwargs["negative_prompt"] = DEFAULT_PLACEHOLDER_PROMPT diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 7981daf519f6..649d2538072d 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -48,6 +48,7 @@ ScenarioConfig, get_model_task_type_for_server_args, get_perf_baseline_update_path, + get_sampling_param_field_names_for_server_args, ) from sglang.multimodal_gen.test.test_utils import ( SGL_TEST_FILES_CI_DATA_REVISION, @@ -118,25 +119,32 @@ def _case_warmup_sampling_params(case: DiffusionTestCase) -> dict[str, Any]: sampling = case.sampling_params if sampling is None: return {} - overrides: dict[str, Any] = {} + sampling_field_names = get_sampling_param_field_names_for_server_args( + case.server_args + ) + overrides = { + name: value + for name, value in sampling.extras.items() + if name in sampling_field_names + } if sampling.output_size: width, height = sampling.output_size.lower().split("x", 1) overrides.update(width=int(width), height=int(height)) - if case.server_args.modality == "image": + if sampling.num_frames is not None: + overrides["num_frames"] = sampling.num_frames + elif case.server_args.modality == "image": # Omni pipelines can expose a video-oriented model default even when # this case calls /v1/images, whose implicit temporal extent is one. overrides["num_frames"] = 1 - elif sampling.num_frames is not None: - overrides["num_frames"] = sampling.num_frames else: # /v1/videos resolves an omitted num_frames from seconds * fps. overrides["num_frames"] = sampling.seconds * (sampling.fps or DEFAULT_FPS) if sampling.fps is not None: overrides["fps"] = sampling.fps - for name in ("num_inference_steps", "guidance_scale", "true_cfg_scale"): - value = sampling.extras.get(name) - if value is not None: - overrides[name] = value + if sampling.num_outputs_per_prompt > 1: + overrides["num_outputs_per_prompt"] = sampling.num_outputs_per_prompt + if isinstance(sampling.image_path, (list, tuple)): + overrides["image_path"] = ["warmup"] * len(sampling.image_path) return overrides diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 981e117eb395..a1c0ef20ecfe 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -23,7 +23,7 @@ import os import shlex import statistics -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, field, fields, replace from functools import lru_cache from pathlib import Path from typing import Any, Sequence @@ -787,22 +787,35 @@ def _get_extra_arg_value(extras: Sequence[str], option_name: str) -> str | None: return None -def get_model_task_type_for_server_args( +def _get_config_classes_for_server_args( server_args: DiffusionServerArgs, -) -> ModelTaskType: +) -> tuple[type, type]: pipeline_class_name = _get_extra_arg_value( server_args.extras, "--pipeline-class-name" ) if pipeline_class_name: config_classes = get_pipeline_config_classes(pipeline_class_name) if config_classes is not None: - pipeline_config_cls, _ = config_classes - return pipeline_config_cls.task_type + return config_classes model_info = get_model_info(server_args.model_path) if model_info is None: raise ValueError(f"Could not resolve model info for {server_args.model_path!r}") - return model_info.pipeline_config_cls.task_type + return model_info.pipeline_config_cls, model_info.sampling_param_cls + + +def get_model_task_type_for_server_args( + server_args: DiffusionServerArgs, +) -> ModelTaskType: + pipeline_config_cls, _ = _get_config_classes_for_server_args(server_args) + return pipeline_config_cls.task_type + + +def get_sampling_param_field_names_for_server_args( + server_args: DiffusionServerArgs, +) -> frozenset[str]: + _, sampling_param_cls = _get_config_classes_for_server_args(server_args) + return frozenset(item.name for item in fields(sampling_param_cls)) def get_default_sampling_params_for_model_task( diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index bb2aca7df48c..1766a081f9c0 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py @@ -883,6 +883,33 @@ def test_server_based_warmup_keeps_ti2i_image_input(self): self.assertEqual(reqs[0].image_path, ["/tmp/warmup.png"]) + def test_server_based_warmup_keeps_image_input_count(self): + server_args = MagicMock() + server_args.warmup_steps = 1 + server_args.enable_cfg_parallel = False + server_args.enable_torch_compile = False + server_args.pipeline_config.task_type = ModelTaskType.TI2I + + with patch( + "sglang.multimodal_gen.runtime.warmup_request_builder.get_model_sampling_defaults", + return_value=SamplingParams( + width=512, + height=512, + image_path=["first.png", "second.png"], + ), + ): + reqs = build_warmup_reqs( + server_args, + warmup_resolutions=None, + warmup_input_path="/tmp/warmup.png", + server_based_warmup=True, + ) + + self.assertEqual( + reqs[0].image_path, + ["/tmp/warmup.png", "/tmp/warmup.png"], + ) + def test_server_based_warmup_keeps_required_image_input(self): server_args = MagicMock() server_args.warmup_steps = 1 diff --git a/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py b/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py index 52f68952be03..d3c77eecdb3b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py +++ b/python/sglang/multimodal_gen/test/unit/test_health_warmup_gate.py @@ -123,20 +123,58 @@ def test_image_case_warmup_uses_one_frame_for_omni_model(self): num_frames=None, fps=None, seconds=1, + num_outputs_per_prompt=1, + image_path=None, extras={"num_inference_steps": 35}, ), ) - self.assertEqual( - _case_warmup_sampling_params(case), - { - "width": 832, - "height": 480, - "num_frames": 1, - "num_inference_steps": 35, - }, + with mock.patch( + "sglang.multimodal_gen.test.server.test_server_common." + "get_sampling_param_field_names_for_server_args", + return_value=frozenset({"num_inference_steps"}), + ): + self.assertEqual( + _case_warmup_sampling_params(case), + { + "width": 832, + "height": 480, + "num_frames": 1, + "num_inference_steps": 35, + }, + ) + + def test_image_case_warmup_keeps_explicit_frames_and_inputs(self): + case = SimpleNamespace( + server_args=SimpleNamespace(modality="image"), + sampling_params=SimpleNamespace( + output_size="1024x1024", + num_frames=4, + fps=None, + seconds=1, + num_outputs_per_prompt=2, + image_path=["first.png", "second.png"], + extras={"enable_upscaling": True, "test_only": "ignored"}, + ), ) + with mock.patch( + "sglang.multimodal_gen.test.server.test_server_common." + "get_sampling_param_field_names_for_server_args", + return_value=frozenset({"enable_upscaling"}), + ): + self.assertEqual( + _case_warmup_sampling_params(case), + { + "width": 1024, + "height": 1024, + "num_frames": 4, + "num_outputs_per_prompt": 2, + "image_path": ["warmup", "warmup"], + "enable_upscaling": True, + }, + ) + def test_video_case_warmup_matches_endpoint_default_fps(self): case = SimpleNamespace( server_args=SimpleNamespace(modality="video"), @@ -145,14 +183,21 @@ def test_video_case_warmup_matches_endpoint_default_fps(self): num_frames=None, fps=None, seconds=2, + num_outputs_per_prompt=1, + image_path=None, extras={}, ), ) - self.assertEqual( - _case_warmup_sampling_params(case), - {"num_frames": 48}, - ) + with mock.patch( + "sglang.multimodal_gen.test.server.test_server_common." + "get_sampling_param_field_names_for_server_args", + return_value=frozenset(), + ): + self.assertEqual( + _case_warmup_sampling_params(case), + {"num_frames": 48}, + ) def test_waits_for_health_after_http_startup(self): manager = ServerManager("test-model", port=11000, wait_deadline=1) From 21743e2ccda5bbb28decdd1ed181eb95c95a314b Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 06:52:02 +0800 Subject: [PATCH 47/90] fix(diffusion): stabilize residency calibration --- .../managers/memory_managers/auto_residency.py | 13 +++++++++---- .../runtime/warmup_request_builder.py | 10 ++++++++++ .../test/unit/test_auto_residency.py | 16 ++++++++++++++-- .../test/unit/test_cfg_parallel_warmup.py | 4 ++-- 4 files changed, 35 insertions(+), 8 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 3b7fe88004fd..369768afa5eb 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -2089,10 +2089,10 @@ def apply_residency_changes( """Apply complete target states transactionally on this rank. Component-offload targets move back to CPU immediately; resident targets - are materialized on their next use by ``ResidentStrategy``. Layerwise - targets may change resident counts and HostPin placement or disable the - streaming hooks entirely. Because targets can move in either direction, - a later calibration round can replace an earlier placement. + materialize on the local device before validation. Layerwise targets may + change resident counts and HostPin placement or disable the streaming + hooks entirely. Because targets can move in either direction, a later + calibration round can replace an earlier placement. Raises on failure after undoing the changes already applied, so a caller observing an exception knows this rank is back on the original @@ -2313,6 +2313,11 @@ def apply_residency_changes( continue if target_mode == RESIDENT: server_args.set_auto_residency_mode(candidate.component_name, RESIDENT) + if not is_fsdp_managed_module(module): + module.to( + current_platform.get_local_torch_device(), + non_blocking=True, + ) elif target_mode == COMPONENT_OFFLOAD: server_args.set_auto_residency_mode( candidate.component_name, COMPONENT_OFFLOAD diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py index 86293a38699e..33e039b5894f 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -57,6 +57,11 @@ SERVER_WARMUP_LTX2_TWO_STAGE_MAX_VIDEO_FRAMES = 25 SERVER_WARMUP_IMAGE_STEPS = 2 SERVER_WARMUP_VIDEO_STEPS = 2 +# Two-step schedulers can have one compile-heavy step and one lower-order +# boundary step, leaving no representative steady-state timing sample. Auto +# residency extrapolates this timing to the default request, so collect four +# steps while retaining the shorter warmup for every non-planning path. +AUTO_RESIDENCY_TIMING_STEPS = 4 def get_model_sampling_defaults(server_args: ServerArgs) -> SamplingParams: @@ -572,6 +577,11 @@ def build_warmup_reqs( and server_based_warmup and auto_residency_args_skip_reason(server_args) is None ) + if collect_auto_residency_metrics and sampling_defaults.num_inference_steps: + warmup_steps = min( + int(sampling_defaults.num_inference_steps), + max(warmup_steps, AUTO_RESIDENCY_TIMING_STEPS), + ) if auto_residency_warmup_shape is not None: width, height, warmup_num_frames = auto_residency_warmup_shape resolutions[0] = (width, height) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index ad7916724056..d40f7cf248fb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -43,6 +43,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, ) +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.warmup_request_builder import ( SERVER_WARMUP_MAX_VIDEO_FRAMES, _resolve_auto_residency_warmup_shape, @@ -2224,8 +2225,17 @@ def test_layerwise_can_become_component_offload_and_rollback(self): assert args.auto_modes == {"text_encoder": LAYERWISE_OFFLOAD} assert module.layerwise_offload_managers[0].pinned_layer_indices() == (0,) - def test_component_offload_promotion_marks_resident_without_moving(self): - module = nn.Linear(4, 4) + def test_component_offload_promotion_materializes_before_validation(self): + class RecordingLinear(nn.Linear): + def __init__(self): + super().__init__(4, 4) + self.to_targets = [] + + def to(self, *args, **kwargs): + self.to_targets.append(args[0]) + return super().to(*args, **kwargs) + + module = RecordingLinear() args = _StubResidencyArgs() applied = apply_residency_changes( plan=_plan_for([_candidate("text_encoder", weight_gib=1)]), @@ -2233,12 +2243,14 @@ def test_component_offload_promotion_marks_resident_without_moving(self): server_args=args, ) assert args.auto_modes == {"text_encoder": RESIDENT} + assert module.to_targets == [current_platform.get_local_torch_device()] assert [p.component_name for p in applied] == ["text_encoder"] rollback_residency_changes( applied=applied, modules={"text_encoder": module}, server_args=args ) assert args.auto_modes == {} + assert module.to_targets[-1] == "cpu" def test_component_resident_can_be_demoted_and_rollback_restores_it(self): module = nn.Linear(4, 4) diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index 1766a081f9c0..44d2289ee3ac 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py @@ -430,9 +430,9 @@ def test_auto_residency_uses_one_full_serving_shape_probe(self): ) self.assertTrue(reqs[0].extra["auto_residency_full_shape_probe"]) self.assertFalse(reqs[0].metrics.suppress_stage_breakdown) - self.assertEqual(reqs[0].num_inference_steps, 2) + self.assertEqual(reqs[0].num_inference_steps, 4) self.assertIn( - "auto residency probe (1280x720x81f, 2/35 steps)", + "auto residency probe (1280x720x81f, 4/35 steps)", format_warmup_req(reqs[0]), ) From 3321551f26135fdc70568bf5081b66dce7aadc2c Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 14:40:52 +0800 Subject: [PATCH 48/90] [diffusion] bound auto residency validation --- .../runtime/entrypoints/control_requests.py | 2 +- .../runtime/managers/gpu_worker.py | 95 +++++++++++- .../memory_managers/auto_residency.py | 9 +- .../runtime/managers/scheduler.py | 2 +- .../multimodal_gen/runtime/server_warmup.py | 131 +++++++++------- .../test/unit/test_auto_residency_warmup.py | 146 ++++++++++++++---- 6 files changed, 289 insertions(+), 96 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py b/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py index 9df19c7f8969..95058ba4429f 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py @@ -57,4 +57,4 @@ class AutoResidencyReq(msgspec.Struct, frozen=True): requests finish and before the server reports ready. """ - action: str = "apply" # "apply" | "rollback" + action: str = "apply" # "apply" | "validate" | "rollback" diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 3d7953fbf2de..8ebd742d1425 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -56,6 +56,7 @@ PLACEMENT_STATUS_ROLLBACK_FAILED, PLACEMENT_STATUS_ROLLED_BACK, PLACEMENT_STATUS_SKIPPED, + PLACEMENT_STATUS_VALIDATED, POST_ADJUSTMENT_REGRESSION_FRACTION, AppliedResidencyChange, AutoResidencyRollbackError, @@ -87,6 +88,7 @@ ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( RESIDENT, + is_dit_component_name, ) from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( HOST_COPY_RESERVE_BYTES, @@ -1232,14 +1234,15 @@ def get_can_stay_resident_components( remaining_gpu_mem_gb -= usage_gb return can_stay_resident - def apply_auto_residency(self) -> OutputBatch: + def apply_auto_residency(self, *, validate_only: bool = False) -> OutputBatch: """Apply one warmup-calibrated residency adjustment round. Every rank executes this handler at the same queue position; the plan is computed from all-gathered rank reports so each rank reaches the - same decision. Post-adjustment warmup can expose more safe headroom, so - the server may call this again with measurements from the new layout. - A failure on any rank rolls every rank back to the previously + same decision. After applying a plan, the server calls this once in + validation-only mode with measurements from the new layout. Validation + may keep or roll back that plan, but never starts another placement + search. A failure on any rank rolls every rank back to the previously calibrated placement. Everything before the first all-gather is fenced into a skip report: @@ -1276,7 +1279,30 @@ def apply_auto_residency(self) -> OutputBatch: skip_reason=describe_error(e), ) reports = self._auto_residency_all_gather(local_report) - if self._auto_residency_round_sizes: + if validate_only: + invalid_report = next( + ( + report + for report in reports + if report.skip_reason is not None + or report.estimated_peak_bytes is None + ), + None, + ) + if invalid_report is not None: + reason = invalid_report.skip_reason or "no usable warmup measurement" + return self._rollback_everywhere( + cause=( + "post-adjustment calibration could not validate rank " + f"{invalid_report.rank}: {reason}" + ), + already_failed=False, + latest_round_only=True, + ) + if ( + self._auto_residency_round_sizes + and not self._latest_auto_residency_round_is_resident_only() + ): regressions = [] for report in reports: reference_ns = report.estimated_request_duration_ns @@ -1302,7 +1328,6 @@ def apply_auto_residency(self) -> OutputBatch: latest_round_only=True, ) plan = plan_auto_residency(reports=reports) - summary = format_plan_summary(plan=plan, workload=workload, records=records) if ( self._auto_residency_round_sizes and plan.current_placement_reserve_shortfall_bytes > 0 @@ -1319,6 +1344,32 @@ def apply_auto_residency(self) -> OutputBatch: already_failed=False, latest_round_only=True, ) + if validate_only: + if self.is_output_rank: + reference_ns = local_report.estimated_request_duration_ns + measured_ns = local_report.measured_request_duration_ns + if self._latest_auto_residency_round_supports_short_validation(): + result = "full-shape one-step memory check passed" + elif reference_ns > 0 and measured_ns > 0: + change = (measured_ns - reference_ns) / reference_ns + result = ( + "calibrated request estimate " + f"{reference_ns / 1e9:.2f}s -> {measured_ns / 1e9:.2f}s " + f"({change:+.1%})" + ) + else: + result = "memory and execution checks passed" + logger.info( + "Auto residency: post-adjustment validation passed (%s); " + "keeping the selected placement.", + result, + ) + return OutputBatch( + output=plan_summary_payload( + plan=plan, status=PLACEMENT_STATUS_VALIDATED + ) + ) + summary = format_plan_summary(plan=plan, workload=workload, records=records) if plan.skip_reason is not None or not plan.changes: if self.is_output_rank: logger.info("%s", summary) @@ -1378,7 +1429,13 @@ def apply_auto_residency(self) -> OutputBatch: logger.info("%s", summary) logger.info("%s", format_applied_changes(plan=plan)) return OutputBatch( - output=plan_summary_payload(plan=plan, status=PLACEMENT_STATUS_ADJUSTED) + output=plan_summary_payload( + plan=plan, + status=PLACEMENT_STATUS_ADJUSTED, + short_validation=( + self._latest_auto_residency_round_supports_short_validation() + ), + ) ) def rollback_auto_residency(self) -> OutputBatch: @@ -1605,6 +1662,30 @@ def _mixed_dtype_residency_components() -> set[str]: return set() return manager.components_with_mixed_use_dtypes() + def _latest_auto_residency_round(self) -> list[AppliedResidencyChange]: + if not self._auto_residency_round_sizes: + return [] + round_size = self._auto_residency_round_sizes[-1] + if round_size <= 0 or round_size > len(self._auto_residency_applied): + return [] + return self._auto_residency_applied[-round_size:] + + def _latest_auto_residency_round_is_resident_only(self) -> bool: + """Whether validation follows only execution-monotonic promotions.""" + changes = self._latest_auto_residency_round() + return all( + adjustment.residency_mode != RESIDENT + and self.server_args.residency_mode(adjustment.component_name) == RESIDENT + for adjustment in changes + ) and bool(changes) + + def _latest_auto_residency_round_supports_short_validation(self) -> bool: + """Whether one full-shape step covers every changed execution path.""" + changes = self._latest_auto_residency_round() + return self._latest_auto_residency_round_is_resident_only() and all( + not is_dit_component_name(change.component_name) for change in changes + ) + def _rollback_applied_residency_changes( self, *, latest_round_only: bool = False ) -> str | None: diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 369768afa5eb..1c5f599dbde4 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -101,6 +101,7 @@ PLACEMENT_STATUS_SKIPPED = "skipped" PLACEMENT_STATUS_ADJUSTED = "adjusted" +PLACEMENT_STATUS_VALIDATED = "validated" PLACEMENT_STATUS_ROLLED_BACK = "rolled_back" PLACEMENT_STATUS_ROLLBACK_FAILED = "rollback_failed" @@ -2568,7 +2569,12 @@ def _format_candidate_summary(candidate: ResidencyTarget) -> str: return f"{candidate.component_name}({target_mode}{details})" -def plan_summary_payload(*, plan: AutoResidencyPlan, status: str) -> dict: +def plan_summary_payload( + *, + plan: AutoResidencyPlan, + status: str, + short_validation: bool = False, +) -> dict: """Minimal decision payload for the warmup orchestrator. The orchestrator only branches on ``status``; the human-readable detail @@ -2577,4 +2583,5 @@ def plan_summary_payload(*, plan: AutoResidencyPlan, status: str) -> dict: return { "status": status, "changed": [candidate.component_name for candidate in plan.changes], + "short_validation": short_validation, } diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 38158348bf4f..d86f76aa7d6c 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -209,7 +209,7 @@ def _handle_auto_residency(self, reqs: List[Any]) -> OutputBatch: return OutputBatch(error="auto residency requires a GPU worker") if req.action == "rollback": return self.worker.rollback_auto_residency() - return self.worker.apply_auto_residency() + return self.worker.apply_auto_residency(validate_only=req.action == "validate") def _handle_set_lora(self, reqs: List[Any]) -> OutputBatch: # TODO: return set status diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index d15fa1141184..58fb943984b8 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -139,6 +139,7 @@ async def maybe_apply_auto_residency( PLACEMENT_STATUS_ADJUSTED, PLACEMENT_STATUS_ROLLBACK_FAILED, PLACEMENT_STATUS_ROLLED_BACK, + PLACEMENT_STATUS_VALIDATED, ) skip_reason = auto_residency_skip_reason(server_args) @@ -156,7 +157,6 @@ async def maybe_apply_auto_residency( # Same fail-open contract as the warmup itself: implicit warmups must # never abort startup, explicit --warmup-resolutions ones must succeed. fail_open = server_args.warmup_resolutions is None - adjusted_any = False async def rollback_and_rewarm(error: Exception) -> None: logger.warning( @@ -172,78 +172,82 @@ async def rollback_and_rewarm(error: Exception) -> None: f"auto residency rollback failed: {rollback.error}" ) from error # Restore warm caches for the previous calibrated placement before - # turning ready, keeping the caller's fail-open contract. + # turning ready. Once placement was mutated, failure to revalidate the + # restored state is unsafe to ignore. await run_async_client_warmup( - server_args, forward, fail_open=fail_open, rewarm=True + server_args, forward, fail_open=False, rewarm=True ) - for _ in range(MAX_AUTO_RESIDENCY_ROUNDS): - try: - response = await forward(AutoResidencyReq(action="apply")) - except Exception as e: - if adjusted_any: - # The worker-side apply is transactional, but an RPC failure - # does not tell us whether its response was lost after commit. - # Do not mistake the last validated round for an unvalidated - # one and roll it back; abort startup rather than serve a - # placement whose calibration state is unknown. - raise RuntimeError( - "auto residency apply failed after a calibrated adjustment" - ) from e - if not fail_open: - raise - logger.warning( - "Auto residency apply request failed; continuing on the original " - "strategy: %s", - e, - ) - return - status = _auto_residency_status(response) - if status == PLACEMENT_STATUS_ROLLBACK_FAILED: - raise RuntimeError(f"auto residency rollback failed: {response.error}") - if response.error is not None: - if status == PLACEMENT_STATUS_ROLLED_BACK: - await run_async_client_warmup( - server_args, forward, fail_open=fail_open, rewarm=True - ) - retained = ( - "the last calibrated placement" - if adjusted_any - else "the original strategy" - ) - logger.warning( - "Auto residency adjustment not applied; continuing on %s: %s", - retained, - response.error, - ) - return - if status != PLACEMENT_STATUS_ADJUSTED: - return - - adjusted_any = True - # This pass physically realizes the selected placement and measures - # phases that overlap under it. The next round can then use that new - # evidence instead of permanently charging components twice. - try: + try: + response = await forward(AutoResidencyReq(action="apply")) + except Exception as e: + if not fail_open: + raise + logger.warning( + "Auto residency apply request failed; continuing on the original " + "strategy: %s", + e, + ) + return + status = _auto_residency_status(response) + if status == PLACEMENT_STATUS_ROLLBACK_FAILED: + raise RuntimeError(f"auto residency rollback failed: {response.error}") + if response.error is not None: + if status == PLACEMENT_STATUS_ROLLED_BACK: await run_async_client_warmup( server_args, forward, fail_open=False, rewarm=True ) - except Exception as e: - await rollback_and_rewarm(e) - return + logger.warning( + "Auto residency adjustment not applied; continuing on the original " + "strategy: %s", + response.error, + ) + return + if status != PLACEMENT_STATUS_ADJUSTED: + return - logger.warning( - "Auto residency stopped after %d adjustment rounds; keeping the last " - "successfully calibrated placement.", - MAX_AUTO_RESIDENCY_ROUNDS, + short_validation = bool( + isinstance(response.output, dict) and response.output.get("short_validation") ) + try: + validation_options = {"step_limit": 1} if short_validation else {} + await run_async_client_warmup( + server_args, + forward, + fail_open=False, + rewarm=True, + **validation_options, + ) + except Exception as e: + await rollback_and_rewarm(e) + return + + try: + validation = await forward(AutoResidencyReq(action="validate")) + except Exception as e: + await rollback_and_rewarm(e) + return + validation_status = _auto_residency_status(validation) + if validation_status == PLACEMENT_STATUS_ROLLBACK_FAILED: + raise RuntimeError(f"auto residency rollback failed: {validation.error}") + if validation_status == PLACEMENT_STATUS_ROLLED_BACK: + await run_async_client_warmup( + server_args, forward, fail_open=False, rewarm=True + ) + return + if validation.error is not None or validation_status != PLACEMENT_STATUS_VALIDATED: + await rollback_and_rewarm( + RuntimeError( + validation.error + or "post-adjustment calibration returned no validation result" + ) + ) # Enough to clear a probe that overshot the card, few enough that a failure # which is not about probe size gives up quickly instead of walking the # workload down to nothing. MAX_WARMUP_DEGRADE_ATTEMPTS = 3 -MAX_AUTO_RESIDENCY_ROUNDS = 8 def _is_out_of_memory(error: Any) -> bool: @@ -304,6 +308,7 @@ def build_client_warmup_reqs( *, warmup_input_path: str | None = None, rewarm: bool = False, + step_limit: int | None = None, ) -> list[Req]: warmup_reqs = build_warmup_reqs( server_args, @@ -316,6 +321,8 @@ def build_client_warmup_reqs( for req in warmup_reqs: if req.is_warmup: req.extra["warmup_total"] = warmup_total + if step_limit is not None: + req.num_inference_steps = min(req.num_inference_steps, step_limit) if rewarm: # a repeat pass after an auto-residency change: keep it out of # the scheduler's warmup progress accounting (already at N/N) @@ -329,6 +336,7 @@ async def run_async_client_warmup( *, fail_open: bool = False, rewarm: bool = False, + step_limit: int | None = None, ) -> None: try: warmup_input_path = None @@ -336,7 +344,10 @@ async def run_async_client_warmup( warmup_input_path = prepare_warmup_image_path(server_args) for req in build_client_warmup_reqs( - server_args, warmup_input_path=warmup_input_path, rewarm=rewarm + server_args, + warmup_input_path=warmup_input_path, + rewarm=rewarm, + step_limit=step_limit, ): response = await forward(req) for _ in range(MAX_WARMUP_DEGRADE_ATTEMPTS): diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py index 934362204e50..9f4ee57ef7ee 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -14,24 +14,41 @@ PLACEMENT_STATUS_ADJUSTED, PLACEMENT_STATUS_ROLLBACK_FAILED, PLACEMENT_STATUS_ROLLED_BACK, - PLACEMENT_STATUS_SKIPPED, + PLACEMENT_STATUS_VALIDATED, AppliedResidencyChange, AutoResidencyPlan, DefaultWorkload, RankResidencyReport, + ResidencyTarget, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( COMPONENT_OFFLOAD, + RESIDENT, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch class TestAutoResidencyWarmup(unittest.TestCase): + def test_rewarm_step_limit_keeps_full_shape_but_shortens_denoising(self): + req = SimpleNamespace(is_warmup=True, extra={}, num_inference_steps=4) + with mock.patch.object(server_warmup, "build_warmup_reqs", return_value=[req]): + result = server_warmup.build_client_warmup_reqs( + SimpleNamespace(warmup_resolutions=None), + rewarm=True, + step_limit=1, + ) + + self.assertEqual(result, [req]) + self.assertEqual(req.num_inference_steps, 1) + self.assertTrue(req.extra["server_warmup_rewarm"]) + def test_worker_rolls_back_a_materially_slower_calibrated_round(self): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 worker.is_output_rank = False - worker.server_args = SimpleNamespace() + worker.server_args = SimpleNamespace( + residency_mode=lambda _name: COMPONENT_OFFLOAD + ) worker._auto_residency_warmup_records = [object()] worker._auto_residency_round_sizes = [1] report = RankResidencyReport( @@ -49,16 +66,19 @@ def test_worker_rolls_back_a_materially_slower_calibrated_round(self): ) worker._rollback_everywhere = mock.Mock(return_value=rolled_back) - with mock.patch.object( - gpu_worker_module, - "resolve_default_workload", - return_value=SimpleNamespace(), - ), mock.patch.object( - gpu_worker_module, - "resolve_measured_default_workload", - return_value=SimpleNamespace(), + with ( + mock.patch.object( + gpu_worker_module, + "resolve_default_workload", + return_value=SimpleNamespace(), + ), + mock.patch.object( + gpu_worker_module, + "resolve_measured_default_workload", + return_value=SimpleNamespace(), + ), ): - response = worker.apply_auto_residency() + response = worker.apply_auto_residency(validate_only=True) self.assertIs(response, rolled_back) worker._rollback_everywhere.assert_called_once() @@ -72,7 +92,9 @@ def test_new_placement_that_spends_the_reserve_rolls_back(self): worker.rank = 0 worker.is_output_rank = False worker.pipeline = SimpleNamespace(modules={}) - worker.server_args = SimpleNamespace() + worker.server_args = SimpleNamespace( + residency_mode=lambda _name: COMPONENT_OFFLOAD + ) worker._auto_residency_warmup_records = [] worker._auto_residency_round_sizes = [1] worker._build_auto_residency_report = mock.Mock( @@ -112,6 +134,71 @@ def test_new_placement_that_spends_the_reserve_rolls_back(self): latest_round_only=True, ) + def test_worker_validation_does_not_apply_a_second_candidate_plan(self): + worker = GPUWorker.__new__(GPUWorker) + worker.rank = 0 + worker.is_output_rank = False + worker.server_args = SimpleNamespace( + residency_mode=lambda _name: COMPONENT_OFFLOAD + ) + worker._auto_residency_warmup_records = [object()] + worker._auto_residency_round_sizes = [1] + report = RankResidencyReport( + rank=0, + budget_bytes=1, + estimated_peak_bytes=1, + estimated_request_duration_ns=10_000_000_000, + measured_request_duration_ns=10_000_000_000, + ) + worker._build_auto_residency_report = mock.Mock(return_value=report) + worker._auto_residency_all_gather = mock.Mock(return_value=[report]) + candidate = ResidencyTarget( + component_name="vae", + residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=1, + h2d_bytes_per_request=1, + ) + + with ( + mock.patch.object( + gpu_worker_module, + "resolve_default_workload", + return_value=SimpleNamespace(), + ), + mock.patch.object( + gpu_worker_module, + "resolve_measured_default_workload", + return_value=SimpleNamespace(), + ), + mock.patch.object( + gpu_worker_module, + "plan_auto_residency", + return_value=AutoResidencyPlan(changes=[candidate]), + ), + mock.patch.object(gpu_worker_module, "apply_residency_changes") as apply, + ): + response = worker.apply_auto_residency(validate_only=True) + + self.assertEqual(response.output["status"], PLACEMENT_STATUS_VALIDATED) + apply.assert_not_called() + + def test_aux_resident_promotion_uses_short_validation_but_dit_does_not(self): + worker = GPUWorker.__new__(GPUWorker) + worker.server_args = SimpleNamespace(residency_mode=lambda _name: RESIDENT) + worker._auto_residency_round_sizes = [1] + + worker._auto_residency_applied = [ + AppliedResidencyChange("text_encoder", COMPONENT_OFFLOAD) + ] + self.assertTrue(worker._latest_auto_residency_round_supports_short_validation()) + + worker._auto_residency_applied = [ + AppliedResidencyChange("transformer", COMPONENT_OFFLOAD) + ] + self.assertFalse( + worker._latest_auto_residency_round_supports_short_validation() + ) + def test_worker_rollback_response_rewarms_the_restored_layout(self): actions = [] @@ -140,16 +227,20 @@ async def forward(req): rewarm.assert_awaited_once_with( server_args, forward, - fail_open=True, + fail_open=False, rewarm=True, ) - def test_replans_until_the_calibrated_layout_reaches_a_fixed_point(self): + def test_successful_adjustment_runs_one_validation_without_replanning(self): responses = iter( [ - OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), - OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), - OutputBatch(output={"status": PLACEMENT_STATUS_SKIPPED}), + OutputBatch( + output={ + "status": PLACEMENT_STATUS_ADJUSTED, + "short_validation": True, + } + ), + OutputBatch(output={"status": PLACEMENT_STATUS_VALIDATED}), ] ) actions = [] @@ -172,8 +263,14 @@ async def forward(req): ): asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) - self.assertEqual(actions, ["apply", "apply", "apply"]) - self.assertEqual(rewarm.await_count, 2) + self.assertEqual(actions, ["apply", "validate"]) + rewarm.assert_awaited_once_with( + server_args, + forward, + fail_open=False, + rewarm=True, + step_limit=1, + ) def test_failed_first_calibration_rolls_back_and_rewarms(self): actions = [] @@ -228,12 +325,13 @@ async def forward(req): ): asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) - def test_apply_rpc_failure_after_calibration_aborts_without_rollback(self): + def test_validation_rpc_failure_rolls_back_and_rewarms(self): actions = [] responses = iter( [ OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), RuntimeError("rpc failed"), + OutputBatch(output={"status": PLACEMENT_STATUS_ROLLED_BACK}), ] ) @@ -255,15 +353,11 @@ async def forward(req): server_warmup, "auto_residency_skip_reason", return_value=None ), mock.patch.object(server_warmup, "run_async_client_warmup", rewarm), - self.assertRaisesRegex( - RuntimeError, - "auto residency apply failed after a calibrated adjustment", - ), ): asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) - self.assertEqual(actions, ["apply", "apply"]) - self.assertEqual(rewarm.await_count, 1) + self.assertEqual(actions, ["apply", "validate", "rollback"]) + self.assertEqual(rewarm.await_count, 2) def test_failed_later_round_keeps_earlier_calibrated_promotions(self): first_round = AppliedResidencyChange("transformer", COMPONENT_OFFLOAD) From 22edea89a8c4bb0997b4783e44bf5c47f8d8f0f6 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 14:44:59 +0800 Subject: [PATCH 49/90] [diffusion] commit validated resident placement --- .../runtime/managers/gpu_worker.py | 23 +++++- .../memory_managers/auto_residency.py | 34 +++++++- .../memory_managers/layerwise_offload.py | 78 +++++++++++++++---- .../test/unit/test_auto_residency.py | 48 ++++++++++++ .../test/unit/test_auto_residency_warmup.py | 8 +- 5 files changed, 170 insertions(+), 21 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 8ebd742d1425..ca4274d148ce 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -65,6 +65,7 @@ WarmupMemoryRecord, apply_residency_changes, collect_residency_targets, + commit_residency_changes, component_current_device_weight_bytes, component_runtime_weight_bytes, describe_error, @@ -1345,10 +1346,30 @@ def apply_auto_residency(self, *, validate_only: bool = False) -> OutputBatch: latest_round_only=True, ) if validate_only: + short_validation = ( + self._latest_auto_residency_round_supports_short_validation() + ) + commit_error = None + try: + commit_residency_changes( + applied=self._latest_auto_residency_round(), + modules=self.pipeline.modules, + server_args=self.server_args, + ) + except Exception as e: + commit_error = describe_error(e) + commit_errors = self._auto_residency_all_gather(commit_error) + if any(error is not None for error in commit_errors): + raise RuntimeError( + "post-validation residency commit failed: " + + next(error for error in commit_errors if error is not None) + ) + self._auto_residency_applied = [] + self._auto_residency_round_sizes = [] if self.is_output_rank: reference_ns = local_report.estimated_request_duration_ns measured_ns = local_report.measured_request_duration_ns - if self._latest_auto_residency_round_supports_short_validation(): + if short_validation: result = "full-shape one-step memory check passed" elif reference_ns > 0 and measured_ns > 0: change = (measured_ns - reference_ns) / reference_ns diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 1c5f599dbde4..ff255171be0a 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -2118,6 +2118,9 @@ def apply_residency_changes( target_modules: dict[str, nn.Module] = {} snapshots: dict[str, AppliedResidencyChange] = {} previous_pins_by_component: dict[str, tuple[tuple[int, ...], ...]] = {} + resident_only = bool(ordered_changes) and all( + candidate.target_mode() == RESIDENT for candidate in ordered_changes + ) for candidate in ordered_changes: if candidate.component_name in target_modules: raise RuntimeError( @@ -2212,6 +2215,7 @@ def apply_residency_changes( candidate for candidate in ordered_changes if snapshots[candidate.component_name].previous_layerwise_configured + if not (resident_only and candidate.target_mode() == RESIDENT) if candidate.target_layerwise_pinned_layers is not None and candidate.target_layerwise_pinned_layers != previous_pins_by_component[candidate.component_name] @@ -2288,9 +2292,10 @@ def apply_residency_changes( f"residency target {candidate.component_name!r} " "did not configure any layerwise groups" ) - module.set_layerwise_pinned_layers( - candidate.target_layerwise_pinned_layers or () - ) + if not (resident_only and target_mode == RESIDENT): + module.set_layerwise_pinned_layers( + candidate.target_layerwise_pinned_layers or () + ) if target_mode == COMPONENT_OFFLOAD: module.remove_layerwise_offload(to_cpu=True) server_args.set_auto_residency_mode( @@ -2301,6 +2306,7 @@ def apply_residency_changes( candidate.component_name, RESIDENT ) if is_layerwise_offloaded_module(module): + module.restore_non_layer_weights() module.disable_offload() else: if not is_layerwise_offloaded_module(module): @@ -2343,6 +2349,28 @@ def apply_residency_changes( return applied +def commit_residency_changes( + *, + applied: Iterable[AppliedResidencyChange], + modules: Mapping[str, object], + server_args: ServerArgs, +) -> None: + """Release rollback-only stores after post-adjustment validation succeeds.""" + for adjustment in applied: + if ( + not adjustment.previous_layerwise_configured + or server_args.residency_mode(adjustment.component_name) != RESIDENT + ): + continue + module = modules.get(adjustment.component_name) + if not isinstance(module, LayerwiseOffloadableModuleMixin): + raise RuntimeError( + f"residency target {adjustment.component_name!r} lost its " + "layerwise offload capability before commit" + ) + module.finalize_resident_layerwise_placement() + + def rollback_residency_changes( *, applied: Iterable[AppliedResidencyChange], 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 7e6b7ebeaeb6..88149e61b136 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 @@ -1263,7 +1263,35 @@ def load_all_layers(self) -> None: for layer_idx in range(self.num_layers): if layer_idx not in self._gpu_layers: - self.prefetch_layer(layer_idx, non_blocking=False) + self.prefetch_layer( + layer_idx, + non_blocking=not bool(self._mapped_cpu_weights.get(layer_idx)), + ) + if self.copy_stream is not None: + torch.get_device_module().current_stream().wait_stream(self.copy_stream) + + def release_host_stores(self) -> None: + """Drop rollback stores after a resident placement is validated.""" + if self.enabled: + raise RuntimeError("cannot release host stores while offload is enabled") + if self._mapped_courier is not None: + self._mapped_courier.close() + self._mapped_courier = None + if self._courier_inflight: + raise RuntimeError( + "cannot release host stores with mapped copies in flight" + ) + + self._pin_budget.release(self.pinned_host_weight_bytes()) + self._consolidated_cpu_weights.clear() + self._strided_cpu_weights.clear() + self._mapped_cpu_weights.clear() + self._mps_cpu_weights.clear() + self._weight_metadata.clear() + self._layer_hosting.clear() + self._prefetch_events.clear() + self._mapped_bytes = 0 + self._configured = False @torch.compiler.disable def sync_layer_to_cpu(self, layer_idx: int) -> None: @@ -2225,22 +2253,26 @@ def disable_offload(self) -> None: """ if self.layerwise_offload_managers is None: return - for manager in self.layerwise_offload_managers: - if not manager.enabled: - continue - manager.remove_forward_hooks() - try: - manager.load_all_layers() - except Exception: - # Loading every layer is the step most likely to OOM. Re-arm - # this manager before re-raising: leaving it hook-less with - # enabled=True would let release_all() swap weights for (1,) - # placeholders that nothing ever streams back in, and - # enable_offload()'s already-armed guard would skip repair. + disabled: list[LayerwiseOffloadManager] = [] + try: + for manager in self.layerwise_offload_managers: + if not manager.enabled: + continue + manager.remove_forward_hooks() + try: + manager.load_all_layers() + except Exception: + manager.release_all() + manager.register_forward_hooks() + raise + manager.enabled = False + disabled.append(manager) + except Exception: + for manager in reversed(disabled): + manager.enabled = True manager.release_all() manager.register_forward_hooks() - raise - manager.enabled = False + raise def enable_offload(self) -> None: """Re-enable layerwise offload: sync weights to CPU, release layers, and restore hooks.""" @@ -2308,6 +2340,22 @@ def remove_layerwise_offload(self, *, to_cpu: bool = False) -> None: self._parked_non_layer_weights.clear() self._park_placeholders.clear() + def finalize_resident_layerwise_placement(self) -> None: + """Discard offload stores once a resident placement has passed validation.""" + managers = list(self.layerwise_offload_managers or ()) + if not managers: + return + if any(manager.enabled for manager in managers): + raise RuntimeError( + "cannot finalize a resident placement while offload is enabled" + ) + for manager in managers: + manager.release_host_stores() + self.layerwise_offload_managers = [] + self._parked_non_layer_weights.clear() + self._park_placeholders.clear() + release_unused_pinned_memory() + def iter_materialized_weights(module: torch.nn.Module): """Yield (name, tensor) pairs with materialized weights, even under offload. diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index d40f7cf248fb..416e3a330649 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -23,6 +23,7 @@ _layerwise_resident_targets, apply_residency_changes, collect_residency_targets, + commit_residency_changes, component_resident_size_bytes, component_runtime_weight_bytes, estimate_candidate_latency_savings_ns, @@ -1624,6 +1625,7 @@ def __init__( self.register_hooks_calls = 0 self.sync_to_cpu_calls = 0 self.release_all_calls = 0 + self.release_host_stores_calls = 0 def iter_cpu_weights(self): yield from self._tensors.items() @@ -1705,6 +1707,12 @@ def sync_all_layers_to_cpu(self): def release_all(self): self.release_all_calls += 1 + def release_host_stores(self): + if self.enabled: + raise RuntimeError("offload is still enabled") + self.release_host_stores_calls += 1 + self._pinned_layers = () + class _FakeLayerwiseDit(LayerwiseOffloadableModuleMixin, nn.Module): def __init__(self, managers: list[_FakeLayerwiseManager]): @@ -2301,6 +2309,30 @@ def test_layerwise_promotion_loads_all_layers_and_rollback_rearms(self): assert manager.register_hooks_calls == 1 assert args.auto_modes == {} + def test_resident_only_promotion_keeps_pins_until_validation_commit(self): + manager = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) + manager._pinned_layers = (0,) + module = _FakeLayerwiseDit([manager]) + args = _StubResidencyArgs() + candidate = _candidate("text_encoder", mode=LAYERWISE_OFFLOAD, weight_gib=1) + + applied = apply_residency_changes( + plan=_plan_for([candidate]), + modules={"text_encoder": module}, + server_args=args, + ) + assert manager.enabled is False + assert manager.pinned_layer_indices() == (0,) + + commit_residency_changes( + applied=applied, + modules={"text_encoder": module}, + server_args=args, + ) + + assert manager.release_host_stores_calls == 1 + assert module.layerwise_offload_managers == [] + def test_partial_layerwise_promotion_restores_exact_group_counts(self): managers = [ _FakeLayerwiseManager( @@ -2617,6 +2649,22 @@ def test_disable_offload_failure_rearms_the_manager(self): assert manager.register_hooks_calls == 1 assert manager.release_all_calls == 1 + def test_disable_offload_failure_rearms_earlier_layer_groups(self): + first = _FakeLayerwiseManager({"layers.0.w": torch.zeros(16)}) + failing = _FakeLayerwiseManager({"layers.1.w": torch.zeros(16)}) + failing.fail_load = True + module = _FakeLayerwiseDit([first, failing]) + + with pytest.raises(RuntimeError, match="out of memory"): + module.disable_offload() + + assert first.enabled is True + assert first.register_hooks_calls == 1 + assert first.release_all_calls == 1 + assert failing.enabled is True + assert failing.register_hooks_calls == 1 + assert failing.release_all_calls == 1 + def test_failed_rollback_raises_rollback_error_with_visible_cause(self): class _BrokenArgs(_StubResidencyArgs): def clear_auto_residency_mode(self, component_name): diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py index 9f4ee57ef7ee..5249f4824e49 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -46,10 +46,12 @@ def test_worker_rolls_back_a_materially_slower_calibrated_round(self): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 worker.is_output_rank = False + worker.pipeline = SimpleNamespace(modules={}) worker.server_args = SimpleNamespace( residency_mode=lambda _name: COMPONENT_OFFLOAD ) worker._auto_residency_warmup_records = [object()] + worker._auto_residency_applied = [] worker._auto_residency_round_sizes = [1] report = RankResidencyReport( rank=0, @@ -59,7 +61,7 @@ def test_worker_rolls_back_a_materially_slower_calibrated_round(self): measured_request_duration_ns=20_000_000_000, ) worker._build_auto_residency_report = mock.Mock(return_value=report) - worker._auto_residency_all_gather = mock.Mock(return_value=[report]) + worker._auto_residency_all_gather = mock.Mock(side_effect=lambda value: [value]) rolled_back = OutputBatch( error="request duration regressed", output={"status": PLACEMENT_STATUS_ROLLED_BACK}, @@ -138,10 +140,12 @@ def test_worker_validation_does_not_apply_a_second_candidate_plan(self): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 worker.is_output_rank = False + worker.pipeline = SimpleNamespace(modules={}) worker.server_args = SimpleNamespace( residency_mode=lambda _name: COMPONENT_OFFLOAD ) worker._auto_residency_warmup_records = [object()] + worker._auto_residency_applied = [] worker._auto_residency_round_sizes = [1] report = RankResidencyReport( rank=0, @@ -151,7 +155,7 @@ def test_worker_validation_does_not_apply_a_second_candidate_plan(self): measured_request_duration_ns=10_000_000_000, ) worker._build_auto_residency_report = mock.Mock(return_value=report) - worker._auto_residency_all_gather = mock.Mock(return_value=[report]) + worker._auto_residency_all_gather = mock.Mock(side_effect=lambda value: [value]) candidate = ResidencyTarget( component_name="vae", residency_mode=COMPONENT_OFFLOAD, From fee9e405f3556e6149aade8f9df18d4fa4a9f408 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 14:50:54 +0800 Subject: [PATCH 50/90] [diffusion] validate residency commit barrier --- .../multimodal_gen/runtime/managers/gpu_worker.py | 13 +++++++++++-- .../managers/memory_managers/auto_residency.py | 2 +- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index ca4274d148ce..d4643a218993 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1346,14 +1346,23 @@ def apply_auto_residency(self, *, validate_only: bool = False) -> OutputBatch: latest_round_only=True, ) if validate_only: + latest_round = self._latest_auto_residency_round() short_validation = ( self._latest_auto_residency_round_supports_short_validation() ) + commit_modules = ( + self.pipeline.modules + if any( + change.previous_layerwise_resident_layers is not None + for change in latest_round + ) + else {} + ) commit_error = None try: commit_residency_changes( - applied=self._latest_auto_residency_round(), - modules=self.pipeline.modules, + applied=latest_round, + modules=commit_modules, server_args=self.server_args, ) except Exception as e: diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index ff255171be0a..ea18ce516630 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -2358,7 +2358,7 @@ def commit_residency_changes( """Release rollback-only stores after post-adjustment validation succeeds.""" for adjustment in applied: if ( - not adjustment.previous_layerwise_configured + adjustment.previous_layerwise_resident_layers is None or server_args.residency_mode(adjustment.component_name) != RESIDENT ): continue From 5f86c0a45b26c952e22728402158ae68e6f23ceb Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 14:54:07 +0800 Subject: [PATCH 51/90] [docs] define auto residency solve boundary --- docs/docs/sglang-diffusion/api/cli.mdx | 2 +- docs/docs/sglang-diffusion/deployment_cookbook.mdx | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index c21c08b7e8c0..ace17a9211bb 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -85,7 +85,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. Multi-process HostPin planning stays inside non-overlapping per-worker shares. The selected placement targets the model-default (or `--warmup-sampling-params`) workload with batch size 1; use explicit residency for larger production shapes or batches. 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 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. Multi-process HostPin planning stays inside non-overlapping per-worker shares. 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. 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. - `--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. Startup impact is model-dependent, so benchmark the target model before deployment. Disabled by default because checkpoint and model weights coexist temporarily, substantially increasing peak GPU memory. 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 diff --git a/docs/docs/sglang-diffusion/deployment_cookbook.mdx b/docs/docs/sglang-diffusion/deployment_cookbook.mdx index e3384126126c..4715180a9569 100644 --- a/docs/docs/sglang-diffusion/deployment_cookbook.mdx +++ b/docs/docs/sglang-diffusion/deployment_cookbook.mdx @@ -128,7 +128,7 @@ See [OpenAI API: Served model name](/docs/sglang-diffusion/api/openai_api#served -With an eligible `sglang serve` warmup, `auto` records allocated-memory peaks for the representative request and their active component layouts. One placement decision then satisfies every rank and execution-phase VRAM constraint together with HostPin and HostPin-transition limits. The target frontier covers component offload, resident components, layerwise resident counts, and pinned layer-store reassignment. A post-adjustment warmup realizes the selected state and validates its memory peak before the server becomes ready. Later rounds may replace an earlier automatic choice rather than only increasing residency. +With an eligible `sglang serve` warmup, `auto` records allocated-memory peaks for the representative request and their active component layouts. 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, layerwise resident counts, and pinned layer-store reassignment. A post-adjustment warmup realizes the selected state and validates its memory peak before the server becomes ready. Auxiliary-only moves to resident use one full-shape denoising step; DiT changes, 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. From 98f2e2d8bb4fab8a90881bbfe2516f9b1b137bb1 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 15:07:23 +0800 Subject: [PATCH 52/90] [diffusion] validate residency without replanning --- .../runtime/managers/gpu_worker.py | 151 ++++++++++++------ .../memory_managers/auto_residency.py | 48 ++++-- .../test/unit/test_auto_residency.py | 11 ++ .../test/unit/test_auto_residency_warmup.py | 13 +- 4 files changed, 158 insertions(+), 65 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index d4643a218993..51e4f74b9fdf 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -59,6 +59,7 @@ PLACEMENT_STATUS_VALIDATED, POST_ADJUSTMENT_REGRESSION_FRACTION, AppliedResidencyChange, + AutoResidencyPlan, AutoResidencyRollbackError, DefaultWorkload, RankResidencyReport, @@ -68,6 +69,7 @@ commit_residency_changes, component_current_device_weight_bytes, component_runtime_weight_bytes, + current_placement_reserve_shortfall_bytes, describe_error, estimate_candidate_latency_savings_ns, estimate_default_workload_peak_bytes, @@ -210,6 +212,7 @@ def __init__( self._auto_residency_warmup_records: list[WarmupMemoryRecord] = [] self._auto_residency_applied: list[AppliedResidencyChange] = [] self._auto_residency_round_sizes: list[int] = [] + self._auto_residency_last_applied_plan: AutoResidencyPlan | None = None # Keep every round on one serving objective. Re-warm measurements # validate and refine memory constraints, but placement-dependent # latency must not rewrite the utility function and cause oscillation. @@ -1260,7 +1263,9 @@ def apply_auto_residency(self, *, validate_only: bool = False) -> OutputBatch: workload = resolve_default_workload(self.server_args) workload = resolve_measured_default_workload(workload, records) local_report = self._build_auto_residency_report( - workload=workload, records=records + workload=workload, + records=records, + include_candidates=not validate_only, ) except Exception as e: logger.warning( @@ -1328,24 +1333,28 @@ def apply_auto_residency(self, *, validate_only: bool = False) -> OutputBatch: already_failed=False, latest_round_only=True, ) - plan = plan_auto_residency(reports=reports) - if ( - self._auto_residency_round_sizes - and plan.current_placement_reserve_shortfall_bytes > 0 - ): - shortfall_gib = plan.current_placement_reserve_shortfall_bytes / GIB_BYTES - if self.is_output_rank: - logger.warning( - "Auto residency calibration exceeded the VRAM reserve by " - "%.1f GiB; rolling back the latest adjustment round.", - shortfall_gib, - ) - return self._rollback_everywhere( - cause=f"VRAM reserve exceeded by {shortfall_gib:.1f} GiB", - already_failed=False, - latest_round_only=True, - ) if validate_only: + shortfall_bytes = current_placement_reserve_shortfall_bytes(reports) + if shortfall_bytes > 0: + shortfall_gib = shortfall_bytes / GIB_BYTES + if self.is_output_rank: + logger.warning( + "Auto residency calibration exceeded the VRAM reserve by " + "%.1f GiB; rolling back the latest adjustment round.", + shortfall_gib, + ) + return self._rollback_everywhere( + cause=f"VRAM reserve exceeded by {shortfall_gib:.1f} GiB", + already_failed=False, + latest_round_only=True, + ) + plan = self._auto_residency_last_applied_plan + if plan is None: + return self._rollback_everywhere( + cause="post-adjustment validation lost the selected plan", + already_failed=False, + latest_round_only=True, + ) latest_round = self._latest_auto_residency_round() short_validation = ( self._latest_auto_residency_round_supports_short_validation() @@ -1375,6 +1384,7 @@ def apply_auto_residency(self, *, validate_only: bool = False) -> OutputBatch: ) self._auto_residency_applied = [] self._auto_residency_round_sizes = [] + self._auto_residency_last_applied_plan = None if self.is_output_rank: reference_ns = local_report.estimated_request_duration_ns measured_ns = local_report.measured_request_duration_ns @@ -1399,6 +1409,23 @@ def apply_auto_residency(self, *, validate_only: bool = False) -> OutputBatch: plan=plan, status=PLACEMENT_STATUS_VALIDATED ) ) + plan = plan_auto_residency(reports=reports) + if ( + self._auto_residency_round_sizes + and plan.current_placement_reserve_shortfall_bytes > 0 + ): + shortfall_gib = plan.current_placement_reserve_shortfall_bytes / GIB_BYTES + if self.is_output_rank: + logger.warning( + "Auto residency calibration exceeded the VRAM reserve by " + "%.1f GiB; rolling back the latest adjustment round.", + shortfall_gib, + ) + return self._rollback_everywhere( + cause=f"VRAM reserve exceeded by {shortfall_gib:.1f} GiB", + already_failed=False, + latest_round_only=True, + ) summary = format_plan_summary(plan=plan, workload=workload, records=records) if plan.skip_reason is not None or not plan.changes: if self.is_output_rank: @@ -1455,6 +1482,7 @@ def apply_auto_residency(self, *, validate_only: bool = False) -> OutputBatch: self._invalidate_component_strategies( [candidate.component_name for candidate in plan.changes] ) + self._auto_residency_last_applied_plan = plan if self.is_output_rank: logger.info("%s", summary) logger.info("%s", format_applied_changes(plan=plan)) @@ -1533,7 +1561,11 @@ def _rollback_everywhere( return OutputBatch(output={"status": PLACEMENT_STATUS_ROLLED_BACK}) def _build_auto_residency_report( - self, *, workload: DefaultWorkload, records: List[WarmupMemoryRecord] + self, + *, + workload: DefaultWorkload, + records: List[WarmupMemoryRecord], + include_candidates: bool = True, ) -> RankResidencyReport: skip_reason = None if self.pipeline is None: @@ -1597,25 +1629,27 @@ def _build_auto_residency_report( test_cap_gib = envs.SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB if test_cap_gib is not None: budget_bytes = min(budget_bytes, int(test_cap_gib * GIB_BYTES)) - candidates = collect_residency_targets( - modules=self.pipeline.modules, - residency_mode_of=self.server_args.residency_mode, - baseline_residency_mode_of=(self.server_args.configured_residency_mode), - explicit_residency_mode_of=self.server_args.explicit_residency_mode, - custom_strategy_names=self.pipeline.component_residency_strategies, - num_inference_steps=workload.num_inference_steps, - allow_host_pin_reallocation=True, - mixed_dtype_components=self._mixed_dtype_residency_components(), - auto_resident_components={ - name - for name in self.pipeline.modules - if self.server_args.auto_residency_mode(name) == RESIDENT - }, - layerwise_tuning_of=lambda name, dit_group: ( - self.server_args.layerwise_tuning_for(name, dit_group=dit_group) - ), - pin_cpu_memory=self.server_args.pin_cpu_memory, - ) + candidates = [] + if include_candidates: + candidates = collect_residency_targets( + modules=self.pipeline.modules, + residency_mode_of=self.server_args.residency_mode, + baseline_residency_mode_of=(self.server_args.configured_residency_mode), + explicit_residency_mode_of=self.server_args.explicit_residency_mode, + custom_strategy_names=self.pipeline.component_residency_strategies, + num_inference_steps=workload.num_inference_steps, + allow_host_pin_reallocation=True, + mixed_dtype_components=self._mixed_dtype_residency_components(), + auto_resident_components={ + name + for name in self.pipeline.modules + if self.server_args.auto_residency_mode(name) == RESIDENT + }, + layerwise_tuning_of=lambda name, dit_group: ( + self.server_args.layerwise_tuning_for(name, dit_group=dit_group) + ), + pin_cpu_memory=self.server_args.pin_cpu_memory, + ) measured_request_duration_ns, _, _ = estimate_default_workload_timing( records=records, target_units=target_units, @@ -1632,9 +1666,13 @@ def _build_auto_residency_report( estimated_request_duration_ns = ( reference_request_duration_ns or measured_request_duration_ns ) - candidate_latency_savings_ns = estimate_candidate_latency_savings_ns( - candidates=candidates, - request_duration_ns=estimated_request_duration_ns, + candidate_latency_savings_ns = ( + estimate_candidate_latency_savings_ns( + candidates=candidates, + request_duration_ns=estimated_request_duration_ns, + ) + if include_candidates + else {} ) local_worker_count = max( 1, self.server_args.num_gpus // self.server_args.nnodes @@ -1655,18 +1693,30 @@ def _build_auto_residency_report( ), current_device_weight_bytes_by_component=( component_current_device_weight_bytes(self.pipeline.modules) + if include_candidates + else {} ), node_rank=self.server_args.node_rank, - pinned_host_bytes=layerwise_pinned_host_bytes(self.pipeline.modules), - host_pin_capacity_bytes=layerwise_host_pin_capacity_bytes( - self.pipeline.modules + pinned_host_bytes=( + layerwise_pinned_host_bytes(self.pipeline.modules) + if include_candidates + else 0 ), - host_transition_headroom_bytes=max( - 0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES - ) - // local_worker_count, - device_transition_allocated_bytes=int( - torch.get_device_module().memory_allocated() + host_pin_capacity_bytes=( + layerwise_host_pin_capacity_bytes(self.pipeline.modules) + if include_candidates + else 0 + ), + host_transition_headroom_bytes=( + max(0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES) + // local_worker_count + if include_candidates + else 0 + ), + device_transition_allocated_bytes=( + int(torch.get_device_module().memory_allocated()) + if include_candidates + else 0 ), estimated_request_duration_ns=estimated_request_duration_ns, measured_request_duration_ns=measured_request_duration_ns, @@ -1757,6 +1807,7 @@ def _rollback_applied_residency_changes( else: self._auto_residency_applied = [] self._auto_residency_round_sizes = [] + self._auto_residency_last_applied_plan = None return None @staticmethod diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index ea18ce516630..956e710d5a06 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -1535,6 +1535,39 @@ def _vram_reserve_bytes(budget_bytes: int, *, target_workload_measured: bool) -> ) +def current_placement_reserve_shortfall_bytes( + reports: Iterable[RankResidencyReport], +) -> int: + """Return the largest reserve deficit of the measured placement. + + This checks the placement that just executed without regenerating a + candidate frontier or running the optimizer. It is therefore suitable for + the single post-adjustment validation pass. + """ + reports = list(reports) + if not reports: + return 0 + target_workload_measured = any( + report.target_workload_measured for report in reports + ) + shortfalls = [] + for report in reports: + if report.estimated_peak_bytes is None: + continue + reserve = _vram_reserve_bytes( + report.budget_bytes, + target_workload_measured=target_workload_measured, + ) + measured_peak = max( + [ + report.estimated_peak_bytes, + *report.estimated_peak_bytes_by_phase.values(), + ] + ) + shortfalls.append(measured_peak + reserve - report.budget_bytes) + return max(0, max(shortfalls, default=0)) + + def _binding_phase_constraints( report: RankResidencyReport, candidate_component_names: set[str], @@ -1785,20 +1818,9 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP for report in reports } reserve = max(reserves_by_rank.values()) - current_placement_reserve_shortfall = max( - ( - max( - [ - report.estimated_peak_bytes or 0, - *report.estimated_peak_bytes_by_phase.values(), - ] - ) - + reserves_by_rank[report.rank] - - report.budget_bytes - ) - for report in reports + current_placement_reserve_shortfall = current_placement_reserve_shortfall_bytes( + reports ) - current_placement_reserve_shortfall = max(0, current_placement_reserve_shortfall) candidates = _consensus_candidates(reports) if not candidates: diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 416e3a330649..11e278661ea5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -26,6 +26,7 @@ commit_residency_changes, component_resident_size_bytes, component_runtime_weight_bytes, + current_placement_reserve_shortfall_bytes, estimate_candidate_latency_savings_ns, estimate_default_workload_peak_bytes, estimate_default_workload_timing, @@ -1162,6 +1163,16 @@ def test_reports_reserve_shortfall_for_a_newly_calibrated_placement(self): assert plan.current_placement_reserve_shortfall_bytes == 2 * GIB_BYTES + def test_validation_checks_measured_placement_without_candidate_frontier(self): + report = _report( + budget_gib=30, + estimated_gib=20, + target_workload_measured=True, + phase_peaks_gib={"decode": 28}, + ) + + assert current_placement_reserve_shortfall_bytes([report]) == 2 * GIB_BYTES + def test_unobserved_later_component_gets_a_conservative_phase(self): first_transformer = _candidate("first_transformer", weight_gib=26, h2d_gib=100) second_transformer = ResidencyTarget( diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py index 5249f4824e49..e8dc23aea733 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -162,6 +162,9 @@ def test_worker_validation_does_not_apply_a_second_candidate_plan(self): target_resident_weight_bytes=1, h2d_bytes_per_request=1, ) + worker._auto_residency_last_applied_plan = AutoResidencyPlan( + changes=[candidate] + ) with ( mock.patch.object( @@ -177,13 +180,19 @@ def test_worker_validation_does_not_apply_a_second_candidate_plan(self): mock.patch.object( gpu_worker_module, "plan_auto_residency", - return_value=AutoResidencyPlan(changes=[candidate]), - ), + ) as plan, mock.patch.object(gpu_worker_module, "apply_residency_changes") as apply, ): response = worker.apply_auto_residency(validate_only=True) self.assertEqual(response.output["status"], PLACEMENT_STATUS_VALIDATED) + self.assertIsNone(worker._auto_residency_last_applied_plan) + worker._build_auto_residency_report.assert_called_once_with( + workload=mock.ANY, + records=[mock.ANY], + include_candidates=False, + ) + plan.assert_not_called() apply.assert_not_called() def test_aux_resident_promotion_uses_short_validation_but_dit_does_not(self): From 888f715b8238021a2c7579387972c220bef8ea77 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 16:18:02 +0800 Subject: [PATCH 53/90] [diffusion] remove obsolete residency retry path --- .../runtime/managers/gpu_worker.py | 16 ------ .../test/unit/test_auto_residency_warmup.py | 49 ------------------- 2 files changed, 65 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 51e4f74b9fdf..e60b4d55e928 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1410,22 +1410,6 @@ def apply_auto_residency(self, *, validate_only: bool = False) -> OutputBatch: ) ) plan = plan_auto_residency(reports=reports) - if ( - self._auto_residency_round_sizes - and plan.current_placement_reserve_shortfall_bytes > 0 - ): - shortfall_gib = plan.current_placement_reserve_shortfall_bytes / GIB_BYTES - if self.is_output_rank: - logger.warning( - "Auto residency calibration exceeded the VRAM reserve by " - "%.1f GiB; rolling back the latest adjustment round.", - shortfall_gib, - ) - return self._rollback_everywhere( - cause=f"VRAM reserve exceeded by {shortfall_gib:.1f} GiB", - already_failed=False, - latest_round_only=True, - ) summary = format_plan_summary(plan=plan, workload=workload, records=records) if plan.skip_reason is not None or not plan.changes: if self.is_output_rank: diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py index e8dc23aea733..5002d37ee8b3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -10,14 +10,12 @@ from sglang.multimodal_gen.runtime.managers import gpu_worker as gpu_worker_module from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( - GIB_BYTES, PLACEMENT_STATUS_ADJUSTED, PLACEMENT_STATUS_ROLLBACK_FAILED, PLACEMENT_STATUS_ROLLED_BACK, PLACEMENT_STATUS_VALIDATED, AppliedResidencyChange, AutoResidencyPlan, - DefaultWorkload, RankResidencyReport, ResidencyTarget, ) @@ -89,53 +87,6 @@ def test_worker_rolls_back_a_materially_slower_calibrated_round(self): self.assertFalse(kwargs["already_failed"]) self.assertTrue(kwargs["latest_round_only"]) - def test_new_placement_that_spends_the_reserve_rolls_back(self): - worker = GPUWorker.__new__(GPUWorker) - worker.rank = 0 - worker.is_output_rank = False - worker.pipeline = SimpleNamespace(modules={}) - worker.server_args = SimpleNamespace( - residency_mode=lambda _name: COMPONENT_OFFLOAD - ) - worker._auto_residency_warmup_records = [] - worker._auto_residency_round_sizes = [1] - worker._build_auto_residency_report = mock.Mock( - return_value=RankResidencyReport( - rank=0, - budget_bytes=30 * GIB_BYTES, - estimated_peak_bytes=28 * GIB_BYTES, - ) - ) - worker._auto_residency_all_gather = lambda value: [value] - expected = OutputBatch(output={"status": PLACEMENT_STATUS_ROLLED_BACK}) - worker._rollback_everywhere = mock.Mock(return_value=expected) - workload = DefaultWorkload(832, 480, 24, 50) - plan = AutoResidencyPlan( - current_placement_reserve_shortfall_bytes=2 * GIB_BYTES - ) - - with ( - mock.patch.object( - gpu_worker_module, "resolve_default_workload", return_value=workload - ), - mock.patch.object( - gpu_worker_module, - "resolve_measured_default_workload", - return_value=workload, - ), - mock.patch.object( - gpu_worker_module, "plan_auto_residency", return_value=plan - ), - ): - response = worker.apply_auto_residency() - - self.assertIs(response, expected) - worker._rollback_everywhere.assert_called_once_with( - cause="VRAM reserve exceeded by 2.0 GiB", - already_failed=False, - latest_round_only=True, - ) - def test_worker_validation_does_not_apply_a_second_candidate_plan(self): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 From 0a12e0f9a557be9b485f6de6d8fa869a8856da99 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 16:45:08 +0800 Subject: [PATCH 54/90] [diffusion] auto-tune compatible custom residency --- .../runtime/managers/gpu_worker.py | 18 +++++--- .../component_residency_strategies.py | 4 ++ .../runtime/pipelines/ltx_2_pipeline.py | 41 ++++++++++++++++++- .../test/unit/test_auto_residency_warmup.py | 18 ++++++++ .../test/unit/test_ltx2_5_config.py | 35 ++++++++++++++++ 5 files changed, 110 insertions(+), 6 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index e60b4d55e928..d9248a412393 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1200,9 +1200,8 @@ def get_can_stay_resident_components( Reuses the auto-residency candidate frontier and benefit ranking, but remains a raw-capacity hint rather than the phase-constrained joint - plan. It omits components driven by a pipeline-custom residency - strategy because user residency flags cannot control them, and it - applies no reserve or activation margin. + plan. It omits components driven by fixed pipeline-custom residency + strategies and applies no reserve or activation margin. """ if not self.pipeline or not self.pipeline.modules: return [] @@ -1215,7 +1214,7 @@ def get_can_stay_resident_components( # chose offload and should learn when the headroom no longer # requires it (automatic adjustment never touches explicit ones). explicit_residency_mode_of=lambda _name: None, - custom_strategy_names=self.pipeline.component_residency_strategies, + custom_strategy_names=self._fixed_custom_residency_strategy_names(), num_inference_steps=( workload.num_inference_steps if workload is not None else 1 ), @@ -1620,7 +1619,7 @@ def _build_auto_residency_report( residency_mode_of=self.server_args.residency_mode, baseline_residency_mode_of=(self.server_args.configured_residency_mode), explicit_residency_mode_of=self.server_args.explicit_residency_mode, - custom_strategy_names=self.pipeline.component_residency_strategies, + custom_strategy_names=self._fixed_custom_residency_strategy_names(), num_inference_steps=workload.num_inference_steps, allow_host_pin_reallocation=True, mixed_dtype_components=self._mixed_dtype_residency_components(), @@ -1726,6 +1725,15 @@ def _mixed_dtype_residency_components() -> set[str]: return set() return manager.components_with_mixed_use_dtypes() + def _fixed_custom_residency_strategy_names(self) -> set[str]: + if self.pipeline is None: + return set() + return { + name + for name, strategy in self.pipeline.component_residency_strategies.items() + if not strategy.supports_auto_residency() + } + def _latest_auto_residency_round(self) -> list[AppliedResidencyChange]: if not self._auto_residency_round_sizes: return [] diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py index ee60f09e1040..cfd05ef4f58e 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py @@ -60,6 +60,10 @@ def is_fsdp_managed_module(module: nn.Module) -> bool: class ComponentResidencyStrategy: """Controls one component's device placement around declared use intervals.""" + def supports_auto_residency(self) -> bool: + """Whether runtime placement may change underneath this strategy.""" + return False + def prepare_for_use( self, module: nn.Module, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py index 988ec73c3064..342476e73c12 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -22,6 +22,7 @@ ComponentResidencyStrategy, ComponentUse, ResidencyState, + build_component_residency_strategy, ) from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, @@ -440,6 +441,12 @@ class LTX2TwoStageResidencyStrategy(ComponentResidencyStrategy): def __init__(self, manager: "LTX2TwoStageResidencyController") -> None: self.manager = manager + self._placement_strategies: dict[ + tuple[str, int, str], ComponentResidencyStrategy + ] = {} + + def supports_auto_residency(self) -> bool: + return True @property def pipeline(self) -> "LTX2TwoStagePipeline": @@ -457,6 +464,19 @@ def _phase(self, use: ComponentUse) -> str: def initialize(self) -> None: pass + def _placement_strategy( + self, module: torch.nn.Module, use: ComponentUse + ) -> ComponentResidencyStrategy: + mode = self.server_args.residency_mode(use.component_name) + key = (use.component_name, id(module), mode) + strategy = self._placement_strategies.get(key) + if strategy is None: + strategy = build_component_residency_strategy( + use.component_name, module, self.server_args + ) + self._placement_strategies[key] = strategy + return strategy + def prepare_for_use( self, module: torch.nn.Module, @@ -466,6 +486,20 @@ def prepare_for_use( phase = self._phase(use) if phase != self.manager._active_phase: self.enter_phase(phase) + self._placement_strategy(module, use).prepare_for_use(module, use, state) + + def prefetch_for_use( + self, + module: torch.nn.Module, + use: ComponentUse, + state: ResidencyState, + ) -> bool: + phase = self._phase(use) + if phase != self.manager._active_phase: + self.enter_phase(phase) + return self._placement_strategy(module, use).prefetch_for_use( + module, use, state + ) def wait_for_use( self, @@ -473,6 +507,7 @@ def wait_for_use( use: ComponentUse, state: ResidencyState, ) -> None: + self._placement_strategy(module, use).wait_for_use(module, use, state) self.ensure_phase_ready(self._phase(use)) def finish_use( @@ -481,6 +516,7 @@ def finish_use( use: ComponentUse, state: ResidencyState, ) -> None: + self._placement_strategy(module, use).finish_use(module, use, state) self.exit_phase(self._phase(use)) def finish_request( @@ -491,8 +527,11 @@ def finish_request( *, preferred: bool, ) -> None: + self._placement_strategy(module, use).finish_request( + module, use, state, preferred=preferred + ) if not preferred: - self.finish_use(module, use, state) + self.exit_phase(self._phase(use)) return phase = self._phase(use) if phase != self.manager._active_phase: diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py index 5002d37ee8b3..212e5ecdcfe6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -23,10 +23,28 @@ COMPONENT_OFFLOAD, RESIDENT, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import ( + ComponentResidencyStrategy, +) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch class TestAutoResidencyWarmup(unittest.TestCase): + def test_worker_skips_only_fixed_custom_residency_strategies(self): + class AutoCompatibleStrategy(ComponentResidencyStrategy): + def supports_auto_residency(self) -> bool: + return True + + worker = GPUWorker.__new__(GPUWorker) + worker.pipeline = SimpleNamespace( + component_residency_strategies={ + "fixed": ComponentResidencyStrategy(), + "dynamic": AutoCompatibleStrategy(), + } + ) + + self.assertEqual(worker._fixed_custom_residency_strategy_names(), {"fixed"}) + def test_rewarm_step_limit_keeps_full_shape_but_shortens_denoising(self): req = SimpleNamespace(is_warmup=True, extra={}, num_inference_steps=4) with mock.patch.object(server_warmup, "build_warmup_reqs", return_value=[req]): diff --git a/python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py b/python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py index 256344885adb..daff47f8426d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py +++ b/python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py @@ -699,5 +699,40 @@ def test_explicit_component_path_wins(self): self.assertEqual(args.component_paths["transformer"], "/some/other/transformer") +class TestLTX2TwoStageAutoResidency(unittest.TestCase): + def test_custom_strategy_delegates_after_runtime_placement_change(self): + from sglang.multimodal_gen.runtime.pipelines import ltx_2_pipeline + + mode = {"value": "resident"} + server_args = SimpleNamespace( + residency_mode=lambda _name: mode["value"], + ) + manager = SimpleNamespace( + pipeline=SimpleNamespace(), + server_args=server_args, + _active_phase="stage1", + ) + strategy = ltx_2_pipeline.LTX2TwoStageResidencyStrategy(manager) + resident = mock.Mock() + offloaded = mock.Mock() + module = mock.Mock() + use = SimpleNamespace(component_name="transformer", phase="stage1") + state = SimpleNamespace() + + with mock.patch.object( + ltx_2_pipeline, + "build_component_residency_strategy", + side_effect=(resident, offloaded), + ) as build: + strategy.prepare_for_use(module, use, state) + mode["value"] = "component-offload" + strategy.prepare_for_use(module, use, state) + + self.assertTrue(strategy.supports_auto_residency()) + self.assertEqual(build.call_count, 2) + resident.prepare_for_use.assert_called_once_with(module, use, state) + offloaded.prepare_for_use.assert_called_once_with(module, use, state) + + if __name__ == "__main__": unittest.main() From f843683bb88af9bbdec18aa1f343af06e0be27db Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 16:53:36 +0800 Subject: [PATCH 55/90] [diffusion] optimize layerwise host pin frontier --- .../memory_managers/auto_residency.py | 117 +++++++++++++++--- .../test/unit/test_auto_residency.py | 41 ++++++ 2 files changed, 142 insertions(+), 16 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 956e710d5a06..ce96100df3bd 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -822,14 +822,16 @@ def _layerwise_pin_targets( current_pinned_layers: tuple[tuple[int, ...], ...], uses_per_streamed_layer: int, ) -> list[tuple[tuple[int, ...], ...]]: - """Bounded useful HostPin frontier for one resident-layer placement. + """Pareto-optimal HostPin targets for one resident-layer placement. - Streamed layers have the highest value per byte because they transfer every - denoise step. Deterministic prefixes keep candidate generation linear in - the layer count; the joint planner still evaluates those targets against - rank VRAM, HostPin, and transition headroom together. + Layers with identical transfer value, host cost and current pin state are + interchangeable, so they are grouped before subset construction. This + keeps repeated transformer blocks linear while retaining non-prefix + packings needed when layer sizes differ. """ - ordered: list[tuple[int, int, int]] = [] + current = [set(indices) for indices in current_pinned_layers] + current_bytes = 0 + grouped: dict[tuple[int, int, int, bool], list[tuple[int, int]]] = {} for manager_index, (manager, resident_count) in enumerate( zip(managers, resident_layers) ): @@ -842,18 +844,101 @@ def _layerwise_pin_targets( ) if not manager.pin_cpu_memory: continue + transfer_bytes = manager.layer_weight_bytes() + host_bytes = manager.layer_host_store_bytes() for layer_idx in manager.pinnable_layer_indices(): uses = uses_per_streamed_layer if layer_idx in streamed else 1 - ordered.append((-uses, manager_index, layer_idx)) - ordered.sort() - - targets: list[tuple[tuple[int, ...], ...]] = [current_pinned_layers] - selected = [set() for _ in managers] - targets.append(tuple(() for _ in managers)) - for _, manager_index, layer_idx in ordered: - selected[manager_index].add(layer_idx) - targets.append(tuple(tuple(sorted(indices)) for indices in selected)) - return list(dict.fromkeys(targets)) + layer_host_bytes = host_bytes.get(layer_idx, 0) + is_current = layer_idx in current[manager_index] + if is_current: + current_bytes += layer_host_bytes + grouped.setdefault( + ( + uses, + transfer_bytes.get(layer_idx, 0), + layer_host_bytes, + is_current, + ), + [], + ).append((manager_index, layer_idx)) + + # target HostPin bytes, avoided pageable-transfer work, unpin scratch, + # pin scratch, selected layer indices. + states = [(0, 0, current_bytes, 0, tuple(() for _ in managers))] + + def dominates(left, right) -> bool: + return ( + left[0] <= right[0] + and left[1] >= right[1] + and left[2] <= right[2] + and left[3] <= right[3] + ) + + def prune(candidates): + best_by_cost = {} + for state in candidates: + key = state[:4] + incumbent = best_by_cost.get(key) + if incumbent is None or state[4] < incumbent[4]: + best_by_cost[key] = state + frontier = [] + for state in sorted( + best_by_cost.values(), + key=lambda item: (item[0], item[2], item[3], -item[1], item[4]), + ): + if any(dominates(existing, state) for existing in frontier): + continue + frontier = [ + existing for existing in frontier if not dominates(state, existing) + ] + frontier.append(state) + return frontier + + for (uses, transfer_bytes, host_bytes, is_current), layers in sorted( + grouped.items(), + key=lambda item: (-item[0][0], -item[0][1], item[0][2], item[0][3]), + ): + layers.sort() + choices = [] + selected = [set() for _ in managers] + for count in range(len(layers) + 1): + if count: + manager_index, layer_idx = layers[count - 1] + selected[manager_index].add(layer_idx) + choices.append( + ( + count * host_bytes, + count * uses * transfer_bytes, + -count * host_bytes if is_current else 0, + 0 if is_current else count * host_bytes, + tuple(tuple(sorted(indices)) for indices in selected), + ) + ) + + expanded = [] + for state in states: + for choice in choices: + expanded.append( + ( + state[0] + choice[0], + state[1] + choice[1], + state[2] + choice[2], + state[3] + choice[3], + tuple( + tuple(sorted(set(left) | set(right))) + for left, right in zip(state[4], choice[4]) + ), + ) + ) + states = prune(expanded) + + return [ + state[4] + for state in sorted( + states, + key=lambda item: (item[0], item[2], item[3], -item[1], item[4]), + ) + ] def _layerwise_resident_targets(managers: Sequence) -> list[tuple[int, ...]]: diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 11e278661ea5..aec9ac598b0a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -20,6 +20,7 @@ RankResidencyReport, ResidencyTarget, WarmupMemoryRecord, + _layerwise_pin_targets, _layerwise_resident_targets, apply_residency_changes, collect_residency_targets, @@ -1835,6 +1836,46 @@ def test_resident_frontier_includes_absolute_and_ratio_cli_states(self): assert (2, 2) in targets assert (1, 2) in targets + def test_host_pin_frontier_includes_non_prefix_packings(self): + manager = SimpleNamespace( + num_layers=3, + residency_policy="leading", + pin_cpu_memory=True, + layer_weight_bytes=lambda: {0: 6, 1: 4, 2: 4}, + layer_host_store_bytes=lambda: {0: 6, 1: 4, 2: 4}, + pinnable_layer_indices=lambda: (0, 1, 2), + ) + + targets = _layerwise_pin_targets( + managers=[manager], + resident_layers=(0,), + current_pinned_layers=((),), + uses_per_streamed_layer=1, + ) + + assert ((0,),) in targets + assert ((1, 2),) in targets + + def test_host_pin_frontier_coalesces_equal_transformer_layers(self): + layer_bytes = {index: 8 for index in range(40)} + manager = SimpleNamespace( + num_layers=40, + residency_policy="leading", + pin_cpu_memory=True, + layer_weight_bytes=lambda: layer_bytes, + layer_host_store_bytes=lambda: layer_bytes, + pinnable_layer_indices=lambda: tuple(layer_bytes), + ) + + targets = _layerwise_pin_targets( + managers=[manager], + resident_layers=(0,), + current_pinned_layers=((),), + uses_per_streamed_layer=20, + ) + + assert len(targets) == 41 + def test_filters_and_sizes(self): te = nn.Linear(8, 8) vae = nn.Linear(4, 4) From b8b2943121587f417b369a669ad26c7e16527b2a Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 17:28:35 +0800 Subject: [PATCH 56/90] [diffusion] calibrate layerwise placement per group --- .../runtime/managers/gpu_worker.py | 35 +++ .../memory_managers/auto_residency.py | 258 ++++++++++++++++-- .../memory_managers/layerwise_offload.py | 53 ++++ .../test/unit/test_auto_residency.py | 111 ++++++++ .../test/unit/test_layerwise_offload.py | 31 +++ 5 files changed, 458 insertions(+), 30 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index d9248a412393..4703ddb6c48f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -74,6 +74,7 @@ estimate_candidate_latency_savings_ns, estimate_default_workload_peak_bytes, estimate_default_workload_timing, + estimate_layerwise_layer_uses, estimate_workload_phase_peaks, format_applied_changes, format_plan_summary, @@ -98,6 +99,7 @@ host_memory_available_bytes, ) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + LayerwiseUsageTracker, configure_layerwise_offload_modules, ) from sglang.multimodal_gen.runtime.managers.memory_managers.memory_occupation_controller import ( @@ -589,6 +591,7 @@ def _execute_forward_common( else None ) warmup_baseline_allocated_bytes = 0 + layerwise_usage_tracker: LayerwiseUsageTracker | None = None try: if measure_server_warmup: # Drop the previous request's allocator pool so each probe @@ -601,6 +604,13 @@ def _execute_forward_common( warmup_baseline_allocated_bytes = ( torch.get_device_module().memory_allocated() ) + if ( + self.server_args.performance_mode == "auto" + and self.pipeline is not None + ): + layerwise_usage_tracker = LayerwiseUsageTracker( + self.pipeline.modules + ) start_time = ( execution_start_time @@ -718,11 +728,17 @@ def _execute_forward_common( # or the estimator would plan from the remaining partial data if measure_server_warmup: assert warmup_workload is not None + layerwise_layer_uses = ( + layerwise_usage_tracker.finish() + if layerwise_usage_tracker is not None + else {} + ) self._record_server_warmup_memory( req=req, workload=warmup_workload, baseline_allocated_bytes=warmup_baseline_allocated_bytes, succeeded=output_batch is not None and output_batch.error is None, + layerwise_layer_uses=layerwise_layer_uses, ) return output_batch @@ -733,6 +749,7 @@ def _record_server_warmup_memory( workload: tuple[int, int, int, int], baseline_allocated_bytes: int, succeeded: bool, + layerwise_layer_uses: dict[str, dict[str, tuple[int, ...]]] | None = None, ) -> None: phase_allocated_peaks: dict[str, int] = {} phase_components: dict[str, tuple[str, ...]] = {} @@ -778,6 +795,7 @@ def _record_server_warmup_memory( phase_full_weight_transition_components=( phase_full_weight_transition_components ), + layerwise_layer_uses=layerwise_layer_uses or {}, num_inference_steps=num_inference_steps, total_duration_ms=( float(metrics.total_duration_ms) if metrics is not None else 0.0 @@ -1602,6 +1620,19 @@ def _build_auto_residency_report( target_units=target_units, component_weight_bytes=runtime_weights_by_component, ) + measured_used_components = { + component_name + for component_names in used_components_by_phase.values() + for component_name in component_names + } + has_component_use_measurement = any( + record.phase_used_components for record in records if record.succeeded + ) + layerwise_layer_uses = estimate_layerwise_layer_uses( + records=records, + target_units=target_units, + target_num_inference_steps=workload.num_inference_steps, + ) # What this process may still grow into: free VRAM plus what its # allocator already reserved. Unlike the raw device total, this # excludes memory held by other processes on a shared GPU. @@ -1632,6 +1663,10 @@ def _build_auto_residency_report( self.server_args.layerwise_tuning_for(name, dit_group=dit_group) ), pin_cpu_memory=self.server_args.pin_cpu_memory, + used_components=( + measured_used_components if has_component_use_measurement else None + ), + layerwise_layer_uses=layerwise_layer_uses, ) measured_request_duration_ns, _, _ = estimate_default_workload_timing( records=records, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index ce96100df3bd..ed084ca5d824 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -25,7 +25,7 @@ from __future__ import annotations import statistics -from itertools import chain +from itertools import chain, product from typing import TYPE_CHECKING, Callable, Iterable, Mapping, Sequence import msgspec @@ -134,6 +134,7 @@ class WarmupMemoryRecord(msgspec.Struct, frozen=True): phase_active_components: dict[str, tuple[str, ...]] = {} phase_used_components: dict[str, tuple[str, ...]] = {} phase_full_weight_transition_components: dict[str, tuple[str, ...]] = {} + layerwise_layer_uses: dict[str, dict[str, tuple[int, ...]]] = {} num_inference_steps: int = 1 total_duration_ms: float = 0.0 stage_duration_ms: dict[str, float] = {} @@ -331,6 +332,55 @@ def resolve_measured_default_workload( ) +def estimate_layerwise_layer_uses( + *, + records: Iterable[WarmupMemoryRecord], + target_units: int | None, + target_num_inference_steps: int, +) -> dict[str, dict[str, tuple[int, ...]]]: + """Estimate per-request layer calls from the same calibration forward. + + A full-shape memory probe deliberately runs only a few denoise steps. A + layer observed at least twice is treated as step-scoped and scaled to the + target step count; a layer observed once remains a one-shot encoder, + refiner, or decoder layer even when it belongs to a DiT component. + """ + successful = [record for record in records if record.succeeded] + if target_units is not None: + covering = [ + record for record in successful if record.workload_units() >= target_units + ] + if covering: + successful = covering + + estimated: dict[str, dict[str, list[int]]] = {} + for record in successful: + source_steps = max(1, record.num_inference_steps) + for component_name, groups in record.layerwise_layer_uses.items(): + component = estimated.setdefault(component_name, {}) + for layer_name, counts in groups.items(): + target = component.setdefault(layer_name, [0] * len(counts)) + if len(target) != len(counts): + continue + for layer_index, count in enumerate(counts): + scaled = count + if ( + is_dit_component_name(component_name) + and count > 1 + and target_num_inference_steps > source_steps + ): + scaled = ( + count * target_num_inference_steps + source_steps - 1 + ) // source_steps + target[layer_index] = max(target[layer_index], scaled) + return { + component_name: { + layer_name: tuple(counts) for layer_name, counts in groups.items() + } + for component_name, groups in estimated.items() + } + + def estimate_default_workload_timing( *, records: Iterable[WarmupMemoryRecord], @@ -789,6 +839,7 @@ def _layerwise_transfer_work_bytes( resident_layers: tuple[int, ...], pinned_layers: tuple[tuple[int, ...], ...], uses_per_streamed_layer: int, + layer_uses: tuple[tuple[int, ...], ...] | None = None, ) -> int: """Relative transfer work for one request under a layerwise placement. @@ -796,9 +847,14 @@ def _layerwise_transfer_work_bytes( internal pinned buffer and the copy cannot run ahead of compute. This is a conservative ordering metric rather than a latency prediction. """ + resolved_uses = _resolve_layerwise_uses( + managers=managers, + layer_uses=layer_uses, + fallback_uses=uses_per_streamed_layer, + ) total = 0 - for manager, resident_count, pinned_indices in zip( - managers, resident_layers, pinned_layers + for manager, resident_count, pinned_indices, manager_uses in zip( + managers, resident_layers, pinned_layers, resolved_uses ): streamed = set( compute_streamed_layers( @@ -809,18 +865,76 @@ def _layerwise_transfer_work_bytes( ) pinned = set(pinned_indices) for layer_idx, weight_bytes in manager.layer_weight_bytes().items(): - uses = uses_per_streamed_layer if layer_idx in streamed else 1 + observed_uses = manager_uses[layer_idx] + uses = observed_uses if layer_idx in streamed else min(1, observed_uses) transfer_multiplier = 1 if layer_idx in pinned else 2 total += uses * transfer_multiplier * weight_bytes return total +def _resolve_layerwise_uses( + *, + managers: Sequence, + layer_uses: tuple[tuple[int, ...], ...] | None, + fallback_uses: int, +) -> tuple[tuple[int, ...], ...]: + if layer_uses is None: + return tuple( + tuple(max(0, fallback_uses) for _ in range(manager.num_layers)) + for manager in managers + ) + if len(layer_uses) != len(managers): + raise ValueError("layerwise usage group count changed") + resolved = [] + for manager, counts in zip(managers, layer_uses): + if len(counts) != manager.num_layers: + raise ValueError("layerwise usage layer count changed") + resolved.append(tuple(max(0, int(count)) for count in counts)) + return tuple(resolved) + + +def _layer_uses_for_managers( + *, + managers: Sequence, + uses_by_layer_name: Mapping[str, tuple[int, ...]] | None, + fallback_uses: int, +) -> tuple[tuple[int, ...], ...] | None: + if uses_by_layer_name is None: + return None + return tuple( + uses_by_layer_name.get( + manager.layers_attr_str, + tuple(max(0, fallback_uses) for _ in range(manager.num_layers)), + ) + for manager in managers + ) + + +def _layerwise_active_peak_device_bytes( + *, + managers: Sequence, + resident_layers: tuple[int, ...], + layer_uses: tuple[tuple[int, ...], ...] | None, +) -> int: + if layer_uses is None: + return sum( + manager.peak_managed_device_weight_bytes(count) + for manager, count in zip(managers, resident_layers) + ) + return sum( + manager.peak_managed_device_weight_bytes(count) + for manager, count, uses in zip(managers, resident_layers, layer_uses) + if any(uses) + ) + + def _layerwise_pin_targets( *, managers: Sequence, resident_layers: tuple[int, ...], current_pinned_layers: tuple[tuple[int, ...], ...], uses_per_streamed_layer: int, + layer_uses: tuple[tuple[int, ...], ...] | None = None, ) -> list[tuple[tuple[int, ...], ...]]: """Pareto-optimal HostPin targets for one resident-layer placement. @@ -829,11 +943,16 @@ def _layerwise_pin_targets( keeps repeated transformer blocks linear while retaining non-prefix packings needed when layer sizes differ. """ + resolved_uses = _resolve_layerwise_uses( + managers=managers, + layer_uses=layer_uses, + fallback_uses=uses_per_streamed_layer, + ) current = [set(indices) for indices in current_pinned_layers] current_bytes = 0 grouped: dict[tuple[int, int, int, bool], list[tuple[int, int]]] = {} - for manager_index, (manager, resident_count) in enumerate( - zip(managers, resident_layers) + for manager_index, (manager, resident_count, manager_uses) in enumerate( + zip(managers, resident_layers, resolved_uses) ): streamed = set( compute_streamed_layers( @@ -847,7 +966,8 @@ def _layerwise_pin_targets( transfer_bytes = manager.layer_weight_bytes() host_bytes = manager.layer_host_store_bytes() for layer_idx in manager.pinnable_layer_indices(): - uses = uses_per_streamed_layer if layer_idx in streamed else 1 + observed_uses = manager_uses[layer_idx] + uses = observed_uses if layer_idx in streamed else min(1, observed_uses) layer_host_bytes = host_bytes.get(layer_idx, 0) is_current = layer_idx in current[manager_index] if is_current: @@ -941,8 +1061,23 @@ def prune(candidates): ] -def _layerwise_resident_targets(managers: Sequence) -> list[tuple[int, ...]]: - """All resident-count tuples expressible by the component-level CLI knob. +def _layerwise_resident_targets( + managers: Sequence, + *, + layer_uses: tuple[tuple[int, ...], ...] | None = None, +) -> list[tuple[int, ...]]: + """Useful resident-count tuples for the measured request. + + With manager-level measurements, each repeatedly executed layer group is + independent. This admits states such as ``(encoder=0, decoder=k)`` or + separate double/single-block DiT allocations. A group used at most once + cannot avoid any per-request H2D bytes through stage residency: its layers + are still loaded once when the component activates. Keep it streamed and + let HostPin or whole-component residency represent its useful choices. + + Without measurements, retain every state expressible by the public + component-level CLI knob. This fallback is used by static/unit callers and + avoids inventing per-group workload semantics before calibration. ``--layerwise-resident-layers`` accepts either one absolute count or one ratio for a component. Those are different paths when a component owns @@ -953,6 +1088,18 @@ def _layerwise_resident_targets(managers: Sequence) -> list[tuple[int, ...]]: explicit interface can represent. """ layer_counts = tuple(manager.num_layers for manager in managers) + if layer_uses is not None: + resolved_uses = _resolve_layerwise_uses( + managers=managers, + layer_uses=layer_uses, + fallback_uses=0, + ) + choices = [ + range(num_layers + 1) if any(count > 1 for count in uses) else (0,) + for num_layers, uses in zip(layer_counts, resolved_uses) + ] + return list(product(*choices)) + targets = {tuple(0 for _ in managers)} for count in range(1, max(layer_counts) + 1): @@ -981,11 +1128,13 @@ class _EstimatedLayerwiseManager: def __init__( self, *, + layers_attr_str: str, layer_bytes: dict[int, int], prefetch_size: int, residency_policy: str, pin_cpu_memory: bool, ) -> None: + self.layers_attr_str = layers_attr_str self._layer_bytes = layer_bytes self.num_layers = len(layer_bytes) self.prefetch_size = min(max(1, prefetch_size), self.num_layers) @@ -1066,6 +1215,7 @@ def _estimate_unconfigured_layerwise_managers( prefetch_size = int(prefetch_value) estimates.append( _EstimatedLayerwiseManager( + layers_attr_str=layer_name, layer_bytes=estimate_layer_weight_bytes(layers), prefetch_size=prefetch_size, residency_policy=residency_policy, @@ -1085,6 +1235,8 @@ def _unconfigured_layerwise_targets( prefetch_value: float, residency_policy: str, pin_cpu_memory: bool, + component_used: bool, + layer_uses_by_name: Mapping[str, tuple[int, ...]] | None, ) -> tuple[list[ResidencyTarget], int]: """Virtual layerwise frontier for a component still using coarse offload. @@ -1107,7 +1259,14 @@ def _unconfigured_layerwise_targets( return [], 0 unmanaged_weight_bytes = max(0, full_weight_bytes - managed_weight_bytes) uses_per_request = ( - max(1, num_inference_steps) if is_dit_component_name(component_name) else 1 + max(1, num_inference_steps) + if component_used and is_dit_component_name(component_name) + else int(component_used) + ) + layer_uses = _layer_uses_for_managers( + managers=managers, + uses_by_layer_name=layer_uses_by_name, + fallback_uses=uses_per_request, ) empty_pins = tuple(() for _ in managers) maximum_transfer_work = _layerwise_transfer_work_bytes( @@ -1115,13 +1274,16 @@ def _unconfigured_layerwise_targets( resident_layers=tuple(0 for _ in managers), pinned_layers=empty_pins, uses_per_streamed_layer=uses_per_request, + layer_uses=layer_uses, ) - coarse_transfer_work = maximum_transfer_work // uses_per_request + coarse_transfer_work = 2 * full_weight_bytes if component_used else 0 coarse_savings = max(0, maximum_transfer_work - coarse_transfer_work) current_resident = current_mode == RESIDENT current_inactive_bytes = full_weight_bytes if current_resident else 0 - initial_managed_bytes = sum( - manager.peak_managed_device_weight_bytes(0) for manager in managers + initial_managed_bytes = _layerwise_active_peak_device_bytes( + managers=managers, + resident_layers=tuple(0 for _ in managers), + layer_uses=layer_uses, ) host_materialize_scratch = max( ( @@ -1132,20 +1294,22 @@ def _unconfigured_layerwise_targets( default=0, ) targets = [] - for resident_layers in _layerwise_resident_targets(managers): + for resident_layers in _layerwise_resident_targets(managers, layer_uses=layer_uses): resident_bytes = sum( manager.resident_weight_bytes(count) for manager, count in zip(managers, resident_layers) ) - active_managed_bytes = sum( - manager.peak_managed_device_weight_bytes(count) - for manager, count in zip(managers, resident_layers) + active_managed_bytes = _layerwise_active_peak_device_bytes( + managers=managers, + resident_layers=resident_layers, + layer_uses=layer_uses, ) transfer_work = _layerwise_transfer_work_bytes( managers=managers, resident_layers=resident_layers, pinned_layers=empty_pins, uses_per_streamed_layer=uses_per_request, + layer_uses=layer_uses, ) targets.append( ResidencyTarget( @@ -1217,6 +1381,8 @@ def collect_residency_targets( auto_resident_components: Iterable[str] = (), layerwise_tuning_of: Callable[[str, bool], tuple[float, float, str]] | None = None, pin_cpu_memory: bool = True, + used_components: Iterable[str] | None = None, + layerwise_layer_uses: Mapping[str, Mapping[str, tuple[int, ...]]] | None = None, ) -> list[ResidencyTarget]: """Build complete target-state frontiers for auto-managed components. @@ -1228,6 +1394,7 @@ def collect_residency_targets( custom_names = set(custom_strategy_names) mixed_dtype_names = set(mixed_dtype_components) auto_resident_names = set(auto_resident_components) + measured_used_names = set(used_components) if used_components is not None else None if baseline_residency_mode_of is None: baseline_residency_mode_of = residency_mode_of candidates = [] @@ -1244,7 +1411,21 @@ def collect_residency_targets( continue if is_fsdp_managed_module(module): continue + component_layer_uses = ( + layerwise_layer_uses.get(name) if layerwise_layer_uses is not None else None + ) + measured_layer_use = bool( + component_layer_uses + and any( + count for counts in component_layer_uses.values() for count in counts + ) + ) current_mode = residency_mode_of(name) + component_used = ( + measured_used_names is None + or name in measured_used_names + or measured_layer_use + ) if current_mode == RESIDENT and name not in auto_resident_names: # A loader or another runtime feature owns this hard requirement. continue @@ -1290,16 +1471,16 @@ def collect_residency_targets( prefetch_value=prefetch_value, residency_policy=residency_policy, pin_cpu_memory=pin_cpu_memory, + component_used=component_used, + layer_uses_by_name=component_layer_uses, ) ) uses_per_request = ( - max(1, num_inference_steps) if is_dit_component_name(name) else 1 - ) - component_transfer_work = ( - maximum_transfer_work // uses_per_request - if maximum_transfer_work - else 0 + max(1, num_inference_steps) + if component_used and is_dit_component_name(name) + else int(component_used) ) + component_transfer_work = 2 * weight_bytes if component_used else 0 candidates.extend( [ ResidencyTarget( @@ -1372,7 +1553,14 @@ def collect_residency_targets( # A layerwise DiT re-streams its layers once per denoise forward; # every other offloaded component transfers once per request. uses_per_request = ( - max(1, num_inference_steps) if is_dit_component_name(name) else 1 + max(1, num_inference_steps) + if component_used and is_dit_component_name(name) + else int(component_used) + ) + layer_uses = _layer_uses_for_managers( + managers=managers, + uses_by_layer_name=component_layer_uses, + fallback_uses=uses_per_request, ) current_resident_layers = tuple(manager.resident_layers for manager in managers) current_pinned_layers = tuple( @@ -1381,7 +1569,11 @@ def collect_residency_targets( current_peak_device_bytes = ( managed_weight_bytes if current_permanent - else sum(manager.peak_managed_device_weight_bytes() for manager in managers) + else _layerwise_active_peak_device_bytes( + managers=managers, + resident_layers=current_resident_layers, + layer_uses=layer_uses, + ) ) current_inactive_device_bytes = managed_weight_bytes if current_permanent else 0 current_pinned_bytes = sum( @@ -1392,10 +1584,11 @@ def collect_residency_targets( resident_layers=tuple(0 for _ in managers), pinned_layers=tuple(() for _ in managers), uses_per_streamed_layer=uses_per_request, + layer_uses=layer_uses, ) full_weight_bytes = _module_weight_bytes(module) unmanaged_weight_bytes = max(0, full_weight_bytes - managed_weight_bytes) - component_transfer_work = 2 * full_weight_bytes + component_transfer_work = 2 * full_weight_bytes if component_used else 0 maximum_transfer_work = max( layerwise_maximum_transfer_work, component_transfer_work ) @@ -1435,14 +1628,17 @@ def collect_residency_targets( ) ) - for target_resident_layers in _layerwise_resident_targets(managers): + for target_resident_layers in _layerwise_resident_targets( + managers, layer_uses=layer_uses + ): target_resident_bytes = sum( manager.resident_weight_bytes(count) for manager, count in zip(managers, target_resident_layers) ) - target_peak_device_bytes = sum( - manager.peak_managed_device_weight_bytes(count) - for manager, count in zip(managers, target_resident_layers) + target_peak_device_bytes = _layerwise_active_peak_device_bytes( + managers=managers, + resident_layers=target_resident_layers, + layer_uses=layer_uses, ) pin_targets = ( _layerwise_pin_targets( @@ -1450,6 +1646,7 @@ def collect_residency_targets( resident_layers=target_resident_layers, current_pinned_layers=current_pinned_layers, uses_per_streamed_layer=uses_per_request, + layer_uses=layer_uses, ) if allow_host_pin_reallocation else [current_pinned_layers] @@ -1467,6 +1664,7 @@ def collect_residency_targets( resident_layers=target_resident_layers, pinned_layers=target_pinned_layers, uses_per_streamed_layer=uses_per_request, + layer_uses=layer_uses, ) unpin_scratch, pin_scratch = _layerwise_host_transition_bytes( managers=managers, 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 88149e61b136..f53614d5000b 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 @@ -2357,6 +2357,59 @@ def finalize_resident_layerwise_placement(self) -> None: release_unused_pinned_memory() +class LayerwiseUsageTracker: + """Temporary per-layer call counters for one calibration request. + + Managers cannot provide this information when layerwise offload has not + been configured yet or was disabled by a resident placement. Observe the + declared layer groups directly and remove every hook before returning, so + ordinary serving requests pay no counter or hook-dispatch cost. + """ + + def __init__(self, modules: Mapping[str, object]) -> None: + self._handles: list[Any] = [] + self._counts: dict[str, dict[str, list[int]]] = {} + for component_name, module in modules.items(): + if not isinstance(module, LayerwiseOffloadableModuleMixin): + continue + named_modules = dict(module.named_modules()) + component_counts: dict[str, list[int]] = {} + for layer_name in module.layer_names: + layers = named_modules.get(layer_name) + if not isinstance(layers, (torch.nn.ModuleList, torch.nn.Sequential)): + continue + counts = [0] * len(layers) + if not counts: + continue + component_counts[layer_name] = counts + for layer_index, layer in enumerate(layers): + + def record_use( + _module, + _inputs, + *, + target_counts=counts, + target_index=layer_index, + ) -> None: + target_counts[target_index] += 1 + + self._handles.append(layer.register_forward_pre_hook(record_use)) + if component_counts: + self._counts[component_name] = component_counts + + def finish(self) -> dict[str, dict[str, tuple[int, ...]]]: + for handle in self._handles: + handle.remove() + self._handles.clear() + return { + component_name: { + layer_name: tuple(counts) + for layer_name, counts in component_counts.items() + } + for component_name, component_counts in self._counts.items() + } + + def iter_materialized_weights(module: torch.nn.Module): """Yield (name, tensor) pairs with materialized weights, even under offload. diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index aec9ac598b0a..472645fe8954 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -31,6 +31,7 @@ estimate_candidate_latency_savings_ns, estimate_default_workload_peak_bytes, estimate_default_workload_timing, + estimate_layerwise_layer_uses, estimate_workload_phase_peaks, format_applied_changes, plan_auto_residency, @@ -68,6 +69,7 @@ def _record( step_duration_ms=(), phase_active_components=None, phase_full_weight_transition_components=None, + layerwise_layer_uses=None, ) -> WarmupMemoryRecord: return WarmupMemoryRecord( width=width, @@ -84,6 +86,7 @@ def _record( phase_full_weight_transition_components=( phase_full_weight_transition_components or {} ), + layerwise_layer_uses=layerwise_layer_uses or {}, ) @@ -195,6 +198,38 @@ def test_measured_placement_keeps_full_cost_of_slower_dit_target(self): assert savings[virtual_layerwise.option_key()] == -5_000_000_000 +class TestEstimateLayerwiseLayerUses: + def test_scales_repeated_groups_but_not_one_shot_groups(self): + record = _record( + num_inference_steps=2, + layerwise_layer_uses={ + "transformer": { + "token_refiner.blocks": (1, 1), + "blocks": (2, 2, 2), + }, + "vae": { + "encoder.down_blocks": (0, 0), + "decoder.up_blocks": (1, 1), + }, + }, + ) + + uses = estimate_layerwise_layer_uses( + records=[record], + target_units=record.workload_units(), + target_num_inference_steps=10, + ) + + assert uses["transformer"] == { + "token_refiner.blocks": (1, 1), + "blocks": (10, 10, 10), + } + assert uses["vae"] == { + "encoder.down_blocks": (0, 0), + "decoder.up_blocks": (1, 1), + } + + class TestEstimateDefaultWorkloadPeak: def test_same_shape_uses_measured_peak(self): record = _record() @@ -1620,6 +1655,7 @@ def __init__( resident_layers: int = 0, event_log: list[str] | None = None, event_name: str = "manager", + layers_attr_str: str = "layers", ): self.enabled = True self._configured = True @@ -1631,6 +1667,7 @@ def __init__( self._pinned_layers: tuple[int, ...] = () self._event_log = event_log self._event_name = event_name + self.layers_attr_str = layers_attr_str self.fail_load = False self.load_all_layers_calls = 0 self.remove_hooks_calls = 0 @@ -1836,6 +1873,41 @@ def test_resident_frontier_includes_absolute_and_ratio_cli_states(self): assert (2, 2) in targets assert (1, 2) in targets + def test_measured_resident_frontier_allocates_repeated_groups_independently(self): + managers = [ + _FakeLayerwiseManager( + {f"first.{index}.w": torch.zeros(16) for index in range(3)} + ), + _FakeLayerwiseManager( + {f"second.{index}.w": torch.zeros(16) for index in range(2)} + ), + ] + + targets = _layerwise_resident_targets( + managers, + layer_uses=((8, 8, 8), (8, 8)), + ) + + assert (0, 2) in targets + assert (3, 0) in targets + + def test_measured_resident_frontier_streams_one_shot_and_unused_groups(self): + managers = [ + _FakeLayerwiseManager( + {f"encoder.{index}.w": torch.zeros(16) for index in range(3)} + ), + _FakeLayerwiseManager( + {f"decoder.{index}.w": torch.zeros(16) for index in range(2)} + ), + ] + + targets = _layerwise_resident_targets( + managers, + layer_uses=((0, 0, 0), (1, 1)), + ) + + assert targets == [(0, 0)] + def test_host_pin_frontier_includes_non_prefix_packings(self): manager = SimpleNamespace( num_layers=3, @@ -2070,6 +2142,45 @@ def test_non_dit_layerwise_frontier_includes_partial_residency(self): assert component.target_layerwise_pinned_layers == ((),) assert component.pinned_host_delta_bytes == 0 + def test_measured_vae_frontier_does_not_spend_on_unused_encoder_layers(self): + encoder = _FakeLayerwiseManager( + {f"encoder.{index}.w": torch.zeros(16) for index in range(2)}, + layers_attr_str="encoder", + ) + decoder = _FakeLayerwiseManager( + {f"decoder.{index}.w": torch.zeros(16) for index in range(3)}, + layers_attr_str="decoder", + ) + module = _FakeLayerwiseDit([encoder, decoder]) + + candidates = collect_residency_targets( + modules={"vae": module}, + residency_mode_of=self._modes({"vae": LAYERWISE_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=20, + used_components={"vae"}, + layerwise_layer_uses={ + "vae": { + "encoder": (0, 0), + "decoder": (1, 1, 1), + } + }, + ) + layerwise = [ + candidate + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + ] + + assert layerwise + assert { + candidate.target_layerwise_resident_layers for candidate in layerwise + } == {(0, 0)} + assert all( + candidate.target_layerwise_pinned_layers[0] == () for candidate in layerwise + ) + def test_component_offload_frontier_can_expand_layerwise_lazily(self): module = _FakeLazyLayerwiseDit(num_layers=4) 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 6496d31f759a..f465a6c5b3f1 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -40,6 +40,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, LayerwiseOffloadManager, + LayerwiseUsageTracker, compute_streamed_layers, configure_layerwise_offload_modules, get_layerwise_offload_component_names_for_pipeline, @@ -53,6 +54,36 @@ from sglang.multimodal_gen.runtime.server_args import ServerArgs +class _UsageTrackedVAE(torch.nn.Module, LayerwiseOffloadableModuleMixin): + layer_names = ["encoder", "decoder"] + + def __init__(self) -> None: + super().__init__() + self.encoder = torch.nn.ModuleList(torch.nn.Identity() for _ in range(2)) + self.decoder = torch.nn.ModuleList(torch.nn.Identity() for _ in range(3)) + + def forward(self, value): + for layer in self.decoder: + value = layer(value) + return value + + +def test_layerwise_usage_tracker_observes_only_executed_groups(): + module = _UsageTrackedVAE() + tracker = LayerwiseUsageTracker({"vae": module}) + + module(torch.zeros(1)) + uses = tracker.finish() + + assert uses == { + "vae": { + "encoder": (0, 0), + "decoder": (1, 1, 1), + } + } + assert all(not layer._forward_pre_hooks for layer in module.decoder) + + class _FakeStream: def wait_stream(self, _stream) -> None: return None From d3878ad81ecd2e23edf23a66250ed66895be69c6 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 17:47:25 +0800 Subject: [PATCH 57/90] [diffusion] scope fixed loading paths per component --- .../runtime/managers/gpu_worker.py | 9 +++- .../runtime/server_args/auto_tune.py | 46 ++++++++++++------- .../test/unit/test_auto_residency.py | 42 +++++++++++++++-- 3 files changed, 75 insertions(+), 22 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 4703ddb6c48f..16dbaeb0f261 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -118,6 +118,9 @@ ) from sglang.multimodal_gen.runtime.realtime.session import RealtimeSessionCache from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs +from sglang.multimodal_gen.runtime.server_args.auto_tune import ( + fixed_loading_residency_components, +) from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch, set_musa_arch from sglang.multimodal_gen.runtime.utils.logging_utils import ( configure_logger, @@ -1763,11 +1766,15 @@ def _mixed_dtype_residency_components() -> set[str]: def _fixed_custom_residency_strategy_names(self) -> set[str]: if self.pipeline is None: return set() - return { + fixed = { name for name, strategy in self.pipeline.component_residency_strategies.items() if not strategy.supports_auto_residency() } + fixed.update( + fixed_loading_residency_components(self.server_args, self.pipeline.modules) + ) + return fixed def _latest_auto_residency_round(self) -> list[AppliedResidencyChange]: if not self._auto_residency_round_sizes: diff --git a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py index 80045e77c13c..1d8d4f1c2956 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Iterable from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import ( @@ -23,6 +23,7 @@ LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP, LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP, LAYERWISE_OFFLOAD_VAE_GROUP, + is_dit_component_name, normalize_layerwise_offload_components, ) from sglang.multimodal_gen.runtime.platforms import current_platform @@ -76,27 +77,38 @@ def auto_residency_args_skip_reason(server_args: ServerArgs) -> str | None: return "cache-dit enabled" if server_args.batching_max_size > 1: return "dynamic batching enabled" - nunchaku = server_args.nunchaku_config - svdquant_enabled = nunchaku is not None and ( - not isinstance(nunchaku, NunchakuSVDQuantArgs) or nunchaku.enable_svdquant - ) - if ( - server_args.quantization is not None - or server_args.component_quantizations - or server_args.transformer_weights_path - or svdquant_enabled - ): - # Explicit and online quantization require a fixed loading path. - # Self-describing pre-quantized checkpoints are safe to move after - # loading: residency changes storage location, not quantized values. - return "quantized checkpoint" - if server_args.direct_gpu_weight_loading: - return "direct GPU weight loading requires a fixed resident placement" if not current_platform.is_cuda(): return "requires CUDA" return None +def fixed_loading_residency_components( + server_args: ServerArgs, component_names: Iterable[str] +) -> set[str]: + """Components whose loader-selected storage must remain unchanged. + + Explicit/online quantization and direct GPU loading constrain their target + DiTs, not the unrelated encoders and VAEs in the same pipeline. Keep those + components outside auto residency while still calibrating the rest. + """ + fixed = set(server_args.component_quantizations) + nunchaku = server_args.nunchaku_config + fixed_dit_loading = ( + server_args.quantization is not None + or server_args.direct_gpu_weight_loading + or ( + nunchaku is not None + and ( + not isinstance(nunchaku, NunchakuSVDQuantArgs) + or nunchaku.enable_svdquant + ) + ) + ) + if fixed_dit_loading: + fixed.update(name for name in component_names if is_dit_component_name(name)) + return fixed + + class ServerArgsAutoTuner: """Auto-tunes the server-arg for the given performance-mode, based on practical deployment experience with different model architectures""" diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 472645fe8954..6717a08c0b29 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -3135,10 +3135,6 @@ def test_manual_performance_mode(self, monkeypatch): # resident aux components on CPU): its peaks are not serving peaks ({"enable_torch_compile": True}, "stripped memory layout"), ({"batching_max_size": 4}, "batching"), - ({"quantization": "fp8"}, "quantized"), - ({"component_quantizations": {"transformer": "fp8"}}, "quantized"), - ({"transformer_weights_path": "/x.safetensors"}, "quantized"), - ({"direct_gpu_weight_loading": True}, "direct GPU"), ], ) def test_excluded_paths(self, monkeypatch, overrides, expected_fragment): @@ -3147,6 +3143,44 @@ def test_excluded_paths(self, monkeypatch, overrides, expected_fragment): reason = self._skip_reason(self._base_args(**overrides)) assert reason is not None and expected_fragment in reason + @pytest.mark.parametrize( + "overrides", + [ + {"quantization": "fp8"}, + {"component_quantizations": {"image_encoder": "fp8"}}, + {"transformer_weights_path": "/x.safetensors"}, + {"direct_gpu_weight_loading": True}, + ], + ) + def test_fixed_loading_paths_still_calibrate_other_components( + self, monkeypatch, overrides + ): + monkeypatch.delenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", raising=False) + monkeypatch.delenv("SGLANG_CACHE_DIT_ENABLED", raising=False) + reason = self._skip_reason(self._base_args(**overrides)) + assert reason is None or reason == "requires CUDA" + + def test_fixed_loading_components_are_excluded_individually(self): + from sglang.multimodal_gen.runtime.server_args.auto_tune import ( + fixed_loading_residency_components, + ) + + component_names = ("transformer", "transformer_2", "image_encoder", "vae") + assert fixed_loading_residency_components( + self._base_args(quantization="fp8"), component_names + ) == {"transformer", "transformer_2"} + assert fixed_loading_residency_components( + self._base_args(component_quantizations={"image_encoder": "fp8"}), + component_names, + ) == {"image_encoder"} + assert ( + fixed_loading_residency_components( + self._base_args(transformer_weights_path="/x.safetensors"), + component_names, + ) + == set() + ) + def test_cache_dit_excluded(self, monkeypatch): monkeypatch.delenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", raising=False) monkeypatch.setenv("SGLANG_CACHE_DIT_ENABLED", "true") From c43a81fccfd69bab9ce57fd00b90da126434f08e Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 17:51:46 +0800 Subject: [PATCH 58/90] [docs] explain calibrated layer-group placement --- docs/docs/sglang-diffusion/api/cli.mdx | 2 +- docs/docs/sglang-diffusion/deployment_cookbook.mdx | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index ace17a9211bb..79bb138c94a8 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -85,7 +85,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. Multi-process HostPin planning stays inside non-overlapping per-worker shares. 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. 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 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. 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. - `--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. Startup impact is model-dependent, so benchmark the target model before deployment. Disabled by default because checkpoint and model weights coexist temporarily, substantially increasing peak GPU memory. 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 diff --git a/docs/docs/sglang-diffusion/deployment_cookbook.mdx b/docs/docs/sglang-diffusion/deployment_cookbook.mdx index 4715180a9569..6e36107bf238 100644 --- a/docs/docs/sglang-diffusion/deployment_cookbook.mdx +++ b/docs/docs/sglang-diffusion/deployment_cookbook.mdx @@ -128,7 +128,7 @@ See [OpenAI API: Served model name](/docs/sglang-diffusion/api/openai_api#served -With an eligible `sglang serve` warmup, `auto` records allocated-memory peaks for the representative request and their active component layouts. 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, layerwise resident counts, and pinned layer-store reassignment. A post-adjustment warmup realizes the selected state and validates its memory peak before the server becomes ready. Auxiliary-only moves to resident use one full-shape denoising step; DiT changes, 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 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. Auxiliary-only moves to resident use one full-shape denoising step; DiT changes, 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. From 5eae402f0247df771236d84151caebdc49d77e6e Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 18:30:20 +0800 Subject: [PATCH 59/90] [diffusion] infer repeated layer usage from warmup stages --- .../memory_managers/auto_residency.py | 18 +++++++++++++- .../test/unit/test_auto_residency.py | 24 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index ed084ca5d824..c551a133f5ad 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -356,6 +356,19 @@ def estimate_layerwise_layer_uses( estimated: dict[str, dict[str, list[int]]] = {} for record in successful: source_steps = max(1, record.num_inference_steps) + denoising_components: set[str] = set() + phase_components = ( + record.phase_used_components or record.phase_active_components + ) + for phase_name, component_names in phase_components.items(): + fields = phase_name.split(":", 2) + if len(fields) < 2 or not fields[0].isdigit(): + continue + stage_name = fields[1] + if stage_name.endswith("DenoisingStage") and not stage_name.endswith( + "BeforeDenoisingStage" + ): + denoising_components.update(component_names) for component_name, groups in record.layerwise_layer_uses.items(): component = estimated.setdefault(component_name, {}) for layer_name, counts in groups.items(): @@ -365,7 +378,10 @@ def estimate_layerwise_layer_uses( for layer_index, count in enumerate(counts): scaled = count if ( - is_dit_component_name(component_name) + ( + is_dit_component_name(component_name) + or component_name in denoising_components + ) and count > 1 and target_num_inference_steps > source_steps ): diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 6717a08c0b29..3eaa67dd62f5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -68,6 +68,7 @@ def _record( stage_duration_ms=None, step_duration_ms=(), phase_active_components=None, + phase_used_components=None, phase_full_weight_transition_components=None, layerwise_layer_uses=None, ) -> WarmupMemoryRecord: @@ -83,6 +84,7 @@ def _record( stage_duration_ms=stage_duration_ms or {}, step_duration_ms=step_duration_ms, phase_active_components=phase_active_components or {}, + phase_used_components=phase_used_components or {}, phase_full_weight_transition_components=( phase_full_weight_transition_components or {} ), @@ -229,6 +231,28 @@ def test_scales_repeated_groups_but_not_one_shot_groups(self): "decoder.up_blocks": (1, 1), } + def test_stage_role_scales_nonstandard_denoiser_but_not_encoder_repeats(self): + record = _record( + num_inference_steps=2, + phase_used_components={ + "0:TextEncodingStage:use:custom_encoder": ("custom_encoder",), + "1:CustomDenoisingStage:use:custom_refiner": ("custom_refiner",), + }, + layerwise_layer_uses={ + "custom_encoder": {"layers": (2, 2)}, + "custom_refiner": {"blocks": (2, 2)}, + }, + ) + + uses = estimate_layerwise_layer_uses( + records=[record], + target_units=record.workload_units(), + target_num_inference_steps=10, + ) + + assert uses["custom_encoder"]["layers"] == (2, 2) + assert uses["custom_refiner"]["blocks"] == (10, 10) + class TestEstimateDefaultWorkloadPeak: def test_same_shape_uses_measured_peak(self): From 6b005b765759690d77946a08f62f4a7e9386f093 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 18:54:04 +0800 Subject: [PATCH 60/90] [diffusion] scale all measured denoising stages --- .../memory_managers/auto_residency.py | 13 ++++----- .../test/unit/test_auto_residency.py | 29 +++++++++++++++++++ 2 files changed, 35 insertions(+), 7 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index c551a133f5ad..58f57eaa9ef3 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -449,13 +449,12 @@ def estimate_default_workload_timing( if is_dit_component_name(component_name) for stage_name in stage_names } - if not denoising_stages: - denoising_stages = { - stage_name - for stage_name in representative.stage_duration_ms - if stage_name.endswith("DenoisingStage") - and not stage_name.endswith("BeforeDenoisingStage") - } + denoising_stages.update( + stage_name + for stage_name in representative.stage_duration_ms + if stage_name.endswith("DenoisingStage") + and not stage_name.endswith("BeforeDenoisingStage") + ) measured_steps = max(1, representative.num_inference_steps) target_steps = max(1, target_num_inference_steps) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 3eaa67dd62f5..476a28bd6d91 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -147,6 +147,35 @@ def test_uses_steady_step_instead_of_scaling_first_step_setup(self): assert stage_ns["DenoisingStage"] == 1_020_000_000 assert request_ns == 1_120_000_000 + def test_scales_standard_and_nonstandard_denoising_stages_together(self): + record = _record( + num_inference_steps=2, + total_duration_ms=1_200, + stage_duration_ms={ + "TextEncodingStage": 100, + "DenoisingStage": 500, + "CustomDenoisingStage": 600, + }, + phase_active_components={ + "0:TextEncodingStage:use:text_encoder": ("text_encoder",), + "1:DenoisingStage:use:transformer": ("transformer",), + "2:CustomDenoisingStage:use:custom_refiner": ("custom_refiner",), + }, + ) + + request_ns, stage_ns, _ = estimate_default_workload_timing( + records=[record], + target_units=record.workload_units(), + target_num_inference_steps=10, + ) + + assert stage_ns == { + "TextEncodingStage": 100_000_000, + "DenoisingStage": 2_500_000_000, + "CustomDenoisingStage": 3_000_000_000, + } + assert request_ns == 5_600_000_000 + def test_transfer_savings_is_capped_by_request_not_component_stage(self): candidates = [ _candidate("text_encoder", weight_gib=10, h2d_gib=10), From 055b9b7b94094da7b9d441ebbddc3c76d14f4597 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 19:52:02 +0800 Subject: [PATCH 61/90] [diffusion] calibrate residency by stage workload --- .../runtime/managers/gpu_worker.py | 56 +++- .../memory_managers/auto_residency.py | 294 ++++++++++++------ .../memory_managers/layerwise_offload.py | 49 ++- .../memory_managers/placement_budget.py | 43 +++ .../runtime/pipelines_core/schedule_batch.py | 26 ++ .../runtime/pipelines_core/stages/base.py | 9 + .../pipelines_core/stages/causal_denoising.py | 3 + .../pipelines_core/stages/denoising.py | 1 + .../pipelines_core/stages/denoising_dmd.py | 1 + .../stages/model_specific_stages/cosmos3.py | 1 + .../model_specific_stages/helios_denoising.py | 4 + .../model_specific_stages/hunyuan3d/paint.py | 43 ++- .../lingbot_world_causal_denoising.py | 1 + .../stages/model_specific_stages/longlive2.py | 3 + .../ltx_2/denoising_av.py | 10 + .../minimax_h3/stages/denoising.py | 1 + .../stages/model_specific_stages/mova.py | 1 + .../model_specific_stages/sana_wm/base.py | 1 + .../model_specific_stages/sana_wm/refiner.py | 2 + .../sana_wm/streaming.py | 2 + .../sana_wm/streaming_refiner.py | 2 + .../progressive_resolution/denoising.py | 1 + .../runtime/utils/perf_logger.py | 24 +- .../test/unit/test_auto_residency.py | 119 +++++++ .../test/unit/test_auto_residency_warmup.py | 9 +- .../test_hunyuan3d_native_texture_models.py | 19 ++ .../test/unit/test_layerwise_offload.py | 17 + .../test/unit/test_performance_metrics.py | 29 ++ .../test/unit/test_placement_budget.py | 41 +++ 29 files changed, 696 insertions(+), 116 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 16dbaeb0f261..f777f200097e 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -595,6 +595,9 @@ def _execute_forward_common( ) warmup_baseline_allocated_bytes = 0 layerwise_usage_tracker: LayerwiseUsageTracker | None = None + layerwise_layer_uses_by_stage: dict[ + str, dict[str, dict[str, tuple[int, ...]]] + ] = {} try: if measure_server_warmup: # Drop the previous request's allocator pool so each probe @@ -612,7 +615,14 @@ def _execute_forward_common( and self.pipeline is not None ): layerwise_usage_tracker = LayerwiseUsageTracker( - self.pipeline.modules + self.pipeline.modules, + stage_name_provider=( + lambda: ( + req.metrics.active_stage_name + if req.metrics is not None + else None + ) + ), ) start_time = ( @@ -731,17 +741,20 @@ def _execute_forward_common( # or the estimator would plan from the remaining partial data if measure_server_warmup: assert warmup_workload is not None - layerwise_layer_uses = ( - layerwise_usage_tracker.finish() - if layerwise_usage_tracker is not None - else {} - ) + if layerwise_usage_tracker is not None: + ( + layerwise_layer_uses, + layerwise_layer_uses_by_stage, + ) = layerwise_usage_tracker.finish_with_stages() + else: + layerwise_layer_uses = {} self._record_server_warmup_memory( req=req, workload=warmup_workload, baseline_allocated_bytes=warmup_baseline_allocated_bytes, succeeded=output_batch is not None and output_batch.error is None, layerwise_layer_uses=layerwise_layer_uses, + layerwise_layer_uses_by_stage=layerwise_layer_uses_by_stage, ) return output_batch @@ -753,6 +766,9 @@ def _record_server_warmup_memory( baseline_allocated_bytes: int, succeeded: bool, layerwise_layer_uses: dict[str, dict[str, tuple[int, ...]]] | None = None, + layerwise_layer_uses_by_stage: ( + dict[str, dict[str, dict[str, tuple[int, ...]]]] | None + ) = None, ) -> None: phase_allocated_peaks: dict[str, int] = {} phase_components: dict[str, tuple[str, ...]] = {} @@ -799,12 +815,24 @@ def _record_server_warmup_memory( phase_full_weight_transition_components ), layerwise_layer_uses=layerwise_layer_uses or {}, + layerwise_layer_uses_by_stage=layerwise_layer_uses_by_stage or {}, num_inference_steps=num_inference_steps, total_duration_ms=( float(metrics.total_duration_ms) if metrics is not None else 0.0 ), stage_duration_ms=(dict(metrics.stages) if metrics is not None else {}), step_duration_ms=(tuple(metrics.steps) if metrics is not None else ()), + step_duration_ms_by_stage=( + { + stage_name: tuple(durations) + for stage_name, durations in metrics.steps_by_stage.items() + } + if metrics is not None + else {} + ), + stage_iterations=( + dict(metrics.stage_iterations) if metrics is not None else {} + ), ) ) @@ -1646,6 +1674,13 @@ def _build_auto_residency_report( test_cap_gib = envs.SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB if test_cap_gib is not None: budget_bytes = min(budget_bytes, int(test_cap_gib * GIB_BYTES)) + local_worker_count = max( + 1, self.server_args.num_gpus // self.server_args.nnodes + ) + host_transition_headroom_bytes = ( + max(0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES) + // local_worker_count + ) candidates = [] if include_candidates: candidates = collect_residency_targets( @@ -1670,6 +1705,7 @@ def _build_auto_residency_report( measured_used_components if has_component_use_measurement else None ), layerwise_layer_uses=layerwise_layer_uses, + host_transition_headroom_bytes=host_transition_headroom_bytes, ) measured_request_duration_ns, _, _ = estimate_default_workload_timing( records=records, @@ -1695,9 +1731,6 @@ def _build_auto_residency_report( if include_candidates else {} ) - local_worker_count = max( - 1, self.server_args.num_gpus // self.server_args.nnodes - ) return RankResidencyReport( rank=self.rank, budget_bytes=budget_bytes, @@ -1729,10 +1762,7 @@ def _build_auto_residency_report( else 0 ), host_transition_headroom_bytes=( - max(0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES) - // local_worker_count - if include_candidates - else 0 + host_transition_headroom_bytes if include_candidates else 0 ), device_transition_allocated_bytes=( int(torch.get_device_module().memory_allocated()) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 58f57eaa9ef3..48a972df2615 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -135,10 +135,13 @@ class WarmupMemoryRecord(msgspec.Struct, frozen=True): phase_used_components: dict[str, tuple[str, ...]] = {} phase_full_weight_transition_components: dict[str, tuple[str, ...]] = {} layerwise_layer_uses: dict[str, dict[str, tuple[int, ...]]] = {} + layerwise_layer_uses_by_stage: dict[str, dict[str, dict[str, tuple[int, ...]]]] = {} num_inference_steps: int = 1 total_duration_ms: float = 0.0 stage_duration_ms: dict[str, float] = {} step_duration_ms: tuple[float, ...] = () + step_duration_ms_by_stage: dict[str, tuple[float, ...]] = {} + stage_iterations: dict[str, tuple[int, int]] = {} def workload_units(self) -> int: return max(1, self.width) * max(1, self.height) * max(1, self.num_frames) @@ -340,10 +343,11 @@ def estimate_layerwise_layer_uses( ) -> dict[str, dict[str, tuple[int, ...]]]: """Estimate per-request layer calls from the same calibration forward. - A full-shape memory probe deliberately runs only a few denoise steps. A - layer observed at least twice is treated as step-scoped and scaled to the - target step count; a layer observed once remains a one-shot encoder, - refiner, or decoder layer even when it belongs to a DiT component. + A full-shape memory probe deliberately runs only a few denoise steps. + Stage-attributed calls use that stage's measured and target iteration + counts, so independent shape, paint, refiner, and chunk loops are not all + multiplied by one request-wide ratio. Legacy records retain the repeated + DiT-layer heuristic. """ successful = [record for record in records if record.succeeded] if target_units is not None: @@ -356,38 +360,65 @@ def estimate_layerwise_layer_uses( estimated: dict[str, dict[str, list[int]]] = {} for record in successful: source_steps = max(1, record.num_inference_steps) - denoising_components: set[str] = set() - phase_components = ( - record.phase_used_components or record.phase_active_components - ) - for phase_name, component_names in phase_components.items(): - fields = phase_name.split(":", 2) - if len(fields) < 2 or not fields[0].isdigit(): - continue - stage_name = fields[1] - if stage_name.endswith("DenoisingStage") and not stage_name.endswith( - "BeforeDenoisingStage" - ): - denoising_components.update(component_names) + component_stages = _component_stages(record) + repeated_stages = _repeated_stages(record, component_stages) for component_name, groups in record.layerwise_layer_uses.items(): component = estimated.setdefault(component_name, {}) for layer_name, counts in groups.items(): target = component.setdefault(layer_name, [0] * len(counts)) if len(target) != len(counts): continue + stage_counts = [ + ( + stage_name, + stage_components[component_name][layer_name], + ) + for stage_name, stage_components in ( + record.layerwise_layer_uses_by_stage.items() + ) + if component_name in stage_components + and layer_name in stage_components[component_name] + and len(stage_components[component_name][layer_name]) == len(counts) + ] for layer_index, count in enumerate(counts): - scaled = count - if ( - ( - is_dit_component_name(component_name) - or component_name in denoising_components + if stage_counts: + measured_by_stage = sum( + per_layer_counts[layer_index] + for _, per_layer_counts in stage_counts ) - and count > 1 - and target_num_inference_steps > source_steps - ): - scaled = ( - count * target_num_inference_steps + source_steps - 1 - ) // source_steps + scaled = max(0, count - measured_by_stage) + for stage_name, per_layer_counts in stage_counts: + measured_iterations, target_iterations = _stage_iterations( + record, + stage_name, + repeated_stages=repeated_stages, + target_num_inference_steps=(target_num_inference_steps), + ) + stage_count = per_layer_counts[layer_index] + if stage_count <= 1: + scaled += stage_count + else: + scaled += ( + stage_count * target_iterations + + measured_iterations + - 1 + ) // measured_iterations + else: + scaled = count + component_is_repeated = is_dit_component_name( + component_name + ) or any( + stage_name in repeated_stages + for stage_name in component_stages.get(component_name, ()) + ) + if ( + component_is_repeated + and count > 1 + and target_num_inference_steps > source_steps + ): + scaled = ( + count * target_num_inference_steps + source_steps - 1 + ) // source_steps target[layer_index] = max(target[layer_index], scaled) return { component_name: { @@ -397,6 +428,64 @@ def estimate_layerwise_layer_uses( } +def _component_stages( + record: WarmupMemoryRecord, + *, + timed_stage_names: set[str] | None = None, +) -> dict[str, set[str]]: + component_stages: dict[str, set[str]] = {} + phase_components = record.phase_used_components or record.phase_active_components + for phase_name, components in phase_components.items(): + fields = phase_name.split(":", 2) + if len(fields) < 2 or not fields[0].isdigit(): + continue + stage_name = fields[1] + if timed_stage_names is not None and stage_name not in timed_stage_names: + continue + for component_name in components: + component_stages.setdefault(component_name, set()).add(stage_name) + return component_stages + + +def _repeated_stages( + record: WarmupMemoryRecord, + component_stages: Mapping[str, set[str]], +) -> set[str]: + stages = { + stage_name + for component_name, stage_names in component_stages.items() + if is_dit_component_name(component_name) + for stage_name in stage_names + } + stages.update(record.stage_iterations) + stages.update( + stage_name + for stage_name in record.stage_duration_ms + if stage_name.endswith("DenoisingStage") + and not stage_name.endswith("BeforeDenoisingStage") + ) + return stages + + +def _stage_iterations( + record: WarmupMemoryRecord, + stage_name: str, + *, + repeated_stages: set[str], + target_num_inference_steps: int, +) -> tuple[int, int]: + explicit = record.stage_iterations.get(stage_name) + if explicit is not None: + return max(1, explicit[0]), max(1, explicit[1]) + measured = max(1, record.num_inference_steps) + target = ( + max(1, target_num_inference_steps) + if stage_name in repeated_stages + else measured + ) + return measured, target + + def estimate_default_workload_timing( *, records: Iterable[WarmupMemoryRecord], @@ -405,10 +494,10 @@ def estimate_default_workload_timing( ) -> tuple[int, dict[str, int], dict[str, tuple[str, ...]]]: """Estimate full-request and stage durations from the warmup workload. - Non-denoise stages run once. Denoising time scales with the requested step - count; the full-shape probe intentionally executes only a few steps, so - using its raw total would make every one-shot encoder transfer look - important relative to a long video request. + Repeated stages scale by their own measured/default iteration counts. The + full-shape probe intentionally executes only a few steps, so using its raw + total would make every one-shot encoder transfer look important relative + to a long video request. """ successful = [record for record in records if record.succeeded] if not successful: @@ -429,57 +518,34 @@ def estimate_default_workload_timing( if representative.total_duration_ms <= 0 or not representative.stage_duration_ms: return 0, {}, {} - component_stages: dict[str, set[str]] = {} - phase_components = ( - representative.phase_used_components or representative.phase_active_components - ) - for phase_name, components in phase_components.items(): - fields = phase_name.split(":", 2) - if len(fields) < 2 or not fields[0].isdigit(): - continue - stage_name = fields[1] - if stage_name not in representative.stage_duration_ms: - continue - for component_name in components: - component_stages.setdefault(component_name, set()).add(stage_name) - - denoising_stages = { - stage_name - for component_name, stage_names in component_stages.items() - if is_dit_component_name(component_name) - for stage_name in stage_names - } - denoising_stages.update( - stage_name - for stage_name in representative.stage_duration_ms - if stage_name.endswith("DenoisingStage") - and not stage_name.endswith("BeforeDenoisingStage") + component_stages = _component_stages( + representative, + timed_stage_names=set(representative.stage_duration_ms), ) - - measured_steps = max(1, representative.num_inference_steps) - target_steps = max(1, target_num_inference_steps) - denoising_scale = target_steps / measured_steps - measured_denoising_ms = sum( - representative.stage_duration_ms[stage_name] for stage_name in denoising_stages - ) - if measured_denoising_ms > 0 and representative.step_duration_ms: - steady_steps = ( - representative.step_duration_ms[1:] - if len(representative.step_duration_ms) > 1 - else representative.step_duration_ms - ) - steady_step_ms = statistics.median(steady_steps) - non_step_denoising_ms = max( - 0.0, - measured_denoising_ms - sum(representative.step_duration_ms), - ) - target_denoising_ms = non_step_denoising_ms + steady_step_ms * target_steps - denoising_scale = target_denoising_ms / measured_denoising_ms + repeated_stages = _repeated_stages(representative, component_stages) stage_duration_ns: dict[str, int] = {} for stage_name, duration_ms in representative.stage_duration_ms.items(): - scale = denoising_scale if stage_name in denoising_stages else 1.0 - stage_duration_ns[stage_name] = max(0, int(duration_ms * scale * 1_000_000)) + measured_iterations, target_iterations = _stage_iterations( + representative, + stage_name, + repeated_stages=repeated_stages, + target_num_inference_steps=target_num_inference_steps, + ) + step_durations = representative.step_duration_ms_by_stage.get(stage_name, ()) + if not step_durations and len(repeated_stages) == 1: + step_durations = representative.step_duration_ms + if stage_name in repeated_stages and step_durations: + steady_steps = ( + step_durations[1:] if len(step_durations) > 1 else step_durations + ) + non_step_ms = max(0.0, duration_ms - sum(step_durations)) + target_duration_ms = non_step_ms + ( + statistics.median(steady_steps) * target_iterations + ) + else: + target_duration_ms = duration_ms * target_iterations / measured_iterations + stage_duration_ns[stage_name] = max(0, int(target_duration_ms * 1_000_000)) measured_stage_ms = sum(representative.stage_duration_ms.values()) untracked_ms = max(0.0, representative.total_duration_ms - measured_stage_ms) @@ -950,6 +1016,7 @@ def _layerwise_pin_targets( current_pinned_layers: tuple[tuple[int, ...], ...], uses_per_streamed_layer: int, layer_uses: tuple[tuple[int, ...], ...] | None = None, + constrain_host_transitions: bool = True, ) -> list[tuple[tuple[int, ...], ...]]: """Pareto-optimal HostPin targets for one resident-layer placement. @@ -1010,6 +1077,38 @@ def dominates(left, right) -> bool: ) def prune(candidates): + if not constrain_host_transitions: + # host transition scratch cannot constrain this solve, so the local + # frontier is exactly pinned bytes versus avoided pageable transfer + best_by_pinned_bytes = {} + for state in candidates: + incumbent = best_by_pinned_bytes.get(state[0]) + if incumbent is None or ( + -state[1], + state[2] + state[3], + state[2], + state[3], + state[4], + ) < ( + -incumbent[1], + incumbent[2] + incumbent[3], + incumbent[2], + incumbent[3], + incumbent[4], + ): + best_by_pinned_bytes[state[0]] = state + frontier = [] + best_utility = -1 + for state in sorted( + best_by_pinned_bytes.values(), + key=lambda item: (item[0], -item[1], item[4]), + ): + if state[1] <= best_utility: + continue + frontier.append(state) + best_utility = state[1] + return frontier + best_by_cost = {} for state in candidates: key = state[:4] @@ -1295,11 +1394,6 @@ def _unconfigured_layerwise_targets( coarse_savings = max(0, maximum_transfer_work - coarse_transfer_work) current_resident = current_mode == RESIDENT current_inactive_bytes = full_weight_bytes if current_resident else 0 - initial_managed_bytes = _layerwise_active_peak_device_bytes( - managers=managers, - resident_layers=tuple(0 for _ in managers), - layer_uses=layer_uses, - ) host_materialize_scratch = max( ( weight_bytes @@ -1344,9 +1438,7 @@ def _unconfigured_layerwise_targets( target_layerwise_pinned_layers=empty_pins, host_materialize_scratch_bytes=host_materialize_scratch, device_transition_delta_bytes=( - unmanaged_weight_bytes - + initial_managed_bytes - - current_inactive_bytes + unmanaged_weight_bytes - current_inactive_bytes ), active_device_delta_bytes=( unmanaged_weight_bytes + active_managed_bytes - full_weight_bytes @@ -1398,6 +1490,7 @@ def collect_residency_targets( pin_cpu_memory: bool = True, used_components: Iterable[str] | None = None, layerwise_layer_uses: Mapping[str, Mapping[str, tuple[int, ...]]] | None = None, + host_transition_headroom_bytes: int | None = None, ) -> list[ResidencyTarget]: """Build complete target-state frontiers for auto-managed components. @@ -1412,6 +1505,32 @@ def collect_residency_targets( measured_used_names = set(used_components) if used_components is not None else None if baseline_residency_mode_of is None: baseline_residency_mode_of = residency_mode_of + constrain_host_transitions = True + if host_transition_headroom_bytes is not None: + seen_managers: set[int] = set() + maximum_unpin_bytes = 0 + maximum_pin_bytes = 0 + for module in modules.values(): + if not isinstance(module, LayerwiseOffloadableModuleMixin): + continue + for manager in module.layerwise_offload_managers: + manager_id = id(manager) + if manager_id in seen_managers: + continue + seen_managers.add(manager_id) + pinned = set(manager.pinned_layer_indices()) + host_bytes = manager.layer_host_store_bytes() + maximum_unpin_bytes += sum( + host_bytes.get(layer_idx, 0) for layer_idx in pinned + ) + maximum_pin_bytes += sum( + host_bytes.get(layer_idx, 0) + for layer_idx in manager.pinnable_layer_indices() + if layer_idx not in pinned + ) + constrain_host_transitions = max(maximum_unpin_bytes, maximum_pin_bytes) > max( + 0, host_transition_headroom_bytes + ) candidates = [] for name in sorted(modules): module = modules[name] @@ -1662,6 +1781,7 @@ def collect_residency_targets( current_pinned_layers=current_pinned_layers, uses_per_streamed_layer=uses_per_request, layer_uses=layer_uses, + constrain_host_transitions=constrain_host_transitions, ) if allow_host_pin_reallocation else [current_pinned_layers] 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 f53614d5000b..a0bac3631678 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 @@ -7,7 +7,7 @@ from collections.abc import Iterable, Mapping, Sequence from contextlib import nullcontext from time import perf_counter -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Any, Callable, Dict, List, Optional, Set, Tuple import torch from torch.distributed.tensor import DTensor @@ -2366,9 +2366,16 @@ class LayerwiseUsageTracker: ordinary serving requests pay no counter or hook-dispatch cost. """ - def __init__(self, modules: Mapping[str, object]) -> None: + def __init__( + self, + modules: Mapping[str, object], + *, + stage_name_provider: Callable[[], str | None] | None = None, + ) -> None: self._handles: list[Any] = [] self._counts: dict[str, dict[str, list[int]]] = {} + self._counts_by_stage: dict[str, dict[str, dict[str, list[int]]]] = {} + self._stage_name_provider = stage_name_provider for component_name, module in modules.items(): if not isinstance(module, LayerwiseOffloadableModuleMixin): continue @@ -2390,24 +2397,60 @@ def record_use( *, target_counts=counts, target_index=layer_index, + target_component_name=component_name, + target_layer_name=layer_name, + target_layer_count=len(layers), ) -> None: target_counts[target_index] += 1 + if self._stage_name_provider is None: + return + stage_name = self._stage_name_provider() + if stage_name is None: + return + stage_counts = self._counts_by_stage.setdefault(stage_name, {}) + component_stage_counts = stage_counts.setdefault( + target_component_name, {} + ) + layer_stage_counts = component_stage_counts.setdefault( + target_layer_name, [0] * target_layer_count + ) + layer_stage_counts[target_index] += 1 self._handles.append(layer.register_forward_pre_hook(record_use)) if component_counts: self._counts[component_name] = component_counts def finish(self) -> dict[str, dict[str, tuple[int, ...]]]: + counts, _ = self.finish_with_stages() + return counts + + def finish_with_stages( + self, + ) -> tuple[ + dict[str, dict[str, tuple[int, ...]]], + dict[str, dict[str, dict[str, tuple[int, ...]]]], + ]: for handle in self._handles: handle.remove() self._handles.clear() - return { + counts = { component_name: { layer_name: tuple(counts) for layer_name, counts in component_counts.items() } for component_name, component_counts in self._counts.items() } + counts_by_stage = { + stage_name: { + component_name: { + layer_name: tuple(counts) + for layer_name, counts in component_counts.items() + } + for component_name, component_counts in stage_counts.items() + } + for stage_name, stage_counts in self._counts_by_stage.items() + } + return counts, counts_by_stage def iter_materialized_weights(module: torch.nn.Module): diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py index 4f1e3fe08560..55e12adbeddc 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py @@ -93,6 +93,49 @@ def _pareto_prune(states: Iterable[_State]) -> list[_State]: _selection_key(state), ), ) + if ordered and not ordered[0][0] and len(ordered[0][2]) == 2: + # with no hard resource dimension, local dominance is a 2D soft-cost + # skyline; a Fenwick sweep avoids the quadratic layer frontier scan + second_costs = sorted({state[2][1] for state in ordered}) + second_cost_indices = { + cost: index + 1 for index, cost in enumerate(second_costs) + } + best_utility: list[int | None] = [None] * (len(second_costs) + 1) + + def prefix_max(index: int) -> int | None: + result = None + while index > 0: + value = best_utility[index] + if value is not None and (result is None or value > result): + result = value + index -= index & -index + return result + + def update(index: int, utility: int) -> None: + while index < len(best_utility): + value = best_utility[index] + if value is None or utility > value: + best_utility[index] = utility + index += index & -index + + frontier = [] + for state in sorted( + ordered, + key=lambda item: ( + item[2][0], + item[2][1], + -item[1], + _selection_key(item), + ), + ): + index = second_cost_indices[state[2][1]] + dominating_utility = prefix_max(index) + if dominating_utility is not None and dominating_utility >= state[1]: + continue + frontier.append(state) + update(index, state[1]) + return frontier + frontier: list[_State] = [] for state in ordered: if any(_dominates(existing, state) for existing in frontier): diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index bc36f0d964d1..1ef848058f8e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -337,6 +337,7 @@ def set_as_warmup(self, warmup_steps: int = 1): self.suppress_logs = True self.metrics.suppress_stage_breakdown = True self.extra["cache_dit_num_inference_steps"] = self.num_inference_steps + self.extra["warmup_target_num_inference_steps"] = self.num_inference_steps self.num_inference_steps = warmup_steps def copy_as_warmup(self, warmup_steps: int = 1) -> Req: @@ -344,6 +345,31 @@ def copy_as_warmup(self, warmup_steps: int = 1) -> Req: req.set_as_warmup(warmup_steps) return req + def record_stage_iterations( + self, + measured_iterations: int, + target_iterations: int | None = None, + ) -> None: + """Record a stage loop against its full default-request work.""" + if not self.is_warmup or self.metrics is None: + return + measured = max(1, int(measured_iterations)) + if target_iterations is None: + measured_request_steps = max(1, int(self.num_inference_steps)) + target_request_steps = int( + self.extra.get( + "stage_target_num_inference_steps", + self.extra.get( + "warmup_target_num_inference_steps", + measured_request_steps, + ), + ) + ) + target_iterations = ( + measured * max(1, target_request_steps) + measured_request_steps - 1 + ) // measured_request_steps + self.metrics.record_stage_iterations(measured, target_iterations) + def validate(self): """Initialize dependent fields after dataclass initialization.""" if getattr(self.sampling_params, "data_type", None) == DataType.ACTION: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index 107c28e43690..68f7ce00fa7e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py @@ -384,7 +384,14 @@ def __call__( # Execute the actual stage logic with unified profiling. previous_batch_is_warmup = self._current_batch_is_warmup + metrics = batch.metrics + warmup_metrics = metrics if batch.is_warmup else None + previous_active_stage = ( + warmup_metrics.active_stage_name if warmup_metrics is not None else None + ) self._current_batch_is_warmup = batch.is_warmup + if warmup_metrics is not None: + warmup_metrics.active_stage_name = stage_name try: with StageProfiler( stage_name, @@ -396,6 +403,8 @@ def __call__( ): result = self.forward(batch, server_args) finally: + if warmup_metrics is not None: + warmup_metrics.active_stage_name = previous_active_stage self._current_batch_is_warmup = previous_batch_is_warmup self._current_use_nvtx = False diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py index 0a9db62d7914..5936eaacd575 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py @@ -1253,6 +1253,9 @@ def forward( block_sizes = [1] + [self.num_frames_per_block] * num_blocks start_index = 0 + total_iterations = len(block_sizes) * len(timesteps) + batch.record_stage_iterations(total_iterations, total_iterations) + def prepare_context_input(current_latents): return current_latents diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 24f36549e40b..6acc2e32596e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -1911,6 +1911,7 @@ def _denoise( # to avoid device-sync caused by timestep comparison timesteps_cpu = ctx.timesteps.cpu() num_timesteps = timesteps_cpu.shape[0] + batch.record_stage_iterations(num_timesteps) # Re-resolve the explicit-range gate so the per-step markers # below honor this request's is_warmup state. Layer hooks are # registered by the residency manager at the use-site. diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py index 808c555599f6..ab192b150ab1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py @@ -80,6 +80,7 @@ def forward( dtype=torch.long, device=get_local_torch_device(), ) + batch.record_stage_iterations(len(timesteps), len(timesteps)) # prepare image_kwargs image_embeds = batch.image_embeds diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index ee496fe86984..a22d207ca22e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -1033,6 +1033,7 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: action_fps = getattr(batch.sampling_params, "action_fps", None) timesteps = batch.timesteps guidance_scale = batch.guidance_scale + batch.record_stage_iterations(len(timesteps)) # Seed the scheduler's stochastic (SDE) noise from the request seed so it # is identical on every sequence-parallel rank; otherwise each rank draws diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py index 1a8796e3506c..2e391efb4fb3 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py @@ -748,5 +748,9 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: # separately to avoid temporal artifacts at chunk boundaries. batch.latent_chunks = chunk_latents_list batch.latents = history_latents[:, :, -total_generated_latent_frames:] + batch.record_stage_iterations( + global_step_offset, + global_step_offset if is_enable_stage2 else None, + ) return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py index ca680fc9afb9..49fa729a1110 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py @@ -246,7 +246,7 @@ def _encode_delight_prompts( return torch.cat([positive, negative, negative]) @torch.no_grad() - def _run_delight(self, image: Image.Image) -> Image.Image: + def _run_delight(self, image: Image.Image, batch: Req) -> Image.Image: assert self.delight_transformer is not None assert self.delight_vae is not None assert self.delight_scheduler is not None @@ -278,7 +278,14 @@ def _run_delight(self, image: Image.Image) -> Image.Image: ) scheduler = self.delight_scheduler - scheduler.set_timesteps(self.config.delight_num_inference_steps, device=device) + target_steps = self.config.delight_num_inference_steps + measured_steps = ( + min(target_steps, max(1, int(batch.num_inference_steps))) + if batch.is_warmup + else target_steps + ) + batch.record_stage_iterations(measured_steps, target_steps) + scheduler.set_timesteps(measured_steps, device=device) generator = torch.Generator(device="cpu").manual_seed(42) latent_channels = self.delight_transformer.config.out_channels latents = randn_tensor( @@ -350,11 +357,11 @@ def _run_delight(self, image: Image.Image) -> Image.Image: (composited.clamp(0, 1).cpu().numpy() * 255).astype(np.uint8) ) - def _prepare_reference_image(self, image_path: str) -> Image.Image: + def _prepare_reference_image(self, image_path: str, batch: Req) -> Image.Image: image = self._load_input_image(image_path) if not self.config.delight_enable: return image - return self._run_delight(image) + return self._run_delight(image, batch) def _render_multiview( self, mesh: Any @@ -404,7 +411,9 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: mesh_future = executor.submit(self._unwrap_mesh, mesh) - image_future = executor.submit(self._prepare_reference_image, image_path) + image_future = executor.submit( + self._prepare_reference_image, image_path, batch + ) paint_mesh = mesh_future.result() delighted_image = image_future.result() @@ -542,15 +551,27 @@ def _camera_index(azimuth: float, elevation: float) -> int: base, divisor = 12, 1 return base + base_index // divisor - def _timesteps(self, device: torch.device) -> torch.Tensor: + def _timesteps( + self, device: torch.device, batch: Req | None = None + ) -> torch.Tensor: + target_steps = ( + 10 + if self.config.paint_turbo_mode + else (self.config.paint_num_inference_steps) + ) + measured_steps = ( + min(target_steps, max(1, int(batch.num_inference_steps))) + if batch is not None and batch.is_warmup + else target_steps + ) + if batch is not None: + batch.record_stage_iterations(measured_steps, target_steps) if not self.config.paint_turbo_mode: - self.scheduler.set_timesteps( - self.config.paint_num_inference_steps, device=device - ) + self.scheduler.set_timesteps(measured_steps, device=device) return self.scheduler.timesteps self.scheduler.set_timesteps( - num_inference_steps=10, + num_inference_steps=measured_steps, original_inference_steps=30, device=device, ) @@ -639,7 +660,7 @@ def _prepare_denoising_inputs(self, batch: Req) -> PaintDenoisingInputs: if position_attention_mask is not None: model_kwargs["position_attn_mask"] = position_attention_mask - timesteps = self._timesteps(device) + timesteps = self._timesteps(device, batch) latent_channels = self.transformer.config.in_channels latent_size = render_size // self.vae_scale_factor latents = randn_tensor( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/lingbot_world_causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/lingbot_world_causal_denoising.py index 1065cee60a2c..4434f23c0e32 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/lingbot_world_causal_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/lingbot_world/lingbot_world_causal_denoising.py @@ -691,6 +691,7 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: "Ensure ImageVAEEncodingStage runs before this stage." ) ctx = self._prepare_causal_dmd_forward_context(batch, server_args) + batch.record_stage_iterations(len(ctx.timesteps), len(ctx.timesteps)) latents = ctx.latents cache_ctx = self._prepare_realtime_causal_caches(batch, server_args, ctx) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longlive2.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longlive2.py index 3716d53205ba..3b22866ddbbb 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longlive2.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longlive2.py @@ -421,6 +421,9 @@ def warm_up(context_input, start_frame): num_blocks = (t - 1) // self.num_frames_per_block block_sizes = [1] + [self.num_frames_per_block] * num_blocks + total_iterations = len(block_sizes) * len(timesteps) + batch.record_stage_iterations(total_iterations, total_iterations) + start_index = 0 self._validate_block_prompt_count(batch, block_sizes) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising_av.py index 57e09ddef8e0..a18b20ba0b26 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising_av.py @@ -394,6 +394,10 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: batch.scheduler = scheduler batch.timesteps = scheduler.timesteps batch.num_inference_steps = num_steps + previous_stage_target_steps = batch.extra.get( + "stage_target_num_inference_steps" + ) + batch.extra["stage_target_num_inference_steps"] = num_steps original_do_cfg = batch.do_classifier_free_guidance batch.do_classifier_free_guidance = False @@ -403,6 +407,12 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: batch.scheduler = original_batch_scheduler batch.timesteps = original_batch_timesteps batch.num_inference_steps = original_batch_num_inference_steps + if previous_stage_target_steps is None: + batch.extra.pop("stage_target_num_inference_steps", None) + else: + batch.extra["stage_target_num_inference_steps"] = ( + previous_stage_target_steps + ) batch.do_classifier_free_guidance = original_do_cfg batch.ltx2_ti2v_clean_latent_background = original_clean_latent_background diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py index 26b817b8b4fb..ea4cddc0faa2 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py @@ -632,6 +632,7 @@ def _run_full_loop(self, batch: Req, server_args: ServerArgs) -> None: device = current_platform.get_local_torch_device() sigmas_video = [float(v) for v in ctx.sigmas["video"]] + batch.record_stage_iterations(len(sigmas_video) - 1) self._maybe_enable_cache_dit_and_torch_compile( len(sigmas_video) - 1, batch, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py index e346b54c0764..15f7103e4e3d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py @@ -474,6 +474,7 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: boundary_ratio = server_args.pipeline_config.boundary_ratio total_steps = paired_timesteps.shape[0] + batch.record_stage_iterations(total_steps) cfg_rank = get_classifier_free_guidance_rank() enable_cfg_parallel = server_args.enable_cfg_parallel diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/base.py index 7689c5ddb555..9962f0cee1f2 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/base.py @@ -955,6 +955,7 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: cfg_parallel, float(getattr(batch, "guidance_scale", 1.0) or 1.0), ) + batch.record_stage_iterations(len(timesteps)) log_sana_wm_tensor_stats("denoise.input_latents", latents) start_time = time.perf_counter() diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/refiner.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/refiner.py index 31e1f7213a3f..87e29cf644d1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/refiner.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/refiner.py @@ -725,6 +725,8 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: return batch batch_size = int(batch.latents.shape[0]) + total_iterations = batch_size * (len(STAGE_2_DISTILLED_SIGMA_VALUES) - 1) + batch.record_stage_iterations(total_iterations, total_iterations) prompts = self._prompts_for_batch(batch, batch_size) fps = float(getattr(batch, "fps", 16) or 16) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py index 3b38547a4789..de6d9f8dbac8 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py @@ -630,6 +630,8 @@ def _fdump(_name, _t): len(explicit_sigmas), do_cfg, ) + total_iterations = num_chunks * len(explicit_sigmas) + batch.record_stage_iterations(total_iterations, total_iterations) for chunk_idx in self.progress_bar(range(num_chunks), batch=batch): chunk_kv, sink_num = self._accumulate_kv_cache( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming_refiner.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming_refiner.py index 6325f0dd79b3..46190432eccb 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming_refiner.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming_refiner.py @@ -671,6 +671,8 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: ) return batch n_blocks = math.ceil(n_active / self.block_size) + total_iterations = n_blocks * (len(STAGE_2_DISTILLED_SIGMA_VALUES) - 1) + batch.record_stage_iterations(total_iterations, total_iterations) self.log_info( "SANA-WM streaming refiner: latent=%s, sink=%d, block=%d, blocks=%d, kv_max=%d, seed=%d", tuple(latents.shape), diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py index 32c3285e07f5..d2a1ca9ded0a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py @@ -516,6 +516,7 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: scheduler = ctx.scheduler n_steps = int(batch.num_inference_steps) + batch.record_stage_iterations(n_steps) timesteps_cpu = ctx.timesteps.cpu() transition_steps = find_transition_steps( diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index 297a0695101e..4f5feb7d4b19 100644 --- a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py +++ b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py @@ -58,6 +58,9 @@ def __init__(self, request_id: str): self.request_id = request_id self.stages: Dict[str, float] = {} self.steps: list[float] = [] + self.steps_by_stage: Dict[str, list[float]] = {} + self.stage_iterations: Dict[str, tuple[int, int]] = {} + self.active_stage_name: str | None = None self.total_duration_ms: float = 0.0 self.suppress_stage_breakdown: bool = False # memory tracking: {checkpoint_name: MemorySnapshot} @@ -77,7 +80,26 @@ def record_step(self, duration_s: float): """Records the duration of a denoising step in execution order.""" if self.suppress_stage_breakdown: return - self.steps.append(duration_s * 1000) + duration_ms = duration_s * 1000 + self.steps.append(duration_ms) + if self.active_stage_name is not None: + self.steps_by_stage.setdefault(self.active_stage_name, []).append( + duration_ms + ) + + def record_stage_iterations( + self, measured_iterations: int, target_iterations: int + ) -> None: + """Record calibration and default-workload iterations for this stage.""" + if self.suppress_stage_breakdown or self.active_stage_name is None: + return + measured = max(1, int(measured_iterations)) + target = max(1, int(target_iterations)) + previous = self.stage_iterations.get(self.active_stage_name, (0, 0)) + self.stage_iterations[self.active_stage_name] = ( + previous[0] + measured, + previous[1] + target, + ) def record_memory_snapshot(self, checkpoint_name: str, snapshot: MemorySnapshot): if self.suppress_stage_breakdown: diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 476a28bd6d91..8eb2a7fe13a8 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -67,10 +67,13 @@ def _record( total_duration_ms=0.0, stage_duration_ms=None, step_duration_ms=(), + step_duration_ms_by_stage=None, + stage_iterations=None, phase_active_components=None, phase_used_components=None, phase_full_weight_transition_components=None, layerwise_layer_uses=None, + layerwise_layer_uses_by_stage=None, ) -> WarmupMemoryRecord: return WarmupMemoryRecord( width=width, @@ -83,12 +86,15 @@ def _record( total_duration_ms=total_duration_ms, stage_duration_ms=stage_duration_ms or {}, step_duration_ms=step_duration_ms, + step_duration_ms_by_stage=step_duration_ms_by_stage or {}, + stage_iterations=stage_iterations or {}, phase_active_components=phase_active_components or {}, phase_used_components=phase_used_components or {}, phase_full_weight_transition_components=( phase_full_weight_transition_components or {} ), layerwise_layer_uses=layerwise_layer_uses or {}, + layerwise_layer_uses_by_stage=layerwise_layer_uses_by_stage or {}, ) @@ -176,6 +182,39 @@ def test_scales_standard_and_nonstandard_denoising_stages_together(self): } assert request_ns == 5_600_000_000 + def test_uses_each_stage_iteration_target(self): + record = _record( + num_inference_steps=4, + total_duration_ms=1_100, + stage_duration_ms={ + "ShapeDenoisingStage": 500, + "PaintStage": 600, + }, + step_duration_ms_by_stage={ + "ShapeDenoisingStage": (200, 100, 100, 100), + }, + stage_iterations={ + "ShapeDenoisingStage": (4, 50), + "PaintStage": (4, 30), + }, + phase_active_components={ + "0:ShapeDenoisingStage:use:hy3dshape_model": ("hy3dshape_model",), + "1:PaintStage:use:paint_transformer": ("paint_transformer",), + }, + ) + + request_ns, stage_ns, _ = estimate_default_workload_timing( + records=[record], + target_units=record.workload_units(), + target_num_inference_steps=50, + ) + + assert stage_ns == { + "ShapeDenoisingStage": 5_100_000_000, + "PaintStage": 4_500_000_000, + } + assert request_ns == 9_600_000_000 + def test_transfer_savings_is_capped_by_request_not_component_stage(self): candidates = [ _candidate("text_encoder", weight_gib=10, h2d_gib=10), @@ -282,6 +321,36 @@ def test_stage_role_scales_nonstandard_denoiser_but_not_encoder_repeats(self): assert uses["custom_encoder"]["layers"] == (2, 2) assert uses["custom_refiner"]["blocks"] == (10, 10) + def test_scales_each_stage_calls_with_its_own_iteration_target(self): + record = _record( + num_inference_steps=4, + phase_used_components={ + "0:ShapeStage:use:transformer": ("transformer",), + "1:PaintStage:use:transformer": ("transformer",), + }, + stage_iterations={ + "ShapeStage": (4, 50), + "PaintStage": (4, 30), + }, + layerwise_layer_uses={ + "transformer": {"blocks": (9, 9), "one_shot": (2, 2)} + }, + layerwise_layer_uses_by_stage={ + "ShapeStage": {"transformer": {"blocks": (4, 4), "one_shot": (1, 1)}}, + "PaintStage": {"transformer": {"blocks": (4, 4), "one_shot": (1, 1)}}, + }, + ) + + uses = estimate_layerwise_layer_uses( + records=[record], + target_units=record.workload_units(), + target_num_inference_steps=50, + ) + + # one untracked one-shot call + 50 shape calls + 30 paint calls + assert uses["transformer"]["blocks"] == (81, 81) + assert uses["transformer"]["one_shot"] == (2, 2) + class TestEstimateDefaultWorkloadPeak: def test_same_shape_uses_measured_peak(self): @@ -2001,6 +2070,36 @@ def test_host_pin_frontier_coalesces_equal_transformer_layers(self): assert len(targets) == 41 + def test_host_pin_frontier_drops_scratch_only_tradeoffs_when_nonbinding(self): + manager = SimpleNamespace( + num_layers=3, + residency_policy="leading", + pin_cpu_memory=True, + layer_weight_bytes=lambda: {0: 6, 1: 4, 2: 4}, + layer_host_store_bytes=lambda: {0: 6, 1: 4, 2: 4}, + pinnable_layer_indices=lambda: (0, 1, 2), + ) + + constrained = _layerwise_pin_targets( + managers=[manager], + resident_layers=(0,), + current_pinned_layers=((0, 1, 2),), + uses_per_streamed_layer=1, + layer_uses=((2, 1, 1),), + ) + unconstrained = _layerwise_pin_targets( + managers=[manager], + resident_layers=(0,), + current_pinned_layers=((0, 1, 2),), + uses_per_streamed_layer=1, + layer_uses=((2, 1, 1),), + constrain_host_transitions=False, + ) + + assert ((1, 2),) in constrained + assert ((1, 2),) not in unconstrained + assert ((0,),) in unconstrained + def test_filters_and_sizes(self): te = nn.Linear(8, 8) vae = nn.Linear(4, 4) @@ -2270,6 +2369,26 @@ def test_component_offload_frontier_can_expand_layerwise_lazily(self): ) assert any(candidate.target_mode() == RESIDENT for candidate in candidates) + def test_virtual_layerwise_prefetch_is_runtime_not_transition_memory(self): + module = _FakeLazyLayerwiseDit(num_layers=3) + candidates = collect_residency_targets( + modules={"transformer": module}, + residency_mode_of=self._modes({"transformer": COMPONENT_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=4, + layerwise_tuning_of=lambda _name, _dit_group: (1.0, 0.0, "leading"), + ) + + streamed = next( + candidate + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + and candidate.target_layerwise_resident_layers == (0,) + ) + assert streamed.device_transition_delta_bytes == 0 + assert streamed.active_device_delta_bytes < 0 + def test_lazy_layerwise_frontier_keeps_buffers_out_of_managed_weights(self): module = _FakeLazyLayerwiseDit(num_layers=2) module.layers[0].register_buffer("cache", torch.zeros(4096)) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py index 212e5ecdcfe6..7201f73ca6c2 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -40,7 +40,14 @@ def supports_auto_residency(self) -> bool: component_residency_strategies={ "fixed": ComponentResidencyStrategy(), "dynamic": AutoCompatibleStrategy(), - } + }, + modules={}, + ) + worker.server_args = SimpleNamespace( + component_quantizations=(), + quantization=None, + direct_gpu_weight_loading=False, + nunchaku_config=None, ) self.assertEqual(worker._fixed_custom_residency_strategy_names(), {"fixed"}) diff --git a/python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py b/python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py index cbfec5d64d0e..4f17edb9f482 100644 --- a/python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py +++ b/python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py @@ -2,6 +2,7 @@ import unittest from types import SimpleNamespace +from unittest.mock import Mock import torch from diffusers import AutoencoderKL as DiffusersAutoencoderKL @@ -222,6 +223,24 @@ def test_uses_standard_lcm_schedule_without_custom_timesteps(self): ) self.assertFalse(stage.scheduler.custom_timesteps) + def test_warmup_caps_paint_steps_and_records_full_target(self): + stage = Hunyuan3DPaintTexGenStage.__new__(Hunyuan3DPaintTexGenStage) + stage.config = Hunyuan3D2PipelineConfig(paint_turbo_mode=True) + stage.scheduler = LCMScheduler( + num_train_timesteps=1000, + original_inference_steps=50, + ) + batch = SimpleNamespace( + is_warmup=True, + num_inference_steps=4, + record_stage_iterations=Mock(), + ) + + timesteps = stage._timesteps(torch.device("cpu"), batch) + + self.assertEqual(len(timesteps), 4) + batch.record_stage_iterations.assert_called_once_with(4, 10) + if __name__ == "__main__": unittest.main() 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 f465a6c5b3f1..add8179fd9fc 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -84,6 +84,23 @@ def test_layerwise_usage_tracker_observes_only_executed_groups(): assert all(not layer._forward_pre_hooks for layer in module.decoder) +def test_layerwise_usage_tracker_attributes_calls_to_stages(): + module = _UsageTrackedVAE() + stage = "encode" + tracker = LayerwiseUsageTracker({"vae": module}, stage_name_provider=lambda: stage) + + module(torch.zeros(1)) + stage = "decode" + module(torch.zeros(1)) + uses, uses_by_stage = tracker.finish_with_stages() + + assert uses["vae"]["decoder"] == (2, 2, 2) + assert uses_by_stage == { + "encode": {"vae": {"decoder": (1, 1, 1)}}, + "decode": {"vae": {"decoder": (1, 1, 1)}}, + } + + class _FakeStream: def wait_stream(self, _stream) -> None: return None diff --git a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py index 2e3cb79dcf9b..519f91996491 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -44,6 +44,27 @@ def _perf_record(memory_snapshots: dict[str, dict]) -> RequestPerfRecord: ) +def test_request_metrics_attributes_steps_and_iterations_to_active_stage(): + metrics = RequestMetrics("request") + metrics.active_stage_name = "ShapeStage" + + metrics.record_step(0.1) + metrics.record_stage_iterations(4, 50) + metrics.active_stage_name = "PaintStage" + metrics.record_step(0.2) + metrics.record_stage_iterations(4, 30) + + assert metrics.steps == [100.0, 200.0] + assert metrics.steps_by_stage == { + "ShapeStage": [100.0], + "PaintStage": [200.0], + } + assert metrics.stage_iterations == { + "ShapeStage": (4, 50), + "PaintStage": (4, 30), + } + + def test_performance_summary_separates_load_and_runtime_peaks(): summary = PerformanceSummary.from_req_perf_record( _perf_record( @@ -201,6 +222,10 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): auto_residency_mode=lambda _name: None, layerwise_tuning_for=lambda _name, *, dit_group: (0.0, 0.0, "leading"), pin_cpu_memory=True, + component_quantizations=(), + quantization=None, + direct_gpu_weight_loading=False, + nunchaku_config=None, ) device_module = Mock() device_module.mem_get_info.return_value = (100, 100) @@ -248,6 +273,10 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): auto_residency_mode=lambda _name: None, layerwise_tuning_for=lambda _name, *, dit_group: (0.0, 0.0, "leading"), pin_cpu_memory=True, + component_quantizations=(), + quantization=None, + direct_gpu_weight_loading=False, + nunchaku_config=None, ) device_module = Mock() device_module.mem_get_info.return_value = (100 * GIB_BYTES, 140 * GIB_BYTES) diff --git a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py index d94f0a99f09d..55b199da57d7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py @@ -610,3 +610,44 @@ def test_branch_and_bound_matches_exhaustive_multiresource_search(): ) assert plan.estimated_latency_savings == expected_utility assert plan.placement_cost_bytes == expected_cost + + +def test_nonbinding_two_cost_sweep_matches_exhaustive_search(): + rng = random.Random(20260828) + + for _ in range(100): + options = [] + for group_index in range(rng.randint(1, 5)): + for option_index in range(rng.randint(1, 6)): + options.append( + PlacementOption( + group_key=f"group-{group_index}", + option_key=f"group-{group_index}:option-{option_index}", + resource_delta_bytes={}, + estimated_latency_savings=rng.randint(0, 50), + placement_cost_bytes=( + rng.randint(0, 20), + rng.randint(0, 20), + ), + ) + ) + tolerance = rng.randint(0, 10) + require_every_group = rng.choice((False, True)) + expected = _exhaustive_plan_key( + options, + resource_budget_bytes={}, + estimated_latency_tolerance=tolerance, + require_selection_from_every_group=require_every_group, + ) + + plan = optimize_placement( + options, + resource_budget_bytes={}, + estimated_latency_tolerance=tolerance, + require_selection_from_every_group=require_every_group, + ) + assert expected is not None + expected_keys, _, expected_utility, expected_cost = expected + assert tuple(option.option_key for option in plan.selections) == expected_keys + assert plan.estimated_latency_savings == expected_utility + assert plan.placement_cost_bytes == expected_cost From 1eb326ae5810a7b5d9975153a8c6df163417b475 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 19:52:11 +0800 Subject: [PATCH 62/90] [diffusion] align 5090 Wan consistency guard --- .../test/server/consistency_thresholds/5090.json | 6 ++++++ python/sglang/multimodal_gen/test/test_utils.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/test/server/consistency_thresholds/5090.json b/python/sglang/multimodal_gen/test/server/consistency_thresholds/5090.json index 9f6f79fa4ca1..f575e698d4c4 100644 --- a/python/sglang/multimodal_gen/test/server/consistency_thresholds/5090.json +++ b/python/sglang/multimodal_gen/test/server/consistency_thresholds/5090.json @@ -5,6 +5,12 @@ "ssim_threshold": 0.86, "psnr_threshold": 19.5, "mean_abs_diff_threshold": 8.5 + }, + "wan2_1_t2v_1.3b": { + "clip_threshold": 0.95, + "ssim_threshold": 0.85, + "psnr_threshold": 25.0, + "mean_abs_diff_threshold": 8.0 } } } diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 5d2dea467d92..7fb3d412183f 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -40,7 +40,7 @@ # NPU/ascend) is read from sgl-project/ci-data-diffusion, where the GT-gen workflows # publish. SGL_TEST_FILES_CI_DATA_REPO = "sgl-project/ci-data-diffusion" -SGL_TEST_FILES_CI_DATA_REVISION = "3d96c172db9dc5d056804bf4b45919e195fa4a2e" +SGL_TEST_FILES_CI_DATA_REVISION = "88b17d4fca324f9994d6f9f58495ed3882700734" # The NPU pin is kept as a separate branch so ascend GT can be bumped independently # when it's regenerated on its own cadence. From e923b79ff2c6e1df9813850640358ad446b1ee9c Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 20:03:01 +0800 Subject: [PATCH 63/90] [diffusion] align residency calibration stage names --- .../sglang/multimodal_gen/runtime/pipelines_core/stages/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index 68f7ce00fa7e..476c4c9181ea 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py @@ -391,7 +391,7 @@ def __call__( ) self._current_batch_is_warmup = batch.is_warmup if warmup_metrics is not None: - warmup_metrics.active_stage_name = stage_name + warmup_metrics.active_stage_name = self._component_stage_name() try: with StageProfiler( stage_name, From d6e0a9e097ccee594a79943ed32b569e4152dd7e Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 20:10:53 +0800 Subject: [PATCH 64/90] [diffusion] use explicit residency warmup contracts --- .../managers/memory_managers/layerwise_offload.py | 4 +--- .../multimodal_gen/runtime/warmup_request_builder.py | 11 +++++------ 2 files changed, 6 insertions(+), 9 deletions(-) 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 a0bac3631678..842aa773826a 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 @@ -2328,9 +2328,7 @@ def remove_layerwise_offload(self, *, to_cpu: bool = False) -> None: parameter = parameters.get(name) if parameter is not None: parameter.data = host_tensor - if current_platform.is_mps() and hasattr( - self, "_mps_cpu_non_layer_parameters" - ): + if current_platform.is_mps(): self.restore_mps_cpu_non_layer_weights() self.to("cpu") else: diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py index 33e039b5894f..4d7917295c8e 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -276,10 +276,9 @@ def _halve_num_frames(server_args: ServerArgs, num_frames: int) -> int: """Halve the latent frame count, keeping the model's frame arithmetic.""" if num_frames <= 1: return num_frames - arch_config = getattr( - getattr(server_args.pipeline_config, "vae_config", None), "arch_config", None + ratio = ( + server_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio ) - ratio = getattr(arch_config, "temporal_compression_ratio", None) if not ratio or ratio <= 1: return max(1, num_frames // 2) latent_frames = (num_frames - 1) // ratio + 1 @@ -299,13 +298,13 @@ def lighten_warmup_req(server_args: ServerArgs, req: Req) -> Req | None: if params is None: return None - num_frames = getattr(params, "num_frames", None) or 1 + num_frames = params.num_frames or 1 lighter_frames = _halve_num_frames(server_args, num_frames) if lighter_frames < num_frames: return _replace_warmup_workload(req, num_frames=lighter_frames) - width = getattr(params, "width", None) - height = getattr(params, "height", None) + width = params.width + height = params.height if not width or not height: return None alignment = _warmup_resolution_alignment(server_args) From ca9155892ca3822de0a9d9d45ff7dd4b1f93c22d Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 20:32:56 +0800 Subject: [PATCH 65/90] [diffusion] shorten resident placement validation --- python/sglang/multimodal_gen/runtime/managers/gpu_worker.py | 6 +----- .../multimodal_gen/test/unit/test_auto_residency_warmup.py | 6 ++---- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index f777f200097e..a477aafcbe6f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -92,7 +92,6 @@ ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( RESIDENT, - is_dit_component_name, ) from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( HOST_COPY_RESERVE_BYTES, @@ -1825,10 +1824,7 @@ def _latest_auto_residency_round_is_resident_only(self) -> bool: def _latest_auto_residency_round_supports_short_validation(self) -> bool: """Whether one full-shape step covers every changed execution path.""" - changes = self._latest_auto_residency_round() - return self._latest_auto_residency_round_is_resident_only() and all( - not is_dit_component_name(change.component_name) for change in changes - ) + return self._latest_auto_residency_round_is_resident_only() def _rollback_applied_residency_changes( self, *, latest_round_only: bool = False diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py index 7201f73ca6c2..9a843daf5585 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -171,7 +171,7 @@ def test_worker_validation_does_not_apply_a_second_candidate_plan(self): plan.assert_not_called() apply.assert_not_called() - def test_aux_resident_promotion_uses_short_validation_but_dit_does_not(self): + def test_resident_promotions_use_short_validation(self): worker = GPUWorker.__new__(GPUWorker) worker.server_args = SimpleNamespace(residency_mode=lambda _name: RESIDENT) worker._auto_residency_round_sizes = [1] @@ -184,9 +184,7 @@ def test_aux_resident_promotion_uses_short_validation_but_dit_does_not(self): worker._auto_residency_applied = [ AppliedResidencyChange("transformer", COMPONENT_OFFLOAD) ] - self.assertFalse( - worker._latest_auto_residency_round_supports_short_validation() - ) + self.assertTrue(worker._latest_auto_residency_round_supports_short_validation()) def test_worker_rollback_response_rewarms_the_restored_layout(self): actions = [] From 713a16a1d91f2b65a20b66d93b816a87cf8911ab Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 20:55:01 +0800 Subject: [PATCH 66/90] [docs] align residency validation policy --- docs/docs/sglang-diffusion/deployment_cookbook.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/docs/sglang-diffusion/deployment_cookbook.mdx b/docs/docs/sglang-diffusion/deployment_cookbook.mdx index 6e36107bf238..337b8abc4c15 100644 --- a/docs/docs/sglang-diffusion/deployment_cookbook.mdx +++ b/docs/docs/sglang-diffusion/deployment_cookbook.mdx @@ -128,7 +128,7 @@ See [OpenAI API: Served model name](/docs/sglang-diffusion/api/openai_api#served -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. Auxiliary-only moves to resident use one full-shape denoising step; DiT changes, 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 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. 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. From 46d9b824c9dd0403c9727c246b69588136f22ee7 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 21:03:50 +0800 Subject: [PATCH 67/90] test(diffusion): remove duplicate residency case --- .../sglang/multimodal_gen/test/unit/test_auto_residency.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 8eb2a7fe13a8..5b2bf017599a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -505,12 +505,6 @@ def test_failure_above_the_target_size_is_dropped(self): ) assert estimate == good.peak_allocated_bytes - def test_unknown_target_disables_estimation(self): - assert ( - estimate_default_workload_peak_bytes(records=[_record()], target_units=None) - is None - ) - def test_no_records_disables_estimation(self): assert ( estimate_default_workload_peak_bytes(records=[], target_units=None) is None From ed21267454938bb2d82340f4e4922ffe61454991 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 21:57:21 +0800 Subject: [PATCH 68/90] [diffusion] bound nonbinding residency optimization --- .../memory_managers/auto_residency.py | 4 +- .../memory_managers/placement_budget.py | 113 ++++++++++++++++++ .../test/unit/test_auto_residency.py | 13 +- 3 files changed, 126 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 48a972df2615..d13b0b603468 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -460,7 +460,9 @@ def _repeated_stages( stages.update(record.stage_iterations) stages.update( stage_name - for stage_name in record.stage_duration_ms + for stage_name in set(record.stage_duration_ms).union( + *(stage_names for stage_names in component_stages.values()) + ) if stage_name.endswith("DenoisingStage") and not stage_name.endswith("BeforeDenoisingStage") ) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py index 55e12adbeddc..e51229542be6 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py @@ -54,6 +54,12 @@ class PlacementPlan(msgspec.Struct, frozen=True): tuple[PlacementOption, ...], ] +_UnconstrainedState = tuple[ + int, + tuple[int, ...], + tuple[PlacementOption, ...], +] + def _selection_key(state: _State) -> tuple[str, ...]: return tuple(option.option_key for option in state[3]) @@ -184,6 +190,93 @@ def _drop_locally_latency_ineligible_options( ] +def _prune_unconstrained_states( + states: Iterable[_UnconstrainedState], +) -> list[_UnconstrainedState]: + """Keep the exact latency-deficit versus lexicographic-cost frontier.""" + best_by_deficit: dict[int, _UnconstrainedState] = {} + for state in states: + incumbent = best_by_deficit.get(state[0]) + if incumbent is None or ( + state[1], + tuple(option.option_key for option in state[2]), + ) < (incumbent[1], tuple(option.option_key for option in incumbent[2])): + best_by_deficit[state[0]] = state + + frontier = [] + best_cost: tuple[int, ...] | None = None + for state in sorted( + best_by_deficit.values(), + key=lambda item: ( + item[0], + item[1], + tuple(option.option_key for option in item[2]), + ), + ): + if best_cost is not None and best_cost <= state[1]: + continue + frontier.append(state) + best_cost = state[1] + return frontier + + +def _optimize_unconstrained( + grouped: Mapping[str, list[PlacementOption]], + *, + placement_cost_dimensions: int, + estimated_latency_tolerance: int, + require_selection_from_every_group: bool, +) -> _State: + """Solve the exact global latency window after every budget is non-binding.""" + zero_cost = (0,) * placement_cost_dimensions + frontier: list[_UnconstrainedState] = [(0, zero_cost, ())] + maximum_utility = 0 + + for group_key in sorted(grouped): + alternatives = list(grouped[group_key]) + group_maximum = max( + [option.estimated_latency_savings for option in alternatives] + + ([] if require_selection_from_every_group else [0]) + ) + maximum_utility += group_maximum + group_states = [ + ( + group_maximum - option.estimated_latency_savings, + option.placement_cost_bytes or zero_cost, + (option,), + ) + for option in alternatives + if group_maximum - option.estimated_latency_savings + <= estimated_latency_tolerance + ] + if ( + not require_selection_from_every_group + and group_maximum <= estimated_latency_tolerance + ): + group_states.append((group_maximum, zero_cost, ())) + group_frontier = _prune_unconstrained_states(group_states) + frontier = _prune_unconstrained_states( + ( + accumulated[0] + choice[0], + tuple(left + right for left, right in zip(accumulated[1], choice[1])), + accumulated[2] + choice[2], + ) + for accumulated in frontier + for choice in group_frontier + if accumulated[0] + choice[0] <= estimated_latency_tolerance + ) + + best = min( + frontier, + key=lambda state: ( + state[1], + state[0], + tuple(option.option_key for option in state[2]), + ), + ) + return (), maximum_utility - best[0], best[1], best[2] + + def optimize_placement( options: Iterable[PlacementOption], *, @@ -259,6 +352,26 @@ def optimize_placement( int(resource_budget_bytes[name]) for name in resource_names ) + if not resource_names: + best = _optimize_unconstrained( + grouped, + placement_cost_dimensions=placement_cost_dimensions, + estimated_latency_tolerance=estimated_latency_tolerance, + require_selection_from_every_group=require_selection_from_every_group, + ) + full_resource_deltas = { + resource_name: sum( + option.resource_delta_bytes.get(resource_name, 0) for option in best[3] + ) + for resource_name in all_resource_names + } + return PlacementPlan( + resource_delta_bytes=full_resource_deltas, + estimated_latency_savings=best[1], + placement_cost_bytes=best[2], + selections=list(best[3]), + ) + # Each local frontier includes the implicit "keep current placement" choice # when selection is optional. Keep the branch representation compact so the # global solve does not materialize and repeatedly Pareto-scan a large diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 5b2bf017599a..9e1933bb57de 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -210,10 +210,10 @@ def test_uses_each_stage_iteration_target(self): ) assert stage_ns == { - "ShapeDenoisingStage": 5_100_000_000, + "ShapeDenoisingStage": 5_000_000_000, "PaintStage": 4_500_000_000, } - assert request_ns == 9_600_000_000 + assert request_ns == 9_500_000_000 def test_transfer_savings_is_capped_by_request_not_component_stage(self): candidates = [ @@ -1967,7 +1967,10 @@ def test_layerwise_runtime_floor_uses_only_the_streaming_window(self): runtime = component_runtime_weight_bytes({"transformer": module}) - assert runtime["transformer"] == manager.peak_managed_device_weight_bytes() + assert runtime["transformer"] == ( + module.layers[0].cache.untyped_storage().nbytes() + + manager.peak_managed_device_weight_bytes() + ) class TestCollectResidencyTargets: @@ -2462,6 +2465,10 @@ def test_layerwise_frontier_accounts_for_immediate_materialization(self): resident.device_transition_delta_bytes == manager.offloaded_weight_bytes() ) module.disable_offload() + module.register_parameter( + "materialized", + nn.Parameter(torch.zeros(manager.offloaded_weight_bytes() // 4)), + ) candidates = collect_residency_targets( modules={"transformer": module}, residency_mode_of=self._modes({"transformer": RESIDENT}), From 07fece53b03fe7ea9a427538357fc92474366992 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 22:01:15 +0800 Subject: [PATCH 69/90] [diffusion] fix parent layerwise fixture --- .../sglang/multimodal_gen/test/unit/test_auto_residency.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 9e1933bb57de..85d379092686 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -1967,10 +1967,7 @@ def test_layerwise_runtime_floor_uses_only_the_streaming_window(self): runtime = component_runtime_weight_bytes({"transformer": module}) - assert runtime["transformer"] == ( - module.layers[0].cache.untyped_storage().nbytes() - + manager.peak_managed_device_weight_bytes() - ) + assert runtime["transformer"] == manager.peak_managed_device_weight_bytes() class TestCollectResidencyTargets: From bf8c6b19c8f780a0a66a19cb2b9906e867aaffe7 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 22:26:14 +0800 Subject: [PATCH 70/90] [diffusion] profile residency planning phases --- .../multimodal_gen/runtime/managers/gpu_worker.py | 11 +++++++++++ .../managers/memory_managers/auto_residency.py | 13 +++++++++++++ 2 files changed, 24 insertions(+) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index a477aafcbe6f..df9cbc89e70e 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1682,6 +1682,7 @@ def _build_auto_residency_report( ) candidates = [] if include_candidates: + candidate_started = time.perf_counter() candidates = collect_residency_targets( modules=self.pipeline.modules, residency_mode_of=self.server_args.residency_mode, @@ -1706,6 +1707,16 @@ def _build_auto_residency_report( layerwise_layer_uses=layerwise_layer_uses, host_transition_headroom_bytes=host_transition_headroom_bytes, ) + candidates_by_component: dict[str, int] = {} + for candidate in candidates: + candidates_by_component[candidate.component_name] = ( + candidates_by_component.get(candidate.component_name, 0) + 1 + ) + logger.debug( + "Auto residency candidate frontier built in %.3fs: %s", + time.perf_counter() - candidate_started, + candidates_by_component, + ) measured_request_duration_ns, _, _ = estimate_default_workload_timing( records=records, target_units=target_units, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index d13b0b603468..993f73e18106 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -25,6 +25,7 @@ from __future__ import annotations import statistics +import time from itertools import chain, product from typing import TYPE_CHECKING, Callable, Iterable, Mapping, Sequence @@ -2370,6 +2371,7 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP if use_latency_utility else 0 ) + option_build_started = time.perf_counter() options = [] for candidate in candidates: resource_deltas: dict[str, int] = {} @@ -2490,6 +2492,13 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP any(candidate.current_placement for candidate in component_candidates) for component_candidates in candidates_by_component.values() ) + logger.debug( + "Auto residency option vectors built in %.3fs: options=%d, resources=%d", + time.perf_counter() - option_build_started, + len(options), + len(resource_budgets), + ) + solve_started = time.perf_counter() try: placement = optimize_placement( options, @@ -2506,6 +2515,10 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP current_placement_reserve_shortfall ), ) + logger.debug( + "Auto residency joint solve completed in %.3fs", + time.perf_counter() - solve_started, + ) changed_candidates = [] for selection in placement.selections: candidate = candidate_by_key[selection.option_key] From 4ea2dec175bc2a41be5e7508554c540d82502885 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 23:04:54 +0800 Subject: [PATCH 71/90] [diffusion] bound residency frontier construction --- .../runtime/managers/gpu_worker.py | 62 +++++----- .../memory_managers/auto_residency.py | 111 +++++++++++++++--- .../memory_managers/placement_budget.py | 4 +- .../test/unit/test_auto_residency.py | 83 +++++++++++++ .../test/unit/test_placement_budget.py | 4 +- 5 files changed, 218 insertions(+), 46 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index df9cbc89e70e..367d5bbf82bd 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1266,7 +1266,9 @@ def get_can_stay_resident_components( num_inference_steps=( workload.num_inference_steps if workload is not None else 1 ), - allow_host_pin_reallocation=self.server_args.num_gpus == 1, + # this hint only consumes permanent-residency targets; HostPin + # alternatives cannot change its answer + allow_host_pin_reallocation=False, mixed_dtype_components=self._mixed_dtype_residency_components(), ) @@ -1663,6 +1665,22 @@ def _build_auto_residency_report( target_units=target_units, target_num_inference_steps=workload.num_inference_steps, ) + measured_request_duration_ns, _, _ = estimate_default_workload_timing( + records=records, + target_units=target_units, + target_num_inference_steps=workload.num_inference_steps, + ) + reference_request_duration_ns = self.__dict__.get( + "_auto_residency_reference_request_duration_ns" + ) + if reference_request_duration_ns is None and measured_request_duration_ns > 0: + reference_request_duration_ns = measured_request_duration_ns + self._auto_residency_reference_request_duration_ns = ( + reference_request_duration_ns + ) + estimated_request_duration_ns = ( + reference_request_duration_ns or measured_request_duration_ns + ) # What this process may still grow into: free VRAM plus what its # allocator already reserved. Unlike the raw device total, this # excludes memory held by other processes on a shared GPU. @@ -1680,6 +1698,16 @@ def _build_auto_residency_report( max(0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES) // local_worker_count ) + pinned_host_bytes = ( + layerwise_pinned_host_bytes(self.pipeline.modules) + if include_candidates + else 0 + ) + host_pin_capacity_bytes = ( + layerwise_host_pin_capacity_bytes(self.pipeline.modules) + if include_candidates + else 0 + ) candidates = [] if include_candidates: candidate_started = time.perf_counter() @@ -1706,6 +1734,10 @@ def _build_auto_residency_report( ), layerwise_layer_uses=layerwise_layer_uses, host_transition_headroom_bytes=host_transition_headroom_bytes, + host_pin_headroom_bytes=max( + 0, host_pin_capacity_bytes - pinned_host_bytes + ), + request_duration_ns=estimated_request_duration_ns, ) candidates_by_component: dict[str, int] = {} for candidate in candidates: @@ -1717,22 +1749,6 @@ def _build_auto_residency_report( time.perf_counter() - candidate_started, candidates_by_component, ) - measured_request_duration_ns, _, _ = estimate_default_workload_timing( - records=records, - target_units=target_units, - target_num_inference_steps=workload.num_inference_steps, - ) - reference_request_duration_ns = self.__dict__.get( - "_auto_residency_reference_request_duration_ns" - ) - if reference_request_duration_ns is None and measured_request_duration_ns > 0: - reference_request_duration_ns = measured_request_duration_ns - self._auto_residency_reference_request_duration_ns = ( - reference_request_duration_ns - ) - estimated_request_duration_ns = ( - reference_request_duration_ns or measured_request_duration_ns - ) candidate_latency_savings_ns = ( estimate_candidate_latency_savings_ns( candidates=candidates, @@ -1761,16 +1777,8 @@ def _build_auto_residency_report( else {} ), node_rank=self.server_args.node_rank, - pinned_host_bytes=( - layerwise_pinned_host_bytes(self.pipeline.modules) - if include_candidates - else 0 - ), - host_pin_capacity_bytes=( - layerwise_host_pin_capacity_bytes(self.pipeline.modules) - if include_candidates - else 0 - ), + pinned_host_bytes=pinned_host_bytes, + host_pin_capacity_bytes=host_pin_capacity_bytes, host_transition_headroom_bytes=( host_transition_headroom_bytes if include_candidates else 0 ), diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index 993f73e18106..d50047cf90b9 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -87,9 +87,10 @@ MAX_VRAM_RESERVE_FRACTION = 0.20 # A feasible placement is not automatically useful. Predictions inside this -# interval are treated as latency-equivalent, then the joint optimizer chooses -# the option with the lowest additional VRAM and HostPin use. The raw estimate -# is already an upper bound: transfer time is capped by the measured request. +# interval are treated as latency-equivalent. The joint optimizer then avoids +# changing strategy, minimizes device memory, preserves the faster estimate, +# and finally minimizes HostPin. The raw estimate is already an upper bound: +# transfer time is capped by the measured request. ESTIMATED_PINNED_H2D_BYTES_PER_SECOND = 24 * GIB_BYTES MIN_LATENCY_EQUIVALENCE_NS = 50_000_000 MAX_LATENCY_EQUIVALENCE_NS = 100_000_000 @@ -1020,6 +1021,7 @@ def _layerwise_pin_targets( uses_per_streamed_layer: int, layer_uses: tuple[tuple[int, ...], ...] | None = None, constrain_host_transitions: bool = True, + maximum_utility_only: bool = False, ) -> list[tuple[tuple[int, ...], ...]]: """Pareto-optimal HostPin targets for one resident-layer placement. @@ -1033,6 +1035,37 @@ def _layerwise_pin_targets( layer_uses=layer_uses, fallback_uses=uses_per_streamed_layer, ) + if maximum_utility_only: + target = [] + for manager, resident_count, manager_uses in zip( + managers, resident_layers, resolved_uses + ): + if not manager.pin_cpu_memory: + target.append(()) + continue + streamed = set( + compute_streamed_layers( + num_layers=manager.num_layers, + resident_layers=resident_count, + policy=manager.residency_policy, + ) + ) + transfer_bytes = manager.layer_weight_bytes() + target.append( + tuple( + layer_idx + for layer_idx in manager.pinnable_layer_indices() + if ( + manager_uses[layer_idx] + if layer_idx in streamed + else min(1, manager_uses[layer_idx]) + ) + * transfer_bytes.get(layer_idx, 0) + > 0 + ) + ) + return [tuple(target)] + current = [set(indices) for indices in current_pinned_layers] current_bytes = 0 grouped: dict[tuple[int, int, int, bool], list[tuple[int, int]]] = {} @@ -1494,6 +1527,8 @@ def collect_residency_targets( used_components: Iterable[str] | None = None, layerwise_layer_uses: Mapping[str, Mapping[str, tuple[int, ...]]] | None = None, host_transition_headroom_bytes: int | None = None, + host_pin_headroom_bytes: int | None = None, + request_duration_ns: int = 0, ) -> list[ResidencyTarget]: """Build complete target-state frontiers for auto-managed components. @@ -1501,6 +1536,12 @@ def collect_residency_targets( but its transfer utility is absolute within the component's frontier. Including the current and lower-memory states lets later calibration rounds replace an earlier choice instead of being limited to monotonic upgrades. + + When both host constraints provably cannot bind and the request-duration + cap cannot flatten transfer utility, every pin subset except the unique + maximum-utility one is dominated under the solver's latency-first ordering. + Only that exact condition permits collapsing the HostPin frontier; otherwise + every Pareto-relevant subset remains available to the joint optimizer. """ custom_names = set(custom_strategy_names) mixed_dtype_names = set(mixed_dtype_components) @@ -1509,7 +1550,11 @@ def collect_residency_targets( if baseline_residency_mode_of is None: baseline_residency_mode_of = residency_mode_of constrain_host_transitions = True - if host_transition_headroom_bytes is not None: + constrain_host_pin = True + if ( + host_transition_headroom_bytes is not None + or host_pin_headroom_bytes is not None + ): seen_managers: set[int] = set() maximum_unpin_bytes = 0 maximum_pin_bytes = 0 @@ -1531,9 +1576,12 @@ def collect_residency_targets( for layer_idx in manager.pinnable_layer_indices() if layer_idx not in pinned ) - constrain_host_transitions = max(maximum_unpin_bytes, maximum_pin_bytes) > max( - 0, host_transition_headroom_bytes - ) + if host_transition_headroom_bytes is not None: + constrain_host_transitions = max( + maximum_unpin_bytes, maximum_pin_bytes + ) > max(0, host_transition_headroom_bytes) + if host_pin_headroom_bytes is not None: + constrain_host_pin = maximum_pin_bytes > max(0, host_pin_headroom_bytes) candidates = [] for name in sorted(modules): module = modules[name] @@ -1729,6 +1777,15 @@ def collect_residency_targets( maximum_transfer_work = max( layerwise_maximum_transfer_work, component_transfer_work ) + pin_utility_cannot_saturate = ( + request_duration_ns > 0 + and int( + maximum_transfer_work + / ESTIMATED_PINNED_H2D_BYTES_PER_SECOND + * 1_000_000_000 + ) + <= request_duration_ns + ) empty_resident_layers = tuple(0 for _ in managers) empty_pinned_layers = tuple(() for _ in managers) @@ -1785,10 +1842,24 @@ def collect_residency_targets( uses_per_streamed_layer=uses_per_request, layer_uses=layer_uses, constrain_host_transitions=constrain_host_transitions, + maximum_utility_only=( + not constrain_host_transitions + and not constrain_host_pin + and pin_utility_cannot_saturate + ), ) if allow_host_pin_reallocation else [current_pinned_layers] ) + if ( + target_resident_layers == current_resident_layers + and current_pinned_layers not in pin_targets + ): + # relative utility is anchored to the measured placement; keep + # that exact state when non-binding host resources allow + # every other pin subset to collapse to its maximum-utility + # representative + pin_targets.append(current_pinned_layers) for target_pinned_layers in pin_targets: target_pinned_bytes = sum( sum( @@ -2459,24 +2530,30 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP resource_deltas[f"hostram:{prefix}:materialize"] = ( rank_candidate.host_materialize_scratch_bytes ) + estimated_latency_savings = ( + min( + report.candidate_latency_savings_ns.get(candidate.option_key(), 0) + for report in timed_reports + ) + if use_latency_utility + else candidate.h2d_bytes_per_request + ) options.append( PlacementOption( group_key=candidate.component_name, option_key=candidate.option_key(), resource_delta_bytes=resource_deltas, - estimated_latency_savings=( - min( - report.candidate_latency_savings_ns.get( - candidate.option_key(), 0 - ) - for report in timed_reports - ) - if use_latency_utility - else candidate.h2d_bytes_per_request - ), + estimated_latency_savings=estimated_latency_savings, placement_cost_bytes=( + ( + 0 + if candidate.current_placement + or candidate.target_mode() == candidate.residency_mode + else 1 + ), candidate.target_device_weight_bytes or max(0, candidate.target_resident_weight_bytes), + -estimated_latency_savings, candidate.target_pinned_host_bytes or max(0, candidate.pinned_host_delta_bytes), ), diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py index e51229542be6..9999252818c8 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/placement_budget.py @@ -76,7 +76,9 @@ def _dominates(candidate: _State, other: _State) -> bool: return ( candidate[1] >= other[1] and all(left <= right for left, right in zip(candidate[0], other[0])) - and all(left <= right for left, right in zip(candidate[2], other[2])) + # soft costs are summed and compared lexicographically by the final + # objective; this order is translation invariant across suffixes + and candidate[2] <= other[2] ) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 85d379092686..d9c1b22ec338 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -903,6 +903,42 @@ def test_long_request_does_not_spend_vram_on_tiny_encoder_gain(self): "transformer" ] + def test_latency_equivalent_plan_avoids_strategy_repacking(self): + current = ResidencyTarget( + component_name="vae", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + current_placement=True, + ) + component_offload = ResidencyTarget( + component_name="vae", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + ) + + plan = plan_auto_residency( + reports=[ + _report( + candidates=[current, component_offload], + estimated_request_duration_ns=1_000_000_000, + candidate_latency_savings_ns={ + current.option_key(): 0, + component_offload.option_key(): 10_000_000, + }, + ) + ] + ) + + assert plan.changes == [] + def test_output_rank_timing_applies_to_untimed_peer_ranks(self): transformer = _candidate( "transformer", @@ -2064,6 +2100,53 @@ def test_host_pin_frontier_coalesces_equal_transformer_layers(self): assert len(targets) == 41 + def test_nonbinding_host_resources_keep_only_maximum_pin_utility(self): + manager = SimpleNamespace( + num_layers=3, + residency_policy="leading", + pin_cpu_memory=True, + layer_weight_bytes=lambda: {0: 6, 1: 4, 2: 4}, + layer_host_store_bytes=lambda: {0: 6, 1: 4, 2: 4}, + pinnable_layer_indices=lambda: (0, 1, 2), + ) + + targets = _layerwise_pin_targets( + managers=[manager], + resident_layers=(1,), + current_pinned_layers=((0, 1, 2),), + uses_per_streamed_layer=1, + layer_uses=((0, 3, 1),), + constrain_host_transitions=False, + maximum_utility_only=True, + ) + + assert targets == [((1, 2),)] + + def test_collapsed_pin_frontier_retains_the_measured_state(self): + manager = _FakeLayerwiseManager( + {f"layers.{index}.w": torch.zeros(16) for index in range(3)} + ) + manager._pinned_layers = (0,) + module = _FakeLayerwiseDit([manager]) + + candidates = collect_residency_targets( + modules={"transformer": module}, + residency_mode_of=lambda _name: LAYERWISE_OFFLOAD, + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=3, + used_components={"transformer"}, + layerwise_layer_uses={"transformer": {"layers": (3, 3, 3)}}, + host_transition_headroom_bytes=GIB_BYTES, + host_pin_headroom_bytes=GIB_BYTES, + request_duration_ns=1_000_000_000, + ) + + current = [candidate for candidate in candidates if candidate.current_placement] + assert len(candidates) <= 8 + assert len(current) == 1 + assert current[0].target_layerwise_pinned_layers == ((0,),) + def test_host_pin_frontier_drops_scratch_only_tradeoffs_when_nonbinding(self): manager = SimpleNamespace( num_layers=3, diff --git a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py index 55b199da57d7..ed4db7134baf 100644 --- a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py @@ -573,7 +573,9 @@ def test_branch_and_bound_matches_exhaustive_multiresource_search(): name: rng.randint(-5, 8) for name in resource_names }, estimated_latency_savings=rng.randint(-5, 20), - placement_cost_bytes=(rng.randint(0, 10), rng.randint(0, 10)), + placement_cost_bytes=tuple( + rng.randint(-10, 10) for _ in range(4) + ), ) ) budgets = {name: rng.randint(-2, 12) for name in resource_names} From 45a5210b88b6a3729b8310d7ed208a99a15b880f Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Aug 2026 00:41:55 +0800 Subject: [PATCH 72/90] [diffusion] preserve LTX2 original DiT residency --- .../multimodal_gen/runtime/server_args/auto_tune.py | 8 +++++--- .../multimodal_gen/test/unit/test_auto_residency.py | 12 ++++++++++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py index 1d8d4f1c2956..031f31cc5aab 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -85,17 +85,19 @@ def auto_residency_args_skip_reason(server_args: ServerArgs) -> str | None: def fixed_loading_residency_components( server_args: ServerArgs, component_names: Iterable[str] ) -> set[str]: - """Components whose loader-selected storage must remain unchanged. + """Components whose selected residency must remain unchanged. Explicit/online quantization and direct GPU loading constrain their target - DiTs, not the unrelated encoders and VAEs in the same pipeline. Keep those - components outside auto residency while still calibrating the rest. + DiTs. LTX-2 two-stage ``original`` mode also relies on its offloaded DiT + while swapping the stage LoRA. Keep those components outside auto + residency while still calibrating unrelated encoders and VAEs. """ fixed = set(server_args.component_quantizations) nunchaku = server_args.nunchaku_config fixed_dit_loading = ( server_args.quantization is not None or server_args.direct_gpu_weight_loading + or server_args.ltx2_two_stage_device_mode == "original" or ( nunchaku is not None and ( diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index d9c1b22ec338..26e5e5d6049b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -3352,6 +3352,7 @@ def _base_args(self, **overrides) -> SimpleNamespace: transformer_weights_path=None, nunchaku_config=None, direct_gpu_weight_loading=False, + ltx2_two_stage_device_mode=None, pipeline_config=SimpleNamespace(task_type=ModelTaskType.T2V), ) for key, value in overrides.items(): @@ -3433,6 +3434,17 @@ def test_fixed_loading_components_are_excluded_individually(self): ) == set() ) + assert fixed_loading_residency_components( + self._base_args(ltx2_two_stage_device_mode="original"), + component_names, + ) == {"transformer", "transformer_2"} + assert ( + fixed_loading_residency_components( + self._base_args(ltx2_two_stage_device_mode="resident"), + component_names, + ) + == set() + ) def test_cache_dit_excluded(self, monkeypatch): monkeypatch.delenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", raising=False) From 08340130a19046ebef816347d0e51727d23f5a15 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Aug 2026 00:43:01 +0800 Subject: [PATCH 73/90] [diffusion] validate realized residency headroom --- docs/docs/sglang-diffusion/api/cli.mdx | 2 +- .../configs/pipeline_configs/base.py | 3 ++ .../configs/pipeline_configs/sana_wm.py | 3 ++ .../runtime/managers/gpu_worker.py | 19 +++++++++---- .../memory_managers/auto_residency.py | 9 +++++- .../runtime/server_args/auto_tune.py | 2 ++ .../test/unit/test_auto_residency.py | 28 ++++++++++++++++++- .../test/unit/test_cfg_parallel_warmup.py | 1 + .../test/unit/test_performance_metrics.py | 1 + 9 files changed, 60 insertions(+), 8 deletions(-) diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 79bb138c94a8..c9e493a31032 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -85,7 +85,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. 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 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. - `--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. Startup impact is model-dependent, so benchmark the target model before deployment. Disabled by default because checkpoint and model weights coexist temporarily, substantially increasing peak GPU memory. 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 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 6bdd0b901b99..334d3978693f 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -200,6 +200,9 @@ class PipelineConfig: task_type: ModelTaskType = ModelTaskType.I2I skip_input_image_preprocess: bool = False + # False when changing component placement after a calibration request is + # known to alter the pipeline's numerical path. + supports_auto_residency: bool = True model_path: str = "" pipeline_config_path: str | None = None diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py index 5826c278960a..f8333dbfada1 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py @@ -56,6 +56,9 @@ class SanaWMPipelineConfig(PipelineConfig): optional 6-DoF camera trajectory, optional Stage-2 LTX-2 refiner).""" task_type: ModelTaskType = ModelTaskType.TI2V + # The current two-stage path is not numerically invariant when its primary + # transformer and connectors are promoted after warmup. + supports_auto_residency: bool = False # SanaWMBeforeDenoisingStage._splice_first_frame handles condition-image # resize + VAE-encode itself, so bypass the framework's generic TI2V diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 367d5bbf82bd..bfd218f1aef4 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -791,8 +791,9 @@ def _record_server_warmup_memory( if request_allocated_peak > max(phase_allocated_peaks.values(), default=0): # Work after the residency-managed stage timeline (for example, # output materialization) must remain a placement constraint. - # A reserved-only increase is reclaimable allocator cache, not a - # second live placement, and is covered by the explicit reserve. + # A reserved-only increase is not a second live placement. Record + # it separately so post-placement validation can still require + # allocator headroom without charging cache to candidate deltas. phase_allocated_peaks["request:untracked"] = request_allocated_peak phase_components["request:untracked"] = untracked_active_components phase_used_components["request:untracked"] = () @@ -807,6 +808,9 @@ def _record_server_warmup_memory( baseline_allocated_bytes=int(baseline_allocated_bytes), peak_allocated_bytes=request_allocated_peak, succeeded=succeeded, + peak_reserved_bytes=int( + torch.get_device_module().max_memory_reserved() + ), phase_peak_allocated_bytes=phase_allocated_peaks, phase_active_components=phase_components, phase_used_components=phase_used_components, @@ -1634,6 +1638,11 @@ def _build_auto_residency_report( ) target_units = workload.workload_units() assert target_units is not None + target_records = [ + record + for record in records + if record.succeeded and record.workload_units() >= target_units + ] runtime_weights_by_component = component_runtime_weight_bytes( self.pipeline.modules ) @@ -1761,9 +1770,9 @@ def _build_auto_residency_report( rank=self.rank, budget_bytes=budget_bytes, estimated_peak_bytes=estimated_peak_bytes, - target_workload_measured=any( - record.succeeded and record.workload_units() >= target_units - for record in records + target_workload_measured=bool(target_records), + observed_reserved_bytes=max( + (record.peak_reserved_bytes for record in target_records), default=0 ), estimated_peak_bytes_by_phase=estimated_phase_peaks, active_components_by_phase=active_components_by_phase, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index d50047cf90b9..c38d90f558a0 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -132,6 +132,7 @@ class WarmupMemoryRecord(msgspec.Struct, frozen=True): baseline_allocated_bytes: int peak_allocated_bytes: int succeeded: bool + peak_reserved_bytes: int = 0 phase_peak_allocated_bytes: dict[str, int] = {} phase_active_components: dict[str, tuple[str, ...]] = {} phase_used_components: dict[str, tuple[str, ...]] = {} @@ -240,6 +241,7 @@ class RankResidencyReport(msgspec.Struct, frozen=True): budget_bytes: int estimated_peak_bytes: int | None target_workload_measured: bool = False + observed_reserved_bytes: int = 0 estimated_peak_bytes_by_phase: dict[str, int] = {} active_components_by_phase: dict[str, tuple[str, ...]] = {} used_components_by_phase: dict[str, tuple[str, ...]] = {} @@ -2056,7 +2058,12 @@ def current_placement_reserve_shortfall_bytes( *report.estimated_peak_bytes_by_phase.values(), ] ) - shortfalls.append(measured_peak + reserve - report.budget_bytes) + # Candidate deltas are intentionally solved against live allocated + # bytes, so reclaimable allocator cache is not charged twice. Once a + # placement has actually run, however, its mapped allocator footprint + # must still leave the same reserve for the next request. + realized_footprint = max(measured_peak, report.observed_reserved_bytes) + shortfalls.append(realized_footprint + reserve - report.budget_bytes) return max(0, max(shortfalls, default=0)) diff --git a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py index 1d8d4f1c2956..062fb738a855 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -63,6 +63,8 @@ def auto_residency_args_skip_reason(server_args: ServerArgs) -> str | None: task_type = server_args.pipeline_config.task_type if not (task_type.is_visual_gen() or task_type.is_mesh_gen()): return "no synthetic server warmup to calibrate from" + if not server_args.pipeline_config.supports_auto_residency: + return "pipeline does not support post-warmup residency changes" if has_realtime_model_adapter(server_args): return "realtime serving has no representative synthetic warmup" if server_args.backend == "diffusers": diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index d9c1b22ec338..809d97e8cdd5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -62,6 +62,7 @@ def _record( num_frames=17, baseline_gib=10, peak_gib=12, + peak_reserved_gib=0, succeeded=True, num_inference_steps=1, total_duration_ms=0.0, @@ -82,6 +83,7 @@ def _record( baseline_allocated_bytes=baseline_gib * GIB_BYTES, peak_allocated_bytes=peak_gib * GIB_BYTES, succeeded=succeeded, + peak_reserved_bytes=peak_reserved_gib * GIB_BYTES, num_inference_steps=num_inference_steps, total_duration_ms=total_duration_ms, stage_duration_ms=stage_duration_ms or {}, @@ -686,6 +688,7 @@ def _report( rank: int = 0, budget_gib: int = 100, estimated_gib: int | None = 50, + observed_reserved_gib: int = 0, candidates: list[ResidencyTarget] | None = None, skip_reason: str | None = None, phase_peaks_gib: dict[str, int] | None = None, @@ -710,6 +713,7 @@ def _report( None if estimated_gib is None else estimated_gib * GIB_BYTES ), target_workload_measured=target_workload_measured, + observed_reserved_bytes=observed_reserved_gib * GIB_BYTES, estimated_peak_bytes_by_phase={ name: value * GIB_BYTES for name, value in (phase_peaks_gib or {}).items() }, @@ -1361,6 +1365,16 @@ def test_validation_checks_measured_placement_without_candidate_frontier(self): assert current_placement_reserve_shortfall_bytes([report]) == 2 * GIB_BYTES + def test_validation_preserves_reserve_beyond_allocator_mapped_footprint(self): + report = _report( + budget_gib=80, + estimated_gib=60, + observed_reserved_gib=78, + target_workload_measured=True, + ) + + assert current_placement_reserve_shortfall_bytes([report]) == 2 * GIB_BYTES + def test_unobserved_later_component_gets_a_conservative_phase(self): first_transformer = _candidate("first_transformer", weight_gib=26, h2d_gib=100) second_transformer = ResidencyTarget( @@ -3352,7 +3366,10 @@ def _base_args(self, **overrides) -> SimpleNamespace: transformer_weights_path=None, nunchaku_config=None, direct_gpu_weight_loading=False, - pipeline_config=SimpleNamespace(task_type=ModelTaskType.T2V), + pipeline_config=SimpleNamespace( + task_type=ModelTaskType.T2V, + supports_auto_residency=True, + ), ) for key, value in overrides.items(): setattr(args, key, value) @@ -3388,6 +3405,15 @@ def test_manual_performance_mode(self, monkeypatch): # resident aux components on CPU): its peaks are not serving peaks ({"enable_torch_compile": True}, "stripped memory layout"), ({"batching_max_size": 4}, "batching"), + ( + { + "pipeline_config": SimpleNamespace( + task_type=ModelTaskType.T2V, + supports_auto_residency=False, + ) + }, + "post-warmup residency changes", + ), ], ) def test_excluded_paths(self, monkeypatch, overrides, expected_fragment): diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index 44d2289ee3ac..4602d2ebd1bf 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py @@ -836,6 +836,7 @@ def test_mesh_pipeline_builds_image_conditioned_warmup(self): server_args.is_arg_explicitly_set.return_value = False server_args.pipeline_config = SimpleNamespace( task_type=ModelTaskType.I2M, + supports_auto_residency=True, vae_stride=None, vae_scale_factor=None, vae_config=None, diff --git a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py index 519f91996491..5233e7622d48 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -155,6 +155,7 @@ def test_server_warmup_preserves_peak_after_managed_stage_timeline(): record = worker._auto_residency_warmup_records[0] assert record.peak_allocated_bytes == 9 + assert record.peak_reserved_bytes == 12 assert (record.width, record.height, record.num_frames) == (128, 96, 9) assert record.num_inference_steps == 2 assert record.phase_peak_allocated_bytes["request:untracked"] == 9 From 052f27f7bb9c52b412356a3c69a80ea842a47105 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Aug 2026 01:40:16 +0800 Subject: [PATCH 74/90] test(diffusion): align residency test fixtures --- .../multimodal_gen/test/unit/test_auto_residency_warmup.py | 1 + .../sglang/multimodal_gen/test/unit/test_performance_metrics.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py index 9a843daf5585..836066662373 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency_warmup.py @@ -47,6 +47,7 @@ def supports_auto_residency(self) -> bool: component_quantizations=(), quantization=None, direct_gpu_weight_loading=False, + ltx2_two_stage_device_mode=None, nunchaku_config=None, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py index 5233e7622d48..4e5c69a5e5e8 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -221,6 +221,7 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): explicit_residency_mode=lambda _name: None, component_residency_requirement=lambda _name: None, auto_residency_mode=lambda _name: None, + ltx2_two_stage_device_mode=None, layerwise_tuning_for=lambda _name, *, dit_group: (0.0, 0.0, "leading"), pin_cpu_memory=True, component_quantizations=(), @@ -272,6 +273,7 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): explicit_residency_mode=lambda _name: None, component_residency_requirement=lambda _name: None, auto_residency_mode=lambda _name: None, + ltx2_two_stage_device_mode=None, layerwise_tuning_for=lambda _name, *, dit_group: (0.0, 0.0, "leading"), pin_cpu_memory=True, component_quantizations=(), From 0df915ce36ff5dba53c26b81e45df9662b905645 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Aug 2026 03:34:02 +0800 Subject: [PATCH 75/90] test(diffusion): update LTX auto residency VRAM baseline --- .../sglang/multimodal_gen/test/server/perf_baselines/h100.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index 5fde9f15c888..8016cf7b1793 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -2796,7 +2796,7 @@ "expected_avg_denoise_ms": 802.06, "expected_median_denoise_ms": 747.62, "load_peak_vram_mb": 56048.0, - "runtime_peak_vram_mb": 66586.0, + "runtime_peak_vram_mb": 79016.0, "estimated_full_test_time_s": 363.2 }, "qwen_image_t2i_cache_dit_scm_config_diffusers_1gpu": { From 024138b0b04415387c2337f2199022875cf078c4 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Aug 2026 04:38:13 +0800 Subject: [PATCH 76/90] fix(diffusion): gate CI workload warmup by residency support --- .../multimodal_gen/test/server/test_server_common.py | 5 ++++- .../sglang/multimodal_gen/test/server/testcase_configs.py | 7 +++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 649d2538072d..b8ecfbc97f0c 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -49,6 +49,7 @@ get_model_task_type_for_server_args, get_perf_baseline_update_path, get_sampling_param_field_names_for_server_args, + supports_auto_residency_for_server_args, ) from sglang.multimodal_gen.test.test_utils import ( SGL_TEST_FILES_CI_DATA_REVISION, @@ -182,7 +183,9 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: has_explicit_warmup_workload = _has_server_option( extra_args, server_args.extras, "--warmup-sampling-params" ) - if not has_explicit_warmup_workload: + if not has_explicit_warmup_workload and supports_auto_residency_for_server_args( + server_args + ): warmup_sampling_params = _case_warmup_sampling_params(case) if warmup_sampling_params: encoded = json.dumps(warmup_sampling_params, separators=(",", ":")) diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index a1c0ef20ecfe..77300ee911d2 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -811,6 +811,13 @@ def get_model_task_type_for_server_args( return pipeline_config_cls.task_type +def supports_auto_residency_for_server_args( + server_args: DiffusionServerArgs, +) -> bool: + pipeline_config_cls, _ = _get_config_classes_for_server_args(server_args) + return pipeline_config_cls.supports_auto_residency + + def get_sampling_param_field_names_for_server_args( server_args: DiffusionServerArgs, ) -> frozenset[str]: From 4549570cb15b00f18d32d7f5bb2531e7778e65fb Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Aug 2026 05:38:56 +0800 Subject: [PATCH 77/90] fix(diffusion): skip calibration hooks for unsupported pipelines --- python/sglang/multimodal_gen/runtime/managers/gpu_worker.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index bfd218f1aef4..cfb93315de90 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -576,10 +576,12 @@ def _execute_forward_common( output_batch = None forward_failed = False # Prewarm reqs (is_warmup=False) run a different offload layout and - # must not contaminate the calibration records. + # must not contaminate the calibration records. Pipelines that cannot + # apply a residency plan also skip the temporary per-layer hooks. measure_server_warmup = ( req.is_warmup and bool(req.extra.get("server_based_warmup")) + and self.server_args.pipeline_config.supports_auto_residency and current_platform.is_cuda() ) warmup_workload = ( From d0207f36ef2ee1a1ffc3c168a8e28341aca14213 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Aug 2026 06:44:48 +0800 Subject: [PATCH 78/90] fix(diffusion): skip residency tracking for unsupported pipelines --- .../memory_managers/component_manager.py | 1 + .../test/unit/test_component_residency.py | 49 ++++++++++++++++--- 2 files changed, 43 insertions(+), 7 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py index 3b8851eaa23c..d28065c7a7e7 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py @@ -224,6 +224,7 @@ def begin_request( ) self._track_warmup_memory = ( self.state.batch_is_warmup + and self.server_args.pipeline_config.supports_auto_residency and current_platform.is_cuda() and torch.get_device_module().is_available() ) diff --git a/python/sglang/multimodal_gen/test/unit/test_component_residency.py b/python/sglang/multimodal_gen/test/unit/test_component_residency.py index 0cbbb8fb74eb..8ff9b8cf1d52 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_residency.py @@ -19,6 +19,15 @@ from sglang.multimodal_gen.runtime.platforms import current_platform +def _server_args(*, supports_auto_residency=True): + return SimpleNamespace( + enable_layerwise_nvtx_marker=False, + pipeline_config=SimpleNamespace( + supports_auto_residency=supports_auto_residency, + ), + ) + + def test_component_offload_releases_preferred_component_after_request(): strategy = ComponentOffloadStrategy() strategy.finish_use = Mock() @@ -137,7 +146,7 @@ def test_group_warmup_state_requires_every_batch_to_be_warmup(): _stage_name_mapping={}, component_residency_strategies={}, ) - server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + server_args = _server_args() manager = ComponentResidencyManager(pipeline, server_args) manager.begin_request( @@ -175,7 +184,7 @@ def test_reports_components_with_mixed_use_dtypes(): }, component_residency_strategies={}, ) - server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + server_args = _server_args() manager = ComponentResidencyManager(pipeline, server_args) manager.begin_request( @@ -203,7 +212,7 @@ def test_warmup_records_use_and_transition_peaks(monkeypatch): _stage_name_mapping={"denoise": stage}, component_residency_strategies={}, ) - server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + server_args = _server_args() manager = ComponentResidencyManager(pipeline, server_args) manager.strategy_for = Mock(return_value=Mock()) manager.refresh_pipeline(pipeline) @@ -234,6 +243,32 @@ def test_warmup_records_use_and_transition_peaks(monkeypatch): ) +def test_warmup_skips_memory_tracking_for_unsupported_pipeline(monkeypatch): + device_module = SimpleNamespace( + is_available=lambda: True, + reset_peak_memory_stats=Mock(), + ) + monkeypatch.setattr(torch, "get_device_module", lambda: device_module) + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + + stage = _Stage() + pipeline = SimpleNamespace( + modules={}, + _stage_name_mapping={"stage": stage}, + component_residency_strategies={}, + ) + server_args = _server_args(supports_auto_residency=False) + manager = ComponentResidencyManager(pipeline, server_args) + + manager.begin_request([stage], SimpleNamespace(is_warmup=True), server_args) + manager.before_stage(stage, 0, SimpleNamespace(is_warmup=True), server_args) + manager.finish_request() + + assert manager._track_warmup_memory is False + assert manager.take_warmup_phase_peaks() == {} + device_module.reset_peak_memory_stats.assert_not_called() + + def test_warmup_records_full_weight_transition_without_preparing(monkeypatch): device_module = SimpleNamespace( is_available=lambda: True, @@ -251,7 +286,7 @@ def test_warmup_records_full_weight_transition_without_preparing(monkeypatch): _stage_name_mapping={"lora_switch": stage}, component_residency_strategies={}, ) - server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + server_args = _server_args() manager = ComponentResidencyManager(pipeline, server_args) manager.strategy_for = Mock() manager.refresh_pipeline(pipeline) @@ -291,7 +326,7 @@ def test_warmup_records_same_component_dtype_prepare_as_transition(monkeypatch): _stage_name_mapping={"stage": stage}, component_residency_strategies={}, ) - server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + server_args = _server_args() manager = ComponentResidencyManager(pipeline, server_args) strategy = Mock() manager.strategy_for = Mock(return_value=strategy) @@ -335,7 +370,7 @@ def test_warmup_attributes_prefetch_peak_to_prefetched_component(monkeypatch): }, component_residency_strategies={}, ) - server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + server_args = _server_args() manager = ComponentResidencyManager(pipeline, server_args) strategy = Mock() strategy.prefetch_for_use.return_value = True @@ -380,7 +415,7 @@ def test_warmup_splits_sequential_component_transition(monkeypatch): _stage_name_mapping={"stage": stage}, component_residency_strategies={}, ) - server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + server_args = _server_args() manager = ComponentResidencyManager(pipeline, server_args) manager.strategy_for = Mock(return_value=Mock()) manager.refresh_pipeline(pipeline) From c2f533e2a36910a00180b3e8f170a2f290b50d73 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Aug 2026 07:49:30 +0800 Subject: [PATCH 79/90] fix(diffusion): keep SANA refiner encoder resident on H100 --- .../configs/pipeline_configs/sana_wm.py | 6 ++++++ .../test/unit/test_server_args.py | 21 +++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py index f8333dbfada1..0b13878dd580 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py @@ -165,6 +165,12 @@ def __post_init__(self): def get_model_deployment_config(self) -> ModelDeploymentConfig: return ModelDeploymentConfig( dit_layerwise_offload_modes=("memory",), + # The two-stage refiner's Gemma-3 encoder is not numerically + # invariant under layerwise offload. Keep text encoders resident + # when there is enough headroom, while preserving the low-memory + # layerwise path on consumer GPUs. + keep_resident_min_available_gb=70, + keep_resident_components=("text_encoder",), # Conservative auto-FSDP gate for the 720p world-model path. Users # can still force FSDP explicitly on smaller cards. fsdp_auto_min_available_memory_gb=60, diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 15e6e678fa7e..9c10eedcc0dd 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -1527,6 +1527,8 @@ def test_pipeline_configs_declare_auto_tune_hints(self): self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60) self.assertEqual(sana_wm_deployment.dit_layerwise_offload_modes, ("memory",)) + self.assertEqual(sana_wm_deployment.keep_resident_min_available_gb, 70) + self.assertEqual(sana_wm_deployment.keep_resident_components, ("text_encoder",)) fast_hunyuan_deployment = FastHunyuanConfig().get_model_deployment_config() self.assertEqual(fast_hunyuan_deployment.keep_resident_min_available_gb, 60) @@ -1569,6 +1571,25 @@ def test_longlive_residency_scales_with_available_memory(self): self.assertIn("text_encoder", constrained_offload) self.assertIn("vae", constrained_offload) + def test_sana_wm_residency_scales_with_available_memory(self): + high_memory_args = self._from_dict_with_pipeline_config( + SanaWMPipelineConfig(), + memory_gb=80, + kwargs={"performance_mode": "auto"}, + ) + high_memory_offload = high_memory_args.layerwise_offload_components or [] + self.assertNotIn("text_encoder", high_memory_offload) + self.assertIn("vae", high_memory_offload) + + constrained_args = self._from_dict_with_pipeline_config( + SanaWMPipelineConfig(), + memory_gb=60, + kwargs={"performance_mode": "auto"}, + ) + constrained_offload = constrained_args.layerwise_offload_components or [] + self.assertIn("text_encoder", constrained_offload) + self.assertIn("vae", constrained_offload) + def test_qwen_ar_generation_residency_scales_with_available_memory(self): pipeline_configs = ( QwenImageLayeredPipelineConfig(), From d4071450ef869276c2b53f3d4cd63b8c68e6347b Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Aug 2026 08:36:11 +0800 Subject: [PATCH 80/90] fix(diffusion): keep SANA VAE resident on H100 --- .../multimodal_gen/configs/pipeline_configs/sana_wm.py | 8 ++++---- .../sglang/multimodal_gen/test/unit/test_server_args.py | 6 ++++-- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py index 0b13878dd580..c88cbc617e8c 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py @@ -165,12 +165,12 @@ def __post_init__(self): def get_model_deployment_config(self) -> ModelDeploymentConfig: return ModelDeploymentConfig( dit_layerwise_offload_modes=("memory",), - # The two-stage refiner's Gemma-3 encoder is not numerically - # invariant under layerwise offload. Keep text encoders resident - # when there is enough headroom, while preserving the low-memory + # The two-stage path is not numerically invariant when its text + # encoders or VAE use layerwise offload. Keep them resident when + # there is enough headroom, while preserving the low-memory # layerwise path on consumer GPUs. keep_resident_min_available_gb=70, - keep_resident_components=("text_encoder",), + keep_resident_components=("text_encoder", "vae"), # Conservative auto-FSDP gate for the 720p world-model path. Users # can still force FSDP explicitly on smaller cards. fsdp_auto_min_available_memory_gb=60, diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 55e982c3ae96..a1fdf2ef910a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -1538,7 +1538,9 @@ def test_pipeline_configs_declare_auto_tune_hints(self): self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60) self.assertEqual(sana_wm_deployment.dit_layerwise_offload_modes, ("memory",)) self.assertEqual(sana_wm_deployment.keep_resident_min_available_gb, 70) - self.assertEqual(sana_wm_deployment.keep_resident_components, ("text_encoder",)) + self.assertEqual( + sana_wm_deployment.keep_resident_components, ("text_encoder", "vae") + ) fast_hunyuan_deployment = FastHunyuanConfig().get_model_deployment_config() self.assertEqual(fast_hunyuan_deployment.keep_resident_min_available_gb, 60) @@ -1589,7 +1591,7 @@ def test_sana_wm_residency_scales_with_available_memory(self): ) high_memory_offload = high_memory_args.layerwise_offload_components or [] self.assertNotIn("text_encoder", high_memory_offload) - self.assertIn("vae", high_memory_offload) + self.assertNotIn("vae", high_memory_offload) constrained_args = self._from_dict_with_pipeline_config( SanaWMPipelineConfig(), From 137286c33d296d77ec092f0e84fc6c3dd1803b0f Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Aug 2026 10:45:48 +0800 Subject: [PATCH 81/90] test(diffusion): pin compatible SANA consistency data --- python/sglang/multimodal_gen/test/test_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 7fb3d412183f..848985256979 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -40,7 +40,7 @@ # NPU/ascend) is read from sgl-project/ci-data-diffusion, where the GT-gen workflows # publish. SGL_TEST_FILES_CI_DATA_REPO = "sgl-project/ci-data-diffusion" -SGL_TEST_FILES_CI_DATA_REVISION = "88b17d4fca324f9994d6f9f58495ed3882700734" +SGL_TEST_FILES_CI_DATA_REVISION = "f8c4481d1fa2fa1b9efdfe6831d9c60dbf386b4a" # The NPU pin is kept as a separate branch so ascend GT can be bumped independently # when it's regenerated on its own cadence. From 1af94a44651c3defa7e619066ceb54ff12346aae Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 30 Aug 2026 22:39:29 +0800 Subject: [PATCH 82/90] Fix transformer fallback test fixture --- .../multimodal_gen/test/unit/test_transformer_loader_fallback.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_loader_fallback.py b/python/sglang/multimodal_gen/test/unit/test_transformer_loader_fallback.py index f237ec06ac43..b412a3ed927a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_loader_fallback.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_loader_fallback.py @@ -40,6 +40,7 @@ def _server_args(*, fsdp_requested=False, **overrides): "dp_size": 1, "use_fsdp_inference": False, "resolve_component_attention_backend": mock.Mock(return_value=(None, None)), + "should_direct_gpu_weight_load_component": mock.Mock(return_value=False), "should_use_fsdp_for_component": mock.Mock(return_value=fsdp_requested), } values.update(overrides) From daf1e456956534d52a778d044589e066873ab305 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 31 Aug 2026 09:56:15 +0800 Subject: [PATCH 83/90] fix(diffusion): contain auto residency regressions --- .../sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py | 3 +++ .../multimodal_gen/configs/pipeline_configs/qwen_image.py | 2 ++ .../multimodal_gen/test/server/perf_baselines/h100.json | 4 ++-- python/sglang/multimodal_gen/test/unit/test_server_args.py | 2 ++ 4 files changed, 9 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py index ccc4d2715b1c..240c23023bf0 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py @@ -723,3 +723,6 @@ def _unpad_and_unpack_latents(self, latents, audio_latents, batch, vae, audio_va @dataclasses.dataclass class LTX23PipelineConfig(LTX2PipelineConfig): """Configuration overrides for LTX-2.3.""" + + # original-mode lora swaps invalidate post-warmup timing calibration + supports_auto_residency: bool = False diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py index 91169c686bda..d076d17b92c2 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py @@ -755,6 +755,8 @@ class QwenImageEditPlus_2511_PipelineConfig(QwenImageEditPlusPipelineConfig): class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig): resolution: int = 640 vae_precision: str = "bf16" + # promoting the auxiliary components regresses first-request latency + supports_auto_residency: bool = False def get_model_deployment_config(self) -> ModelDeploymentConfig: return ModelDeploymentConfig( diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index 8016cf7b1793..4cada4fc2202 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -906,7 +906,7 @@ "expected_avg_denoise_ms": 666.89, "expected_median_denoise_ms": 673.29, "load_peak_vram_mb": 35324.0, - "runtime_peak_vram_mb": 45124.0, + "runtime_peak_vram_mb": 56582.0, "estimated_full_test_time_s": 117.6 }, "qwen_image_t2i_cache_dit_enabled": { @@ -2254,7 +2254,7 @@ "expected_avg_denoise_ms": 32.68, "expected_median_denoise_ms": 30.2, "load_peak_vram_mb": 876.0, - "runtime_peak_vram_mb": 6094.0, + "runtime_peak_vram_mb": 15686.0, "estimated_full_test_time_s": 420.1 }, "wan2_1_t2v_1.3b_frame_interp_2x": { diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index dd2b6af3af46..c155ac252bd5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -1585,6 +1585,8 @@ def test_pipeline_configs_declare_auto_tune_hints(self): self.assertFalse( ltx23_config.dit_config.arch_config.enable_packed_qkv_input_a2a ) + self.assertFalse(ltx23_config.supports_auto_residency) + self.assertFalse(QwenImageLayeredPipelineConfig().supports_auto_residency) self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60) self.assertEqual(sana_wm_deployment.dit_layerwise_offload_modes, ("memory",)) From d54ea114364144f6ee1385925b49804c624af38d Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 31 Aug 2026 13:01:36 +0800 Subject: [PATCH 84/90] fix(diffusion): refresh 2-gpu residency baselines --- .../multimodal_gen/test/server/perf_baselines/h100.json | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index 4cada4fc2202..22a2064035eb 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -172,7 +172,7 @@ "expected_avg_denoise_ms": 192.9, "expected_median_denoise_ms": 190.02, "load_peak_vram_mb": 41840.0, - "runtime_peak_vram_mb": 46938.0, + "runtime_peak_vram_mb": 58542.0, "estimated_full_test_time_s": 133.2 }, "ideogram4_fp8_t2i": { @@ -670,7 +670,7 @@ "expected_avg_denoise_ms": 79.98, "expected_median_denoise_ms": 81.01, "load_peak_vram_mb": 23284.0, - "runtime_peak_vram_mb": 28148.0, + "runtime_peak_vram_mb": 32294.0, "estimated_full_test_time_s": 125.5 }, "zimage_image_t2i": { @@ -1638,7 +1638,7 @@ "expected_avg_denoise_ms": 653.89, "expected_median_denoise_ms": 656.68, "load_peak_vram_mb": 35350.0, - "runtime_peak_vram_mb": 42674.0, + "runtime_peak_vram_mb": 49998.0, "estimated_full_test_time_s": 243.0 }, "wan2_1_i2v_14b_720P_2gpu": { @@ -1766,7 +1766,7 @@ "expected_avg_denoise_ms": 2029.88, "expected_median_denoise_ms": 2020.71, "load_peak_vram_mb": 6440.0, - "runtime_peak_vram_mb": 21524.0, + "runtime_peak_vram_mb": 72524.0, "estimated_full_test_time_s": 204.3 }, "wan2_1_t2v_14b_2gpu": { From 20c12cb78ff00cbd02c25bb57b125a5973bc93da Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 31 Aug 2026 16:01:00 +0800 Subject: [PATCH 85/90] test(diffusion): include component precision in selector fixtures --- .../test/unit/test_attention_backend_selector.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py b/python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py index 5ab486125de3..56faa26542d5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py +++ b/python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py @@ -333,6 +333,7 @@ def load_customized(self, *_args): return object() class _Args: + component_precisions = {} component_quantizations = {} @staticmethod @@ -411,6 +412,7 @@ def load_customized(self, *_args): return object() class _Args: + component_precisions = {} component_quantizations = {} @staticmethod @@ -475,6 +477,7 @@ def load_native(self, *_args): return object() class _Args: + component_precisions = {} component_quantizations = {} pipeline_config = SimpleNamespace(native_only_components=()) @@ -525,6 +528,7 @@ def load_native(self, *_args): return object() class _Args: + component_precisions = {} component_quantizations = {} pipeline_config = SimpleNamespace(native_only_components=()) From 992c36ba334bd840473d902514ec71f9fbfad51f Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 31 Aug 2026 22:01:09 +0800 Subject: [PATCH 86/90] test(diffusion): refresh auto residency VRAM baselines --- .../test/server/perf_baselines/h100.json | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index 22a2064035eb..326854423d1e 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -1270,7 +1270,7 @@ "expected_avg_denoise_ms": 118.58, "expected_median_denoise_ms": 109.14, "load_peak_vram_mb": 43920.0, - "runtime_peak_vram_mb": 57528.0, + "runtime_peak_vram_mb": 72438.0, "estimated_full_test_time_s": 345.4 }, "wan2_2_ti2v_5b": { @@ -1568,7 +1568,7 @@ "expected_avg_denoise_ms": 1689.73, "expected_median_denoise_ms": 1714.58, "load_peak_vram_mb": 6440.0, - "runtime_peak_vram_mb": 26322.0, + "runtime_peak_vram_mb": 75468.0, "estimated_full_test_time_s": 264.1 }, "wan2_1_i2v_14b_480P_2gpu": { @@ -1892,7 +1892,7 @@ "expected_avg_denoise_ms": 1437.82, "expected_median_denoise_ms": 1416.85, "load_peak_vram_mb": 6784.0, - "runtime_peak_vram_mb": 21450.0, + "runtime_peak_vram_mb": 72396.0, "estimated_full_test_time_s": 181.8 }, "wan2_1_t2v_1_3b_lora_1gpu": { @@ -2030,7 +2030,7 @@ "expected_avg_denoise_ms": 1913.15, "expected_median_denoise_ms": 1922.21, "load_peak_vram_mb": 36266.0, - "runtime_peak_vram_mb": 51132.0, + "runtime_peak_vram_mb": 58868.0, "estimated_full_test_time_s": 248.2 }, "flux_2_image_t2i_2_gpus": { @@ -2185,7 +2185,7 @@ "expected_avg_denoise_ms": 82.26, "expected_median_denoise_ms": 81.75, "load_peak_vram_mb": 7176.0, - "runtime_peak_vram_mb": 12634.0, + "runtime_peak_vram_mb": 24338.0, "estimated_full_test_time_s": 122.7 }, "hunyuan3d_shape_gen": { From 29083a389f2cc9014d9b7174a3ef090906fcd7f8 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 1 Sep 2026 02:34:07 +0800 Subject: [PATCH 87/90] fix(diffusion): contain auto residency regressions --- python/sglang/multimodal_gen/runtime/server_args/auto_tune.py | 2 ++ .../multimodal_gen/test/server/perf_baselines/h100.json | 2 +- python/sglang/multimodal_gen/test/unit/test_auto_residency.py | 4 ++++ 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py index fe4c71451c78..e2bca48aebd0 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -55,6 +55,8 @@ def auto_residency_args_skip_reason(server_args: ServerArgs) -> str | None: return "disabled via SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY" if server_args.performance_mode != "auto": return f"performance_mode={server_args.performance_mode}" + if server_args.ltx2_two_stage_device_mode == "original": + return "LTX-2 original two-stage placement" if ( server_args.warmup_mode != "server" or server_args.disagg_role != RoleType.MONOLITHIC diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index 326854423d1e..6ae73c3fd030 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -1516,7 +1516,7 @@ "stages_ms": { "InputValidationStage": 15.59, "ImageVAEEncodingStage": 1484.2, - "DecodingStage": 1114.63, + "DecodingStage": 2450.0, "LatentPreparationStage": 0.21, "TextEncodingStage": 327.94, "TimestepPreparationStage": 2.47, diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 7a02e7a57128..3d074c2d1d38 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -3401,6 +3401,10 @@ def test_manual_performance_mode(self, monkeypatch): ({"warmup_mode": "request"}, "server warmup"), ({"disagg_role": "denoiser"}, "server warmup"), ({"backend": "diffusers"}, "diffusers"), + ( + {"ltx2_two_stage_device_mode": "original"}, + "LTX-2 original two-stage placement", + ), ({"enable_breakable_cuda_graph": True}, "CUDA graph"), # compile warmup strips the memory layout (layerwise DiT + # resident aux components on CPU): its peaks are not serving peaks From a16fc476cdcdac2d8e6441d87b7151d478cf66ac Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 1 Sep 2026 05:33:53 +0800 Subject: [PATCH 88/90] fix(diffusion): preserve legacy LTX-2 placement --- .../sglang/multimodal_gen/runtime/server_args/auto_tune.py | 5 +++++ .../sglang/multimodal_gen/test/unit/test_auto_residency.py | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py index e2bca48aebd0..0df676dee6e6 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -55,6 +55,11 @@ def auto_residency_args_skip_reason(server_args: ServerArgs) -> str | None: return "disabled via SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY" if server_args.performance_mode != "auto": return f"performance_mode={server_args.performance_mode}" + if ( + server_args.pipeline_class_name == "LTX2TwoStagePipeline" + and server_args.ltx2_two_stage_device_mode is None + ): + return "legacy LTX-2 two-stage placement" if server_args.ltx2_two_stage_device_mode == "original": return "LTX-2 original two-stage placement" if ( diff --git a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py index 3d074c2d1d38..b5d548675037 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -3367,6 +3367,7 @@ def _base_args(self, **overrides) -> SimpleNamespace: nunchaku_config=None, direct_gpu_weight_loading=False, ltx2_two_stage_device_mode=None, + pipeline_class_name=None, pipeline_config=SimpleNamespace( task_type=ModelTaskType.T2V, supports_auto_residency=True, @@ -3405,6 +3406,10 @@ def test_manual_performance_mode(self, monkeypatch): {"ltx2_two_stage_device_mode": "original"}, "LTX-2 original two-stage placement", ), + ( + {"pipeline_class_name": "LTX2TwoStagePipeline"}, + "legacy LTX-2 two-stage placement", + ), ({"enable_breakable_cuda_graph": True}, "CUDA graph"), # compile warmup strips the memory layout (layerwise DiT + # resident aux components on CPU): its peaks are not serving peaks From 56596e0570b3ef7346d6cacb45883ef1492b3a82 Mon Sep 17 00:00:00 2001 From: mickqian Date: Thu, 3 Sep 2026 20:28:01 +0800 Subject: [PATCH 89/90] Fix Ruff formatting after main merge --- .../runtime/managers/memory_managers/auto_residency.py | 4 +--- .../sglang/multimodal_gen/runtime/warmup_request_builder.py | 2 +- .../multimodal_gen/test/unit/test_cfg_parallel_warmup.py | 2 +- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py index c38d90f558a0..cf72fc6ef3a3 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/auto_residency.py @@ -3120,9 +3120,7 @@ def _format_residency_change(candidate: ResidencyTarget) -> str: f"{candidate.target_layerwise_resident_layers}, pinned layers=" f"{pin_counts}" ) - return ( - f"{candidate.component_name}: {candidate.residency_mode} -> " f"{target_mode}" - ) + return f"{candidate.component_name}: {candidate.residency_mode} -> {target_mode}" def _format_candidate_summary(candidate: ResidencyTarget) -> str: diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py index 4d7917295c8e..532771cc0744 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -100,7 +100,7 @@ def _apply_warmup_sampling_overrides( unknown = value.keys() - field_names if unknown: raise ValueError( - "invalid --warmup-sampling-params fields: " f"{', '.join(sorted(unknown))}" + f"invalid --warmup-sampling-params fields: {', '.join(sorted(unknown))}" ) updated = copy(defaults) for name, field_value in value.items(): diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index 4602d2ebd1bf..8447f29e1374 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py @@ -119,7 +119,7 @@ def test_sampling_workload_override_accepts_json(self): ) server_args = SimpleNamespace( warmup_sampling_params=( - '{"width":832,"height":480,"num_frames":9,' '"num_inference_steps":4}' + '{"width":832,"height":480,"num_frames":9,"num_inference_steps":4}' ) ) From 5c9b3d0bdb83a5bb4f8022b0c2c1f96da99d95b7 Mon Sep 17 00:00:00 2001 From: Mick Qian Date: Fri, 11 Sep 2026 22:52:29 +0800 Subject: [PATCH 90/90] Let a pipeline without a residency manager enter the LoRA offload context The snapshot-offload tests build the pipeline as a bare namespace; the transition wrapper must not require the attribute. --- .../multimodal_gen/runtime/pipelines_core/lora/pipeline.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py index 8f87304844eb..fe7e0b663814 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py @@ -304,9 +304,10 @@ def _temporarily_disable_offload( # Record every target, including coarse component offload. The current # coarse path can already materialize the full component here, and the # planner must preserve that phase if it later chooses layerwise mode. + residency_manager = getattr(self, "component_residency_manager", None) residency_transition = ( - self.component_residency_manager.full_weight_transition(module_names) - if self.component_residency_manager is not None + residency_manager.full_weight_transition(module_names) + if residency_manager is not None else nullcontext() )