From 66f2864fe5725029bfb5a66e9fbdcbd6436bd1b3 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 22:10:11 +0800 Subject: [PATCH 01/70] [diffusion] plan reversible residency before warmup Inventory selected component weights before loading and use them to seed a conservative resident set. The seed changes only ordinary component residency, so post-load planning can replace it before the first serving forward. Run the same exact multi-resource placement solver after components are sized and again after one representative warmup. Each target state jointly obeys per-rank execution-phase VRAM, placement-transition scratch, per-worker HostPin, and host repacking headroom. The complete target frontier includes resident, component offload, and partial layerwise placements for DiTs, encoders, VAEs, lazy components, and two-stage pipelines. Keep structural loading decisions such as explicit FSDP outside the reversible planner, cap normal calibration at one adjustment and one verification warmup, and retain bounded OOM recovery. Extend offline startup to use the same sequence and add exhaustive solver, lifecycle, rollback, multi-rank, custom-pipeline, and server-args coverage. --- .../runtime/entrypoints/control_requests.py | 4 +- .../entrypoints/diffusion_generator.py | 8 + .../runtime/entrypoints/http_server.py | 32 +- .../loader/component_loaders/bridge_loader.py | 1 + .../component_loaders/component_loader.py | 14 + .../component_loaders/transformer_loader.py | 1 + .../runtime/managers/gpu_worker.py | 384 +++++++-- .../memory_managers/auto_residency.py | 589 +++++++++---- .../component_loading_order.py | 104 +-- .../memory_managers/component_manager.py | 34 +- .../component_residency_strategies.py | 34 +- .../component_weight_inventory.py | 229 ++++++ .../memory_managers/initial_residency.py | 152 ++++ .../memory_managers/layerwise_offload.py | 87 +- .../memory_managers/placement_budget.py | 26 +- .../runtime/managers/scheduler.py | 2 +- .../runtime/pipelines/hunyuan3d_pipeline.py | 115 +++ .../runtime/pipelines/ltx_2_pipeline.py | 48 +- .../runtime/pipelines/sana_wm_pipeline.py | 21 + .../pipelines_core/composed_pipeline_base.py | 120 ++- .../executors/pipeline_executor.py | 3 + .../model_specific_stages/sana_wm/refiner.py | 36 +- .../runtime/server_args/auto_tune.py | 87 +- .../runtime/server_args/server_args.py | 8 + .../multimodal_gen/runtime/server_warmup.py | 110 ++- .../test/server/test_server_utils.py | 1 + .../test/unit/sana_wm/test_pipeline_config.py | 87 +- .../test/unit/test_auto_residency.py | 771 +++++++++++++++++- .../test/unit/test_auto_residency_warmup.py | 282 ++++++- .../test/unit/test_cfg_parallel_warmup.py | 26 +- .../test/unit/test_component_loading_order.py | 155 +++- .../test/unit/test_component_residency.py | 125 ++- .../test_hunyuan3d_native_texture_models.py | 71 ++ .../test/unit/test_initial_residency.py | 182 +++++ .../test/unit/test_layerwise_offload.py | 17 + .../test/unit/test_ltx2_5_config.py | 43 + .../test/unit/test_performance_metrics.py | 74 ++ .../test/unit/test_placement_budget.py | 75 +- .../test/unit/test_server_args.py | 186 +++-- .../test/unit/test_server_case_requests.py | 72 ++ 40 files changed, 3821 insertions(+), 595 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/managers/memory_managers/component_weight_inventory.py create mode 100644 python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_initial_residency.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_server_case_requests.py diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py b/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py index 9df19c7f8969..51843e2274e7 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py @@ -51,10 +51,10 @@ class GetDisaggStatsReq(msgspec.Struct): class AutoResidencyReq(msgspec.Struct, frozen=True): - """Apply or roll back the warmup-calibrated residency promotion. + """Plan, apply, or roll back automatic component residency. Sent by the server warmup orchestrator after the synthetic warmup requests finish and before the server reports ready. """ - action: str = "apply" # "apply" | "rollback" + action: str = "apply" # "apply" | "apply_static" | "rollback" diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 712ba7d1bea6..ce3fcf08d5e9 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -40,7 +40,10 @@ from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs from sglang.multimodal_gen.runtime.server_warmup import ( run_sync_client_warmup, + run_sync_startup_warmup, + should_apply_pre_warmup_auto_residency, should_run_explicit_client_warmup, + should_run_synthetic_server_warmup, ) from sglang.multimodal_gen.runtime.utils.logging_utils import ( GREEN, @@ -191,6 +194,11 @@ def _start_local_server_if_needed( return processes def _run_client_warmup_if_needed(self) -> None: + if should_apply_pre_warmup_auto_residency( + self.server_args + ) or should_run_synthetic_server_warmup(self.server_args): + run_sync_startup_warmup(self.server_args, sync_scheduler_client.forward) + return if not should_run_explicit_client_warmup(self.server_args): return diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index 42a73d2b22f4..79c625e622dd 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -36,7 +36,9 @@ from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args from sglang.multimodal_gen.runtime.server_warmup import ( maybe_apply_auto_residency, + maybe_apply_pre_warmup_auto_residency, run_async_client_warmup, + should_apply_pre_warmup_auto_residency, should_run_synthetic_server_warmup, ) from sglang.multimodal_gen.runtime.utils.logging_utils import ( @@ -84,21 +86,31 @@ async def _run_server_warmup_after_http_live( server_args: ServerArgs, warmup_done: asyncio.Event ) -> None: try: - if not should_run_synthetic_server_warmup(server_args): + run_static_planner = should_apply_pre_warmup_auto_residency(server_args) + run_warmup = should_run_synthetic_server_warmup(server_args) + if not run_static_planner and not run_warmup: warmup_done.set() return await _wait_until_http_live(server_args) - await run_async_client_warmup( - server_args, - async_scheduler_client.forward, - fail_open=server_args.warmup_resolutions is None, - ) - # Freeze the auto residency decision before the server reports ready. - # Raises only when a rollback failed, which aborts startup below. - await maybe_apply_auto_residency(server_args, async_scheduler_client.forward) - logger.info("The server is fired up and ready to roll!") + if run_static_planner: + await maybe_apply_pre_warmup_auto_residency( + server_args, async_scheduler_client.forward + ) + + if run_warmup: + await run_async_client_warmup( + server_args, + async_scheduler_client.forward, + fail_open=server_args.warmup_resolutions is None, + ) + # Freeze the auto residency decision before the server reports ready. + # Raises only when a rollback failed, which aborts startup below. + await maybe_apply_auto_residency( + server_args, async_scheduler_client.forward + ) + logger.info("The server is fired up and ready to roll!") warmup_done.set() except asyncio.CancelledError: raise diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py index d648b00a8a48..af1564e4858b 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py @@ -24,6 +24,7 @@ class BridgeLoader(PlainStateDictComponentLoader): """Loader for MOVA dual tower bridge with FSDP support.""" pipeline_bridge_config_attr: str = "bridge_config" + supports_fsdp_loading = True component_names = ["dual_tower_bridge"] expected_library = "diffusers" diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py index ac3360d81a58..3104d5814077 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py @@ -133,6 +133,7 @@ class ComponentLoader(ABC): # Gates only --component-quantizations.. Quantization declared by a # checkpoint is discovered and admitted by the component's normal loader. supports_online_quantization_override = False + supports_fsdp_loading = False _loaders_registered = False @@ -640,6 +641,19 @@ class PipelineComponentLoader: Utility class for loading the components in a pipeline. """ + @staticmethod + def supports_fsdp_loading( + component_name: str, + transformers_or_diffusers: str, + component_architecture: str | None = None, + ) -> bool: + loader = ComponentLoader.for_component_type( + component_name, + transformers_or_diffusers, + component_architecture, + ) + return loader.supports_fsdp_loading + @staticmethod def load_component( component_name: str, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index 4b4050c7d05b..860b83ddb50f 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py @@ -154,6 +154,7 @@ class TransformerLoader(ComponentLoader): allow_global_attention_backend_fallback = False supports_online_quantization_override = True + supports_fsdp_loading = True component_names = [ "transformer", diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 3d7953fbf2de..829bfa2aeff0 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -51,16 +51,18 @@ ) from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( GIB_BYTES, - MAX_LATENCY_EQUIVALENCE_NS, + MIN_POST_ADJUSTMENT_REGRESSION_NS, PLACEMENT_STATUS_ADJUSTED, PLACEMENT_STATUS_ROLLBACK_FAILED, PLACEMENT_STATUS_ROLLED_BACK, PLACEMENT_STATUS_SKIPPED, POST_ADJUSTMENT_REGRESSION_FRACTION, AppliedResidencyChange, + AutoResidencyPlan, AutoResidencyRollbackError, DefaultWorkload, RankResidencyReport, + ResidencyTarget, WarmupMemoryRecord, apply_residency_changes, collect_residency_targets, @@ -75,6 +77,7 @@ format_plan_summary, layerwise_host_pin_capacity_bytes, layerwise_pinned_host_bytes, + measured_failed_workload_phase_peaks, plan_auto_residency, plan_summary_payload, rank_candidates_by_h2d_savings, @@ -85,9 +88,6 @@ 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, @@ -95,6 +95,9 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( configure_layerwise_offload_modules, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( + is_dit_component_name, +) from sglang.multimodal_gen.runtime.managers.memory_managers.memory_occupation_controller import ( MemoryOccupationController, ) @@ -207,10 +210,14 @@ 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_budget_correction_bytes = 0 + 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. self._auto_residency_reference_request_duration_ns: int | None = None + self._auto_residency_reference_stage_duration_ns: dict[str, int] = {} + self._auto_residency_reference_component_stages: dict[str, tuple[str, ...]] = {} # default workload resolved once for the per-request residency hint self._cached_default_workload: DefaultWorkload | None = None self._cached_default_workload_failed = False @@ -411,8 +418,13 @@ def do_mem_analysis(self, output_batch: OutputBatch): peak_reserved_gb = peak_reserved_bytes / (1024**3) peak_allocated_gb = peak_allocated_bytes / (1024**3) - remaining_gpu_mem_gb = ( - current_platform.get_device_total_memory() / (1024**3) - peak_reserved_gb + # Use the same process-visible capacity as auto residency. Physical + # card memory can be larger than this worker's allocator cap or can be + # partly occupied by another process; either case would make a hint + # based on total device memory unsafe. + remaining_gpu_mem_gb = max( + 0.0, + self._auto_residency_budget_bytes() / GIB_BYTES - peak_reserved_gb, ) try: can_stay_resident = self.get_can_stay_resident_components( @@ -757,6 +769,16 @@ def _record_server_warmup_memory( phase_used_components["request:untracked"] = () phase_full_weight_transition_components["request:untracked"] = () metrics = req.metrics + stage_duration_ms = {} + if metrics is not None: + profile_to_component_stage = { + stage._active_profile_stage_name(): stage._component_stage_name() + for stage in self.pipeline.stages + } + stage_duration_ms = { + profile_to_component_stage.get(stage_name, stage_name): duration_ms + for stage_name, duration_ms in metrics.stages.items() + } width, height, num_frames, num_inference_steps = workload self._auto_residency_warmup_records.append( WarmupMemoryRecord( @@ -776,7 +798,7 @@ def _record_server_warmup_memory( 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 {}), + stage_duration_ms=stage_duration_ms, step_duration_ms=(tuple(metrics.steps) if metrics is not None else ()), ) ) @@ -1198,12 +1220,15 @@ def get_can_stay_resident_components( 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: + if not self.pipeline: return [] + modules = self._auto_residency_modules() + if not modules: + return [] workload = self._default_workload_for_hint() candidates = collect_residency_targets( - modules=self.pipeline.modules, + modules=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 @@ -1232,7 +1257,7 @@ 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, *, pre_warmup: bool = False) -> OutputBatch: """Apply one warmup-calibrated residency adjustment round. Every rank executes this handler at the same queue position; the plan @@ -1255,8 +1280,12 @@ 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 + local_report = ( + self._build_pre_warmup_auto_residency_report(workload=workload) + if pre_warmup + else self._build_auto_residency_report( + workload=workload, records=records + ) ) except Exception as e: logger.warning( @@ -1276,13 +1305,18 @@ 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: + recovering_from_oom = any(report.warmup_oom for report in reports) + if ( + not pre_warmup + and self._auto_residency_round_sizes + and not recovering_from_oom + ): 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, + MIN_POST_ADJUSTMENT_REGRESSION_NS, int(reference_ns * POST_ADJUSTMENT_REGRESSION_FRACTION), ) if reference_ns > 0 and measured_ns > reference_ns + tolerance_ns: @@ -1304,9 +1338,14 @@ 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 + not pre_warmup + and self._auto_residency_round_sizes + and not recovering_from_oom and plan.current_placement_reserve_shortfall_bytes > 0 ): + self._record_auto_residency_budget_correction( + plan.current_placement_reserve_shortfall_bytes + ) shortfall_gib = plan.current_placement_reserve_shortfall_bytes / GIB_BYTES if self.is_output_rank: logger.warning( @@ -1318,6 +1357,7 @@ def apply_auto_residency(self) -> OutputBatch: cause=f"VRAM reserve exceeded by {shortfall_gib:.1f} GiB", already_failed=False, latest_round_only=True, + retryable=True, ) if plan.skip_reason is not None or not plan.changes: if self.is_output_rank: @@ -1335,13 +1375,14 @@ def apply_auto_residency(self) -> OutputBatch: try: newly_applied = apply_residency_changes( plan=plan, - modules=self.pipeline.modules, + modules=self._auto_residency_modules(), server_args=self.server_args, ) if not newly_applied: raise RuntimeError("placement plan applied no residency changes") self._auto_residency_applied.extend(newly_applied) self._auto_residency_round_sizes[-1] = len(newly_applied) + self._auto_residency_last_applied_plan = plan except AutoResidencyRollbackError as e: logger.error( "Auto residency adjustment failed on rank %d and the rank " @@ -1393,6 +1434,7 @@ def _rollback_everywhere( cause: str | None, already_failed: bool, latest_round_only: bool = False, + retryable: bool = False, ) -> OutputBatch: """Roll this rank back and gather the replica-wide outcome. @@ -1429,7 +1471,10 @@ def _rollback_everywhere( error=( f"auto residency adjustment failed; {restored} restored: {cause}" ), - output={"status": PLACEMENT_STATUS_ROLLED_BACK}, + output={ + "status": PLACEMENT_STATUS_ROLLED_BACK, + "retryable": retryable, + }, ) if self.is_output_rank: if self._auto_residency_applied: @@ -1445,64 +1490,33 @@ def _rollback_everywhere( ) return OutputBatch(output={"status": PLACEMENT_STATUS_ROLLED_BACK}) - def _build_auto_residency_report( - self, *, workload: DefaultWorkload, records: List[WarmupMemoryRecord] - ) -> RankResidencyReport: - 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 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 increasing residency 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, - budget_bytes=0, - estimated_peak_bytes=None, - 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=target_units, - constant_weight_bytes=max(runtime_weights_by_component.values(), default=0), + def _record_auto_residency_budget_correction(self, shortfall_bytes: int) -> None: + """Exclude a layout whose calibrated peak exceeded its VRAM model.""" + previous_plan = self._auto_residency_last_applied_plan + if previous_plan is None: + return + runtime_resources = [ + resource_name + for resource_name in previous_plan.resource_budget_bytes + if resource_name.startswith("gpu:") + and not resource_name.endswith(":placement-transition") + ] + if not runtime_resources: + return + minimum_slack = min( + previous_plan.resource_budget_bytes[resource_name] + - previous_plan.resource_delta_bytes.get(resource_name, 0) + for resource_name in runtime_resources ) - ( - 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, - component_weight_bytes=runtime_weights_by_component, + correction = max( + shortfall_bytes, + max(0, minimum_slack) + 64 * 1024**2, ) - # 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. + self._auto_residency_budget_correction_bytes += correction + + def _auto_residency_budget_bytes(self) -> int: + # Free VRAM plus this process's reserved allocator pool excludes memory + # held by unrelated processes while preserving reusable local capacity. free_bytes, _ = torch.get_device_module().mem_get_info() budget_bytes = int(free_bytes) + int( torch.get_device_module().memory_reserved() @@ -1510,8 +1524,15 @@ 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, + return budget_bytes + + def _collect_auto_residency_targets( + self, workload: DefaultWorkload + ) -> list[ResidencyTarget]: + assert self.pipeline is not None + modules = self._auto_residency_modules() + return collect_residency_targets( + modules=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, @@ -1519,43 +1540,223 @@ def _build_auto_residency_report( num_inference_steps=workload.num_inference_steps, allow_host_pin_reallocation=True, mixed_dtype_components=self._mixed_dtype_residency_components(), - auto_resident_components={ + required_resident_components={ name - for name in self.pipeline.modules - if self.server_args.auto_residency_mode(name) == RESIDENT + for name in modules + if self.server_args.component_residency_requirement(name) is not None }, 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( + + def _build_pre_warmup_auto_residency_report( + self, *, workload: DefaultWorkload + ) -> RankResidencyReport: + """Build exact weight lower bounds before the first serving forward.""" + if self.pipeline is None: + return RankResidencyReport( + rank=self.rank, + budget_bytes=0, + estimated_peak_bytes=None, + candidates=[], + skip_reason="pipeline not initialized", + ) + + candidates = self._collect_auto_residency_targets(workload) + if not candidates: + return RankResidencyReport( + rank=self.rank, + budget_bytes=self._auto_residency_budget_bytes(), + estimated_peak_bytes=None, + candidates=[], + skip_reason="no eligible residency alternatives", + ) + modules = self._auto_residency_modules() + active_weights = component_runtime_weight_bytes(modules) + current_device_weights = component_current_device_weight_bytes(modules) + permanent_components = { + candidate.component_name + for candidate in candidates + if candidate.current_placement and candidate.permanent_residency + } + baseline_allocated = int(torch.get_device_module().memory_allocated()) + phase_peaks = {} + phase_active = {} + phase_used = {} + for component_name in sorted({c.component_name for c in candidates}): + phase_name = f"static:{component_name}" + phase_peaks[phase_name] = baseline_allocated + ( + 0 + if component_name in permanent_components + else active_weights.get(component_name, 0) + ) + phase_active[phase_name] = tuple( + sorted(permanent_components | {component_name}) + ) + phase_used[phase_name] = (component_name,) + + local_worker_count = max( + 1, self.server_args.num_gpus // self.server_args.nnodes + ) + pin_budget = self.server_args.host_pin_budget() + return RankResidencyReport( + rank=self.rank, + budget_bytes=self._auto_residency_budget_bytes(), + estimated_peak_bytes=max(phase_peaks.values()), + estimated_peak_bytes_by_phase=phase_peaks, + active_components_by_phase=phase_active, + used_components_by_phase=phase_used, + current_device_weight_bytes_by_component=current_device_weights, + current_active_weight_bytes_by_component=active_weights, + node_rank=self.server_args.node_rank, + pinned_host_bytes=layerwise_pinned_host_bytes( + modules, pin_budget=pin_budget + ), + host_pin_capacity_bytes=layerwise_host_pin_capacity_bytes( + modules, pin_budget=pin_budget + ), + host_transition_headroom_bytes=max( + 0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES + ) + // local_worker_count, + device_transition_allocated_bytes=baseline_allocated, + candidates=candidates, + require_feasible_placement=True, + ) + + def _build_auto_residency_report( + self, *, workload: DefaultWorkload, records: List[WarmupMemoryRecord] + ) -> RankResidencyReport: + 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, + ) + modules = self._auto_residency_modules() + target_units = workload.workload_units() + assert target_units is not None + warmup_oom = any( + not record.succeeded and record.workload_units() <= target_units + for record in records + ) + runtime_weights_by_component = component_runtime_weight_bytes(modules) + if warmup_oom: + ( + estimated_peak_bytes, + estimated_phase_peaks, + active_components_by_phase, + used_components_by_phase, + full_weight_transition_components_by_phase, + ) = measured_failed_workload_phase_peaks( + records=records, + target_units=target_units, + ) + else: + estimated_peak_bytes = estimate_default_workload_peak_bytes( + records=records, + target_units=target_units, + constant_weight_bytes=max( + runtime_weights_by_component.values(), default=0 + ), + ) + ( + 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, + component_weight_bytes=runtime_weights_by_component, + ) + budget_bytes = self._auto_residency_budget_bytes() + candidates = self._collect_auto_residency_targets(workload) + if warmup_oom: + logger.debug( + "Auto residency OOM report on rank %d: phase_peaks_gib=%s, " + "active=%s, used=%s, candidate_deltas_mib=%s", + self.rank, + { + phase: round(peak / GIB_BYTES, 3) + for phase, peak in estimated_phase_peaks.items() + }, + active_components_by_phase, + used_components_by_phase, + [ + ( + candidate.option_key(), + round(candidate.active_device_delta_bytes / 1024**2), + round(candidate.inactive_device_delta_bytes / 1024**2), + round(candidate.device_transition_delta_bytes / 1024**2), + candidate.current_placement, + ) + for candidate in candidates + if is_dit_component_name(candidate.component_name) + ], + ) + ( + measured_request_duration_ns, + measured_stage_duration_ns, + measured_component_stages, + ) = 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" + reference_request_duration_ns = ( + self._auto_residency_reference_request_duration_ns ) - if reference_request_duration_ns is None and measured_request_duration_ns > 0: + if ( + not warmup_oom + and 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 ) + self._auto_residency_reference_stage_duration_ns = ( + measured_stage_duration_ns + ) + self._auto_residency_reference_component_stages = measured_component_stages 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, + stage_duration_ns=( + self._auto_residency_reference_stage_duration_ns + or measured_stage_duration_ns + ), + component_stages=( + self._auto_residency_reference_component_stages + or measured_component_stages + ), ) local_worker_count = max( 1, self.server_args.num_gpus // self.server_args.nnodes ) + pin_budget = self.server_args.host_pin_budget() return RankResidencyReport( rank=self.rank, budget_bytes=budget_bytes, estimated_peak_bytes=estimated_peak_bytes, + planning_headroom_correction_bytes=( + self._auto_residency_budget_correction_bytes + ), target_workload_measured=any( record.succeeded and record.workload_units() >= target_units for record in records @@ -1567,12 +1768,15 @@ def _build_auto_residency_report( full_weight_transition_components_by_phase ), current_device_weight_bytes_by_component=( - component_current_device_weight_bytes(self.pipeline.modules) + component_current_device_weight_bytes(modules) ), + current_active_weight_bytes_by_component=runtime_weights_by_component, node_rank=self.server_args.node_rank, - pinned_host_bytes=layerwise_pinned_host_bytes(self.pipeline.modules), + pinned_host_bytes=layerwise_pinned_host_bytes( + modules, pin_budget=pin_budget + ), host_pin_capacity_bytes=layerwise_host_pin_capacity_bytes( - self.pipeline.modules + modules, pin_budget=pin_budget ), host_transition_headroom_bytes=max( 0, host_memory_available_bytes() - HOST_COPY_RESERVE_BYTES @@ -1585,6 +1789,7 @@ def _build_auto_residency_report( measured_request_duration_ns=measured_request_duration_ns, candidate_latency_savings_ns=candidate_latency_savings_ns, candidates=candidates, + warmup_oom=warmup_oom, skip_reason=skip_reason, ) @@ -1605,6 +1810,13 @@ def _mixed_dtype_residency_components() -> set[str]: return set() return manager.components_with_mixed_use_dtypes() + def _auto_residency_modules(self) -> dict[str, object]: + assert self.pipeline is not None + manager = peek_global_component_residency_manager() + if manager is None: + return dict(self.pipeline.modules) + return manager.placement_modules() + def _rollback_applied_residency_changes( self, *, latest_round_only: bool = False ) -> str | None: @@ -1626,7 +1838,7 @@ def _rollback_applied_residency_changes( try: rollback_residency_changes( applied=applied, - modules=self.pipeline.modules, + modules=self._auto_residency_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 f9b64481ad48..8669e286baaf 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,10 +16,11 @@ 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. +auto policy provides the initial state. This module then optimizes the +long-lived serving state and validates the transition with a post-placement +warmup. A failed default-workload probe can demote that initial state and retry; +successful probes may promote or rebalance it. It does not force a single +placement to serve two different lifecycle objectives. """ from __future__ import annotations @@ -80,23 +81,29 @@ MEASURED_VRAM_RESERVE_FRACTION = 0.05 EXTRAPOLATED_VRAM_RESERVE_FRACTION = 0.10 MIN_VRAM_RESERVE_BYTES = 4 * GIB_BYTES +# A failed or not-yet-run target workload still needs a useful activation +# allowance on small cards. Three GiB covered the measured image-model decode +# and denoise gap without applying the datacenter-card 4 GiB floor verbatim. +MIN_UNCALIBRATED_VRAM_RESERVE_BYTES = 3 * GIB_BYTES # 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 +MAX_MEASURED_VRAM_RESERVE_FRACTION = 0.10 -# 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. +# The transfer model ranks feasible placements; the post-adjustment warmup is +# the authority on whether a selected placement actually helped. Do not apply +# a request-duration-relative utility tolerance here: on long video requests it +# can hide an entire encoder transfer and leave useful VRAM idle. ESTIMATED_PINNED_H2D_BYTES_PER_SECOND = 24 * GIB_BYTES -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. +# Pageable CUDA copies first stage through an internal pinned buffer and do not +# overlap layer compute like the explicit pinned path. H200 measurements put +# their effective layer-streaming cost near 3x the pinned path. +PAGEABLE_H2D_COST_MULTIPLIER = 3 +AUTO_PLACEMENT_LATENCY_TOLERANCE_NS = 0 +# Allow normal measurement noise, but undo a round whose calibrated request is +# materially slower than the original layout. +MIN_POST_ADJUSTMENT_REGRESSION_NS = 100_000_000 POST_ADJUSTMENT_REGRESSION_FRACTION = 0.05 PLACEMENT_STATUS_SKIPPED = "skipped" @@ -232,12 +239,17 @@ class RankResidencyReport(msgspec.Struct, frozen=True): rank: int budget_bytes: int estimated_peak_bytes: int | None + # Conservative margin learned from a failed candidate. It constrains only + # future growth: the current placement was measured against the physical + # budget and must not become retroactively infeasible after a rollback. + planning_headroom_correction_bytes: int = 0 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, ...]] = {} full_weight_transition_components_by_phase: dict[str, tuple[str, ...]] = {} current_device_weight_bytes_by_component: dict[str, int] = {} + current_active_weight_bytes_by_component: dict[str, int] = {} node_rank: int = 0 pinned_host_bytes: int = 0 host_pin_capacity_bytes: int = 0 @@ -247,6 +259,12 @@ class RankResidencyReport(msgspec.Struct, frozen=True): measured_request_duration_ns: int = 0 candidate_latency_savings_ns: dict[str, int] = {} candidates: list[ResidencyTarget] = [] + # The current placement could not execute the default workload. Its phase + # budgets are intentionally allowed to be negative so the joint solver + # must select a lower-memory complete target state instead of treating the + # measured zero-delta placement as feasible. + warmup_oom: bool = False + require_feasible_placement: bool = False skip_reason: str | None = None @@ -257,7 +275,9 @@ class AutoResidencyPlan(msgspec.Struct, frozen=True): reserve_bytes: int = 0 budget_bytes: int = 0 resource_budget_bytes: dict[str, int] = {} + resource_delta_bytes: dict[str, int] = {} changes: list[ResidencyTarget] = [] + recovering_from_oom: bool = False skip_reason: str | None = None current_placement_reserve_shortfall_bytes: int = 0 @@ -368,7 +388,11 @@ def estimate_default_workload_timing( ) for phase_name, components in phase_components.items(): fields = phase_name.split(":", 2) - if len(fields) < 2 or not fields[0].isdigit(): + if ( + len(fields) < 3 + or not fields[0].isdigit() + or not fields[2].startswith("use:") + ): continue stage_name = fields[1] if stage_name not in representative.stage_duration_ms: @@ -435,25 +459,89 @@ def estimate_candidate_latency_savings_ns( *, candidates: Iterable[ResidencyTarget], request_duration_ns: int, + stage_duration_ns: Mapping[str, int] | None = None, + component_stages: Mapping[str, tuple[str, ...]] | None = None, ) -> dict[str, int]: - """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. + """Estimate latency removed by each complete placement option. + + Transfer work remains the ordering signal within one component frontier. + Across components, cap the frontier's total benefit by the stages where + that component is used. Otherwise two large DiTs in different stages both + saturate at the full request duration and the solver cannot distinguish + which stage dominates latency. Scaling the intermediate options against + the frontier maximum preserves their ordering without claiming that every + byte of asynchronous transfer stalls the critical path. """ + candidates = list(candidates) + stage_duration_ns = stage_duration_ns or {} + component_stages = component_stages or {} + maximum_h2d_bytes: dict[str, int] = {} + for candidate in candidates: + maximum_h2d_bytes[candidate.component_name] = max( + maximum_h2d_bytes.get(candidate.component_name, 0), + candidate.h2d_bytes_per_request, + ) + estimates: dict[str, int] = {} for candidate in candidates: - transfer_ns = int( - candidate.h2d_bytes_per_request + component_maximum_bytes = maximum_h2d_bytes[candidate.component_name] + if component_maximum_bytes <= 0: + estimates[candidate.option_key()] = int( + candidate.h2d_bytes_per_request + / ESTIMATED_PINNED_H2D_BYTES_PER_SECOND + * 1_000_000_000 + ) + continue + + transfer_upper_bound_ns = int( + component_maximum_bytes / ESTIMATED_PINNED_H2D_BYTES_PER_SECOND * 1_000_000_000 ) - estimates[candidate.option_key()] = ( - min(transfer_ns, request_duration_ns) - if transfer_ns > 0 and request_duration_ns > 0 - else transfer_ns + component_stage_duration_ns = sum( + stage_duration_ns.get(stage_name, 0) + for stage_name in component_stages.get(candidate.component_name, ()) ) + # Stage profiling is intentionally asynchronous by default. That is a + # useful low-overhead measurement for one-shot encoder/decoder work, + # but it records mostly CPU launch time for a repeatedly streamed DiT + # and can understate layerwise H2D stalls by orders of magnitude. The + # transfer model is already an upper bound, so constrain DiT savings by + # the synchronized request wall time instead. + latency_upper_bound_ns = ( + request_duration_ns + if is_dit_component_name(candidate.component_name) + else component_stage_duration_ns or request_duration_ns + ) + if latency_upper_bound_ns > 0: + transfer_upper_bound_ns = min( + transfer_upper_bound_ns, latency_upper_bound_ns + ) + if is_dit_component_name(candidate.component_name): + # Model the remaining layer traffic, not a linearly scaled share + # of the worst placement. When an all-pageable frontier exceeds + # the request wall time, linear scaling compresses every marginal + # resident/pinned-layer gain and can discard hundreds of + # milliseconds of useful HostPin placement. + remaining_transfer_bytes = max( + 0, component_maximum_bytes - candidate.h2d_bytes_per_request + ) + remaining_transfer_ns = int( + remaining_transfer_bytes + / ESTIMATED_PINNED_H2D_BYTES_PER_SECOND + * 1_000_000_000 + ) + estimates[candidate.option_key()] = max( + 0, + transfer_upper_bound_ns + - min(remaining_transfer_ns, transfer_upper_bound_ns), + ) + else: + estimates[candidate.option_key()] = int( + transfer_upper_bound_ns + * candidate.h2d_bytes_per_request + / component_maximum_bytes + ) return estimates @@ -646,6 +734,80 @@ def estimate_workload_phase_peaks( ) +def measured_failed_workload_phase_peaks( + *, + records: Iterable[WarmupMemoryRecord], + target_units: int, +) -> tuple[ + int | None, + dict[str, int], + dict[str, tuple[str, ...]], + dict[str, tuple[str, ...]], + dict[str, tuple[str, ...]], +]: + """Preserve the binding phases reached before a target-workload OOM. + + OOM records cannot support activation extrapolation, but their allocated + peaks and component ownership still identify which complete target states + release memory. Repeated phase names with different layouts remain + separate constraints, matching the successful-measurement path. + """ + failed = [ + record + for record in records + if not record.succeeded and record.workload_units() <= target_units + ] + if not failed: + return None, {}, {}, {}, {} + + grouped: dict[ + tuple[str, tuple[str, ...], tuple[str, ...], tuple[str, ...]], int + ] = {} + for record in failed: + used_by_phase = record.phase_used_components or record.phase_active_components + phases = record.phase_peak_allocated_bytes or { + "request:failed": record.peak_allocated_bytes + } + for phase_name, phase_peak in phases.items(): + active = tuple(sorted(record.phase_active_components.get(phase_name, ()))) + used = tuple(sorted(used_by_phase.get(phase_name, ()))) + transitions = tuple( + sorted( + record.phase_full_weight_transition_components.get(phase_name, ()) + ) + ) + key = (phase_name, active, used, transitions) + grouped[key] = max(grouped.get(key, 0), phase_peak) + + layout_counts: dict[str, int] = {} + for phase_name, _, _, _ in grouped: + layout_counts[phase_name] = layout_counts.get(phase_name, 0) + 1 + + peaks: dict[str, int] = {} + active_components: dict[str, tuple[str, ...]] = {} + used_components: dict[str, tuple[str, ...]] = {} + transition_components: dict[str, tuple[str, ...]] = {} + layout_indices: dict[str, int] = {} + for (phase_name, active, used, transitions), phase_peak in sorted(grouped.items()): + output_name = phase_name + if layout_counts[phase_name] > 1: + index = layout_indices.get(phase_name, 0) + layout_indices[phase_name] = index + 1 + output_name = f"{phase_name}:layout:{index}" + peaks[output_name] = phase_peak + active_components[output_name] = active + used_components[output_name] = used + transition_components[output_name] = transitions + + return ( + max(record.peak_allocated_bytes for record in failed), + peaks, + active_components, + used_components, + transition_components, + ) + + def _module_weight_bytes(module: nn.Module) -> int: """Full weight+buffer footprint, reading through layerwise CPU buffers.""" return tensor_storage_bytes( @@ -696,7 +858,9 @@ def component_current_device_weight_bytes( return result -def layerwise_pinned_host_bytes(modules: Mapping[str, object]) -> int: +def layerwise_pinned_host_bytes( + modules: Mapping[str, object], *, pin_budget=None +) -> int: """Pinned layer-store bytes in this process, without alias double counts.""" seen_managers: set[int] = set() total = 0 @@ -709,10 +873,14 @@ def layerwise_pinned_host_bytes(modules: Mapping[str, object]) -> int: continue seen_managers.add(manager_id) total += manager.pinned_host_weight_bytes() + if pin_budget is not None: + total = max(total, pin_budget.committed_bytes) return total -def layerwise_host_pin_capacity_bytes(modules: Mapping[str, object]) -> int: +def layerwise_host_pin_capacity_bytes( + modules: Mapping[str, object], *, pin_budget=None +) -> int: """This process's non-overlapping HostPin allowance.""" seen_budgets: set[int] = set() total = 0 @@ -726,6 +894,8 @@ def layerwise_host_pin_capacity_bytes(modules: Mapping[str, object]) -> int: continue seen_budgets.add(budget_id) total += max(0, budget.available_bytes - budget.reserve_bytes) + if pin_budget is not None and id(pin_budget) not in seen_budgets: + total += max(0, pin_budget.available_bytes - pin_budget.reserve_bytes) return total @@ -773,7 +943,9 @@ 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 - transfer_multiplier = 1 if layer_idx in pinned else 2 + transfer_multiplier = ( + 1 if layer_idx in pinned else PAGEABLE_H2D_COST_MULTIPLIER + ) total += uses * transfer_multiplier * weight_bytes return total @@ -963,6 +1135,7 @@ def _unconfigured_layerwise_targets( prefetch_value: float, residency_policy: str, pin_cpu_memory: bool, + allow_host_pin_reallocation: bool, ) -> tuple[list[ResidencyTarget], int]: """Virtual layerwise frontier for a component still using coarse offload. @@ -998,9 +1171,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 = sum( - manager.peak_managed_device_weight_bytes(0) for manager in managers - ) host_materialize_scratch = max( ( weight_bytes @@ -1019,46 +1189,69 @@ def _unconfigured_layerwise_targets( 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, - # 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, - 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, + pin_targets = ( + _layerwise_pin_targets( + managers=managers, + resident_layers=resident_layers, + current_pinned_layers=empty_pins, + uses_per_streamed_layer=uses_per_request, ) + if allow_host_pin_reallocation + else [empty_pins] ) + for pinned_layers in pin_targets: + 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, pinned_layers) + ) + transfer_work = _layerwise_transfer_work_bytes( + managers=managers, + resident_layers=resident_layers, + pinned_layers=pinned_layers, + uses_per_streamed_layer=uses_per_request, + ) + targets.append( + ResidencyTarget( + component_name=component_name, + residency_mode=current_mode, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=resident_bytes, + # 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=pinned_layers, + pinned_host_delta_bytes=pinned_bytes, + host_pin_scratch_bytes=pinned_bytes, + host_materialize_scratch_bytes=host_materialize_scratch, + device_transition_delta_bytes=( + unmanaged_weight_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, + target_pinned_host_bytes=pinned_bytes, + ) + ) return targets, maximum_transfer_work @@ -1092,7 +1285,7 @@ def collect_residency_targets( num_inference_steps: int, allow_host_pin_reallocation: bool = True, mixed_dtype_components: Iterable[str] = (), - auto_resident_components: Iterable[str] = (), + required_resident_components: Iterable[str] = (), layerwise_tuning_of: Callable[[str, bool], tuple[float, float, str]] | None = None, pin_cpu_memory: bool = True, ) -> list[ResidencyTarget]: @@ -1105,7 +1298,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) + required_resident_names = set(required_resident_components) if baseline_residency_mode_of is None: baseline_residency_mode_of = residency_mode_of candidates = [] @@ -1116,15 +1309,16 @@ def collect_residency_targets( if name in custom_names: continue baseline_mode = baseline_residency_mode_of(name) - if baseline_mode not in (COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD): + if baseline_mode not in (COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD, RESIDENT): continue if explicit_residency_mode_of(name) is not None: continue if is_fsdp_managed_module(module): continue 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. + if current_mode == RESIDENT and name in required_resident_names: + # A runtime feature owns this hard requirement. Unset startup + # residency is only an initial placement and remains tunable. continue has_layerwise_managers = isinstance( @@ -1135,7 +1329,14 @@ def collect_residency_targets( # 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 + # A configured manager exposes its measured layerwise frontier. Without + # managers, the module is still in its ordinary current placement even + # when the baseline policy requested layerwise offload. Build a virtual + # frontier from that real state so an initial resident seed can still + # be demoted after an OOM. + frontier_mode = ( + LAYERWISE_OFFLOAD if has_layerwise_managers else COMPONENT_OFFLOAD + ) if frontier_mode == COMPONENT_OFFLOAD: if current_mode not in (COMPONENT_OFFLOAD, RESIDENT): @@ -1168,6 +1369,7 @@ def collect_residency_targets( prefetch_value=prefetch_value, residency_policy=residency_policy, pin_cpu_memory=pin_cpu_memory, + allow_host_pin_reallocation=allow_host_pin_reallocation, ) ) uses_per_request = ( @@ -1182,7 +1384,7 @@ def collect_residency_targets( [ ResidencyTarget( component_name=name, - residency_mode=baseline_mode, + residency_mode=current_mode, target_residency_mode=COMPONENT_OFFLOAD, target_resident_weight_bytes=0, h2d_bytes_per_request=max( @@ -1205,7 +1407,7 @@ def collect_residency_targets( ), ResidencyTarget( component_name=name, - residency_mode=baseline_mode, + residency_mode=current_mode, target_residency_mode=RESIDENT, target_resident_weight_bytes=weight_bytes, h2d_bytes_per_request=(maximum_transfer_work or weight_bytes), @@ -1473,9 +1675,13 @@ def rank_candidates_by_h2d_savings( def _skip_plan( - reason: str, *, current_placement_reserve_shortfall_bytes: int = 0 + reason: str, + *, + recovering_from_oom: bool = False, + current_placement_reserve_shortfall_bytes: int = 0, ) -> AutoResidencyPlan: return AutoResidencyPlan( + recovering_from_oom=recovering_from_oom, skip_reason=reason, current_placement_reserve_shortfall_bytes=( current_placement_reserve_shortfall_bytes @@ -1489,13 +1695,23 @@ def _vram_reserve_bytes(budget_bytes: int, *, target_workload_measured: bool) -> if target_workload_measured else EXTRAPOLATED_VRAM_RESERVE_FRACTION ) - return max( - int(budget_bytes * reserve_fraction), - min( - MIN_VRAM_RESERVE_BYTES, - int(budget_bytes * MAX_VRAM_RESERVE_FRACTION), + floor_fraction = ( + MAX_MEASURED_VRAM_RESERVE_FRACTION + if target_workload_measured + else MAX_VRAM_RESERVE_FRACTION + ) + minimum_floor = ( + 0 if target_workload_measured else MIN_UNCALIBRATED_VRAM_RESERVE_BYTES + ) + capped_floor = min( + MIN_VRAM_RESERVE_BYTES, + budget_bytes, + max( + minimum_floor, + int(budget_bytes * floor_fraction), ), ) + return max(int(budget_bytes * reserve_fraction), capped_floor) def _binding_phase_constraints( @@ -1535,6 +1751,7 @@ def _binding_phase_constraints( else max(phase_peaks.values()) ) current_device_bytes = dict(report.current_device_weight_bytes_by_component) + current_active_bytes = dict(report.current_active_weight_bytes_by_component) for candidate in report.candidates: if ( candidate.current_placement @@ -1544,11 +1761,17 @@ def _binding_phase_constraints( candidate.target_device_weight_bytes ) steady_state_components = { + candidate.component_name + for candidate in report.candidates + if candidate.current_placement and candidate.permanent_residency + } + steady_state_components.update( 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 - } + if component_name not in candidate_component_names + ) binding: dict[ tuple[ @@ -1559,13 +1782,45 @@ def _binding_phase_constraints( ], tuple[str, int], ] = {} + attributed_baselines = [] for phase_name, phase_peak in phase_peaks.items(): measured_components = set(report.active_components_by_phase.get(phase_name, ())) + used = tuple(sorted(report.used_components_by_phase.get(phase_name, ()))) + full_weight_transitions = tuple( + sorted( + set( + report.full_weight_transition_components_by_phase.get( + phase_name, () + ) + ) + & candidate_component_names + ) + ) carried_components = steady_state_components - measured_components phase_peak += sum( current_device_bytes.get(component_name, 0) for component_name in carried_components ) + attributed_baselines.append( + max( + 0, + phase_peak + - sum( + current_active_bytes.get(component_name, 0) + for component_name in ( + (measured_components - steady_state_components) + | set(used) + | set(full_weight_transitions) + ) + ), + ) + ) + # An OOM can unwind before the component manager publishes its phase + # snapshot. The resulting request:untracked peak is useful as a + # conservative baseline for the synthetic component phases below, but + # no placement can release an empty-layout constraint directly. + if report.warmup_oom and not used and not full_weight_transitions: + continue present = tuple( sorted( (measured_components | steady_state_components) @@ -1573,17 +1828,6 @@ def _binding_phase_constraints( ) ) measured = tuple(sorted(measured_components & candidate_component_names)) - used = tuple(sorted(report.used_components_by_phase.get(phase_name, ()))) - full_weight_transitions = tuple( - sorted( - set( - report.full_weight_transition_components_by_phase.get( - phase_name, () - ) - ) - & candidate_component_names - ) - ) layout = (present, measured, used, full_weight_transitions) current = binding.get(layout) if current is None or (phase_peak, phase_name) > (current[1], current[0]): @@ -1593,20 +1837,43 @@ def _binding_phase_constraints( 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()), + # Separate the stable/request activation baseline from transient component + # weights before constructing phases that a short failed probe did not + # reach. Reusing the largest component-prefetch peak verbatim would charge + # those weights to every later component and can make every lower-memory + # target look infeasible. The post-placement warmup remains the authority + # for previously unobserved activation memory. + synthetic_base_peak = ( + max(attributed_baselines) if attributed_baselines else request_peak + ) + synthetic_components = ( + candidate_component_names + if report.warmup_oom or report.require_feasible_placement + else candidate_component_names - observed_components ) - for component_name in sorted(candidate_component_names - observed_components): + for component_name in sorted(synthetic_components): present = tuple( sorted( (steady_state_components | {component_name}) & candidate_component_names ) ) - binding[(present, (), (component_name,), ())] = ( - f"unobserved:{component_name}", - steady_request_peak, + layout = (present, (), (component_name,), ()) + synthetic = ( + ( + f"lower-bound:{component_name}" + if component_name in observed_components + else f"unobserved:{component_name}" + ), + synthetic_base_peak + + ( + 0 + if component_name in steady_state_components + else current_active_bytes.get(component_name, 0) + ), ) + current = binding.get(layout) + if current is None or synthetic[1] > current[1]: + binding[layout] = synthetic return [ ( f"gpu:rank{report.rank}:{phase_name}", @@ -1726,18 +1993,25 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP """Turn gathered rank reports into one deterministic placement plan.""" if not reports: return _skip_plan("no rank reports") + recovering_from_oom = any(report.warmup_oom for report in reports) for report in reports: if report.skip_reason is not None: - return _skip_plan(f"rank {report.rank}: {report.skip_reason}") + return _skip_plan( + f"rank {report.rank}: {report.skip_reason}", + recovering_from_oom=recovering_from_oom, + ) if report.estimated_peak_bytes is None: - return _skip_plan(f"rank {report.rank}: no usable warmup measurement") + return _skip_plan( + f"rank {report.rank}: no usable warmup measurement", + recovering_from_oom=recovering_from_oom, + ) 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( + target_workload_measured = not recovering_from_oom and any( report.target_workload_measured for report in reports ) reserves_by_rank = { @@ -1767,6 +2041,7 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP if not candidates: return _skip_plan( "no eligible residency alternatives", + recovering_from_oom=recovering_from_oom, current_placement_reserve_shortfall_bytes=( current_placement_reserve_shortfall ), @@ -1799,9 +2074,21 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP # 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], + available = ( + report.budget_bytes + - phase_peak + - reserves_by_rank[report.rank] + - report.planning_headroom_correction_bytes + ) + # A successful measured placement remains a feasible zero-delta + # baseline even when it consumed the reserve. An OOM placement is + # different: preserving the negative shortfall forces the complete + # target-state solve to release enough memory before retrying. + enforce_rank_feasibility = ( + report.warmup_oom or report.require_feasible_placement + ) + resource_budgets[resource_name] = ( + available if enforce_rank_feasibility else max(0, available) ) has_device_transition_options = any( candidate.device_transition_delta_bytes != 0 for candidate in candidates @@ -1875,22 +2162,6 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP and report.candidate_latency_savings_ns ] use_latency_utility = bool(timed_reports) - latency_equivalence_ns = ( - max( - 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 - else 0 - ) options = [] for candidate in candidates: resource_deltas: dict[str, int] = {} @@ -1978,24 +2249,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 - ), - placement_cost_bytes=( + estimated_latency_savings=estimated_latency_savings, + preference_cost=( + ( + 0 + if candidate.current_placement + or candidate.target_residency_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), ), @@ -2015,14 +2292,37 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP placement = optimize_placement( options, resource_budget_bytes=resource_budgets, - estimated_latency_tolerance=( - latency_equivalence_ns if use_latency_utility else 0 - ), + estimated_latency_tolerance=AUTO_PLACEMENT_LATENCY_TOLERANCE_NS, require_selection_from_every_group=complete_state_frontier, ) except NoFeasiblePlacementError as error: + diagnostic_options = [ + option + for option in options + if candidate_by_key[option.option_key].current_placement + or ( + is_dit_component_name( + candidate_by_key[option.option_key].component_name + ) + and candidate_by_key[option.option_key].target_layerwise_resident_layers + == (0,) + ) + ] + logger.debug( + "Auto residency infeasible budgets_mib=%s diagnostic_options=%s", + {name: round(value / 1024**2) for name, value in resource_budgets.items()}, + { + option.option_key: { + name: round(value / 1024**2) + for name, value in option.resource_delta_bytes.items() + if value + } + for option in diagnostic_options + }, + ) return _skip_plan( str(error), + recovering_from_oom=recovering_from_oom, current_placement_reserve_shortfall_bytes=( current_placement_reserve_shortfall ), @@ -2039,7 +2339,9 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP reserve_bytes=reserve, budget_bytes=budget, resource_budget_bytes=resource_budgets, + resource_delta_bytes=placement.resource_delta_bytes, changes=changes, + recovering_from_oom=recovering_from_oom, current_placement_reserve_shortfall_bytes=(current_placement_reserve_shortfall), ) @@ -2442,7 +2744,8 @@ def format_plan_summary( ) -> str: """One-line decision summary for the startup log.""" if plan.skip_reason is not None: - return f"Auto residency: skipped ({plan.skip_reason})" + action = "recovery stopped" if plan.recovering_from_oom else "skipped" + return f"Auto residency: {action} ({plan.skip_reason})" changes = ( ", ".join(_format_candidate_summary(candidate) for candidate in plan.changes) or "none" @@ -2454,8 +2757,9 @@ def format_plan_summary( if record.succeeded ) measured_part = f"measured_allocated=[{measured}], " if measured else "" + prefix = "Auto residency recovery" if plan.recovering_from_oom else "Auto residency" return ( - f"Auto residency: target={workload.describe()} " + f"{prefix}: target={workload.describe()} " f"steps={workload.num_inference_steps}, " f"{measured_part}" f"estimated_peak={plan.estimated_peak_bytes / GIB_BYTES:.1f} GiB, " @@ -2536,4 +2840,5 @@ def plan_summary_payload(*, plan: AutoResidencyPlan, status: str) -> dict: return { "status": status, "changed": [candidate.component_name for candidate in plan.changes], + "recovering_from_oom": plan.recovering_from_oom, } diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_loading_order.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_loading_order.py index 890eaf92e74e..dc28790d9d57 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_loading_order.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_loading_order.py @@ -6,11 +6,11 @@ module only ranks already-selected load specs. """ -import glob -import json -import os from dataclasses import dataclass +from sglang.multimodal_gen.runtime.managers.memory_managers.component_weight_inventory import ( + infer_component_weight_size_bytes, +) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( is_dit_component_name, is_image_encoder_component_name, @@ -31,9 +31,6 @@ class ComponentLoadSpec: index: int -_WEIGHT_FILE_SUFFIXES = (".bin", ".pt", ".pth") - - def _component_base_name(component_name: str) -> str: prefix, separator, suffix = component_name.rpartition("_") if separator and suffix.isdigit(): @@ -62,101 +59,6 @@ def component_load_risk_rank(component_name: str) -> int: return 10 -def _safe_file_size(file_path: str) -> int | None: - try: - return os.path.getsize(file_path) - except OSError: - return None - - -def _safetensors_payload_size_bytes(file_path: str) -> int | None: - try: - with open(file_path, "rb") as f: - header_size_bytes = f.read(8) - if len(header_size_bytes) != 8: - return _safe_file_size(file_path) - header_size = int.from_bytes(header_size_bytes, "little") - header = json.loads(f.read(header_size)) - except (OSError, json.JSONDecodeError, ValueError): - return _safe_file_size(file_path) - - payload_size = 0 - for tensor_name, tensor_info in header.items(): - if tensor_name == "__metadata__": - continue - offsets = tensor_info.get("data_offsets") - if not isinstance(offsets, list) or len(offsets) != 2: - return _safe_file_size(file_path) - payload_size += offsets[1] - offsets[0] - return payload_size - - -def _safetensors_files_from_index(component_model_path: str) -> list[str]: - indexed_files: set[str] = set() - index_paths = sorted( - glob.glob(os.path.join(component_model_path, "*.safetensors.index.json")) - ) - for index_path in index_paths: - try: - with open(index_path) as f: - weight_map = json.load(f).get("weight_map", {}) - except (OSError, json.JSONDecodeError): - continue - for shard_name in weight_map.values(): - shard_path = os.path.join(component_model_path, shard_name) - if os.path.isfile(shard_path): - indexed_files.add(shard_path) - return sorted(indexed_files) - - -def _list_component_safetensors_files(component_model_path: str) -> list[str]: - if os.path.isfile(component_model_path): - if component_model_path.endswith(".safetensors"): - return [component_model_path] - return [] - if not os.path.isdir(component_model_path): - return [] - - indexed_files = _safetensors_files_from_index(component_model_path) - if indexed_files: - return indexed_files - return sorted(glob.glob(os.path.join(component_model_path, "*.safetensors"))) - - -def infer_component_weight_size_bytes(component_model_path: str) -> int | None: - """Infer checkpoint payload size from safetensors without materializing tensors.""" - safetensors_files = _list_component_safetensors_files(component_model_path) - if safetensors_files: - sizes = [ - size - for size in ( - _safetensors_payload_size_bytes(file_path) - for file_path in safetensors_files - ) - if size is not None - ] - return sum(sizes) if sizes else None - - if os.path.isfile(component_model_path): - if component_model_path.endswith(_WEIGHT_FILE_SUFFIXES): - return _safe_file_size(component_model_path) - return None - if not os.path.isdir(component_model_path): - return None - - weight_files = [] - for suffix in _WEIGHT_FILE_SUFFIXES: - weight_files.extend(glob.glob(os.path.join(component_model_path, f"*{suffix}"))) - if not weight_files: - return None - sizes = [ - size - for size in (_safe_file_size(file_path) for file_path in weight_files) - if size is not None - ] - return sum(sizes) if sizes else None - - def order_component_load_specs( component_specs: list[ComponentLoadSpec], ) -> list[ComponentLoadSpec]: 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..d7731c5402bc 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 @@ -143,6 +143,8 @@ def __init__( self._warmup_phase_used_components: tuple[str, ...] = () self._warmup_phase_full_weight_transition_components: tuple[str, ...] = () self._warmup_phase_peaks: dict[str, WarmupPhasePeak] = {} + self._warmup_phase_occurrences: dict[str, int] = {} + self._failed_warmup_phase_peaks: dict[str, WarmupPhasePeak] = {} self._completed_warmup_phase_peaks: dict[str, WarmupPhasePeak] = {} def refresh_pipeline(self, pipeline: ComponentResidencyPipeline) -> None: @@ -232,6 +234,8 @@ def begin_request( self._warmup_phase_used_components = () self._warmup_phase_full_weight_transition_components = () self._warmup_phase_peaks = {} + self._warmup_phase_occurrences = {} + self._failed_warmup_phase_peaks = {} self._completed_warmup_phase_peaks = {} if self._track_warmup_memory: # GPUWorker reset the request peak before entering the pipeline. @@ -298,6 +302,7 @@ def full_weight_transition(self, component_names: Iterable[str]) -> Iterator[Non components=components, used_components=used_components, full_weight_transition_components=transitions, + new_occurrence=False, ) def begin_stage(self) -> None: @@ -657,10 +662,15 @@ def _begin_warmup_phase( components: tuple[str, ...], used_components: tuple[str, ...], full_weight_transition_components: tuple[str, ...] = (), + new_occurrence: bool = True, ) -> None: if not self._track_warmup_memory: return self._record_warmup_phase_peak() + if new_occurrence: + occurrence = self._warmup_phase_occurrences.get(key, 0) + self._warmup_phase_occurrences[key] = occurrence + 1 + key = key if occurrence == 0 else f"{key}:occurrence:{occurrence}" self._warmup_phase_key = key self._warmup_phase_components = tuple(sorted(set(components))) self._warmup_phase_used_components = tuple(sorted(set(used_components))) @@ -762,10 +772,19 @@ def take_warmup_phase_peaks( self, ) -> dict[str, WarmupPhasePeak]: """Return and clear the most recently completed warmup phase peaks.""" - peaks = self._completed_warmup_phase_peaks + peaks = dict(self._failed_warmup_phase_peaks) + peaks.update(self._completed_warmup_phase_peaks) + self._failed_warmup_phase_peaks = {} self._completed_warmup_phase_peaks = {} return peaks + def capture_failed_warmup_phase(self) -> None: + """Preserve failure attribution before request cleanup changes phase.""" + if not self._track_warmup_memory: + return + self._record_warmup_phase_peak() + self._failed_warmup_phase_peaks = dict(self._warmup_phase_peaks) + def current_device_components(self) -> tuple[str, ...]: """Components whose complete module is currently on the device. @@ -775,6 +794,19 @@ def current_device_components(self) -> tuple[str, ...]: """ return self._warmup_active_components() + def placement_modules(self) -> dict[str, object]: + """Return every module available to post-request placement planning. + + Some stages materialize request-dependent components lazily, so those + modules are absent from ``pipeline.modules`` until they are first used. + ``_modules_seen`` retains the concrete module passed through the + residency manager for the current request. Pipeline-owned modules win + if a component was replaced after it was observed. + """ + modules = dict(self._modules_seen) + modules.update(self.pipeline.modules) + return modules + 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/component_residency_strategies.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py index ee60f09e1040..0fb74f0df85e 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 @@ -130,7 +130,24 @@ def prepare_for_use( use: ComponentUse, state: ResidencyState, ) -> None: - _module_to_local_device(module, dtype=use.target_dtype) + try: + _module_to_local_device(module, dtype=use.target_dtype) + except BaseException: + self._restore_cpu_after_failed_transfer(module) + raise + + def _restore_cpu_after_failed_transfer( + self, module: nn.Module, *, stream: object | None = None + ) -> None: + """Undo a partial ``Module.to(device)`` after allocation failure.""" + if stream is None: + module.to("cpu", non_blocking=False) + else: + with torch.get_device_module().stream(stream): + module.to("cpu", non_blocking=False) + stream.synchronize() + if current_platform.is_cuda(): + torch.get_device_module().empty_cache() def wait_for_use( self, @@ -158,10 +175,17 @@ def prefetch_for_use( self._prefetch_stream = torch.get_device_module().Stream( device=get_local_torch_device() ) - with torch.get_device_module().stream(self._prefetch_stream): - _module_to_local_device(module, dtype=use.target_dtype) - event = torch.get_device_module().Event() - event.record(self._prefetch_stream) + try: + with torch.get_device_module().stream(self._prefetch_stream): + _module_to_local_device(module, dtype=use.target_dtype) + event = torch.get_device_module().Event() + event.record(self._prefetch_stream) + except BaseException: + self._ready_events.pop(use.component_name, None) + self._restore_cpu_after_failed_transfer( + module, stream=self._prefetch_stream + ) + raise self._ready_events[use.component_name] = event return True diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_weight_inventory.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_weight_inventory.py new file mode 100644 index 000000000000..555d4072f158 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_weight_inventory.py @@ -0,0 +1,229 @@ +"""Pre-load weight inventory for pipeline placement planning.""" + +import glob +import json +import math +import os +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ComponentWeightSource: + """One physical checkpoint source selected by the pipeline.""" + + component_name: str + component_model_path: str + checkpoint_bytes: int | None = None + parameter_count: int | None = None + target_element_size: int | None = None + supports_fsdp_loading: bool = False + + +@dataclass(frozen=True) +class ComponentWeightEstimate: + component_name: str + component_model_path: str + checkpoint_bytes: int | None + parameter_count: int | None + target_element_size: int | None + supports_fsdp_loading: bool = False + + def materialized_bytes(self) -> int | None: + if self.parameter_count is not None and self.target_element_size is not None: + return self.parameter_count * self.target_element_size + return self.checkpoint_bytes + + +_WEIGHT_FILE_SUFFIXES = (".bin", ".pt", ".pth", ".ckpt", ".gguf") + + +def _safe_file_size(file_path: str) -> int | None: + try: + return os.path.getsize(file_path) + except OSError: + return None + + +def _read_safetensors_header(file_path: str) -> dict | None: + try: + with open(file_path, "rb") as f: + header_size_bytes = f.read(8) + if len(header_size_bytes) != 8: + return None + return json.loads(f.read(int.from_bytes(header_size_bytes, "little"))) + except (OSError, json.JSONDecodeError, ValueError): + return None + + +def _tensor_stats(tensor_info: dict) -> tuple[int, int] | None: + offsets = tensor_info.get("data_offsets") + shape = tensor_info.get("shape") + if ( + not isinstance(offsets, list) + or len(offsets) != 2 + or not isinstance(shape, list) + or any(not isinstance(dimension, int) for dimension in shape) + ): + return None + return offsets[1] - offsets[0], math.prod(shape) + + +def _safetensors_payload_stats(file_path: str) -> tuple[int, int | None] | None: + header = _read_safetensors_header(file_path) + if header is None: + size = _safe_file_size(file_path) + return None if size is None else (size, None) + + checkpoint_bytes = 0 + parameter_count = 0 + for tensor_name, tensor_info in header.items(): + if tensor_name == "__metadata__": + continue + stats = _tensor_stats(tensor_info) + if stats is None: + size = _safe_file_size(file_path) + return None if size is None else (size, None) + checkpoint_bytes += stats[0] + parameter_count += stats[1] + return checkpoint_bytes, parameter_count + + +def infer_safetensors_weight_stats_by_prefix( + file_path: str, +) -> dict[str, tuple[int, int]] | None: + """Return payload bytes and parameter count by first tensor-key segment.""" + header = _read_safetensors_header(file_path) + if header is None: + return None + + result: dict[str, tuple[int, int]] = {} + for tensor_name, tensor_info in header.items(): + if tensor_name == "__metadata__": + continue + stats = _tensor_stats(tensor_info) + if stats is None: + return None + prefix = tensor_name.split(".", 1)[0] + previous = result.get(prefix, (0, 0)) + result[prefix] = (previous[0] + stats[0], previous[1] + stats[1]) + return result + + +def _safetensors_files_from_index(component_model_path: str) -> list[str]: + indexed_files: set[str] = set() + for index_path in sorted( + glob.glob(os.path.join(component_model_path, "*.safetensors.index.json")) + ): + try: + with open(index_path) as f: + weight_map = json.load(f).get("weight_map", {}) + except (OSError, json.JSONDecodeError): + continue + for shard_name in weight_map.values(): + shard_path = os.path.join(component_model_path, shard_name) + if os.path.isfile(shard_path): + indexed_files.add(shard_path) + return sorted(indexed_files) + + +def _list_component_safetensors_files(component_model_path: str) -> list[str]: + if os.path.isfile(component_model_path): + return ( + [component_model_path] + if component_model_path.endswith(".safetensors") + else [] + ) + if not os.path.isdir(component_model_path): + return [] + + indexed_files = _safetensors_files_from_index(component_model_path) + return indexed_files or sorted( + glob.glob(os.path.join(component_model_path, "*.safetensors")) + ) + + +def infer_component_weight_stats( + component_model_path: str, +) -> tuple[int | None, int | None]: + """Infer checkpoint bytes and parameter count without loading tensors.""" + safetensors_files = _list_component_safetensors_files(component_model_path) + if safetensors_files: + stats = [ + stat + for file_path in safetensors_files + if (stat := _safetensors_payload_stats(file_path)) is not None + ] + if not stats: + return None, None + parameter_counts = [parameters for _, parameters in stats] + return ( + sum(checkpoint_bytes for checkpoint_bytes, _ in stats), + ( + sum(parameter_counts) + if all(parameters is not None for parameters in parameter_counts) + else None + ), + ) + + if os.path.isfile(component_model_path): + if component_model_path.endswith(_WEIGHT_FILE_SUFFIXES): + return _safe_file_size(component_model_path), None + return None, None + if not os.path.isdir(component_model_path): + return None, None + + weight_files = [ + file_path + for suffix in _WEIGHT_FILE_SUFFIXES + for file_path in glob.glob(os.path.join(component_model_path, f"*{suffix}")) + ] + if not weight_files: + return 0, 0 + sizes = [ + size + for file_path in weight_files + if (size := _safe_file_size(file_path)) is not None + ] + return (sum(sizes), None) if sizes else (None, None) + + +def infer_component_weight_size_bytes(component_model_path: str) -> int | None: + return infer_component_weight_stats(component_model_path)[0] + + +def estimate_component_weight_inventory( + sources: list[ComponentWeightSource], +) -> list[ComponentWeightEstimate]: + """Resolve the selected checkpoint sources without loading tensors.""" + seen_names: set[str] = set() + inventory = [] + for source in sources: + if source.component_name in seen_names: + raise ValueError( + f"duplicate component weight source {source.component_name!r}" + ) + seen_names.add(source.component_name) + inferred_bytes, inferred_parameters = ( + infer_component_weight_stats(source.component_model_path) + if source.checkpoint_bytes is None + else (None, None) + ) + inventory.append( + ComponentWeightEstimate( + component_name=source.component_name, + component_model_path=source.component_model_path, + checkpoint_bytes=( + source.checkpoint_bytes + if source.checkpoint_bytes is not None + else inferred_bytes + ), + parameter_count=( + source.parameter_count + if source.parameter_count is not None + else inferred_parameters + ), + target_element_size=source.target_element_size, + supports_fsdp_loading=source.supports_fsdp_loading, + ) + ) + return inventory diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py new file mode 100644 index 000000000000..ef27ac2eaaf1 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Choose a reversible residency seed before component loading.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from sglang.multimodal_gen import envs +from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( + RESIDENT, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_weight_inventory import ( + ComponentWeightEstimate, +) +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 ( + 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: + from sglang.multimodal_gen.runtime.server_args import ServerArgs + +logger = init_logger(__name__) + +GIB_BYTES = 1024**3 +MIN_INITIAL_RESIDENCY_RESERVE_BYTES = 4 * GIB_BYTES +INITIAL_RESIDENCY_RESERVE_FRACTION = 0.20 + + +def _default_denoising_steps(server_args: ServerArgs) -> int: + from sglang.multimodal_gen.registry import get_model_info + + model_info = get_model_info( + server_args.model_path, + server_args.backend, + server_args.model_id, + ) + if model_info is None: + return 1 + return max(1, model_info.sampling_param_cls().num_inference_steps or 1) + + +def choose_initial_resident_components( + server_args: ServerArgs, + inventory: list[ComponentWeightEstimate], + *, + available_bytes: int, + denoising_steps: int | None = None, +) -> set[str]: + """Choose resident startup components without removing runtime options. + + This is a load-feasibility seed, not a second serving-placement planner. + Every selected ordinary module can still be demoted to component or + layerwise offload by warmup calibration. Unknown weights and explicit + component choices remain on their configured load path. + """ + reserve_bytes = max( + MIN_INITIAL_RESIDENCY_RESERVE_BYTES, + int(available_bytes * INITIAL_RESIDENCY_RESERVE_FRACTION), + ) + fixed_resident = [ + item + for item in inventory + if server_args.residency_mode(item.component_name) == RESIDENT + ] + fixed_resident_sizes = [item.materialized_bytes() for item in fixed_resident] + if any(weight_bytes is None for weight_bytes in fixed_resident_sizes): + return set() + fixed_resident_bytes = sum( + weight_bytes for weight_bytes in fixed_resident_sizes if weight_bytes + ) + budget_bytes = available_bytes - reserve_bytes - fixed_resident_bytes + if budget_bytes <= 0: + return set() + + denoising_steps = denoising_steps or _default_denoising_steps(server_args) + options = [] + for item in inventory: + weight_bytes = item.materialized_bytes() + if ( + weight_bytes is None + or weight_bytes <= 0 + or server_args.residency_mode(item.component_name) == RESIDENT + or server_args.explicit_residency_mode(item.component_name) is not None + ): + continue + uses = denoising_steps if is_dit_component_name(item.component_name) else 1 + options.append( + PlacementOption( + group_key=item.component_name, + option_key=f"{item.component_name}:resident", + resource_delta_bytes={"gpu:load": weight_bytes}, + estimated_latency_savings=weight_bytes * uses, + preference_cost=(weight_bytes,), + ) + ) + if not options: + return set() + + plan = optimize_placement( + options, + resource_budget_bytes={"gpu:load": budget_bytes}, + ) + return {selection.group_key for selection in plan.selections} + + +def maybe_seed_initial_residency( + server_args: ServerArgs, + inventory: list[ComponentWeightEstimate], +) -> None: + """Apply a reversible resident seed for warmup-calibrated auto mode.""" + from sglang.multimodal_gen.runtime.server_args.auto_tune import ( + auto_residency_static_skip_reason, + ) + + if ( + auto_residency_static_skip_reason(server_args) is not None + or server_args.use_fsdp_inference + or not current_platform.is_cuda() + ): + return + + available_gib = current_platform.get_available_gpu_memory( + distributed=server_args.num_gpus > 1, + empty_cache=False, + ) + if envs.SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB is not None: + available_gib = min( + available_gib, + envs.SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB, + ) + selected = choose_initial_resident_components( + server_args, + inventory, + available_bytes=int(available_gib * GIB_BYTES), + ) + if not selected: + return + + for component_name in selected: + server_args.set_auto_residency_mode(component_name, RESIDENT) + logger.info( + "Initial auto residency: resident=%s (minimum free VRAM=%.1f GiB); " + "warmup may rebalance this load-safe seed.", + sorted(selected), + available_gib, + ) 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..9b8f69db8722 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 @@ -1682,52 +1682,59 @@ def iter_cpu_weights(self): cpu_buffer = self._consolidated_cpu_weights[layer_idx][dtype] yield name, cpu_buffer[offset : offset + numel].reshape(shape) + def prepare_layer_for_forward(self, layer_idx: int) -> None: + """Materialize one managed layer and schedule its successors. + + Forward hooks normally own this lifecycle. Model-specific forwards that + intentionally call a block's internals instead of ``block(...)`` must + invoke this method explicitly, because PyTorch cannot run that block's + hooks when its ``__call__`` is bypassed. + """ + if layer_idx == 0: + self._activate_residency() + self.prepare_for_next_req(non_blocking=False) + if layer_idx not in self._gpu_layers: + # LTX audio VAE traverses decoder.up in reverse order. + self.prefetch_layer(layer_idx, non_blocking=False) + if layer_idx in self._prefetch_events and self.copy_stream is not None: + torch.get_device_module().current_stream().wait_event( + self._prefetch_events[layer_idx] + ) + + if self.residency_policy == RESIDENCY_POLICY_STRIDED: + # Top up the stream at every layer rather than in bursts of + # prefetch_size. Under `strided` the next streamed layer can be + # several layers away, so index-based bursts would issue it late. + for layer_to_prefetch in self._next_streamed( + after=layer_idx, count=self.prefetch_size + ): + self.prefetch_layer(layer_to_prefetch, non_blocking=True) + elif self.prefetch_size and layer_idx % self.prefetch_size == 0: + for index in range( + layer_idx + self.prefetch_size, + layer_idx + 2 * self.prefetch_size, + ): + self.prefetch_layer(index % self.num_layers, non_blocking=True) + + def finish_layer_forward(self, layer_idx: int) -> None: + """Release a streamed layer after its forward completes.""" + self.release_layer(layer_idx) + def register_forward_hooks(self) -> None: if not self.enabled: return layers = dict(self.model.named_modules())[self.layers_attr_str] - def make_pre_hook(i): + def make_pre_hook(layer_idx): def hook(module, input): - if i == 0: - self._activate_residency() - self.prepare_for_next_req(non_blocking=False) - if i not in self._gpu_layers: - # LTX audio VAE traverses decoder.up in reverse order - self.prefetch_layer(i, non_blocking=False) - if i in self._prefetch_events and self.copy_stream is not None: - torch.get_device_module().current_stream().wait_event( - self._prefetch_events[i] - ) - - if self.residency_policy == RESIDENCY_POLICY_STRIDED: - # Top up the stream at every layer rather than in bursts of - # prefetch_size. Under `strided` the next streamed layer can - # be several layers away, so a burst schedule keyed on index - # arithmetic would either skip it or issue it late; asking - # for "the next N streamed layers" is the same request every - # layer and prefetch_layer is idempotent, so the repeats are - # free. This is what buys the wider hiding window: the - # transfer is issued as soon as the previous streamed layer - # is done with, not one layer before it is needed. - for layer_to_prefetch in self._next_streamed( - after=i, count=self.prefetch_size - ): - self.prefetch_layer(layer_to_prefetch, non_blocking=True) - # trigger batch prefetch (i + prefetch_size ~ i + 2 * prefetch_size) if needed - elif self.prefetch_size and i % self.prefetch_size == 0: - for j in range(i + self.prefetch_size, i + 2 * self.prefetch_size): - layer_to_prefetch = j % self.num_layers - self.prefetch_layer(layer_to_prefetch, non_blocking=True) + self.prepare_layer_for_forward(layer_idx) return hook - def make_post_hook(i): + def make_post_hook(layer_idx): def hook(module, input, output): - # previous, we wait here, until the copy stream for next layer is finished, - # now with any prefetch_size, only wait for the copy stream, when the copy stream is for the next layer - self.release_layer(i) + self.finish_layer_forward(layer_idx) return hook @@ -2453,6 +2460,16 @@ def configure_layerwise_offload_modules( normalized_component_names, ) ) + # Startup auto residency may replace the configured layerwise baseline + # with a reversible resident seed before this post-load setup runs. + # Configure only components whose effective placement is still + # layerwise; warmup can create managers lazily if it later demotes one. + selected_pipeline_component_names = [ + component_name + for component_name in selected_pipeline_component_names + if server_args.auto_residency_mode(component_name) is None + or server_args.residency_mode(component_name) == LAYERWISE_OFFLOAD + ] if ( warn_missing 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..dcf2d9686552 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 @@ -36,13 +36,13 @@ class PlacementOption(msgspec.Struct, frozen=True): 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, ...] = () + preference_cost: tuple[int, ...] = () class PlacementPlan(msgspec.Struct, frozen=True): resource_delta_bytes: dict[str, int] = {} estimated_latency_savings: int = 0 - placement_cost_bytes: tuple[int, ...] = () + preference_cost: tuple[int, ...] = () selections: list[PlacementOption] = [] @@ -166,15 +166,15 @@ def optimize_placement( 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 + preference_dimensions = max( + (len(option.preference_cost) for option in options), default=0 ) for option in options: - if option.placement_cost_bytes and ( - len(option.placement_cost_bytes) != placement_cost_dimensions + if option.preference_cost and ( + len(option.preference_cost) != preference_dimensions ): raise ValueError( - "all non-empty placement costs must have the same dimensions" + "all non-empty preference costs must have the same dimensions" ) grouped: dict[str, list[PlacementOption]] = {} @@ -236,7 +236,7 @@ def optimize_placement( ( (0,) * len(resource_names), 0, - (0,) * placement_cost_dimensions, + (0,) * preference_dimensions, (), ) ) @@ -247,7 +247,7 @@ def optimize_placement( for resource_name in resource_names ), option.estimated_latency_savings, - option.placement_cost_bytes or (0,) * placement_cost_dimensions, + option.preference_cost or (0,) * preference_dimensions, (option,), ) for option in group_options @@ -257,9 +257,7 @@ def optimize_placement( 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) - ] + minimum_cost = [(0,) * preference_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( @@ -393,7 +391,7 @@ def _minimize_cost( 0, (0,) * len(resource_names), 0, - (0,) * placement_cost_dimensions, + (0,) * preference_dimensions, (), ) assert best is not None @@ -406,6 +404,6 @@ def _minimize_cost( return PlacementPlan( resource_delta_bytes=full_resource_deltas, estimated_latency_savings=best[1], - placement_cost_bytes=best[2], + preference_cost=best[2], selections=list(best[3]), ) diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 38158348bf4f..49dcba5e338a 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(pre_warmup=req.action == "apply_static") def _handle_set_lora(self, reqs: List[Any]) -> OutputBatch: # TODO: return set status diff --git a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py index a2769085747d..784133e8e714 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py @@ -38,6 +38,10 @@ set_default_torch_dtype, ) from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping +from sglang.multimodal_gen.runtime.managers.memory_managers.component_weight_inventory import ( + ComponentWeightSource, + infer_safetensors_weight_stats_by_prefix, +) from sglang.multimodal_gen.runtime.models.dits.hunyuan3d_paint import ( Hunyuan3DPaintUNet, ) @@ -491,6 +495,99 @@ def _load_texture_components( return components + @classmethod + def _texture_component_weight_sources( + cls, + server_args: ServerArgs, + config: Hunyuan3D2PipelineConfig, + ) -> list[ComponentWeightSource]: + paint_dir = cls._resolve_model_subfolder( + server_args.model_path, + config.paint_subfolder, + ( + "vae/config.json", + "unet/config.json", + "scheduler/scheduler_config.json", + ), + ) + target_element_size = PRECISION_TO_TYPE[config.dit_precision].itemsize + sources = [ + ComponentWeightSource( + "paint_vae", + os.path.join(paint_dir, "vae"), + target_element_size=target_element_size, + ), + ComponentWeightSource( + "paint_transformer", + os.path.join(paint_dir, "unet"), + target_element_size=target_element_size, + ), + ] + if not config.delight_enable: + return sources + + delight_dir = cls._resolve_model_subfolder( + server_args.model_path, + config.delight_subfolder, + ( + "vae/config.json", + "unet/config.json", + "scheduler/scheduler_config.json", + "text_encoder/config.json", + "tokenizer/tokenizer_config.json", + ), + ) + sources.extend( + [ + ComponentWeightSource( + "delight_vae", + os.path.join(delight_dir, "vae"), + target_element_size=target_element_size, + ), + ComponentWeightSource( + "delight_transformer", + os.path.join(delight_dir, "unet"), + target_element_size=target_element_size, + ), + ComponentWeightSource( + "delight_text_encoder", + os.path.join(delight_dir, "text_encoder"), + target_element_size=target_element_size, + ), + ] + ) + return sources + + @staticmethod + def _shape_component_weight_sources( + ckpt_path: str, use_safetensors: bool, target_element_size: int + ) -> list[ComponentWeightSource]: + prefix_stats = ( + infer_safetensors_weight_stats_by_prefix(ckpt_path) + if use_safetensors + else None + ) + if prefix_stats is None: + return [ComponentWeightSource("hy3dshape_checkpoint", ckpt_path)] + + component_names = { + "model": "hy3dshape_model", + "vae": "hy3dshape_vae", + "conditioner": "hy3dshape_conditioner", + } + return [ + ComponentWeightSource( + component_name=component_names.get(prefix, prefix), + component_model_path=ckpt_path, + checkpoint_bytes=checkpoint_bytes, + parameter_count=parameter_count, + target_element_size=target_element_size, + ) + for prefix, (checkpoint_bytes, parameter_count) in sorted( + prefix_stats.items() + ) + ] + # Module loading override def load_modules( self, @@ -515,6 +612,24 @@ def load_modules( config.shape_variant, ) + self.prepare_component_weight_inventory( + server_args, + self._shape_component_weight_sources( + ckpt_path, + config.shape_use_safetensors, + ( + torch.bfloat16.itemsize + if config.shape_variant and "bf16" in config.shape_variant + else torch.float16.itemsize + ), + ) + + ( + self._texture_component_weight_sources(server_args, config) + if config.paint_enable + else [] + ), + ) + with open(config_path, "r") as f: model_config = yaml.safe_load(f) 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..bfe45e2a024a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -23,6 +23,9 @@ ComponentUse, ResidencyState, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_weight_inventory import ( + ComponentWeightSource, +) from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) @@ -56,6 +59,7 @@ _normalize_ltx2_two_stage_device_mode, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision logger = init_logger(__name__) @@ -655,6 +659,46 @@ def __init__(self, *args, **kwargs): self._ltx2_residency.strategy ) + def additional_component_weight_sources( + self, server_args: ServerArgs + ) -> list[ComponentWeightSource]: + server_args.component_paths = _resolve_ltx2_two_stage_component_paths( + self.model_path, server_args.component_paths + ) + sources = [] + for component_name in ("spatial_upsampler", "distilled_lora"): + component_path = server_args.component_paths.get(component_name) + if component_path: + sources.append( + ComponentWeightSource( + component_name=component_name, + component_model_path=component_path, + target_element_size=( + resolve_component_precision(server_args, component_name) + or resolve_component_precision(server_args, "transformer") + ).itemsize, + ) + ) + + if ( + server_args.ltx2_two_stage_device_mode == "resident" + and self._should_merge_stage2_distilled_lora(server_args) + and server_args.lora_path is None + ): + sources.append( + ComponentWeightSource( + component_name="transformer_2", + component_model_path=self._resolve_component_path( + server_args, "transformer", "transformer" + ), + target_element_size=resolve_component_precision( + server_args, "transformer_2" + ).itemsize, + supports_fsdp_loading=True, + ) + ) + return sources + @staticmethod def _should_merge_stage2_distilled_lora(server_args: ServerArgs) -> bool: return is_ltx23_native_variant( @@ -754,8 +798,8 @@ def _can_merge_stage1_distilled_into_base(self, server_args: ServerArgs) -> bool and self._should_merge_stage2_distilled_lora(server_args) and self._stage1_lora_path is None and float(self.STAGE_1_DISTILLED_LORA_STRENGTH) != 0.0 - and not bool(getattr(server_args, "use_fsdp_inference", False)) - and getattr(server_args, "quantization", None) is None + and not server_args.should_use_fsdp_for_component("transformer") + and server_args.quantization is None ) def _maybe_merge_stage1_distilled_into_base(self, server_args: ServerArgs) -> None: diff --git a/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_pipeline.py index 97e08d0ba4f5..08d512e5d0d5 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_pipeline.py @@ -5,6 +5,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import SanaWMPipelineConfig from sglang.multimodal_gen.configs.sample.sana_wm import SanaWMSamplingParams from sglang.multimodal_gen.runtime.loader.utils import get_memory_usage_of_component +from sglang.multimodal_gen.runtime.managers.memory_managers.component_weight_inventory import ( + ComponentWeightSource, +) from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, @@ -170,6 +173,24 @@ def initialize_pipeline(self, server_args: ServerArgs) -> None: return self._load_refiner_modules(server_args) + def additional_component_weight_sources( + self, server_args: ServerArgs + ) -> list[ComponentWeightSource]: + if sana_wm_skip_refiner_enabled(): + return [] + target_element_size = default_sana_wm_refiner_dtype(server_args).itemsize + return [ + ComponentWeightSource( + component_name=module_name, + component_model_path=self._resolve_refiner_component_path( + server_args, module_name, subpath + ), + target_element_size=target_element_size, + ) + for module_name, subpath in self._REFINER_SUB_MODULES + if module_name != "tokenizer_2" + ] + def _resolve_refiner_paths(self, server_args: ServerArgs) -> tuple[str, str]: component_paths = getattr(server_args, "component_paths", {}) or {} refiner_root = component_paths.get( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 3f0ada45fb28..f69f837c0ea6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -30,6 +30,17 @@ ComponentResidencyStrategy, get_global_component_residency_manager, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_weight_inventory import ( + ComponentWeightEstimate, + ComponentWeightSource, + estimate_component_weight_inventory, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( + is_image_encoder_component_name, + is_legacy_dit_offload_component_name, + is_text_encoder_component_name, + is_vae_component_name, +) from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import ( PipelineExecutor, ) @@ -57,6 +68,7 @@ verify_model_config_and_directory, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision logger = init_logger(__name__) @@ -142,6 +154,24 @@ def __init__( # [module_name, gpu memory usage] self.memory_usages: dict[str, float] = {} + self.component_weight_inventory: list[ComponentWeightEstimate] = [] + if loaded_modules is None and os.path.isfile(self.model_path): + # Registered single-file pipelines contain a standalone DiT and + # bypass the normal model-index component loop. + target_dtype = resolve_component_precision(server_args, "transformer") + self.prepare_component_weight_inventory( + server_args, + [ + ComponentWeightSource( + "transformer", + self.model_path, + target_element_size=( + target_dtype.itemsize if target_dtype is not None else None + ), + supports_fsdp_loading=True, + ) + ], + ) # Load modules directly in initialization logger.info("Loading pipeline modules...") self.modules = self.load_modules(server_args, loaded_modules) @@ -360,6 +390,83 @@ def _resolve_component_path( logger.debug("Resolved component path: %s", component_model_path) return component_model_path + def additional_component_weight_sources( + self, server_args: ServerArgs + ) -> list[ComponentWeightSource]: + """Selected weight sources loaded outside the base component loop.""" + return [] + + @staticmethod + def _component_target_element_size( + component_name: str, server_args: ServerArgs + ) -> int | None: + target_dtype = resolve_component_precision(server_args, component_name) + if target_dtype is None: + fallback_component = None + if is_legacy_dit_offload_component_name(component_name): + fallback_component = "transformer" + elif is_vae_component_name(component_name): + fallback_component = "vae" + elif is_image_encoder_component_name(component_name): + fallback_component = "image_encoder" + elif is_text_encoder_component_name(component_name): + fallback_component = "text_encoder" + if fallback_component is not None: + target_dtype = resolve_component_precision( + server_args, fallback_component + ) + return target_dtype.itemsize if target_dtype is not None else None + + @staticmethod + def _component_weight_source( + spec: ComponentLoadSpec, server_args: ServerArgs + ) -> ComponentWeightSource: + weights_path = server_args.component_weights_paths.get(spec.load_module_name) + if ( + weights_path is None + and spec.load_module_name == "transformer" + and server_args.transformer_weights_path is not None + ): + weights_path = server_args.transformer_weights_path + return ComponentWeightSource( + component_name=spec.module_name, + component_model_path=weights_path or spec.component_model_path, + target_element_size=ComposedPipelineBase._component_target_element_size( + spec.load_module_name, server_args + ), + supports_fsdp_loading=PipelineComponentLoader.supports_fsdp_loading( + spec.module_name, + spec.transformers_or_diffusers, + spec.architecture, + ), + ) + + def prepare_component_weight_inventory( + self, server_args: ServerArgs, sources: list[ComponentWeightSource] + ) -> None: + """Resolve selected checkpoint weights before component loading.""" + self.component_weight_inventory = estimate_component_weight_inventory(sources) + known = { + item.component_name: round(item.checkpoint_bytes / 1024**3, 2) + for item in self.component_weight_inventory + if item.checkpoint_bytes is not None and item.checkpoint_bytes > 0 + } + unknown = [ + item.component_name + for item in self.component_weight_inventory + if item.checkpoint_bytes is None + ] + logger.debug( + "Pre-load component weights (GiB): known=%s unknown=%s", + known, + unknown, + ) + from sglang.multimodal_gen.runtime.managers.memory_managers.initial_residency import ( + maybe_seed_initial_residency, + ) + + maybe_seed_initial_residency(server_args, self.component_weight_inventory) + def load_modules( self, server_args: ServerArgs, @@ -528,10 +635,17 @@ def load_modules( ) ) - # reorder loading order to avoid OOM - component_load_specs: ComponentLoadSpec = order_component_load_specs( - component_load_specs + self.prepare_component_weight_inventory( + server_args, + [ + self._component_weight_source(spec, server_args) + for spec in component_load_specs + ] + + self.additional_component_weight_sources(server_args), ) + + # reorder loading order to avoid OOM + component_load_specs = order_component_load_specs(component_load_specs) logger.info( "Memory-aware component load order: %s", [spec.module_name for spec in component_load_specs], diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py index bcd324efca39..0421c7cce0fe 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py @@ -85,6 +85,9 @@ def _component_residency_request( self.begin_component_residency_request(stages, payload, server_args) try: yield + except Exception: + self.component_residency_manager.capture_failed_warmup_phase() + raise finally: self.finish_component_residency_request() 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..94134f4b8ea3 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 @@ -239,17 +239,35 @@ def _forward_diffusers_video_only( batch_size, -1, hidden_states.size(-1) ) - for block in transformer.transformer_blocks: - hidden_states = _forward_diffusers_video_block( - block=block, - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - temb=temb, - video_rotary_emb=video_rotary_emb, - encoder_attention_mask=encoder_attention_mask, - n_context_tokens=n_context_tokens, + layerwise_manager = None + if isinstance(transformer, OfficialDiffusersLTX2RefinerModule): + layerwise_manager = next( + ( + manager + for manager in transformer.layerwise_offload_managers + if manager.layers_attr_str == "module.transformer_blocks" + and manager.enabled + ), + None, ) + for block_idx, block in enumerate(transformer.transformer_blocks): + if layerwise_manager is not None: + layerwise_manager.prepare_layer_for_forward(block_idx) + try: + hidden_states = _forward_diffusers_video_block( + block=block, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + video_rotary_emb=video_rotary_emb, + encoder_attention_mask=encoder_attention_mask, + n_context_tokens=n_context_tokens, + ) + finally: + if layerwise_manager is not None: + layerwise_manager.finish_layer_forward(block_idx) + scale_shift_values = ( transformer.scale_shift_table[None, None] + embedded_timestep[:, :, None] ) 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..8c7d11be5c47 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -48,17 +48,14 @@ 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.""" +def auto_residency_static_skip_reason(server_args: ServerArgs) -> str | None: + """Return why args cannot use pre-warmup static residency planning.""" 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" + if server_args.disagg_role != RoleType.MONOLITHIC: + return "disaggregated role" 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" @@ -83,11 +80,11 @@ def auto_residency_args_skip_reason(server_args: ServerArgs) -> str | None: 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 + # A transformer weight-path override alone is not a quantization mode. + # Self-describing pre-quantized checkpoints are also safe to move after # loading: residency changes storage location, not quantized values. return "quantized checkpoint" if server_args.direct_gpu_weight_loading: @@ -97,6 +94,16 @@ def auto_residency_args_skip_reason(server_args: ServerArgs) -> str | None: return None +def auto_residency_args_skip_reason(server_args: ServerArgs) -> str | None: + """Return why args cannot use warmup-calibrated residency refinement.""" + reason = auto_residency_static_skip_reason(server_args) + if reason is not None: + return reason + if server_args.warmup_mode != "server": + return "no synthetic server warmup to calibrate from" + return None + + class ServerArgsAutoTuner: """Auto-tunes the server-arg for the given performance-mode, based on practical deployment experience with different model architectures""" @@ -107,9 +114,9 @@ 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 _uses_auto_residency_planner(self) -> bool: + """Whether the static planner replaces coarse deployment defaults.""" + return auto_residency_static_skip_reason(self.server_args) is None def _resolve_keep_resident_min_available_gb( self, deployment_config: ModelDeploymentConfig @@ -183,10 +190,10 @@ 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. + if self._uses_auto_residency_planner(): + # Keep the load-safe placement until the post-load static planner + # replaces model/card thresholds with component sizes and the + # default workload. A server warmup may refine that estimate. return # Explicit placement is component-scoped; unmatched components still @@ -314,8 +321,20 @@ def _maybe_keep_ltx23_resident_aux_components_resident(self) -> None: def maybe_adjust_auto_fsdp_with_offload_enabled(self) -> None: args = self.server_args + if ( + args.use_fsdp_inference + and args.is_arg_explicitly_set("use_fsdp_inference") + and args.explicit_residency_mode("transformer") is None + ): + # An explicit FSDP request must manage the DiT unless the user also + # supplied a component-scoped placement. Auto's load-safe default + # otherwise leaves the transformer offloaded and silently turns + # --use-fsdp-inference into a no-op. + args.dit_cpu_offload = False + if ( args.performance_mode == "auto" + and not self._uses_auto_residency_planner() and args.num_gpus >= 2 and not self._explicit_dit_residency and self._auto_uses_dit_offload() @@ -352,20 +371,28 @@ def maybe_adjust_auto_default_layerwise_offload(self) -> None: return min_available_gb = self._get_min_available_device_memory_gb() - logger.info( - "Auto memory policy for %s: %s of free device memory selects " - "layerwise offload for %s. Explicit placement flags always win, " - "and the per-component lines below say where each one's weights " - "landed -- see the model's cookbook page for what to expect from " - "your memory budget.", - args.pipeline_config.__class__.__name__, - ( - f"{min_available_gb:.1f} GiB" - if min_available_gb is not None - else "an unknown amount" - ), - ", ".join(layerwise_components), + uses_planner = self._uses_auto_residency_planner() + model_name = args.pipeline_config.__class__.__name__ + available = ( + f"{min_available_gb:.1f} GiB" if min_available_gb is not None else "unknown" ) + if uses_planner: + logger.info( + "Auto memory planner for %s: layerwise candidates=%s, " + "free VRAM=%s. Final placement follows component sizing and " + "warmup calibration; explicit placement flags always win.", + model_name, + ", ".join(layerwise_components), + available, + ) + else: + logger.info( + "Auto memory policy for %s: free VRAM=%s selects layerwise " + "offload for %s. Explicit placement flags always win.", + model_name, + available, + ", ".join(layerwise_components), + ) args.layerwise_offload_components = layerwise_components self._warn_if_resident_dit_contradicts_its_own_threshold(layerwise_components) @@ -582,7 +609,7 @@ 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(): + if self._uses_auto_residency_planner(): return components deployment_config = self._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 d056c5db09b1..6848a4e87bec 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -891,6 +891,14 @@ def _adjust_ltx2_two_stage_device_mode(self): configured, ) + if mode == "resident": + # This structural mode owns both DiTs and is excluded from the + # generic residency planner. Keep the ordinary component flags in + # agreement so loading does not take an offload path that no later + # planner can correct. + self.dit_cpu_offload = False + self.dit_layerwise_offload = False + self.ltx2_two_stage_device_mode = mode def _resolve_default_ltx2_two_stage_device_mode(self) -> str: diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index d15fa1141184..5a2d6c150869 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -1,6 +1,7 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 +import asyncio import os import tempfile from typing import Any, Awaitable, Callable @@ -17,6 +18,7 @@ from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.server_args.auto_tune import ( auto_residency_args_skip_reason, + auto_residency_static_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 @@ -115,12 +117,42 @@ def auto_residency_skip_reason(server_args: ServerArgs) -> str | None: return None +def should_apply_pre_warmup_auto_residency(server_args: ServerArgs) -> bool: + return auto_residency_static_skip_reason(server_args) is 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_pre_warmup_auto_residency( + server_args: ServerArgs, + forward: Callable[[Req], Awaitable[OutputBatch]], +) -> None: + """Choose a weight-feasible placement before the first serving forward.""" + from sglang.multimodal_gen.runtime.entrypoints.control_requests import ( + AutoResidencyReq, + ) + from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( + PLACEMENT_STATUS_ROLLBACK_FAILED, + ) + + if not should_apply_pre_warmup_auto_residency(server_args): + return + response = await forward(AutoResidencyReq(action="apply_static")) + 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: + logger.warning( + "Pre-warmup residency placement was not applied; calibrating the " + "original placement: %s", + response.error, + ) + + async def maybe_apply_auto_residency( server_args: ServerArgs, forward: Callable[[Req], Awaitable[OutputBatch]], @@ -157,6 +189,7 @@ async def maybe_apply_auto_residency( # never abort startup, explicit --warmup-resolutions ones must succeed. fail_open = server_args.warmup_resolutions is None adjusted_any = False + recovering_from_oom = False async def rollback_and_rewarm(error: Exception) -> None: logger.warning( @@ -199,13 +232,26 @@ async def rollback_and_rewarm(error: Exception) -> None: ) return status = _auto_residency_status(response) + response_recovering_from_oom = bool( + isinstance(response.output, dict) + and response.output.get("recovering_from_oom") + ) if status == PLACEMENT_STATUS_ROLLBACK_FAILED: raise RuntimeError(f"auto residency rollback failed: {response.error}") if response.error is not None: + retryable = bool( + isinstance(response.output, dict) and response.output.get("retryable") + ) if status == PLACEMENT_STATUS_ROLLED_BACK: await run_async_client_warmup( server_args, forward, fail_open=fail_open, rewarm=True ) + if retryable: + logger.warning( + "Auto residency calibration corrected its VRAM model; " + "retrying with the restored placement." + ) + continue retained = ( "the last calibrated placement" if adjusted_any @@ -218,9 +264,15 @@ async def rollback_and_rewarm(error: Exception) -> None: ) return if status != PLACEMENT_STATUS_ADJUSTED: + if response_recovering_from_oom: + raise RuntimeError( + "default workload exceeds the VRAM budget and auto " + "residency found no feasible placement" + ) return adjusted_any = True + recovering_from_oom = response_recovering_from_oom # 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. @@ -229,21 +281,35 @@ async def rollback_and_rewarm(error: Exception) -> None: server_args, forward, fail_open=False, rewarm=True ) except Exception as e: + if recovering_from_oom and _is_out_of_memory(e): + logger.warning( + "Default workload still exceeds the VRAM budget after an " + "auto-residency recovery step; keeping the lower-memory " + "placement and replanning." + ) + continue await rollback_and_rewarm(e) return + return - logger.warning( - "Auto residency stopped after %d adjustment rounds; keeping the last " - "successfully calibrated placement.", - MAX_AUTO_RESIDENCY_ROUNDS, - ) + if recovering_from_oom: + raise RuntimeError( + "default workload still exceeds the VRAM budget after " + f"{MAX_AUTO_RESIDENCY_ROUNDS} auto-residency corrections" + ) + else: + logger.warning( + "Auto residency stopped after %d adjustment 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 +MAX_AUTO_RESIDENCY_ROUNDS = 2 def _is_out_of_memory(error: Any) -> bool: @@ -331,6 +397,7 @@ async def run_async_client_warmup( rewarm: bool = False, ) -> None: try: + auto_residency_handles_oom = auto_residency_skip_reason(server_args) is None warmup_input_path = None if should_include_warmup_image(server_args, server_based_warmup=True): warmup_input_path = prepare_warmup_image_path(server_args) @@ -342,6 +409,12 @@ async def run_async_client_warmup( for _ in range(MAX_WARMUP_DEGRADE_ATTEMPTS): if response.error is None or not _is_out_of_memory(response.error): break + # The residency planner needs the first failed target-shape + # measurement. Retrying smaller requests cannot fix a weight- + # dominated OOM and can retain failed-forward allocations that + # contaminate the phase model used by the planner. + if auto_residency_handles_oom: + break lighter = _degrade_after_oom(server_args, req) if lighter is None: break @@ -382,6 +455,31 @@ def run_sync_client_warmup( raise RuntimeError(response.error) +def run_sync_startup_warmup( + server_args: ServerArgs, + forward: Callable[[Req], OutputBatch], +) -> None: + """Run the server-style startup sequence through a synchronous client. + + Offline ``DiffGenerator`` launches the same scheduler and workers as the + HTTP entrypoint, but its client API is synchronous. Adapting that call into + the shared async orchestration keeps static placement, one warmup probe, + and measured refinement identical across both entrypoints. + """ + + async def async_forward(req: Req) -> OutputBatch: + return forward(req) + + async def run() -> None: + await maybe_apply_pre_warmup_auto_residency(server_args, async_forward) + if not should_run_synthetic_server_warmup(server_args): + return + await run_async_client_warmup(server_args, async_forward, fail_open=True) + await maybe_apply_auto_residency(server_args, async_forward) + + asyncio.run(run()) + + def prepare_warmup_image_path(server_args: ServerArgs) -> str: if server_args.input_save_path is not None: uploads_dir = server_args.input_save_path 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 3fb8cac07898..349025eb1368 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -1415,6 +1415,7 @@ def generate_text_url_image_to_video(case_id, client) -> tuple[str, bytes]: "reference_url": sampling_params.image_path, "fps": sampling_params.fps, "num_frames": sampling_params.num_frames, + **extra_body, }, ) diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py index c6d145515373..d6e539ea44e9 100644 --- a/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py @@ -5,7 +5,7 @@ import types import unittest from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, call, patch import torch @@ -50,6 +50,7 @@ OfficialGemma3TextEncoderModule, SanaWMLTX2RefinerStage, SanaWMRefinerDecodingStage, + _forward_diffusers_video_only, _refiner_config_value, _streaming_diffusers_self_attention, _uses_diffusers_ltx2_refiner, @@ -421,6 +422,33 @@ def test_refiner_modules_are_loaded_from_official_subtrees(self) -> None: ), ) + @patch( + "sglang.multimodal_gen.runtime.pipelines.sana_wm_pipeline." + "sana_wm_skip_refiner_enabled", + return_value=False, + ) + def test_refiner_weights_are_included_in_preload_inventory(self, _skip) -> None: + pipeline = object.__new__(SanaWMTwoStagePipeline) + pipeline.model_path = "/models/sana-wm" + sources = pipeline.additional_component_weight_sources( + SimpleNamespace( + component_paths={}, + pipeline_config=SimpleNamespace(dit_precision="bf16"), + ) + ) + + self.assertEqual( + [ + (source.component_name, source.component_model_path) + for source in sources + ], + [ + ("transformer_2", "/models/sana-wm/refiner/transformer"), + ("connectors", "/models/sana-wm/refiner/connectors"), + ("text_encoder_2", "/models/sana-wm/refiner/text_encoder"), + ], + ) + class TestSanaWMPipeline(unittest.TestCase): def test_validate_parallelism_rejects_tensor_parallelism(self) -> None: @@ -1255,6 +1283,63 @@ def test_camera_chunk_scan_matches_recurrent_scan(self) -> None: class TestSanaWMRefinerStage(_GlobalStageArgsMixin, unittest.TestCase): + def test_manual_diffusers_forward_runs_layerwise_lifecycle(self) -> None: + class _Rope: + def prepare_video_coords(self, *args, **kwargs): + return None + + def __call__(self, *args, **kwargs): + return None, None + + class _TimeEmbed(torch.nn.Module): + def forward(self, timestep, *, batch_size, hidden_dtype): + embedded = torch.zeros( + timestep.numel(), 4, dtype=hidden_dtype, device=timestep.device + ) + return embedded, embedded + + class LTX2VideoTransformer3DModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.rope = _Rope() + self.proj_in = torch.nn.Identity() + self.time_embed = _TimeEmbed() + self.caption_projection = torch.nn.Identity() + self.transformer_blocks = torch.nn.ModuleList( + [torch.nn.Identity(), torch.nn.Identity()] + ) + self.scale_shift_table = torch.nn.Parameter(torch.zeros(2, 4)) + self.norm_out = torch.nn.Identity() + self.proj_out = torch.nn.Identity() + + transformer = OfficialDiffusersLTX2RefinerModule(LTX2VideoTransformer3DModel()) + manager = Mock() + manager.layers_attr_str = "module.transformer_blocks" + manager.enabled = True + transformer.layerwise_offload_managers = [manager] + hidden_states = torch.zeros(1, 2, 4) + + with patch( + f"{_SANA_WM_REFINER_STAGE_MODULE}._forward_diffusers_video_block", + side_effect=lambda **kwargs: kwargs["hidden_states"], + ): + result = _forward_diffusers_video_only( + transformer, + hidden_states=hidden_states, + encoder_hidden_states=hidden_states, + timestep=torch.zeros(1, 2), + encoder_attention_mask=None, + num_frames=1, + height=1, + width=2, + fps=24.0, + n_context_tokens=1, + ) + + self.assertEqual(result.shape, hidden_states.shape) + manager.prepare_layer_for_forward.assert_has_calls([call(0), call(1)]) + manager.finish_layer_forward.assert_has_calls([call(0), call(1)]) + def test_diffusers_refiner_detection_uses_official_class_name(self) -> None: class LTX2VideoTransformer3DModel(torch.nn.Module): pass 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..cc0a6e114c9a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -13,6 +13,7 @@ ACTIVATION_EXTRAPOLATION_MARGIN, GIB_BYTES, MIN_VRAM_RESERVE_BYTES, + PAGEABLE_H2D_COST_MULTIPLIER, AppliedResidencyChange, AutoResidencyPlan, AutoResidencyRollbackError, @@ -30,6 +31,7 @@ estimate_default_workload_timing, estimate_workload_phase_peaks, format_applied_changes, + measured_failed_workload_phase_peaks, plan_auto_residency, rank_candidates_by_h2d_savings, resolve_measured_default_workload, @@ -114,6 +116,25 @@ def test_scales_only_denoising_with_target_steps(self): "transformer": ("DenoisingStage",), } + def test_prefetch_phase_does_not_claim_stage_latency(self): + record = _record( + num_inference_steps=2, + total_duration_ms=1_000, + stage_duration_ms={"encode": 100, "denoise": 900}, + phase_active_components={ + "0:encode:prefetch:transformer": ("transformer",), + "1:denoise:use:transformer": ("transformer",), + }, + ) + + _, _, component_stages = estimate_default_workload_timing( + records=[record], + target_units=record.workload_units(), + target_num_inference_steps=2, + ) + + assert component_stages == {"transformer": ("denoise",)} + def test_uses_steady_step_instead_of_scaling_first_step_setup(self): record = _record( num_inference_steps=2, @@ -138,7 +159,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_request_not_component_stage(self): + def test_transfer_savings_uses_component_stage_as_cross_component_cap(self): candidates = [ _candidate("text_encoder", weight_gib=10, h2d_gib=10), _candidate("transformer", weight_gib=10, h2d_gib=10), @@ -148,12 +169,103 @@ def test_transfer_savings_is_capped_by_request_not_component_stage(self): savings = estimate_candidate_latency_savings_ns( candidates=candidates, request_duration_ns=10_000_000_000, + stage_duration_ns={"encode": 100_000_000, "denoise": 4_000_000_000}, + component_stages={ + "text_encoder": ("encode",), + "transformer": ("denoise",), + }, ) - assert savings[candidates[0].option_key()] > 400_000_000 + assert savings[candidates[0].option_key()] == 100_000_000 assert savings[candidates[1].option_key()] > 400_000_000 assert savings[candidates[2].option_key()] < 0 + def test_stage_cap_preserves_order_inside_one_shot_component_frontier(self): + partial = ResidencyTarget( + component_name="text_encoder", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=10 * GIB_BYTES, + h2d_bytes_per_request=5 * GIB_BYTES, + target_layerwise_resident_layers=(5,), + target_layerwise_pinned_layers=((),), + ) + resident = ResidencyTarget( + component_name="text_encoder", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=10 * GIB_BYTES, + target_layerwise_resident_layers=(10,), + target_layerwise_pinned_layers=((),), + ) + + savings = estimate_candidate_latency_savings_ns( + candidates=[partial, resident], + request_duration_ns=10_000_000_000, + stage_duration_ns={"encode": 100_000_000}, + component_stages={"text_encoder": ("encode",)}, + ) + + assert savings[partial.option_key()] == 50_000_000 + assert savings[resident.option_key()] == 100_000_000 + + def test_async_stage_time_does_not_cap_repeated_dit_transfers(self): + partial = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=10 * GIB_BYTES, + h2d_bytes_per_request=5 * GIB_BYTES, + target_layerwise_resident_layers=(5,), + target_layerwise_pinned_layers=((),), + ) + resident = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=10 * GIB_BYTES, + target_layerwise_resident_layers=(10,), + target_layerwise_pinned_layers=((),), + ) + + savings = estimate_candidate_latency_savings_ns( + candidates=[partial, resident], + request_duration_ns=10_000_000_000, + stage_duration_ns={"denoise": 100_000_000}, + component_stages={"transformer": ("denoise",)}, + ) + + assert savings[partial.option_key()] > 100_000_000 + assert savings[resident.option_key()] > savings[partial.option_key()] + + def test_dit_request_cap_does_not_compress_marginal_transfer_savings(self): + partial = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=95 * GIB_BYTES, + h2d_bytes_per_request=95 * GIB_BYTES, + target_layerwise_resident_layers=(95,), + target_layerwise_pinned_layers=((),), + ) + resident = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=100 * GIB_BYTES, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(100,), + target_layerwise_pinned_layers=((),), + ) + + savings = estimate_candidate_latency_savings_ns( + candidates=[partial, resident], + request_duration_ns=1_000_000_000, + stage_duration_ns={"denoise": 100_000_000}, + component_stages={"transformer": ("denoise",)}, + ) + + assert savings[resident.option_key()] == 1_000_000_000 + assert ( + savings[resident.option_key()] - savings[partial.option_key()] > 200_000_000 + ) + class TestEstimateDefaultWorkloadPeak: def test_same_shape_uses_measured_peak(self): @@ -353,6 +465,43 @@ def test_phase_estimation_preserves_active_component_membership(self): assert active == {"denoise": ("transformer",)} assert used == active + def test_failed_probe_preserves_phase_layouts_for_recovery(self): + first = WarmupMemoryRecord( + width=512, + height=512, + num_frames=1, + baseline_allocated_bytes=20 * GIB_BYTES, + peak_allocated_bytes=30 * GIB_BYTES, + succeeded=False, + phase_peak_allocated_bytes={"denoise": 29 * GIB_BYTES}, + phase_active_components={"denoise": ("transformer",)}, + phase_used_components={"denoise": ("transformer",)}, + ) + second = WarmupMemoryRecord( + width=512, + height=512, + num_frames=1, + baseline_allocated_bytes=20 * GIB_BYTES, + peak_allocated_bytes=31 * GIB_BYTES, + succeeded=False, + phase_peak_allocated_bytes={"denoise": 31 * GIB_BYTES}, + phase_active_components={"denoise": ("transformer_2",)}, + phase_used_components={"denoise": ("transformer_2",)}, + ) + + peak, phases, active, used, _ = measured_failed_workload_phase_peaks( + records=[first, second], + target_units=first.workload_units(), + ) + + assert peak == 31 * GIB_BYTES + assert phases == { + "denoise:layout:0": 29 * GIB_BYTES, + "denoise:layout:1": 31 * GIB_BYTES, + } + assert set(active.values()) == {("transformer",), ("transformer_2",)} + assert used == active + def test_phase_estimation_uses_allocated_peak(self): record = WarmupMemoryRecord( width=1024, @@ -493,15 +642,17 @@ def _candidate( def _report( *, rank: int = 0, - budget_gib: int = 100, - estimated_gib: int | None = 50, + budget_gib: float = 100, + estimated_gib: float | None = 50, + planning_headroom_correction_gib: int = 0, candidates: list[ResidencyTarget] | None = None, skip_reason: str | None = None, - phase_peaks_gib: dict[str, int] | None = None, + phase_peaks_gib: dict[str, float] | 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, + component_active_weight_gib: dict[str, int] | None = None, node_rank: int = 0, pinned_host_gib: int = 0, host_pin_capacity_gib: int = 0, @@ -511,16 +662,22 @@ def _report( estimated_request_duration_ns: int = 0, measured_request_duration_ns: int = 0, candidate_latency_savings_ns: dict[str, int] | None = None, + warmup_oom: bool = False, + require_feasible_placement: bool = False, ) -> RankResidencyReport: return RankResidencyReport( rank=rank, - budget_bytes=budget_gib * GIB_BYTES, + budget_bytes=int(budget_gib * GIB_BYTES), estimated_peak_bytes=( - None if estimated_gib is None else estimated_gib * GIB_BYTES + None if estimated_gib is None else int(estimated_gib * GIB_BYTES) + ), + planning_headroom_correction_bytes=( + planning_headroom_correction_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() + name: int(value * GIB_BYTES) + for name, value in (phase_peaks_gib or {}).items() }, active_components_by_phase=( phase_present_components @@ -535,6 +692,10 @@ def _report( name: value * GIB_BYTES for name, value in (component_weight_gib or {}).items() }, + current_active_weight_bytes_by_component={ + name: value * GIB_BYTES + for name, value in (component_active_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, @@ -544,6 +705,8 @@ def _report( 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")], + warmup_oom=warmup_oom, + require_feasible_placement=require_feasible_placement, skip_reason=skip_reason, ) @@ -603,11 +766,195 @@ def test_does_not_replace_explicit_default_shape(self): class TestPlanAutoResidency: + def test_lower_latency_utility_keeps_the_measured_placement(self): + resident = ResidencyTarget( + component_name="transformer", + residency_mode=RESIDENT, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=20 * GIB_BYTES, + permanent_residency=True, + target_device_weight_bytes=20 * GIB_BYTES, + current_placement=True, + ) + offloaded = ResidencyTarget( + component_name="transformer", + residency_mode=RESIDENT, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=19 * GIB_BYTES, + device_transition_delta_bytes=-20 * GIB_BYTES, + inactive_device_delta_bytes=-20 * GIB_BYTES, + present_device_delta_bytes=-20 * GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=100, + estimated_gib=30, + candidates=[resident, offloaded], + estimated_request_duration_ns=1_000_000_000, + candidate_latency_savings_ns={ + resident.option_key(): 500_000_000, + offloaded.option_key(): 475_000_000, + }, + ) + ] + ) + + assert plan.changes == [] + + def test_material_latency_gain_can_replace_the_measured_placement(self): + resident = ResidencyTarget( + component_name="transformer", + residency_mode=RESIDENT, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=20 * GIB_BYTES, + permanent_residency=True, + target_device_weight_bytes=20 * GIB_BYTES, + current_placement=True, + ) + offloaded = ResidencyTarget( + component_name="transformer", + residency_mode=RESIDENT, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=19 * GIB_BYTES, + device_transition_delta_bytes=-20 * GIB_BYTES, + inactive_device_delta_bytes=-20 * GIB_BYTES, + present_device_delta_bytes=-20 * GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=100, + estimated_gib=30, + candidates=[resident, offloaded], + estimated_request_duration_ns=1_000_000_000, + candidate_latency_savings_ns={ + resident.option_key(): 460_000_000, + offloaded.option_key(): 500_000_000, + }, + ) + ] + ) + + assert plan.changes == [offloaded] + + def test_free_host_pin_gain_is_selected(self): + pageable = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=5 * GIB_BYTES, + h2d_bytes_per_request=10 * GIB_BYTES, + target_layerwise_resident_layers=(5,), + target_layerwise_pinned_layers=((),), + target_device_weight_bytes=5 * GIB_BYTES, + current_placement=True, + ) + pinned = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=5 * GIB_BYTES, + h2d_bytes_per_request=11 * GIB_BYTES, + target_layerwise_resident_layers=(5,), + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=GIB_BYTES, + target_device_weight_bytes=5 * GIB_BYTES, + target_pinned_host_bytes=GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=100, + estimated_gib=30, + candidates=[pageable, pinned], + host_pin_capacity_gib=10, + estimated_request_duration_ns=1_000_000_000, + candidate_latency_savings_ns={ + pageable.option_key(): 100_000_000, + pinned.option_key(): 110_000_000, + }, + ) + ] + ) + + assert plan.changes == [pinned] + + def test_oom_phase_forces_a_lower_memory_complete_target(self): + resident = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=20 * GIB_BYTES, + permanent_residency=True, + target_device_weight_bytes=20 * GIB_BYTES, + current_placement=True, + ) + layerwise = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=2 * GIB_BYTES, + h2d_bytes_per_request=10 * GIB_BYTES, + target_layerwise_resident_layers=(2,), + target_layerwise_pinned_layers=((),), + device_transition_delta_bytes=-15 * GIB_BYTES, + active_device_delta_bytes=-15 * GIB_BYTES, + present_device_delta_bytes=-15 * GIB_BYTES, + inactive_device_delta_bytes=-20 * GIB_BYTES, + target_device_weight_bytes=2 * GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=30, + estimated_gib=29, + warmup_oom=True, + phase_peaks_gib={"denoise": 29}, + phase_components={"denoise": ("transformer",)}, + phase_present_components={"denoise": ("transformer",)}, + component_weight_gib={"transformer": 20}, + candidates=[resident, layerwise], + ) + ] + ) + + assert plan.recovering_from_oom + assert plan.resource_budget_bytes["gpu:rank0:denoise"] == -3 * GIB_BYTES + assert [candidate.option_key() for candidate in plan.changes] == [ + layerwise.option_key() + ] + 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" + + def test_oom_recovery_state_survives_an_infeasible_plan(self): + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=30, + estimated_gib=29, + warmup_oom=True, + phase_peaks_gib={"denoise": 29}, + candidates=[_candidate("vae")], + ) + ] + ) + + assert plan.recovering_from_oom + assert plan.skip_reason == "no placement satisfies all resource budgets" assert not plan.changes def test_missing_estimate_skips(self): @@ -629,11 +976,57 @@ 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_uncalibrated_small_card_reserve_has_three_gib_floor(self): + plan = plan_auto_residency(reports=[_report(budget_gib=12, estimated_gib=5)]) + assert plan.reserve_bytes == 3 * GIB_BYTES + + def test_pre_warmup_report_enforces_weight_feasibility_without_oom(self): + resident = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=11 * GIB_BYTES, + h2d_bytes_per_request=11 * GIB_BYTES, + permanent_residency=True, + target_device_weight_bytes=11 * GIB_BYTES, + 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=5 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + active_device_delta_bytes=-9 * GIB_BYTES, + inactive_device_delta_bytes=-10 * GIB_BYTES, + target_device_weight_bytes=GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=12, + estimated_gib=11, + phase_peaks_gib={"static:transformer": 11}, + phase_components={"static:transformer": ("transformer",)}, + candidates=[resident, layerwise], + require_feasible_placement=True, + ) + ] + ) + + assert not plan.recovering_from_oom + assert [candidate.target_mode() for candidate in plan.changes] == [ + LAYERWISE_OFFLOAD + ] + 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. + # Keep enough room for unmeasured activations without applying the + # datacenter-card 4 GiB floor verbatim. plan = plan_auto_residency(reports=[_report(budget_gib=10, estimated_gib=4)]) - assert plan.reserve_bytes == 2 * GIB_BYTES + assert plan.reserve_bytes == 3 * GIB_BYTES def test_target_workload_measurement_uses_tighter_reserve(self): plan = plan_auto_residency( @@ -647,6 +1040,18 @@ def test_target_workload_measurement_uses_tighter_reserve(self): ) assert plan.reserve_bytes == 5 * GIB_BYTES + def test_measured_small_card_reserve_drops_to_ten_percent(self): + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=12, + estimated_gib=10, + target_workload_measured=True, + ) + ] + ) + assert plan.reserve_bytes == int(1.2 * GIB_BYTES) + def test_target_workload_measurement_is_replica_wide(self): candidate = _candidate("vae", weight_gib=1) plan = plan_auto_residency( @@ -684,7 +1089,7 @@ def test_optimizes_promotion_by_h2d_savings(self): ) assert [c.component_name for c in plan.changes] == ["dit"] - def test_long_request_does_not_spend_vram_on_tiny_encoder_gain(self): + def test_long_request_uses_available_vram_for_encoder_gain(self): transformer = _candidate( "transformer", mode=LAYERWISE_OFFLOAD, @@ -709,7 +1114,8 @@ def test_long_request_does_not_spend_vram_on_tiny_encoder_gain(self): ) assert [candidate.component_name for candidate in plan.changes] == [ - "transformer" + "transformer", + "text_encoder", ] def test_output_rank_timing_applies_to_untimed_peer_ranks(self): @@ -746,7 +1152,8 @@ def test_output_rank_timing_applies_to_untimed_peer_ranks(self): ) assert [candidate.component_name for candidate in plan.changes] == [ - "transformer" + "transformer", + "text_encoder", ] def test_smaller_component_still_fits_after_skipping_a_big_one(self): @@ -1110,6 +1517,44 @@ def test_materialized_layerwise_transition_does_not_double_count_weights(self): LAYERWISE_OFFLOAD ] + def test_untracked_oom_peak_constrains_synthetic_component_phase(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=10 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + active_device_delta_bytes=-10 * GIB_BYTES, + inactive_device_delta_bytes=GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=12, + estimated_gib=11.75, + phase_peaks_gib={"request:untracked": 11.75}, + phase_present_components={"request:untracked": ("vae",)}, + candidates=[component_offload, layerwise], + warmup_oom=True, + ) + ] + ) + + 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=[ @@ -1122,7 +1567,26 @@ def test_reports_reserve_shortfall_for_a_newly_calibrated_placement(self): ] ) - assert plan.current_placement_reserve_shortfall_bytes == 2 * GIB_BYTES + assert plan.current_placement_reserve_shortfall_bytes == GIB_BYTES + + def test_failed_candidate_margin_only_constrains_future_growth(self): + candidate = _candidate("transformer", weight_gib=5, h2d_gib=5) + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=40, + estimated_gib=30, + planning_headroom_correction_gib=3, + target_workload_measured=True, + candidates=[candidate], + phase_peaks_gib={"denoise": 30}, + ) + ] + ) + + assert plan.current_placement_reserve_shortfall_bytes == 0 + assert plan.resource_budget_bytes["gpu:rank0:denoise"] == 3 * GIB_BYTES + assert plan.changes == [] def test_unobserved_later_component_gets_a_conservative_phase(self): first_transformer = _candidate("first_transformer", weight_gib=26, h2d_gib=100) @@ -1155,6 +1619,69 @@ def test_unobserved_later_component_gets_a_conservative_phase(self): ] assert "gpu:rank0:unobserved:second_transformer" in plan.resource_budget_bytes + def test_unobserved_phase_excludes_another_components_transient_weights(self): + transformer_component = 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, + ) + transformer_layerwise = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_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=((),), + active_device_delta_bytes=-10 * GIB_BYTES, + inactive_device_delta_bytes=GIB_BYTES, + ) + vae = ResidencyTarget( + component_name="vae", + residency_mode=RESIDENT, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=GIB_BYTES, + current_placement=True, + permanent_residency=True, + target_device_weight_bytes=GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=12, + estimated_gib=11.75, + warmup_oom=True, + candidates=[transformer_component, transformer_layerwise, vae], + phase_peaks_gib={ + "setup": 5, + "prefetch:transformer": 11.75, + }, + phase_components={ + "prefetch:transformer": ("transformer",), + }, + phase_present_components={ + "setup": ("vae",), + "prefetch:transformer": ("transformer", "vae"), + }, + component_weight_gib={"vae": 1}, + component_active_weight_gib={"transformer": 11, "vae": 1}, + ) + ] + ) + + assert plan.skip_reason is None + assert [candidate.target_mode() for candidate in plan.changes] == [ + LAYERWISE_OFFLOAD + ] + assert plan.resource_budget_bytes["gpu:rank0:unobserved:vae"] == ( + 12 * GIB_BYTES - 5 * GIB_BYTES - plan.reserve_bytes + ) + def test_redundant_phase_constraints_collapse_to_the_highest_peak(self): plan = plan_auto_residency( reports=[ @@ -1815,14 +2342,18 @@ def test_filters_and_sizes(self): for candidate in candidates if candidate.permanent_residency } - # resident vae, explicit placement, custom strategy, non-module: out - assert set(permanent_candidates) == {"text_encoder", "transformer"} + # Unset resident components remain tunable. Explicit placement, + # custom strategy, and non-module objects stay outside the frontier. + assert set(permanent_candidates) == {"text_encoder", "transformer", "vae"} + assert permanent_candidates["vae"].current_placement assert permanent_candidates["text_encoder"].h2d_bytes_per_request == ( permanent_candidates["text_encoder"].target_resident_weight_bytes ) - # layerwise DiT re-streams its layers once per denoise step + # The all-pageable layerwise DiT re-streams its layers every step. assert permanent_candidates["transformer"].h2d_bytes_per_request == ( - permanent_candidates["transformer"].target_resident_weight_bytes * 100 + permanent_candidates["transformer"].target_resident_weight_bytes + * 50 + * PAGEABLE_H2D_COST_MULTIPLIER ) assert any( candidate.component_name == "transformer" @@ -1973,6 +2504,58 @@ def test_non_dit_layerwise_frontier_includes_partial_residency(self): assert component.target_layerwise_pinned_layers == ((),) assert component.pinned_host_delta_bytes == 0 + @pytest.mark.parametrize( + "component_name", + ( + "text_encoder", + "text_encoder_2", + "image_encoder", + "image_encoder_2", + "transformer", + "transformer_2", + "video_dit", + "video_dit_2", + "audio_dit", + "audio_dit_2", + "connectors", + "dual_tower_bridge", + "vae", + "vae_2", + "video_vae", + "audio_vae", + "condition_image_encoder", + ), + ) + @pytest.mark.parametrize("startup_mode", (COMPONENT_OFFLOAD, RESIDENT)) + def test_every_serving_component_group_gets_a_complete_frontier( + self, component_name, startup_mode + ): + module = _FakeLazyLayerwiseDit(num_layers=4) + + candidates = collect_residency_targets( + modules={component_name: module}, + residency_mode_of=self._modes({component_name: startup_mode}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=8, + layerwise_tuning_of=lambda _name, _dit_group: ( + 0.0, + 0.0, + "leading", + ), + ) + + assert {candidate.target_mode() for candidate in candidates} == { + COMPONENT_OFFLOAD, + LAYERWISE_OFFLOAD, + RESIDENT, + } + assert any( + candidate.target_mode() == LAYERWISE_OFFLOAD + and candidate.target_layerwise_resident_layers == (2,) + for candidate in candidates + ) + def test_component_offload_frontier_can_expand_layerwise_lazily(self): module = _FakeLazyLayerwiseDit(num_layers=4) @@ -2046,16 +2629,14 @@ def test_component_weight_footprint_deduplicates_tied_storage(self): "transformer": backing.untyped_storage().nbytes() } - def test_auto_resident_component_keeps_offload_as_a_target(self): + def test_default_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} @@ -2065,6 +2646,120 @@ def test_auto_resident_component_keeps_offload_as_a_target(self): assert by_mode[COMPONENT_OFFLOAD].device_transition_delta_bytes < 0 assert by_mode[COMPONENT_OFFLOAD].host_materialize_scratch_bytes > 0 + def test_required_resident_component_has_no_tunable_frontier(self): + candidates = collect_residency_targets( + modules={"transformer_2": nn.Linear(4, 4)}, + residency_mode_of=self._modes({"transformer_2": RESIDENT}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=1, + required_resident_components={"transformer_2"}, + ) + + assert candidates == [] + + def test_default_resident_dit_exposes_complete_layerwise_frontier(self): + module = _FakeLazyLayerwiseDit(num_layers=3) + candidates = collect_residency_targets( + modules={"transformer_2": module}, + residency_mode_of=self._modes({"transformer_2": RESIDENT}), + 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"), + ) + + assert {candidate.target_mode() for candidate in candidates} == { + COMPONENT_OFFLOAD, + LAYERWISE_OFFLOAD, + RESIDENT, + } + assert { + candidate.target_layerwise_resident_layers + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + } == {(0,), (1,), (2,), (3,)} + assert all( + candidate.residency_mode == RESIDENT + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + ) + + def test_resident_seed_preserves_unconfigured_layerwise_frontier(self): + module = _FakeLazyLayerwiseDit(num_layers=3) + candidates = collect_residency_targets( + modules={"text_encoder": module}, + residency_mode_of=self._modes({"text_encoder": RESIDENT}), + baseline_residency_mode_of=self._modes({"text_encoder": LAYERWISE_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=1, + layerwise_tuning_of=lambda _name, _dit_group: (1.0, 0.0, "leading"), + ) + + assert {candidate.target_mode() for candidate in candidates} == { + COMPONENT_OFFLOAD, + LAYERWISE_OFFLOAD, + RESIDENT, + } + assert all(candidate.residency_mode == RESIDENT for candidate in candidates) + assert next( + candidate for candidate in candidates if candidate.target_mode() == RESIDENT + ).current_placement + + 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_virtual_layerwise_frontier_includes_host_pin_prefixes(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 = [ + candidate + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + and candidate.target_layerwise_resident_layers == (0,) + ] + assert {candidate.target_layerwise_pinned_layers for candidate in streamed} == { + ((),), + ((0,),), + ((0, 1),), + ((0, 1, 2),), + } + fully_pinned = next( + candidate + for candidate in streamed + if candidate.target_layerwise_pinned_layers == ((0, 1, 2),) + ) + assert fully_pinned.pinned_host_delta_bytes > 0 + assert ( + fully_pinned.host_pin_scratch_bytes == fully_pinned.pinned_host_delta_bytes + ) + def test_layerwise_frontier_accounts_for_immediate_materialization(self): manager = _FakeLayerwiseManager( { @@ -2095,7 +2790,6 @@ def test_layerwise_frontier_accounts_for_immediate_materialization(self): explicit_residency_mode_of=lambda _name: None, custom_strategy_names=(), num_inference_steps=10, - auto_resident_components={"transformer"}, ) streamed = next( candidate @@ -2124,6 +2818,7 @@ def test_lazy_layerwise_configuration_rolls_back_to_component_offload(self): for candidate in candidates if candidate.target_mode() == LAYERWISE_OFFLOAD and candidate.target_layerwise_resident_layers == (2,) + and candidate.target_layerwise_pinned_layers == ((0, 1, 2, 3),) ) applied = apply_residency_changes( @@ -2134,6 +2829,12 @@ 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 module.layerwise_offload_managers[0].pinned_layer_indices() == ( + 0, + 1, + 2, + 3, + ) assert args.auto_modes == {"transformer": LAYERWISE_OFFLOAD} manager = module.layerwise_offload_managers[0] @@ -2853,6 +3554,19 @@ def _skip_reason(self, args): return auto_residency_skip_reason(args) + def _static_skip_reason(self, args): + from sglang.multimodal_gen.runtime.server_args.auto_tune import ( + auto_residency_static_skip_reason, + ) + + return auto_residency_static_skip_reason(args) + + def test_static_planning_does_not_require_warmup(self, monkeypatch): + monkeypatch.delenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", raising=False) + monkeypatch.delenv("SGLANG_CACHE_DIT_ENABLED", raising=False) + assert self._static_skip_reason(self._base_args(warmup_mode="off")) is None + assert "server warmup" in self._skip_reason(self._base_args(warmup_mode="off")) + def test_env_kill_switch(self, monkeypatch): monkeypatch.setenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", "1") reason = self._skip_reason(self._base_args()) @@ -2869,7 +3583,7 @@ def test_manual_performance_mode(self, monkeypatch): "overrides, expected_fragment", [ ({"warmup_mode": "request"}, "server warmup"), - ({"disagg_role": "denoiser"}, "server warmup"), + ({"disagg_role": "denoiser"}, "disaggregated role"), ({"backend": "diffusers"}, "diffusers"), ({"enable_breakable_cuda_graph": True}, "CUDA graph"), # compile warmup strips the memory layout (layerwise DiT + @@ -2878,7 +3592,6 @@ def test_manual_performance_mode(self, monkeypatch): ({"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"), ], ) @@ -2914,3 +3627,11 @@ def test_eligible_path_reaches_platform_gate(self, monkeypatch): 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" + + def test_transformer_weight_override_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(transformer_weights_path="/x.safetensors") + ) + assert reason is None or reason == "requires CUDA" 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..973c1ac4291f 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 @@ -27,6 +27,156 @@ class TestAutoResidencyWarmup(unittest.TestCase): + def test_residency_hint_uses_effective_device_budget(self): + worker = GPUWorker.__new__(GPUWorker) + worker._auto_residency_budget_bytes = mock.Mock(return_value=12 * GIB_BYTES) + worker.get_can_stay_resident_components = mock.Mock(return_value=[]) + output = SimpleNamespace(metrics=None) + snapshot = SimpleNamespace( + peak_reserved_mb=11 * 1024, + peak_allocated_mb=10 * 1024, + ) + + with mock.patch.object( + gpu_worker_module, "capture_memory_snapshot", return_value=snapshot + ): + worker.do_mem_analysis(output) + + worker.get_can_stay_resident_components.assert_called_once_with(1.0) + self.assertEqual(output.peak_memory_mb, 11 * 1024) + + def test_sync_startup_uses_shared_auto_residency_sequence(self): + server_args = SimpleNamespace() + forward = mock.Mock() + + with ( + mock.patch.object( + server_warmup, "maybe_apply_pre_warmup_auto_residency" + ) as apply_static, + mock.patch.object(server_warmup, "run_async_client_warmup") as warmup, + mock.patch.object(server_warmup, "maybe_apply_auto_residency") as refine, + mock.patch.object( + server_warmup, + "should_run_synthetic_server_warmup", + return_value=True, + ), + ): + server_warmup.run_sync_startup_warmup(server_args, forward) + + apply_static.assert_awaited_once() + warmup.assert_awaited_once() + refine.assert_awaited_once() + self.assertIs(apply_static.await_args.args[0], server_args) + self.assertIs(warmup.await_args.args[0], server_args) + self.assertIs(refine.await_args.args[0], server_args) + self.assertTrue(warmup.await_args.kwargs["fail_open"]) + self.assertIs(apply_static.await_args.args[1], warmup.await_args.args[1]) + self.assertIs(apply_static.await_args.args[1], refine.await_args.args[1]) + + def test_sync_startup_static_planning_does_not_require_warmup(self): + with ( + mock.patch.object( + server_warmup, "maybe_apply_pre_warmup_auto_residency" + ) as apply_static, + mock.patch.object( + server_warmup, + "should_run_synthetic_server_warmup", + return_value=False, + ), + mock.patch.object(server_warmup, "run_async_client_warmup") as warmup, + mock.patch.object(server_warmup, "maybe_apply_auto_residency") as refine, + ): + server_warmup.run_sync_startup_warmup(SimpleNamespace(), mock.Mock()) + + apply_static.assert_awaited_once() + warmup.assert_not_awaited() + refine.assert_not_awaited() + + def test_pre_warmup_planner_uses_static_action(self): + actions = [] + + async def forward(req): + actions.append(req.action) + return OutputBatch(output={"status": PLACEMENT_STATUS_SKIPPED}) + + with mock.patch.object( + server_warmup, + "should_apply_pre_warmup_auto_residency", + return_value=True, + ): + asyncio.run( + server_warmup.maybe_apply_pre_warmup_auto_residency( + SimpleNamespace(), forward + ) + ) + + self.assertEqual(actions, ["apply_static"]) + + def test_auto_residency_uses_first_oom_without_degrading_probe(self): + req = SimpleNamespace() + forward = mock.AsyncMock(return_value=OutputBatch(error="CUDA out of memory")) + + with ( + mock.patch.object( + server_warmup, "auto_residency_skip_reason", return_value=None + ), + mock.patch.object( + server_warmup, + "should_include_warmup_image", + return_value=False, + ), + mock.patch.object( + server_warmup, "build_client_warmup_reqs", return_value=[req] + ), + mock.patch.object(server_warmup, "_degrade_after_oom") as degrade, + ): + asyncio.run( + server_warmup.run_async_client_warmup( + SimpleNamespace(), forward, fail_open=True + ) + ) + + forward.assert_awaited_once_with(req) + degrade.assert_not_called() + + def test_non_auto_warmup_keeps_oom_probe_degradation(self): + req = SimpleNamespace() + lighter = SimpleNamespace() + forward = mock.AsyncMock( + side_effect=[ + OutputBatch(error="CUDA out of memory"), + OutputBatch(), + ] + ) + + with ( + mock.patch.object( + server_warmup, + "auto_residency_skip_reason", + return_value="not auto", + ), + mock.patch.object( + server_warmup, + "should_include_warmup_image", + return_value=False, + ), + mock.patch.object( + server_warmup, "build_client_warmup_reqs", return_value=[req] + ), + mock.patch.object( + server_warmup, "_degrade_after_oom", return_value=lighter + ) as degrade, + ): + asyncio.run( + server_warmup.run_async_client_warmup(SimpleNamespace(), forward) + ) + + self.assertEqual( + [call.args[0] for call in forward.await_args_list], + [req, lighter], + ) + degrade.assert_called_once_with(mock.ANY, req) + def test_worker_rolls_back_a_materially_slower_calibrated_round(self): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 @@ -49,14 +199,17 @@ 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() @@ -75,6 +228,11 @@ def test_new_placement_that_spends_the_reserve_rolls_back(self): worker.server_args = SimpleNamespace() worker._auto_residency_warmup_records = [] worker._auto_residency_round_sizes = [1] + worker._auto_residency_budget_correction_bytes = 0 + worker._auto_residency_last_applied_plan = AutoResidencyPlan( + resource_budget_bytes={"gpu:rank0:denoise": 8 * GIB_BYTES}, + resource_delta_bytes={"gpu:rank0:denoise": 7 * GIB_BYTES}, + ) worker._build_auto_residency_report = mock.Mock( return_value=RankResidencyReport( rank=0, @@ -110,6 +268,11 @@ def test_new_placement_that_spends_the_reserve_rolls_back(self): cause="VRAM reserve exceeded by 2.0 GiB", already_failed=False, latest_round_only=True, + retryable=True, + ) + self.assertEqual( + worker._auto_residency_budget_correction_bytes, + 2 * GIB_BYTES, ) def test_worker_rollback_response_rewarms_the_restored_layout(self): @@ -144,18 +307,22 @@ async def forward(req): rewarm=True, ) - def test_replans_until_the_calibrated_layout_reaches_a_fixed_point(self): + def test_retryable_rollback_rewarms_and_replans(self): responses = iter( [ - OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), - OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), + OutputBatch( + error="VRAM reserve exceeded", + output={ + "status": PLACEMENT_STATUS_ROLLED_BACK, + "retryable": True, + }, + ), OutputBatch(output={"status": PLACEMENT_STATUS_SKIPPED}), ] ) actions = [] async def forward(req): - self.assertIsInstance(req, AutoResidencyReq) actions.append(req.action) return next(responses) @@ -172,8 +339,76 @@ 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", "apply"]) + rewarm.assert_awaited_once_with( + server_args, + forward, + fail_open=True, + rewarm=True, + ) + + def test_oom_recovery_keeps_demotion_and_replans(self): + responses = iter( + [ + OutputBatch( + output={ + "status": PLACEMENT_STATUS_ADJUSTED, + "recovering_from_oom": True, + } + ), + OutputBatch(output={"status": PLACEMENT_STATUS_SKIPPED}), + ] + ) + actions = [] + + async def forward(req): + actions.append(req.action) + return next(responses) + + rewarm = mock.AsyncMock(side_effect=RuntimeError("CUDA out of memory")) + 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"]) + rewarm.assert_awaited_once_with( + server_args, + forward, + fail_open=False, + rewarm=True, + ) + + def test_stops_after_one_successful_calibrated_adjustment(self): + actions = [] + + async def forward(req): + self.assertIsInstance(req, AutoResidencyReq) + actions.append(req.action) + return OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}) + + 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() def test_failed_first_calibration_rolls_back_and_rewarms(self): actions = [] @@ -228,22 +463,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_successful_calibration_does_not_start_another_apply_round(self): actions = [] - responses = iter( - [ - OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), - 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 + return OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}) rewarm = mock.AsyncMock() server_args = SimpleNamespace( @@ -255,14 +481,10 @@ 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(actions, ["apply"]) self.assertEqual(rewarm.await_count, 1) def test_failed_later_round_keeps_earlier_calibrated_promotions(self): 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..ba5e04e6442c 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 @@ -16,7 +16,7 @@ import unittest from collections import deque from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import torch @@ -346,6 +346,30 @@ def test_diff_generator_runs_explicit_warmup_through_scheduler_client(self): self.assertTrue(req.extra["server_based_warmup"]) self.assertEqual(req.extra["warmup_total"], 1) + def test_diff_generator_runs_server_style_startup_warmup(self): + generator = object.__new__(DiffGenerator) + generator.server_args = server_args = MagicMock() + + with ( + patch( + "sglang.multimodal_gen.runtime.entrypoints.diffusion_generator.should_run_synthetic_server_warmup", + return_value=True, + ), + patch( + "sglang.multimodal_gen.runtime.entrypoints.diffusion_generator.run_sync_startup_warmup" + ) as startup_warmup, + patch( + "sglang.multimodal_gen.runtime.entrypoints.diffusion_generator.run_sync_client_warmup" + ) as request_warmup, + ): + generator._run_client_warmup_if_needed() + + startup_warmup.assert_called_once_with( + server_args, + ANY, + ) + request_warmup.assert_not_called() + def test_server_based_warmup_uses_model_default_negative_prompt(self): server_args = MagicMock() server_args.warmup_steps = 1 diff --git a/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py b/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py index d0f0f7294139..30083a8237c3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py @@ -1,4 +1,5 @@ import json +from types import SimpleNamespace from sglang.multimodal_gen.runtime.managers.memory_managers.component_loading_order import ( ComponentLoadSpec, @@ -6,6 +7,14 @@ infer_component_weight_size_bytes, order_component_load_specs, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_weight_inventory import ( + ComponentWeightSource, + estimate_component_weight_inventory, + infer_safetensors_weight_stats_by_prefix, +) +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) def _spec( @@ -15,7 +24,9 @@ def _spec( module_name=component_name, load_module_name=component_name, component_model_path=component_model_path, - transformers_or_diffusers="diffusers", + transformers_or_diffusers=( + "transformers" if component_name.startswith("text_encoder") else "diffusers" + ), architecture=None, index=index, ) @@ -128,3 +139,145 @@ def test_component_load_order_prefers_inferred_safetensors_size(tmp_path): assert ordered_names == ["text_encoder", "transformer", "scheduler"] assert infer_component_weight_size_bytes(str(large_encoder_path)) == 64 + + +def test_safetensors_inventory_splits_bundled_components_by_prefix(tmp_path): + checkpoint = tmp_path / "model.safetensors" + header = { + "model.block.weight": { + "dtype": "F16", + "shape": [8], + "data_offsets": [0, 16], + }, + "vae.decoder.weight": { + "dtype": "F16", + "shape": [4], + "data_offsets": [16, 24], + }, + "model.proj.weight": { + "dtype": "F16", + "shape": [2], + "data_offsets": [24, 28], + }, + } + header_bytes = json.dumps(header).encode("utf-8") + checkpoint.write_bytes( + len(header_bytes).to_bytes(8, "little") + header_bytes + b"\0" * 28 + ) + + assert infer_safetensors_weight_stats_by_prefix(str(checkpoint)) == { + "model": (20, 10), + "vae": (8, 4), + } + + +def test_component_weight_inventory_uses_explicit_bundled_size(tmp_path): + component_dir = tmp_path / "transformer" + component_dir.mkdir() + _write_safetensors(component_dir / "model.safetensors", 64) + + inventory = estimate_component_weight_inventory( + [ + ComponentWeightSource( + "transformer", + str(component_dir), + supports_fsdp_loading=True, + ), + ComponentWeightSource( + "bundled_vae", + "/shared/model", + checkpoint_bytes=8, + parameter_count=2, + target_element_size=4, + ), + ] + ) + + assert [ + (item.component_name, item.checkpoint_bytes, item.parameter_count) + for item in inventory + ] == [ + ("transformer", 64, 32), + ("bundled_vae", 8, 2), + ] + assert inventory[0].materialized_bytes() == 64 + assert inventory[1].materialized_bytes() == 8 + assert inventory[0].supports_fsdp_loading + assert not inventory[1].supports_fsdp_loading + + +def test_component_weight_inventory_distinguishes_config_only_and_unknown(tmp_path): + config_only = tmp_path / "tokenizer" + config_only.mkdir() + (config_only / "tokenizer_config.json").write_text("{}") + + inventory = estimate_component_weight_inventory( + [ + ComponentWeightSource("tokenizer", str(config_only)), + ComponentWeightSource("transformer", "remote/model"), + ] + ) + + assert [ + (item.component_name, item.checkpoint_bytes, item.parameter_count) + for item in inventory + ] == [ + ("tokenizer", 0, 0), + ("transformer", None, None), + ] + + +def test_preload_inventory_uses_the_actual_weight_override(): + server_args = SimpleNamespace( + component_weights_paths={"text_encoder": "/weights/text.safetensors"}, + transformer_weights_path="/weights/dit.gguf", + pipeline_config=SimpleNamespace( + dit_precision="bf16", + vae_precision="fp32", + image_encoder_precision="fp16", + text_encoder_precisions=("bf16",), + ), + ) + + text_source = ComposedPipelineBase._component_weight_source( + _spec("text_encoder", 0, "/model/text_encoder"), server_args + ) + transformer_source = ComposedPipelineBase._component_weight_source( + _spec("transformer", 1, "/model/transformer"), server_args + ) + secondary_source = ComposedPipelineBase._component_weight_source( + _spec("transformer_2", 2, "/model/transformer_2"), server_args + ) + + assert text_source.component_model_path == "/weights/text.safetensors" + assert transformer_source.component_model_path == "/weights/dit.gguf" + assert secondary_source.component_model_path == "/model/transformer_2" + assert transformer_source.target_element_size == 2 + assert secondary_source.target_element_size == 2 + assert not text_source.supports_fsdp_loading + assert transformer_source.supports_fsdp_loading + assert secondary_source.supports_fsdp_loading + + +def test_preload_inventory_resolves_group_precision_fallbacks(): + server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + dit_precision="bf16", + vae_precision="fp32", + image_encoder_precision="fp16", + text_encoder_precisions=("bf16",), + ) + ) + + assert ( + ComposedPipelineBase._component_target_element_size( + "unconditional_transformer", server_args + ) + == 2 + ) + assert ( + ComposedPipelineBase._component_target_element_size( + "condition_image_encoder", server_args + ) + == 4 + ) 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..596fb1560c5b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_residency.py @@ -1,8 +1,11 @@ from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, patch import torch +from sglang.multimodal_gen.runtime.managers.memory_managers import ( + component_residency_strategies, +) from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( ComponentResidencyManager, ComponentUse, @@ -55,6 +58,27 @@ def test_component_offload_keeps_preferred_component_after_warmup(): strategy.finish_use.assert_not_called() +def test_component_offload_restores_cpu_after_failed_device_transfer(): + strategy = ComponentOffloadStrategy() + module = Mock() + use = ComponentUse("DenoiseStage", "transformer") + error = RuntimeError("CUDA out of memory") + + with patch.object( + component_residency_strategies, + "_module_to_local_device", + side_effect=error, + ): + try: + strategy.prepare_for_use(module, use, ResidencyState()) + except RuntimeError as caught: + assert caught is error + else: + raise AssertionError("expected the simulated device transfer to fail") + + module.to.assert_called_once_with("cpu", non_blocking=False) + + def test_request_tail_uses_dynamic_component_instance(): pipeline = SimpleNamespace( modules={}, @@ -79,6 +103,30 @@ def test_request_tail_uses_dynamic_component_instance(): ) +def test_placement_modules_include_lazy_request_components(): + pipeline_module = torch.nn.Linear(2, 2) + stale_seen_module = torch.nn.Linear(2, 2) + lazy_module = torch.nn.Linear(2, 2) + pipeline = SimpleNamespace( + modules={"transformer": pipeline_module}, + _stage_name_mapping={}, + component_residency_strategies={}, + ) + manager = ComponentResidencyManager( + pipeline, + SimpleNamespace(enable_layerwise_nvtx_marker=False), + ) + manager._modules_seen = { + "transformer": stale_seen_module, + "condition_image_encoder": lazy_module, + } + + assert manager.placement_modules() == { + "transformer": pipeline_module, + "condition_image_encoder": lazy_module, + } + + def test_strategy_cache_replaces_stale_component_instance(): server_args = SimpleNamespace( enable_layerwise_nvtx_marker=False, @@ -234,6 +282,41 @@ def test_warmup_records_use_and_transition_peaks(monkeypatch): ) +def test_failed_warmup_phase_survives_cleanup(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) + + use = ComponentUse("denoise", "transformer") + stage = _Stage(use) + pipeline = SimpleNamespace( + modules={"transformer": torch.nn.Linear(2, 2)}, + _stage_name_mapping={"denoise": stage}, + component_residency_strategies={}, + ) + server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + manager = ComponentResidencyManager(pipeline, server_args) + 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_warmup_prefetch(use) + + manager.capture_failed_warmup_phase() + manager._warmup_phase_peaks = {} + manager._completed_warmup_phase_peaks = {"request:cleanup": WarmupPhasePeak((), 2)} + + peaks = manager.take_warmup_phase_peaks() + assert peaks["0:denoise:prefetch:transformer"] == WarmupPhasePeak( + ("transformer",), 7, used_components=("transformer",) + ) + assert peaks["request:cleanup"] == WarmupPhasePeak((), 2) + + def test_warmup_records_full_weight_transition_without_preparing(monkeypatch): device_module = SimpleNamespace( is_available=lambda: True, @@ -399,6 +482,46 @@ def test_warmup_splits_sequential_component_transition(monkeypatch): ) +def test_warmup_preserves_repeated_phase_layouts(monkeypatch): + device_module = SimpleNamespace( + is_available=lambda: True, + reset_peak_memory_stats=Mock(), + max_memory_allocated=lambda: 7, + ) + 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 = SimpleNamespace(enable_layerwise_nvtx_marker=False) + manager = ComponentResidencyManager(pipeline, server_args) + manager.refresh_pipeline(pipeline) + manager.begin_request([stage], SimpleNamespace(is_warmup=True), server_args) + + manager._begin_warmup_phase( + key="0:stage:between", + components=("text_encoder",), + used_components=(), + ) + manager._begin_warmup_phase( + key="0:stage:between", + components=("transformer",), + used_components=(), + ) + manager._record_warmup_phase_peak() + + assert manager._warmup_phase_peaks["0:stage:between"] == WarmupPhasePeak( + ("text_encoder",), 7 + ) + assert manager._warmup_phase_peaks[ + "0:stage:between:occurrence:1" + ] == WarmupPhasePeak(("transformer",), 7) + + def _manager_for_stage(stage, modules): pipeline = SimpleNamespace( modules=modules, 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..fa6869f7b35b 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 @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 +import json import unittest from types import SimpleNamespace @@ -205,6 +206,76 @@ def test_layerwise_texture_component_starts_on_cpu(self): ) +def test_hunyuan3d_shape_inventory_splits_bundled_checkpoint(tmp_path): + checkpoint = tmp_path / "model.safetensors" + header = { + "conditioner.encoder.weight": { + "dtype": "F16", + "shape": [3], + "data_offsets": [0, 6], + }, + "model.block.weight": { + "dtype": "F16", + "shape": [5], + "data_offsets": [6, 16], + }, + "vae.decoder.weight": { + "dtype": "F16", + "shape": [2], + "data_offsets": [16, 20], + }, + } + encoded_header = json.dumps(header).encode() + checkpoint.write_bytes( + len(encoded_header).to_bytes(8, "little") + encoded_header + b"\0" * 20 + ) + + sources = Hunyuan3D2Pipeline._shape_component_weight_sources( + str(checkpoint), True, 2 + ) + + assert [ + (source.component_name, source.checkpoint_bytes, source.parameter_count) + for source in sources + ] == [ + ("hy3dshape_conditioner", 6, 3), + ("hy3dshape_model", 10, 5), + ("hy3dshape_vae", 4, 2), + ] + + +def test_hunyuan3d_texture_inventory_includes_every_weighted_component(tmp_path): + config = Hunyuan3D2PipelineConfig( + paint_subfolder="paint", + delight_subfolder="delight", + ) + for relative_path in ( + "paint/vae/config.json", + "paint/unet/config.json", + "paint/scheduler/scheduler_config.json", + "delight/vae/config.json", + "delight/unet/config.json", + "delight/scheduler/scheduler_config.json", + "delight/text_encoder/config.json", + "delight/tokenizer/tokenizer_config.json", + ): + path = tmp_path / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("{}") + + sources = Hunyuan3D2Pipeline._texture_component_weight_sources( + SimpleNamespace(model_path=str(tmp_path)), config + ) + + assert [source.component_name for source in sources] == [ + "paint_vae", + "paint_transformer", + "delight_vae", + "delight_transformer", + "delight_text_encoder", + ] + + class TestHunyuan3DPaintTurboSchedule(unittest.TestCase): def test_uses_standard_lcm_schedule_without_custom_timesteps(self): stage = Hunyuan3DPaintTexGenStage.__new__(Hunyuan3DPaintTexGenStage) diff --git a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py new file mode 100644 index 000000000000..33894a5c4cb6 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py @@ -0,0 +1,182 @@ +from unittest.mock import Mock, patch + +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_weight_inventory import ( + ComponentWeightEstimate, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.initial_residency import ( + GIB_BYTES, + choose_initial_resident_components, + maybe_seed_initial_residency, +) + + +def _weight(component_name: str, gib: int) -> ComponentWeightEstimate: + return ComponentWeightEstimate( + component_name=component_name, + component_model_path=f"/{component_name}", + checkpoint_bytes=gib * GIB_BYTES, + parameter_count=None, + target_element_size=None, + ) + + +class _Args: + model_path = "model" + backend = "sglang" + model_id = None + num_gpus = 1 + use_fsdp_inference = False + + def __init__(self, modes=None, explicit=()): + self.modes = modes or {} + self.explicit = set(explicit) + self.selected = {} + + def residency_mode(self, component_name): + return self.selected.get( + component_name, + self.modes.get(component_name, COMPONENT_OFFLOAD), + ) + + def explicit_residency_mode(self, component_name): + return ( + self.modes.get(component_name) if component_name in self.explicit else None + ) + + def set_auto_residency_mode(self, component_name, mode): + self.selected[component_name] = mode + + +def test_initial_seed_prefers_reused_dit_under_load_budget(): + args = _Args() + selected = choose_initial_resident_components( + args, + [_weight("transformer", 20), _weight("text_encoder", 20)], + available_bytes=30 * GIB_BYTES, + denoising_steps=8, + ) + + assert selected == {"transformer"} + + +def test_initial_seed_accounts_for_fixed_resident_weights(): + args = _Args(modes={"vae": RESIDENT}) + selected = choose_initial_resident_components( + args, + [_weight("vae", 8), _weight("transformer", 20)], + available_bytes=32 * GIB_BYTES, + denoising_steps=8, + ) + + assert selected == set() + + +def test_initial_seed_stays_conservative_for_unknown_fixed_resident_weights(): + args = _Args(modes={"custom_encoder": RESIDENT}) + unknown = ComponentWeightEstimate( + component_name="custom_encoder", + component_model_path="/custom_encoder", + checkpoint_bytes=None, + parameter_count=None, + target_element_size=None, + ) + + selected = choose_initial_resident_components( + args, + [unknown, _weight("transformer", 8)], + available_bytes=40 * GIB_BYTES, + denoising_steps=8, + ) + + assert selected == set() + + +def test_initial_seed_does_not_override_explicit_component_mode(): + args = _Args( + modes={"transformer": LAYERWISE_OFFLOAD}, + explicit={"transformer"}, + ) + selected = choose_initial_resident_components( + args, + [_weight("transformer", 8), _weight("vae", 2)], + available_bytes=40 * GIB_BYTES, + denoising_steps=8, + ) + + assert selected == {"vae"} + + +def test_initial_seed_applies_one_reversible_override(): + args = _Args() + inventory = [_weight("transformer", 8)] + with ( + patch( + "sglang.multimodal_gen.runtime.server_args.auto_tune." + "auto_residency_static_skip_reason", + return_value=None, + ), + patch( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.current_platform" + ) as platform, + ): + platform.is_cuda.return_value = True + platform.get_available_gpu_memory.return_value = 40 + maybe_seed_initial_residency(args, inventory) + + assert args.selected == {"transformer": RESIDENT} + platform.get_available_gpu_memory.assert_called_once_with( + distributed=False, + empty_cache=False, + ) + + +def test_initial_seed_honors_the_test_allocator_cap(): + args = _Args() + inventory = [_weight("transformer", 20)] + with ( + patch( + "sglang.multimodal_gen.runtime.server_args.auto_tune." + "auto_residency_static_skip_reason", + return_value=None, + ), + patch( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.current_platform" + ) as platform, + patch( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.envs.SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB", + 24, + ), + ): + platform.is_cuda.return_value = True + platform.get_available_gpu_memory.return_value = 140 + maybe_seed_initial_residency(args, inventory) + + assert args.selected == {} + + +def test_initial_seed_skips_structural_fsdp_path(): + args = _Args() + args.use_fsdp_inference = True + with ( + patch( + "sglang.multimodal_gen.runtime.server_args.auto_tune." + "auto_residency_static_skip_reason", + return_value=None, + ), + patch( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.current_platform", + Mock(is_cuda=Mock(return_value=True)), + ), + ): + maybe_seed_initial_residency(args, [_weight("transformer", 8)]) + + assert args.selected == {} 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..a8f0b00fe396 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -641,6 +641,23 @@ def test_auto_layerwise_skips_unsupported_component(monkeypatch): assert server_args.residency_mode("text_encoder") == COMPONENT_OFFLOAD +def test_initial_auto_resident_seed_skips_configured_layerwise_setup(): + module = _NestedEncoderDummyModel() + server_args = _server_args( + layerwise_offload_components=["text_encoder"], + _auto_residency_modes={"text_encoder": RESIDENT}, + ) + + configured = configure_layerwise_offload_modules( + {"text_encoder": module}, + server_args, + component_names=["text_encoder"], + ) + + assert configured == [] + assert not is_layerwise_offloaded_module(module) + + def test_canonical_selector_does_not_make_auto_layerwise_selection_strict( monkeypatch, ): 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..8bbb502ec5fc 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 @@ -7,6 +7,7 @@ """ import json +import os import tempfile import unittest from types import SimpleNamespace @@ -27,6 +28,48 @@ LTX25_DISTILLED_SIGMA_VALUES, LTX25PipelineConfig, ) +from sglang.multimodal_gen.runtime.pipelines.ltx_2_pipeline import ( + LTX2TwoStagePipeline, +) + + +class TestLTX2TwoStageComponentInventory(unittest.TestCase): + def test_resident_mode_declares_every_extra_weight_source(self): + with tempfile.TemporaryDirectory() as model_path: + for filename in ( + "ltx-2.3-spatial-upscaler-x2-1.0.safetensors", + "ltx-2.3-20b-distilled-lora-384.safetensors", + ): + open(os.path.join(model_path, filename), "wb").close() + + pipeline = object.__new__(LTX2TwoStagePipeline) + pipeline.model_path = model_path + server_args = SimpleNamespace( + component_paths={}, + ltx2_two_stage_device_mode="resident", + lora_path=None, + pipeline_config=SimpleNamespace( + dit_precision="bf16", + vae_precision="bf16", + text_encoder_precisions=("bf16",), + ), + ) + with mock.patch.object( + pipeline, + "_should_merge_stage2_distilled_lora", + return_value=True, + ): + sources = pipeline.additional_component_weight_sources(server_args) + + self.assertEqual( + [source.component_name for source in sources], + ["spatial_upsampler", "distilled_lora", "transformer_2"], + ) + self.assertEqual( + sources[-1].component_model_path, + os.path.join(model_path, "transformer"), + ) + self.assertTrue(sources[-1].supports_fsdp_loading) class TestLTX25DiTConfig(unittest.TestCase): 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..d69158b5b4cb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -15,6 +15,9 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( WarmupPhasePeak, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + HostPinBudget, +) 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 ( @@ -101,6 +104,7 @@ def test_worker_records_replica_load_and_runtime_peaks(): def test_server_warmup_preserves_peak_after_managed_stage_timeline(): worker = GPUWorker.__new__(GPUWorker) worker._auto_residency_warmup_records = [] + worker.pipeline = SimpleNamespace(stages=[]) residency_manager = Mock() residency_manager.take_warmup_phase_peaks.return_value = { "0:denoise:use:transformer": WarmupPhasePeak(("transformer",), 8) @@ -143,6 +147,7 @@ def test_server_warmup_preserves_peak_after_managed_stage_timeline(): def test_server_warmup_does_not_treat_allocator_cache_as_untracked_live_memory(): worker = GPUWorker.__new__(GPUWorker) worker._auto_residency_warmup_records = [] + worker.pipeline = SimpleNamespace(stages=[]) residency_manager = Mock() residency_manager.take_warmup_phase_peaks.return_value = { "0:denoise:use:transformer": WarmupPhasePeak(("transformer",), 8) @@ -178,9 +183,68 @@ def test_server_warmup_does_not_treat_allocator_cache_as_untracked_live_memory() assert "request:untracked" not in record.phase_peak_allocated_bytes +def test_server_warmup_normalizes_profile_stage_names_for_residency_timing(): + worker = GPUWorker.__new__(GPUWorker) + worker._auto_residency_warmup_records = [] + stage = Mock() + stage._active_profile_stage_name.return_value = "RefinerStage" + stage._component_stage_name.return_value = "refiner" + worker.pipeline = SimpleNamespace(stages=[stage]) + residency_manager = Mock() + residency_manager.take_warmup_phase_peaks.return_value = { + "0:refiner:use:transformer_2": WarmupPhasePeak( + ("transformer_2",), + 8, + used_components=("transformer_2",), + ) + } + residency_manager.current_device_components.return_value = ("transformer_2",) + device_module = Mock() + device_module.max_memory_allocated.return_value = 8 + req = SimpleNamespace( + metrics=SimpleNamespace( + total_duration_ms=12.0, + stages={"RefinerStage": 10.0}, + steps=[], + ) + ) + + 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, + workload=(64, 64, 1, 1), + baseline_allocated_bytes=3, + succeeded=True, + ) + + record = worker._auto_residency_warmup_records[0] + assert record.stage_duration_ms == {"refiner": 10.0} + _, stage_duration_ns, component_stages = ( + gpu_worker_module.estimate_default_workload_timing( + records=[record], + target_units=record.workload_units(), + target_num_inference_steps=1, + ) + ) + assert stage_duration_ns == {"refiner": 10_000_000} + assert component_stages == {"transformer_2": ("refiner",)} + + def test_loaded_prequantized_checkpoint_can_use_auto_residency(): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 + worker.is_output_rank = True + worker._auto_residency_budget_correction_bytes = 0 + worker._auto_residency_reference_request_duration_ns = None + worker._auto_residency_reference_stage_duration_ns = {} + worker._auto_residency_reference_component_stages = {} quantized_module = torch.nn.Module() quantized_module.register_parameter( "weight", @@ -190,6 +254,7 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): modules={"transformer": quantized_module}, component_residency_strategies={}, ) + host_pin_budget = HostPinBudget(available_bytes=0, reserve_bytes=0) worker.server_args = SimpleNamespace( num_gpus=1, nnodes=1, @@ -201,6 +266,7 @@ 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, + host_pin_budget=lambda: host_pin_budget, ) device_module = Mock() device_module.mem_get_info.return_value = (100, 100) @@ -233,10 +299,16 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 + worker.is_output_rank = True + worker._auto_residency_budget_correction_bytes = 2 * GIB_BYTES + worker._auto_residency_reference_request_duration_ns = None + worker._auto_residency_reference_stage_duration_ns = {} + worker._auto_residency_reference_component_stages = {} worker.pipeline = SimpleNamespace( modules={"transformer": torch.nn.Linear(1, 1)}, component_residency_strategies={}, ) + host_pin_budget = HostPinBudget(available_bytes=0, reserve_bytes=0) worker.server_args = SimpleNamespace( num_gpus=1, nnodes=1, @@ -248,6 +320,7 @@ 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, + host_pin_budget=lambda: host_pin_budget, ) device_module = Mock() device_module.mem_get_info.return_value = (100 * GIB_BYTES, 140 * GIB_BYTES) @@ -280,6 +353,7 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): ) assert report.budget_bytes == 80 * GIB_BYTES + assert report.planning_headroom_correction_bytes == 2 * GIB_BYTES assert report.device_transition_allocated_bytes == 60 * 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..7cf6b09c7060 100644 --- a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py @@ -63,15 +63,11 @@ def _exhaustive_plan_key( 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 + (len(option.preference_cost) for option in options), default=0 ) cost = tuple( sum( - ( - option.placement_cost_bytes[index] - if option.placement_cost_bytes - else 0 - ) + (option.preference_cost[index] if option.preference_cost else 0) for option in selected ) for index in range(cost_dimensions) @@ -121,6 +117,53 @@ def test_joint_plan_beats_independent_greedy_allocations(): assert plan.estimated_latency_savings == 109 +def test_one_selection_vector_obeys_load_runtime_and_hostpin_budgets(): + plan = optimize_placement( + [ + PlacementOption( + group_key="text_encoder", + option_key="text-encoder-resident", + resource_delta_bytes={"gpu:load": 8, "gpu:runtime": 2}, + estimated_latency_savings=60, + ), + PlacementOption( + group_key="text_encoder", + option_key="text-encoder-pinned", + resource_delta_bytes={"gpu:load": 1, "hostpin": 6}, + estimated_latency_savings=50, + ), + PlacementOption( + group_key="transformer", + option_key="transformer-resident", + resource_delta_bytes={"gpu:load": 2, "gpu:runtime": 8}, + estimated_latency_savings=100, + ), + PlacementOption( + group_key="transformer", + option_key="transformer-pinned", + resource_delta_bytes={"gpu:runtime": 1, "hostpin": 6}, + estimated_latency_savings=20, + ), + ], + resource_budget_bytes={ + "gpu:load": 9, + "gpu:runtime": 8, + "hostpin": 6, + }, + require_selection_from_every_group=True, + ) + + assert [option.option_key for option in plan.selections] == [ + "text-encoder-pinned", + "transformer-resident", + ] + assert plan.resource_delta_bytes == { + "gpu:load": 3, + "gpu:runtime": 8, + "hostpin": 6, + } + + def test_selects_at_most_one_option_per_group(): plan = optimize_placement( [ @@ -222,14 +265,14 @@ def test_latency_equivalent_plan_prefers_lower_soft_memory_cost(): option_key="transformer-resident", resource_delta_bytes={"gpu:rank0:runtime": 40}, estimated_latency_savings=1_000, - placement_cost_bytes=(40, 0), + preference_cost=(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), + preference_cost=(10, 0), ), ], resource_budget_bytes={"gpu:rank0:runtime": 100}, @@ -238,7 +281,7 @@ def test_latency_equivalent_plan_prefers_lower_soft_memory_cost(): 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) + assert plan.preference_cost == (40, 0) def test_latency_tolerance_is_global_not_per_option(): @@ -249,14 +292,14 @@ def test_latency_tolerance_is_global_not_per_option(): option_key="a-resident", resource_delta_bytes={}, estimated_latency_savings=60, - placement_cost_bytes=(1,), + preference_cost=(1,), ), PlacementOption( group_key="b", option_key="b-resident", resource_delta_bytes={}, estimated_latency_savings=60, - placement_cost_bytes=(1,), + preference_cost=(1,), ), ], resource_budget_bytes={}, @@ -531,21 +574,21 @@ def test_nonbinding_pin_prefixes_keep_only_globally_latency_eligible_states(): option_key="dit-pin-none", resource_delta_bytes={"gpu": 5, "hostpin": 0}, estimated_latency_savings=800, - placement_cost_bytes=(5, 0), + preference_cost=(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), + preference_cost=(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), + preference_cost=(5, 100), ), ], resource_budget_bytes={"gpu": 5, "hostpin": 100}, @@ -573,7 +616,7 @@ 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)), + preference_cost=(rng.randint(0, 10), rng.randint(0, 10)), ) ) budgets = {name: rng.randint(-2, 12) for name in resource_names} @@ -609,4 +652,4 @@ def test_branch_and_bound_matches_exhaustive_multiresource_search(): == expected_resources ) assert plan.estimated_latency_savings == expected_utility - assert plan.placement_cost_bytes == expected_cost + assert plan.preference_cost == expected_cost 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..ea6fcb263198 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -1419,7 +1419,7 @@ def test_cpu_offload_component_selector_keeps_unmatched_auto_defaults(self): ) self.assertEqual(args.residency_mode("vae"), COMPONENT_OFFLOAD) - self.assertEqual(args.residency_mode("transformer"), RESIDENT) + self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD) self.assertEqual(args.residency_mode("text_encoder"), LAYERWISE_OFFLOAD) def test_explicit_dit_layerwise_component_wins_over_auto_residency(self): @@ -1445,15 +1445,15 @@ def test_server_warmup_defers_high_memory_residency_to_calibration(self): 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): + def test_no_warmup_keeps_load_safe_static_planner_seed(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) + self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD) + self.assertEqual(args.residency_mode("vae"), LAYERWISE_OFFLOAD) def test_realtime_server_keeps_model_residency_fallback(self): args = self._from_dict_with_pipeline_config( @@ -1550,15 +1550,15 @@ def test_pipeline_configs_declare_auto_tune_hints(self): self.assertEqual(qwen_deployment.keep_resident_components, ("vae",)) self.assertIsNone(qwen_deployment.keep_resident_min_available_gb) - def test_longlive_residency_scales_with_available_memory(self): + def test_longlive_planner_candidates_do_not_use_a_coarse_vram_threshold(self): high_memory_args = self._from_dict_with_pipeline_config( LongLive2T2VConfig(), 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.assertNotIn("vae", high_memory_offload) + self.assertIn("text_encoder", high_memory_offload) + self.assertIn("vae", high_memory_offload) constrained_args = self._from_dict_with_pipeline_config( LongLive2T2VConfig(), @@ -1569,7 +1569,7 @@ def test_longlive_residency_scales_with_available_memory(self): self.assertIn("text_encoder", constrained_offload) self.assertIn("vae", constrained_offload) - def test_qwen_ar_generation_residency_scales_with_available_memory(self): + def test_qwen_ar_planner_candidates_do_not_use_a_coarse_vram_threshold(self): pipeline_configs = ( QwenImageLayeredPipelineConfig(), LongCatImagePipelineConfig(), @@ -1581,7 +1581,7 @@ def test_qwen_ar_generation_residency_scales_with_available_memory(self): memory_gb=80, kwargs={"performance_mode": "auto"}, ) - self.assertNotIn( + self.assertIn( "text_encoder", high_memory_args.layerwise_offload_components or [] ) self.assertFalse(high_memory_args.text_encoder_cpu_offload) @@ -1595,7 +1595,7 @@ def test_qwen_ar_generation_residency_scales_with_available_memory(self): "text_encoder", constrained_args.layerwise_offload_components or [] ) - def test_auto_multi_gpu_sana_wm_prefers_fsdp_and_cfg_parallel(self): + def test_auto_multi_gpu_sana_wm_defers_fsdp_and_enables_cfg_parallel(self): args = self._from_dict_with_pipeline_config( SanaWMPipelineConfig(), kwargs={ @@ -1605,7 +1605,22 @@ def test_auto_multi_gpu_sana_wm_prefers_fsdp_and_cfg_parallel(self): }, ) - self.assertTrue(args.use_fsdp_inference) + self.assertFalse(args.use_fsdp_inference) + self.assertTrue(args.enable_cfg_parallel) + + def test_server_warmup_keeps_sana_wm_dit_available_to_residency_planner(self): + args = self._from_dict_with_pipeline_config( + SanaWMPipelineConfig(), + kwargs={ + "model_path": "Efficient-Large-Model/SANA-WM_bidirectional", + "num_gpus": 2, + "performance_mode": "auto", + "warmup_mode": "server", + }, + ) + + self.assertFalse(args.use_fsdp_inference) + self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD) self.assertTrue(args.enable_cfg_parallel) def test_cache_dit_rejects_explicit_fsdp(self): @@ -1684,7 +1699,7 @@ def test_manual_mode_preserves_unset_performance_args(self): self.assertIsNone(args.image_encoder_cpu_offload) self.assertFalse(args.enable_cfg_parallel) - def test_default_auto_keeps_image_vae_resident_when_memory_allows(self): + def test_default_auto_exposes_all_implicit_image_placement_candidates(self): args = self._from_dict_with_pipeline_config( QwenImagePipelineConfig(), kwargs={"model_path": "Qwen/Qwen-Image"}, @@ -1692,12 +1707,10 @@ def test_default_auto_keeps_image_vae_resident_when_memory_allows(self): self.assertEqual(args.performance_mode, "auto") self.assertFalse(args.use_fsdp_inference) - # 80gb > image threshold (45gb): vae and dit stay resident, while the - # large encoders use layerwise offload. - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder"], + ["text_encoder", "image_encoder", "vae"], ) self.assertFalse(args.vae_cpu_offload) @@ -1716,7 +1729,7 @@ def test_auto_image_offloads_aux_below_resident_threshold(self): ["text_encoder", "image_encoder", "vae"], ) - def test_auto_zimage_keeps_dit_resident_on_5090(self): + def test_auto_zimage_defers_residency_until_weights_are_known(self): args = self._from_dict_with_pipeline_config( ZImagePipelineConfig(), memory_gb=32, @@ -1727,10 +1740,10 @@ def test_auto_zimage_keeps_dit_resident_on_5090(self): }, ) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder"], + ["text_encoder", "image_encoder", "vae"], ) def test_auto_lingbot_keeps_dit_resident_on_h100(self): @@ -1763,7 +1776,7 @@ def test_auto_image_preserves_explicit_dit_cpu_offload(self): self.assertTrue(args.dit_cpu_offload) - def test_auto_ltx_original_replaces_component_cpu_offload( + def test_auto_ltx_original_defers_component_residency( self, ): args = self._from_dict_with_pipeline_config( @@ -1777,7 +1790,7 @@ def test_auto_ltx_original_replaces_component_cpu_offload( ) self.assertEqual(args.ltx2_two_stage_device_mode, "original") - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertTrue(args.layerwise_offload_components) self.assertFalse(args.text_encoder_cpu_offload) self.assertFalse(args.image_encoder_cpu_offload) @@ -1786,7 +1799,7 @@ def test_auto_ltx_original_replaces_component_cpu_offload( ["text_encoder", "image_encoder", "vae"], ) - def test_auto_wan_keeps_single_dit_resident_on_h100(self): + def test_auto_wan_defers_single_dit_residency(self): args = self._from_dict_with_pipeline_config( WanT2V480PConfig(), kwargs={"performance_mode": "auto"}, @@ -1794,7 +1807,7 @@ def test_auto_wan_keeps_single_dit_resident_on_h100(self): self.assertTrue(args.layerwise_offload_components) self.assertFalse(args.use_fsdp_inference) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertFalse(args.text_encoder_cpu_offload) self.assertFalse(args.image_encoder_cpu_offload) self.assertEqual( @@ -1841,7 +1854,7 @@ def test_auto_wan2_2_a14b_layerwise_offload_adds_dit(self): ["dit", "text_encoder", "image_encoder", "vae"], ) - def test_auto_wan2_1_14b_keeps_dit_resident_on_h100(self): + def test_auto_wan2_1_14b_defers_dit_residency(self): for pipeline_config, model_path in ( (WanT2V720PConfig(), "Wan-AI/Wan2.1-T2V-14B-Diffusers"), (WanI2V480PConfig(), "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"), @@ -1857,7 +1870,7 @@ def test_auto_wan2_1_14b_keeps_dit_resident_on_h100(self): ) self.assertTrue(args.layerwise_offload_components) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertEqual(args.dit_offload_prefetch_size, 0.0) self.assertEqual( args.layerwise_offload_components, @@ -1908,6 +1921,23 @@ def test_auto_wan_layerwise_offload_does_not_disable_explicit_fsdp(self): ["text_encoder", "image_encoder", "vae"], ) self.assertTrue(args.use_fsdp_inference) + self.assertEqual(args.residency_mode("transformer"), RESIDENT) + self.assertTrue(args.should_use_fsdp_for_component("transformer")) + + def test_explicit_component_offload_still_scopes_explicit_fsdp(self): + args = self._from_dict_with_pipeline_config( + WanT2V480PConfig(), + kwargs={ + "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "num_gpus": 2, + "performance_mode": "auto", + "use_fsdp_inference": True, + "component_residency": {"transformer": COMPONENT_OFFLOAD}, + }, + ) + + self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD) + self.assertFalse(args.should_use_fsdp_for_component("transformer")) def test_auto_wan_layerwise_offload_preserves_explicit_dit_cpu_offload(self): args = self._from_dict_with_pipeline_config( @@ -1940,7 +1970,7 @@ def test_auto_mova_layerwise_offload_adds_dit_below_memory_threshold(self): ["dit", "text_encoder", "image_encoder", "vae"], ) - def test_auto_mova_keeps_dit_resident_at_memory_threshold(self): + def test_auto_mova_keeps_dit_layerwise_candidate_at_high_memory(self): args = self._from_dict_with_pipeline_config( MOVAPipelineConfig(), memory_gb=140, @@ -1950,13 +1980,13 @@ def test_auto_mova_keeps_dit_resident_at_memory_threshold(self): }, ) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder"], + ["dit", "text_encoder", "image_encoder", "vae"], ) - def test_auto_cosmos3_keeps_dit_resident_on_high_memory_gpu(self): + def test_auto_cosmos3_defers_high_memory_residency(self): args = self._from_dict_with_pipeline_config( Cosmos3Config(model_path="nvidia/Cosmos3-Nano"), available_memory_gb=139, @@ -1966,11 +1996,11 @@ def test_auto_cosmos3_keeps_dit_resident_on_high_memory_gpu(self): }, ) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertFalse(args.vae_cpu_offload) self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder"], + ["text_encoder", "image_encoder", "vae"], ) def test_auto_cosmos3_offloads_dit_below_resident_threshold(self): @@ -1986,9 +2016,7 @@ def test_auto_cosmos3_offloads_dit_below_resident_threshold(self): self.assertTrue(args.dit_cpu_offload) self.assertFalse(args.vae_cpu_offload) - def test_auto_cosmos3_super_keeps_dit_resident_on_high_memory_gpu(self): - # Super is a single-DiT pipeline like Nano, so above the threshold the - # component-offload round trip is pure per-request copy cost. + def test_auto_cosmos3_super_defers_high_memory_residency(self): args = self._from_dict_with_pipeline_config( Cosmos3Config(model_path="nvidia/Cosmos3-Super"), available_memory_gb=139, @@ -1998,7 +2026,7 @@ def test_auto_cosmos3_super_keeps_dit_resident_on_high_memory_gpu(self): }, ) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertFalse(args.vae_cpu_offload) def test_auto_cosmos3_super_offloads_dit_below_resident_threshold(self): @@ -2028,7 +2056,7 @@ def test_memory_sana_wm_layerwise_offload_adds_dit(self): ["dit", "text_encoder", "image_encoder", "vae"], ) - def test_auto_fastwan_keeps_dit_resident_on_h100(self): + def test_auto_fastwan_defers_dit_residency(self): args = self._from_dict_with_pipeline_config( FastWan2_2_TI2V_5B_Config(), available_memory_gb=72, @@ -2038,7 +2066,7 @@ def test_auto_fastwan_keeps_dit_resident_on_h100(self): }, ) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertEqual( args.layerwise_offload_components, ["text_encoder", "image_encoder", "vae"], @@ -2056,7 +2084,7 @@ def test_auto_fastwan_offloads_dit_below_resident_threshold(self): self.assertTrue(args.dit_cpu_offload) - def test_auto_fast_hunyuan_keeps_dit_resident_on_h100(self): + def test_auto_fast_hunyuan_defers_residency(self): args = self._from_dict_with_pipeline_config( FastHunyuanConfig(), available_memory_gb=72, @@ -2066,7 +2094,7 @@ def test_auto_fast_hunyuan_keeps_dit_resident_on_h100(self): }, ) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertFalse(args.vae_cpu_offload) def test_auto_fast_hunyuan_offloads_dit_below_resident_threshold(self): @@ -2081,7 +2109,7 @@ def test_auto_fast_hunyuan_offloads_dit_below_resident_threshold(self): self.assertTrue(args.dit_cpu_offload) - def test_auto_turbo_wan_keeps_dit_resident_on_h100(self): + def test_auto_turbo_wan_defers_dit_residency(self): args = self._from_dict_with_pipeline_config( TurboWanT2V480PConfig(), kwargs={ @@ -2090,7 +2118,7 @@ def test_auto_turbo_wan_keeps_dit_resident_on_h100(self): }, ) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertEqual( args.layerwise_offload_components, ["text_encoder", "image_encoder", "vae"], @@ -2108,7 +2136,7 @@ def test_explicit_fastwan_dit_layerwise_still_selects_dit_group(self): self.assertTrue(args.dit_cpu_offload) self.assertEqual(args.layerwise_offload_components, ["dit"]) - def test_auto_multi_gpu_wan_uses_layerwise_offload_without_cfg(self): + def test_auto_multi_gpu_wan_defers_residency_without_cfg(self): with patch.object(ServerArgs, "_model_default_uses_cfg", return_value=False): args = self._from_dict_with_pipeline_config( WanT2V480PConfig(), @@ -2121,7 +2149,7 @@ def test_auto_multi_gpu_wan_uses_layerwise_offload_without_cfg(self): self.assertFalse(args.use_fsdp_inference) self.assertFalse(args.enable_cfg_parallel) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertTrue(args.layerwise_offload_components) self.assertFalse(args.text_encoder_cpu_offload) self.assertFalse(args.image_encoder_cpu_offload) @@ -2147,7 +2175,7 @@ def test_explicit_multi_gpu_dit_layerwise_only_selects_dit_group(self): self.assertTrue(args.image_encoder_cpu_offload) self.assertEqual(args.layerwise_offload_components, ["dit"]) - def test_auto_multi_gpu_ltx_replaces_component_cpu_offload_with_resident_dit(self): + def test_auto_multi_gpu_ltx_defers_component_residency(self): args = self._from_dict_with_pipeline_config( LTX2PipelineConfig(), available_memory_gb=76, @@ -2159,7 +2187,7 @@ def test_auto_multi_gpu_ltx_replaces_component_cpu_offload_with_resident_dit(sel ) self.assertFalse(args.use_fsdp_inference) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertTrue(args.layerwise_offload_components) self.assertFalse(args.text_encoder_cpu_offload) self.assertFalse(args.image_encoder_cpu_offload) @@ -2168,7 +2196,7 @@ def test_auto_multi_gpu_ltx_replaces_component_cpu_offload_with_resident_dit(sel ["text_encoder", "image_encoder", "vae"], ) - def test_auto_high_memory_ltx23_resident_keeps_aux_components_resident(self): + def test_auto_high_memory_ltx23_resident_defers_aux_placement(self): args = self._from_dict_with_pipeline_config( LTX2PipelineConfig(), memory_gb=140, @@ -2183,10 +2211,15 @@ def test_auto_high_memory_ltx23_resident_keeps_aux_components_resident(self): self.assertEqual(args.ltx2_two_stage_device_mode, "resident") self.assertFalse(args.use_fsdp_inference) self.assertFalse(args.dit_cpu_offload) + self.assertEqual(args.residency_mode("transformer"), RESIDENT) + self.assertEqual(args.residency_mode("transformer_2"), RESIDENT) self.assertFalse(args.text_encoder_cpu_offload) self.assertFalse(args.image_encoder_cpu_offload) self.assertFalse(args.vae_cpu_offload) - self.assertIsNone(args.layerwise_offload_components) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder", "vae"], + ) def test_auto_high_memory_ltx23_original_keeps_default_layerwise_components(self): args = self._from_dict_with_pipeline_config( @@ -2236,7 +2269,7 @@ def test_explicit_layerwise_components_preserved_in_ltx23_resident(self): self.assertEqual(args.ltx2_two_stage_device_mode, "resident") self.assertEqual(args.layerwise_offload_components, ["text_encoder"]) - def test_auto_multi_gpu_qwen_keeps_vae_resident_with_cfg(self): + def test_auto_multi_gpu_qwen_defers_residency_with_cfg(self): args = self._from_dict_with_pipeline_config( QwenImagePipelineConfig(), kwargs={ @@ -2248,16 +2281,14 @@ def test_auto_multi_gpu_qwen_keeps_vae_resident_with_cfg(self): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) - # 80gb > image threshold (45gb): vae and dit stay resident, while the - # large encoders use layerwise offload. - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder"], + ["text_encoder", "image_encoder", "vae"], ) self.assertFalse(args.vae_cpu_offload) - def test_auto_multi_gpu_zimage_base_prefers_fsdp(self): + def test_auto_multi_gpu_zimage_defers_fsdp_to_weight_inventory(self): args = self._from_dict_with_pipeline_config( ZImagePipelineConfig(), kwargs={ @@ -2267,7 +2298,22 @@ def test_auto_multi_gpu_zimage_base_prefers_fsdp(self): }, ) - self.assertTrue(args.use_fsdp_inference) + self.assertFalse(args.use_fsdp_inference) + self.assertTrue(args.enable_cfg_parallel) + + def test_server_warmup_keeps_zimage_dit_available_to_residency_planner(self): + args = self._from_dict_with_pipeline_config( + ZImagePipelineConfig(), + kwargs={ + "model_path": "Tongyi-MAI/Z-Image", + "num_gpus": 2, + "performance_mode": "auto", + "warmup_mode": "server", + }, + ) + + self.assertFalse(args.use_fsdp_inference) + self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD) self.assertTrue(args.enable_cfg_parallel) def test_auto_multi_gpu_zimage_turbo_skips_fsdp(self): @@ -2296,15 +2342,13 @@ def test_auto_multi_gpu_qwen_preserves_explicit_fsdp_false(self): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) - # Explicit FSDP selection must not freeze unrelated, implicit DiT - # residency decisions on a high-memory GPU. - self.assertFalse(args.dit_cpu_offload) + # Explicit FSDP selection must not freeze unrelated, implicit + # residency decisions before component weights are known. + self.assertTrue(args.dit_cpu_offload) self.assertFalse(args.vae_cpu_offload) - # The layerwise filter still drops VAE (kept resident); encoders stay - # offloaded. self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder"], + ["text_encoder", "image_encoder", "vae"], ) def test_auto_multi_gpu_qwen_skips_fsdp_when_available_memory_is_low(self): @@ -2320,12 +2364,10 @@ def test_auto_multi_gpu_qwen_skips_fsdp_when_available_memory_is_low(self): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) - # 50gb still > image threshold (45gb): vae and dit stay resident, while - # the encoders remain offloaded; qwen does not opt into auto fsdp. - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder"], + ["text_encoder", "image_encoder", "vae"], ) self.assertFalse(args.vae_cpu_offload) @@ -2344,7 +2386,7 @@ def test_auto_multi_gpu_qwen_uses_selected_gpu_min_available_memory(self): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) - def test_auto_multi_gpu_qwen_keeps_vae_resident_with_headroom(self): + def test_auto_multi_gpu_qwen_defers_residency_with_headroom(self): args = self._from_dict_with_pipeline_config( QwenImagePipelineConfig(), available_memory_gb={1: 72, 2: 80}, @@ -2358,16 +2400,14 @@ def test_auto_multi_gpu_qwen_keeps_vae_resident_with_headroom(self): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) - # min available across selected gpus is 72gb > image threshold (45gb): - # vae and dit stay resident, while the encoders remain offloaded. - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder"], + ["text_encoder", "image_encoder", "vae"], ) self.assertFalse(args.vae_cpu_offload) - def test_auto_minimax_h3_keeps_large_components_resident_with_headroom(self): + def test_auto_minimax_h3_defers_large_component_residency(self): args = self._from_dict_with_pipeline_config( MiniMaxH3PipelineConfig(), memory_gb=141, @@ -2380,11 +2420,11 @@ def test_auto_minimax_h3_keeps_large_components_resident_with_headroom(self): }, ) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertFalse(args.text_encoder_cpu_offload) self.assertFalse(args.vae_cpu_offload) - self.assertNotIn("text_encoder", args.layerwise_offload_components or []) - self.assertNotIn("vae", args.layerwise_offload_components or []) + self.assertIn("text_encoder", args.layerwise_offload_components or []) + self.assertIn("vae", args.layerwise_offload_components or []) def test_auto_minimax_h3_keeps_memory_policy_below_residency_threshold(self): args = self._from_dict_with_pipeline_config( diff --git a/python/sglang/multimodal_gen/test/unit/test_server_case_requests.py b/python/sglang/multimodal_gen/test/unit/test_server_case_requests.py new file mode 100644 index 000000000000..6be45088d351 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_server_case_requests.py @@ -0,0 +1,72 @@ +from types import SimpleNamespace + +from sglang.multimodal_gen.test.server import test_server_utils +from sglang.multimodal_gen.test.server.testcase_configs import ( + DiffusionSamplingParams, +) + + +class _FakeVideos: + def __init__(self): + self.create_kwargs = None + + def create(self, **kwargs): + self.create_kwargs = kwargs + return SimpleNamespace(id="request-id") + + def list(self): + return SimpleNamespace( + data=[SimpleNamespace(id="request-id", status="completed")] + ) + + def download_content(self, *, video_id): + assert video_id == "request-id" + return SimpleNamespace(read=lambda: b"video") + + +def test_direct_url_video_request_preserves_workload_extras(monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(test_server_utils, "validate_openai_video", lambda _data: None) + monkeypatch.setattr( + test_server_utils, + "validate_video_file", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + test_server_utils, + "upload_file_to_slack", + lambda **_kwargs: None, + ) + videos = _FakeVideos() + client = SimpleNamespace(videos=videos) + sampling_params = DiffusionSamplingParams( + prompt="prompt", + image_path="https://example.com/input.png", + direct_url_test=True, + output_size="384x640", + num_frames=17, + fps=16, + extras={ + "num_inference_steps": 12, + "guidance_scale": 4.5, + "seed": 0, + }, + ) + + generate = test_server_utils.get_generate_fn( + model_path="test/model", + modality="video", + sampling_params=sampling_params, + ) + request_id, content = generate("case", client) + + assert request_id == "request-id" + assert content == b"video" + assert videos.create_kwargs["extra_body"] == { + "reference_url": "https://example.com/input.png", + "fps": 16, + "num_frames": 17, + "num_inference_steps": 12, + "guidance_scale": 4.5, + "seed": 0, + } From df6cfa2ce3fb03535c92b68c51919ae22cf6c103 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 22:46:20 +0800 Subject: [PATCH 02/70] [diffusion] resolve NVFP4 weights before residency planning --- .../runtime/pipelines/flux_2_nvfp4.py | 8 +++++++ .../test/unit/test_transformer_quant.py | 24 +++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/flux_2_nvfp4.py b/python/sglang/multimodal_gen/runtime/pipelines/flux_2_nvfp4.py index 9a7df8426d22..d1c60dbd79a1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/flux_2_nvfp4.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/flux_2_nvfp4.py @@ -5,6 +5,9 @@ from functools import lru_cache from typing import Any, cast +from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( + resolve_transformer_safetensors_to_load, +) from sglang.multimodal_gen.runtime.pipelines.flux_2 import Flux2Pipeline from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( @@ -40,6 +43,11 @@ def _resolve_nvfp4_transformer_weights_path( server_args: ServerArgs, model_path: str ) -> str: if server_args.transformer_weights_path is not None: + override_files = resolve_transformer_safetensors_to_load( + server_args, model_path + ) + if len(override_files) == 1: + return override_files[0] return server_args.transformer_weights_path local_nvfp4_path = maybe_download_model(model_path) 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 b4185e5bb0e9..4fde2ccb3901 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -107,6 +107,9 @@ def _loads(input_str, _flags=None): from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan from sglang.multimodal_gen.runtime.models.dits.flux import FluxSingleTransformerBlock from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import MiniMaxH3DiTModel +from sglang.multimodal_gen.runtime.pipelines.flux_2_nvfp4 import ( + _resolve_nvfp4_transformer_weights_path, +) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability from sglang.multimodal_gen.runtime.utils.quantization_utils import ( @@ -792,6 +795,27 @@ def test_resolve_transformer_safetensors_to_load_prefers_mixed_export( self.assertEqual(resolved, [mixed]) + @patch( + "sglang.multimodal_gen.runtime.pipelines.flux_2_nvfp4.resolve_transformer_safetensors_to_load" + ) + def test_flux2_nvfp4_resolves_override_before_preload_inventory( + self, mock_resolve_override + ): + with tempfile.TemporaryDirectory() as cached_dir: + mixed = os.path.join(cached_dir, "flux2-dev-nvfp4-mixed.safetensors") + open(mixed, "a").close() + mock_resolve_override.return_value = [mixed] + repo_id = "black-forest-labs/FLUX.2-dev-NVFP4" + server_args = SimpleNamespace(transformer_weights_path=repo_id) + + resolved = _resolve_nvfp4_transformer_weights_path( + server_args, + "/unused/model/path", + ) + + self.assertEqual(resolved, mixed) + mock_resolve_override.assert_called_once_with(server_args, "/unused/model/path") + @patch( "sglang.multimodal_gen.runtime.loader.transformer_load_utils.snapshot_download", ) From 034b80462376b9fd7825be6aebd9d85b7f7e7df4 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 23:41:05 +0800 Subject: [PATCH 03/70] [diffusion] preserve load-safe residency admission --- .../component_weight_inventory.py | 39 ++++++++------- .../memory_managers/initial_residency.py | 17 +++++-- .../pipelines_core/composed_pipeline_base.py | 23 +++++++++ .../runtime/server_args/auto_tune.py | 28 +++++++---- .../test/unit/test_component_loading_order.py | 30 ++++++++++++ .../test/unit/test_initial_residency.py | 47 ++++++++++++++++--- 6 files changed, 150 insertions(+), 34 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_weight_inventory.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_weight_inventory.py index 555d4072f158..371a49001b7f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_weight_inventory.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_weight_inventory.py @@ -88,6 +88,28 @@ def _safetensors_payload_stats(file_path: str) -> tuple[int, int | None] | None: return checkpoint_bytes, parameter_count +def infer_safetensors_weight_stats( + file_paths: list[str], +) -> tuple[int | None, int | None]: + """Infer aggregate stats for the exact safetensors files selected to load.""" + stats = [ + stat + for file_path in file_paths + if (stat := _safetensors_payload_stats(file_path)) is not None + ] + if not stats: + return None, None + parameter_counts = [parameters for _, parameters in stats] + return ( + sum(checkpoint_bytes for checkpoint_bytes, _ in stats), + ( + sum(parameter_counts) + if all(parameters is not None for parameters in parameter_counts) + else None + ), + ) + + def infer_safetensors_weight_stats_by_prefix( file_path: str, ) -> dict[str, tuple[int, int]] | None: @@ -148,22 +170,7 @@ def infer_component_weight_stats( """Infer checkpoint bytes and parameter count without loading tensors.""" safetensors_files = _list_component_safetensors_files(component_model_path) if safetensors_files: - stats = [ - stat - for file_path in safetensors_files - if (stat := _safetensors_payload_stats(file_path)) is not None - ] - if not stats: - return None, None - parameter_counts = [parameters for _, parameters in stats] - return ( - sum(checkpoint_bytes for checkpoint_bytes, _ in stats), - ( - sum(parameter_counts) - if all(parameters is not None for parameters in parameter_counts) - else None - ), - ) + return infer_safetensors_weight_stats(safetensors_files) if os.path.isfile(component_model_path): if component_model_path.endswith(_WEIGHT_FILE_SUFFIXES): diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py index ef27ac2eaaf1..73811867c778 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py @@ -20,6 +20,10 @@ optimize_placement, ) from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.server_args.auto_tune import ( + auto_residency_static_skip_reason, + resolve_keep_resident_min_available_gb, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger if TYPE_CHECKING: @@ -114,10 +118,6 @@ def maybe_seed_initial_residency( inventory: list[ComponentWeightEstimate], ) -> None: """Apply a reversible resident seed for warmup-calibrated auto mode.""" - from sglang.multimodal_gen.runtime.server_args.auto_tune import ( - auto_residency_static_skip_reason, - ) - if ( auto_residency_static_skip_reason(server_args) is not None or server_args.use_fsdp_inference @@ -134,6 +134,15 @@ def maybe_seed_initial_residency( available_gib, envs.SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB, ) + admission_gib = resolve_keep_resident_min_available_gb(server_args) + if available_gib < admission_gib: + logger.debug( + "Keeping configured load placement: %.1f GiB is below the %.1f GiB " + "initial-residency admission threshold.", + available_gib, + admission_gib, + ) + return selected = choose_initial_resident_components( server_args, inventory, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index f69f837c0ea6..a1a7c27bb2d5 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -7,6 +7,7 @@ This module defines the base class for pipelines that are composed of multiple stages. """ +import copy import os from abc import ABC, abstractmethod from typing import Any, Callable, Iterator, Literal, cast @@ -21,6 +22,10 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PipelineComponentLoader, ) +from sglang.multimodal_gen.runtime.loader.gguf_weights import names_gguf_checkpoint +from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( + resolve_transformer_safetensors_to_load, +) from sglang.multimodal_gen.runtime.managers.memory_managers.component_loading_order import ( ComponentLoadSpec, order_component_load_specs, @@ -34,6 +39,7 @@ ComponentWeightEstimate, ComponentWeightSource, estimate_component_weight_inventory, + infer_safetensors_weight_stats, ) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( is_image_encoder_component_name, @@ -428,9 +434,26 @@ def _component_weight_source( and server_args.transformer_weights_path is not None ): weights_path = server_args.transformer_weights_path + checkpoint_bytes = None + parameter_count = None + if ( + weights_path is not None + and is_legacy_dit_offload_component_name(spec.load_module_name) + and not names_gguf_checkpoint(weights_path) + ): + component_server_args = copy.copy(server_args) + component_server_args.transformer_weights_path = weights_path + selected_files = resolve_transformer_safetensors_to_load( + component_server_args, spec.component_model_path + ) + checkpoint_bytes, parameter_count = infer_safetensors_weight_stats( + selected_files + ) return ComponentWeightSource( component_name=spec.module_name, component_model_path=weights_path or spec.component_model_path, + checkpoint_bytes=checkpoint_bytes, + parameter_count=parameter_count, target_element_size=ComposedPipelineBase._component_target_element_size( spec.load_module_name, server_args ), 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 8c7d11be5c47..3ab17e928cbc 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -48,6 +48,22 @@ DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB = 120.0 +def resolve_keep_resident_min_available_gb( + server_args: ServerArgs, + deployment_config: ModelDeploymentConfig | None = None, +) -> float: + """Resolve the established device-memory admission threshold for residency.""" + deployment_config = ( + deployment_config or server_args.pipeline_config.get_model_deployment_config() + ) + explicit = deployment_config.keep_resident_min_available_gb + if explicit is not None: + return explicit + if server_args.pipeline_config.task_type.is_image_gen(): + return IMAGE_GEN_KEEP_RESIDENT_MIN_AVAILABLE_GB + return DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB + + def auto_residency_static_skip_reason(server_args: ServerArgs) -> str | None: """Return why args cannot use pre-warmup static residency planning.""" if envs.SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY: @@ -120,14 +136,10 @@ def _uses_auto_residency_planner(self) -> bool: def _resolve_keep_resident_min_available_gb( self, deployment_config: ModelDeploymentConfig - ) -> float | None: - # explicit per-model > task-type default > global default - explicit = deployment_config.keep_resident_min_available_gb - if explicit is not None: - return explicit - if self.server_args.pipeline_config.task_type.is_image_gen(): - return IMAGE_GEN_KEEP_RESIDENT_MIN_AVAILABLE_GB - return DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB + ) -> float: + return resolve_keep_resident_min_available_gb( + self.server_args, deployment_config + ) def adjust_based_on_performance_mode(self) -> None: """Adjust the server args based on the performance mode""" diff --git a/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py b/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py index 30083a8237c3..0220aa93ebb6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py @@ -1,5 +1,6 @@ import json from types import SimpleNamespace +from unittest.mock import patch from sglang.multimodal_gen.runtime.managers.memory_managers.component_loading_order import ( ComponentLoadSpec, @@ -259,6 +260,35 @@ def test_preload_inventory_uses_the_actual_weight_override(): assert secondary_source.supports_fsdp_loading +def test_preload_inventory_uses_selected_transformer_safetensors(tmp_path): + mixed_path = tmp_path / "flux2-dev-nvfp4-mixed.safetensors" + _write_safetensors(mixed_path, 64) + server_args = SimpleNamespace( + component_weights_paths={}, + transformer_weights_path="owner/flux2-nvfp4", + revision=None, + pipeline_config=SimpleNamespace( + dit_precision="bf16", + vae_precision="fp32", + image_encoder_precision="fp16", + text_encoder_precisions=("bf16",), + ), + ) + + with patch( + "sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base." + "resolve_transformer_safetensors_to_load", + return_value=[str(mixed_path)], + ): + source = ComposedPipelineBase._component_weight_source( + _spec("transformer", 0, "/model/transformer"), server_args + ) + + assert source.component_model_path == "owner/flux2-nvfp4" + assert source.checkpoint_bytes == 64 + assert source.parameter_count == 32 + + def test_preload_inventory_resolves_group_precision_fallbacks(): server_args = SimpleNamespace( pipeline_config=SimpleNamespace( diff --git a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py index 33894a5c4cb6..463b49b681b4 100644 --- a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py @@ -116,10 +116,15 @@ def test_initial_seed_applies_one_reversible_override(): inventory = [_weight("transformer", 8)] with ( patch( - "sglang.multimodal_gen.runtime.server_args.auto_tune." - "auto_residency_static_skip_reason", + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.auto_residency_static_skip_reason", return_value=None, ), + patch( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.resolve_keep_resident_min_available_gb", + return_value=30, + ), patch( "sglang.multimodal_gen.runtime.managers.memory_managers." "initial_residency.current_platform" @@ -141,10 +146,15 @@ def test_initial_seed_honors_the_test_allocator_cap(): inventory = [_weight("transformer", 20)] with ( patch( - "sglang.multimodal_gen.runtime.server_args.auto_tune." - "auto_residency_static_skip_reason", + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.auto_residency_static_skip_reason", return_value=None, ), + patch( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.resolve_keep_resident_min_available_gb", + return_value=30, + ), patch( "sglang.multimodal_gen.runtime.managers.memory_managers." "initial_residency.current_platform" @@ -167,8 +177,8 @@ def test_initial_seed_skips_structural_fsdp_path(): args.use_fsdp_inference = True with ( patch( - "sglang.multimodal_gen.runtime.server_args.auto_tune." - "auto_residency_static_skip_reason", + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.auto_residency_static_skip_reason", return_value=None, ), patch( @@ -180,3 +190,28 @@ def test_initial_seed_skips_structural_fsdp_path(): maybe_seed_initial_residency(args, [_weight("transformer", 8)]) assert args.selected == {} + + +def test_initial_seed_keeps_configured_load_path_below_admission_threshold(): + args = _Args() + with ( + patch( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.auto_residency_static_skip_reason", + return_value=None, + ), + patch( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.resolve_keep_resident_min_available_gb", + return_value=30, + ), + patch( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.current_platform" + ) as platform, + ): + platform.is_cuda.return_value = True + platform.get_available_gpu_memory.return_value = 29.8 + maybe_seed_initial_residency(args, [_weight("transformer", 8)]) + + assert args.selected == {} From eff6b09a58c0bdd4abdab1d9c684672529fee060 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 27 Aug 2026 23:58:46 +0800 Subject: [PATCH 04/70] [diffusion] plan HostPin residency per node --- .../runtime/managers/gpu_worker.py | 1 + .../memory_managers/auto_residency.py | 95 +++++++--- .../memory_managers/host_memory_budget.py | 22 +++ .../memory_managers/initial_residency.py | 16 +- .../test/unit/test_auto_residency.py | 168 ++++++++++++++++-- .../test/unit/test_host_memory_budget.py | 16 ++ .../test/unit/test_initial_residency.py | 20 ++- 7 files changed, 297 insertions(+), 41 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 829bfa2aeff0..cd4eb0177f07 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1377,6 +1377,7 @@ def apply_auto_residency(self, *, pre_warmup: bool = False) -> OutputBatch: plan=plan, modules=self._auto_residency_modules(), server_args=self.server_args, + rank=self.rank, ) if not newly_applied: raise RuntimeError("placement plan applied no residency changes") 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 8669e286baaf..c70370176a35 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 @@ -276,6 +276,7 @@ class AutoResidencyPlan(msgspec.Struct, frozen=True): budget_bytes: int = 0 resource_budget_bytes: dict[str, int] = {} resource_delta_bytes: dict[str, int] = {} + host_pin_target_bytes_by_rank: dict[int, int] = {} changes: list[ResidencyTarget] = [] recovering_from_oom: bool = False skip_reason: str | None = None @@ -290,6 +291,7 @@ class AppliedResidencyChange(msgspec.Struct, frozen=True): previous_layerwise_offload_enabled: bool = True previous_layerwise_configured: bool = True previous_auto_residency_mode: str | None = None + previous_host_pin_budget_state: tuple[int, int] | None = None pinned_host_changed: bool = False applied_device_delta_bytes: int = 0 @@ -881,7 +883,7 @@ def layerwise_pinned_host_bytes( def layerwise_host_pin_capacity_bytes( modules: Mapping[str, object], *, pin_budget=None ) -> int: - """This process's non-overlapping HostPin allowance.""" + """This process's non-overlapping share of the HostPin planner budget.""" seen_budgets: set[int] = set() total = 0 for module in modules.values(): @@ -893,9 +895,9 @@ def layerwise_host_pin_capacity_bytes( if budget_id in seen_budgets: continue seen_budgets.add(budget_id) - total += max(0, budget.available_bytes - budget.reserve_bytes) + total += budget.planning_capacity_bytes if pin_budget is not None and id(pin_budget) not in seen_budgets: - total += max(0, pin_budget.available_bytes - pin_budget.reserve_bytes) + total += pin_budget.planning_capacity_bytes return total @@ -1139,9 +1141,9 @@ def _unconfigured_layerwise_targets( ) -> 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. + Estimated layer stores expose the complete resident-layer and HostPin + prefix frontier before managers exist. The selected state configures the + real managers lazily, and the single validation warmup checks that estimate. """ managers = _estimate_unconfigured_layerwise_managers( module=module, @@ -2109,21 +2111,22 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP for candidate in candidates ) if has_host_resources: + reports_by_node: dict[int, list[RankResidencyReport]] = {} for report in reports: - prefix = f"node{report.node_rank}:rank{report.rank}" + reports_by_node.setdefault(report.node_rank, []).append(report) + for node_rank, node_reports in reports_by_node.items(): + prefix = f"node{node_rank}" resource_budgets[f"hostpin:{prefix}"] = max( 0, - report.host_pin_capacity_bytes - report.pinned_host_bytes, + sum(report.host_pin_capacity_bytes for report in node_reports) + - sum(report.pinned_host_bytes for report in node_reports), ) - resource_budgets[f"hostram:{prefix}:unpin"] = ( - report.host_transition_headroom_bytes - ) - resource_budgets[f"hostram:{prefix}:pin"] = ( - report.host_transition_headroom_bytes - ) - resource_budgets[f"hostram:{prefix}:materialize"] = ( - report.host_transition_headroom_bytes + transition_headroom = sum( + report.host_transition_headroom_bytes for report in node_reports ) + resource_budgets[f"hostram:{prefix}:unpin"] = transition_headroom + resource_budgets[f"hostram:{prefix}:pin"] = transition_headroom + resource_budgets[f"hostram:{prefix}:materialize"] = transition_headroom candidate_by_key = {candidate.option_key(): candidate for candidate in candidates} report_candidates = [ @@ -2234,20 +2237,26 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP rank_candidate.device_transition_delta_bytes ) if has_host_resources: - prefix = f"node{report.node_rank}:rank{report.rank}" + prefix = f"node{report.node_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:{prefix}:unpin"] = ( - rank_candidate.host_unpin_scratch_bytes + unpin_resource = f"hostram:{prefix}:unpin" + resource_deltas[unpin_resource] = ( + resource_deltas.get(unpin_resource, 0) + + rank_candidate.host_unpin_scratch_bytes ) - resource_deltas[f"hostram:{prefix}:pin"] = ( - rank_candidate.host_pin_scratch_bytes + pin_resource = f"hostram:{prefix}:pin" + resource_deltas[pin_resource] = ( + resource_deltas.get(pin_resource, 0) + + rank_candidate.host_pin_scratch_bytes ) - resource_deltas[f"hostram:{prefix}:materialize"] = ( - rank_candidate.host_materialize_scratch_bytes + materialize_resource = f"hostram:{prefix}:materialize" + resource_deltas[materialize_resource] = ( + resource_deltas.get(materialize_resource, 0) + + rank_candidate.host_materialize_scratch_bytes ) estimated_latency_savings = ( min( @@ -2333,6 +2342,17 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP if not candidate.current_placement: changed_candidates.append(candidate) changes = rank_candidates_by_h2d_savings(changed_candidates) + host_pin_targets = {} + if has_host_resources: + for report, rank_candidates in zip(reports, report_candidates): + host_pin_targets[report.rank] = max( + 0, + report.pinned_host_bytes + + sum( + rank_candidates[selection.option_key].pinned_host_delta_bytes + for selection in placement.selections + ), + ) return AutoResidencyPlan( estimated_peak_bytes=estimated_peak, @@ -2340,6 +2360,7 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP budget_bytes=budget, resource_budget_bytes=resource_budgets, resource_delta_bytes=placement.resource_delta_bytes, + host_pin_target_bytes_by_rank=host_pin_targets, changes=changes, recovering_from_oom=recovering_from_oom, current_placement_reserve_shortfall_bytes=(current_placement_reserve_shortfall), @@ -2351,6 +2372,7 @@ def apply_residency_changes( plan: AutoResidencyPlan, modules: Mapping[str, object], server_args: ServerArgs, + rank: int | None = None, ) -> list[AppliedResidencyChange]: """Apply complete target states transactionally on this rank. @@ -2367,6 +2389,7 @@ def apply_residency_changes( must abort). """ applied: list[AppliedResidencyChange] = [] + previous_host_pin_budget_state = None try: ordered_changes = sorted( plan.changes, @@ -2383,6 +2406,12 @@ def apply_residency_changes( target_modules: dict[str, nn.Module] = {} snapshots: dict[str, AppliedResidencyChange] = {} previous_pins_by_component: dict[str, tuple[tuple[int, ...], ...]] = {} + if rank is not None and rank in plan.host_pin_target_bytes_by_rank: + previous_host_pin_budget_state = ( + server_args.host_pin_budget().set_spendable_capacity( + plan.host_pin_target_bytes_by_rank[rank] + ) + ) for candidate in ordered_changes: if candidate.component_name in target_modules: raise RuntimeError( @@ -2454,6 +2483,7 @@ def apply_residency_changes( previous_layerwise_offload_enabled=enabled.pop(), previous_layerwise_configured=previous_configured, previous_auto_residency_mode=previous_auto_mode, + previous_host_pin_budget_state=previous_host_pin_budget_state, pinned_host_changed=( previous_configured and candidate.target_layerwise_pinned_layers != previous_pinned @@ -2471,6 +2501,7 @@ def apply_residency_changes( component_name=candidate.component_name, residency_mode=candidate.residency_mode, previous_auto_residency_mode=previous_auto_mode, + previous_host_pin_budget_state=previous_host_pin_budget_state, applied_device_delta_bytes=applied_device_delta_bytes, ) pinning_changes = [ @@ -2594,6 +2625,10 @@ def apply_residency_changes( rollback_residency_changes( applied=applied, modules=modules, server_args=server_args ) + if previous_host_pin_budget_state is not None and not applied: + server_args.host_pin_budget().restore_capacity( + previous_host_pin_budget_state + ) except Exception as rollback_error: raise AutoResidencyRollbackError( f"residency adjustment failed ({describe_error(apply_error)}) and rollback " @@ -2617,6 +2652,20 @@ def rollback_residency_changes( """ applied = list(applied) errors: list[str] = [] + previous_host_pin_budget_states = { + adjustment.previous_host_pin_budget_state + for adjustment in applied + if adjustment.previous_host_pin_budget_state is not None + } + if len(previous_host_pin_budget_states) > 1: + errors.append("host pin capacity history is inconsistent") + elif previous_host_pin_budget_states: + try: + server_args.host_pin_budget().restore_capacity( + previous_host_pin_budget_states.pop() + ) + except Exception as e: + errors.append(f"host pin capacity: {describe_error(e)}") pinning_changes = [ adjustment for adjustment in applied 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 7edac38d1cd9..e122c19473e5 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 @@ -211,6 +211,10 @@ def __init__( if reserve_bytes is None else reserve_bytes ) + # This worker's non-overlapping share of the node-wide planner budget. + # Auto placement may later assign a different execution quota, but the + # next joint solve must still see the original node capacity. + self.planning_capacity_bytes = max(0, self.available_bytes - self.reserve_bytes) self.committed_bytes = 0 @classmethod @@ -228,6 +232,24 @@ def for_local_worker(cls, local_worker_count: int) -> "HostPinBudget": def spendable_bytes(self) -> int: return max(0, self.available_bytes - self.reserve_bytes - self.committed_bytes) + def set_spendable_capacity(self, capacity_bytes: int) -> tuple[int, int]: + """Replace this worker's quota and return its previous budget state. + + The startup default is an equal, non-overlapping node share. Once all + ranks have reported their real layer-store sizes, the joint placement + planner can safely replace that provisional share with an asymmetric + quota whose node-wide sum stays within the same allowance. + """ + if capacity_bytes < 0: + raise ValueError("host pin capacity must be non-negative") + previous = (self.available_bytes, self.reserve_bytes) + self.available_bytes = capacity_bytes + self.reserve_bytes = 0 + return previous + + def restore_capacity(self, state: tuple[int, int]) -> None: + self.available_bytes, self.reserve_bytes = state + def request(self, *, component_name: str, weight_bytes: int) -> bool: """Whether `component_name` may pin `weight_bytes`, and book it if so. diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py index 73811867c778..13531817a252 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py @@ -7,6 +7,7 @@ from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( + LAYERWISE_OFFLOAD, RESIDENT, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_weight_inventory import ( @@ -59,9 +60,10 @@ def choose_initial_resident_components( """Choose resident startup components without removing runtime options. This is a load-feasibility seed, not a second serving-placement planner. - Every selected ordinary module can still be demoted to component or - layerwise offload by warmup calibration. Unknown weights and explicit - component choices remain on their configured load path. + Native DiTs using the coarse residency mechanism can still be demoted to + component offload by warmup calibration. Auxiliary modules, layerwise + components, unknown weights, and explicit choices retain their configured + loading semantics. """ reserve_bytes = max( MIN_INITIAL_RESIDENCY_RESERVE_BYTES, @@ -89,17 +91,21 @@ def choose_initial_resident_components( if ( weight_bytes is None or weight_bytes <= 0 + or not is_dit_component_name(item.component_name) or server_args.residency_mode(item.component_name) == RESIDENT + or ( + server_args.configured_residency_mode(item.component_name) + == LAYERWISE_OFFLOAD + ) or server_args.explicit_residency_mode(item.component_name) is not None ): continue - uses = denoising_steps if is_dit_component_name(item.component_name) else 1 options.append( PlacementOption( group_key=item.component_name, option_key=f"{item.component_name}:resident", resource_delta_bytes={"gpu:load": weight_bytes}, - estimated_latency_savings=weight_bytes * uses, + estimated_latency_savings=weight_bytes * denoising_steps, preference_cost=(weight_bytes,), ) ) 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 cc0a6e114c9a..18e36500c581 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -42,6 +42,9 @@ LAYERWISE_OFFLOAD, RESIDENT, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + HostPinBudget, +) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, ) @@ -1869,9 +1872,9 @@ def test_node_hostpin_is_optimized_with_the_same_selection_vector(self): "hot_dit", "cold_encoder", ] - assert plan.resource_budget_bytes["hostpin:node0:rank0"] == 0 + assert plan.resource_budget_bytes["hostpin:node0"] == 0 - def test_dynamic_hostpin_is_constrained_for_every_worker(self): + def test_dynamic_hostpin_uses_one_node_budget_and_assigns_rank_quotas(self): pin_more = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, @@ -1899,8 +1902,86 @@ def test_dynamic_hostpin_is_constrained_for_every_worker(self): 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 + assert plan.resource_budget_bytes["hostpin:node0"] == 40 * GIB_BYTES + assert plan.host_pin_target_bytes_by_rank == { + 0: 10 * GIB_BYTES, + 1: 10 * GIB_BYTES, + } + + def test_node_hostpin_can_assign_asymmetric_rank_quotas(self): + rank0 = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=30 * GIB_BYTES, + ) + rank1 = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=5 * GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + rank=0, + host_pin_capacity_gib=20, + candidates=[rank0], + ), + _report( + rank=1, + host_pin_capacity_gib=20, + candidates=[rank1], + ), + ] + ) + + assert [candidate.component_name for candidate in plan.changes] == [ + "transformer" + ] + assert plan.host_pin_target_bytes_by_rank == { + 0: 30 * GIB_BYTES, + 1: 5 * GIB_BYTES, + } + + def test_hostpin_capacity_is_scoped_per_node(self): + pin_more = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=30 * GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + rank=0, + node_rank=0, + host_pin_capacity_gib=20, + candidates=[pin_more], + ), + _report( + rank=1, + node_rank=1, + host_pin_capacity_gib=20, + candidates=[pin_more], + ), + ] + ) + + assert plan.changes == [] + assert plan.resource_budget_bytes["hostpin:node0"] == 20 * GIB_BYTES + assert plan.resource_budget_bytes["hostpin:node1"] == 20 * GIB_BYTES def test_hostpin_repack_must_fit_transition_headroom(self): pin_more = ResidencyTarget( @@ -1924,7 +2005,7 @@ def test_hostpin_repack_must_fit_transition_headroom(self): ) assert plan.changes == [] - assert plan.resource_budget_bytes["hostram:node0:rank0:pin"] == 5 * GIB_BYTES + assert plan.resource_budget_bytes["hostram:node0:pin"] == 5 * GIB_BYTES def test_component_demotion_must_fit_host_materialization_headroom(self): demote = ResidencyTarget( @@ -1946,10 +2027,7 @@ def test_component_demotion_must_fit_host_materialization_headroom(self): ) assert plan.changes == [] - assert ( - plan.resource_budget_bytes["hostram:node0:rank0:materialize"] - == 5 * GIB_BYTES - ) + assert plan.resource_budget_bytes["hostram:node0:materialize"] == 5 * GIB_BYTES def test_layerwise_materialization_must_fit_transition_vram(self): resident = ResidencyTarget( @@ -2229,8 +2307,9 @@ def configure_layerwise_offload( class _StubResidencyArgs: """Duck-typed stand-in for the two ServerArgs hooks adjustments use.""" - def __init__(self): + def __init__(self, host_pin_budget=None): self.auto_modes: dict[str, str] = {} + self._host_pin_budget = host_pin_budget def auto_residency_mode(self, component_name): return self.auto_modes.get(component_name) @@ -2245,7 +2324,7 @@ def clear_auto_residency_mode(self, component_name): self.auto_modes.pop(component_name, None) def host_pin_budget(self): - return None + return self._host_pin_budget def _plan_for(candidates: list[ResidencyTarget]) -> AutoResidencyPlan: @@ -2802,6 +2881,73 @@ def test_layerwise_frontier_accounts_for_immediate_materialization(self): class TestApplyAndRollback: + def test_apply_and_rollback_restore_rank_hostpin_quota(self): + budget = HostPinBudget( + available_bytes=40 * GIB_BYTES, + reserve_bytes=2 * GIB_BYTES, + ) + args = _StubResidencyArgs(host_pin_budget=budget) + target = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=GIB_BYTES, + ) + plan = AutoResidencyPlan( + host_pin_target_bytes_by_rank={0: 12 * GIB_BYTES}, + changes=[target], + ) + + applied = apply_residency_changes( + plan=plan, + modules={"text_encoder": nn.Linear(2, 2)}, + server_args=args, + rank=0, + ) + assert budget.available_bytes == 12 * GIB_BYTES + assert budget.reserve_bytes == 0 + + rollback_residency_changes( + applied=applied, + modules={"text_encoder": nn.Linear(2, 2)}, + server_args=args, + ) + assert budget.available_bytes == 40 * GIB_BYTES + assert budget.reserve_bytes == 2 * GIB_BYTES + assert budget.planning_capacity_bytes == 38 * GIB_BYTES + + def test_apply_failure_restores_rank_hostpin_quota(self): + budget = HostPinBudget( + available_bytes=40 * GIB_BYTES, + reserve_bytes=2 * GIB_BYTES, + ) + args = _StubResidencyArgs(host_pin_budget=budget) + plan = AutoResidencyPlan( + host_pin_target_bytes_by_rank={0: 12 * GIB_BYTES}, + changes=[ + ResidencyTarget( + component_name="missing", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=GIB_BYTES, + ) + ], + ) + + with pytest.raises(RuntimeError, match="is missing"): + apply_residency_changes( + plan=plan, + modules={}, + server_args=args, + rank=0, + ) + + assert budget.available_bytes == 40 * GIB_BYTES + assert budget.reserve_bytes == 2 * GIB_BYTES + assert budget.planning_capacity_bytes == 38 * GIB_BYTES + def test_lazy_layerwise_configuration_rolls_back_to_component_offload(self): module = _FakeLazyLayerwiseDit(num_layers=4) args = _StubResidencyArgs() 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 1395628b8a6b..d5fc3ff0472d 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 @@ -250,6 +250,22 @@ def test_a_weightless_component_needs_no_budget(self): budget = HostPinBudget(available_bytes=0) assert budget.request(component_name="scheduler", weight_bytes=0) + def test_planner_quota_can_be_applied_and_restored(self): + budget = HostPinBudget( + available_bytes=40 * GIB_BYTES, + reserve_bytes=2 * GIB_BYTES, + ) + + previous = budget.set_spendable_capacity(12 * GIB_BYTES) + assert budget.available_bytes == 12 * GIB_BYTES + assert budget.reserve_bytes == 0 + assert budget.planning_capacity_bytes == 38 * GIB_BYTES + + budget.restore_capacity(previous) + assert budget.available_bytes == 40 * GIB_BYTES + assert budget.reserve_bytes == 2 * GIB_BYTES + assert budget.planning_capacity_bytes == 38 * GIB_BYTES + class TestModuleWeightBytes: def test_parameters_and_buffers_are_counted(self): diff --git a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py index 463b49b681b4..ca750d89bcb0 100644 --- a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py @@ -43,6 +43,9 @@ def residency_mode(self, component_name): self.modes.get(component_name, COMPONENT_OFFLOAD), ) + def configured_residency_mode(self, component_name): + return self.modes.get(component_name, COMPONENT_OFFLOAD) + def explicit_residency_mode(self, component_name): return ( self.modes.get(component_name) if component_name in self.explicit else None @@ -96,7 +99,7 @@ def test_initial_seed_stays_conservative_for_unknown_fixed_resident_weights(): assert selected == set() -def test_initial_seed_does_not_override_explicit_component_mode(): +def test_initial_seed_does_not_override_explicit_or_auxiliary_components(): args = _Args( modes={"transformer": LAYERWISE_OFFLOAD}, explicit={"transformer"}, @@ -108,7 +111,20 @@ def test_initial_seed_does_not_override_explicit_component_mode(): denoising_steps=8, ) - assert selected == {"vae"} + assert selected == set() + + +def test_initial_seed_keeps_layerwise_dit_on_its_configured_load_path(): + args = _Args(modes={"transformer": LAYERWISE_OFFLOAD}) + + selected = choose_initial_resident_components( + args, + [_weight("transformer", 8)], + available_bytes=40 * GIB_BYTES, + denoising_steps=8, + ) + + assert selected == set() def test_initial_seed_applies_one_reversible_override(): From 781964818c402e5bbc6ec79ea85126649cd4bdfb Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 01:05:24 +0800 Subject: [PATCH 05/70] [diffusion] preserve residency phase semantics --- .../runtime/managers/gpu_worker.py | 66 +++++-- .../memory_managers/auto_residency.py | 186 +++++++++++++----- .../memory_managers/component_manager.py | 77 +++++++- .../test/unit/test_auto_residency.py | 146 +++++++++++++- .../test/unit/test_component_residency.py | 27 ++- 5 files changed, 419 insertions(+), 83 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index cd4eb0177f07..1c8489dd9706 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -86,6 +86,7 @@ rollback_residency_changes, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( + get_global_component_residency_manager, peek_global_component_residency_manager, ) from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( @@ -743,6 +744,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_prefetched_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() @@ -751,6 +753,7 @@ 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_prefetched_components[phase_name] = peak.prefetched_components phase_full_weight_transition_components[phase_name] = ( peak.full_weight_transition_components ) @@ -767,6 +770,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_prefetched_components["request:untracked"] = () phase_full_weight_transition_components["request:untracked"] = () metrics = req.metrics stage_duration_ms = {} @@ -791,6 +795,7 @@ def _record_server_warmup_memory( phase_peak_allocated_bytes=phase_allocated_peaks, phase_active_components=phase_components, phase_used_components=phase_used_components, + phase_prefetched_components=phase_prefetched_components, phase_full_weight_transition_components=( phase_full_weight_transition_components ), @@ -1609,6 +1614,7 @@ def _build_pre_warmup_auto_residency_report( estimated_peak_bytes_by_phase=phase_peaks, active_components_by_phase=phase_active, used_components_by_phase=phase_used, + prefetched_components_by_phase={}, current_device_weight_bytes_by_component=current_device_weights, current_active_weight_bytes_by_component=active_weights, node_rank=self.server_args.node_rank, @@ -1659,6 +1665,7 @@ def _build_auto_residency_report( estimated_phase_peaks, active_components_by_phase, used_components_by_phase, + prefetched_components_by_phase, full_weight_transition_components_by_phase, ) = measured_failed_workload_phase_peaks( records=records, @@ -1676,6 +1683,7 @@ def _build_auto_residency_report( estimated_phase_peaks, active_components_by_phase, used_components_by_phase, + prefetched_components_by_phase, full_weight_transition_components_by_phase, ) = estimate_workload_phase_peaks( records=records, @@ -1685,9 +1693,38 @@ def _build_auto_residency_report( budget_bytes = self._auto_residency_budget_bytes() candidates = self._collect_auto_residency_targets(workload) if warmup_oom: + dit_candidate_deltas: dict[str, list[tuple[int, int, int]]] = {} + for candidate in candidates: + if not is_dit_component_name(candidate.component_name): + continue + dit_candidate_deltas.setdefault(candidate.component_name, []).append( + ( + round(candidate.active_device_delta_bytes / 1024**2), + round(candidate.inactive_device_delta_bytes / 1024**2), + round(candidate.device_transition_delta_bytes / 1024**2), + ) + ) + candidate_summary = { + name: { + "count": len(deltas), + "active_mib": ( + min(d[0] for d in deltas), + max(d[0] for d in deltas), + ), + "inactive_mib": ( + min(d[1] for d in deltas), + max(d[1] for d in deltas), + ), + "transition_mib": ( + min(d[2] for d in deltas), + max(d[2] for d in deltas), + ), + } + for name, deltas in dit_candidate_deltas.items() + } logger.debug( "Auto residency OOM report on rank %d: phase_peaks_gib=%s, " - "active=%s, used=%s, candidate_deltas_mib=%s", + "active=%s, used=%s, prefetched=%s, candidate_deltas_mib=%s", self.rank, { phase: round(peak / GIB_BYTES, 3) @@ -1695,17 +1732,8 @@ def _build_auto_residency_report( }, active_components_by_phase, used_components_by_phase, - [ - ( - candidate.option_key(), - round(candidate.active_device_delta_bytes / 1024**2), - round(candidate.inactive_device_delta_bytes / 1024**2), - round(candidate.device_transition_delta_bytes / 1024**2), - candidate.current_placement, - ) - for candidate in candidates - if is_dit_component_name(candidate.component_name) - ], + prefetched_components_by_phase, + candidate_summary, ) ( measured_request_duration_ns, @@ -1765,6 +1793,7 @@ 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, + prefetched_components_by_phase=prefetched_components_by_phase, full_weight_transition_components_by_phase=( full_weight_transition_components_by_phase ), @@ -1804,12 +1833,15 @@ 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: + def _mixed_dtype_residency_components(self) -> set[str]: + if self.pipeline is None: return set() - return manager.components_with_mixed_use_dtypes() + manager = get_global_component_residency_manager( + self.pipeline, self.server_args + ) + return manager.components_with_mixed_use_dtypes( + self.pipeline.stages, self.server_args + ) def _auto_residency_modules(self) -> dict[str, object]: assert self.pipeline is not 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 c70370176a35..fe7a915332b2 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 @@ -139,6 +139,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_prefetched_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 @@ -178,6 +179,10 @@ class ResidencyTarget(msgspec.Struct, frozen=True): # 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 + # Delta for a prefetched component while another component owns the phase. + # Component offload suppresses that overlap; resident and layerwise targets + # use the same footprint as an ordinary active use. + concurrent_prefetch_device_delta_bytes: int | None = None 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 @@ -197,6 +202,11 @@ def target_mode(self) -> str: return LAYERWISE_OFFLOAD return RESIDENT + def concurrent_prefetch_delta(self) -> int: + if self.concurrent_prefetch_device_delta_bytes is not None: + return self.concurrent_prefetch_device_delta_bytes + return self.active_device_delta_bytes + def option_key(self) -> str: target_mode = self.target_mode() if target_mode == COMPONENT_OFFLOAD: @@ -247,6 +257,7 @@ 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, ...]] = {} + prefetched_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] = {} current_active_weight_bytes_by_component: dict[str, int] = {} @@ -646,6 +657,7 @@ def estimate_workload_phase_peaks( dict[str, tuple[str, ...]], dict[str, tuple[str, ...]], dict[str, tuple[str, ...]], + dict[str, tuple[str, ...]], ]: """Estimate each measured execution phase at the target workload. @@ -665,7 +677,13 @@ def estimate_workload_phase_peaks( successful = covering grouped: dict[ - tuple[str, tuple[str, ...], tuple[str, ...], tuple[str, ...]], + tuple[ + str, + tuple[str, ...], + tuple[str, ...], + tuple[str, ...], + tuple[str, ...], + ], list[WarmupMemoryRecord], ] = {} for record in successful: @@ -673,28 +691,33 @@ def estimate_workload_phase_peaks( 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, ()))) + prefetched = tuple( + sorted(record.phase_prefetched_components.get(phase_name, ())) + ) full_weight_transitions = tuple( sorted( record.phase_full_weight_transition_components.get(phase_name, ()) ) ) grouped.setdefault( - (phase_name, active, used, full_weight_transitions), [] + (phase_name, active, used, prefetched, 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, ...]] = {} + prefetched_components: dict[str, tuple[str, ...]] = {} full_weight_transition_components: dict[str, tuple[str, ...]] = {} layout_indices: dict[str, int] = {} for ( phase_name, active, used, + prefetched, full_weight_transitions, ), phase_records in sorted(grouped.items()): output_name = phase_name @@ -727,11 +750,13 @@ def estimate_workload_phase_peaks( estimated_peaks[output_name] = estimate active_components[output_name] = active used_components[output_name] = used + prefetched_components[output_name] = prefetched full_weight_transition_components[output_name] = full_weight_transitions return ( estimated_peaks, active_components, used_components, + prefetched_components, full_weight_transition_components, ) @@ -746,6 +771,7 @@ def measured_failed_workload_phase_peaks( dict[str, tuple[str, ...]], dict[str, tuple[str, ...]], dict[str, tuple[str, ...]], + dict[str, tuple[str, ...]], ]: """Preserve the binding phases reached before a target-workload OOM. @@ -760,10 +786,17 @@ def measured_failed_workload_phase_peaks( if not record.succeeded and record.workload_units() <= target_units ] if not failed: - return None, {}, {}, {}, {} + return None, {}, {}, {}, {}, {} grouped: dict[ - tuple[str, tuple[str, ...], tuple[str, ...], tuple[str, ...]], int + tuple[ + str, + tuple[str, ...], + tuple[str, ...], + tuple[str, ...], + tuple[str, ...], + ], + int, ] = {} for record in failed: used_by_phase = record.phase_used_components or record.phase_active_components @@ -773,24 +806,34 @@ def measured_failed_workload_phase_peaks( for phase_name, phase_peak in phases.items(): active = tuple(sorted(record.phase_active_components.get(phase_name, ()))) used = tuple(sorted(used_by_phase.get(phase_name, ()))) + prefetched = tuple( + sorted(record.phase_prefetched_components.get(phase_name, ())) + ) transitions = tuple( sorted( record.phase_full_weight_transition_components.get(phase_name, ()) ) ) - key = (phase_name, active, used, transitions) + key = (phase_name, active, used, prefetched, transitions) grouped[key] = max(grouped.get(key, 0), phase_peak) layout_counts: dict[str, int] = {} - for phase_name, _, _, _ in grouped: + for phase_name, _, _, _, _ in grouped: layout_counts[phase_name] = layout_counts.get(phase_name, 0) + 1 peaks: dict[str, int] = {} active_components: dict[str, tuple[str, ...]] = {} used_components: dict[str, tuple[str, ...]] = {} + prefetched_components: dict[str, tuple[str, ...]] = {} transition_components: dict[str, tuple[str, ...]] = {} layout_indices: dict[str, int] = {} - for (phase_name, active, used, transitions), phase_peak in sorted(grouped.items()): + for ( + phase_name, + active, + used, + prefetched, + transitions, + ), phase_peak in sorted(grouped.items()): output_name = phase_name if layout_counts[phase_name] > 1: index = layout_indices.get(phase_name, 0) @@ -799,6 +842,7 @@ def measured_failed_workload_phase_peaks( peaks[output_name] = phase_peak active_components[output_name] = active used_components[output_name] = used + prefetched_components[output_name] = prefetched transition_components[output_name] = transitions return ( @@ -806,6 +850,7 @@ def measured_failed_workload_phase_peaks( peaks, active_components, used_components, + prefetched_components, transition_components, ) @@ -1242,6 +1287,11 @@ def _unconfigured_layerwise_targets( + active_managed_bytes - full_weight_bytes ), + concurrent_prefetch_device_delta_bytes=( + unmanaged_weight_bytes + + active_managed_bytes + - (full_weight_bytes if current_resident else 0) + ), present_device_delta_bytes=( unmanaged_weight_bytes + active_managed_bytes @@ -1357,6 +1407,7 @@ def collect_residency_targets( if ( isinstance(module, LayerwiseOffloadableModuleMixin) and layerwise_tuning_of is not None + and name not in mixed_dtype_names ): prefetch_value, _, residency_policy = layerwise_tuning_of( name, module.layerwise_offload_dit_group_enabled @@ -1405,6 +1456,9 @@ def collect_residency_targets( present_device_delta_bytes=( -weight_bytes if current_resident else 0 ), + concurrent_prefetch_device_delta_bytes=( + -weight_bytes if current_resident else 0 + ), current_placement=not current_resident, ), ResidencyTarget( @@ -1420,6 +1474,9 @@ def collect_residency_targets( 0 if current_resident else weight_bytes ), present_device_delta_bytes=0, + concurrent_prefetch_device_delta_bytes=( + 0 if current_resident else weight_bytes + ), current_placement=current_resident, target_device_weight_bytes=weight_bytes, ), @@ -1489,33 +1546,37 @@ def collect_residency_targets( 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, + if name not in mixed_dtype_names: + 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=( + managed_weight_bytes - current_peak_device_bytes + ), + present_device_delta_bytes=( + managed_weight_bytes - current_peak_device_bytes + ), + concurrent_prefetch_device_delta_bytes=( + -unmanaged_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( @@ -1577,6 +1638,9 @@ def collect_residency_targets( active_device_delta_bytes=( target_peak_device_bytes - current_peak_device_bytes ), + concurrent_prefetch_device_delta_bytes=( + target_peak_device_bytes - current_peak_device_bytes + ), inactive_device_delta_bytes=-current_inactive_device_bytes, present_device_delta_bytes=-current_inactive_device_bytes, current_placement=( @@ -1591,11 +1655,10 @@ def collect_residency_targets( ) ) - # 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. + # Layerwise stores have one fixed dtype, while both coarse strategies + # cast at every declared use. A mixed-dtype component may be tuned within + # its current strategy family, but crossing that boundary would change + # its numerical path. if name in mixed_dtype_names: continue @@ -1648,6 +1711,9 @@ def collect_residency_targets( present_device_delta_bytes=( managed_weight_bytes - current_inactive_device_bytes ), + concurrent_prefetch_device_delta_bytes=( + managed_weight_bytes - current_peak_device_bytes + ), current_placement=( current_permanent and permanent_pins == current_pinned_layers ), @@ -1727,6 +1793,7 @@ def _binding_phase_constraints( tuple[str, ...], tuple[str, ...], tuple[str, ...], + tuple[str, ...], ] ]: """Keep the binding measured and conservative unobserved phases. @@ -1781,6 +1848,7 @@ def _binding_phase_constraints( tuple[str, ...], tuple[str, ...], tuple[str, ...], + tuple[str, ...], ], tuple[str, int], ] = {} @@ -1788,6 +1856,9 @@ def _binding_phase_constraints( for phase_name, phase_peak in phase_peaks.items(): measured_components = set(report.active_components_by_phase.get(phase_name, ())) used = tuple(sorted(report.used_components_by_phase.get(phase_name, ()))) + prefetched = tuple( + sorted(report.prefetched_components_by_phase.get(phase_name, ())) + ) full_weight_transitions = tuple( sorted( set( @@ -1821,7 +1892,12 @@ def _binding_phase_constraints( # snapshot. The resulting request:untracked peak is useful as a # conservative baseline for the synthetic component phases below, but # no placement can release an empty-layout constraint directly. - if report.warmup_oom and not used and not full_weight_transitions: + if ( + report.warmup_oom + and not used + and not prefetched + and not full_weight_transitions + ): continue present = tuple( sorted( @@ -1830,14 +1906,16 @@ def _binding_phase_constraints( ) ) measured = tuple(sorted(measured_components & candidate_component_names)) - layout = (present, measured, used, full_weight_transitions) + layout = (present, measured, used, prefetched, 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 component_name in set(used) | set(full_weight_transitions) + for (_, _, used, prefetched, full_weight_transitions) in binding + for component_name in ( + set(used) | set(prefetched) | set(full_weight_transitions) + ) } # Separate the stable/request activation baseline from transient component # weights before constructing phases that a short failed probe did not @@ -1859,7 +1937,7 @@ def _binding_phase_constraints( (steady_state_components | {component_name}) & candidate_component_names ) ) - layout = (present, (), (component_name,), ()) + layout = (present, (), (component_name,), (), ()) synthetic = ( ( f"lower-bound:{component_name}" @@ -1883,9 +1961,10 @@ def _binding_phase_constraints( present, measured, used, + prefetched, full_weight_transitions, ) - for (present, measured, used, full_weight_transitions), ( + for (present, measured, used, prefetched, full_weight_transitions), ( phase_name, phase_peak, ) in sorted(binding.items(), key=lambda item: item[1][0]) @@ -1973,6 +2052,10 @@ def _consensus_candidates( candidate.present_device_delta_bytes for candidate in rank_candidates ), + concurrent_prefetch_device_delta_bytes=max( + candidate.concurrent_prefetch_delta() + for candidate in rank_candidates + ), inactive_device_delta_bytes=max( candidate.inactive_device_delta_bytes for candidate in rank_candidates @@ -2061,7 +2144,9 @@ 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, ) @@ -2072,7 +2157,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. @@ -2189,6 +2274,7 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP present_components, measured_components, used_components, + prefetched_components, full_weight_transition_components, ) in phase_constraints[report.rank]: # ``present_components`` also includes request-end placement @@ -2218,6 +2304,12 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP 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 prefetched_components: + phase_cost = ( + rank_candidate.concurrent_prefetch_delta() + if used_components + else 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: 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 d7731c5402bc..646cb0967ec8 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, ...] = () + prefetched_components: tuple[str, ...] = () full_weight_transition_components: tuple[str, ...] = () @@ -141,6 +142,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_prefetched_components: tuple[str, ...] = () self._warmup_phase_full_weight_transition_components: tuple[str, ...] = () self._warmup_phase_peaks: dict[str, WarmupPhasePeak] = {} self._warmup_phase_occurrences: dict[str, int] = {} @@ -180,7 +182,26 @@ 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]: + def declared_stage_uses( + self, + stages: Sequence[ComponentResidencyStage], + server_args: ServerArgs, + ) -> tuple[tuple[ComponentUse, ...], ...]: + """Resolve the stage timeline without starting a request. + + Static placement runs before the first forward, so it cannot rely on + ``begin_request`` having populated the cached timeline already. + """ + return tuple( + tuple(stage.component_uses(server_args, self.stage_name(stage))) + for stage in stages + ) + + def components_with_mixed_use_dtypes( + self, + stages: Sequence[ComponentResidencyStage] | None = None, + server_args: ServerArgs | None = None, + ) -> set[str]: """Components whose stages request more than one runtime dtype. Layerwise offload keeps its managed weight stores in one fixed dtype, @@ -188,8 +209,19 @@ def components_with_mixed_use_dtypes(self) -> set[str]: Switching such a component between those strategies would therefore change its numerical path, even when both placements fit. """ + if stages is None: + ordered_uses = self._ordered_uses + else: + resolved_server_args = ( + self.server_args if server_args is None else server_args + ) + ordered_uses = tuple( + use + for uses in self.declared_stage_uses(stages, resolved_server_args) + for use in uses + ) dtypes_by_component: dict[str, set[torch.dtype | None]] = {} - for use in self._ordered_uses: + for use in ordered_uses: dtypes_by_component.setdefault(use.component_name, set()).add( use.target_dtype ) @@ -217,10 +249,7 @@ def begin_request( self._prefetched_use_keys.clear() self._uses_seen.clear() self._modules_seen.clear() - self._stage_uses_by_index = [ - tuple(stage.component_uses(server_args, self.stage_name(stage))) - for stage in stages - ] + self._stage_uses_by_index = list(self.declared_stage_uses(stages, server_args)) self._ordered_uses = tuple( use for uses in self._stage_uses_by_index for use in uses ) @@ -232,6 +261,7 @@ def begin_request( self._warmup_phase_key = None self._warmup_phase_components = () self._warmup_phase_used_components = () + self._warmup_phase_prefetched_components = () self._warmup_phase_full_weight_transition_components = () self._warmup_phase_peaks = {} self._warmup_phase_occurrences = {} @@ -278,6 +308,7 @@ def full_weight_transition(self, component_names: Iterable[str]) -> Iterator[Non self._warmup_phase_key, self._warmup_phase_components, self._warmup_phase_used_components, + self._warmup_phase_prefetched_components, self._warmup_phase_full_weight_transition_components, ) self._begin_warmup_phase( @@ -292,7 +323,13 @@ def full_weight_transition(self, component_names: Iterable[str]) -> Iterator[Non try: yield finally: - previous_key, components, used_components, transitions = previous_phase + ( + previous_key, + components, + used_components, + prefetched_components, + transitions, + ) = previous_phase if previous_key is None: self._record_warmup_phase_peak() self._warmup_phase_key = None @@ -301,6 +338,7 @@ def full_weight_transition(self, component_names: Iterable[str]) -> Iterator[Non key=previous_key, components=components, used_components=used_components, + prefetched_components=prefetched_components, full_weight_transition_components=transitions, new_occurrence=False, ) @@ -661,6 +699,7 @@ def _begin_warmup_phase( key: str, components: tuple[str, ...], used_components: tuple[str, ...], + prefetched_components: tuple[str, ...] = (), full_weight_transition_components: tuple[str, ...] = (), new_occurrence: bool = True, ) -> None: @@ -674,6 +713,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_prefetched_components = tuple( + sorted(set(prefetched_components)) + ) self._warmup_phase_full_weight_transition_components = tuple( sorted(set(full_weight_transition_components)) ) @@ -716,10 +758,20 @@ def _begin_warmup_between_uses(self) -> None: def _begin_warmup_prefetch(self, use: ComponentUse) -> None: phase = use.phase or use.component_name + active_uses = tuple( + active_use + for active_use in (self._active_use, use) + if active_use is not 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,), + components=self._warmup_active_components(active_uses), + used_components=( + (self._active_use.component_name,) + if self._active_use is not None + else () + ), + prefetched_components=(use.component_name,), ) def _warmup_active_components( @@ -742,6 +794,7 @@ 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, + prefetched_components=self._warmup_phase_prefetched_components, full_weight_transition_components=( self._warmup_phase_full_weight_transition_components ), @@ -760,6 +813,12 @@ def _record_warmup_phase_peak(self) -> None: used_components=tuple( sorted(set(previous.used_components) & set(peak.used_components)) ), + prefetched_components=tuple( + sorted( + set(previous.prefetched_components) + & set(peak.prefetched_components) + ) + ), full_weight_transition_components=tuple( sorted( set(previous.full_weight_transition_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 18e36500c581..619884cfefd7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -458,7 +458,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}, @@ -492,7 +492,7 @@ def test_failed_probe_preserves_phase_layouts_for_recovery(self): phase_used_components={"denoise": ("transformer_2",)}, ) - peak, phases, active, used, _ = measured_failed_workload_phase_peaks( + peak, phases, active, used, _, _ = measured_failed_workload_phase_peaks( records=[first, second], target_units=first.workload_units(), ) @@ -517,7 +517,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}, @@ -537,7 +537,7 @@ def test_phase_estimation_preserves_full_weight_transition_components(self): phase_full_weight_transition_components={"lora_switch": ("transformer",)}, ) - peaks, _, _, transitions = estimate_workload_phase_peaks( + peaks, _, _, _, transitions = estimate_workload_phase_peaks( records=[record], target_units=record.workload_units(), component_weight_bytes={"transformer": 2 * GIB_BYTES}, @@ -568,7 +568,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}, @@ -600,7 +600,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={ @@ -619,6 +619,53 @@ def test_phase_estimation_keeps_distinct_active_layouts_separate(self): } assert used == active + def test_phase_estimation_keeps_distinct_prefetch_layouts_separate(self): + without_prefetch = 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={"encode": 20 * GIB_BYTES}, + phase_active_components={"encode": ("text_encoder",)}, + phase_used_components={"encode": ("text_encoder",)}, + ) + with_prefetch = WarmupMemoryRecord( + width=768, + height=512, + num_frames=25, + baseline_allocated_bytes=2 * GIB_BYTES, + peak_allocated_bytes=30 * GIB_BYTES, + succeeded=True, + phase_peak_allocated_bytes={"encode": 30 * GIB_BYTES}, + phase_active_components={ + "encode": ("text_encoder", "transformer"), + }, + phase_used_components={"encode": ("text_encoder",)}, + phase_prefetched_components={"encode": ("transformer",)}, + ) + + peaks, active, used, prefetched, _ = estimate_workload_phase_peaks( + records=[without_prefetch, with_prefetch], + target_units=without_prefetch.workload_units(), + component_weight_bytes={ + "transformer": 10 * GIB_BYTES, + "text_encoder": 8 * GIB_BYTES, + }, + ) + + assert peaks == { + "encode:layout:0": 20 * GIB_BYTES, + "encode:layout:1": 30 * GIB_BYTES, + } + assert set(active.values()) == { + ("text_encoder",), + ("text_encoder", "transformer"), + } + assert set(used.values()) == {("text_encoder",)} + assert set(prefetched.values()) == {(), ("transformer",)} + def _candidate( name: str, @@ -653,6 +700,7 @@ def _report( phase_peaks_gib: dict[str, float] | None = None, phase_components: dict[str, tuple[str, ...]] | None = None, phase_present_components: dict[str, tuple[str, ...]] | None = None, + phase_prefetched_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, component_active_weight_gib: dict[str, int] | None = None, @@ -688,6 +736,7 @@ def _report( else phase_components or {} ), used_components_by_phase=phase_components or {}, + prefetched_components_by_phase=phase_prefetched_components or {}, full_weight_transition_components_by_phase=( phase_full_weight_transition_components or {} ), @@ -1246,6 +1295,9 @@ def test_async_prefetch_presence_is_not_charged_twice(self): "text_and_prefetch": ("text_encoder", "transformer"), "denoise": ("transformer",), }, + phase_prefetched_components={ + "text_and_prefetch": ("transformer",), + }, candidates=[offload, resident], ) ] @@ -1255,6 +1307,56 @@ def test_async_prefetch_presence_is_not_charged_twice(self): "transformer:resident" ] + def test_component_offload_can_release_an_overlapped_layerwise_prefetch(self): + layerwise = ResidencyTarget( + component_name="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=((),), + current_placement=True, + ) + component_offload = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + concurrent_prefetch_device_delta_bytes=-10 * GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=18, + estimated_gib=18, + warmup_oom=True, + candidates=[layerwise, component_offload], + phase_peaks_gib={"encode_prefetch": 18}, + phase_components={"encode_prefetch": ("text_encoder",)}, + phase_present_components={ + "encode_prefetch": ("text_encoder", "transformer"), + }, + phase_prefetched_components={ + "encode_prefetch": ("transformer",), + }, + component_active_weight_gib={ + "text_encoder": 8, + "transformer": 10, + }, + ) + ] + ) + + assert plan.skip_reason is None + assert [candidate.target_mode() for candidate in plan.changes] == [ + COMPONENT_OFFLOAD + ] + def test_request_end_residency_is_carried_into_the_next_request(self): text_offload = ResidencyTarget( component_name="text_encoder", @@ -1665,12 +1767,15 @@ def test_unobserved_phase_excludes_another_components_transient_weights(self): "prefetch:transformer": 11.75, }, phase_components={ - "prefetch:transformer": ("transformer",), + "prefetch:transformer": (), }, phase_present_components={ "setup": ("vae",), "prefetch:transformer": ("transformer", "vae"), }, + phase_prefetched_components={ + "prefetch:transformer": ("transformer",), + }, component_weight_gib={"vae": 1}, component_active_weight_gib={"transformer": 11, "vae": 1}, ) @@ -2554,7 +2659,32 @@ def test_mixed_dtype_layerwise_component_is_not_permanently_promoted(self): mixed_dtype_components={"vae"}, ) - assert not any(candidate.permanent_residency for candidate in candidates) + assert candidates + assert {candidate.target_mode() for candidate in candidates} == { + LAYERWISE_OFFLOAD + } + + def test_mixed_dtype_coarse_component_does_not_add_virtual_layerwise(self): + module = _FakeLazyLayerwiseDit() + + candidates = collect_residency_targets( + modules={"vae": module}, + residency_mode_of=self._modes({"vae": COMPONENT_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=10, + mixed_dtype_components={"vae"}, + layerwise_tuning_of=lambda _name, _dit_group: ( + 1.0, + 0.0, + "leading", + ), + ) + + assert {candidate.target_mode() for candidate in candidates} == { + COMPONENT_OFFLOAD, + RESIDENT, + } def test_non_dit_layerwise_frontier_includes_partial_residency(self): manager = _FakeLayerwiseManager( 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 596fb1560c5b..b672752e494f 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_residency.py @@ -233,6 +233,24 @@ def test_reports_components_with_mixed_use_dtypes(): assert manager.components_with_mixed_use_dtypes() == {"vae"} +def test_reports_mixed_use_dtypes_before_first_request(): + encode = _Stage(ComponentUse("encode", "vae", target_dtype=torch.float32)) + decode = _Stage(ComponentUse("decode", "vae", target_dtype=torch.bfloat16)) + pipeline = SimpleNamespace( + modules={}, + _stage_name_mapping={"encode": encode, "decode": decode}, + component_residency_strategies={}, + ) + server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + manager = ComponentResidencyManager(pipeline, server_args) + manager.refresh_pipeline(pipeline) + + assert manager.components_with_mixed_use_dtypes([encode, decode], server_args) == { + "vae" + } + assert manager._ordered_uses == () + + def test_warmup_records_use_and_transition_peaks(monkeypatch): device_module = SimpleNamespace( is_available=lambda: True, @@ -312,7 +330,7 @@ def test_failed_warmup_phase_survives_cleanup(monkeypatch): peaks = manager.take_warmup_phase_peaks() assert peaks["0:denoise:prefetch:transformer"] == WarmupPhasePeak( - ("transformer",), 7, used_components=("transformer",) + ("transformer",), 7, prefetched_components=("transformer",) ) assert peaks["request:cleanup"] == WarmupPhasePeak((), 2) @@ -437,7 +455,12 @@ def test_warmup_attributes_prefetch_peak_to_prefetched_component(monkeypatch): assert manager._warmup_phase_peaks[ "0:encode:prefetch:transformer" - ] == WarmupPhasePeak(("transformer",), 7, used_components=("transformer",)) + ] == WarmupPhasePeak( + ("text_encoder", "transformer"), + 7, + used_components=("text_encoder",), + prefetched_components=("transformer",), + ) assert manager._warmup_phase_peaks["0:encode:between"] == WarmupPhasePeak((), 7) From f25c7a02927000638afad2b99377751791136be5 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 01:16:06 +0800 Subject: [PATCH 06/70] [diffusion] preserve implicit prefetch deltas --- .../managers/memory_managers/auto_residency.py | 13 +++++++++---- 1 file changed, 9 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 fe7a915332b2..ba4281a5c8f5 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 @@ -2014,6 +2014,14 @@ def _consensus_candidates( or len(current) != 1 ): continue + concurrent_prefetch_delta = None + if any( + candidate.concurrent_prefetch_device_delta_bytes is not None + for candidate in rank_candidates + ): + concurrent_prefetch_delta = max( + candidate.concurrent_prefetch_delta() for candidate in rank_candidates + ) merged.append( ResidencyTarget( component_name=component_names.pop(), @@ -2052,10 +2060,7 @@ def _consensus_candidates( candidate.present_device_delta_bytes for candidate in rank_candidates ), - concurrent_prefetch_device_delta_bytes=max( - candidate.concurrent_prefetch_delta() - for candidate in rank_candidates - ), + concurrent_prefetch_device_delta_bytes=concurrent_prefetch_delta, inactive_device_delta_bytes=max( candidate.inactive_device_delta_bytes for candidate in rank_candidates From c30aa9d396992adf257f5344b5f7ed7a0400bef2 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 01:32:30 +0800 Subject: [PATCH 07/70] [diffusion] complete residency pipeline fixtures --- .../sglang/multimodal_gen/test/unit/test_performance_metrics.py | 2 ++ 1 file changed, 2 insertions(+) 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 d69158b5b4cb..87682d535bf9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -252,6 +252,7 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): ) worker.pipeline = SimpleNamespace( modules={"transformer": quantized_module}, + _stage_name_mapping={}, component_residency_strategies={}, ) host_pin_budget = HostPinBudget(available_bytes=0, reserve_bytes=0) @@ -306,6 +307,7 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): worker._auto_residency_reference_component_stages = {} worker.pipeline = SimpleNamespace( modules={"transformer": torch.nn.Linear(1, 1)}, + _stage_name_mapping={}, component_residency_strategies={}, ) host_pin_budget = HostPinBudget(available_bytes=0, reserve_bytes=0) From 4c067be58a7b6e6f52b0bb75cf91273853eff0e0 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 01:34:17 +0800 Subject: [PATCH 08/70] [diffusion] complete auto residency test pipelines --- .../sglang/multimodal_gen/test/unit/test_performance_metrics.py | 2 ++ 1 file changed, 2 insertions(+) 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 87682d535bf9..05ba31074529 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -252,6 +252,7 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): ) worker.pipeline = SimpleNamespace( modules={"transformer": quantized_module}, + stages=[], _stage_name_mapping={}, component_residency_strategies={}, ) @@ -307,6 +308,7 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): worker._auto_residency_reference_component_stages = {} worker.pipeline = SimpleNamespace( modules={"transformer": torch.nn.Linear(1, 1)}, + stages=[], _stage_name_mapping={}, component_residency_strategies={}, ) From b55c0ea995179cd57a6176d88a3152369159d806 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 01:41:23 +0800 Subject: [PATCH 09/70] [diffusion] distinguish runtime residency transitions --- .../memory_managers/auto_residency.py | 27 +++---------------- .../test/unit/test_auto_residency.py | 5 ++-- 2 files changed, 6 insertions(+), 26 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 ba4281a5c8f5..095518d89970 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 @@ -2916,34 +2916,13 @@ 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. - """ + """Describe runtime transitions without claiming startup equivalence.""" changes = "; ".join( _format_residency_change(candidate) for candidate in plan.changes ) - 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" - ) - if has_auto_only_change: - equivalent += " (partial layer/HostPin placement remains auto-only)" 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." + f"Auto residency: adjusted {changes}. Startup flags may use a different " + f"load path; disable with SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY=1." ) 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 619884cfefd7..6a9e131a61e2 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -3626,7 +3626,7 @@ def test_post_request_hint_order_matches_promotion_order(self): class TestAppliedChangesLog: - def test_reports_equivalent_server_args_and_kill_switch(self): + def test_reports_runtime_transition_and_kill_switch(self): plan = _plan_for( [ _candidate("text_encoder", mode=LAYERWISE_OFFLOAD, weight_gib=7), @@ -3636,7 +3636,8 @@ 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 "vae: component-offload -> resident" in message - assert "--component-residency text_encoder=resident vae=resident" in message + assert "Startup flags may use a different load path" in message + assert "--component-residency" not in message assert "SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY=1" in message From 7ef8ea8e6caadc5dad61163fe6aba620c3051c11 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 02:07:40 +0800 Subject: [PATCH 10/70] [diffusion] avoid FSDP fallback load collectives --- .../runtime/loader/fsdp_load.py | 83 ++++++++++++++++--- .../test/unit/test_fsdp_load.py | 66 ++++++++++++++- 2 files changed, 138 insertions(+), 11 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index e1161aceab2f..00518209ffe2 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -16,7 +16,6 @@ import torch.distributed.tensor as dist_tensor from torch import nn from torch.distributed import DeviceMesh, init_device_mesh -from torch.distributed._tensor import distribute_tensor from torch.distributed.fsdp import ( CPUOffloadPolicy, FSDPModule, @@ -60,6 +59,75 @@ _DTYPE_MISMATCH_EXAMPLE_LIMIT = 3 +def _local_dtensor_from_full_tensor( + full_tensor: torch.Tensor, + template: dist_tensor.DTensor, +) -> dist_tensor.DTensor: + """Build this rank's DTensor shard without a load-time collective. + + Every rank has already read the checkpoint tensor on this fallback path. + Slice the final local storage directly instead of asking rank 0 to + redistribute another copy through ``distribute_tensor``. + """ + if tuple(full_tensor.shape) != tuple(template.shape): + raise ValueError( + "FSDP checkpoint tensor shape mismatch: " + f"loaded={tuple(full_tensor.shape)}, expected={tuple(template.shape)}" + ) + local_shape, global_offset = ( + dist_tensor._utils.compute_local_shape_and_global_offset( + template.shape, + template.device_mesh, + template.placements, + ) + ) + slices = tuple( + slice(offset, offset + size) + for offset, size in zip(global_offset, local_shape, strict=True) + ) + local_tensor = full_tensor[slices] + if tuple(local_shape) != tuple(full_tensor.shape): + # Do not retain the full checkpoint storage through a shard view. + local_tensor = local_tensor.clone(memory_format=torch.contiguous_format) + return dist_tensor.DTensor.from_local( + local_tensor, + template.device_mesh, + template.placements, + run_check=False, + shape=template.shape, + stride=template.stride(), + ) + + +def _initialize_local_dtensor( + template: dist_tensor.DTensor, + *, + fill_value: int, + device: torch.device, + dtype: torch.dtype, +) -> dist_tensor.DTensor: + """Initialize only this rank's final FSDP storage.""" + local_shape, _ = dist_tensor._utils.compute_local_shape_and_global_offset( + template.shape, + template.device_mesh, + template.placements, + ) + local_tensor = torch.full( + tuple(local_shape), + fill_value, + device=device, + dtype=dtype, + ) + return dist_tensor.DTensor.from_local( + local_tensor, + template.device_mesh, + template.placements, + run_check=False, + shape=template.shape, + stride=template.stride(), + ) + + def _is_bitsandbytes_quant_config(quant_config: Any | None) -> bool: if quant_config is None: return False @@ -810,10 +878,9 @@ def load_model_from_full_model_state_dict( f"param_cls={type(actual_param).__name__}" ) from exc full_tensor = temp_param.data - sharded_tensor = distribute_tensor( + sharded_tensor = _local_dtensor_from_full_tensor( full_tensor, - meta_sharded_param.device_mesh, - meta_sharded_param.placements, + meta_sharded_param, ) if cpu_offload: sharded_tensor = sharded_tensor.to("cpu") @@ -944,16 +1011,12 @@ def load_model_from_full_model_state_dict( if cpu_offload and not is_fsdp_model: sharded_tensor = sharded_tensor.cpu() else: - full_tensor = init_like( + sharded_tensor = _initialize_local_dtensor( meta_sharded_param, + fill_value=1 if init_like is torch.ones_like else 0, device=checkpoint_load_device, dtype=meta_sharded_param_dtype, ) - sharded_tensor = distribute_tensor( - full_tensor, - meta_sharded_param.device_mesh, - meta_sharded_param.placements, - ) if cpu_offload: sharded_tensor = sharded_tensor.cpu() sharded_sd[new_param_name] = nn.Parameter(sharded_tensor) diff --git a/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py b/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py index 11ff7fa0c18a..28b4e7f7e0af 100644 --- a/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py +++ b/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py @@ -3,7 +3,8 @@ import tempfile import unittest from pathlib import Path -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import patch, sentinel import torch from safetensors.torch import safe_open, save_file @@ -366,3 +367,66 @@ def test_reads_tp_local_merged_row_slice(self): self.assertEqual( rank_local_checkpoint.tp_local_shape(sources, 1, 2), (8, 2) ) + + +class TestRankLocalDTensorConstruction(unittest.TestCase): + @staticmethod + def _template() -> SimpleNamespace: + return SimpleNamespace( + shape=torch.Size((6, 4)), + device_mesh=sentinel.mesh, + placements=(sentinel.shard,), + stride=lambda: (4, 1), + ) + + def test_full_checkpoint_fallback_slices_locally_without_collective(self): + full_tensor = torch.arange(24, dtype=torch.float32).reshape(6, 4) + with ( + patch.object( + fsdp_load.dist_tensor._utils, + "compute_local_shape_and_global_offset", + return_value=((2, 4), (2, 0)), + ), + patch.object(fsdp_load.dist_tensor, "DTensor") as dtensor_cls, + ): + dtensor_cls.from_local.return_value = sentinel.dtensor + result = fsdp_load._local_dtensor_from_full_tensor( + full_tensor, + self._template(), + ) + + self.assertIs(result, sentinel.dtensor) + local_tensor = dtensor_cls.from_local.call_args.args[0] + torch.testing.assert_close(local_tensor, full_tensor[2:4]) + self.assertEqual(local_tensor.untyped_storage().nbytes(), local_tensor.nbytes) + dtensor_cls.from_local.assert_called_once_with( + local_tensor, + sentinel.mesh, + (sentinel.shard,), + run_check=False, + shape=torch.Size((6, 4)), + stride=(4, 1), + ) + + def test_missing_parameter_initializes_only_local_storage(self): + with ( + patch.object( + fsdp_load.dist_tensor._utils, + "compute_local_shape_and_global_offset", + return_value=((2, 4), (2, 0)), + ), + patch.object(fsdp_load.dist_tensor, "DTensor") as dtensor_cls, + ): + dtensor_cls.from_local.return_value = sentinel.dtensor + result = fsdp_load._initialize_local_dtensor( + self._template(), + fill_value=1, + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + + self.assertIs(result, sentinel.dtensor) + local_tensor = dtensor_cls.from_local.call_args.args[0] + self.assertEqual(tuple(local_tensor.shape), (2, 4)) + self.assertEqual(local_tensor.dtype, torch.bfloat16) + torch.testing.assert_close(local_tensor, torch.ones_like(local_tensor)) From d95f560f6dd50f8cb6e427070074a79afd362506 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 02:25:58 +0800 Subject: [PATCH 11/70] [diffusion] preserve LTX two-stage load placement --- .../managers/memory_managers/initial_residency.py | 5 +++++ .../runtime/pipelines/ltx_2_pipeline.py | 3 +++ .../pipelines_core/composed_pipeline_base.py | 10 +++++++++- .../test/unit/test_initial_residency.py | 14 ++++++++++++++ 4 files changed, 31 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py index 13531817a252..fc064dc11e61 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py @@ -56,6 +56,7 @@ def choose_initial_resident_components( *, available_bytes: int, denoising_steps: int | None = None, + excluded_components: frozenset[str] = frozenset(), ) -> set[str]: """Choose resident startup components without removing runtime options. @@ -91,6 +92,7 @@ def choose_initial_resident_components( if ( weight_bytes is None or weight_bytes <= 0 + or item.component_name in excluded_components or not is_dit_component_name(item.component_name) or server_args.residency_mode(item.component_name) == RESIDENT or ( @@ -122,6 +124,8 @@ def choose_initial_resident_components( def maybe_seed_initial_residency( server_args: ServerArgs, inventory: list[ComponentWeightEstimate], + *, + excluded_components: frozenset[str] = frozenset(), ) -> None: """Apply a reversible resident seed for warmup-calibrated auto mode.""" if ( @@ -153,6 +157,7 @@ def maybe_seed_initial_residency( server_args, inventory, available_bytes=int(available_gib * GIB_BYTES), + excluded_components=excluded_components, ) if not selected: return 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 bfe45e2a024a..dc4f1818443b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -638,6 +638,9 @@ def _sync_refinement_stage_transformer(self, phase: str) -> None: class LTX2TwoStagePipeline(_BaseLTX2Pipeline): pipeline_name = "LTX2TwoStagePipeline" + # The built-in stage LoRA establishes the DiT's runtime representation. + # Preserve that initialization path, then let post-load placement tune it. + preload_residency_excluded_components = frozenset(("transformer", "transformer_2")) STAGE_2_DISTILLED_SIGMA_VALUES = list(_SHARED_STAGE_2_DISTILLED_SIGMA_VALUES) STAGE_1_DISTILLED_LORA_STRENGTH = 0.0 STAGE_2_DISTILLED_LORA_STRENGTH = 1.0 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index a1a7c27bb2d5..9012b8e5c43f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -92,6 +92,10 @@ class ComposedPipelineBase(ABC): # should contains only the modules to be loaded _required_config_modules: list[str] = [] _extra_config_module_map: dict[str, str] = {} + # Components whose initialization mutates or aliases weights in a way that + # depends on their load placement. The post-load planner may still tune + # them after those semantics have been established. + preload_residency_excluded_components: frozenset[str] = frozenset() server_args: ServerArgs | None = None modules: dict[str, Any] = {} executor: PipelineExecutor | None = None @@ -488,7 +492,11 @@ def prepare_component_weight_inventory( maybe_seed_initial_residency, ) - maybe_seed_initial_residency(server_args, self.component_weight_inventory) + maybe_seed_initial_residency( + server_args, + self.component_weight_inventory, + excluded_components=self.preload_residency_excluded_components, + ) def load_modules( self, diff --git a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py index ca750d89bcb0..dcb8ea0668f7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py @@ -127,6 +127,20 @@ def test_initial_seed_keeps_layerwise_dit_on_its_configured_load_path(): assert selected == set() +def test_initial_seed_honors_pipeline_load_placement_exclusions(): + args = _Args() + + selected = choose_initial_resident_components( + args, + [_weight("transformer", 8), _weight("transformer_2", 8)], + available_bytes=40 * GIB_BYTES, + denoising_steps=8, + excluded_components=frozenset(("transformer", "transformer_2")), + ) + + assert selected == set() + + def test_initial_seed_applies_one_reversible_override(): args = _Args() inventory = [_weight("transformer", 8)] From 5ef01cfac40e6808bbb3ba3748e02ee9725fbeaf Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 02:59:16 +0800 Subject: [PATCH 12/70] [diffusion] preserve residency regression cost --- .../memory_managers/auto_residency.py | 42 +++++++++++++++++++ .../test/unit/test_auto_residency.py | 38 +++++++++++++++++ 2 files changed, 80 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 095518d89970..92e455c39dc5 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 @@ -488,15 +488,57 @@ def estimate_candidate_latency_savings_ns( candidates = list(candidates) stage_duration_ns = stage_duration_ns or {} component_stages = component_stages or {} + candidates_by_component: dict[str, list[ResidencyTarget]] = {} maximum_h2d_bytes: dict[str, int] = {} for candidate in candidates: + candidates_by_component.setdefault(candidate.component_name, []).append( + candidate + ) maximum_h2d_bytes[candidate.component_name] = max( maximum_h2d_bytes.get(candidate.component_name, 0), candidate.h2d_bytes_per_request, ) estimates: dict[str, int] = {} + for component_name, component_candidates in candidates_by_component.items(): + current = next( + ( + candidate + for candidate in component_candidates + if candidate.current_placement + ), + None, + ) + if current is None: + continue + + component_stage_duration_ns = sum( + stage_duration_ns.get(stage_name, 0) + for stage_name in component_stages.get(component_name, ()) + ) + latency_upper_bound_ns = ( + request_duration_ns + if is_dit_component_name(component_name) + else component_stage_duration_ns or request_duration_ns + ) + 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 + ) + # A measured stage bounds how much latency a placement can remove, + # but not how much an unmeasured mechanism can add. Clipping both + # directions made a repeatedly streamed DiT appear merely neutral + # once its transfer time exceeded the whole request; another small + # component gain could then select a several-times-slower layout. + if relative_savings_ns > 0 and latency_upper_bound_ns > 0: + relative_savings_ns = min(relative_savings_ns, latency_upper_bound_ns) + estimates[candidate.option_key()] = relative_savings_ns + for candidate in candidates: + if candidate.option_key() in estimates: + continue component_maximum_bytes = maximum_h2d_bytes[candidate.component_name] if component_maximum_bytes <= 0: estimates[candidate.option_key()] = 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 6a9e131a61e2..52d8ec6c91a3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -269,6 +269,44 @@ def test_dit_request_cap_does_not_compress_marginal_transfer_savings(self): savings[resident.option_key()] - savings[partial.option_key()] > 200_000_000 ) + def test_measured_placement_keeps_the_full_cost_of_a_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, + stage_duration_ns={"denoise": 1_000_000_000}, + component_stages={"transformer": ("denoise",)}, + ) + + 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 76b2ab16dca4bdf99695ea42b0c8b710725ad4f2 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 03:10:21 +0800 Subject: [PATCH 13/70] [diffusion] keep low-memory preload placement conservative --- .../runtime/managers/gpu_worker.py | 22 ++++++++++-- .../test/unit/test_auto_residency.py | 36 +++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 1c8489dd9706..a8f56c64eab2 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -115,6 +115,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 ( + resolve_keep_resident_min_available_gb, +) 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, @@ -1570,11 +1573,26 @@ def _build_pre_warmup_auto_residency_report( skip_reason="pipeline not initialized", ) + budget_bytes = self._auto_residency_budget_bytes() + admission_gib = resolve_keep_resident_min_available_gb(self.server_args) + budget_gib = budget_bytes / GIB_BYTES + if budget_gib < admission_gib: + return RankResidencyReport( + rank=self.rank, + budget_bytes=budget_bytes, + estimated_peak_bytes=None, + candidates=[], + skip_reason=( + f"{budget_gib:.1f} GiB VRAM budget is below the " + f"{admission_gib:.1f} GiB uncalibrated residency threshold" + ), + ) + candidates = self._collect_auto_residency_targets(workload) if not candidates: return RankResidencyReport( rank=self.rank, - budget_bytes=self._auto_residency_budget_bytes(), + budget_bytes=budget_bytes, estimated_peak_bytes=None, candidates=[], skip_reason="no eligible residency alternatives", @@ -1609,7 +1627,7 @@ def _build_pre_warmup_auto_residency_report( pin_budget = self.server_args.host_pin_budget() return RankResidencyReport( rank=self.rank, - budget_bytes=self._auto_residency_budget_bytes(), + budget_bytes=budget_bytes, estimated_peak_bytes=max(phase_peaks.values()), estimated_peak_bytes_by_phase=phase_peaks, active_components_by_phase=phase_active, 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 52d8ec6c91a3..badf4f624aa6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -1112,6 +1112,42 @@ def test_pre_warmup_report_enforces_weight_feasibility_without_oom(self): LAYERWISE_OFFLOAD ] + def test_pre_warmup_report_skips_uncalibrated_promotion_below_threshold( + self, monkeypatch + ): + from sglang.multimodal_gen.runtime.managers import gpu_worker + + worker = SimpleNamespace( + rank=0, + pipeline=object(), + server_args=object(), + _auto_residency_budget_bytes=lambda: 30 * GIB_BYTES, + _collect_auto_residency_targets=lambda _workload: pytest.fail( + "candidate collection should be skipped" + ), + ) + monkeypatch.setattr( + gpu_worker, "resolve_keep_resident_min_available_gb", lambda _args: 60.0 + ) + + report = gpu_worker.GPUWorker._build_pre_warmup_auto_residency_report( + worker, + workload=DefaultWorkload( + width=832, + height=480, + num_frames=24, + num_inference_steps=50, + ), + ) + + assert report.budget_bytes == 30 * GIB_BYTES + assert report.estimated_peak_bytes is None + assert report.candidates == [] + assert report.skip_reason == ( + "30.0 GiB VRAM budget is below the 60.0 GiB uncalibrated residency " + "threshold" + ) + def test_reserve_floor_is_capped_on_a_small_card(self): # Keep enough room for unmeasured activations without applying the # datacenter-card 4 GiB floor verbatim. From dcc26aef69d37eddfb076f16d735b5f8e287cbbb Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 04:03:53 +0800 Subject: [PATCH 14/70] [diffusion] tune layerwise residency schedule --- .../runtime/managers/gpu_worker.py | 5 + .../memory_managers/auto_residency.py | 489 ++++++++++++------ .../memory_managers/layerwise_offload.py | 95 +++- .../runtime/server_args/server_args.py | 18 +- .../test/unit/test_auto_residency.py | 126 ++++- .../test/unit/test_layerwise_offload.py | 41 ++ 6 files changed, 588 insertions(+), 186 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index a8f56c64eab2..86beb71ff1e4 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1557,6 +1557,11 @@ def _collect_auto_residency_targets( layerwise_tuning_of=lambda name, dit_group: ( self.server_args.layerwise_tuning_for(name, dit_group=dit_group) ), + layerwise_policy_is_explicit=lambda name, dit_group: ( + self.server_args.is_layerwise_residency_policy_explicit( + name, dit_group=dit_group + ) + ), pin_cpu_memory=self.server_args.pin_cpu_memory, ) 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 84c537e1c1ff..aeb6a51a294e 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 @@ -55,6 +55,8 @@ release_unused_pinned_memory, ) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( + RESIDENCY_POLICIES, + RESIDENCY_POLICY_STRIDED, is_dit_component_name, ) from sglang.multimodal_gen.runtime.managers.memory_managers.placement_budget import ( @@ -161,6 +163,7 @@ class ResidencyTarget(msgspec.Struct, frozen=True): # Layerwise candidates jointly choose stage-scoped GPU residency and host # pinning. None is used by ordinary component placement. target_layerwise_resident_layers: tuple[int, ...] | None = None + target_layerwise_residency_policies: tuple[str, ...] | None = None target_layerwise_pinned_layers: tuple[tuple[int, ...], ...] | None = None pinned_host_delta_bytes: int = 0 host_unpin_scratch_bytes: int = 0 @@ -216,12 +219,20 @@ def option_key(self) -> str: layer_counts = ",".join( str(count) for count in self.target_layerwise_resident_layers ) + policies = ( + ":policies=" + ",".join(self.target_layerwise_residency_policies) + if self.target_layerwise_residency_policies is not None + else "" + ) 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}" + return ( + f"{self.component_name}:{permanence}:layers={layer_counts}" + f"{policies}:pins={pinned}" + ) class DefaultWorkload(msgspec.Struct, frozen=True): @@ -298,6 +309,7 @@ class AppliedResidencyChange(msgspec.Struct, frozen=True): component_name: str residency_mode: str previous_layerwise_resident_layers: tuple[int, ...] | None = None + previous_layerwise_residency_policies: tuple[str, ...] | None = None previous_layerwise_pinned_layers: tuple[tuple[int, ...], ...] | None = None previous_layerwise_offload_enabled: bool = True previous_layerwise_configured: bool = True @@ -1004,10 +1016,23 @@ def component_resident_size_bytes(module: nn.Module, residency_mode: str) -> int return _module_weight_bytes(module) +def _resolve_layerwise_policies( + managers: Sequence, + residency_policies: tuple[str, ...] | None, +) -> tuple[str, ...]: + policies = residency_policies or tuple( + manager.residency_policy for manager in managers + ) + if len(policies) != len(managers): + raise ValueError("layerwise residency-policy group count changed") + return policies + + def _layerwise_transfer_work_bytes( *, managers: Sequence, resident_layers: tuple[int, ...], + residency_policies: tuple[str, ...] | None = None, pinned_layers: tuple[tuple[int, ...], ...], uses_per_streamed_layer: int, ) -> int: @@ -1017,15 +1042,16 @@ 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. """ + policies = _resolve_layerwise_policies(managers, residency_policies) total = 0 - for manager, resident_count, pinned_indices in zip( - managers, resident_layers, pinned_layers + for manager, resident_count, policy, pinned_indices in zip( + managers, resident_layers, policies, pinned_layers ): streamed = set( compute_streamed_layers( num_layers=manager.num_layers, resident_layers=resident_count, - policy=manager.residency_policy, + policy=policy, ) ) pinned = set(pinned_indices) @@ -1042,6 +1068,7 @@ def _layerwise_pin_targets( *, managers: Sequence, resident_layers: tuple[int, ...], + residency_policies: tuple[str, ...] | None = None, current_pinned_layers: tuple[tuple[int, ...], ...], uses_per_streamed_layer: int, ) -> list[tuple[tuple[int, ...], ...]]: @@ -1052,15 +1079,16 @@ def _layerwise_pin_targets( the layer count; the joint planner still evaluates those targets against rank VRAM, HostPin, and transition headroom together. """ + policies = _resolve_layerwise_policies(managers, residency_policies) ordered: list[tuple[int, int, int]] = [] - for manager_index, (manager, resident_count) in enumerate( - zip(managers, resident_layers) + for manager_index, (manager, resident_count, policy) in enumerate( + zip(managers, resident_layers, policies) ): streamed = set( compute_streamed_layers( num_layers=manager.num_layers, resident_layers=resident_count, - policy=manager.residency_policy, + policy=policy, ) ) if not manager.pin_cpu_memory: @@ -1113,6 +1141,37 @@ def _layerwise_resident_targets(managers: Sequence) -> list[tuple[int, ...]]: return sorted(targets) +def _layerwise_policy_targets( + *, + managers: Sequence, + resident_layers: tuple[int, ...], + tune_policy: bool, +) -> list[tuple[str, ...]]: + """Useful component-level policy choices for one resident-count target.""" + current = tuple(manager.residency_policy for manager in managers) + if not tune_policy or all( + count in (0, manager.num_layers) + for manager, count in zip(managers, resident_layers) + ): + return [current] + return list( + dict.fromkeys( + [current] + + [tuple(policy for _ in managers) for policy in RESIDENCY_POLICIES] + ) + ) + + +def _policy_affects_layerwise_layout( + managers: Sequence, + resident_layers: tuple[int, ...], +) -> bool: + return any( + 0 < count < manager.num_layers + for manager, count in zip(managers, resident_layers) + ) + + class _EstimatedLayerwiseManager: """Size-only layer group used before a real offload manager exists.""" @@ -1148,13 +1207,18 @@ def pinnable_layer_indices(self) -> tuple[int, ...]: return () return tuple(self._layer_bytes) - def resident_weight_bytes(self, resident_layers: int | None = None) -> int: + def resident_weight_bytes( + self, + resident_layers: int | None = None, + residency_policy: str | None = None, + ) -> int: count = self.resident_layers if resident_layers is None else resident_layers + policy = self.residency_policy if residency_policy is None else residency_policy streamed = set( compute_streamed_layers( num_layers=self.num_layers, resident_layers=count, - policy=self.residency_policy, + policy=policy, ) ) return sum( @@ -1164,15 +1228,18 @@ def resident_weight_bytes(self, resident_layers: int | None = None) -> int: ) def peak_managed_device_weight_bytes( - self, resident_layers: int | None = None + self, + resident_layers: int | None = None, + residency_policy: str | None = None, ) -> int: count = self.resident_layers if resident_layers is None else resident_layers + policy = self.residency_policy if residency_policy is None else residency_policy streamed = compute_streamed_layers( num_layers=self.num_layers, resident_layers=count, - policy=self.residency_policy, + policy=policy, ) - resident_bytes = self.resident_weight_bytes(count) + resident_bytes = self.resident_weight_bytes(count, policy) copy_window = min(self.prefetch_size, len(streamed)) streamed_window_bytes = sum( sorted( @@ -1224,6 +1291,7 @@ def _unconfigured_layerwise_targets( residency_policy: str, pin_cpu_memory: bool, allow_host_pin_reallocation: bool, + tune_residency_policy: bool, ) -> tuple[list[ResidencyTarget], int]: """Virtual layerwise frontier for a component still using coarse offload. @@ -1269,82 +1337,99 @@ def _unconfigured_layerwise_targets( ) 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) - ) - pin_targets = ( - _layerwise_pin_targets( - managers=managers, - resident_layers=resident_layers, - current_pinned_layers=empty_pins, - uses_per_streamed_layer=uses_per_request, - ) - if allow_host_pin_reallocation - else [empty_pins] - ) - for pinned_layers in pin_targets: - 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, pinned_layers) + for policies in _layerwise_policy_targets( + managers=managers, + resident_layers=resident_layers, + tune_policy=tune_residency_policy, + ): + resident_bytes = sum( + manager.resident_weight_bytes(count, policy) + for manager, count, policy in zip(managers, resident_layers, policies) ) - transfer_work = _layerwise_transfer_work_bytes( - managers=managers, - resident_layers=resident_layers, - pinned_layers=pinned_layers, - uses_per_streamed_layer=uses_per_request, + active_managed_bytes = sum( + manager.peak_managed_device_weight_bytes(count, policy) + for manager, count, policy in zip(managers, resident_layers, policies) ) - targets.append( - ResidencyTarget( - component_name=component_name, - residency_mode=current_mode, - target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=resident_bytes, - # 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=pinned_layers, - pinned_host_delta_bytes=pinned_bytes, - host_pin_scratch_bytes=pinned_bytes, - host_materialize_scratch_bytes=host_materialize_scratch, - device_transition_delta_bytes=( - unmanaged_weight_bytes - current_inactive_bytes - ), - active_device_delta_bytes=( - unmanaged_weight_bytes - + active_managed_bytes - - full_weight_bytes - ), - concurrent_prefetch_device_delta_bytes=( - unmanaged_weight_bytes - + active_managed_bytes - - (full_weight_bytes if current_resident else 0) - ), - 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, - target_pinned_host_bytes=pinned_bytes, + pin_targets = ( + _layerwise_pin_targets( + managers=managers, + resident_layers=resident_layers, + residency_policies=policies, + current_pinned_layers=empty_pins, + uses_per_streamed_layer=uses_per_request, ) + if allow_host_pin_reallocation + else [empty_pins] ) + for pinned_layers in pin_targets: + 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, pinned_layers) + ) + transfer_work = _layerwise_transfer_work_bytes( + managers=managers, + resident_layers=resident_layers, + residency_policies=policies, + pinned_layers=pinned_layers, + uses_per_streamed_layer=uses_per_request, + ) + targets.append( + ResidencyTarget( + component_name=component_name, + residency_mode=current_mode, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=resident_bytes, + # 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_residency_policies=( + policies + if tune_residency_policy + and _policy_affects_layerwise_layout( + managers, resident_layers + ) + else None + ), + target_layerwise_pinned_layers=pinned_layers, + pinned_host_delta_bytes=pinned_bytes, + host_pin_scratch_bytes=pinned_bytes, + host_materialize_scratch_bytes=host_materialize_scratch, + device_transition_delta_bytes=( + unmanaged_weight_bytes - current_inactive_bytes + ), + active_device_delta_bytes=( + unmanaged_weight_bytes + + active_managed_bytes + - full_weight_bytes + ), + concurrent_prefetch_device_delta_bytes=( + unmanaged_weight_bytes + + active_managed_bytes + - (full_weight_bytes if current_resident else 0) + ), + 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 + ), + target_pinned_host_bytes=pinned_bytes, + ) + ) return targets, maximum_transfer_work @@ -1380,6 +1465,7 @@ def collect_residency_targets( mixed_dtype_components: Iterable[str] = (), required_resident_components: Iterable[str] = (), layerwise_tuning_of: Callable[[str, bool], tuple[float, float, str]] | None = None, + layerwise_policy_is_explicit: Callable[[str, bool], bool] | None = None, pin_cpu_memory: bool = True, ) -> list[ResidencyTarget]: """Build complete target-state frontiers for auto-managed components. @@ -1414,6 +1500,16 @@ def collect_residency_targets( # residency is only an initial placement and remains tunable. continue + dit_group = ( + isinstance(module, LayerwiseOffloadableModuleMixin) + and module.layerwise_offload_dit_group_enabled + ) + tune_residency_policy = bool( + is_dit_component_name(name) + and layerwise_policy_is_explicit is not None + and not layerwise_policy_is_explicit(name, dit_group) + ) + has_layerwise_managers = isinstance( module, LayerwiseOffloadableModuleMixin ) and bool(module.layerwise_offload_managers) @@ -1451,7 +1547,7 @@ def collect_residency_targets( and name not in mixed_dtype_names ): prefetch_value, _, residency_policy = layerwise_tuning_of( - name, module.layerwise_offload_dit_group_enabled + name, dit_group ) virtual_targets, maximum_transfer_work = ( _unconfigured_layerwise_targets( @@ -1464,6 +1560,7 @@ def collect_residency_targets( residency_policy=residency_policy, pin_cpu_memory=pin_cpu_memory, allow_host_pin_reallocation=allow_host_pin_reallocation, + tune_residency_policy=tune_residency_policy, ) ) uses_per_request = ( @@ -1555,6 +1652,9 @@ def collect_residency_targets( 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_residency_policies = tuple( + manager.residency_policy for manager in managers + ) current_pinned_layers = tuple( manager.pinned_layer_indices() for manager in managers ) @@ -1620,81 +1720,105 @@ def collect_residency_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) - ) - target_peak_device_bytes = sum( - manager.peak_managed_device_weight_bytes(count) - for manager, count in zip(managers, target_resident_layers) - ) - 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: - target_pinned_bytes = sum( - sum( - manager.layer_host_store_bytes().get(layer_idx, 0) - for layer_idx in pinned_indices + for target_policies in _layerwise_policy_targets( + managers=managers, + resident_layers=target_resident_layers, + tune_policy=tune_residency_policy, + ): + target_resident_bytes = sum( + manager.resident_weight_bytes(count, policy) + for manager, count, policy in zip( + managers, target_resident_layers, target_policies ) - 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, + target_peak_device_bytes = sum( + manager.peak_managed_device_weight_bytes(count, policy) + for manager, count, policy in zip( + managers, target_resident_layers, target_policies + ) ) - candidates.append( - ResidencyTarget( - component_name=name, - residency_mode=frontier_mode, - target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=target_resident_bytes, - h2d_bytes_per_request=( - maximum_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, - device_transition_delta_bytes=( - -managed_weight_bytes if current_permanent else 0 - ), - active_device_delta_bytes=( - target_peak_device_bytes - current_peak_device_bytes - ), - concurrent_prefetch_device_delta_bytes=( - target_peak_device_bytes - current_peak_device_bytes - ), - 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=( - unmanaged_weight_bytes + target_resident_bytes - ), - target_pinned_host_bytes=target_pinned_bytes, + pin_targets = ( + _layerwise_pin_targets( + managers=managers, + resident_layers=target_resident_layers, + residency_policies=target_policies, + 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: + 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, + residency_policies=target_policies, + 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( + ResidencyTarget( + component_name=name, + residency_mode=frontier_mode, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=target_resident_bytes, + h2d_bytes_per_request=( + maximum_transfer_work - target_transfer_work + ), + target_layerwise_resident_layers=target_resident_layers, + target_layerwise_residency_policies=( + target_policies + if tune_residency_policy + and _policy_affects_layerwise_layout( + managers, target_resident_layers + ) + else None + ), + 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, + device_transition_delta_bytes=( + -managed_weight_bytes if current_permanent else 0 + ), + active_device_delta_bytes=( + target_peak_device_bytes - current_peak_device_bytes + ), + concurrent_prefetch_device_delta_bytes=( + target_peak_device_bytes - current_peak_device_bytes + ), + 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_policies == current_residency_policies + and target_pinned_layers == current_pinned_layers + ), + target_device_weight_bytes=( + unmanaged_weight_bytes + target_resident_bytes + ), + target_pinned_host_bytes=target_pinned_bytes, + ) + ) # Layerwise stores have one fixed dtype, while both coarse strategies # cast at every declared use. A mixed-dtype component may be tuned within @@ -2038,6 +2162,10 @@ def _consensus_candidates( targets = { candidate.target_layerwise_resident_layers for candidate in rank_candidates } + policy_targets = { + candidate.target_layerwise_residency_policies + for candidate in rank_candidates + } pinned_targets = { candidate.target_layerwise_pinned_layers for candidate in rank_candidates } @@ -2049,6 +2177,7 @@ def _consensus_candidates( if ( len(component_names) != 1 or len(targets) != 1 + or len(policy_targets) != 1 or len(pinned_targets) != 1 or len(permanent) != 1 or len(target_modes) != 1 @@ -2075,6 +2204,7 @@ def _consensus_candidates( candidate.h2d_bytes_per_request for candidate in rank_candidates ), target_layerwise_resident_layers=targets.pop(), + target_layerwise_residency_policies=policy_targets.pop(), target_layerwise_pinned_layers=pinned_targets.pop(), pinned_host_delta_bytes=max( candidate.pinned_host_delta_bytes for candidate in rank_candidates @@ -2120,6 +2250,20 @@ def _consensus_candidates( return merged +def _layerwise_policy_preference_cost(candidate: ResidencyTarget) -> int: + """Prefer the validated partial-DiT schedule only after latency ties. + + Same-placement H200 A/Bs on Flux and LTX2.3 showed about 3% lower E2E for + strided residency without changing output. This soft cost follows the + strategy-transition cost, so it cannot make layerwise beat a latency- + equivalent current resident/component-offload strategy. + """ + policies = candidate.target_layerwise_residency_policies + if policies is None: + return 0 + return int(any(policy != RESIDENCY_POLICY_STRIDED for policy in policies)) + + def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyPlan: """Turn gathered rank reports into one deterministic placement plan.""" if not reports: @@ -2417,6 +2561,7 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP or candidate.target_residency_mode == candidate.residency_mode else 1 ), + _layerwise_policy_preference_cost(candidate), candidate.target_device_weight_bytes or max(0, candidate.target_resident_weight_bytes), -estimated_latency_savings, @@ -2599,6 +2744,9 @@ def apply_residency_changes( if previous_configured else () ) + previous_policies = ( + module.layerwise_residency_policies() if previous_configured else () + ) previous_pinned = ( module.layerwise_pinned_layers() if previous_configured else () ) @@ -2617,6 +2765,7 @@ def apply_residency_changes( component_name=candidate.component_name, residency_mode=candidate.residency_mode, previous_layerwise_resident_layers=previous, + previous_layerwise_residency_policies=previous_policies, previous_layerwise_pinned_layers=previous_pinned, previous_layerwise_offload_enabled=enabled.pop(), previous_layerwise_configured=previous_configured, @@ -2742,8 +2891,19 @@ def apply_residency_changes( server_args.set_auto_residency_mode( candidate.component_name, LAYERWISE_OFFLOAD ) - module.set_layerwise_resident_layer_counts( - candidate.target_layerwise_resident_layers + target_resident_layers = candidate.target_layerwise_resident_layers + if target_resident_layers is None: + raise RuntimeError( + f"layerwise placement {candidate.option_key()!r} has no " + "resident-layer target" + ) + target_policies = ( + candidate.target_layerwise_residency_policies + or module.layerwise_residency_policies() + ) + module.set_layerwise_residency_layout( + target_resident_layers, + target_policies, ) continue if target_mode == RESIDENT: @@ -2890,8 +3050,13 @@ def rollback_residency_changes( adjustment.component_name, adjustment.previous_auto_residency_mode, ) - module.restore_layerwise_resident_layers( - adjustment.previous_layerwise_resident_layers + previous_policies = ( + adjustment.previous_layerwise_residency_policies + or module.layerwise_residency_policies() + ) + module.restore_layerwise_residency_layout( + adjustment.previous_layerwise_resident_layers, + previous_policies, ) if adjustment.previous_layerwise_pinned_layers is None: raise RuntimeError("lost previous layerwise host placement") @@ -2979,7 +3144,8 @@ def _format_residency_change(candidate: ResidencyTarget) -> str: return ( f"{candidate.component_name}: layerwise resident layers=" f"{candidate.target_layerwise_resident_layers}, pinned layers=" - f"{pin_counts}" + f"{pin_counts}, policies=" + f"{candidate.target_layerwise_residency_policies or 'unchanged'}" ) return ( f"{candidate.component_name}: {candidate.residency_mode} -> " f"{target_mode}" @@ -2992,7 +3158,8 @@ def _format_candidate_summary(candidate: ResidencyTarget) -> str: 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 ())}" + f"pins={tuple(len(indices) for indices in candidate.target_layerwise_pinned_layers or ())}, " + f"policies={candidate.target_layerwise_residency_policies or 'unchanged'}" ) return f"{candidate.component_name}({target_mode}{details})" 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 9b8f69db8722..c66190e54007 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 @@ -113,8 +113,9 @@ def compute_streamed_layers( ) -> tuple[int, ...]: """Which layer indices are streamed rather than held on the GPU. - Both policies stream the same *count* of layers, so they cost the same - memory and move the same bytes. They differ only in when those bytes move: + Both policies stream the same *count* of layers. Their byte footprints are + usually close, but can differ when layer sizes are nonuniform. Their main + difference is when those bytes move: ``leading`` keeps layers ``0..r-1`` and streams the tail. Every streamed layer sits next to another streamed layer, so the transfers @@ -122,8 +123,8 @@ def compute_streamed_layers( step, and each has exactly one layer of compute to hide behind. ``strided`` spreads the streamed layers evenly across the whole step, so - the same bytes move over ``n`` layers instead of ``n-r`` and - the peak concurrent traffic drops by ``n/(n-r)``. + traffic moves over ``n`` layers instead of ``n-r`` and the + peak concurrent traffic drops by ``n/(n-r)``. What that buys is contention, not bandwidth and not stalls. Profiling the two policies on an 8-GPU run shows the same HtoD volume to within 0.1%, the @@ -1582,20 +1583,27 @@ def set_pinned_layers(self, pinned_layers: Iterable[int]) -> tuple[int, ...]: def pin_budget(self) -> HostPinBudget: return self._pin_budget - def resident_weight_bytes(self, resident_layers: int | None = None) -> int: + def resident_weight_bytes( + self, + resident_layers: int | None = None, + residency_policy: str | 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 + policy = self.residency_policy if residency_policy is None else residency_policy streamed = compute_streamed_layers( num_layers=self.num_layers, resident_layers=count, - policy=self.residency_policy, + policy=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 + self, + resident_layers: int | None = None, + residency_policy: str | None = None, ) -> int: """Conservative managed-weight working set during this component's use. @@ -1606,10 +1614,11 @@ def peak_managed_device_weight_bytes( the measured peak already contains the current copy window. """ count = self.resident_layers if resident_layers is None else resident_layers + policy = self.residency_policy if residency_policy is None else residency_policy streamed = compute_streamed_layers( num_layers=self.num_layers, resident_layers=count, - policy=self.residency_policy, + policy=policy, ) layer_bytes = self.layer_weight_bytes() resident = set(range(self.num_layers)) - set(streamed) @@ -1623,23 +1632,31 @@ def peak_managed_device_weight_bytes( 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. + def set_residency_layout( + self, resident_layers: int, residency_policy: str + ) -> tuple[int, str]: + """Change the stage-scoped resident set 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: + if residency_policy not in RESIDENCY_POLICIES: + raise ValueError( + f"unknown residency policy {residency_policy!r}; expected one of " + f"{RESIDENCY_POLICIES}" + ) + previous = (self.resident_layers, self.residency_policy) + if (target, residency_policy) == previous: return previous self.release_all() self.resident_layers = target + self.residency_policy = residency_policy self._streamed_order = compute_streamed_layers( num_layers=self.num_layers, resident_layers=target, - policy=self.residency_policy, + policy=residency_policy, ) self._resident_set = frozenset(range(self.num_layers)) - set( self._streamed_order @@ -1647,6 +1664,11 @@ def set_resident_layers(self, resident_layers: int) -> int: self._residency_active = False return previous + def set_resident_layers(self, resident_layers: int) -> int: + """Change only the stage-scoped resident-layer count.""" + previous, _ = self.set_residency_layout(resident_layers, self.residency_policy) + return previous + def iter_cpu_weights(self): """Yield (name, tensor) pairs from consolidated CPU buffers. @@ -2144,21 +2166,48 @@ def set_layerwise_resident_layer_counts( self, resident_layers: Sequence[int] ) -> tuple[int, ...]: """Apply an exact resident-layer count to every managed layer group.""" + previous, _ = self.set_layerwise_residency_layout( + resident_layers, + self.layerwise_residency_policies(), + ) + return previous + + def layerwise_residency_policies(self) -> tuple[str, ...]: + return tuple( + manager.residency_policy for manager in self.layerwise_offload_managers + ) + + def set_layerwise_residency_layout( + self, + resident_layers: Sequence[int], + residency_policies: Sequence[str], + ) -> tuple[tuple[int, ...], tuple[str, ...]]: + """Apply exact resident counts and policies to every managed group.""" if len(resident_layers) != len(self.layerwise_offload_managers): raise ValueError("layerwise resident-layer group count changed") + if len(residency_policies) != len(self.layerwise_offload_managers): + raise ValueError("layerwise residency-policy group count changed") previous = tuple( - manager.resident_layers for manager in self.layerwise_offload_managers + (manager.resident_layers, manager.residency_policy) + 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) + for manager, count, policy in zip( + self.layerwise_offload_managers, + resident_layers, + residency_policies, + ): + manager.set_residency_layout(count, policy) updated.append(manager) except Exception: - for manager, count in zip(updated, previous): - manager.set_resident_layers(count) + for manager, (count, policy) in zip(updated, previous): + manager.set_residency_layout(count, policy) raise - return previous + return ( + tuple(count for count, _ in previous), + tuple(policy for _, policy in previous), + ) def layerwise_pinned_host_bytes(self) -> int: return sum( @@ -2219,6 +2268,14 @@ def restore_layerwise_resident_layers(self, resident_layers: Sequence[int]) -> N raise ValueError("layerwise resident-layer group count changed") self.set_layerwise_resident_layer_counts(resident_layers) + def restore_layerwise_residency_layout( + self, + resident_layers: Sequence[int], + residency_policies: Sequence[str], + ) -> None: + """Restore exact per-group counts and policies after a failed adjustment.""" + self.set_layerwise_residency_layout(resident_layers, residency_policies) + def disable_offload(self) -> None: """Disable layerwise offload: load all layers to GPU and remove 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 6848a4e87bec..1043f87c418e 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -1113,6 +1113,18 @@ def _pick(mapping: dict[str, str], group_default, aux_default): ) return prefetch, resident, policy + def is_layerwise_residency_policy_explicit( + self, component_name: str, *, dit_group: bool + ) -> bool: + """Whether auto placement must preserve this component's policy.""" + policy_map = self._parse_component_value_map( + self.layerwise_residency_policy, + option="--layerwise-residency-policy", + ) + return component_name in policy_map or ( + dit_group and self.is_arg_explicitly_set("dit_layerwise_residency_policy") + ) + @staticmethod def _parse_component_attention_backend_map( value: dict[str, str] | str | None, @@ -2480,9 +2492,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: help="Which layers --dit-layerwise-resident-layers keeps resident. " "'leading' (default) keeps the first N, which crams the whole " "weight stream into the tail of each step. 'strided' spreads the " - "resident layers evenly over the stack so the same bytes move over " - "the whole step instead: same VRAM, same bytes, only a different " - "schedule. Worth trying when weight streaming overlaps " + "resident layers evenly over the stack so traffic moves over the " + "whole step instead. The layer count is unchanged; exact bytes can " + "differ when layer sizes are nonuniform. Worth trying when streaming overlaps " "memory-bound compute -- the transfers stop competing with it for " "L2 and DRAM bandwidth, which is where the gain comes from.", ) 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 badf4f624aa6..e2f9c117fc2e 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -48,6 +48,10 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import ( + RESIDENCY_POLICY_LEADING, + RESIDENCY_POLICY_STRIDED, +) from sglang.multimodal_gen.runtime.warmup_request_builder import ( SERVER_WARMUP_MAX_VIDEO_FRAMES, _resolve_auto_residency_warmup_shape, @@ -856,6 +860,42 @@ def test_does_not_replace_explicit_default_shape(self): class TestPlanAutoResidency: + def test_latency_equivalent_partial_dit_prefers_strided_schedule(self): + common = dict( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=4 * GIB_BYTES, + h2d_bytes_per_request=10 * GIB_BYTES, + target_layerwise_resident_layers=(8,), + target_layerwise_pinned_layers=((),), + target_device_weight_bytes=4 * GIB_BYTES, + ) + leading = ResidencyTarget( + **common, + target_layerwise_residency_policies=(RESIDENCY_POLICY_LEADING,), + current_placement=True, + ) + strided = ResidencyTarget( + **common, + target_layerwise_residency_policies=(RESIDENCY_POLICY_STRIDED,), + ) + + plan = plan_auto_residency( + reports=[ + _report( + estimated_gib=30, + candidates=[leading, strided], + target_workload_measured=True, + ) + ] + ) + + assert len(plan.changes) == 1 + assert plan.changes[0].target_layerwise_residency_policies == ( + RESIDENCY_POLICY_STRIDED, + ) + def test_lower_latency_utility_keeps_the_measured_placement(self): resident = ResidencyTarget( component_name="transformer", @@ -2380,7 +2420,8 @@ def offloaded_weight_bytes(self): tensor.numel() * tensor.element_size() for tensor in self._tensors.values() ) - def resident_weight_bytes(self, resident_layers=None): + def resident_weight_bytes(self, resident_layers=None, residency_policy=None): + del residency_policy count = self.resident_layers if resident_layers is None else resident_layers return ( self.offloaded_weight_bytes() @@ -2388,7 +2429,10 @@ def resident_weight_bytes(self, resident_layers=None): // self.num_layers ) - def peak_managed_device_weight_bytes(self, resident_layers=None): + def peak_managed_device_weight_bytes( + self, resident_layers=None, residency_policy=None + ): + del residency_policy count = self.resident_layers if resident_layers is None else resident_layers return ( self.offloaded_weight_bytes() @@ -2427,8 +2471,13 @@ def set_pinned_layers(self, pinned_layers): return previous def set_resident_layers(self, resident_layers): - previous = self.resident_layers + previous, _ = self.set_residency_layout(resident_layers, self.residency_policy) + return previous + + def set_residency_layout(self, resident_layers, residency_policy): + previous = (self.resident_layers, self.residency_policy) self.resident_layers = min(max(0, resident_layers), self.num_layers) + self.residency_policy = residency_policy return previous def load_all_layers(self): @@ -2665,6 +2714,45 @@ def test_layerwise_dit_can_release_existing_resident_layers(self): assert target.active_device_delta_bytes < 0 assert target.h2d_bytes_per_request > 0 + def test_auto_dit_frontier_tunes_policy_unless_explicit(self): + manager = _FakeLayerwiseManager( + {f"layers.{index}.w": torch.zeros(16) for index in range(4)}, + resident_layers=2, + ) + module = _FakeLayerwiseDit([manager]) + common = dict( + 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, + ) + + candidates = collect_residency_targets( + **common, + layerwise_policy_is_explicit=lambda _name, _dit_group: False, + ) + policies = { + candidate.target_layerwise_residency_policies + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + and candidate.target_layerwise_resident_layers == (2,) + and candidate.target_layerwise_pinned_layers == ((),) + } + assert policies == { + (RESIDENCY_POLICY_LEADING,), + (RESIDENCY_POLICY_STRIDED,), + } + + explicit_candidates = collect_residency_targets( + **common, + layerwise_policy_is_explicit=lambda _name, _dit_group: True, + ) + assert all( + candidate.target_layerwise_residency_policies is None + for candidate in explicit_candidates + ) + def test_fully_pinned_component_can_release_hostpin_for_a_hotter_component(self): manager = _FakeLayerwiseManager( { @@ -3347,6 +3435,38 @@ def test_partial_layerwise_promotion_restores_exact_group_counts(self): ) assert [manager.resident_layers for manager in managers] == [1, 0] + def test_partial_layerwise_policy_change_rolls_back_exact_layout(self): + manager = _FakeLayerwiseManager( + {f"layers.{index}.w": torch.zeros(16) for index in range(4)}, + resident_layers=2, + ) + module = _FakeLayerwiseDit([manager]) + args = _StubResidencyArgs() + candidate = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=manager.resident_weight_bytes(2), + h2d_bytes_per_request=GIB_BYTES, + target_layerwise_resident_layers=(2,), + target_layerwise_residency_policies=(RESIDENCY_POLICY_STRIDED,), + target_layerwise_pinned_layers=((),), + ) + + applied = apply_residency_changes( + plan=_plan_for([candidate]), + modules={"transformer": module}, + server_args=args, + ) + assert manager.residency_policy == RESIDENCY_POLICY_STRIDED + + rollback_residency_changes( + applied=applied, + modules={"transformer": module}, + server_args=args, + ) + assert manager.resident_layers == 2 + assert manager.residency_policy == RESIDENCY_POLICY_LEADING + def test_permanent_layerwise_can_be_demoted_and_rollback_restores_it(self): manager = _FakeLayerwiseManager( { 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 a8f0b00fe396..c90a8dd92826 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -202,6 +202,9 @@ class _TestServerArgs(SimpleNamespace): ) _parse_component_value_map = staticmethod(ServerArgs._parse_component_value_map) layerwise_tuning_for = ServerArgs.layerwise_tuning_for + is_layerwise_residency_policy_explicit = ( + ServerArgs.is_layerwise_residency_policy_explicit + ) def _server_args(**kwargs): @@ -1043,6 +1046,44 @@ def test_configure_resolves_residency_policy(monkeypatch): ) +def test_explicit_residency_policy_is_scoped_to_the_selected_component(): + args = _server_args( + _explicit_arg_names={"dit_layerwise_residency_policy"}, + dit_layerwise_residency_policy=RESIDENCY_POLICY_STRIDED, + layerwise_residency_policy={ + "text_encoder": RESIDENCY_POLICY_STRIDED, + }, + ) + + assert args.is_layerwise_residency_policy_explicit("transformer", dit_group=True) + assert args.is_layerwise_residency_policy_explicit("text_encoder", dit_group=False) + assert not args.is_layerwise_residency_policy_explicit("vae", dit_group=False) + + +def test_residency_layout_switches_policy_without_rebuilding_host_stores(monkeypatch): + _patch_fake_device(monkeypatch) + manager = _resident_manager( + _MultiBlockModel(8), + num_layers=8, + prefetch_size=1, + resident_layers=2, + residency_policy=RESIDENCY_POLICY_LEADING, + ) + host_stores = manager._consolidated_cpu_weights + + previous = manager.set_residency_layout(3, RESIDENCY_POLICY_STRIDED) + + assert previous == (2, RESIDENCY_POLICY_LEADING) + assert manager.resident_layers == 3 + assert manager.residency_policy == RESIDENCY_POLICY_STRIDED + assert manager._streamed_order == compute_streamed_layers( + num_layers=8, + resident_layers=3, + policy=RESIDENCY_POLICY_STRIDED, + ) + assert manager._consolidated_cpu_weights is host_stores + + def test_configure_logs_component_start_and_completion(monkeypatch): _patch_fake_device(monkeypatch) logs = [] From 60f5f69a4673e3ca415d0bff9c6e7572aadcb564 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 05:00:37 +0800 Subject: [PATCH 15/70] [diffusion] align residency performance guards --- .../multimodal_gen/test/server/perf_baselines/h100.json | 4 ++-- .../multimodal_gen/test/unit/test_performance_metrics.py | 2 ++ 2 files changed, 4 insertions(+), 2 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 5fde9f15c888..fc436ed69adc 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -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": 48504.0, + "runtime_peak_vram_mb": 50112.0, "estimated_full_test_time_s": 77.0 }, "wan2_2_i2v_a14b_2gpu": { @@ -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": 58158.0, + "runtime_peak_vram_mb": 61652.0, "estimated_full_test_time_s": 248.9 }, "wan2_2_t2v_a14b_2gpu": { 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 05ba31074529..d9c5045b6525 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -267,6 +267,7 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): component_residency_requirement=lambda _name: None, auto_residency_mode=lambda _name: None, layerwise_tuning_for=lambda _name, *, dit_group: (0.0, 0.0, "leading"), + is_layerwise_residency_policy_explicit=lambda _name, *, dit_group: False, pin_cpu_memory=True, host_pin_budget=lambda: host_pin_budget, ) @@ -323,6 +324,7 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): component_residency_requirement=lambda _name: None, auto_residency_mode=lambda _name: None, layerwise_tuning_for=lambda _name, *, dit_group: (0.0, 0.0, "leading"), + is_layerwise_residency_policy_explicit=lambda _name, *, dit_group: False, pin_cpu_memory=True, host_pin_budget=lambda: host_pin_budget, ) From 9fe11c208e0125727a525ae0a88adef739afeb4b Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 06:21:32 +0800 Subject: [PATCH 16/70] fix(diffusion): validate automatic residency placement --- .../runtime/entrypoints/control_requests.py | 2 +- .../runtime/managers/gpu_worker.py | 58 ++++++-- .../memory_managers/auto_residency.py | 49 ++++++- .../runtime/managers/scheduler.py | 5 +- .../multimodal_gen/runtime/server_warmup.py | 34 ++++- .../test/unit/test_auto_residency.py | 107 +++++++++++++- .../test/unit/test_auto_residency_warmup.py | 132 ++++++++++++++++-- .../test/unit/test_dp_routing.py | 6 + 8 files changed, 368 insertions(+), 25 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py b/python/sglang/multimodal_gen/runtime/entrypoints/control_requests.py index 51843e2274e7..e41cc87f74bd 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" | "apply_static" | "rollback" + action: str = "apply" # "apply" | "apply_static" | "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 86beb71ff1e4..fcfb51bf1625 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, AutoResidencyPlan, @@ -80,6 +81,7 @@ measured_failed_workload_phase_peaks, plan_auto_residency, plan_summary_payload, + pre_warmup_residency_targets, rank_candidates_by_h2d_savings, resolve_default_workload, resolve_measured_default_workload, @@ -1265,20 +1267,25 @@ def get_can_stay_resident_components( remaining_gpu_mem_gb -= usage_gb return can_stay_resident - def apply_auto_residency(self, *, pre_warmup: bool = False) -> OutputBatch: + def apply_auto_residency( + self, *, pre_warmup: bool = False, 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: an uncaught raise there would leave the peer ranks parked in the collective until the group timeout. """ + if pre_warmup and validate_only: + raise ValueError("static placement cannot be validation-only") records = list(self._auto_residency_warmup_records) # Each round must describe one placement. Intersecting phase ownership # across old and newly adjusted layouts would double-count weights that @@ -1314,6 +1321,26 @@ def apply_auto_residency(self, *, pre_warmup: bool = False) -> OutputBatch: ) reports = self._auto_residency_all_gather(local_report) recovering_from_oom = any(report.warmup_oom for report in reports) + 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 ( not pre_warmup and self._auto_residency_round_sizes @@ -1344,7 +1371,6 @@ def apply_auto_residency(self, *, pre_warmup: bool = False) -> OutputBatch: latest_round_only=True, ) plan = plan_auto_residency(reports=reports) - summary = format_plan_summary(plan=plan, workload=workload, records=records) if ( not pre_warmup and self._auto_residency_round_sizes @@ -1367,6 +1393,18 @@ def apply_auto_residency(self, *, pre_warmup: bool = False) -> OutputBatch: latest_round_only=True, retryable=True, ) + if validate_only: + if self.is_output_rank: + logger.info( + "Auto residency: post-adjustment calibration passed; " + "keeping the selected placement." + ) + 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) @@ -1593,7 +1631,11 @@ def _build_pre_warmup_auto_residency_report( ), ) - candidates = self._collect_auto_residency_targets(workload) + all_candidates = self._collect_auto_residency_targets(workload) + candidates = pre_warmup_residency_targets( + all_candidates, + excluded_components=(self.pipeline.preload_residency_excluded_components), + ) if not candidates: return RankResidencyReport( rank=self.rank, @@ -1607,14 +1649,14 @@ def _build_pre_warmup_auto_residency_report( current_device_weights = component_current_device_weight_bytes(modules) permanent_components = { candidate.component_name - for candidate in candidates + for candidate in all_candidates if candidate.current_placement and candidate.permanent_residency } baseline_allocated = int(torch.get_device_module().memory_allocated()) phase_peaks = {} phase_active = {} phase_used = {} - for component_name in sorted({c.component_name for c in candidates}): + for component_name in sorted(active_weights): phase_name = f"static:{component_name}" phase_peaks[phase_name] = baseline_allocated + ( 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 aeb6a51a294e..9de5f6e41189 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 @@ -27,7 +27,7 @@ import statistics from itertools import chain -from typing import TYPE_CHECKING, Callable, Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Callable, Collection, Iterable, Mapping, Sequence import msgspec import torch @@ -110,6 +110,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" @@ -308,6 +309,10 @@ class AutoResidencyPlan(msgspec.Struct, frozen=True): class AppliedResidencyChange(msgspec.Struct, frozen=True): component_name: str residency_mode: str + # Keep the exact object modified by the transaction. Lazy stage-owned + # components are not guaranteed to remain discoverable by name after a + # failed request, but rollback must still restore that same module. + module_ref: object | None = None previous_layerwise_resident_layers: tuple[int, ...] | None = None previous_layerwise_residency_policies: tuple[str, ...] | None = None previous_layerwise_pinned_layers: tuple[tuple[int, ...], ...] | None = None @@ -329,6 +334,34 @@ def residency_device_growth_bytes(candidate: ResidencyTarget) -> int: ) +def pre_warmup_residency_targets( + candidates: Iterable[ResidencyTarget], + *, + excluded_components: Collection[str] = (), +) -> list[ResidencyTarget]: + """Keep only weight-safe choices before activation memory is measured. + + Static planning knows complete component weights but not the runtime + working set. It may demote an unsafe configured layout or redistribute + HostPin at the same device footprint, but promoting more weights here can + consume activation headroom that only the first shape probe can reveal. + """ + candidates = list(candidates) + current_device_bytes = { + candidate.component_name: candidate.target_device_weight_bytes + for candidate in candidates + if candidate.current_placement + } + return [ + candidate + for candidate in candidates + if candidate.component_name not in excluded_components + and candidate.component_name in current_device_bytes + and candidate.target_device_weight_bytes + <= current_device_bytes[candidate.component_name] + ] + + def resolve_default_workload(server_args: ServerArgs) -> DefaultWorkload: """Resolve the default request shape the planner is optimized for.""" from sglang.multimodal_gen.runtime.warmup_request_builder import ( @@ -2764,6 +2797,7 @@ def apply_residency_changes( snapshots[candidate.component_name] = AppliedResidencyChange( component_name=candidate.component_name, residency_mode=candidate.residency_mode, + module_ref=module, previous_layerwise_resident_layers=previous, previous_layerwise_residency_policies=previous_policies, previous_layerwise_pinned_layers=previous_pinned, @@ -2787,6 +2821,7 @@ def apply_residency_changes( snapshots[candidate.component_name] = AppliedResidencyChange( component_name=candidate.component_name, residency_mode=candidate.residency_mode, + module_ref=module, previous_auto_residency_mode=previous_auto_mode, previous_host_pin_budget_state=previous_host_pin_budget_state, applied_device_delta_bytes=applied_device_delta_bytes, @@ -2972,7 +3007,11 @@ def rollback_residency_changes( ] for adjustment in pinning_changes: try: - module = modules.get(adjustment.component_name) + module = ( + adjustment.module_ref + if adjustment.module_ref is not None + else modules.get(adjustment.component_name) + ) if not isinstance(module, LayerwiseOffloadableModuleMixin): raise RuntimeError("lost layerwise offload capability") if not module.layerwise_offload_managers: @@ -3003,7 +3042,11 @@ def rollback_residency_changes( ): torch.get_device_module().empty_cache() try: - module = modules.get(adjustment.component_name) + module = ( + adjustment.module_ref + if adjustment.module_ref is not None + else 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") diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 49dcba5e338a..3b7e137711a0 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -209,7 +209,10 @@ 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(pre_warmup=req.action == "apply_static") + return self.worker.apply_auto_residency( + pre_warmup=req.action == "apply_static", + 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 5a2d6c150869..d510c0793f94 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -171,6 +171,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) @@ -205,9 +206,11 @@ 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 even when the initial synthetic + # warmup itself was implicit. 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): @@ -244,7 +247,7 @@ async def rollback_and_rewarm(error: Exception) -> None: ) if status == PLACEMENT_STATUS_ROLLED_BACK: await run_async_client_warmup( - server_args, forward, fail_open=fail_open, rewarm=True + server_args, forward, fail_open=False, rewarm=True ) if retryable: logger.warning( @@ -290,6 +293,31 @@ async def rollback_and_rewarm(error: Exception) -> None: continue 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" + ) + ) + return return if recovering_from_oom: 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 e2f9c117fc2e..6f1c83570066 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -33,6 +33,7 @@ format_applied_changes, measured_failed_workload_phase_peaks, plan_auto_residency, + pre_warmup_residency_targets, rank_candidates_by_h2d_savings, resolve_measured_default_workload, rollback_residency_changes, @@ -1152,6 +1153,27 @@ def test_pre_warmup_report_enforces_weight_feasibility_without_oom(self): LAYERWISE_OFFLOAD ] + def test_pre_warmup_targets_defer_device_weight_promotion(self): + current = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=8 * GIB_BYTES, + target_device_weight_bytes=0, + current_placement=True, + ) + resident = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=8 * GIB_BYTES, + h2d_bytes_per_request=0, + target_device_weight_bytes=8 * GIB_BYTES, + ) + + assert pre_warmup_residency_targets([current, resident]) == [current] + def test_pre_warmup_report_skips_uncalibrated_promotion_below_threshold( self, monkeypatch ): @@ -1188,6 +1210,87 @@ def test_pre_warmup_report_skips_uncalibrated_promotion_below_threshold( "threshold" ) + def test_pre_warmup_report_keeps_excluded_components_fixed(self, monkeypatch): + from sglang.multimodal_gen.runtime.managers import gpu_worker + + transformer = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=20 * GIB_BYTES, + permanent_residency=True, + ) + vae = ResidencyTarget( + component_name="vae", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=GIB_BYTES, + permanent_residency=True, + target_device_weight_bytes=GIB_BYTES, + current_placement=True, + ) + modules = {"transformer": object(), "vae": object()} + worker = SimpleNamespace( + rank=0, + pipeline=SimpleNamespace( + preload_residency_excluded_components=frozenset(("transformer",)) + ), + server_args=SimpleNamespace( + num_gpus=1, + nnodes=1, + node_rank=0, + host_pin_budget=lambda: None, + ), + _auto_residency_budget_bytes=lambda: 80 * GIB_BYTES, + _collect_auto_residency_targets=lambda _workload: [transformer, vae], + _auto_residency_modules=lambda: modules, + ) + monkeypatch.setattr( + gpu_worker, "resolve_keep_resident_min_available_gb", lambda _args: 0 + ) + monkeypatch.setattr( + gpu_worker, + "component_runtime_weight_bytes", + lambda _modules: {"transformer": 20 * GIB_BYTES, "vae": GIB_BYTES}, + ) + monkeypatch.setattr( + gpu_worker, + "component_current_device_weight_bytes", + lambda _modules: {"transformer": 0, "vae": 0}, + ) + monkeypatch.setattr( + gpu_worker, "layerwise_pinned_host_bytes", lambda *a, **k: 0 + ) + monkeypatch.setattr( + gpu_worker, "layerwise_host_pin_capacity_bytes", lambda *a, **k: 0 + ) + monkeypatch.setattr( + gpu_worker, "host_memory_available_bytes", lambda: 100 * GIB_BYTES + ) + monkeypatch.setattr( + gpu_worker.torch, + "get_device_module", + lambda: SimpleNamespace(memory_allocated=lambda: 2 * GIB_BYTES), + ) + + report = gpu_worker.GPUWorker._build_pre_warmup_auto_residency_report( + worker, + workload=DefaultWorkload( + width=1920, + height=1088, + num_frames=24, + num_inference_steps=15, + ), + ) + + assert [candidate.component_name for candidate in report.candidates] == ["vae"] + assert report.estimated_peak_bytes_by_phase["static:transformer"] == ( + 22 * GIB_BYTES + ) + assert report.used_components_by_phase["static:transformer"] == ("transformer",) + def test_reserve_floor_is_capped_on_a_small_card(self): # Keep enough room for unmeasured activations without applying the # datacenter-card 4 GiB floor verbatim. @@ -3320,7 +3423,9 @@ def test_layerwise_can_become_component_offload_and_rollback(self): rollback_residency_changes( applied=applied, - modules={"text_encoder": module}, + # Lazy stage-owned components can disappear from the manager's + # per-request name index before validation rollback. + modules={}, server_args=args, ) assert len(module.layerwise_offload_managers) == 1 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 973c1ac4291f..4ff27015612d 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,10 +15,12 @@ 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, @@ -211,7 +213,7 @@ def test_worker_rolls_back_a_materially_slower_calibrated_round(self): 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() @@ -220,6 +222,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_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() + 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, + ) + plan = AutoResidencyPlan(changes=[candidate]) + + 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=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) + apply.assert_not_called() + def test_new_placement_that_spends_the_reserve_rolls_back(self): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 @@ -303,7 +350,7 @@ async def forward(req): rewarm.assert_awaited_once_with( server_args, forward, - fail_open=True, + fail_open=False, rewarm=True, ) @@ -343,10 +390,37 @@ async def forward(req): rewarm.assert_awaited_once_with( server_args, forward, - fail_open=True, + fail_open=False, rewarm=True, ) + def test_failed_restored_placement_warmup_aborts_startup(self): + responses = iter( + [ + OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), + OutputBatch(output={"status": PLACEMENT_STATUS_ROLLED_BACK}), + ] + ) + + async def forward(_req): + return next(responses) + + rewarm = mock.AsyncMock( + side_effect=[RuntimeError("candidate OOM"), RuntimeError("restore 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, "restore OOM"), + ): + asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) + def test_oom_recovery_keeps_demotion_and_replans(self): responses = iter( [ @@ -386,13 +460,18 @@ async def forward(req): rewarm=True, ) - def test_stops_after_one_successful_calibrated_adjustment(self): + def test_successful_adjustment_is_validated_without_another_warmup(self): actions = [] async def forward(req): self.assertIsInstance(req, AutoResidencyReq) actions.append(req.action) - return OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}) + status = ( + PLACEMENT_STATUS_ADJUSTED + if req.action == "apply" + else PLACEMENT_STATUS_VALIDATED + ) + return OutputBatch(output={"status": status}) rewarm = mock.AsyncMock() server_args = SimpleNamespace( @@ -407,9 +486,41 @@ async def forward(req): ): asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) - self.assertEqual(actions, ["apply"]) + self.assertEqual(actions, ["apply", "validate"]) rewarm.assert_awaited_once() + def test_regressed_validation_rolls_back_and_rewarms_original_placement(self): + responses = iter( + [ + OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), + OutputBatch( + error="request duration regressed", + output={"status": PLACEMENT_STATUS_ROLLED_BACK}, + ), + ] + ) + actions = [] + + async def forward(req): + 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", "validate"]) + self.assertEqual(rewarm.await_count, 2) + def test_failed_first_calibration_rolls_back_and_rewarms(self): actions = [] @@ -469,7 +580,12 @@ def test_successful_calibration_does_not_start_another_apply_round(self): async def forward(req): self.assertIsInstance(req, AutoResidencyReq) actions.append(req.action) - return OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}) + status = ( + PLACEMENT_STATUS_ADJUSTED + if req.action == "apply" + else PLACEMENT_STATUS_VALIDATED + ) + return OutputBatch(output={"status": status}) rewarm = mock.AsyncMock() server_args = SimpleNamespace( @@ -484,7 +600,7 @@ async def forward(req): ): asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) - self.assertEqual(actions, ["apply"]) + self.assertEqual(actions, ["apply", "validate"]) self.assertEqual(rewarm.await_count, 1) def test_failed_later_round_keeps_earlier_calibrated_promotions(self): 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 fc3438c54dfb..bcb4461e2054 100644 --- a/python/sglang/multimodal_gen/test/unit/test_dp_routing.py +++ b/python/sglang/multimodal_gen/test/unit/test_dp_routing.py @@ -95,6 +95,12 @@ def test_auto_residency_fanout_reports_successful_rollback(): assert _merge_auto_residency_results([skipped, restored]) is restored +def test_auto_residency_validation_never_hides_a_replica_rollback(): + validated = OutputBatch(output={"status": "validated"}) + restored = OutputBatch(output={"status": "rolled_back"}) + assert _merge_auto_residency_results([validated, restored]) is restored + + def test_scheduler_endpoints_one_per_replica(): args = ServerArgs.__new__(ServerArgs) args.host = "localhost" From 8016510c7900de777ef09232a38ab27d786a86a9 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 06:52:02 +0800 Subject: [PATCH 17/70] 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 9de5f6e41189..738f9a4cc232 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 @@ -2693,10 +2693,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 @@ -2943,6 +2943,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 038420365f54..b594330ad9a6 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 6f1c83570066..9ec39dade07c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -53,6 +53,7 @@ RESIDENCY_POLICY_LEADING, RESIDENCY_POLICY_STRIDED, ) +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, @@ -3432,8 +3433,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)]), @@ -3441,12 +3451,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 ba5e04e6442c..a981b1e4212f 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 @@ -454,9 +454,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 e67255f25ad8019636d19f90372ff48a4b916af5 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 07:26:03 +0800 Subject: [PATCH 18/70] test(diffusion): align residency consistency guards --- .../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 0a0d396bfc46136a8467353bc0fbcb5923bb18a2 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 07:51:30 +0800 Subject: [PATCH 19/70] fix(diffusion): preload model-default layerwise DiTs --- .../managers/memory_managers/initial_residency.py | 13 ++++--------- .../test/unit/test_initial_residency.py | 4 ++-- 2 files changed, 6 insertions(+), 11 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py index fc064dc11e61..a00764847f39 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py @@ -7,7 +7,6 @@ from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( - LAYERWISE_OFFLOAD, RESIDENT, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_weight_inventory import ( @@ -61,10 +60,10 @@ def choose_initial_resident_components( """Choose resident startup components without removing runtime options. This is a load-feasibility seed, not a second serving-placement planner. - Native DiTs using the coarse residency mechanism can still be demoted to - component offload by warmup calibration. Auxiliary modules, layerwise - components, unknown weights, and explicit choices retain their configured - loading semantics. + Native DiTs can still be demoted by warmup calibration. Auxiliary modules, + unknown weights, pipeline exclusions, and explicit choices retain their + configured loading semantics. A model-default layerwise placement remains + eligible because avoiding that initialization is the point of this seed. """ reserve_bytes = max( MIN_INITIAL_RESIDENCY_RESERVE_BYTES, @@ -95,10 +94,6 @@ def choose_initial_resident_components( or item.component_name in excluded_components or not is_dit_component_name(item.component_name) or server_args.residency_mode(item.component_name) == RESIDENT - or ( - server_args.configured_residency_mode(item.component_name) - == LAYERWISE_OFFLOAD - ) or server_args.explicit_residency_mode(item.component_name) is not None ): continue diff --git a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py index dcb8ea0668f7..3ec4d05f8f2b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py @@ -114,7 +114,7 @@ def test_initial_seed_does_not_override_explicit_or_auxiliary_components(): assert selected == set() -def test_initial_seed_keeps_layerwise_dit_on_its_configured_load_path(): +def test_initial_seed_can_bypass_model_default_layerwise_initialization(): args = _Args(modes={"transformer": LAYERWISE_OFFLOAD}) selected = choose_initial_resident_components( @@ -124,7 +124,7 @@ def test_initial_seed_keeps_layerwise_dit_on_its_configured_load_path(): denoising_steps=8, ) - assert selected == set() + assert selected == {"transformer"} def test_initial_seed_honors_pipeline_load_placement_exclusions(): From 3ceeba3081079e35b1de2f16ac5c3510d8ca4345 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 08:46:12 +0800 Subject: [PATCH 20/70] test(diffusion): align LTX residency VRAM guard --- .../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 fc436ed69adc..8aeed65e2bc9 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": 68712.0, "estimated_full_test_time_s": 363.2 }, "qwen_image_t2i_cache_dit_scm_config_diffusers_1gpu": { From 31f8138a0e68bb0432e737e2d6fb358ae7f56f77 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 09:31:31 +0800 Subject: [PATCH 21/70] [diffusion] refine auto residency admission --- docs/docs/sglang-diffusion/api/cli.mdx | 2 +- .../sglang-diffusion/deployment_cookbook.mdx | 4 +- .../runtime/managers/gpu_worker.py | 40 ++++++------ .../memory_managers/initial_residency.py | 20 ++---- .../pipelines_core/composed_pipeline_base.py | 21 +++++-- .../test/unit/test_auto_residency.py | 61 ++++++++++++++----- .../test/unit/test_auto_residency_warmup.py | 52 +++++++++++++++- .../test/unit/test_initial_residency.py | 40 ++++++------ 8 files changed, 166 insertions(+), 74 deletions(-) diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index c21c08b7e8c0..e861eb7c0d64 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 node-wide HostPin limits. These are constraints on one placement decision, not independently spendable budgets. Multi-process workers begin with non-overlapping HostPin shares; the joint solve may reassign them asymmetrically without exceeding the node 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. Layerwise prefetch depth remains explicit because its latency effect depends on runtime overlap rather than weight size alone. Explicit residency policies also remain authoritative; otherwise native policy candidates may be tuned. Offline, realtime, and no-warmup paths retain the 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..54e1326e2dfd 100644 --- a/docs/docs/sglang-diffusion/deployment_cookbook.mdx +++ b/docs/docs/sglang-diffusion/deployment_cookbook.mdx @@ -128,9 +128,9 @@ 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 node-wide HostPin and HostPin-transition limits. Load and runtime VRAM are lifecycle constraints on the same physical memory, not independently spendable pools; every constraint and HostPin uses the same placement selection. The target frontier covers component offload, resident components, layerwise resident counts, implicit DiT residency policies, and pinned layer-store reassignment. A post-adjustment warmup realizes the selected state and validates its memory peak before the server becomes ready; it may keep or roll back that one measured adjustment but does not start an open-ended tuning loop. -This optimization targets the model-default workload, batch size 1, unless `--warmup-sampling-params` supplies a representative request. Larger production shapes or batches should set that override or use explicit `--component-residency`. Layerwise prefetch depth and residency policy remain user-controlled because their latency effect depends on measured copy/compute overlap; the planner does not infer them from weight bytes. Multi-process runs use non-overlapping per-worker HostPin shares, which prevents node overcommit but can be more conservative than a node-global asymmetric allocation. +This optimization targets the model-default workload, 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 remains user-controlled because its latency effect depends on measured copy/compute overlap. An explicitly selected residency policy also remains authoritative; otherwise the planner may choose among the native policy candidates. Multi-process runs begin with non-overlapping per-worker HostPin shares. The joint placement solve may then assign asymmetric per-rank quotas while keeping their sum within the same node-wide allowance, so HostPin is optimized with VRAM rather than by a separate planner. Before warmup, `auto` still uses load-safe model/card defaults. It checks selected GPU memory before applying FSDP, uses the least available memory across selected GPUs, and only enables FSDP automatically on validated paths where it can replace DiT offload. Offline, realtime, no-warmup, and unsupported serving paths retain those static defaults. When the model default uses CFG and the user did not set a parallelism policy, `auto` may also enable CFG parallelism. `memory` favors lower memory, while `speed` favors latency/throughput and accepts OOM risk. `speed` keeps `torch.compile` disabled by default because its effect varies by model and workload; `--enable-torch-compile true` always opts in explicitly. diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index fcfb51bf1625..3d766a7bd187 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -91,6 +91,9 @@ get_global_component_residency_manager, 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, @@ -117,9 +120,6 @@ ) 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 ( - resolve_keep_resident_min_available_gb, -) 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, @@ -1345,6 +1345,7 @@ def apply_auto_residency( not pre_warmup and self._auto_residency_round_sizes and not recovering_from_oom + and not self._latest_auto_residency_round_is_resident_only() ): regressions = [] for report in reports: @@ -1617,20 +1618,6 @@ def _build_pre_warmup_auto_residency_report( ) budget_bytes = self._auto_residency_budget_bytes() - admission_gib = resolve_keep_resident_min_available_gb(self.server_args) - budget_gib = budget_bytes / GIB_BYTES - if budget_gib < admission_gib: - return RankResidencyReport( - rank=self.rank, - budget_bytes=budget_bytes, - estimated_peak_bytes=None, - candidates=[], - skip_reason=( - f"{budget_gib:.1f} GiB VRAM budget is below the " - f"{admission_gib:.1f} GiB uncalibrated residency threshold" - ), - ) - all_candidates = self._collect_auto_residency_targets(workload) candidates = pre_warmup_residency_targets( all_candidates, @@ -1915,6 +1902,25 @@ def _auto_residency_modules(self) -> dict[str, object]: return dict(self.pipeline.modules) return manager.placement_modules() + def _latest_auto_residency_round_is_resident_only(self) -> bool: + """Whether validation follows only execution-monotonic promotions. + + Moving an offloaded component to resident removes transfer/hooks from + its existing compute path. Its post-adjustment warmup must still prove + VRAM safety, but one noisy duration sample must not undo that promotion. + Layerwise layout and lower-memory transitions retain latency rollback. + """ + if not self._auto_residency_round_sizes: + return False + round_size = self._auto_residency_round_sizes[-1] + if round_size <= 0 or round_size > len(self._auto_residency_applied): + return False + return all( + adjustment.residency_mode != RESIDENT + and self.server_args.residency_mode(adjustment.component_name) == RESIDENT + for adjustment in self._auto_residency_applied[-round_size:] + ) + def _rollback_applied_residency_changes( self, *, latest_round_only: bool = False ) -> str | None: diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py index a00764847f39..f38b70f58187 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py @@ -22,7 +22,6 @@ from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args.auto_tune import ( auto_residency_static_skip_reason, - resolve_keep_resident_min_available_gb, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -60,8 +59,8 @@ def choose_initial_resident_components( """Choose resident startup components without removing runtime options. This is a load-feasibility seed, not a second serving-placement planner. - Native DiTs can still be demoted by warmup calibration. Auxiliary modules, - unknown weights, pipeline exclusions, and explicit choices retain their + Native weighted components can still be demoted by warmup calibration. + Unknown weights, pipeline exclusions, and explicit choices retain their configured loading semantics. A model-default layerwise placement remains eligible because avoiding that initialization is the point of this seed. """ @@ -92,17 +91,19 @@ def choose_initial_resident_components( weight_bytes is None or weight_bytes <= 0 or item.component_name in excluded_components - or not is_dit_component_name(item.component_name) or server_args.residency_mode(item.component_name) == RESIDENT or server_args.explicit_residency_mode(item.component_name) is not None ): continue + uses_per_request = ( + denoising_steps if is_dit_component_name(item.component_name) else 1 + ) options.append( PlacementOption( group_key=item.component_name, option_key=f"{item.component_name}:resident", resource_delta_bytes={"gpu:load": weight_bytes}, - estimated_latency_savings=weight_bytes * denoising_steps, + estimated_latency_savings=weight_bytes * uses_per_request, preference_cost=(weight_bytes,), ) ) @@ -139,15 +140,6 @@ def maybe_seed_initial_residency( available_gib, envs.SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB, ) - admission_gib = resolve_keep_resident_min_available_gb(server_args) - if available_gib < admission_gib: - logger.debug( - "Keeping configured load placement: %.1f GiB is below the %.1f GiB " - "initial-residency admission threshold.", - available_gib, - admission_gib, - ) - return selected = choose_initial_resident_components( server_args, inventory, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 9012b8e5c43f..e48bce347936 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -474,17 +474,30 @@ def prepare_component_weight_inventory( """Resolve selected checkpoint weights before component loading.""" self.component_weight_inventory = estimate_component_weight_inventory(sources) known = { - item.component_name: round(item.checkpoint_bytes / 1024**3, 2) + item.component_name: ( + ( + round(item.checkpoint_bytes / 1024**3, 2) + if item.checkpoint_bytes is not None + else None + ), + ( + round(materialized_bytes / 1024**3, 2) + if (materialized_bytes := item.materialized_bytes()) is not None + else None + ), + ) for item in self.component_weight_inventory - if item.checkpoint_bytes is not None and item.checkpoint_bytes > 0 + if item.checkpoint_bytes is not None + or item.materialized_bytes() is not None } unknown = [ item.component_name for item in self.component_weight_inventory - if item.checkpoint_bytes is None + if item.checkpoint_bytes is None and item.materialized_bytes() is None ] logger.debug( - "Pre-load component weights (GiB): known=%s unknown=%s", + "Pre-load component weights (checkpoint/materialized GiB): " + "known=%s unknown=%s", known, unknown, ) 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 9ec39dade07c..eed6ebbda551 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -1175,22 +1175,57 @@ def test_pre_warmup_targets_defer_device_weight_promotion(self): assert pre_warmup_residency_targets([current, resident]) == [current] - def test_pre_warmup_report_skips_uncalibrated_promotion_below_threshold( + def test_pre_warmup_report_uses_weight_model_below_legacy_threshold( self, monkeypatch ): from sglang.multimodal_gen.runtime.managers import gpu_worker + current = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=8 * GIB_BYTES, + target_device_weight_bytes=0, + current_placement=True, + ) + modules = {"transformer": object()} worker = SimpleNamespace( rank=0, - pipeline=object(), - server_args=object(), - _auto_residency_budget_bytes=lambda: 30 * GIB_BYTES, - _collect_auto_residency_targets=lambda _workload: pytest.fail( - "candidate collection should be skipped" + pipeline=SimpleNamespace(preload_residency_excluded_components=frozenset()), + server_args=SimpleNamespace( + num_gpus=1, + nnodes=1, + node_rank=0, + host_pin_budget=lambda: None, ), + _auto_residency_budget_bytes=lambda: 30 * GIB_BYTES, + _collect_auto_residency_targets=lambda _workload: [current], + _auto_residency_modules=lambda: modules, + ) + monkeypatch.setattr( + gpu_worker, + "component_runtime_weight_bytes", + lambda _modules: {"transformer": 8 * GIB_BYTES}, + ) + monkeypatch.setattr( + gpu_worker, + "component_current_device_weight_bytes", + lambda _modules: {"transformer": 0}, + ) + monkeypatch.setattr( + gpu_worker, "layerwise_pinned_host_bytes", lambda *a, **k: 0 + ) + monkeypatch.setattr( + gpu_worker, "layerwise_host_pin_capacity_bytes", lambda *a, **k: 0 ) monkeypatch.setattr( - gpu_worker, "resolve_keep_resident_min_available_gb", lambda _args: 60.0 + gpu_worker, "host_memory_available_bytes", lambda: 100 * GIB_BYTES + ) + monkeypatch.setattr( + gpu_worker.torch, + "get_device_module", + lambda: SimpleNamespace(memory_allocated=lambda: 2 * GIB_BYTES), ) report = gpu_worker.GPUWorker._build_pre_warmup_auto_residency_report( @@ -1204,12 +1239,9 @@ def test_pre_warmup_report_skips_uncalibrated_promotion_below_threshold( ) assert report.budget_bytes == 30 * GIB_BYTES - assert report.estimated_peak_bytes is None - assert report.candidates == [] - assert report.skip_reason == ( - "30.0 GiB VRAM budget is below the 60.0 GiB uncalibrated residency " - "threshold" - ) + assert report.estimated_peak_bytes == 10 * GIB_BYTES + assert report.candidates == [current] + assert report.skip_reason is None def test_pre_warmup_report_keeps_excluded_components_fixed(self, monkeypatch): from sglang.multimodal_gen.runtime.managers import gpu_worker @@ -1248,9 +1280,6 @@ def test_pre_warmup_report_keeps_excluded_components_fixed(self, monkeypatch): _collect_auto_residency_targets=lambda _workload: [transformer, vae], _auto_residency_modules=lambda: modules, ) - monkeypatch.setattr( - gpu_worker, "resolve_keep_resident_min_available_gb", lambda _args: 0 - ) monkeypatch.setattr( gpu_worker, "component_runtime_weight_bytes", 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 4ff27015612d..157e81a8a91c 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 @@ -24,6 +24,7 @@ ) 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 @@ -183,8 +184,13 @@ 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_applied = [ + AppliedResidencyChange("transformer", COMPONENT_OFFLOAD) + ] worker._auto_residency_round_sizes = [1] report = RankResidencyReport( rank=0, @@ -222,6 +228,50 @@ def test_worker_rolls_back_a_materially_slower_calibrated_round(self): self.assertFalse(kwargs["already_failed"]) self.assertTrue(kwargs["latest_round_only"]) + def test_worker_keeps_resident_only_round_despite_noisy_duration(self): + worker = GPUWorker.__new__(GPUWorker) + worker.rank = 0 + worker.is_output_rank = False + worker.server_args = SimpleNamespace(residency_mode=lambda _name: RESIDENT) + worker._auto_residency_warmup_records = [object()] + worker._auto_residency_applied = [ + AppliedResidencyChange("text_encoder", COMPONENT_OFFLOAD), + AppliedResidencyChange("vae", COMPONENT_OFFLOAD), + ] + worker._auto_residency_round_sizes = [2] + 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]) + worker._rollback_everywhere = mock.Mock() + + 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(), + ), + ): + response = worker.apply_auto_residency(validate_only=True) + + self.assertEqual(response.output["status"], PLACEMENT_STATUS_VALIDATED) + worker._rollback_everywhere.assert_not_called() + def test_worker_validation_does_not_apply_a_second_candidate_plan(self): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 diff --git a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py index 3ec4d05f8f2b..f7db1cde02d6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py @@ -99,7 +99,7 @@ def test_initial_seed_stays_conservative_for_unknown_fixed_resident_weights(): assert selected == set() -def test_initial_seed_does_not_override_explicit_or_auxiliary_components(): +def test_initial_seed_keeps_explicit_placement_and_can_select_auxiliary_components(): args = _Args( modes={"transformer": LAYERWISE_OFFLOAD}, explicit={"transformer"}, @@ -111,7 +111,24 @@ def test_initial_seed_does_not_override_explicit_or_auxiliary_components(): denoising_steps=8, ) - assert selected == set() + assert selected == {"vae"} + + +def test_initial_seed_selects_all_known_components_when_they_fit(): + args = _Args(modes={"transformer": LAYERWISE_OFFLOAD}) + + selected = choose_initial_resident_components( + args, + [ + _weight("transformer", 8), + _weight("text_encoder", 4), + _weight("vae", 2), + ], + available_bytes=40 * GIB_BYTES, + denoising_steps=8, + ) + + assert selected == {"transformer", "text_encoder", "vae"} def test_initial_seed_can_bypass_model_default_layerwise_initialization(): @@ -150,11 +167,6 @@ def test_initial_seed_applies_one_reversible_override(): "initial_residency.auto_residency_static_skip_reason", return_value=None, ), - patch( - "sglang.multimodal_gen.runtime.managers.memory_managers." - "initial_residency.resolve_keep_resident_min_available_gb", - return_value=30, - ), patch( "sglang.multimodal_gen.runtime.managers.memory_managers." "initial_residency.current_platform" @@ -180,11 +192,6 @@ def test_initial_seed_honors_the_test_allocator_cap(): "initial_residency.auto_residency_static_skip_reason", return_value=None, ), - patch( - "sglang.multimodal_gen.runtime.managers.memory_managers." - "initial_residency.resolve_keep_resident_min_available_gb", - return_value=30, - ), patch( "sglang.multimodal_gen.runtime.managers.memory_managers." "initial_residency.current_platform" @@ -222,7 +229,7 @@ def test_initial_seed_skips_structural_fsdp_path(): assert args.selected == {} -def test_initial_seed_keeps_configured_load_path_below_admission_threshold(): +def test_initial_seed_uses_inventory_budget_below_legacy_threshold(): args = _Args() with ( patch( @@ -230,11 +237,6 @@ def test_initial_seed_keeps_configured_load_path_below_admission_threshold(): "initial_residency.auto_residency_static_skip_reason", return_value=None, ), - patch( - "sglang.multimodal_gen.runtime.managers.memory_managers." - "initial_residency.resolve_keep_resident_min_available_gb", - return_value=30, - ), patch( "sglang.multimodal_gen.runtime.managers.memory_managers." "initial_residency.current_platform" @@ -244,4 +246,4 @@ def test_initial_seed_keeps_configured_load_path_below_admission_threshold(): platform.get_available_gpu_memory.return_value = 29.8 maybe_seed_initial_residency(args, [_weight("transformer", 8)]) - assert args.selected == {} + assert args.selected == {"transformer": RESIDENT} From 4602a063b828c790816c00fef6e4b2b61fefb33c Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 09:37:30 +0800 Subject: [PATCH 22/70] test(diffusion): complete auto residency worker fixtures --- .../multimodal_gen/test/unit/test_auto_residency_warmup.py | 2 ++ 1 file changed, 2 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 157e81a8a91c..a528de401a9b 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 @@ -278,6 +278,7 @@ def test_worker_validation_does_not_apply_a_second_candidate_plan(self): worker.is_output_rank = False worker.server_args = SimpleNamespace() worker._auto_residency_warmup_records = [object()] + worker._auto_residency_applied = [] worker._auto_residency_round_sizes = [1] report = RankResidencyReport( rank=0, @@ -324,6 +325,7 @@ def test_new_placement_that_spends_the_reserve_rolls_back(self): worker.pipeline = SimpleNamespace(modules={}) worker.server_args = SimpleNamespace() worker._auto_residency_warmup_records = [] + worker._auto_residency_applied = [] worker._auto_residency_round_sizes = [1] worker._auto_residency_budget_correction_bytes = 0 worker._auto_residency_last_applied_plan = AutoResidencyPlan( From f8ac41c9a24b5f7dce2746c4077c4674290024de Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 10:01:16 +0800 Subject: [PATCH 23/70] test(diffusion): enforce rank-local layerwise DTensors --- .../memory_managers/layerwise_offload.py | 5 ++- .../test/unit/test_layerwise_offload.py | 31 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) 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 c66190e54007..86eee4aa0d06 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 @@ -579,7 +579,10 @@ def _wrap_for_target( ) -> torch.Tensor: if isinstance(target, DTensor): return DTensor.from_local( - local_tensor, target.device_mesh, target.placements + local_tensor, + target.device_mesh, + target.placements, + run_check=False, ) return local_tensor 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 c90a8dd92826..47f08a2439f2 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -90,6 +90,37 @@ def stream(_stream): return nullcontext() +def test_dtensor_layer_store_remains_rank_local(monkeypatch): + local_weight = torch.arange(4) + replacement = torch.arange(4) + 10 + wraps = [] + + class FakeDTensor: + def __init__(self, local_tensor): + self.local_tensor = local_tensor + self.device_mesh = "mesh" + self.placements = ("shard",) + + def to_local(self): + return self.local_tensor + + @classmethod + def from_local(cls, local_tensor, device_mesh, placements, **kwargs): + wraps.append((local_tensor, device_mesh, placements, kwargs)) + return cls(local_tensor) + + monkeypatch.setattr(layerwise_offload_mod, "DTensor", FakeDTensor) + target = FakeDTensor(local_weight) + + assert LayerwiseOffloadManager._to_local_tensor(target) is local_weight + wrapped = LayerwiseOffloadManager._wrap_for_target( + object.__new__(LayerwiseOffloadManager), target, replacement + ) + + assert wrapped.local_tensor is replacement + assert wraps == [(replacement, "mesh", ("shard",), {"run_check": False})] + + class _DummyBlock(torch.nn.Module): def __init__(self) -> None: super().__init__() From 66f51ad8f98b7d99c8ab51eb9b4dc27f1d52b16f Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 10:31:39 +0800 Subject: [PATCH 24/70] [diffusion] bound auto residency calibration --- .../runtime/managers/gpu_worker.py | 17 ++++++++- .../memory_managers/auto_residency.py | 29 +++++++++++++++ .../memory_managers/initial_residency.py | 20 ++++++++++- .../memory_managers/placement_budget.py | 21 ++++++++++- .../test/unit/test_auto_residency.py | 35 +++++++++++++++++++ .../test/unit/test_initial_residency.py | 19 ++++++++++ .../test/unit/test_placement_budget.py | 25 +++++++++++++ 7 files changed, 163 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 3d766a7bd187..775e42dfbae3 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -70,6 +70,7 @@ component_current_device_weight_bytes, component_runtime_weight_bytes, describe_error, + estimate_allocator_headroom_bytes, estimate_candidate_latency_savings_ns, estimate_default_workload_peak_bytes, estimate_default_workload_timing, @@ -601,6 +602,7 @@ def _execute_forward_common( else None ) warmup_baseline_allocated_bytes = 0 + warmup_baseline_reserved_bytes = 0 try: if measure_server_warmup: # Drop the previous request's allocator pool so each probe @@ -613,6 +615,9 @@ def _execute_forward_common( warmup_baseline_allocated_bytes = ( torch.get_device_module().memory_allocated() ) + warmup_baseline_reserved_bytes = ( + torch.get_device_module().memory_reserved() + ) start_time = ( execution_start_time @@ -734,6 +739,7 @@ def _execute_forward_common( req=req, workload=warmup_workload, baseline_allocated_bytes=warmup_baseline_allocated_bytes, + baseline_reserved_bytes=warmup_baseline_reserved_bytes, succeeded=output_batch is not None and output_batch.error is None, ) return output_batch @@ -744,6 +750,7 @@ def _record_server_warmup_memory( req: Req, workload: tuple[int, int, int, int], baseline_allocated_bytes: int, + baseline_reserved_bytes: int, succeeded: bool, ) -> None: phase_allocated_peaks: dict[str, int] = {} @@ -797,6 +804,10 @@ def _record_server_warmup_memory( baseline_allocated_bytes=int(baseline_allocated_bytes), peak_allocated_bytes=request_allocated_peak, succeeded=succeeded, + baseline_reserved_bytes=int(baseline_reserved_bytes), + 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, @@ -1711,6 +1722,10 @@ def _build_auto_residency_report( for record in records ) runtime_weights_by_component = component_runtime_weight_bytes(modules) + allocator_headroom_bytes = estimate_allocator_headroom_bytes( + records=records, + target_units=target_units, + ) if warmup_oom: ( estimated_peak_bytes, @@ -1836,7 +1851,7 @@ def _build_auto_residency_report( budget_bytes=budget_bytes, estimated_peak_bytes=estimated_peak_bytes, planning_headroom_correction_bytes=( - self._auto_residency_budget_correction_bytes + self._auto_residency_budget_correction_bytes + allocator_headroom_bytes ), target_workload_measured=any( record.succeeded and record.workload_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 738f9a4cc232..4eddb8a6212c 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 @@ -139,6 +139,8 @@ class WarmupMemoryRecord(msgspec.Struct, frozen=True): baseline_allocated_bytes: int peak_allocated_bytes: int succeeded: bool + baseline_reserved_bytes: int = 0 + 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, ...]] = {} @@ -733,6 +735,33 @@ def estimate_default_workload_peak_bytes( return max(estimates) +def estimate_allocator_headroom_bytes( + *, records: Iterable[WarmupMemoryRecord], target_units: int | None +) -> int: + """Return observed allocator capacity above live tensor allocations. + + Placement deltas are expressed in live weight bytes, but CUDA must fit the + allocator pool that served the measured workload too. Prefer covering + target-shape records; otherwise retain the largest observed gap without + pretending it scales linearly with pixels or frames. + """ + records = list(records) + if target_units is not None: + covering = [ + record for record in records if record.workload_units() >= target_units + ] + if covering: + records = covering + return max( + ( + max(record.peak_reserved_bytes, record.peak_allocated_bytes) + - record.peak_allocated_bytes + for record in records + ), + default=0, + ) + + def estimate_workload_phase_peaks( *, records: Iterable[WarmupMemoryRecord], diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py index f38b70f58187..2090dd9ee470 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py @@ -33,6 +33,7 @@ GIB_BYTES = 1024**3 MIN_INITIAL_RESIDENCY_RESERVE_BYTES = 4 * GIB_BYTES INITIAL_RESIDENCY_RESERVE_FRACTION = 0.20 +MULTI_DIT_INITIAL_RESIDENCY_RESERVE_FRACTION = 0.30 def _default_denoising_steps(server_args: ServerArgs) -> int: @@ -64,9 +65,26 @@ def choose_initial_resident_components( configured loading semantics. A model-default layerwise placement remains eligible because avoiding that initialization is the point of this seed. """ + dit_count = sum( + is_dit_component_name(item.component_name) + and ( + server_args.explicit_residency_mode(item.component_name) is None + or server_args.residency_mode(item.component_name) == RESIDENT + ) + for item in inventory + ) + # Before the first target-shape probe there is no activation/workspace + # measurement. Multiple denoisers can otherwise consume nearly the whole + # device with static weights even though only one is active at a time. Keep + # a larger temporary probe allowance; measured planning can spend it again. + reserve_fraction = ( + MULTI_DIT_INITIAL_RESIDENCY_RESERVE_FRACTION + if dit_count > 1 + else INITIAL_RESIDENCY_RESERVE_FRACTION + ) reserve_bytes = max( MIN_INITIAL_RESIDENCY_RESERVE_BYTES, - int(available_bytes * INITIAL_RESIDENCY_RESERVE_FRACTION), + int(available_bytes * reserve_fraction), ) fixed_resident = [ item 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 dcf2d9686552..d8430749ad32 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 @@ -197,7 +197,7 @@ def optimize_placement( # 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( + binding_resource_names = tuple( resource_name for resource_name in all_resource_names if sum( @@ -212,6 +212,25 @@ def optimize_placement( ) > resource_budget_bytes[resource_name] ) + + # Phase-level VRAM accounting commonly produces many constraints with the + # exact same cost for every option. Such columns describe the same left-hand + # side, so only the smallest budget can bind. Collapsing them before local + # Pareto construction is exact and avoids quadratic skyline work over dozens + # of duplicate stage dimensions. + ordered_options = sorted(options, key=lambda option: option.option_key) + resource_by_cost_vector: dict[tuple[int, ...], str] = {} + for resource_name in binding_resource_names: + cost_vector = tuple( + option.resource_delta_bytes.get(resource_name, 0) + for option in ordered_options + ) + incumbent = resource_by_cost_vector.get(cost_vector) + if incumbent is None or ( + resource_budget_bytes[resource_name] < resource_budget_bytes[incumbent] + ): + resource_by_cost_vector[cost_vector] = resource_name + resource_names = tuple(sorted(resource_by_cost_vector.values())) resource_budgets = tuple( int(resource_budget_bytes[name]) for name in resource_names ) 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 eed6ebbda551..6f3631d1c9e8 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 @@ collect_residency_targets, component_resident_size_bytes, component_runtime_weight_bytes, + estimate_allocator_headroom_bytes, estimate_candidate_latency_savings_ns, estimate_default_workload_peak_bytes, estimate_default_workload_timing, @@ -68,6 +69,7 @@ def _record( num_frames=17, baseline_gib=10, peak_gib=12, + peak_reserved_gib=None, succeeded=True, num_inference_steps=1, total_duration_ms=0.0, @@ -83,6 +85,7 @@ def _record( baseline_allocated_bytes=baseline_gib * GIB_BYTES, peak_allocated_bytes=peak_gib * GIB_BYTES, succeeded=succeeded, + peak_reserved_bytes=int((peak_reserved_gib or peak_gib) * GIB_BYTES), num_inference_steps=num_inference_steps, total_duration_ms=total_duration_ms, stage_duration_ms=stage_duration_ms or {}, @@ -94,6 +97,38 @@ def _record( ) +class TestEstimateAllocatorHeadroom: + def test_uses_covering_workload_allocator_gap(self): + small = _record( + num_frames=9, + peak_gib=10, + peak_reserved_gib=11, + ) + target = _record( + num_frames=17, + peak_gib=20, + peak_reserved_gib=26, + ) + + assert ( + estimate_allocator_headroom_bytes( + records=[small, target], target_units=target.workload_units() + ) + == 6 * GIB_BYTES + ) + + def test_keeps_largest_observed_gap_without_covering_workload(self): + first = _record(peak_gib=10, peak_reserved_gib=12) + second = _record(peak_gib=15, peak_reserved_gib=19) + + assert ( + estimate_allocator_headroom_bytes( + records=[first, second], target_units=second.workload_units() * 2 + ) + == 4 * GIB_BYTES + ) + + class TestEstimateDefaultWorkloadTiming: def test_scales_only_denoising_with_target_steps(self): record = _record( diff --git a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py index f7db1cde02d6..f3a01a304f7b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py @@ -131,6 +131,25 @@ def test_initial_seed_selects_all_known_components_when_they_fit(): assert selected == {"transformer", "text_encoder", "vae"} +def test_initial_seed_keeps_probe_headroom_for_multiple_dits(): + args = _Args() + + selected = choose_initial_resident_components( + args, + [ + _weight("transformer", 27), + _weight("transformer_2", 27), + _weight("text_encoder", 10), + _weight("vae", 1), + ], + available_bytes=75 * GIB_BYTES, + denoising_steps=40, + ) + + assert len(selected & {"transformer", "transformer_2"}) == 1 + assert {"text_encoder", "vae"} <= selected + + def test_initial_seed_can_bypass_model_default_layerwise_initialization(): args = _Args(modes={"transformer": LAYERWISE_OFFLOAD}) 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 7cf6b09c7060..dd33a19715f7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py @@ -653,3 +653,28 @@ def test_branch_and_bound_matches_exhaustive_multiresource_search(): ) assert plan.estimated_latency_savings == expected_utility assert plan.preference_cost == expected_cost + + +def test_identical_resource_columns_keep_only_the_strictest_budget(): + options = [ + PlacementOption( + group_key="dit", + option_key="dit:small", + resource_delta_bytes={"phase-a": 4, "phase-b": 4}, + estimated_latency_savings=4, + ), + PlacementOption( + group_key="dit", + option_key="dit:large", + resource_delta_bytes={"phase-a": 8, "phase-b": 8}, + estimated_latency_savings=8, + ), + ] + + plan = optimize_placement( + options, + resource_budget_bytes={"phase-a": 10, "phase-b": 6}, + ) + + assert [option.option_key for option in plan.selections] == ["dit:small"] + assert plan.resource_delta_bytes == {"phase-a": 4, "phase-b": 4} From bc9eb0a8421d891b5818fb5eed34bf26d8cd28e9 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 10:44:58 +0800 Subject: [PATCH 25/70] [diffusion] avoid speculative preload retuning --- .../runtime/managers/gpu_worker.py | 8 +++ .../memory_managers/auto_residency.py | 8 +-- .../memory_managers/placement_budget.py | 7 ++- .../test/unit/test_auto_residency.py | 53 +++++++++++++++++++ .../test/unit/test_placement_budget.py | 26 +++++++++ 5 files changed, 95 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 775e42dfbae3..4eba9bb4fb15 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1694,6 +1694,14 @@ def _build_pre_warmup_auto_residency_report( device_transition_allocated_bytes=baseline_allocated, candidates=candidates, require_feasible_placement=True, + # Static weight accounting establishes feasibility, but contains no + # latency observation. Tie utility so the solver preserves the + # current placement whenever it fits and only makes the minimum + # necessary demotion before the first target-shape probe. + estimated_request_duration_ns=1, + candidate_latency_savings_ns={ + candidate.option_key(): 0 for candidate in candidates + }, ) 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 4eddb8a6212c..33c1f0319f13 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 @@ -2617,12 +2617,8 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP resource_delta_bytes=resource_deltas, estimated_latency_savings=estimated_latency_savings, preference_cost=( - ( - 0 - if candidate.current_placement - or candidate.target_residency_mode == candidate.residency_mode - else 1 - ), + int(not candidate.current_placement) + + int(candidate.target_residency_mode != candidate.residency_mode), _layerwise_policy_preference_cost(candidate), candidate.target_device_weight_bytes or max(0, candidate.target_resident_weight_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 d8430749ad32..78066cf6f1e2 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 @@ -70,7 +70,12 @@ 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])) + # Preference costs are summed and compared lexicographically by the + # final objective. Lexicographic order is translation invariant, so a + # lexicographically better prefix remains better after any suffix is + # added; requiring every individual dimension to improve only retains + # dominated HostPin-prefix states. + 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 6f3631d1c9e8..27584251655d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -1189,6 +1189,57 @@ def test_pre_warmup_report_enforces_weight_feasibility_without_oom(self): LAYERWISE_OFFLOAD ] + def test_pre_warmup_plan_keeps_feasible_current_placement(self): + current = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=5 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + target_device_weight_bytes=GIB_BYTES, + current_placement=True, + ) + component_offload = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=11 * GIB_BYTES, + target_device_weight_bytes=0, + ) + less_layerwise = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=7 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + target_device_weight_bytes=0, + ) + candidates = [current, less_layerwise, component_offload] + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=12, + estimated_gib=2, + phase_peaks_gib={"static:transformer": 2}, + phase_components={"static:transformer": ("transformer",)}, + candidates=candidates, + require_feasible_placement=True, + estimated_request_duration_ns=1, + candidate_latency_savings_ns={ + candidate.option_key(): 0 for candidate in candidates + }, + ) + ] + ) + + assert plan.changes == [] + def test_pre_warmup_targets_defer_device_weight_promotion(self): current = ResidencyTarget( component_name="text_encoder", @@ -1277,6 +1328,8 @@ def test_pre_warmup_report_uses_weight_model_below_legacy_threshold( assert report.estimated_peak_bytes == 10 * GIB_BYTES assert report.candidates == [current] assert report.skip_reason is None + assert report.estimated_request_duration_ns == 1 + assert report.candidate_latency_savings_ns == {current.option_key(): 0} def test_pre_warmup_report_keeps_excluded_components_fixed(self, monkeypatch): from sglang.multimodal_gen.runtime.managers import gpu_worker 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 dd33a19715f7..1562651c82ff 100644 --- a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py @@ -313,6 +313,32 @@ def test_latency_tolerance_is_global_not_per_option(): assert plan.estimated_latency_savings == 60 +def test_lexicographic_preference_prunes_irrelevant_later_dimension(): + plan = optimize_placement( + [ + PlacementOption( + group_key="dit", + option_key="partial-pin", + resource_delta_bytes={"gpu": 1}, + estimated_latency_savings=90, + preference_cost=(0, -90, 5), + ), + PlacementOption( + group_key="dit", + option_key="full-pin", + resource_delta_bytes={"gpu": 1}, + estimated_latency_savings=100, + preference_cost=(0, -100, 10), + ), + ], + resource_budget_bytes={"gpu": 1}, + estimated_latency_tolerance=20, + require_selection_from_every_group=True, + ) + + assert [option.option_key for option in plan.selections] == ["full-pin"] + + def test_one_static_placement_satisfies_all_observed_phase_constraints(): plan = optimize_placement( [ From 7baddfcd89f5159f57ba0782842e53e9dd638cc5 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 10:58:07 +0800 Subject: [PATCH 26/70] [diffusion] scope preload placement preference --- .../multimodal_gen/runtime/managers/gpu_worker.py | 10 ++++++---- .../runtime/managers/memory_managers/auto_residency.py | 8 ++++++-- .../multimodal_gen/test/unit/test_auto_residency.py | 5 +++-- 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 4eba9bb4fb15..f5904ad88fa6 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1695,12 +1695,14 @@ def _build_pre_warmup_auto_residency_report( candidates=candidates, require_feasible_placement=True, # Static weight accounting establishes feasibility, but contains no - # latency observation. Tie utility so the solver preserves the - # current placement whenever it fits and only makes the minimum - # necessary demotion before the first target-shape probe. + # latency observation. Give each exact current state one synthetic + # unit so the solver preserves the maximum number of current + # placements and only demotes what is necessary before the first + # target-shape probe. estimated_request_duration_ns=1, candidate_latency_savings_ns={ - candidate.option_key(): 0 for candidate in candidates + candidate.option_key(): int(candidate.current_placement) + for candidate in 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 33c1f0319f13..4eddb8a6212c 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 @@ -2617,8 +2617,12 @@ def plan_auto_residency(*, reports: list[RankResidencyReport]) -> AutoResidencyP resource_delta_bytes=resource_deltas, estimated_latency_savings=estimated_latency_savings, preference_cost=( - int(not candidate.current_placement) - + int(candidate.target_residency_mode != candidate.residency_mode), + ( + 0 + if candidate.current_placement + or candidate.target_residency_mode == candidate.residency_mode + else 1 + ), _layerwise_policy_preference_cost(candidate), candidate.target_device_weight_bytes or max(0, candidate.target_resident_weight_bytes), 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 27584251655d..5f2f71b577dd 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -1232,7 +1232,8 @@ def test_pre_warmup_plan_keeps_feasible_current_placement(self): require_feasible_placement=True, estimated_request_duration_ns=1, candidate_latency_savings_ns={ - candidate.option_key(): 0 for candidate in candidates + candidate.option_key(): int(candidate.current_placement) + for candidate in candidates }, ) ] @@ -1329,7 +1330,7 @@ def test_pre_warmup_report_uses_weight_model_below_legacy_threshold( assert report.candidates == [current] assert report.skip_reason is None assert report.estimated_request_duration_ns == 1 - assert report.candidate_latency_savings_ns == {current.option_key(): 0} + assert report.candidate_latency_savings_ns == {current.option_key(): 1} def test_pre_warmup_report_keeps_excluded_components_fixed(self, monkeypatch): from sglang.multimodal_gen.runtime.managers import gpu_worker From 6314c5e03bec123893812be9547145f0eb55000d Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 11:06:29 +0800 Subject: [PATCH 27/70] [diffusion] defer automatic layerwise setup --- .../memory_managers/initial_residency.py | 42 ++++++++---------- .../test/unit/test_initial_residency.py | 43 +++++++++++++++++-- 2 files changed, 59 insertions(+), 26 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py index 2090dd9ee470..870db295291d 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py @@ -7,6 +7,8 @@ from sglang.multimodal_gen import envs 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_weight_inventory import ( @@ -33,7 +35,6 @@ GIB_BYTES = 1024**3 MIN_INITIAL_RESIDENCY_RESERVE_BYTES = 4 * GIB_BYTES INITIAL_RESIDENCY_RESERVE_FRACTION = 0.20 -MULTI_DIT_INITIAL_RESIDENCY_RESERVE_FRACTION = 0.30 def _default_denoising_steps(server_args: ServerArgs) -> int: @@ -65,26 +66,9 @@ def choose_initial_resident_components( configured loading semantics. A model-default layerwise placement remains eligible because avoiding that initialization is the point of this seed. """ - dit_count = sum( - is_dit_component_name(item.component_name) - and ( - server_args.explicit_residency_mode(item.component_name) is None - or server_args.residency_mode(item.component_name) == RESIDENT - ) - for item in inventory - ) - # Before the first target-shape probe there is no activation/workspace - # measurement. Multiple denoisers can otherwise consume nearly the whole - # device with static weights even though only one is active at a time. Keep - # a larger temporary probe allowance; measured planning can spend it again. - reserve_fraction = ( - MULTI_DIT_INITIAL_RESIDENCY_RESERVE_FRACTION - if dit_count > 1 - else INITIAL_RESIDENCY_RESERVE_FRACTION - ) reserve_bytes = max( MIN_INITIAL_RESIDENCY_RESERVE_BYTES, - int(available_bytes * reserve_fraction), + int(available_bytes * INITIAL_RESIDENCY_RESERVE_FRACTION), ) fixed_resident = [ item @@ -164,14 +148,26 @@ def maybe_seed_initial_residency( available_bytes=int(available_gib * GIB_BYTES), excluded_components=excluded_components, ) - if not selected: - return + coarse_offload = { + item.component_name + for item in inventory + if item.component_name not in selected + and item.component_name not in excluded_components + and server_args.explicit_residency_mode(item.component_name) is None + and server_args.residency_mode(item.component_name) == LAYERWISE_OFFLOAD + } for component_name in selected: server_args.set_auto_residency_mode(component_name, RESIDENT) + for component_name in coarse_offload: + server_args.set_auto_residency_mode(component_name, COMPONENT_OFFLOAD) + if not selected and not coarse_offload: + return logger.info( - "Initial auto residency: resident=%s (minimum free VRAM=%.1f GiB); " - "warmup may rebalance this load-safe seed.", + "Initial auto residency: resident=%s, component-offload=%s " + "(minimum free VRAM=%.1f GiB); warmup may rebalance this " + "load-safe seed.", sorted(selected), + sorted(coarse_offload), available_gib, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py index f3a01a304f7b..8386b79fdf7d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py @@ -131,7 +131,7 @@ def test_initial_seed_selects_all_known_components_when_they_fit(): assert selected == {"transformer", "text_encoder", "vae"} -def test_initial_seed_keeps_probe_headroom_for_multiple_dits(): +def test_initial_seed_packs_reused_dits_under_the_common_reserve(): args = _Args() selected = choose_initial_resident_components( @@ -146,8 +146,7 @@ def test_initial_seed_keeps_probe_headroom_for_multiple_dits(): denoising_steps=40, ) - assert len(selected & {"transformer", "transformer_2"}) == 1 - assert {"text_encoder", "vae"} <= selected + assert selected == {"transformer", "transformer_2", "vae"} def test_initial_seed_can_bypass_model_default_layerwise_initialization(): @@ -202,6 +201,44 @@ def test_initial_seed_applies_one_reversible_override(): ) +def test_initial_seed_defers_unselected_layerwise_setup(): + args = _Args( + modes={ + "transformer": LAYERWISE_OFFLOAD, + "transformer_2": LAYERWISE_OFFLOAD, + "text_encoder": LAYERWISE_OFFLOAD, + "vae": LAYERWISE_OFFLOAD, + } + ) + inventory = [ + _weight("transformer", 27), + _weight("transformer_2", 27), + _weight("text_encoder", 10), + _weight("vae", 1), + ] + with ( + patch( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.auto_residency_static_skip_reason", + return_value=None, + ), + patch( + "sglang.multimodal_gen.runtime.managers.memory_managers." + "initial_residency.current_platform" + ) as platform, + ): + platform.is_cuda.return_value = True + platform.get_available_gpu_memory.return_value = 75 + maybe_seed_initial_residency(args, inventory) + + assert args.selected == { + "transformer": RESIDENT, + "transformer_2": RESIDENT, + "text_encoder": COMPONENT_OFFLOAD, + "vae": RESIDENT, + } + + def test_initial_seed_honors_the_test_allocator_cap(): args = _Args() inventory = [_weight("transformer", 20)] From ca5d5649fa96b5865646f70841b96bc3c8ce36b5 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 11:31:54 +0800 Subject: [PATCH 28/70] [diffusion] keep load-safe layerwise fallback --- .../memory_managers/auto_residency.py | 19 ++++++++++--- .../test/unit/test_auto_residency.py | 27 +++++++++++++++++++ 2 files changed, 43 insertions(+), 3 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 4eddb8a6212c..7f24c11f69a6 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 @@ -345,8 +345,15 @@ def pre_warmup_residency_targets( Static planning knows complete component weights but not the runtime working set. It may demote an unsafe configured layout or redistribute - HostPin at the same device footprint, but promoting more weights here can + HostPin at the same device footprint, but a pure residency promotion can consume activation headroom that only the first shape probe can reveal. + + A layerwise layout can require a small persistent non-layer footprint even + when the coarse component-offload state keeps every weight on the host. + Keep that alternative when it reduces the component's active working set: + the joint solver still checks its inactive, transition, and HostPin costs. + Excluding it here would leave no pre-warmup escape when the full component + fits neither the device nor the representative request. """ candidates = list(candidates) current_device_bytes = { @@ -359,8 +366,14 @@ def pre_warmup_residency_targets( for candidate in candidates if candidate.component_name not in excluded_components and candidate.component_name in current_device_bytes - and candidate.target_device_weight_bytes - <= current_device_bytes[candidate.component_name] + and ( + candidate.target_device_weight_bytes + <= current_device_bytes[candidate.component_name] + or ( + candidate.target_residency_mode == LAYERWISE_OFFLOAD + and candidate.active_device_delta_bytes < 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 5f2f71b577dd..e9ab04ef8c59 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -1262,6 +1262,33 @@ def test_pre_warmup_targets_defer_device_weight_promotion(self): assert pre_warmup_residency_targets([current, resident]) == [current] + def test_pre_warmup_targets_keep_active_working_set_demotion(self): + current = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + target_device_weight_bytes=0, + current_placement=True, + ) + layerwise = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + target_device_weight_bytes=GIB_BYTES, + active_device_delta_bytes=-16 * GIB_BYTES, + inactive_device_delta_bytes=GIB_BYTES, + device_transition_delta_bytes=GIB_BYTES, + ) + + assert pre_warmup_residency_targets([current, layerwise]) == [ + current, + layerwise, + ] + def test_pre_warmup_report_uses_weight_model_below_legacy_threshold( self, monkeypatch ): From ed0be88ad6d1fb2fc2891500206987bece22ed23 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 11:43:24 +0800 Subject: [PATCH 29/70] [diffusion] preserve auxiliary load semantics --- .../memory_managers/initial_residency.py | 31 +++++-------------- .../test/unit/test_initial_residency.py | 14 ++++----- 2 files changed, 13 insertions(+), 32 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py index 870db295291d..01d4ed2e5b7c 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/initial_residency.py @@ -7,8 +7,6 @@ from sglang.multimodal_gen import envs 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_weight_inventory import ( @@ -61,8 +59,8 @@ def choose_initial_resident_components( """Choose resident startup components without removing runtime options. This is a load-feasibility seed, not a second serving-placement planner. - Native weighted components can still be demoted by warmup calibration. - Unknown weights, pipeline exclusions, and explicit choices retain their + Native DiTs can still be demoted by warmup calibration. Auxiliary modules, + unknown weights, pipeline exclusions, and explicit choices retain their configured loading semantics. A model-default layerwise placement remains eligible because avoiding that initialization is the point of this seed. """ @@ -93,19 +91,17 @@ def choose_initial_resident_components( weight_bytes is None or weight_bytes <= 0 or item.component_name in excluded_components + or not is_dit_component_name(item.component_name) or server_args.residency_mode(item.component_name) == RESIDENT or server_args.explicit_residency_mode(item.component_name) is not None ): continue - uses_per_request = ( - denoising_steps if is_dit_component_name(item.component_name) else 1 - ) options.append( PlacementOption( group_key=item.component_name, option_key=f"{item.component_name}:resident", resource_delta_bytes={"gpu:load": weight_bytes}, - estimated_latency_savings=weight_bytes * uses_per_request, + estimated_latency_savings=weight_bytes * denoising_steps, preference_cost=(weight_bytes,), ) ) @@ -148,26 +144,13 @@ def maybe_seed_initial_residency( available_bytes=int(available_gib * GIB_BYTES), excluded_components=excluded_components, ) - coarse_offload = { - item.component_name - for item in inventory - if item.component_name not in selected - and item.component_name not in excluded_components - and server_args.explicit_residency_mode(item.component_name) is None - and server_args.residency_mode(item.component_name) == LAYERWISE_OFFLOAD - } - for component_name in selected: server_args.set_auto_residency_mode(component_name, RESIDENT) - for component_name in coarse_offload: - server_args.set_auto_residency_mode(component_name, COMPONENT_OFFLOAD) - if not selected and not coarse_offload: + if not selected: return logger.info( - "Initial auto residency: resident=%s, component-offload=%s " - "(minimum free VRAM=%.1f GiB); warmup may rebalance this " - "load-safe seed.", + "Initial auto residency: resident=%s (minimum free VRAM=%.1f GiB); " + "warmup may rebalance this load-safe seed.", sorted(selected), - sorted(coarse_offload), available_gib, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py index 8386b79fdf7d..7c0bf44d5d69 100644 --- a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py @@ -99,7 +99,7 @@ def test_initial_seed_stays_conservative_for_unknown_fixed_resident_weights(): assert selected == set() -def test_initial_seed_keeps_explicit_placement_and_can_select_auxiliary_components(): +def test_initial_seed_keeps_explicit_placement_and_auxiliary_load_semantics(): args = _Args( modes={"transformer": LAYERWISE_OFFLOAD}, explicit={"transformer"}, @@ -111,10 +111,10 @@ def test_initial_seed_keeps_explicit_placement_and_can_select_auxiliary_componen denoising_steps=8, ) - assert selected == {"vae"} + assert selected == set() -def test_initial_seed_selects_all_known_components_when_they_fit(): +def test_initial_seed_selects_only_dits_when_all_components_fit(): args = _Args(modes={"transformer": LAYERWISE_OFFLOAD}) selected = choose_initial_resident_components( @@ -128,7 +128,7 @@ def test_initial_seed_selects_all_known_components_when_they_fit(): denoising_steps=8, ) - assert selected == {"transformer", "text_encoder", "vae"} + assert selected == {"transformer"} def test_initial_seed_packs_reused_dits_under_the_common_reserve(): @@ -146,7 +146,7 @@ def test_initial_seed_packs_reused_dits_under_the_common_reserve(): denoising_steps=40, ) - assert selected == {"transformer", "transformer_2", "vae"} + assert selected == {"transformer", "transformer_2"} def test_initial_seed_can_bypass_model_default_layerwise_initialization(): @@ -201,7 +201,7 @@ def test_initial_seed_applies_one_reversible_override(): ) -def test_initial_seed_defers_unselected_layerwise_setup(): +def test_initial_seed_preserves_unselected_layerwise_setup(): args = _Args( modes={ "transformer": LAYERWISE_OFFLOAD, @@ -234,8 +234,6 @@ def test_initial_seed_defers_unselected_layerwise_setup(): assert args.selected == { "transformer": RESIDENT, "transformer_2": RESIDENT, - "text_encoder": COMPONENT_OFFLOAD, - "vae": RESIDENT, } From d4b2364b3ba161d5c14afd2c6461980a501c24a0 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 12:02:47 +0800 Subject: [PATCH 30/70] [diffusion] align warmup allocator fixtures --- .../multimodal_gen/test/unit/test_performance_metrics.py | 4 ++++ 1 file changed, 4 insertions(+) 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 d9c5045b6525..9283005cbbf9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -133,6 +133,7 @@ def test_server_warmup_preserves_peak_after_managed_stage_timeline(): req=req, workload=(128, 96, 9, 2), baseline_allocated_bytes=3, + baseline_reserved_bytes=3, succeeded=True, ) @@ -176,6 +177,7 @@ def test_server_warmup_does_not_treat_allocator_cache_as_untracked_live_memory() req=req, workload=(64, 64, 1, 1), baseline_allocated_bytes=3, + baseline_reserved_bytes=12, succeeded=True, ) @@ -201,6 +203,7 @@ def test_server_warmup_normalizes_profile_stage_names_for_residency_timing(): residency_manager.current_device_components.return_value = ("transformer_2",) device_module = Mock() device_module.max_memory_allocated.return_value = 8 + device_module.max_memory_reserved.return_value = 8 req = SimpleNamespace( metrics=SimpleNamespace( total_duration_ms=12.0, @@ -221,6 +224,7 @@ def test_server_warmup_normalizes_profile_stage_names_for_residency_timing(): req=req, workload=(64, 64, 1, 1), baseline_allocated_bytes=3, + baseline_reserved_bytes=3, succeeded=True, ) From 60e81a1331fb6b401b6214f0cc62ef030d725381 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 12:31:46 +0800 Subject: [PATCH 31/70] [diffusion] refresh preload VRAM baseline --- .../sglang/multimodal_gen/test/server/perf_baselines/5090.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 2012eaee4829..521a983b7afd 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json @@ -60,7 +60,7 @@ "expected_e2e_ms": 18255.39, "expected_avg_denoise_ms": 338.78, "expected_median_denoise_ms": 345.95, - "load_peak_vram_mb": 1126.0, + "load_peak_vram_mb": 8522.0, "runtime_peak_vram_mb": 20064.0, "estimated_full_test_time_s": 94.0 }, From 8faf3ee5f86528bc69c32b3d22e51adab609fec1 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 12:33:36 +0800 Subject: [PATCH 32/70] [diffusion] document serving placement scope --- 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 54e1326e2dfd..1893268e8797 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 node-wide HostPin and HostPin-transition limits. Load and runtime VRAM are lifecycle constraints on the same physical memory, not independently spendable pools; every constraint and HostPin uses the same placement selection. The target frontier covers component offload, resident components, layerwise resident counts, implicit DiT residency policies, and pinned layer-store reassignment. A post-adjustment warmup realizes the selected state and validates its memory peak before the server becomes ready; it may keep or roll back that one measured adjustment but does not start an open-ended tuning loop. +With an eligible `sglang serve` warmup, `auto` records allocated-memory peaks for the representative request and their active component layouts. One serving-placement decision then satisfies every rank and execution-phase VRAM constraint together with node-wide HostPin and HostPin-transition limits. These are constraints on one decision vector, not independently spendable pools. The earlier preload plan is only a reversible load-feasibility seed: after exact modules exist, the serving planner rebuilds the complete frontier and may replace every seeded choice. This preserves the optimum of the supplied serving frontier, but does not claim to jointly minimize startup and serving latency over every possible lifecycle. The target frontier covers component offload, resident components, layerwise resident counts, implicit DiT residency policies, and pinned layer-store reassignment. A post-adjustment warmup realizes the selected state and validates its memory peak before the server becomes ready; it may keep or roll back that one measured adjustment but does not start an open-ended tuning loop. 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 remains user-controlled because its latency effect depends on measured copy/compute overlap. An explicitly selected residency policy also remains authoritative; otherwise the planner may choose among the native policy candidates. Multi-process runs begin with non-overlapping per-worker HostPin shares. The joint placement solve may then assign asymmetric per-rank quotas while keeping their sum within the same node-wide allowance, so HostPin is optimized with VRAM rather than by a separate planner. From 05efe7f4a4f7c1bf813ffe94c8fe942f6366a3ec Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 13:02:49 +0800 Subject: [PATCH 33/70] [diffusion] clarify joint placement invariants --- .../sglang-diffusion/deployment_cookbook.mdx | 2 ++ .../test/unit/test_auto_residency.py | 16 ++++++++++++++-- .../test/unit/test_placement_budget.py | 4 ++-- 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/docs/docs/sglang-diffusion/deployment_cookbook.mdx b/docs/docs/sglang-diffusion/deployment_cookbook.mdx index 1893268e8797..03054545bcbd 100644 --- a/docs/docs/sglang-diffusion/deployment_cookbook.mdx +++ b/docs/docs/sglang-diffusion/deployment_cookbook.mdx @@ -191,6 +191,8 @@ In this example, `auto` will not re-enable FSDP. The same applies to parallelism **FSDP** shards DiT weights across multiple GPUs and all-gathers weights during forward. It can reduce DiT CPU offload cost on multi-GPU deployments, especially for validated Wan I2V workloads. +Checkpoint loading does not use that forward-time collective as a loading primitive. Each rank reads the checkpoint independently and constructs its final local DTensor storage directly; the fallback for a full checkpoint tensor slices locally instead of redistributing from another rank. Layerwise offload follows the same rule: its CPU layer stores contain that rank's local DTensor shards and wrap them back with `DTensor.from_local`, without a load-time all-gather. FSDP forward execution still performs the collectives required by its sharding policy. + FSDP sharding granularity matters. SGLang Diffusion prefers sharding direct repeated transformer block entries such as `transformer_blocks.0` or `blocks.0`. Coarser sharding lowers wrapper count but can increase all-gather peak memory; finer sharding can reduce transient memory but adds communication and scheduling overhead. If a model does not define an explicit sharding rule, the loader falls back to repeated block class names and common direct numbered block paths. **CFG parallelism** splits positive and negative CFG branches across GPUs. For Qwen/Wan workloads with normal step counts, this is the most reliable multi-GPU speedup observed so far. 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 e9ab04ef8c59..b935a2f0cd3f 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -3146,11 +3146,23 @@ def test_non_dit_layerwise_frontier_includes_partial_residency(self): "condition_image_encoder", ), ) - @pytest.mark.parametrize("startup_mode", (COMPONENT_OFFLOAD, RESIDENT)) + @pytest.mark.parametrize( + "startup_mode", (COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD, RESIDENT) + ) def test_every_serving_component_group_gets_a_complete_frontier( self, component_name, startup_mode ): - module = _FakeLazyLayerwiseDit(num_layers=4) + module = ( + _FakeLayerwiseDit( + [ + _FakeLayerwiseManager( + {f"layers.{index}.w": torch.zeros(16) for index in range(4)} + ) + ] + ) + if startup_mode == LAYERWISE_OFFLOAD + else _FakeLazyLayerwiseDit(num_layers=4) + ) candidates = collect_residency_targets( modules={component_name: 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 1562651c82ff..d7adbe1dcacf 100644 --- a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py @@ -417,7 +417,7 @@ def test_lifecycle_plan_jointly_constrains_load_transition_runtime_and_hostpin() ), PlacementOption( group_key="dit-lifecycle", - option_key="dit-sharded-load-to-layerwise", + option_key="dit-rank-local-load-to-layerwise", resource_delta_bytes={ "gpu:rank0:load": 4, "gpu:rank0:transition": 3, @@ -460,7 +460,7 @@ def test_lifecycle_plan_jointly_constrains_load_transition_runtime_and_hostpin() # 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", + "dit-rank-local-load-to-layerwise", "encoder-gpu-load-to-resident", ] assert plan.resource_delta_bytes == { From e29d77a75e592c90f0353b04277b88515addbf02 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 13:22:02 +0800 Subject: [PATCH 34/70] [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 b594330ad9a6..33e039b5894f 100644 --- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py +++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py @@ -611,7 +611,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 a981b1e4212f..102e5372d10d 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 @@ -907,6 +907,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 9ebe2fc404f1f703a782f10890168e02a645af9f Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 13:31:19 +0800 Subject: [PATCH 35/70] [diffusion] shorten resident-only residency validation --- .../runtime/managers/gpu_worker.py | 31 +++++-- .../memory_managers/auto_residency.py | 12 ++- .../multimodal_gen/runtime/server_warmup.py | 25 +++++- .../test/unit/test_auto_residency_warmup.py | 80 ++++++++++++++++++- 4 files changed, 132 insertions(+), 16 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index f5904ad88fa6..000445343fcc 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1478,7 +1478,13 @@ def apply_auto_residency( 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: @@ -1927,6 +1933,14 @@ def _auto_residency_modules(self) -> dict[str, object]: return dict(self.pipeline.modules) return manager.placement_modules() + 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. @@ -1935,15 +1949,18 @@ def _latest_auto_residency_round_is_resident_only(self) -> bool: VRAM safety, but one noisy duration sample must not undo that promotion. Layerwise layout and lower-memory transitions retain latency rollback. """ - if not self._auto_residency_round_sizes: - return False - round_size = self._auto_residency_round_sizes[-1] - if round_size <= 0 or round_size > len(self._auto_residency_applied): - return False + 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 self._auto_residency_applied[-round_size:] + for adjustment in changes + ) and bool(changes) + + def _latest_auto_residency_round_supports_short_validation(self) -> bool: + """Whether one full-shape step covers the changed execution paths.""" + 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( 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 7f24c11f69a6..646c73c0be4e 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 @@ -3254,14 +3254,20 @@ 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 - lives in the logged ``format_plan_summary`` line. + The orchestrator branches on ``status`` and whether validation only needs + a full-shape memory check; human-readable detail stays in the plan log. """ return { "status": status, "changed": [candidate.component_name for candidate in plan.changes], "recovering_from_oom": plan.recovering_from_oom, + "short_validation": short_validation, } diff --git a/python/sglang/multimodal_gen/runtime/server_warmup.py b/python/sglang/multimodal_gen/runtime/server_warmup.py index d510c0793f94..f2a17e6faeb6 100644 --- a/python/sglang/multimodal_gen/runtime/server_warmup.py +++ b/python/sglang/multimodal_gen/runtime/server_warmup.py @@ -276,12 +276,22 @@ async def rollback_and_rewarm(error: Exception) -> None: adjusted_any = True recovering_from_oom = response_recovering_from_oom + short_validation = bool( + isinstance(response.output, dict) + and response.output.get("short_validation") + ) # 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. + # phases that overlap under it. Resident-only changes need one + # full-shape step for memory safety; other changes retain the longer + # timing sample used by regression validation. try: + validation_options = {"step_limit": 1} if short_validation else {} await run_async_client_warmup( - server_args, forward, fail_open=False, rewarm=True + server_args, + forward, + fail_open=False, + rewarm=True, + **validation_options, ) except Exception as e: if recovering_from_oom and _is_out_of_memory(e): @@ -398,6 +408,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, @@ -410,6 +421,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) @@ -423,6 +436,7 @@ async def run_async_client_warmup( *, fail_open: bool = False, rewarm: bool = False, + step_limit: int | None = None, ) -> None: try: auto_residency_handles_oom = auto_residency_skip_reason(server_args) is None @@ -431,7 +445,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 a528de401a9b..3d1780d7e41c 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 @@ -30,6 +30,28 @@ 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.assertEqual(req.extra["warmup_total"], 1) + self.assertTrue(req.extra["server_warmup_rewarm"]) + def test_residency_hint_uses_effective_device_budget(self): worker = GPUWorker.__new__(GPUWorker) worker._auto_residency_budget_bytes = mock.Mock(return_value=12 * GIB_BYTES) @@ -271,6 +293,20 @@ def test_worker_keeps_resident_only_round_despite_noisy_duration(self): self.assertEqual(response.output["status"], PLACEMENT_STATUS_VALIDATED) worker._rollback_everywhere.assert_not_called() + self.assertTrue(worker._latest_auto_residency_round_supports_short_validation()) + + def test_dit_residency_keeps_multi_step_validation(self): + worker = GPUWorker.__new__(GPUWorker) + worker.server_args = SimpleNamespace(residency_mode=lambda _name: RESIDENT) + worker._auto_residency_applied = [ + AppliedResidencyChange("transformer", COMPONENT_OFFLOAD) + ] + worker._auto_residency_round_sizes = [1] + + self.assertTrue(worker._latest_auto_residency_round_is_resident_only()) + self.assertFalse( + worker._latest_auto_residency_round_supports_short_validation() + ) def test_worker_validation_does_not_apply_a_second_candidate_plan(self): worker = GPUWorker.__new__(GPUWorker) @@ -523,7 +559,12 @@ async def forward(req): if req.action == "apply" else PLACEMENT_STATUS_VALIDATED ) - return OutputBatch(output={"status": status}) + return OutputBatch( + output={ + "status": status, + "short_validation": req.action == "apply", + } + ) rewarm = mock.AsyncMock() server_args = SimpleNamespace( @@ -539,7 +580,42 @@ async def forward(req): asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) self.assertEqual(actions, ["apply", "validate"]) - rewarm.assert_awaited_once() + rewarm.assert_awaited_once_with( + server_args, + forward, + fail_open=False, + rewarm=True, + step_limit=1, + ) + + def test_non_resident_adjustment_keeps_timing_validation_steps(self): + async def forward(req): + status = ( + PLACEMENT_STATUS_ADJUSTED + if req.action == "apply" + else PLACEMENT_STATUS_VALIDATED + ) + return OutputBatch(output={"status": status}) + + 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)) + + rewarm.assert_awaited_once_with( + server_args, + forward, + fail_open=False, + rewarm=True, + ) def test_regressed_validation_rolls_back_and_rewarms_original_placement(self): responses = iter( From 63b8b320d60e9456922979caf2614a010a6a8988 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 13:35:22 +0800 Subject: [PATCH 36/70] [docs] define auto residency optimization boundary --- docs/docs/sglang-diffusion/api/cli.mdx | 2 +- docs/docs/sglang-diffusion/deployment_cookbook.mdx | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index e861eb7c0d64..668059f2e2cf 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 node-wide HostPin limits. These are constraints on one placement decision, not independently spendable budgets. Multi-process workers begin with non-overlapping HostPin shares; the joint solve may reassign them asymmetrically without exceeding the node 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. Layerwise prefetch depth remains explicit because its latency effect depends on runtime overlap rather than weight size alone. Explicit residency policies also remain authoritative; otherwise native policy candidates may be tuned. Offline, realtime, and no-warmup paths retain the model/card fallback policy. It uses FSDP only for validated DiT-offload replacement paths. `speed` keeps `torch.compile` disabled unless a model-specific deployment config opts in after validation; pass `--enable-torch-compile true` to enable it explicitly. Use `manual` to keep performance-related server args under explicit user control. Explicit offload, FSDP, and parallelism flags take precedence in all modes. +- `--performance-mode {manual|auto|speed|memory}` / `--mode`: preset for latency/throughput and memory defaults. `auto` is the default. With an eligible synthetic server warmup, it starts from a load-safe checkpoint-weight placement, measures the representative workload, and jointly selects component/layer residency against every rank and execution-phase VRAM constraint plus node-wide HostPin limits. These are constraints on one placement decision, not independently spendable budgets. Multi-process workers begin with non-overlapping HostPin shares; the joint solve may reassign them asymmetrically without exceeding the node allowance. The selected placement targets the model-default (or `--warmup-sampling-params`) workload with one output per prompt; use explicit residency or representative warmup parameters for larger production shapes, batches, or output counts. Layerwise prefetch depth remains explicit because its latency effect depends on runtime overlap rather than weight size alone. Explicit residency policies also remain authoritative; otherwise native policy candidates may be tuned. Offline, realtime, and no-warmup paths retain the static fallback policy. It uses FSDP only for validated DiT-offload replacement paths. `speed` keeps `torch.compile` disabled unless a model-specific deployment config opts in after validation; pass `--enable-torch-compile true` to enable it explicitly. Use `manual` to keep performance-related server args under explicit user control. Explicit offload, FSDP, and parallelism flags take precedence in all modes. - `--direct-gpu-weight-loading {true|false}`: opt into direct GPU loading for an unquantized, GPU-resident, TP=1 DiT by materializing its complete checkpoint state dict on GPU. 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 03054545bcbd..2643666aa354 100644 --- a/docs/docs/sglang-diffusion/deployment_cookbook.mdx +++ b/docs/docs/sglang-diffusion/deployment_cookbook.mdx @@ -128,11 +128,11 @@ 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 serving-placement decision then satisfies every rank and execution-phase VRAM constraint together with node-wide HostPin and HostPin-transition limits. These are constraints on one decision vector, not independently spendable pools. The earlier preload plan is only a reversible load-feasibility seed: after exact modules exist, the serving planner rebuilds the complete frontier and may replace every seeded choice. This preserves the optimum of the supplied serving frontier, but does not claim to jointly minimize startup and serving latency over every possible lifecycle. The target frontier covers component offload, resident components, layerwise resident counts, implicit DiT residency policies, and pinned layer-store reassignment. A post-adjustment warmup realizes the selected state and validates its memory peak before the server becomes ready; it may keep or roll back that one measured adjustment but does not start an open-ended tuning loop. +With an eligible `sglang serve` warmup, `auto` records allocated-memory peaks for the representative request and their active component layouts. One serving-placement decision then satisfies every rank and execution-phase VRAM constraint together with node-wide HostPin and HostPin-transition limits. These are constraints on one decision vector, not independently spendable pools. The earlier preload plan is only a reversible load-feasibility seed: after exact modules exist, the serving planner rebuilds the complete frontier and may replace every seeded choice. This preserves the optimum of the supplied serving frontier, but does not claim to jointly minimize startup and serving latency over every possible lifecycle. The target frontier covers component offload, resident components, layerwise resident counts, implicit DiT residency policies, 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 for this validation; DiT changes, demotions, and layerwise or HostPin rebalancing retain the longer timing sample. Validation may keep or roll back that one measured adjustment but never starts an open-ended tuning loop. -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 remains user-controlled because its latency effect depends on measured copy/compute overlap. An explicitly selected residency policy also remains authoritative; otherwise the planner may choose among the native policy candidates. Multi-process runs begin with non-overlapping per-worker HostPin shares. The joint placement solve may then assign asymmetric per-rank quotas while keeping their sum within the same node-wide allowance, so HostPin is optimized with VRAM rather than by a separate planner. +This optimization targets the model-default workload and one output per prompt unless `--warmup-sampling-params` supplies a representative request. Larger production shapes, batches, or output counts should set that override or use explicit `--component-residency`. Layerwise prefetch depth remains user-controlled because its latency effect depends on measured copy/compute overlap. An explicitly selected residency policy also remains authoritative; otherwise the planner may choose among the native policy candidates. Multi-process runs begin with non-overlapping per-worker HostPin shares. The joint placement solve may then assign asymmetric per-rank quotas while keeping their sum within the same node-wide allowance, so HostPin is optimized with VRAM rather than by a separate planner. -Before warmup, `auto` still uses load-safe model/card defaults. It checks selected GPU memory before applying FSDP, uses the least available memory across selected GPUs, and only enables FSDP automatically on validated paths where it can replace DiT offload. Offline, realtime, no-warmup, and unsupported serving paths retain those static defaults. When the model default uses CFG and the user did not set a parallelism policy, `auto` may also enable CFG parallelism. `memory` favors lower memory, while `speed` favors latency/throughput and accepts OOM risk. `speed` keeps `torch.compile` disabled by default because its effect varies by model and workload; `--enable-torch-compile true` always opts in explicitly. +Before warmup, `auto` inventories the selected checkpoint weights on every rank and may seed native DiTs as resident while keeping auxiliary components on their configured loading path. This seed reserves activation headroom and is not a second serving optimizer: weight size alone cannot safely promote encoders or VAEs before the target shape is measured. Each rank reads its checkpoint independently; neither this seed nor layerwise initialization introduces a load-time weight all-gather. Auto parallelism still uses the least available memory across selected GPUs and only enables FSDP on validated paths where it can replace DiT offload. Offline, realtime, no-warmup, and unsupported serving paths retain static defaults. When the model default uses CFG and the user did not set a parallelism policy, `auto` may also enable CFG parallelism. `memory` favors lower memory, while `speed` favors latency/throughput and accepts OOM risk. `speed` keeps `torch.compile` disabled by default because its effect varies by model and workload; `--enable-torch-compile true` always opts in explicitly. The modes tune native pipeline components declared to the component residency manager. DiTs, text/image encoders, VAEs, vocoders, adapters, and upsamplers can use layerwise offload when their native module declares its executable layer structure. Explicitly selecting an unsupported component fails at startup instead of falling back to another residency mode. From eeeb378e6595d48b4a5372b59db6505f4119fb99 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 13:48:33 +0800 Subject: [PATCH 37/70] [diffusion] report residency validation timing --- .../runtime/managers/gpu_worker.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 000445343fcc..9c42f2a7526f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1407,9 +1407,23 @@ def apply_auto_residency( ) 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 calibration passed; " - "keeping the selected placement." + "Auto residency: post-adjustment validation passed (%s); " + "keeping the selected placement.", + result, ) return OutputBatch( output=plan_summary_payload( From ad73c4df8d8fe0fb0ef91df3bceeb50a36024f83 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 14:25:20 +0800 Subject: [PATCH 38/70] [diffusion] avoid redundant residency host copies --- .../runtime/managers/gpu_worker.py | 23 ++++- .../memory_managers/auto_residency.py | 38 +++++++- .../memory_managers/layerwise_offload.py | 94 ++++++++++++++++--- .../test/unit/test_auto_residency.py | 50 ++++++++++ 4 files changed, 186 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 9c42f2a7526f..97d1ce1d6737 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -67,6 +67,7 @@ WarmupMemoryRecord, apply_residency_changes, collect_residency_targets, + commit_residency_changes, component_current_device_weight_bytes, component_runtime_weight_bytes, describe_error, @@ -1406,10 +1407,30 @@ def apply_auto_residency( retryable=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._auto_residency_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 646c73c0be4e..88ebc50aa83e 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 @@ -2764,6 +2764,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 + ) if rank is not None and rank in plan.host_pin_target_bytes_by_rank: previous_host_pin_budget_state = ( server_args.host_pin_budget().set_spendable_capacity( @@ -2872,6 +2875,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] @@ -2948,9 +2952,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( @@ -2961,6 +2966,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): @@ -3018,6 +3024,32 @@ 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 = ( + adjustment.module_ref + if adjustment.module_ref is not None + else 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 86eee4aa0d06..0aca7630110a 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 @@ -1267,7 +1267,45 @@ 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) + # Anonymous host stores can fill the copy stream without a + # per-layer host wait. Checkpoint mappings still use the + # synchronous path: the mapped courier has a bounded slot ring + # intended to overlap one forward, not materialize a whole + # model at once. + 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. + + The real device tensors must already be materialized and the manager + disabled. Repacking pinned stores as pageable here would copy the full + checkpoint for data that will never be streamed again. + """ + 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: @@ -2292,22 +2330,32 @@ 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: + # 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. + manager.release_all() + manager.register_forward_hooks() + raise + manager.enabled = False + disabled.append(manager) + except Exception: + # A module may own multiple layer groups. Do not leave earlier + # groups resident and later groups streaming when one group fails. + 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.""" @@ -2375,6 +2423,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 b935a2f0cd3f..07c7a91c4efb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -24,6 +24,7 @@ _layerwise_resident_targets, apply_residency_changes, collect_residency_targets, + commit_residency_changes, component_resident_size_bytes, component_runtime_weight_bytes, estimate_allocator_headroom_bytes, @@ -2660,6 +2661,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() @@ -2750,6 +2752,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]): @@ -3666,6 +3674,32 @@ 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,) + assert manager.release_host_stores_calls == 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( @@ -4014,6 +4048,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): From 393841283c61c4a974cedaf2a1825641a9918df4 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 14:49:52 +0800 Subject: [PATCH 39/70] [diffusion] validate residency commit barrier --- .../multimodal_gen/runtime/managers/gpu_worker.py | 13 +++++++++++-- .../managers/memory_managers/auto_residency.py | 2 +- .../test/unit/test_auto_residency_warmup.py | 9 +++++---- 3 files changed, 17 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 97d1ce1d6737..be88d1de7a6c 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1407,14 +1407,23 @@ def apply_auto_residency( retryable=True, ) if validate_only: + latest_round = self._latest_auto_residency_round() short_validation = ( self._latest_auto_residency_round_supports_short_validation() ) + commit_modules = ( + self._auto_residency_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._auto_residency_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 88ebc50aa83e..4cb731e8a5d1 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 @@ -3033,7 +3033,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 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 3d1780d7e41c..49c248fe2698 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 @@ -222,7 +222,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}, @@ -269,8 +269,9 @@ def test_worker_keeps_resident_only_round_despite_noisy_duration(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]) worker._rollback_everywhere = mock.Mock() + self.assertTrue(worker._latest_auto_residency_round_supports_short_validation()) with ( mock.patch.object( @@ -293,7 +294,7 @@ def test_worker_keeps_resident_only_round_despite_noisy_duration(self): self.assertEqual(response.output["status"], PLACEMENT_STATUS_VALIDATED) worker._rollback_everywhere.assert_not_called() - self.assertTrue(worker._latest_auto_residency_round_supports_short_validation()) + self.assertEqual(worker._auto_residency_round_sizes, []) def test_dit_residency_keeps_multi_step_validation(self): worker = GPUWorker.__new__(GPUWorker) @@ -324,7 +325,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 46f75630b6ece2e9e626b94e3d3b18628c3f0de2 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 18:19:29 +0800 Subject: [PATCH 40/70] [diffusion] tune layerwise policies per group --- docs/docs/sglang-diffusion/api/cli.mdx | 2 +- .../sglang-diffusion/deployment_cookbook.mdx | 2 +- .../memory_managers/auto_residency.py | 14 +++++--- .../test/unit/test_auto_residency.py | 33 +++++++++++++++++++ 4 files changed, 44 insertions(+), 7 deletions(-) diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index ec505be45458..db4d3087ccbf 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 checkpoint-weight placement, measures the representative workload, and jointly selects component/layer residency against every rank and execution-phase VRAM constraint plus node-wide HostPin limits. These are simultaneous constraints on one placement decision, not independently spendable budgets. Warmup also observes each declared layer group's call count, so repeated groups can receive independent resident-layer allocations while one-shot or unused groups use HostPin or whole-component residency instead. Multi-process workers begin with non-overlapping HostPin shares; the joint solve may reassign them asymmetrically without exceeding the node allowance. Explicit/online quantization and direct GPU loading fix only their affected components; eligible encoders, VAEs, and other components in the same pipeline still participate. The selected placement targets the model-default (or `--warmup-sampling-params`) workload with one output per prompt; use explicit residency or representative warmup parameters for larger production shapes, batches, or output counts. Layerwise prefetch depth remains explicit because its latency effect depends on runtime overlap rather than weight size alone. Explicit residency policies also remain authoritative; otherwise native policy candidates may be tuned. Offline, realtime, and no-warmup paths retain the static fallback policy. It uses FSDP only for validated DiT-offload replacement paths. `speed` keeps `torch.compile` disabled unless a model-specific deployment config opts in after validation; pass `--enable-torch-compile true` to enable it explicitly. Use `manual` to keep performance-related server args under explicit user control. Explicit offload, FSDP, and parallelism flags take precedence in all modes. +- `--performance-mode {manual|auto|speed|memory}` / `--mode`: preset for latency/throughput and memory defaults. `auto` is the default. With an eligible synthetic server warmup, it starts from a load-safe checkpoint-weight placement, measures the representative workload, and jointly selects component/layer residency against every rank and execution-phase VRAM constraint plus node-wide HostPin limits. These are simultaneous constraints on one placement decision, not independently spendable budgets. Warmup also observes each declared layer group's call count, so repeated groups can receive independent resident-layer allocations and implicit residency policies while one-shot or unused groups use HostPin or whole-component residency instead. Multi-process workers begin with non-overlapping HostPin shares; the joint solve may reassign them asymmetrically without exceeding the node allowance. Explicit/online quantization and direct GPU loading fix only their affected components; eligible encoders, VAEs, and other components in the same pipeline still participate. The selected placement targets the model-default (or `--warmup-sampling-params`) workload with one output per prompt; use explicit residency or representative warmup parameters for larger production shapes, batches, or output counts. Layerwise prefetch depth remains explicit because its latency effect depends on runtime overlap rather than weight size alone. Explicit residency policies also remain authoritative; otherwise native policy candidates may be tuned. Offline, realtime, and no-warmup paths retain the static fallback policy. It uses FSDP only for validated DiT-offload replacement paths. `speed` keeps `torch.compile` disabled unless a model-specific deployment config opts in after validation; pass `--enable-torch-compile true` to enable it explicitly. Use `manual` to keep performance-related server args under explicit user control. Explicit offload, FSDP, and parallelism flags take precedence in all modes. - `--direct-gpu-weight-loading {true|false}`: opt into direct GPU loading for an unquantized, GPU-resident, TP=1 DiT by materializing its complete checkpoint state dict on GPU. 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 915b1b2578d9..d9bb71afd13f 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 serving-placement decision then satisfies every rank and execution-phase VRAM constraint together with node-wide HostPin and HostPin-transition limits. These are simultaneous constraints on one decision vector, not separate planners or independently spendable pools. The earlier preload plan is only a replaceable load-feasibility seed: after exact modules exist, the serving planner rebuilds the complete frontier and may replace every seeded choice. This preserves the optimum of the supplied serving frontier, but does not claim to jointly minimize startup and serving latency over every possible lifecycle. The target frontier covers component offload, resident components, independent resident counts for repeatedly executed layer groups, implicit DiT residency policies, 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 for this validation; 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 serving-placement decision then satisfies every rank and execution-phase VRAM constraint together with node-wide HostPin and HostPin-transition limits. These are simultaneous constraints on one decision vector, not separate planners or independently spendable pools. The earlier preload plan is only a replaceable load-feasibility seed: after exact modules exist, the serving planner rebuilds the complete frontier and may replace every seeded choice. This preserves the optimum of the supplied serving frontier, but does not claim to jointly minimize startup and serving latency over every possible lifecycle. The target frontier covers component offload, resident components, independent resident counts and implicit residency policies for repeatedly executed layer groups, and pinned layer-store reassignment; one-shot or unused groups do not consume stage-resident VRAM. Explicit/online quantization and direct GPU loading keep only their affected components fixed, so unrelated encoders and VAEs remain eligible. A post-adjustment warmup realizes the selected state and validates its memory peak before the server becomes ready. Auxiliary-only moves to resident use one full-shape denoising step for this validation; 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 and one output per prompt unless `--warmup-sampling-params` supplies a representative request. Larger production shapes, batches, or output counts should set that override or use explicit `--component-residency`. Layerwise prefetch depth remains user-controlled because its latency effect depends on measured copy/compute overlap. An explicitly selected residency policy also remains authoritative; otherwise the planner may choose among the native policy candidates. Multi-process runs begin with non-overlapping per-worker HostPin shares. The joint placement solve may then assign asymmetric per-rank quotas while keeping their sum within the same node-wide allowance, so HostPin is optimized with VRAM rather than by a separate planner. 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 6f9d8092ec63..500eba662a82 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 @@ -1451,19 +1451,23 @@ def _layerwise_policy_targets( resident_layers: tuple[int, ...], tune_policy: bool, ) -> list[tuple[str, ...]]: - """Useful component-level policy choices for one resident-count target.""" + """Useful per-group policy choices for one resident-count target.""" current = tuple(manager.residency_policy for manager in managers) if not tune_policy or all( count in (0, manager.num_layers) for manager, count in zip(managers, resident_layers) ): return [current] - return list( - dict.fromkeys( - [current] - + [tuple(policy for _ in managers) for policy in RESIDENCY_POLICIES] + + choices = tuple( + ( + RESIDENCY_POLICIES + if 0 < count < manager.num_layers + else (manager.residency_policy,) ) + for manager, count in zip(managers, resident_layers) ) + return list(dict.fromkeys([current, *product(*choices)])) def _policy_affects_layerwise_layout( 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 61c811e7a3b2..b6668e98d869 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -3145,6 +3145,39 @@ def test_auto_dit_frontier_tunes_policy_unless_explicit(self): for candidate in explicit_candidates ) + def test_auto_dit_frontier_tunes_each_layer_group_policy(self): + managers = [ + _FakeLayerwiseManager( + {f"layers.{index}.w": torch.zeros(16) for index in range(4)}, + resident_layers=2, + ) + for _ in range(2) + ] + module = _FakeLayerwiseDit(managers) + + 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, + layerwise_policy_is_explicit=lambda _name, _dit_group: False, + ) + + policies = { + candidate.target_layerwise_residency_policies + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + and candidate.target_layerwise_resident_layers == (2, 2) + and candidate.target_layerwise_pinned_layers == ((), ()) + } + assert policies == { + (RESIDENCY_POLICY_LEADING, RESIDENCY_POLICY_LEADING), + (RESIDENCY_POLICY_LEADING, RESIDENCY_POLICY_STRIDED), + (RESIDENCY_POLICY_STRIDED, RESIDENCY_POLICY_LEADING), + (RESIDENCY_POLICY_STRIDED, RESIDENCY_POLICY_STRIDED), + } + def test_fully_pinned_component_can_release_hostpin_for_a_hotter_component(self): manager = _FakeLayerwiseManager( { From 8addcc8b71aafdf0b19a1938a583e7d5353b7335 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 18:27:56 +0800 Subject: [PATCH 41/70] [diffusion] tune repeated component layer policies --- .../memory_managers/auto_residency.py | 3 +- .../test/unit/test_auto_residency.py | 29 +++++++++++++++++++ 2 files changed, 30 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 500eba662a82..9846b4979387 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 @@ -1856,8 +1856,7 @@ def collect_residency_targets( and module.layerwise_offload_dit_group_enabled ) tune_residency_policy = bool( - is_dit_component_name(name) - and layerwise_policy_is_explicit is not None + layerwise_policy_is_explicit is not None and not layerwise_policy_is_explicit(name, dit_group) ) 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 b6668e98d869..f056f2520257 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -3178,6 +3178,35 @@ def test_auto_dit_frontier_tunes_each_layer_group_policy(self): (RESIDENCY_POLICY_STRIDED, RESIDENCY_POLICY_STRIDED), } + def test_repeated_non_dit_group_can_tune_policy(self): + manager = _FakeLayerwiseManager( + {f"layers.{index}.w": torch.zeros(16) for index in range(4)}, + resident_layers=2, + ) + module = _FakeLayerwiseDit([manager]) + + candidates = collect_residency_targets( + modules={"custom_refiner": module}, + residency_mode_of=self._modes({"custom_refiner": LAYERWISE_OFFLOAD}), + explicit_residency_mode_of=lambda _name: None, + custom_strategy_names=(), + num_inference_steps=10, + layerwise_policy_is_explicit=lambda _name, _dit_group: False, + layerwise_layer_uses={"custom_refiner": {"layers": (4, 4, 4, 4)}}, + ) + + policies = { + candidate.target_layerwise_residency_policies + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + and candidate.target_layerwise_resident_layers == (2,) + and candidate.target_layerwise_pinned_layers == ((),) + } + assert policies == { + (RESIDENCY_POLICY_LEADING,), + (RESIDENCY_POLICY_STRIDED,), + } + def test_fully_pinned_component_can_release_hostpin_for_a_hotter_component(self): manager = _FakeLayerwiseManager( { From 5b94aa271ecb85cdc02acd1e35059c8ec49aa397 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 18:37:47 +0800 Subject: [PATCH 42/70] [diffusion] classify repeated residency components by stage --- .../runtime/managers/gpu_worker.py | 22 +++++++++++++- .../memory_managers/auto_residency.py | 16 ++++++++-- .../test/unit/test_auto_residency.py | 29 +++++++++++++++++++ .../test/unit/test_auto_residency_warmup.py | 11 +++++++ 4 files changed, 74 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 256f55b21545..5cfb1ef98566 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -231,6 +231,7 @@ def __init__( self._auto_residency_reference_request_duration_ns: int | None = None self._auto_residency_reference_stage_duration_ns: dict[str, int] = {} self._auto_residency_reference_component_stages: dict[str, tuple[str, ...]] = {} + self._auto_residency_repeated_components: set[str] = set() # default workload resolved once for the per-request residency hint self._cached_default_workload: DefaultWorkload | None = None self._cached_default_workload_failed = False @@ -1845,6 +1846,11 @@ def _build_auto_residency_report( target_units=target_units, target_num_inference_steps=workload.num_inference_steps, ) + repeated_components = { + component_name + for component_name in layerwise_layer_uses + if is_dit_component_name(component_name) + } budget_bytes = self._auto_residency_budget_bytes() candidates = ( self._collect_auto_residency_targets( @@ -1909,6 +1915,17 @@ def _build_auto_residency_report( target_units=target_units, target_num_inference_steps=workload.num_inference_steps, ) + repeated_components.update( + component_name + for component_name, stage_names in measured_component_stages.items() + if any( + stage_name.endswith("DenoisingStage") + and not stage_name.endswith("BeforeDenoisingStage") + for stage_name in stage_names + ) + ) + if include_candidates: + self._auto_residency_repeated_components = repeated_components reference_request_duration_ns = ( self._auto_residency_reference_request_duration_ns ) @@ -1940,6 +1957,7 @@ def _build_auto_residency_report( self._auto_residency_reference_component_stages or measured_component_stages ), + repeated_components=repeated_components, ) if include_candidates else {} @@ -2067,7 +2085,9 @@ def _latest_auto_residency_round_supports_short_validation(self) -> bool: """Whether one full-shape step covers the changed execution paths.""" 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 + not is_dit_component_name(change.component_name) + and change.component_name not in self._auto_residency_repeated_components + for change in changes ) def _rollback_applied_residency_changes( 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 10f8807d0eae..1b438027bb57 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 @@ -599,6 +599,7 @@ def estimate_candidate_latency_savings_ns( request_duration_ns: int, stage_duration_ns: Mapping[str, int] | None = None, component_stages: Mapping[str, tuple[str, ...]] | None = None, + repeated_components: Iterable[str] = (), ) -> dict[str, int]: """Estimate each complete placement relative to the measured one. @@ -612,6 +613,7 @@ def estimate_candidate_latency_savings_ns( candidates = list(candidates) stage_duration_ns = stage_duration_ns or {} component_stages = component_stages or {} + repeated_component_names = set(repeated_components) candidates_by_component: dict[str, list[ResidencyTarget]] = {} maximum_h2d_bytes: dict[str, int] = {} for candidate in candidates: @@ -640,9 +642,13 @@ def estimate_candidate_latency_savings_ns( stage_duration_ns.get(stage_name, 0) for stage_name in component_stages.get(component_name, ()) ) + repeatedly_streamed = ( + is_dit_component_name(component_name) + or component_name in repeated_component_names + ) latency_upper_bound_ns = ( request_duration_ns - if is_dit_component_name(component_name) + if repeatedly_streamed else component_stage_duration_ns or request_duration_ns ) for candidate in component_candidates: @@ -687,16 +693,20 @@ def estimate_candidate_latency_savings_ns( # and can understate layerwise H2D stalls by orders of magnitude. The # transfer model is already an upper bound, so constrain DiT savings by # the synchronized request wall time instead. + repeatedly_streamed = ( + is_dit_component_name(candidate.component_name) + or candidate.component_name in repeated_component_names + ) latency_upper_bound_ns = ( request_duration_ns - if is_dit_component_name(candidate.component_name) + if repeatedly_streamed else component_stage_duration_ns or request_duration_ns ) if latency_upper_bound_ns > 0: transfer_upper_bound_ns = min( transfer_upper_bound_ns, latency_upper_bound_ns ) - if is_dit_component_name(candidate.component_name): + if repeatedly_streamed: # Model the remaining layer traffic, not a linearly scaled share # of the worst placement. When an all-pageable frontier exceeds # the request wall time, linear scaling compresses every marginal 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 20df02b33bc6..35f0c0798f13 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -288,6 +288,35 @@ def test_async_stage_time_does_not_cap_repeated_dit_transfers(self): assert savings[partial.option_key()] > 100_000_000 assert savings[resident.option_key()] > savings[partial.option_key()] + def test_repeated_nonstandard_component_uses_request_latency_cap(self): + partial = ResidencyTarget( + component_name="custom_refiner", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=10 * GIB_BYTES, + h2d_bytes_per_request=5 * GIB_BYTES, + target_layerwise_resident_layers=(5,), + target_layerwise_pinned_layers=((),), + ) + resident = ResidencyTarget( + component_name="custom_refiner", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=10 * GIB_BYTES, + target_layerwise_resident_layers=(10,), + target_layerwise_pinned_layers=((),), + ) + + savings = estimate_candidate_latency_savings_ns( + candidates=[partial, resident], + request_duration_ns=10_000_000_000, + stage_duration_ns={"denoise": 100_000_000}, + component_stages={"custom_refiner": ("denoise",)}, + repeated_components={"custom_refiner"}, + ) + + assert savings[partial.option_key()] > 100_000_000 + assert savings[resident.option_key()] > savings[partial.option_key()] + def test_dit_request_cap_does_not_compress_marginal_transfer_savings(self): partial = ResidencyTarget( component_name="transformer", 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 a077aee90077..cb304c4ba6e2 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 @@ -336,6 +336,7 @@ def test_worker_keeps_resident_only_round_despite_noisy_duration(self): AppliedResidencyChange("text_encoder", COMPONENT_OFFLOAD), AppliedResidencyChange("vae", COMPONENT_OFFLOAD), ] + worker._auto_residency_repeated_components = set() worker._auto_residency_round_sizes = [2] worker._auto_residency_last_applied_plan = AutoResidencyPlan() report = RankResidencyReport( @@ -375,6 +376,7 @@ def test_dit_residency_keeps_multi_step_validation(self): AppliedResidencyChange("transformer", COMPONENT_OFFLOAD) ] worker._auto_residency_round_sizes = [1] + worker._auto_residency_repeated_components = set() self.assertTrue(worker._latest_auto_residency_round_is_resident_only()) self.assertFalse( @@ -493,6 +495,7 @@ 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_repeated_components = set() worker._auto_residency_applied = [ AppliedResidencyChange("text_encoder", COMPONENT_OFFLOAD) @@ -506,6 +509,14 @@ def test_aux_resident_promotion_uses_short_validation_but_dit_does_not(self): worker._latest_auto_residency_round_supports_short_validation() ) + worker._auto_residency_applied = [ + AppliedResidencyChange("custom_refiner", COMPONENT_OFFLOAD) + ] + worker._auto_residency_repeated_components = {"custom_refiner"} + self.assertFalse( + worker._latest_auto_residency_round_supports_short_validation() + ) + def test_worker_rollback_response_rewarms_the_restored_layout(self): actions = [] From 1ab90f5c5bbde5abb8b3bbdeda05c86d1bd0f750 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 18:42:11 +0800 Subject: [PATCH 43/70] [diffusion] clarify repeated component utility --- .../managers/memory_managers/auto_residency.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 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 1b438027bb57..447dc63898db 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 @@ -605,7 +605,7 @@ def estimate_candidate_latency_savings_ns( Transfer work remains the ordering signal within a component frontier. Removable latency is capped by the measured component stage for one-shot - components and by request wall time for repeatedly streamed DiTs. Added + components and by request wall time for repeatedly streamed components. Added transfer time is not capped: an unmeasured mechanism may be slower than the whole request. Callers without a measured current option retain the historical absolute frontier estimate. @@ -659,7 +659,7 @@ def estimate_candidate_latency_savings_ns( ) # A measured stage bounds how much latency a placement can remove, # but not how much an unmeasured mechanism can add. Clipping both - # directions made a repeatedly streamed DiT appear merely neutral + # directions made a repeatedly streamed component appear merely neutral # once its transfer time exceeded the whole request; another small # component gain could then select a several-times-slower layout. if relative_savings_ns > 0 and latency_upper_bound_ns > 0: @@ -689,10 +689,10 @@ def estimate_candidate_latency_savings_ns( ) # Stage profiling is intentionally asynchronous by default. That is a # useful low-overhead measurement for one-shot encoder/decoder work, - # but it records mostly CPU launch time for a repeatedly streamed DiT - # and can understate layerwise H2D stalls by orders of magnitude. The - # transfer model is already an upper bound, so constrain DiT savings by - # the synchronized request wall time instead. + # but it records mostly CPU launch time for a repeatedly streamed + # component and can understate layerwise H2D stalls by orders of + # magnitude. The transfer model is already an upper bound, so constrain + # repeated-component savings by the synchronized request wall time. repeatedly_streamed = ( is_dit_component_name(candidate.component_name) or candidate.component_name in repeated_component_names From 2203893aa9a5339a10ef489a38027bdf051634ff Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 20:03:47 +0800 Subject: [PATCH 44/70] [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 15051c5758528cfad997ca57220d665442ab7be7 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 20:20:16 +0800 Subject: [PATCH 45/70] [diffusion] account resident layerwise host stores --- .../managers/memory_managers/auto_residency.py | 16 +++++++++------- .../test/unit/test_auto_residency.py | 16 ++++++++++++++++ 2 files changed, 25 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 5ec9cb5aa735..41b0012174ee 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 @@ -1775,13 +1775,15 @@ 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 - host_materialize_scratch = max( - ( - weight_bytes - for manager in managers - for weight_bytes in manager.layer_host_store_bytes().values() - ), - default=0, + layer_host_store_bytes = [ + weight_bytes + for manager in managers + for weight_bytes in manager.layer_host_store_bytes().values() + ] + host_materialize_scratch = ( + sum(layer_host_store_bytes) + if current_resident + else max(layer_host_store_bytes, default=0) ) targets = [] for resident_layers in _layerwise_resident_targets(managers, layer_uses=layer_uses): 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 ec9dde999f79..21b37f2a15cf 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -3761,6 +3761,14 @@ def test_default_resident_dit_exposes_complete_layerwise_frontier(self): for candidate in candidates if candidate.target_mode() == LAYERWISE_OFFLOAD ) + resident = next( + candidate for candidate in candidates if candidate.target_mode() == RESIDENT + ) + assert { + candidate.host_materialize_scratch_bytes + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + } == {resident.target_device_weight_bytes} def test_resident_seed_preserves_unconfigured_layerwise_frontier(self): module = _FakeLazyLayerwiseDit(num_layers=3) @@ -3803,6 +3811,14 @@ def test_virtual_layerwise_prefetch_is_runtime_not_transition_memory(self): ) assert streamed.device_transition_delta_bytes == 0 assert streamed.active_device_delta_bytes < 0 + resident = next( + candidate for candidate in candidates if candidate.target_mode() == RESIDENT + ) + assert ( + 0 + < streamed.host_materialize_scratch_bytes + < resident.target_device_weight_bytes + ) def test_virtual_layerwise_frontier_includes_host_pin_prefixes(self): module = _FakeLazyLayerwiseDit(num_layers=3) From 49b1f37139c2da4699d3279e4e0f84c46b29b277 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 20:22:38 +0800 Subject: [PATCH 46/70] [diffusion] account residency transition VRAM --- .../runtime/managers/memory_managers/auto_residency.py | 4 +++- .../multimodal_gen/test/unit/test_auto_residency.py | 8 +++++++- 2 files changed, 10 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 41b0012174ee..0ab8293807d2 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 @@ -2105,7 +2105,9 @@ def collect_residency_targets( target_resident_weight_bytes=weight_bytes, h2d_bytes_per_request=(maximum_transfer_work or weight_bytes), permanent_residency=True, - device_transition_delta_bytes=0, + device_transition_delta_bytes=( + 0 if current_resident else weight_bytes + ), active_device_delta_bytes=0, inactive_device_delta_bytes=( 0 if current_resident else weight_bytes 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 21b37f2a15cf..c81731e69223 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -3647,7 +3647,13 @@ def test_component_offload_frontier_can_expand_layerwise_lazily(self): candidate.target_layerwise_resident_layers == (2,) for candidate in candidates ) - assert any(candidate.target_mode() == RESIDENT for candidate in candidates) + resident = next( + candidate for candidate in candidates if candidate.target_mode() == RESIDENT + ) + assert ( + resident.device_transition_delta_bytes + == resident.target_device_weight_bytes + ) def test_virtual_layerwise_prefetch_is_runtime_not_transition_memory(self): module = _FakeLazyLayerwiseDit(num_layers=3) From 19f32615e23d8fbcdbbc8bd9c1856e3c765bc51b Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 20:29:33 +0800 Subject: [PATCH 47/70] [diffusion] align residency lifecycle footprints --- .../memory_managers/auto_residency.py | 28 +++---------------- .../test/unit/test_auto_residency.py | 27 +++++++++++++++++- 2 files changed, 30 insertions(+), 25 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 0ab8293807d2..84c644ab4ad7 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 @@ -190,8 +190,8 @@ class ResidencyTarget(msgspec.Struct, frozen=True): # but is not the semantic owner of this phase. present_device_delta_bytes: int = 0 # Delta for a prefetched component while another component owns the phase. - # Component offload suppresses that overlap; resident and layerwise targets - # use the same footprint as an ordinary active use. + # The ordinary strategies use their active footprint during prefetch, so + # None reuses active_device_delta_bytes. concurrent_prefetch_device_delta_bytes: int | None = None inactive_device_delta_bytes: int = 0 # None preserves the historical derived target for hand-built callers: @@ -1865,11 +1865,6 @@ def _unconfigured_layerwise_targets( + active_managed_bytes - full_weight_bytes ), - concurrent_prefetch_device_delta_bytes=( - unmanaged_weight_bytes - + active_managed_bytes - - (full_weight_bytes if current_resident else 0) - ), present_device_delta_bytes=( unmanaged_weight_bytes + active_managed_bytes @@ -2093,9 +2088,6 @@ def collect_residency_targets( present_device_delta_bytes=( -weight_bytes if current_resident else 0 ), - concurrent_prefetch_device_delta_bytes=( - -weight_bytes if current_resident else 0 - ), current_placement=not current_resident, ), ResidencyTarget( @@ -2113,9 +2105,6 @@ def collect_residency_targets( 0 if current_resident else weight_bytes ), present_device_delta_bytes=0, - concurrent_prefetch_device_delta_bytes=( - 0 if current_resident else weight_bytes - ), current_placement=current_resident, target_device_weight_bytes=weight_bytes, ), @@ -2218,7 +2207,7 @@ def collect_residency_targets( 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 + -full_weight_bytes if current_permanent else 0 ), active_device_delta_bytes=( managed_weight_bytes - current_peak_device_bytes @@ -2226,9 +2215,6 @@ def collect_residency_targets( present_device_delta_bytes=( managed_weight_bytes - current_peak_device_bytes ), - concurrent_prefetch_device_delta_bytes=( - -unmanaged_weight_bytes - current_peak_device_bytes - ), inactive_device_delta_bytes=-current_inactive_device_bytes, target_device_weight_bytes=0, target_pinned_host_bytes=0, @@ -2321,9 +2307,6 @@ def collect_residency_targets( active_device_delta_bytes=( target_peak_device_bytes - current_peak_device_bytes ), - concurrent_prefetch_device_delta_bytes=( - target_peak_device_bytes - current_peak_device_bytes - ), inactive_device_delta_bytes=( -current_inactive_device_bytes ), @@ -2383,7 +2366,7 @@ 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 + 0 if current_permanent else full_weight_bytes ), permanent_residency=True, active_device_delta_bytes=( @@ -2397,9 +2380,6 @@ def collect_residency_targets( present_device_delta_bytes=( managed_weight_bytes - current_inactive_device_bytes ), - concurrent_prefetch_device_delta_bytes=( - managed_weight_bytes - current_peak_device_bytes - ), current_placement=( current_permanent and permanent_pins == current_pinned_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 c81731e69223..cc7f6ae89a3d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -3056,6 +3056,7 @@ def test_layerwise_runtime_floor_uses_only_the_streaming_window(self): } ) module = _FakeLayerwiseDit([manager]) + module.register_parameter("unmanaged", nn.Parameter(torch.zeros(8))) runtime = component_runtime_weight_bytes({"transformer": module}) @@ -3654,6 +3655,9 @@ def test_component_offload_frontier_can_expand_layerwise_lazily(self): resident.device_transition_delta_bytes == resident.target_device_weight_bytes ) + assert ( + resident.concurrent_prefetch_delta() == resident.active_device_delta_bytes + ) def test_virtual_layerwise_prefetch_is_runtime_not_transition_memory(self): module = _FakeLazyLayerwiseDit(num_layers=3) @@ -3674,6 +3678,9 @@ def test_virtual_layerwise_prefetch_is_runtime_not_transition_memory(self): ) assert streamed.device_transition_delta_bytes == 0 assert streamed.active_device_delta_bytes < 0 + assert ( + streamed.concurrent_prefetch_delta() == streamed.active_device_delta_bytes + ) def test_lazy_layerwise_frontier_keeps_buffers_out_of_managed_weights(self): module = _FakeLazyLayerwiseDit(num_layers=2) @@ -3878,8 +3885,17 @@ def test_layerwise_frontier_accounts_for_immediate_materialization(self): resident = next( candidate for candidate in candidates if candidate.permanent_residency ) + component = next( + candidate + for candidate in candidates + if candidate.target_mode() == COMPONENT_OFFLOAD + ) + assert ( + component.concurrent_prefetch_delta() == component.active_device_delta_bytes + ) assert ( - resident.device_transition_delta_bytes == manager.offloaded_weight_bytes() + resident.device_transition_delta_bytes + == resident.target_device_weight_bytes ) module.disable_offload() candidates = collect_residency_targets( @@ -3898,6 +3914,15 @@ def test_layerwise_frontier_accounts_for_immediate_materialization(self): assert ( streamed.device_transition_delta_bytes == -manager.offloaded_weight_bytes() ) + component = next( + candidate + for candidate in candidates + if candidate.target_mode() == COMPONENT_OFFLOAD + ) + assert ( + component.device_transition_delta_bytes + == -resident.target_device_weight_bytes + ) class TestApplyAndRollback: From 5cd449a73d2542065cb8863da556db8be47c1dd8 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 21:02:03 +0800 Subject: [PATCH 48/70] test(diffusion): restore residency frontier coverage --- .../multimodal_gen/test/unit/test_auto_residency.py | 8 +------- 1 file changed, 1 insertion(+), 7 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 cc7f6ae89a3d..3d373d1709c0 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -691,12 +691,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 @@ -3805,7 +3799,7 @@ def test_resident_seed_preserves_unconfigured_layerwise_frontier(self): candidate for candidate in candidates if candidate.target_mode() == RESIDENT ).current_placement - def test_virtual_layerwise_prefetch_is_runtime_not_transition_memory(self): + def test_virtual_layerwise_reuses_component_offload_host_weights(self): module = _FakeLazyLayerwiseDit(num_layers=3) candidates = collect_residency_targets( modules={"transformer": module}, From bc593b3c111fd53166e53eaae67ac8ac4e12ea3c Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 21:53:39 +0800 Subject: [PATCH 49/70] fix(diffusion): reconcile residency calibration tests --- .../runtime/managers/gpu_worker.py | 8 +++-- .../memory_managers/auto_residency.py | 34 ++++++++++++++----- .../test/unit/test_auto_residency.py | 5 ++- .../test/unit/test_auto_residency_warmup.py | 6 ++-- .../test/unit/test_performance_metrics.py | 2 +- 5 files changed, 39 insertions(+), 16 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index ea29f0691fa2..bf7acf1b0072 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -864,13 +864,17 @@ def _record_server_warmup_memory( step_duration_ms_by_stage=( { stage_name: tuple(durations) - for stage_name, durations in metrics.steps_by_stage.items() + for stage_name, durations in getattr( + metrics, "steps_by_stage", {} + ).items() } if metrics is not None else {} ), stage_iterations=( - dict(metrics.stage_iterations) if metrics is not None else {} + dict(getattr(metrics, "stage_iterations", {})) + if metrics is not None + else {} ), ) ) 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 84c644ab4ad7..834f014aca01 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 @@ -553,6 +553,13 @@ def _repeated_stages( if is_dit_component_name(component_name) for stage_name in stage_names } + stages.update( + stage_name + for stage_names in component_stages.values() + for stage_name in stage_names + if stage_name.endswith("DenoisingStage") + and not stage_name.endswith("BeforeDenoisingStage") + ) stages.update(record.stage_iterations) stages.update( stage_name @@ -628,17 +635,28 @@ def estimate_default_workload_timing( repeated_stages=repeated_stages, target_num_inference_steps=target_num_inference_steps, ) - step_durations = representative.step_duration_ms_by_stage.get(stage_name, ()) + stage_step_durations = representative.step_duration_ms_by_stage.get( + stage_name, () + ) + step_durations = stage_step_durations 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 - ) + if stage_step_durations and len(step_durations) > 1: + target_duration_ms = ( + non_step_ms + + step_durations[0] + + statistics.median(step_durations[1:]) + * max(0, target_iterations - 1) + ) + else: + steady_steps = ( + step_durations[1:] if len(step_durations) > 1 else 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)) @@ -2176,7 +2194,7 @@ def collect_residency_targets( uses_per_streamed_layer=uses_per_request, layer_uses=layer_uses, ) - full_weight_bytes = _module_weight_bytes(module) + full_weight_bytes = max(_module_weight_bytes(module), managed_weight_bytes) unmanaged_weight_bytes = max(0, full_weight_bytes - managed_weight_bytes) component_transfer_work = ( PAGEABLE_H2D_COST_MULTIPLIER * full_weight_bytes if component_used else 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 3d373d1709c0..d510fd7b08d1 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -3054,7 +3054,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.unmanaged.numel() * module.unmanaged.element_size() + + manager.peak_managed_device_weight_bytes() + ) class TestCollectResidencyTargets: 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 0aeed812ef5e..4a8ddace4da6 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 @@ -376,7 +376,7 @@ def test_worker_keeps_resident_only_round_despite_noisy_duration(self): worker._rollback_everywhere.assert_not_called() self.assertEqual(worker._auto_residency_round_sizes, []) - def test_dit_residency_keeps_multi_step_validation(self): + def test_dit_residency_supports_short_validation(self): worker = GPUWorker.__new__(GPUWorker) worker.server_args = SimpleNamespace(residency_mode=lambda _name: RESIDENT) worker._auto_residency_applied = [ @@ -386,9 +386,7 @@ def test_dit_residency_keeps_multi_step_validation(self): worker._auto_residency_repeated_components = set() self.assertTrue(worker._latest_auto_residency_round_is_resident_only()) - self.assertFalse( - worker._latest_auto_residency_round_supports_short_validation() - ) + self.assertTrue(worker._latest_auto_residency_round_supports_short_validation()) def test_validation_rolls_back_when_selected_placement_spends_the_reserve(self): worker = GPUWorker.__new__(GPUWorker) 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 8376e0fa91e8..ec1e97706528 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -332,7 +332,6 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 worker.is_output_rank = True - worker._auto_residency_budget_correction_bytes = 2 * GIB_BYTES worker._auto_residency_reference_request_duration_ns = None worker._auto_residency_reference_stage_duration_ns = {} worker._auto_residency_reference_component_stages = {} @@ -386,6 +385,7 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): num_frames=1, baseline_allocated_bytes=1, peak_allocated_bytes=2, + peak_reserved_bytes=2 * GIB_BYTES + 2, succeeded=True, ) ], From fecc30983c3ebc73af8d96ceaebe1fe938314d75 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 21:59:07 +0800 Subject: [PATCH 50/70] [diffusion] align preload planner tests --- .../test/unit/test_auto_residency_warmup.py | 14 -------------- .../test/unit/test_performance_metrics.py | 6 ++++-- 2 files changed, 4 insertions(+), 16 deletions(-) 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 0aeed812ef5e..22b43b3cfe42 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 @@ -376,20 +376,6 @@ def test_worker_keeps_resident_only_round_despite_noisy_duration(self): worker._rollback_everywhere.assert_not_called() self.assertEqual(worker._auto_residency_round_sizes, []) - def test_dit_residency_keeps_multi_step_validation(self): - worker = GPUWorker.__new__(GPUWorker) - worker.server_args = SimpleNamespace(residency_mode=lambda _name: RESIDENT) - worker._auto_residency_applied = [ - AppliedResidencyChange("transformer", COMPONENT_OFFLOAD) - ] - worker._auto_residency_round_sizes = [1] - worker._auto_residency_repeated_components = set() - - self.assertTrue(worker._latest_auto_residency_round_is_resident_only()) - self.assertFalse( - worker._latest_auto_residency_round_supports_short_validation() - ) - def test_validation_rolls_back_when_selected_placement_spends_the_reserve(self): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 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 8376e0fa91e8..a30cff0ef035 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -230,6 +230,8 @@ def test_server_warmup_normalizes_profile_stage_names_for_residency_timing(): total_duration_ms=12.0, stages={"RefinerStage": 10.0}, steps=[], + steps_by_stage={}, + stage_iterations={}, ) ) @@ -332,7 +334,6 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 worker.is_output_rank = True - worker._auto_residency_budget_correction_bytes = 2 * GIB_BYTES worker._auto_residency_reference_request_duration_ns = None worker._auto_residency_reference_stage_duration_ns = {} worker._auto_residency_reference_component_stages = {} @@ -385,7 +386,8 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): height=64, num_frames=1, baseline_allocated_bytes=1, - peak_allocated_bytes=2, + peak_allocated_bytes=60 * GIB_BYTES, + peak_reserved_bytes=62 * GIB_BYTES, succeeded=True, ) ], From 39ebb602207b30b941e47dc69c08bf5a3c96f240 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 28 Aug 2026 22:10:49 +0800 Subject: [PATCH 51/70] [diffusion] preserve preload preference cost --- .../managers/memory_managers/placement_budget.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 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 d5896b78baea..f031ca206e0a 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 @@ -199,12 +199,12 @@ def _prune_unconstrained_states( def _optimize_unconstrained( grouped: Mapping[str, list[PlacementOption]], *, - placement_cost_dimensions: int, + preference_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 + zero_cost = (0,) * preference_dimensions frontier: list[_UnconstrainedState] = [(0, zero_cost, ())] maximum_utility = 0 @@ -218,7 +218,7 @@ def _optimize_unconstrained( group_states = [ ( group_maximum - option.estimated_latency_savings, - option.placement_cost_bytes or zero_cost, + option.preference_cost or zero_cost, (option,), ) for option in alternatives @@ -350,7 +350,7 @@ def optimize_placement( if not resource_names: best = _optimize_unconstrained( grouped, - placement_cost_dimensions=placement_cost_dimensions, + preference_dimensions=preference_dimensions, estimated_latency_tolerance=estimated_latency_tolerance, require_selection_from_every_group=require_selection_from_every_group, ) @@ -363,7 +363,7 @@ def optimize_placement( return PlacementPlan( resource_delta_bytes=full_resource_deltas, estimated_latency_savings=best[1], - placement_cost_bytes=best[2], + preference_cost=best[2], selections=list(best[3]), ) From a7430f5f160bba187bfdbb09383d732c12b5314a Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Aug 2026 00:02:09 +0800 Subject: [PATCH 52/70] [diffusion] defer HostPin tuning until calibration --- .../multimodal_gen/runtime/managers/gpu_worker.py | 9 +++++++-- .../managers/memory_managers/auto_residency.py | 6 +++--- .../test/unit/test_auto_residency.py | 14 ++++++++++++-- 3 files changed, 22 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 d20a42524b43..fd91b8141d9e 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1675,6 +1675,7 @@ def _collect_auto_residency_targets( self, workload: DefaultWorkload, *, + allow_host_pin_reallocation: bool = True, used_components: set[str] | None = None, layerwise_layer_uses: dict[str, dict[str, tuple[int, ...]]] | None = None, host_pin_headroom_bytes: int | None = None, @@ -1693,7 +1694,7 @@ def _collect_auto_residency_targets( explicit_residency_mode_of=self.server_args.explicit_residency_mode, custom_strategy_names=self._fixed_custom_residency_strategy_names(), num_inference_steps=workload.num_inference_steps, - allow_host_pin_reallocation=True, + allow_host_pin_reallocation=allow_host_pin_reallocation, mixed_dtype_components=self._mixed_dtype_residency_components(), required_resident_components={ name @@ -1734,7 +1735,11 @@ def _build_pre_warmup_auto_residency_report( ) budget_bytes = self._auto_residency_budget_bytes() - all_candidates = self._collect_auto_residency_targets(workload) + all_candidates = self._collect_auto_residency_targets( + workload, + # static weights cannot score pin placement; keep the current layout + allow_host_pin_reallocation=False, + ) candidates = pre_warmup_residency_targets( all_candidates, excluded_components=(self.pipeline.preload_residency_excluded_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 ce77d010ca66..68cbc0a422f7 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 @@ -349,9 +349,9 @@ def pre_warmup_residency_targets( """Keep only weight-safe choices before activation memory is measured. Static planning knows complete component weights but not the runtime - working set. It may demote an unsafe configured layout or redistribute - HostPin at the same device footprint, but a pure residency promotion can - consume activation headroom that only the first shape probe can reveal. + working set. It may demote an unsafe configured layout, but a pure + residency promotion can consume activation headroom that only the first + shape probe can reveal. A layerwise layout can require a small persistent non-layer footprint even when the coarse component-offload state keeps every weight on the host. 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 9cfa6f716bcd..469552ad4412 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -1487,6 +1487,12 @@ def test_pre_warmup_report_uses_weight_model_below_legacy_threshold( current_placement=True, ) modules = {"transformer": object()} + collect_calls = [] + + def collect_targets(_workload, *, allow_host_pin_reallocation=True): + collect_calls.append(allow_host_pin_reallocation) + return [current] + worker = SimpleNamespace( rank=0, pipeline=SimpleNamespace(preload_residency_excluded_components=frozenset()), @@ -1497,7 +1503,7 @@ def test_pre_warmup_report_uses_weight_model_below_legacy_threshold( host_pin_budget=lambda: None, ), _auto_residency_budget_bytes=lambda: 30 * GIB_BYTES, - _collect_auto_residency_targets=lambda _workload: [current], + _collect_auto_residency_targets=collect_targets, _auto_residency_modules=lambda: modules, ) monkeypatch.setattr( @@ -1541,6 +1547,7 @@ def test_pre_warmup_report_uses_weight_model_below_legacy_threshold( assert report.skip_reason is None assert report.estimated_request_duration_ns == 1 assert report.candidate_latency_savings_ns == {current.option_key(): 1} + assert collect_calls == [False] def test_pre_warmup_report_keeps_excluded_components_fixed(self, monkeypatch): from sglang.multimodal_gen.runtime.managers import gpu_worker @@ -1576,7 +1583,10 @@ def test_pre_warmup_report_keeps_excluded_components_fixed(self, monkeypatch): host_pin_budget=lambda: None, ), _auto_residency_budget_bytes=lambda: 80 * GIB_BYTES, - _collect_auto_residency_targets=lambda _workload: [transformer, vae], + _collect_auto_residency_targets=lambda _workload, **_kwargs: [ + transformer, + vae, + ], _auto_residency_modules=lambda: modules, ) monkeypatch.setattr( From 5e6db41b6da56941ff915fc6f054f70eacf4460a Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Aug 2026 01:34:39 +0800 Subject: [PATCH 53/70] test(diffusion): align merged residency fixtures --- .../test/unit/test_auto_residency_warmup.py | 1 + .../test/unit/test_initial_residency.py | 1 + .../test/unit/test_performance_metrics.py | 2 ++ .../test/unit/test_server_args.py | 19 ++----------------- 4 files changed, 6 insertions(+), 17 deletions(-) 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 22b43b3cfe42..cc402ebf3d9f 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 @@ -51,6 +51,7 @@ def supports_auto_residency(self) -> bool: quantization=None, direct_gpu_weight_loading=False, nunchaku_config=None, + ltx2_two_stage_device_mode=None, ) self.assertEqual(worker._fixed_custom_residency_strategy_names(), {"fixed"}) diff --git a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py index 7942c03f5837..bf6c244da7db 100644 --- a/python/sglang/multimodal_gen/test/unit/test_initial_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_initial_residency.py @@ -40,6 +40,7 @@ def __init__(self, modes=None, explicit=()): self.component_quantizations = {} self.nunchaku_config = None self.direct_gpu_weight_loading = False + self.ltx2_two_stage_device_mode = None def residency_mode(self, component_name): return self.selected.get( 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 ecc249ba0d7e..dabb5c99f8b2 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -302,6 +302,7 @@ def test_loaded_prequantized_checkpoint_can_use_auto_residency(): quantization=None, direct_gpu_weight_loading=False, nunchaku_config=None, + ltx2_two_stage_device_mode=None, ) device_module = Mock() device_module.mem_get_info.return_value = (100, 100) @@ -362,6 +363,7 @@ def test_auto_residency_budget_respects_test_device_memory_cap(monkeypatch): quantization=None, direct_gpu_weight_loading=False, nunchaku_config=None, + ltx2_two_stage_device_mode=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_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index ea6fcb263198..adc2f93baffc 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -1595,7 +1595,7 @@ def test_qwen_ar_planner_candidates_do_not_use_a_coarse_vram_threshold(self): "text_encoder", constrained_args.layerwise_offload_components or [] ) - def test_auto_multi_gpu_sana_wm_defers_fsdp_and_enables_cfg_parallel(self): + def test_auto_multi_gpu_sana_wm_prefers_fsdp_and_cfg_parallel(self): args = self._from_dict_with_pipeline_config( SanaWMPipelineConfig(), kwargs={ @@ -1605,22 +1605,7 @@ def test_auto_multi_gpu_sana_wm_defers_fsdp_and_enables_cfg_parallel(self): }, ) - self.assertFalse(args.use_fsdp_inference) - self.assertTrue(args.enable_cfg_parallel) - - def test_server_warmup_keeps_sana_wm_dit_available_to_residency_planner(self): - args = self._from_dict_with_pipeline_config( - SanaWMPipelineConfig(), - kwargs={ - "model_path": "Efficient-Large-Model/SANA-WM_bidirectional", - "num_gpus": 2, - "performance_mode": "auto", - "warmup_mode": "server", - }, - ) - - self.assertFalse(args.use_fsdp_inference) - self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD) + self.assertTrue(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) def test_cache_dit_rejects_explicit_fsdp(self): From 6613ad6722b02e74a11a59331b2616f5a8553bfc Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 30 Aug 2026 02:27:58 +0800 Subject: [PATCH 54/70] fix: use current transformer checkpoint resolver --- .../runtime/pipelines_core/composed_pipeline_base.py | 6 +++--- .../test/unit/test_component_loading_order.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index d39560af0fbc..307dc8c074fc 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -24,7 +24,7 @@ ) from sglang.multimodal_gen.runtime.loader.gguf_weights import names_gguf_checkpoint from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( - resolve_transformer_safetensors_to_load, + resolve_transformer_checkpoint_files, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_loading_order import ( ComponentLoadSpec, @@ -453,9 +453,9 @@ def _component_weight_source( ): component_server_args = copy.copy(server_args) component_server_args.transformer_weights_path = weights_path - selected_files = resolve_transformer_safetensors_to_load( + selected_files = resolve_transformer_checkpoint_files( component_server_args, spec.component_model_path - ) + ).safetensors checkpoint_bytes, parameter_count = infer_safetensors_weight_stats( selected_files ) diff --git a/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py b/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py index 0220aa93ebb6..6a4ba8e40891 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py @@ -277,8 +277,8 @@ def test_preload_inventory_uses_selected_transformer_safetensors(tmp_path): with patch( "sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base." - "resolve_transformer_safetensors_to_load", - return_value=[str(mixed_path)], + "resolve_transformer_checkpoint_files", + return_value=SimpleNamespace(safetensors=(str(mixed_path),)), ): source = ComposedPipelineBase._component_weight_source( _spec("transformer", 0, "/model/transformer"), server_args From f53b0eebd1c7b2ae4e9c27211d2c1f128535bb20 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 31 Aug 2026 11:29:53 +0800 Subject: [PATCH 55/70] fix(diffusion): align residency planner test scope --- .../test/unit/test_server_args.py | 38 ++++++++----------- 1 file changed, 16 insertions(+), 22 deletions(-) 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 6543a205260e..669b24034b53 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -1655,31 +1655,25 @@ def test_sana_wm_residency_scales_with_available_memory(self): self.assertIn("text_encoder", constrained_offload) self.assertIn("vae", constrained_offload) - def test_qwen_ar_planner_candidates_do_not_use_a_coarse_vram_threshold(self): - pipeline_configs = ( - QwenImageLayeredPipelineConfig(), + def test_longcat_ar_planner_candidates_do_not_use_a_coarse_vram_threshold(self): + high_memory_args = self._from_dict_with_pipeline_config( LongCatImagePipelineConfig(), + memory_gb=80, + kwargs={"performance_mode": "auto"}, ) + self.assertIn( + "text_encoder", high_memory_args.layerwise_offload_components or [] + ) + self.assertFalse(high_memory_args.text_encoder_cpu_offload) - for pipeline_config in pipeline_configs: - high_memory_args = self._from_dict_with_pipeline_config( - pipeline_config, - memory_gb=80, - kwargs={"performance_mode": "auto"}, - ) - self.assertIn( - "text_encoder", high_memory_args.layerwise_offload_components or [] - ) - self.assertFalse(high_memory_args.text_encoder_cpu_offload) - - constrained_args = self._from_dict_with_pipeline_config( - pipeline_config, - memory_gb=60, - kwargs={"performance_mode": "auto"}, - ) - self.assertIn( - "text_encoder", constrained_args.layerwise_offload_components or [] - ) + constrained_args = self._from_dict_with_pipeline_config( + LongCatImagePipelineConfig(), + memory_gb=60, + kwargs={"performance_mode": "auto"}, + ) + self.assertIn( + "text_encoder", constrained_args.layerwise_offload_components or [] + ) def test_auto_multi_gpu_sana_wm_prefers_fsdp_and_cfg_parallel(self): args = self._from_dict_with_pipeline_config( From 404541082c8dd1978cc4dffffcf208a329ec0677 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 31 Aug 2026 16:08:20 +0800 Subject: [PATCH 56/70] test(diffusion): complete loader metadata fixtures --- .../test/unit/test_component_loader_identity.py | 1 + .../multimodal_gen/test/unit/test_component_loading_order.py | 5 ++++- python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py | 1 + 3 files changed, 6 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py b/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py index a70472e57ddf..7243e2d8785c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py @@ -90,6 +90,7 @@ def test_declared_alias_loads_by_exact_key_and_structural_source(self): } server_args = SimpleNamespace( component_paths={}, + component_weights_paths={}, component_direct_gpu_weight_loading=set(), resolve_component_attention_backend=lambda *_names: (None, None), ) diff --git a/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py b/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py index 6a4ba8e40891..c66ebf764161 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_loading_order.py @@ -230,6 +230,7 @@ def test_component_weight_inventory_distinguishes_config_only_and_unknown(tmp_pa def test_preload_inventory_uses_the_actual_weight_override(): server_args = SimpleNamespace( + component_precisions={}, component_weights_paths={"text_encoder": "/weights/text.safetensors"}, transformer_weights_path="/weights/dit.gguf", pipeline_config=SimpleNamespace( @@ -264,6 +265,7 @@ def test_preload_inventory_uses_selected_transformer_safetensors(tmp_path): mixed_path = tmp_path / "flux2-dev-nvfp4-mixed.safetensors" _write_safetensors(mixed_path, 64) server_args = SimpleNamespace( + component_precisions={}, component_weights_paths={}, transformer_weights_path="owner/flux2-nvfp4", revision=None, @@ -291,12 +293,13 @@ def test_preload_inventory_uses_selected_transformer_safetensors(tmp_path): def test_preload_inventory_resolves_group_precision_fallbacks(): server_args = SimpleNamespace( + component_precisions={}, pipeline_config=SimpleNamespace( dit_precision="bf16", vae_precision="fp32", image_encoder_precision="fp16", text_encoder_precisions=("bf16",), - ) + ), ) assert ( 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 1289ce762c95..16cc08d55769 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 @@ -46,6 +46,7 @@ def test_resident_mode_declares_every_extra_weight_source(self): pipeline.model_path = model_path server_args = SimpleNamespace( component_paths={}, + component_precisions={}, ltx2_two_stage_device_mode="resident", lora_path=None, pipeline_config=SimpleNamespace( From 9b5a05280ed82adf927db41220fe90987be6cb33 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 31 Aug 2026 17:37:14 +0800 Subject: [PATCH 57/70] test(diffusion): complete component loader fixtures --- .../test/unit/test_attention_backend_selector.py | 4 ++++ .../test/unit/test_component_loader_identity.py | 1 + 2 files changed, 5 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 56faa26542d5..244fdc16cba8 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 @@ -335,6 +335,7 @@ def load_customized(self, *_args): class _Args: component_precisions = {} component_quantizations = {} + component_weights_paths = {} @staticmethod def requested_component_attention_backend(_component_name): @@ -414,6 +415,7 @@ def load_customized(self, *_args): class _Args: component_precisions = {} component_quantizations = {} + component_weights_paths = {} @staticmethod def requested_component_attention_backend(_component_name): @@ -479,6 +481,7 @@ def load_native(self, *_args): class _Args: component_precisions = {} component_quantizations = {} + component_weights_paths = {} pipeline_config = SimpleNamespace(native_only_components=()) @staticmethod @@ -530,6 +533,7 @@ def load_native(self, *_args): class _Args: component_precisions = {} component_quantizations = {} + component_weights_paths = {} pipeline_config = SimpleNamespace(native_only_components=()) @staticmethod diff --git a/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py b/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py index 7243e2d8785c..979cfc323b55 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py @@ -90,6 +90,7 @@ def test_declared_alias_loads_by_exact_key_and_structural_source(self): } server_args = SimpleNamespace( component_paths={}, + component_precisions={}, component_weights_paths={}, component_direct_gpu_weight_loading=set(), resolve_component_attention_backend=lambda *_names: (None, None), From 8714e70b026e0a74ad0b731ec2f440def9012f8e Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 1 Sep 2026 05:39:18 +0800 Subject: [PATCH 58/70] fix(diffusion): scope LTX warmup exclusions --- .../runtime/server_args/auto_tune.py | 14 +++++++------- .../test/unit/test_component_loader_identity.py | 1 + 2 files changed, 8 insertions(+), 7 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 0b58ccf4bb65..d64139b943bb 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -71,13 +71,6 @@ def auto_residency_static_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 server_args.disagg_role != RoleType.MONOLITHIC: return "disaggregated role" task_type = server_args.pipeline_config.task_type @@ -109,6 +102,13 @@ def auto_residency_args_skip_reason(server_args: ServerArgs) -> str | None: reason = auto_residency_static_skip_reason(server_args) if reason is not None: return reason + 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 server_args.warmup_mode != "server": return "no synthetic server warmup to calibrate from" return None diff --git a/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py b/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py index 979cfc323b55..476c7369519c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py @@ -93,6 +93,7 @@ def test_declared_alias_loads_by_exact_key_and_structural_source(self): component_precisions={}, component_weights_paths={}, component_direct_gpu_weight_loading=set(), + pipeline_config=SimpleNamespace(), resolve_component_attention_backend=lambda *_names: (None, None), ) From 0fa3496816632500a9e4788a7c6e88604a92653e Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 1 Sep 2026 06:57:38 +0800 Subject: [PATCH 59/70] test(diffusion): opt out loader fixture from residency planning --- .../multimodal_gen/test/unit/test_component_loader_identity.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py b/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py index 476c7369519c..edcfd94afe89 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py @@ -89,6 +89,7 @@ def test_declared_alias_loads_by_exact_key_and_structural_source(self): "scheduler": ["diffusers", "Scheduler"], } server_args = SimpleNamespace( + performance_mode="manual", component_paths={}, component_precisions={}, component_weights_paths={}, From a4614c0a5672c19ab21d14c34ac6bc7c57f6ede3 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 1 Sep 2026 08:45:25 +0800 Subject: [PATCH 60/70] fix(diffusion): contain preload residency regressions --- .../runtime/pipelines/qwen_image.py | 16 ++++++++++++++ .../test/server/perf_baselines/h100.json | 4 ++-- .../test/unit/test_auto_residency.py | 22 +++++++++++++++++++ 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py b/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py index 5b08bf644dc3..b0b7a9bc657e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py @@ -3,6 +3,9 @@ # SPDX-License-Identifier: Apache-2.0 from diffusers.image_processor import VaeImageProcessor +from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( + QwenImageEditPlus_2511_PipelineConfig, +) from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( @@ -84,6 +87,19 @@ def create_pipeline_stages(self, server_args: ServerArgs): class QwenImageEditPlusPipeline(QwenImageEditPipeline): pipeline_name = "QwenImageEditPlusPipeline" + @property + def preload_residency_excluded_components(self) -> frozenset[str]: + # The 2511 DiT fits on an H100, but making it resident before warmup + # causes reversible-frontier construction to exceed server startup. + # Keep its proven load path; post-warmup calibration can still promote + # it once the real request shape is known. + if isinstance( + self.server_args.pipeline_config, + QwenImageEditPlus_2511_PipelineConfig, + ): + return frozenset(("transformer",)) + return frozenset() + def prepare_mu_layered(batch: Req, server_args: ServerArgs): base_seqlen = 256 * 256 / 16 / 16 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 6090320b1834..ea84093726d4 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -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": 49998.0, + "runtime_peak_vram_mb": 52830.0, "estimated_full_test_time_s": 243.0 }, "wan2_1_i2v_14b_720P_2gpu": { @@ -1765,7 +1765,7 @@ "expected_e2e_ms": 83297.78, "expected_avg_denoise_ms": 2029.88, "expected_median_denoise_ms": 2020.71, - "load_peak_vram_mb": 6440.0, + "load_peak_vram_mb": 60038.0, "runtime_peak_vram_mb": 72524.0, "estimated_full_test_time_s": 204.3 }, 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 ca584ad69f5f..48d6b2c0db8e 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,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType from sglang.multimodal_gen.configs.pipeline_configs.longlive2 import LongLive2T2VConfig +from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( + QwenImageEditPlus_2511_PipelineConfig, + QwenImageEditPlusPipelineConfig, +) from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( ACTIVATION_EXTRAPOLATION_MARGIN, GIB_BYTES, @@ -111,6 +115,24 @@ def _record( ) +def test_qwen_image_2511_excludes_transformer_only_during_preload(): + from sglang.multimodal_gen.runtime.pipelines.qwen_image import ( + QwenImageEditPlusPipeline, + ) + + pipeline = QwenImageEditPlusPipeline.__new__(QwenImageEditPlusPipeline) + + pipeline.server_args = SimpleNamespace( + pipeline_config=QwenImageEditPlus_2511_PipelineConfig() + ) + assert pipeline.preload_residency_excluded_components == frozenset(("transformer",)) + + pipeline.server_args = SimpleNamespace( + pipeline_config=QwenImageEditPlusPipelineConfig() + ) + assert pipeline.preload_residency_excluded_components == frozenset() + + class TestEstimateAllocatorHeadroom: def test_uses_covering_workload_allocator_gap(self): small = _record( From a5ba54e3233044215e4399565bd4e4b20ef366db Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 1 Sep 2026 10:11:11 +0800 Subject: [PATCH 61/70] test(diffusion): align residency CI workloads --- .../multimodal_gen/test/server/perf_baselines/h100.json | 4 ++-- python/sglang/multimodal_gen/test/server/testcase_configs.py | 3 ++- 2 files changed, 4 insertions(+), 3 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 ea84093726d4..4beb929d5333 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -1718,7 +1718,7 @@ "LatentPreparationStage": 0.14, "TimestepPreparationStage": 1.36, "DenoisingStage": 81747.14, - "DecodingStage": 784.99 + "DecodingStage": 2146.55 }, "denoise_step_ms": { "0": 1453.25, @@ -2794,7 +2794,7 @@ }, "expected_e2e_ms": 15974.21, "expected_avg_denoise_ms": 802.06, - "expected_median_denoise_ms": 747.62, + "expected_median_denoise_ms": 946.4, "load_peak_vram_mb": 56048.0, "runtime_peak_vram_mb": 68712.0, "estimated_full_test_time_s": 363.2 diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 77300ee911d2..725f4d198602 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -733,7 +733,8 @@ def from_req_perf_record( direct_url_test=True, output_size="384x640", num_frames=17, - extras={"num_inference_steps": 12, "seed": 0, "guidance_scale": 4.5}, + # Keep the URL request explicit while matching the canonical GT defaults. + extras={"num_inference_steps": 20, "seed": 42, "guidance_scale": 4.5}, ) LONGLIVE2_T2V_CI_sampling_params = replace( From d6f9883a7080a74168aea82d0cc5904c378ddf87 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 1 Sep 2026 14:45:32 +0800 Subject: [PATCH 62/70] fix(diffusion): bound auto residency frontier search --- .../memory_managers/auto_residency.py | 379 +++++++++++++----- .../test/unit/test_auto_residency.py | 65 +++ 2 files changed, 334 insertions(+), 110 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 bac9c42f9104..49922e2374a0 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,6 +109,15 @@ MIN_POST_ADJUSTMENT_REGRESSION_NS = 100_000_000 POST_ADJUSTMENT_REGRESSION_FRACTION = 0.05 +# Layerwise residency is a startup-time search, not an exhaustive knapsack +# benchmark. Large heterogeneous models can otherwise create millions of +# measured resident/HostPin combinations and miss the server startup timeout +# before the optimizer runs. Keep a deterministic, resource-safe sampling of +# the frontier; every retained state still uses its exact byte accounting. +MAX_LAYERWISE_RESIDENT_TARGETS = 64 +MAX_LAYERWISE_PIN_TARGETS = 64 +MAX_LAYERWISE_COMPONENT_TARGETS = 4096 + PLACEMENT_STATUS_SKIPPED = "skipped" PLACEMENT_STATUS_ADJUSTED = "adjusted" PLACEMENT_STATUS_VALIDATED = "validated" @@ -1361,13 +1370,16 @@ def _layerwise_pin_targets( layer_uses: tuple[tuple[int, ...], ...] | None = None, constrain_host_transitions: bool = True, maximum_utility_only: bool = False, + max_targets: int | None = None, ) -> list[tuple[tuple[int, ...], ...]]: - """Pareto-optimal HostPin targets for one resident-layer placement. + """Useful HostPin targets for one resident-layer placement. 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. + packings needed when layer sizes differ. Large heterogeneous frontiers are + deterministically sampled so startup work remains bounded; retained states + keep their exact resource costs and transfer utility. """ policies = _resolve_layerwise_policies(managers, residency_policies) resolved_uses = _resolve_layerwise_uses( @@ -1452,7 +1464,57 @@ def dominates(left, right) -> bool: and left[3] <= right[3] ) + def limit(states): + if max_targets is None or len(states) <= max_targets: + return states + + limit_count = max(1, max_targets) + orderings = ( + sorted( + states, + key=lambda item: (item[0], item[2], item[3], -item[1], item[4]), + ), + sorted( + states, + key=lambda item: ( + max(item[2], item[3]), + item[2] + item[3], + item[0], + -item[1], + item[4], + ), + ), + sorted( + states, + key=lambda item: (-item[1], item[0], item[2], item[3], item[4]), + ), + ) + selected = [] + selected_targets = set() + + def add(state): + if state[4] in selected_targets or len(selected) >= limit_count: + return + selected.append(state) + selected_targets.add(state[4]) + + for ordering in orderings: + add(ordering[0]) + add(ordering[-1]) + + sample_count = max(2, limit_count) + for sample_index in range(sample_count): + for ordering in orderings: + index = round(sample_index * (len(ordering) - 1) / (sample_count - 1)) + add(ordering[index]) + return selected + def prune(candidates): + # Exact multidimensional dominance is quadratic in the frontier size. + # Pre-limit only after the raw expansion is already much larger than + # the requested search budget; small frontiers remain fully exact. + if max_targets is not None and len(candidates) > max_targets * 8: + candidates = limit(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 @@ -1483,7 +1545,7 @@ def prune(candidates): continue frontier.append(state) best_utility = state[1] - return frontier + return limit(frontier) best_by_cost = {} for state in candidates: @@ -1502,7 +1564,7 @@ def prune(candidates): existing for existing in frontier if not dominates(state, existing) ] frontier.append(state) - return frontier + return limit(frontier) for (uses, transfer_bytes, host_bytes, is_current), layers in sorted( grouped.items(), @@ -1555,6 +1617,7 @@ def _layerwise_resident_targets( managers: Sequence, *, layer_uses: tuple[tuple[int, ...], ...] | None = None, + current_resident_layers: tuple[int, ...] | None = None, ) -> list[tuple[int, ...]]: """Useful resident-count tuples for the measured request. @@ -1572,6 +1635,32 @@ def _layerwise_resident_targets( explicit interface can represent. """ layer_counts = tuple(manager.num_layers for manager in managers) + + def limit( + targets: Collection[tuple[int, ...]], + required: Sequence[tuple[int, ...]] = (), + ) -> list[tuple[int, ...]]: + ordered = sorted(set(targets) | set(required)) + if len(ordered) <= MAX_LAYERWISE_RESIDENT_TARGETS: + return ordered + selected = [] + seen = set() + + def add(target): + if target in seen or len(selected) >= MAX_LAYERWISE_RESIDENT_TARGETS: + return + selected.append(target) + seen.add(target) + + for target in required: + add(target) + sample_count = max(2, MAX_LAYERWISE_RESIDENT_TARGETS) + for index in range(sample_count): + add(ordered[round(index * (len(ordered) - 1) / (sample_count - 1))]) + for target in ordered: + add(target) + return sorted(selected) + if layer_uses is not None: resolved_uses = _resolve_layerwise_uses( managers=managers, @@ -1582,7 +1671,67 @@ def _layerwise_resident_targets( 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)) + target_count = 1 + for group_choices in choices: + target_count *= len(group_choices) + if target_count <= MAX_LAYERWISE_RESIDENT_TARGETS: + targets = list(product(*choices)) + required = ( + (current_resident_layers,) + if current_resident_layers is not None + else () + ) + return limit(targets, required) + + repeated = tuple(len(group_choices) > 1 for group_choices in choices) + targets = {tuple(0 for _ in managers)} + maximum_count = max( + ( + num_layers + for num_layers, active in zip(layer_counts, repeated) + if active + ), + default=0, + ) + for count in range(1, maximum_count + 1): + targets.add( + tuple( + min(count, num_layers) if active else 0 + for num_layers, active in zip(layer_counts, repeated) + ) + ) + for manager_index, (num_layers, active) in enumerate( + zip(layer_counts, repeated) + ): + if active: + targets.add( + tuple( + num_layers if index == manager_index else 0 + for index in range(len(managers)) + ) + ) + if current_resident_layers is not None: + targets.add(current_resident_layers) + required = [tuple(0 for _ in managers)] + if current_resident_layers is not None: + required.append(current_resident_layers) + required.append( + tuple( + num_layers if active else 0 + for num_layers, active in zip(layer_counts, repeated) + ) + ) + required.extend( + tuple( + num_layers if index == manager_index else 0 + for index in range(len(managers)) + ) + for manager_index, (num_layers, active) in enumerate( + zip(layer_counts, repeated) + ) + if active + ) + return limit(targets, required) targets = {tuple(0 for _ in managers)} @@ -1603,7 +1752,10 @@ def _layerwise_resident_targets( tuple(max(1, int(round(ratio * num_layers))) for num_layers in layer_counts) ) - return sorted(targets) + required = [tuple(0 for _ in managers), layer_counts] + if current_resident_layers is not None: + required.insert(1, current_resident_layers) + return limit(targets, required) def _layerwise_policy_targets( @@ -1959,7 +2111,7 @@ def collect_residency_targets( request_duration_ns: int = 0, latency_upper_bound_ns_by_component: Mapping[str, int] | None = None, ) -> list[ResidencyTarget]: - """Build complete target-state frontiers for auto-managed components. + """Build bounded 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. @@ -1969,8 +2121,8 @@ def collect_residency_targets( When both host constraints provably cannot bind and the component latency 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. + Other large layerwise frontiers retain bounded samples across device, host, + transition-scratch, and transfer-utility tradeoffs. """ custom_names = set(custom_strategy_names) mixed_dtype_names = set(mixed_dtype_components) @@ -2287,122 +2439,129 @@ def collect_residency_targets( ) ) - for target_resident_layers in _layerwise_resident_targets( - managers, layer_uses=layer_uses - ): + layerwise_layouts = [ + (target_resident_layers, target_policies) + for target_resident_layers in _layerwise_resident_targets( + managers, + layer_uses=layer_uses, + current_resident_layers=current_resident_layers, + ) for target_policies in _layerwise_policy_targets( managers=managers, resident_layers=target_resident_layers, tune_policy=tune_residency_policy, + ) + ] + max_pin_targets = min( + MAX_LAYERWISE_PIN_TARGETS, + max(1, MAX_LAYERWISE_COMPONENT_TARGETS // len(layerwise_layouts)), + ) + for target_resident_layers, target_policies in layerwise_layouts: + target_resident_bytes = sum( + manager.resident_weight_bytes(count, policy) + for manager, count, policy in zip( + managers, target_resident_layers, target_policies + ) + ) + target_peak_device_bytes = _layerwise_active_peak_device_bytes( + managers=managers, + resident_layers=target_resident_layers, + residency_policies=target_policies, + layer_uses=layer_uses, + ) + pin_targets = ( + _layerwise_pin_targets( + managers=managers, + resident_layers=target_resident_layers, + residency_policies=target_policies, + current_pinned_layers=current_pinned_layers, + 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 + ), + max_targets=max_pin_targets, + ) + if allow_host_pin_reallocation + else [current_pinned_layers] + ) + if ( + target_resident_layers == current_resident_layers + and target_policies == current_residency_policies + and current_pinned_layers not in pin_targets ): - target_resident_bytes = sum( - manager.resident_weight_bytes(count, policy) - for manager, count, policy in zip( - managers, target_resident_layers, target_policies + # relative utility is anchored to the measured placement; + # keep that exact state when every other pin subset folds + # into its maximum-utility representative + pin_targets.append(current_pinned_layers) + for target_pinned_layers in pin_targets: + 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_peak_device_bytes = _layerwise_active_peak_device_bytes( + target_transfer_work = _layerwise_transfer_work_bytes( managers=managers, resident_layers=target_resident_layers, residency_policies=target_policies, + pinned_layers=target_pinned_layers, + uses_per_streamed_layer=uses_per_request, layer_uses=layer_uses, ) - pin_targets = ( - _layerwise_pin_targets( - managers=managers, - resident_layers=target_resident_layers, - residency_policies=target_policies, - current_pinned_layers=current_pinned_layers, - 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 + unpin_scratch, pin_scratch = _layerwise_host_transition_bytes( + managers=managers, + current_pinned_layers=current_pinned_layers, + target_pinned_layers=target_pinned_layers, + ) + candidates.append( + ResidencyTarget( + component_name=name, + residency_mode=frontier_mode, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=target_resident_bytes, + h2d_bytes_per_request=( + maximum_transfer_work - target_transfer_work ), + target_layerwise_resident_layers=target_resident_layers, + target_layerwise_residency_policies=( + target_policies + if tune_residency_policy + and _policy_affects_layerwise_layout( + managers, target_resident_layers + ) + else None + ), + 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, + device_transition_delta_bytes=( + -managed_weight_bytes if current_permanent else 0 + ), + active_device_delta_bytes=( + target_peak_device_bytes - current_peak_device_bytes + ), + 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_policies == current_residency_policies + and target_pinned_layers == current_pinned_layers + ), + target_device_weight_bytes=( + unmanaged_weight_bytes + target_resident_bytes + ), + target_pinned_host_bytes=target_pinned_bytes, ) - if allow_host_pin_reallocation - else [current_pinned_layers] ) - if ( - target_resident_layers == current_resident_layers - and target_policies == current_residency_policies - and current_pinned_layers not in pin_targets - ): - # relative utility is anchored to the measured placement; - # keep that exact state when every other pin subset folds - # into its maximum-utility representative - pin_targets.append(current_pinned_layers) - for target_pinned_layers in pin_targets: - 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, - residency_policies=target_policies, - 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, - current_pinned_layers=current_pinned_layers, - target_pinned_layers=target_pinned_layers, - ) - candidates.append( - ResidencyTarget( - component_name=name, - residency_mode=frontier_mode, - target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=target_resident_bytes, - h2d_bytes_per_request=( - maximum_transfer_work - target_transfer_work - ), - target_layerwise_resident_layers=target_resident_layers, - target_layerwise_residency_policies=( - target_policies - if tune_residency_policy - and _policy_affects_layerwise_layout( - managers, target_resident_layers - ) - else None - ), - 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, - device_transition_delta_bytes=( - -managed_weight_bytes if current_permanent else 0 - ), - active_device_delta_bytes=( - target_peak_device_bytes - current_peak_device_bytes - ), - 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_policies == current_residency_policies - and target_pinned_layers == current_pinned_layers - ), - target_device_weight_bytes=( - unmanaged_weight_bytes + target_resident_bytes - ), - target_pinned_host_bytes=target_pinned_bytes, - ) - ) # Layerwise stores have one fixed dtype, while both coarse strategies # cast at every declared use. A mixed-dtype component may be tuned within 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 48d6b2c0db8e..53b0c01d2de6 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 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( ACTIVATION_EXTRAPOLATION_MARGIN, GIB_BYTES, + MAX_LAYERWISE_RESIDENT_TARGETS, MIN_VRAM_RESERVE_BYTES, PAGEABLE_H2D_COST_MULTIPLIER, AppliedResidencyChange, @@ -3194,6 +3195,47 @@ def test_measured_resident_frontier_streams_one_shot_and_unused_groups(self): assert targets == [(0, 0)] + def test_measured_resident_frontier_keeps_current_one_shot_layout(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)), + current_resident_layers=(2, 1), + ) + + assert targets == [(0, 0), (2, 1)] + + def test_large_measured_resident_frontier_is_bounded_and_keeps_anchors(self): + managers = [ + _FakeLayerwiseManager( + {f"first.{index}.w": torch.zeros(16) for index in range(70)} + ), + _FakeLayerwiseManager( + {f"second.{index}.w": torch.zeros(16) for index in range(70)} + ), + ] + + targets = _layerwise_resident_targets( + managers, + layer_uses=(tuple([8] * 70), tuple([8] * 70)), + current_resident_layers=(13, 17), + ) + + assert len(targets) <= MAX_LAYERWISE_RESIDENT_TARGETS + assert (0, 0) in targets + assert (70, 70) in targets + assert (70, 0) in targets + assert (0, 70) in targets + assert (13, 17) in targets + def test_host_pin_frontier_includes_non_prefix_packings(self): manager = SimpleNamespace( num_layers=3, @@ -3234,6 +3276,29 @@ def test_host_pin_frontier_coalesces_equal_transformer_layers(self): assert len(targets) == 41 + def test_large_host_pin_frontier_is_bounded_and_keeps_extremes(self): + layer_bytes = {index: 1 << index for index in range(24)} + manager = SimpleNamespace( + num_layers=24, + 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=1, + max_targets=16, + ) + + assert len(targets) <= 16 + assert ((),) in targets + assert (tuple(layer_bytes),) in targets + def test_nonbinding_host_resources_keep_only_maximum_pin_utility(self): manager = SimpleNamespace( num_layers=3, From 80183ad1d33b6409d75dd14a0920b6953e593c2e Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 1 Sep 2026 19:10:58 +0800 Subject: [PATCH 63/70] fix(diffusion): bound residency policy search --- .../memory_managers/auto_residency.py | 41 ++++++++++++++++++- .../test/server/perf_baselines/h100.json | 2 +- .../test/unit/test_auto_residency.py | 18 ++++++++ 3 files changed, 59 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 49922e2374a0..c60d8dfcd0b9 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 @@ -115,6 +115,7 @@ # before the optimizer runs. Keep a deterministic, resource-safe sampling of # the frontier; every retained state still uses its exact byte accounting. MAX_LAYERWISE_RESIDENT_TARGETS = 64 +MAX_LAYERWISE_POLICY_TARGETS = 64 MAX_LAYERWISE_PIN_TARGETS = 64 MAX_LAYERWISE_COMPONENT_TARGETS = 4096 @@ -1780,7 +1781,45 @@ def _layerwise_policy_targets( ) for manager, count in zip(managers, resident_layers) ) - return list(dict.fromkeys([current, *product(*choices)])) + tunable_indices = tuple( + index for index, policy_choices in enumerate(choices) if len(policy_choices) > 1 + ) + target_count = 1 << len(tunable_indices) + if target_count <= MAX_LAYERWISE_POLICY_TARGETS: + return list(dict.fromkeys([current, *product(*choices)])) + + # A component may own many independently managed layer stacks. Expanding + # every leading/strided combination is exponential and can consume the + # whole server-startup timeout before the bounded resident and HostPin + # frontiers are considered. Sample the binary policy space directly by + # mask, retaining both extremes and evenly spaced deterministic layouts. + denominator = MAX_LAYERWISE_POLICY_TARGETS - 1 + masks = { + (sample_index * (target_count - 1) + denominator // 2) // denominator + for sample_index in range(MAX_LAYERWISE_POLICY_TARGETS) + } + targets = [] + seen = set() + + def add(target: tuple[str, ...]) -> None: + if target in seen or len(targets) >= MAX_LAYERWISE_POLICY_TARGETS: + return + targets.append(target) + seen.add(target) + + add(current) + ordered_masks = [0, target_count - 1, *sorted(masks - {0, target_count - 1})] + for mask in ordered_masks: + bit_index = 0 + target = [] + for policy_choices in choices: + if len(policy_choices) == 1: + target.append(policy_choices[0]) + continue + target.append(policy_choices[(mask >> bit_index) & 1]) + bit_index += 1 + add(tuple(target)) + return targets def _policy_affects_layerwise_layout( 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 4beb929d5333..27ca22000f10 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -2253,7 +2253,7 @@ "expected_e2e_ms": 237254.68, "expected_avg_denoise_ms": 32.68, "expected_median_denoise_ms": 30.2, - "load_peak_vram_mb": 876.0, + "load_peak_vram_mb": 8264.0, "runtime_peak_vram_mb": 15686.0, "estimated_full_test_time_s": 420.1 }, 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 53b0c01d2de6..af2efedd6d7f 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 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.auto_residency import ( ACTIVATION_EXTRAPOLATION_MARGIN, GIB_BYTES, + MAX_LAYERWISE_POLICY_TARGETS, MAX_LAYERWISE_RESIDENT_TARGETS, MIN_VRAM_RESERVE_BYTES, PAGEABLE_H2D_COST_MULTIPLIER, @@ -27,6 +28,7 @@ ResidencyTarget, WarmupMemoryRecord, _layerwise_pin_targets, + _layerwise_policy_targets, _layerwise_resident_targets, apply_residency_changes, collect_residency_targets, @@ -3236,6 +3238,22 @@ def test_large_measured_resident_frontier_is_bounded_and_keeps_anchors(self): assert (0, 70) in targets assert (13, 17) in targets + def test_large_policy_frontier_is_bounded_and_keeps_extremes(self): + managers = [ + SimpleNamespace(num_layers=4, residency_policy=RESIDENCY_POLICY_LEADING) + for _ in range(80) + ] + + targets = _layerwise_policy_targets( + managers=managers, + resident_layers=tuple(2 for _ in managers), + tune_policy=True, + ) + + assert len(targets) == MAX_LAYERWISE_POLICY_TARGETS + assert tuple(RESIDENCY_POLICY_LEADING for _ in managers) in targets + assert tuple(RESIDENCY_POLICY_STRIDED for _ in managers) in targets + def test_host_pin_frontier_includes_non_prefix_packings(self): manager = SimpleNamespace( num_layers=3, From a6d8d3edd6e04efcd44ab7d9e8f43005bfa4c33b Mon Sep 17 00:00:00 2001 From: mickqian Date: Thu, 3 Sep 2026 11:15:30 +0800 Subject: [PATCH 64/70] 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 c60d8dfcd0b9..0a5d42af1129 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 @@ -4012,9 +4012,7 @@ def _format_residency_change(candidate: ResidencyTarget) -> str: f"{pin_counts}, policies=" f"{candidate.target_layerwise_residency_policies or 'unchanged'}" ) - 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 b6754726f754..999de938df6a 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 b3be6177d7cb2a9fa7af80c397db2ebd4cbfd21d Mon Sep 17 00:00:00 2001 From: mickqian Date: Sat, 5 Sep 2026 20:58:05 +0800 Subject: [PATCH 65/70] ci: 5090 turbo_wan load peak follows the load-time resident seed The load-safe seed keeps the 1.3B DiT resident from load on a 32 GiB card, so the load peak moves from 928 MiB to 3810 MiB (measured on the 5090 canary). --- .../sglang/multimodal_gen/test/server/perf_baselines/5090.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 0b2551c3a1ed..e59f4b2d779a 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json @@ -107,7 +107,7 @@ "expected_e2e_ms": 2155.59, "expected_avg_denoise_ms": 143.62, "expected_median_denoise_ms": 160.64, - "load_peak_vram_mb": 928.0, + "load_peak_vram_mb": 3810.0, "runtime_peak_vram_mb": 14216.0, "load_peak_host_anon_mb": 6759.33, "runtime_peak_host_anon_mb": 6759.33, From b2851fae2526a44ec126d06633cea1ecab1289d1 Mon Sep 17 00:00:00 2001 From: Mick Qian Date: Sun, 6 Sep 2026 21:09:55 +0800 Subject: [PATCH 66/70] fix(diffusion): use the local checkpoint in plain-loader server args --- .../test/unit/test_plain_component_loader.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/sglang/multimodal_gen/test/unit/test_plain_component_loader.py b/python/sglang/multimodal_gen/test/unit/test_plain_component_loader.py index 4974a8e4ed5e..07dee46f694d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_plain_component_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_plain_component_loader.py @@ -72,9 +72,10 @@ def checkpoint(tmp_path): @pytest.fixture -def server_args(): +def server_args(checkpoint): + component, _ = checkpoint return ServerArgs( - model_path="x", + model_path=str(component.parent), component_residency={"projection": COMPONENT_OFFLOAD}, ) @@ -108,7 +109,7 @@ def test_pipeline_loads_strict_state_dict(checkpoint, server_args, sharded, prec server_args.component_precisions["projection"] = precision pipeline = object.__new__(_ProjectionPipeline) - pipeline.model_path = str(component.parent) + pipeline.model_path = server_args.model_path pipeline.server_args = server_args pipeline._disagg_role = RoleType.MONOLITHIC pipeline.memory_usages = {} From 93f585e9d594eedabb53a8ef9e0e65c0bfc2b650 Mon Sep 17 00:00:00 2001 From: mickqian Date: Mon, 7 Sep 2026 00:11:52 +0800 Subject: [PATCH 67/70] tests: rebuild the residency test modules on the folded base --- .../test/unit/test_auto_residency.py | 4202 +++++++++-------- .../test/unit/test_auto_residency_warmup.py | 884 ++-- .../test/unit/test_cfg_parallel_warmup.py | 54 +- .../test/unit/test_component_residency.py | 281 +- .../test/unit/test_layerwise_offload.py | 86 - .../test/unit/test_performance_metrics.py | 58 - .../test/unit/test_placement_budget.py | 146 +- 7 files changed, 2788 insertions(+), 2923 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 af2efedd6d7f..e49586d73bee 100644 --- a/python/sglang/multimodal_gen/test/unit/test_auto_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_auto_residency.py @@ -118,56 +118,6 @@ def _record( ) -def test_qwen_image_2511_excludes_transformer_only_during_preload(): - from sglang.multimodal_gen.runtime.pipelines.qwen_image import ( - QwenImageEditPlusPipeline, - ) - - pipeline = QwenImageEditPlusPipeline.__new__(QwenImageEditPlusPipeline) - - pipeline.server_args = SimpleNamespace( - pipeline_config=QwenImageEditPlus_2511_PipelineConfig() - ) - assert pipeline.preload_residency_excluded_components == frozenset(("transformer",)) - - pipeline.server_args = SimpleNamespace( - pipeline_config=QwenImageEditPlusPipelineConfig() - ) - assert pipeline.preload_residency_excluded_components == frozenset() - - -class TestEstimateAllocatorHeadroom: - def test_uses_covering_workload_allocator_gap(self): - small = _record( - num_frames=9, - peak_gib=10, - peak_reserved_gib=11, - ) - target = _record( - num_frames=17, - peak_gib=20, - peak_reserved_gib=26, - ) - - assert ( - estimate_allocator_headroom_bytes( - records=[small, target], target_units=target.workload_units() - ) - == 6 * GIB_BYTES - ) - - def test_keeps_largest_observed_gap_without_covering_workload(self): - first = _record(peak_gib=10, peak_reserved_gib=12) - second = _record(peak_gib=15, peak_reserved_gib=19) - - assert ( - estimate_allocator_headroom_bytes( - records=[first, second], target_units=second.workload_units() * 2 - ) - == 4 * GIB_BYTES - ) - - class TestEstimateDefaultWorkloadTiming: def test_scales_only_denoising_with_target_steps(self): record = _record( @@ -199,25 +149,6 @@ def test_scales_only_denoising_with_target_steps(self): "transformer": ("DenoisingStage",), } - def test_prefetch_phase_does_not_claim_stage_latency(self): - record = _record( - num_inference_steps=2, - total_duration_ms=1_000, - stage_duration_ms={"encode": 100, "denoise": 900}, - phase_active_components={ - "0:encode:prefetch:transformer": ("transformer",), - "1:denoise:use:transformer": ("transformer",), - }, - ) - - _, _, component_stages = estimate_default_workload_timing( - records=[record], - target_units=record.workload_units(), - target_num_inference_steps=2, - ) - - assert component_stages == {"transformer": ("denoise",)} - def test_uses_steady_step_instead_of_scaling_first_step_setup(self): record = _record( num_inference_steps=2, @@ -304,55 +235,6 @@ def test_uses_each_stage_iteration_target(self): } assert request_ns == 9_500_000_000 - def test_transfer_savings_uses_component_stage_as_cross_component_cap(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={"encode": 100_000_000, "denoise": 4_000_000_000}, - component_stages={ - "text_encoder": ("encode",), - "transformer": ("denoise",), - }, - ) - - assert savings[candidates[0].option_key()] == 100_000_000 - assert savings[candidates[1].option_key()] > 400_000_000 - assert savings[candidates[2].option_key()] < 0 - - def test_stage_cap_preserves_order_inside_one_shot_component_frontier(self): - partial = ResidencyTarget( - component_name="text_encoder", - residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=10 * GIB_BYTES, - h2d_bytes_per_request=5 * GIB_BYTES, - target_layerwise_resident_layers=(5,), - target_layerwise_pinned_layers=((),), - ) - resident = ResidencyTarget( - component_name="text_encoder", - residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=20 * GIB_BYTES, - h2d_bytes_per_request=10 * GIB_BYTES, - target_layerwise_resident_layers=(10,), - target_layerwise_pinned_layers=((),), - ) - - savings = estimate_candidate_latency_savings_ns( - candidates=[partial, resident], - request_duration_ns=10_000_000_000, - stage_duration_ns={"encode": 100_000_000}, - component_stages={"text_encoder": ("encode",)}, - ) - - assert savings[partial.option_key()] == 50_000_000 - assert savings[resident.option_key()] == 100_000_000 - def test_async_stage_time_does_not_cap_repeated_dit_transfers(self): partial = ResidencyTarget( component_name="transformer", @@ -381,35 +263,6 @@ def test_async_stage_time_does_not_cap_repeated_dit_transfers(self): assert savings[partial.option_key()] > 100_000_000 assert savings[resident.option_key()] > savings[partial.option_key()] - def test_repeated_nonstandard_component_uses_request_latency_cap(self): - partial = ResidencyTarget( - component_name="custom_refiner", - residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=10 * GIB_BYTES, - h2d_bytes_per_request=5 * GIB_BYTES, - target_layerwise_resident_layers=(5,), - target_layerwise_pinned_layers=((),), - ) - resident = ResidencyTarget( - component_name="custom_refiner", - residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=20 * GIB_BYTES, - h2d_bytes_per_request=10 * GIB_BYTES, - target_layerwise_resident_layers=(10,), - target_layerwise_pinned_layers=((),), - ) - - savings = estimate_candidate_latency_savings_ns( - candidates=[partial, resident], - request_duration_ns=10_000_000_000, - stage_duration_ns={"denoise": 100_000_000}, - component_stages={"custom_refiner": ("denoise",)}, - repeated_components={"custom_refiner"}, - ) - - assert savings[partial.option_key()] > 100_000_000 - assert savings[resident.option_key()] > savings[partial.option_key()] - def test_dit_request_cap_does_not_compress_marginal_transfer_savings(self): partial = ResidencyTarget( component_name="transformer", @@ -478,6 +331,103 @@ def test_measured_placement_keeps_the_full_cost_of_a_slower_dit_target(self): assert savings[resident.option_key()] == 416_666_666 assert savings[virtual_layerwise.option_key()] == -5_000_000_000 + def test_prefetch_phase_does_not_claim_stage_latency(self): + record = _record( + num_inference_steps=2, + total_duration_ms=1_000, + stage_duration_ms={"encode": 100, "denoise": 900}, + phase_active_components={ + "0:encode:prefetch:transformer": ("transformer",), + "1:denoise:use:transformer": ("transformer",), + }, + ) + + _, _, component_stages = estimate_default_workload_timing( + records=[record], + target_units=record.workload_units(), + target_num_inference_steps=2, + ) + + assert component_stages == {"transformer": ("denoise",)} + + def test_repeated_nonstandard_component_uses_request_latency_cap(self): + partial = ResidencyTarget( + component_name="custom_refiner", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=10 * GIB_BYTES, + h2d_bytes_per_request=5 * GIB_BYTES, + target_layerwise_resident_layers=(5,), + target_layerwise_pinned_layers=((),), + ) + resident = ResidencyTarget( + component_name="custom_refiner", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=10 * GIB_BYTES, + target_layerwise_resident_layers=(10,), + target_layerwise_pinned_layers=((),), + ) + + savings = estimate_candidate_latency_savings_ns( + candidates=[partial, resident], + request_duration_ns=10_000_000_000, + stage_duration_ns={"denoise": 100_000_000}, + component_stages={"custom_refiner": ("denoise",)}, + repeated_components={"custom_refiner"}, + ) + + assert savings[partial.option_key()] > 100_000_000 + assert savings[resident.option_key()] > savings[partial.option_key()] + + def test_stage_cap_preserves_order_inside_one_shot_component_frontier(self): + partial = ResidencyTarget( + component_name="text_encoder", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=10 * GIB_BYTES, + h2d_bytes_per_request=5 * GIB_BYTES, + target_layerwise_resident_layers=(5,), + target_layerwise_pinned_layers=((),), + ) + resident = ResidencyTarget( + component_name="text_encoder", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=10 * GIB_BYTES, + target_layerwise_resident_layers=(10,), + target_layerwise_pinned_layers=((),), + ) + + savings = estimate_candidate_latency_savings_ns( + candidates=[partial, resident], + request_duration_ns=10_000_000_000, + stage_duration_ns={"encode": 100_000_000}, + component_stages={"text_encoder": ("encode",)}, + ) + + assert savings[partial.option_key()] == 50_000_000 + assert savings[resident.option_key()] == 100_000_000 + + def test_transfer_savings_uses_component_stage_as_cross_component_cap(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={"encode": 100_000_000, "denoise": 4_000_000_000}, + component_stages={ + "text_encoder": ("encode",), + "transformer": ("denoise",), + }, + ) + + assert savings[candidates[0].option_key()] == 100_000_000 + assert savings[candidates[1].option_key()] > 400_000_000 + assert savings[candidates[2].option_key()] < 0 + class TestEstimateLayerwiseLayerUses: def test_scales_repeated_groups_but_not_one_shot_groups(self): @@ -755,47 +705,10 @@ def test_phase_estimation_preserves_active_component_membership(self): assert active == {"denoise": ("transformer",)} assert used == active - def test_failed_probe_preserves_phase_layouts_for_recovery(self): - first = WarmupMemoryRecord( - width=512, - height=512, - num_frames=1, - baseline_allocated_bytes=20 * GIB_BYTES, - peak_allocated_bytes=30 * GIB_BYTES, - succeeded=False, - phase_peak_allocated_bytes={"denoise": 29 * GIB_BYTES}, - phase_active_components={"denoise": ("transformer",)}, - phase_used_components={"denoise": ("transformer",)}, - ) - second = WarmupMemoryRecord( - width=512, - height=512, - num_frames=1, - baseline_allocated_bytes=20 * GIB_BYTES, - peak_allocated_bytes=31 * GIB_BYTES, - succeeded=False, - phase_peak_allocated_bytes={"denoise": 31 * GIB_BYTES}, - phase_active_components={"denoise": ("transformer_2",)}, - phase_used_components={"denoise": ("transformer_2",)}, - ) - - peak, phases, active, used, _, _ = measured_failed_workload_phase_peaks( - records=[first, second], - target_units=first.workload_units(), - ) - - assert peak == 31 * GIB_BYTES - assert phases == { - "denoise:layout:0": 29 * GIB_BYTES, - "denoise:layout:1": 31 * GIB_BYTES, - } - assert set(active.values()) == {("transformer",), ("transformer_2",)} - assert used == active - - def test_phase_estimation_uses_allocated_peak(self): - record = WarmupMemoryRecord( - width=1024, - height=1024, + 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, @@ -906,6 +819,43 @@ def test_phase_estimation_keeps_distinct_active_layouts_separate(self): } assert used == active + def test_failed_probe_preserves_phase_layouts_for_recovery(self): + first = WarmupMemoryRecord( + width=512, + height=512, + num_frames=1, + baseline_allocated_bytes=20 * GIB_BYTES, + peak_allocated_bytes=30 * GIB_BYTES, + succeeded=False, + phase_peak_allocated_bytes={"denoise": 29 * GIB_BYTES}, + phase_active_components={"denoise": ("transformer",)}, + phase_used_components={"denoise": ("transformer",)}, + ) + second = WarmupMemoryRecord( + width=512, + height=512, + num_frames=1, + baseline_allocated_bytes=20 * GIB_BYTES, + peak_allocated_bytes=31 * GIB_BYTES, + succeeded=False, + phase_peak_allocated_bytes={"denoise": 31 * GIB_BYTES}, + phase_active_components={"denoise": ("transformer_2",)}, + phase_used_components={"denoise": ("transformer_2",)}, + ) + + peak, phases, active, used, _, _ = measured_failed_workload_phase_peaks( + records=[first, second], + target_units=first.workload_units(), + ) + + assert peak == 31 * GIB_BYTES + assert phases == { + "denoise:layout:0": 29 * GIB_BYTES, + "denoise:layout:1": 31 * GIB_BYTES, + } + assert set(active.values()) == {("transformer",), ("transformer_2",)} + assert used == active + def test_phase_estimation_keeps_distinct_prefetch_layouts_separate(self): without_prefetch = WarmupMemoryRecord( width=768, @@ -1107,1272 +1057,1404 @@ def test_does_not_replace_explicit_default_shape(self): class TestPlanAutoResidency: - def test_latency_equivalent_partial_dit_prefers_strided_schedule(self): - common = dict( - component_name="transformer", - residency_mode=LAYERWISE_OFFLOAD, - target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=4 * GIB_BYTES, - h2d_bytes_per_request=10 * GIB_BYTES, - target_layerwise_resident_layers=(8,), - target_layerwise_pinned_layers=((),), - target_device_weight_bytes=4 * GIB_BYTES, - ) - leading = ResidencyTarget( - **common, - target_layerwise_residency_policies=(RESIDENCY_POLICY_LEADING,), - current_placement=True, - ) - strided = ResidencyTarget( - **common, - target_layerwise_residency_policies=(RESIDENCY_POLICY_STRIDED,), + 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" + + 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( - estimated_gib=30, - candidates=[leading, strided], - target_workload_measured=True, - ) + _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 - assert len(plan.changes) == 1 - assert plan.changes[0].target_layerwise_residency_policies == ( - RESIDENCY_POLICY_STRIDED, - ) + 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_lower_latency_utility_keeps_the_measured_placement(self): - resident = ResidencyTarget( - component_name="transformer", - residency_mode=RESIDENT, - target_residency_mode=RESIDENT, - target_resident_weight_bytes=20 * GIB_BYTES, - h2d_bytes_per_request=20 * GIB_BYTES, - permanent_residency=True, - target_device_weight_bytes=20 * GIB_BYTES, - current_placement=True, - ) - offloaded = ResidencyTarget( - component_name="transformer", - residency_mode=RESIDENT, - target_residency_mode=COMPONENT_OFFLOAD, - target_resident_weight_bytes=0, - h2d_bytes_per_request=19 * GIB_BYTES, - device_transition_delta_bytes=-20 * GIB_BYTES, - inactive_device_delta_bytes=-20 * GIB_BYTES, - present_device_delta_bytes=-20 * GIB_BYTES, - ) + def test_reserve_floor_is_capped_on_a_small_card(self): + # Keep enough room for unmeasured activations without applying the + # datacenter-card 4 GiB floor verbatim. + plan = plan_auto_residency(reports=[_report(budget_gib=10, estimated_gib=4)]) + assert plan.reserve_bytes == 3 * GIB_BYTES + def test_target_workload_measurement_uses_tighter_reserve(self): plan = plan_auto_residency( reports=[ _report( budget_gib=100, - estimated_gib=30, - candidates=[resident, offloaded], - estimated_request_duration_ns=1_000_000_000, - candidate_latency_savings_ns={ - resident.option_key(): 500_000_000, - offloaded.option_key(): 475_000_000, - }, + estimated_gib=50, + target_workload_measured=True, ) ] ) + assert plan.reserve_bytes == 5 * GIB_BYTES - assert plan.changes == [] - - def test_material_latency_gain_can_replace_the_measured_placement(self): - resident = ResidencyTarget( - component_name="transformer", - residency_mode=RESIDENT, - target_residency_mode=RESIDENT, - target_resident_weight_bytes=20 * GIB_BYTES, - h2d_bytes_per_request=20 * GIB_BYTES, - permanent_residency=True, - target_device_weight_bytes=20 * GIB_BYTES, - current_placement=True, - ) - offloaded = ResidencyTarget( - component_name="transformer", - residency_mode=RESIDENT, - target_residency_mode=COMPONENT_OFFLOAD, - target_resident_weight_bytes=0, - h2d_bytes_per_request=19 * GIB_BYTES, - device_transition_delta_bytes=-20 * GIB_BYTES, - inactive_device_delta_bytes=-20 * GIB_BYTES, - present_device_delta_bytes=-20 * 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=30, - candidates=[resident, offloaded], - estimated_request_duration_ns=1_000_000_000, - candidate_latency_savings_ns={ - resident.option_key(): 460_000_000, - offloaded.option_key(): 500_000_000, - }, - ) + estimated_gib=50, + candidates=[candidate], + target_workload_measured=True, + ), + _report( + rank=1, + budget_gib=140, + estimated_gib=70, + candidates=[candidate], + ), ] ) - assert plan.changes == [offloaded] + 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_free_host_pin_gain_is_selected(self): - pageable = ResidencyTarget( - component_name="transformer", - residency_mode=LAYERWISE_OFFLOAD, - target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=5 * GIB_BYTES, - h2d_bytes_per_request=10 * GIB_BYTES, - target_layerwise_resident_layers=(5,), - target_layerwise_pinned_layers=((),), - target_device_weight_bytes=5 * GIB_BYTES, - current_placement=True, + 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). + 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)] ) - pinned = ResidencyTarget( - component_name="transformer", - residency_mode=LAYERWISE_OFFLOAD, - target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=5 * GIB_BYTES, - h2d_bytes_per_request=11 * GIB_BYTES, - target_layerwise_resident_layers=(5,), - target_layerwise_pinned_layers=((0,),), - pinned_host_delta_bytes=GIB_BYTES, - target_device_weight_bytes=5 * GIB_BYTES, - target_pinned_host_bytes=GIB_BYTES, + assert [c.component_name for c in plan.changes] == ["dit"] + + 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( - budget_gib=100, - estimated_gib=30, - candidates=[pageable, pinned], - host_pin_capacity_gib=10, - estimated_request_duration_ns=1_000_000_000, - candidate_latency_savings_ns={ - pageable.option_key(): 100_000_000, - pinned.option_key(): 110_000_000, - }, - ) - ] - ) - - assert plan.changes == [pinned] - - def test_oom_phase_forces_a_lower_memory_complete_target(self): - resident = ResidencyTarget( - component_name="transformer", - residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=RESIDENT, - target_resident_weight_bytes=20 * GIB_BYTES, - h2d_bytes_per_request=20 * GIB_BYTES, - permanent_residency=True, - target_device_weight_bytes=20 * GIB_BYTES, - current_placement=True, - ) - layerwise = ResidencyTarget( - component_name="transformer", - residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=2 * GIB_BYTES, - h2d_bytes_per_request=10 * GIB_BYTES, - target_layerwise_resident_layers=(2,), - target_layerwise_pinned_layers=((),), - device_transition_delta_bytes=-15 * GIB_BYTES, - active_device_delta_bytes=-15 * GIB_BYTES, - present_device_delta_bytes=-15 * GIB_BYTES, - inactive_device_delta_bytes=-20 * GIB_BYTES, - target_device_weight_bytes=2 * GIB_BYTES, - ) - - plan = plan_auto_residency( - reports=[ + rank=0, + budget_gib=80, + estimated_gib=20, + candidates=candidates, + estimated_request_duration_ns=66_000_000_000, + candidate_latency_savings_ns=timing, + ), _report( - budget_gib=30, - estimated_gib=29, - warmup_oom=True, - phase_peaks_gib={"denoise": 29}, - phase_components={"denoise": ("transformer",)}, - phase_present_components={"denoise": ("transformer",)}, - component_weight_gib={"transformer": 20}, - candidates=[resident, layerwise], - ) + rank=1, + budget_gib=80, + estimated_gib=20, + candidates=candidates, + ), ] ) - assert plan.recovering_from_oom - assert plan.resource_budget_bytes["gpu:rank0:denoise"] == -3 * GIB_BYTES - assert [candidate.option_key() for candidate in plan.changes] == [ - layerwise.option_key() + assert [candidate.component_name for candidate in plan.changes] == [ + "transformer", + "text_encoder", ] - def test_rank_skip_reason_propagates(self): + 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(), _report(rank=1, skip_reason="no measurements")] + reports=[_report(budget_gib=100, estimated_gib=50, candidates=candidates)] ) - assert plan.skip_reason == "rank 1: no measurements" + assert [c.component_name for c in plan.changes] == ["vae"] - def test_oom_recovery_state_survives_an_infeasible_plan(self): + def test_consensus_requires_component_on_every_rank(self): plan = plan_auto_residency( reports=[ _report( - budget_gib=30, - estimated_gib=29, - warmup_oom=True, - phase_peaks_gib={"denoise": 29}, - candidates=[_candidate("vae")], - ) + 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.changes] + assert names == ["vae"] + # worst-case size across ranks + assert plan.changes[0].target_resident_weight_bytes == 12 * GIB_BYTES - assert plan.recovering_from_oom - assert plan.skip_reason == "no placement satisfies all resource budgets" - assert not plan.changes - - def test_missing_estimate_skips(self): - plan = plan_auto_residency(reports=[_report(estimated_gib=None)]) + def test_no_candidates_skips(self): + plan = plan_auto_residency(reports=[_report(candidates=[])]) assert plan.skip_reason is not None - def test_worst_case_aggregation_across_ranks(self): + 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(rank=0, budget_gib=100, estimated_gib=40), - _report(rank=1, budget_gib=80, estimated_gib=60), + _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 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_uncalibrated_small_card_reserve_has_three_gib_floor(self): - plan = plan_auto_residency(reports=[_report(budget_gib=12, estimated_gib=5)]) - assert plan.reserve_bytes == 3 * GIB_BYTES + assert [candidate.component_name for candidate in plan.changes] == [ + "transformer" + ] - def test_pre_warmup_report_enforces_weight_feasibility_without_oom(self): - resident = ResidencyTarget( + def test_async_prefetch_presence_is_not_charged_twice(self): + offload = ResidencyTarget( component_name="transformer", residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=RESIDENT, - target_resident_weight_bytes=11 * GIB_BYTES, - h2d_bytes_per_request=11 * GIB_BYTES, - permanent_residency=True, - target_device_weight_bytes=11 * GIB_BYTES, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, current_placement=True, ) - layerwise = ResidencyTarget( + resident = ResidencyTarget( component_name="transformer", residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=GIB_BYTES, - h2d_bytes_per_request=5 * GIB_BYTES, - target_layerwise_resident_layers=(0,), - target_layerwise_pinned_layers=((),), - active_device_delta_bytes=-9 * GIB_BYTES, - inactive_device_delta_bytes=-10 * GIB_BYTES, - target_device_weight_bytes=GIB_BYTES, + 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=12, - estimated_gib=11, - phase_peaks_gib={"static:transformer": 11}, - phase_components={"static:transformer": ("transformer",)}, - candidates=[resident, layerwise], - require_feasible_placement=True, + 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",), + }, + phase_prefetched_components={ + "text_and_prefetch": ("transformer",), + }, + candidates=[offload, resident], ) ] ) - assert not plan.recovering_from_oom - assert [candidate.target_mode() for candidate in plan.changes] == [ - LAYERWISE_OFFLOAD + assert [candidate.option_key() for candidate in plan.changes] == [ + "transformer:resident" ] - def test_pre_warmup_plan_keeps_feasible_current_placement(self): - current = ResidencyTarget( - component_name="transformer", + 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=GIB_BYTES, - h2d_bytes_per_request=5 * GIB_BYTES, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, target_layerwise_resident_layers=(0,), target_layerwise_pinned_layers=((),), - target_device_weight_bytes=GIB_BYTES, current_placement=True, ) - component_offload = ResidencyTarget( - component_name="transformer", - residency_mode=LAYERWISE_OFFLOAD, - target_residency_mode=COMPONENT_OFFLOAD, - target_resident_weight_bytes=0, - h2d_bytes_per_request=11 * GIB_BYTES, - target_device_weight_bytes=0, - ) - less_layerwise = ResidencyTarget( - component_name="transformer", + text_resident = ResidencyTarget( + component_name="text_encoder", residency_mode=LAYERWISE_OFFLOAD, - target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=0, - h2d_bytes_per_request=7 * GIB_BYTES, - target_layerwise_resident_layers=(0,), + 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=((),), - target_device_weight_bytes=0, + active_device_delta_bytes=20 * GIB_BYTES, + inactive_device_delta_bytes=20 * GIB_BYTES, + target_device_weight_bytes=20 * GIB_BYTES, ) - candidates = [current, less_layerwise, component_offload] plan = plan_auto_residency( reports=[ _report( - budget_gib=12, - estimated_gib=2, - phase_peaks_gib={"static:transformer": 2}, - phase_components={"static:transformer": ("transformer",)}, - candidates=candidates, - require_feasible_placement=True, - estimated_request_duration_ns=1, - candidate_latency_savings_ns={ - candidate.option_key(): int(candidate.current_placement) - for candidate in candidates + 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_pre_warmup_targets_defer_device_weight_promotion(self): - current = ResidencyTarget( - component_name="text_encoder", - residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=COMPONENT_OFFLOAD, - target_resident_weight_bytes=0, - h2d_bytes_per_request=8 * GIB_BYTES, - target_device_weight_bytes=0, - current_placement=True, - ) - resident = ResidencyTarget( + 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.changes + + def test_inactive_phase_charges_full_permanent_target(self): + layerwise = ResidencyTarget( component_name="text_encoder", - residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=RESIDENT, - target_resident_weight_bytes=8 * GIB_BYTES, + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, h2d_bytes_per_request=0, - target_device_weight_bytes=8 * GIB_BYTES, + 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, ) - assert pre_warmup_residency_targets([current, resident]) == [current] + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=80, + estimated_gib=55, + target_workload_measured=True, + phase_peaks_gib={"denoise": 55}, + candidates=[layerwise, resident], + ) + ] + ) - def test_pre_warmup_targets_keep_active_working_set_demotion(self): - current = ResidencyTarget( - component_name="text_encoder", + # 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_unmaterialized_weight_update_charges_full_layerwise_target(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, - target_device_weight_bytes=0, current_placement=True, ) layerwise = ResidencyTarget( - component_name="text_encoder", + component_name="transformer", residency_mode=COMPONENT_OFFLOAD, target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=0, - h2d_bytes_per_request=0, - target_device_weight_bytes=GIB_BYTES, - active_device_delta_bytes=-16 * GIB_BYTES, + 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, - device_transition_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, ) - assert pre_warmup_residency_targets([current, layerwise]) == [ - current, - layerwise, - ] + 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": (), + "idle": ("transformer",), + }, + phase_full_weight_transition_components={ + "lora_switch": ("transformer",) + }, + candidates=[component_offload, layerwise, resident], + ) + ] + ) - def test_pre_warmup_report_uses_weight_model_below_legacy_threshold( - self, monkeypatch - ): - from sglang.multimodal_gen.runtime.managers import gpu_worker + assert plan.changes == [] - current = ResidencyTarget( + 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=8 * GIB_BYTES, - target_device_weight_bytes=0, + h2d_bytes_per_request=0, current_placement=True, ) - modules = {"transformer": object()} - collect_calls = [] - - def collect_targets(_workload, *, allow_host_pin_reallocation=True): - collect_calls.append(allow_host_pin_reallocation) - return [current] - - worker = SimpleNamespace( - rank=0, - pipeline=SimpleNamespace(preload_residency_excluded_components=frozenset()), - server_args=SimpleNamespace( - num_gpus=1, - nnodes=1, - node_rank=0, - host_pin_budget=lambda: None, - ), - _auto_residency_budget_bytes=lambda: 30 * GIB_BYTES, - _collect_auto_residency_targets=collect_targets, - _auto_residency_modules=lambda: modules, - ) - monkeypatch.setattr( - gpu_worker, - "component_runtime_weight_bytes", - lambda _modules: {"transformer": 8 * GIB_BYTES}, - ) - monkeypatch.setattr( - gpu_worker, - "component_current_device_weight_bytes", - lambda _modules: {"transformer": 0}, - ) - monkeypatch.setattr( - gpu_worker, "layerwise_pinned_host_bytes", lambda *a, **k: 0 - ) - monkeypatch.setattr( - gpu_worker, "layerwise_host_pin_capacity_bytes", lambda *a, **k: 0 - ) - monkeypatch.setattr( - gpu_worker, "host_memory_available_bytes", lambda: 100 * GIB_BYTES + 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, ) - monkeypatch.setattr( - gpu_worker.torch, - "get_device_module", - lambda: SimpleNamespace(memory_allocated=lambda: 2 * 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, ) - report = gpu_worker.GPUWorker._build_pre_warmup_auto_residency_report( - worker, - workload=DefaultWorkload( - width=832, - height=480, - num_frames=24, - num_inference_steps=50, - ), + 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 report.budget_bytes == 30 * GIB_BYTES - assert report.estimated_peak_bytes == 10 * GIB_BYTES - assert report.candidates == [current] - assert report.skip_reason is None - assert report.estimated_request_duration_ns == 1 - assert report.candidate_latency_savings_ns == {current.option_key(): 1} - assert collect_calls == [False] - - def test_pre_warmup_report_keeps_excluded_components_fixed(self, monkeypatch): - from sglang.multimodal_gen.runtime.managers import gpu_worker + assert [candidate.target_mode() for candidate in plan.changes] == [ + LAYERWISE_OFFLOAD + ] - transformer = ResidencyTarget( + def test_materialized_layerwise_transition_does_not_double_count_weights(self): + current_layerwise = ResidencyTarget( component_name="transformer", - residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=RESIDENT, - target_resident_weight_bytes=20 * GIB_BYTES, - h2d_bytes_per_request=20 * GIB_BYTES, - permanent_residency=True, + 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, ) - vae = ResidencyTarget( - component_name="vae", - residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=RESIDENT, + layerwise = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, target_resident_weight_bytes=GIB_BYTES, - h2d_bytes_per_request=GIB_BYTES, - permanent_residency=True, + 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, - current_placement=True, - ) - modules = {"transformer": object(), "vae": object()} - worker = SimpleNamespace( - rank=0, - pipeline=SimpleNamespace( - preload_residency_excluded_components=frozenset(("transformer",)) - ), - server_args=SimpleNamespace( - num_gpus=1, - nnodes=1, - node_rank=0, - host_pin_budget=lambda: None, - ), - _auto_residency_budget_bytes=lambda: 80 * GIB_BYTES, - _collect_auto_residency_targets=lambda _workload, **_kwargs: [ - transformer, - vae, - ], - _auto_residency_modules=lambda: modules, - ) - monkeypatch.setattr( - gpu_worker, - "component_runtime_weight_bytes", - lambda _modules: {"transformer": 20 * GIB_BYTES, "vae": GIB_BYTES}, ) - monkeypatch.setattr( - gpu_worker, - "component_current_device_weight_bytes", - lambda _modules: {"transformer": 0, "vae": 0}, + 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, ) - monkeypatch.setattr( - gpu_worker, "layerwise_pinned_host_bytes", lambda *a, **k: 0 + + 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=[current_layerwise, layerwise, resident], + ) + ] ) - monkeypatch.setattr( - gpu_worker, "layerwise_host_pin_capacity_bytes", lambda *a, **k: 0 + + 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}, + ) + ] ) - monkeypatch.setattr( - gpu_worker, "host_memory_available_bytes", lambda: 100 * GIB_BYTES + + assert plan.current_placement_reserve_shortfall_bytes == 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}, ) - monkeypatch.setattr( - gpu_worker.torch, - "get_device_module", - lambda: SimpleNamespace(memory_allocated=lambda: 2 * GIB_BYTES), + + assert current_placement_reserve_shortfall_bytes([report]) == 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, ) - report = gpu_worker.GPUWorker._build_pre_warmup_auto_residency_report( - worker, - workload=DefaultWorkload( - width=1920, - height=1088, - num_frames=24, - num_inference_steps=15, - ), + 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( + component_name="second_transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_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, ) - - assert [candidate.component_name for candidate in report.candidates] == ["vae"] - assert report.estimated_peak_bytes_by_phase["static:transformer"] == ( - 22 * GIB_BYTES + 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 report.used_components_by_phase["static:transformer"] == ("transformer",) - def test_reserve_floor_is_capped_on_a_small_card(self): - # Keep enough room for unmeasured activations without applying the - # datacenter-card 4 GiB floor verbatim. - plan = plan_auto_residency(reports=[_report(budget_gib=10, estimated_gib=4)]) - assert plan.reserve_bytes == 3 * GIB_BYTES + assert [candidate.component_name for candidate in plan.changes] == [ + "first_transformer" + ] + assert "gpu:rank0:unobserved:second_transformer" in plan.resource_budget_bytes - def test_target_workload_measurement_uses_tighter_reserve(self): + def test_redundant_phase_constraints_collapse_to_the_highest_peak(self): plan = plan_auto_residency( reports=[ _report( budget_gib=100, - estimated_gib=50, - target_workload_measured=True, + 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 plan.reserve_bytes == 5 * GIB_BYTES - def test_measured_small_card_reserve_drops_to_ten_percent(self): + 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_validated_baseline_with_no_reserve_headroom_accepts_zero_delta(self): plan = plan_auto_residency( reports=[ _report( - budget_gib=12, - estimated_gib=10, - target_workload_measured=True, + budget_gib=80, + estimated_gib=76, + candidates=[_candidate("transformer", weight_gib=10)], + phase_peaks_gib={"denoise": 76}, + phase_components={"denoise": ("transformer",)}, ) ] ) - assert plan.reserve_bytes == int(1.2 * GIB_BYTES) - def test_target_workload_measurement_is_replica_wide(self): + 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) plan = plan_auto_residency( reports=[ _report( rank=0, - budget_gib=100, - estimated_gib=50, + budget_gib=40, + estimated_gib=20, candidates=[candidate], - target_workload_measured=True, + phase_peaks_gib={"decode": 20}, + phase_components={"decode": ("vae",)}, ), _report( rank=1, budget_gib=140, - estimated_gib=70, + estimated_gib=100, candidates=[candidate], + phase_peaks_gib={"decode": 100}, + phase_components={"decode": ("vae",)}, ), ] ) - 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 + 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_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). - candidates = [ - _candidate("text_encoder", weight_gib=30), - _candidate("dit", mode=LAYERWISE_OFFLOAD, weight_gib=25, h2d_gib=25 * 50), - ] + def test_partial_layer_residency_is_selected_when_full_dit_does_not_fit(self): + full = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_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 = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_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=50, candidates=candidates)] + reports=[ + _report( + budget_gib=100, + estimated_gib=70, + candidates=[full, partial], + phase_peaks_gib={"denoise": 70}, + phase_components={"denoise": ("transformer",)}, + ) + ] ) - assert [c.component_name for c in plan.changes] == ["dit"] - def test_long_request_uses_available_vram_for_encoder_gain(self): - transformer = _candidate( - "transformer", - mode=LAYERWISE_OFFLOAD, - weight_gib=40, - h2d_gib=40 * 40, + 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 = ResidencyTarget( + component_name="transformer_2", + residency_mode=LAYERWISE_OFFLOAD, + 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, ) - text_encoder = _candidate("text_encoder", weight_gib=8, h2d_gib=8) - candidates = [transformer, text_encoder] + keep_hot_encoder = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_resident_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=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, + 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.changes] == [ - "transformer", + assert {candidate.component_name for candidate in plan.changes} == { + "transformer_2", "text_encoder", - ] + } - def test_equal_latency_plan_avoids_strategy_repacking(self): - current = ResidencyTarget( - component_name="vae", + def test_node_hostpin_is_optimized_with_the_same_selection_vector(self): + cold_pageable = ResidencyTarget( + component_name="cold_encoder", residency_mode=LAYERWISE_OFFLOAD, - target_residency_mode=LAYERWISE_OFFLOAD, target_resident_weight_bytes=0, - h2d_bytes_per_request=0, + h2d_bytes_per_request=-GIB_BYTES, target_layerwise_resident_layers=(0,), target_layerwise_pinned_layers=((),), - current_placement=True, + pinned_host_delta_bytes=-10 * GIB_BYTES, ) - component_offload = ResidencyTarget( - component_name="vae", + hot_pinned = ResidencyTarget( + component_name="hot_dit", residency_mode=LAYERWISE_OFFLOAD, - target_residency_mode=COMPONENT_OFFLOAD, target_resident_weight_bytes=0, - h2d_bytes_per_request=GIB_BYTES, + h2d_bytes_per_request=100 * GIB_BYTES, target_layerwise_resident_layers=(0,), - target_layerwise_pinned_layers=((),), + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=10 * GIB_BYTES, ) - plan = plan_auto_residency( reports=[ _report( - candidates=[current, component_offload], + 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={ - current.option_key(): 0, - component_offload.option_key(): 0, + cold_pageable.option_key(): 10_000_000, + hot_pinned.option_key(): 500_000_000, }, ) ] ) - assert plan.changes == [] + # 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.changes] == [ + "hot_dit", + "cold_encoder", + ] + assert plan.resource_budget_bytes["hostpin:node0"] == 0 - def test_output_rank_timing_applies_to_untimed_peer_ranks(self): - transformer = _candidate( - "transformer", - mode=LAYERWISE_OFFLOAD, - weight_gib=40, - h2d_gib=40 * 40, + def test_hostpin_repack_must_fit_transition_headroom(self): + pin_more = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_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], + ) + ] ) - 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, - } + assert plan.changes == [] + assert plan.resource_budget_bytes["hostram:node0: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( - 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, - ), + host_transition_headroom_gib=5, + candidates=[demote], + ) ] ) - assert [candidate.component_name for candidate in plan.changes] == [ - "transformer", - "text_encoder", - ] + assert plan.changes == [] + assert plan.resource_budget_bytes["hostram:node0:materialize"] == 5 * GIB_BYTES - 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)] + 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, ) - assert [c.component_name for c in plan.changes] == ["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)]), + budget_gib=100, + estimated_gib=50, + device_transition_allocated_gib=90, + target_workload_measured=True, + candidates=[resident], + ) ] ) - names = [c.component_name for c in plan.changes] - assert names == ["vae"] - # worst-case size across ranks - assert plan.changes[0].target_resident_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 + assert plan.changes == [] + assert ( + plan.resource_budget_bytes["gpu:rank0:placement-transition"] + == 5 * GIB_BYTES + ) - 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. + 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=139, - estimated_gib=66, - candidates=[_candidate("transformer", weight_gib=60)], - phase_peaks_gib={"denoise": 66, "idle": 4}, - phase_components={"denoise": ("transformer",), "idle": ()}, + 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] == [ - "transformer" - ] + assert {candidate.component_name for candidate in plan.changes} == { + "cold_transformer", + "hot_transformer", + } - def test_async_prefetch_presence_is_not_charged_twice(self): - offload = ResidencyTarget( + 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, ) - resident = ResidencyTarget( - component_name="transformer", + encoder_resident = ResidencyTarget( + component_name="text_encoder", 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, + 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=95, - estimated_gib=50, + budget_gib=100, + estimated_gib=90, target_workload_measured=True, phase_peaks_gib={ - "text_and_prefetch": 50, - "denoise": 50, + "transformer_use": 70, + "encoder_use": 90, }, phase_components={ - "text_and_prefetch": ("text_encoder",), - "denoise": ("transformer",), + "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", + } + + def test_component_offload_can_release_an_overlapped_layerwise_prefetch(self): + layerwise = ResidencyTarget( + component_name="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=((),), + current_placement=True, + ) + component_offload = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((),), + concurrent_prefetch_device_delta_bytes=-10 * GIB_BYTES, + ) + + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=18, + estimated_gib=18, + warmup_oom=True, + candidates=[layerwise, component_offload], + phase_peaks_gib={"encode_prefetch": 18}, + phase_components={"encode_prefetch": ("text_encoder",)}, phase_present_components={ - "text_and_prefetch": ("text_encoder", "transformer"), - "denoise": ("transformer",), + "encode_prefetch": ("text_encoder", "transformer"), }, phase_prefetched_components={ - "text_and_prefetch": ("transformer",), + "encode_prefetch": ("transformer",), + }, + component_active_weight_gib={ + "text_encoder": 8, + "transformer": 10, }, - candidates=[offload, resident], ) ] ) - assert [candidate.option_key() for candidate in plan.changes] == [ - "transformer:resident" + assert plan.skip_reason is None + assert [candidate.target_mode() for candidate in plan.changes] == [ + COMPONENT_OFFLOAD ] - def test_component_offload_can_release_an_overlapped_layerwise_prefetch(self): - layerwise = ResidencyTarget( + def test_dynamic_hostpin_uses_one_node_budget_and_assigns_rank_quotas(self): + pin_more = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, + target_resident_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 [candidate.component_name for candidate in plan.changes] == [ + "transformer" + ] + assert plan.resource_budget_bytes["hostpin:node0"] == 40 * GIB_BYTES + assert plan.host_pin_target_bytes_by_rank == { + 0: 10 * GIB_BYTES, + 1: 10 * GIB_BYTES, + } + + def test_equal_latency_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=10 * GIB_BYTES, + h2d_bytes_per_request=0, target_layerwise_resident_layers=(0,), target_layerwise_pinned_layers=((),), current_placement=True, ) component_offload = ResidencyTarget( - component_name="transformer", + component_name="vae", residency_mode=LAYERWISE_OFFLOAD, target_residency_mode=COMPONENT_OFFLOAD, target_resident_weight_bytes=0, - h2d_bytes_per_request=0, + h2d_bytes_per_request=GIB_BYTES, target_layerwise_resident_layers=(0,), target_layerwise_pinned_layers=((),), - concurrent_prefetch_device_delta_bytes=-10 * GIB_BYTES, ) plan = plan_auto_residency( reports=[ _report( - budget_gib=18, - estimated_gib=18, - warmup_oom=True, - candidates=[layerwise, component_offload], - phase_peaks_gib={"encode_prefetch": 18}, - phase_components={"encode_prefetch": ("text_encoder",)}, - phase_present_components={ - "encode_prefetch": ("text_encoder", "transformer"), - }, - phase_prefetched_components={ - "encode_prefetch": ("transformer",), - }, - component_active_weight_gib={ - "text_encoder": 8, - "transformer": 10, + candidates=[current, component_offload], + estimated_request_duration_ns=1_000_000_000, + candidate_latency_savings_ns={ + current.option_key(): 0, + component_offload.option_key(): 0, }, ) ] ) - assert plan.skip_reason is None - assert [candidate.target_mode() for candidate in plan.changes] == [ - COMPONENT_OFFLOAD - ] + assert plan.changes == [] - def test_request_end_residency_is_carried_into_the_next_request(self): - text_offload = ResidencyTarget( - component_name="text_encoder", + def test_failed_candidate_margin_only_constrains_future_growth(self): + candidate = _candidate("transformer", weight_gib=5, h2d_gib=5) + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=40, + estimated_gib=30, + planning_headroom_correction_gib=3, + target_workload_measured=True, + candidates=[candidate], + phase_peaks_gib={"denoise": 30}, + ) + ] + ) + + assert plan.current_placement_reserve_shortfall_bytes == 0 + assert plan.resource_budget_bytes["gpu:rank0:denoise"] == 3 * GIB_BYTES + assert plan.changes == [] + + def test_free_host_pin_gain_is_selected(self): + pageable = 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_resident_weight_bytes=5 * GIB_BYTES, + h2d_bytes_per_request=10 * GIB_BYTES, + target_layerwise_resident_layers=(5,), target_layerwise_pinned_layers=((),), + target_device_weight_bytes=5 * GIB_BYTES, current_placement=True, ) - text_resident = ResidencyTarget( - component_name="text_encoder", + pinned = ResidencyTarget( + component_name="transformer", 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, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=5 * GIB_BYTES, + h2d_bytes_per_request=11 * GIB_BYTES, + target_layerwise_resident_layers=(5,), + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=GIB_BYTES, + target_device_weight_bytes=5 * GIB_BYTES, + target_pinned_host_bytes=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",), + budget_gib=100, + estimated_gib=30, + candidates=[pageable, pinned], + host_pin_capacity_gib=10, + estimated_request_duration_ns=1_000_000_000, + candidate_latency_savings_ns={ + pageable.option_key(): 100_000_000, + pinned.option_key(): 110_000_000, }, - 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 == [] + assert plan.changes == [pinned] + + def test_hostpin_capacity_is_scoped_per_node(self): + pin_more = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=100 * GIB_BYTES, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=30 * GIB_BYTES, + ) - 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": ()}, - ) + rank=0, + node_rank=0, + host_pin_capacity_gib=20, + candidates=[pin_more], + ), + _report( + rank=1, + node_rank=1, + host_pin_capacity_gib=20, + candidates=[pin_more], + ), ] ) - assert not plan.changes + assert plan.changes == [] + assert plan.resource_budget_bytes["hostpin:node0"] == 20 * GIB_BYTES + assert plan.resource_budget_bytes["hostpin:node1"] == 20 * GIB_BYTES - def test_inactive_phase_charges_full_permanent_target(self): - layerwise = ResidencyTarget( - component_name="text_encoder", + def test_latency_equivalent_partial_dit_prefers_strided_schedule(self): + common = dict( + 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_resident_weight_bytes=4 * GIB_BYTES, + h2d_bytes_per_request=10 * GIB_BYTES, + target_layerwise_resident_layers=(8,), target_layerwise_pinned_layers=((),), + target_device_weight_bytes=4 * GIB_BYTES, + ) + leading = ResidencyTarget( + **common, + target_layerwise_residency_policies=(RESIDENCY_POLICY_LEADING,), 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, + strided = ResidencyTarget( + **common, + target_layerwise_residency_policies=(RESIDENCY_POLICY_STRIDED,), ) plan = plan_auto_residency( reports=[ _report( - budget_gib=80, - estimated_gib=55, + estimated_gib=30, + candidates=[leading, strided], 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 == [] + assert len(plan.changes) == 1 + assert plan.changes[0].target_layerwise_residency_policies == ( + RESIDENCY_POLICY_STRIDED, + ) - def test_unmaterialized_weight_update_charges_full_layerwise_target(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, + def test_long_request_uses_available_vram_for_encoder_gain(self): + transformer = _candidate( + "transformer", + mode=LAYERWISE_OFFLOAD, + weight_gib=40, + h2d_gib=40 * 40, ) - 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, + 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.changes] == [ + "transformer", + "text_encoder", + ] + + def test_lower_latency_utility_keeps_the_measured_placement(self): resident = ResidencyTarget( component_name="transformer", - residency_mode=COMPONENT_OFFLOAD, + residency_mode=RESIDENT, target_residency_mode=RESIDENT, - target_resident_weight_bytes=40 * GIB_BYTES, - h2d_bytes_per_request=0, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=20 * GIB_BYTES, permanent_residency=True, - inactive_device_delta_bytes=40 * GIB_BYTES, - target_device_weight_bytes=40 * GIB_BYTES, + target_device_weight_bytes=20 * GIB_BYTES, + current_placement=True, + ) + offloaded = ResidencyTarget( + component_name="transformer", + residency_mode=RESIDENT, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=19 * GIB_BYTES, + device_transition_delta_bytes=-20 * GIB_BYTES, + inactive_device_delta_bytes=-20 * GIB_BYTES, + present_device_delta_bytes=-20 * 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": (), - "idle": ("transformer",), - }, - phase_full_weight_transition_components={ - "lora_switch": ("transformer",) + budget_gib=100, + estimated_gib=30, + candidates=[resident, offloaded], + estimated_request_duration_ns=1_000_000_000, + candidate_latency_savings_ns={ + resident.option_key(): 500_000_000, + offloaded.option_key(): 475_000_000, }, - candidates=[component_offload, layerwise, resident], ) ] ) 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, - ) + def test_material_latency_gain_can_replace_the_measured_placement(self): resident = ResidencyTarget( component_name="transformer", - residency_mode=COMPONENT_OFFLOAD, + residency_mode=RESIDENT, target_residency_mode=RESIDENT, - target_resident_weight_bytes=40 * GIB_BYTES, - h2d_bytes_per_request=0, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=20 * GIB_BYTES, permanent_residency=True, - inactive_device_delta_bytes=40 * GIB_BYTES, - target_device_weight_bytes=40 * GIB_BYTES, + target_device_weight_bytes=20 * GIB_BYTES, + current_placement=True, + ) + offloaded = ResidencyTarget( + component_name="transformer", + residency_mode=RESIDENT, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=19 * GIB_BYTES, + device_transition_delta_bytes=-20 * GIB_BYTES, + inactive_device_delta_bytes=-20 * GIB_BYTES, + present_device_delta_bytes=-20 * 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",) + budget_gib=100, + estimated_gib=30, + candidates=[resident, offloaded], + estimated_request_duration_ns=1_000_000_000, + candidate_latency_savings_ns={ + resident.option_key(): 460_000_000, + offloaded.option_key(): 500_000_000, }, - candidates=[component_offload, layerwise, resident], ) ] ) - assert [candidate.target_mode() for candidate in plan.changes] == [ - LAYERWISE_OFFLOAD - ] + assert plan.changes == [offloaded] - def test_materialized_layerwise_transition_does_not_double_count_weights(self): - current_layerwise = ResidencyTarget( + def test_measured_small_card_reserve_drops_to_ten_percent(self): + plan = plan_auto_residency( + reports=[ + _report( + budget_gib=12, + estimated_gib=10, + target_workload_measured=True, + ) + ] + ) + assert plan.reserve_bytes == int(1.2 * GIB_BYTES) + + def test_node_hostpin_can_assign_asymmetric_rank_quotas(self): + rank0 = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, - target_residency_mode=LAYERWISE_OFFLOAD, target_resident_weight_bytes=0, - h2d_bytes_per_request=0, + h2d_bytes_per_request=100 * GIB_BYTES, target_layerwise_resident_layers=(0,), - target_layerwise_pinned_layers=((),), - current_placement=True, + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=30 * GIB_BYTES, ) - layerwise = ResidencyTarget( + rank1 = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, - target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=GIB_BYTES, + target_resident_weight_bytes=0, 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, + target_layerwise_resident_layers=(0,), + target_layerwise_pinned_layers=((0,),), + pinned_host_delta_bytes=5 * 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=[current_layerwise, layerwise, resident], - ) + rank=0, + host_pin_capacity_gib=20, + candidates=[rank0], + ), + _report( + rank=1, + host_pin_capacity_gib=20, + candidates=[rank1], + ), ] ) - assert [candidate.target_mode() for candidate in plan.changes] == [ - LAYERWISE_OFFLOAD + assert [candidate.component_name for candidate in plan.changes] == [ + "transformer" ] + assert plan.host_pin_target_bytes_by_rank == { + 0: 30 * GIB_BYTES, + 1: 5 * GIB_BYTES, + } - def test_untracked_oom_peak_constrains_synthetic_component_phase(self): - component_offload = ResidencyTarget( + def test_oom_phase_forces_a_lower_memory_complete_target(self): + resident = ResidencyTarget( component_name="transformer", residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=COMPONENT_OFFLOAD, - target_resident_weight_bytes=0, - h2d_bytes_per_request=0, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=20 * GIB_BYTES, + permanent_residency=True, + target_device_weight_bytes=20 * GIB_BYTES, current_placement=True, ) layerwise = ResidencyTarget( component_name="transformer", residency_mode=COMPONENT_OFFLOAD, target_residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=GIB_BYTES, + target_resident_weight_bytes=2 * GIB_BYTES, h2d_bytes_per_request=10 * GIB_BYTES, - target_layerwise_resident_layers=(0,), + target_layerwise_resident_layers=(2,), target_layerwise_pinned_layers=((),), - active_device_delta_bytes=-10 * GIB_BYTES, - inactive_device_delta_bytes=GIB_BYTES, - ) - - plan = plan_auto_residency( - reports=[ - _report( - budget_gib=12, - estimated_gib=11.75, - phase_peaks_gib={"request:untracked": 11.75}, - phase_present_components={"request:untracked": ("vae",)}, - candidates=[component_offload, layerwise], - warmup_oom=True, - ) - ] + device_transition_delta_bytes=-15 * GIB_BYTES, + active_device_delta_bytes=-15 * GIB_BYTES, + present_device_delta_bytes=-15 * GIB_BYTES, + inactive_device_delta_bytes=-20 * GIB_BYTES, + target_device_weight_bytes=2 * GIB_BYTES, ) - 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 == GIB_BYTES - - def test_failed_candidate_margin_only_constrains_future_growth(self): - candidate = _candidate("transformer", weight_gib=5, h2d_gib=5) - plan = plan_auto_residency( - reports=[ - _report( - budget_gib=40, - estimated_gib=30, - planning_headroom_correction_gib=3, - target_workload_measured=True, - candidates=[candidate], - phase_peaks_gib={"denoise": 30}, + estimated_gib=29, + warmup_oom=True, + phase_peaks_gib={"denoise": 29}, + phase_components={"denoise": ("transformer",)}, + phase_present_components={"denoise": ("transformer",)}, + component_weight_gib={"transformer": 20}, + candidates=[resident, layerwise], ) - ] - ) - - assert plan.current_placement_reserve_shortfall_bytes == 0 - assert plan.resource_budget_bytes["gpu:rank0:denoise"] == 3 * GIB_BYTES - assert plan.changes == [] - - 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]) == GIB_BYTES - - def test_validation_keeps_conservative_reserve_after_warmup_oom(self): - report = _report( - budget_gib=30, - estimated_gib=20, - target_workload_measured=True, - phase_peaks_gib={"decode": 28}, - warmup_oom=True, - ) - - 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 + assert plan.recovering_from_oom + assert plan.resource_budget_bytes["gpu:rank0:denoise"] == -3 * GIB_BYTES + assert [candidate.option_key() for candidate in plan.changes] == [ + layerwise.option_key() + ] - def test_unobserved_later_component_gets_a_conservative_phase(self): - first_transformer = _candidate("first_transformer", weight_gib=26, h2d_gib=100) - second_transformer = ResidencyTarget( - component_name="second_transformer", - residency_mode=LAYERWISE_OFFLOAD, - target_resident_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, - ) + def test_oom_recovery_state_survives_an_infeasible_plan(self): 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",), - }, + budget_gib=30, + estimated_gib=29, + warmup_oom=True, + phase_peaks_gib={"denoise": 29}, + candidates=[_candidate("vae")], ) ] ) - assert [candidate.component_name for candidate in plan.changes] == [ - "first_transformer" - ] - assert "gpu:rank0:unobserved:second_transformer" in plan.resource_budget_bytes + assert plan.recovering_from_oom + assert plan.skip_reason == "no placement satisfies all resource budgets" + assert not plan.changes + + def test_uncalibrated_small_card_reserve_has_three_gib_floor(self): + plan = plan_auto_residency(reports=[_report(budget_gib=12, estimated_gib=5)]) + assert plan.reserve_bytes == 3 * GIB_BYTES def test_unobserved_phase_excludes_another_components_transient_weights(self): transformer_component = ResidencyTarget( @@ -2440,486 +2522,354 @@ def test_unobserved_phase_excludes_another_components_transient_weights(self): 12 * GIB_BYTES - 5 * GIB_BYTES - plan.reserve_bytes ) - 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_validated_baseline_with_no_reserve_headroom_accepts_zero_delta(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 [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) - 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 = ResidencyTarget( + def test_untracked_oom_peak_constrains_synthetic_component_phase(self): + component_offload = ResidencyTarget( component_name="transformer", - residency_mode=LAYERWISE_OFFLOAD, - target_resident_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, + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + current_placement=True, ) - partial = ResidencyTarget( + layerwise = ResidencyTarget( component_name="transformer", - residency_mode=LAYERWISE_OFFLOAD, - target_resident_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.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 = ResidencyTarget( - component_name="transformer_2", - residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=-20 * GIB_BYTES, - h2d_bytes_per_request=-20 * GIB_BYTES, + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=10 * GIB_BYTES, target_layerwise_resident_layers=(0,), target_layerwise_pinned_layers=((),), - active_device_delta_bytes=-20 * GIB_BYTES, - ) - keep_hot_encoder = ResidencyTarget( - component_name="text_encoder", - residency_mode=COMPONENT_OFFLOAD, - target_resident_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, + active_device_delta_bytes=-10 * GIB_BYTES, + inactive_device_delta_bytes=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",), - }, + budget_gib=12, + estimated_gib=11.75, + phase_peaks_gib={"request:untracked": 11.75}, + phase_present_components={"request:untracked": ("vae",)}, + candidates=[component_offload, layerwise], + warmup_oom=True, ) ] ) - assert {candidate.component_name for candidate in plan.changes} == { - "transformer_2", - "text_encoder", - } + assert [candidate.target_mode() for candidate in plan.changes] == [ + LAYERWISE_OFFLOAD + ] - def test_node_hostpin_is_optimized_with_the_same_selection_vector(self): - cold_pageable = ResidencyTarget( - component_name="cold_encoder", - residency_mode=LAYERWISE_OFFLOAD, - 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 = ResidencyTarget( - component_name="hot_dit", - residency_mode=LAYERWISE_OFFLOAD, - target_resident_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], - 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, - }, - ) - ] + def test_validation_keeps_conservative_reserve_after_warmup_oom(self): + report = _report( + budget_gib=30, + estimated_gib=20, + target_workload_measured=True, + phase_peaks_gib={"decode": 28}, + warmup_oom=True, ) - # 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.changes] == [ - "hot_dit", - "cold_encoder", - ] - assert plan.resource_budget_bytes["hostpin:node0"] == 0 + assert current_placement_reserve_shortfall_bytes([report]) == 2 * GIB_BYTES - def test_dynamic_hostpin_uses_one_node_budget_and_assigns_rank_quotas(self): - pin_more = ResidencyTarget( + def test_pre_warmup_report_enforces_weight_feasibility_without_oom(self): + resident = ResidencyTarget( + component_name="transformer", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=11 * GIB_BYTES, + h2d_bytes_per_request=11 * GIB_BYTES, + permanent_residency=True, + target_device_weight_bytes=11 * GIB_BYTES, + current_placement=True, + ) + layerwise = ResidencyTarget( component_name="transformer", - residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=0, - h2d_bytes_per_request=100 * GIB_BYTES, + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=5 * GIB_BYTES, target_layerwise_resident_layers=(0,), - target_layerwise_pinned_layers=((0,),), - pinned_host_delta_bytes=10 * GIB_BYTES, + target_layerwise_pinned_layers=((),), + active_device_delta_bytes=-9 * GIB_BYTES, + inactive_device_delta_bytes=-10 * GIB_BYTES, + target_device_weight_bytes=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], - ), + budget_gib=12, + estimated_gib=11, + phase_peaks_gib={"static:transformer": 11}, + phase_components={"static:transformer": ("transformer",)}, + candidates=[resident, layerwise], + require_feasible_placement=True, + ) ] ) - assert [candidate.component_name for candidate in plan.changes] == [ - "transformer" + assert not plan.recovering_from_oom + assert [candidate.target_mode() for candidate in plan.changes] == [ + LAYERWISE_OFFLOAD ] - assert plan.resource_budget_bytes["hostpin:node0"] == 40 * GIB_BYTES - assert plan.host_pin_target_bytes_by_rank == { - 0: 10 * GIB_BYTES, - 1: 10 * GIB_BYTES, - } - def test_node_hostpin_can_assign_asymmetric_rank_quotas(self): - rank0 = ResidencyTarget( + def test_pre_warmup_plan_keeps_feasible_current_placement(self): + current = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=0, - h2d_bytes_per_request=100 * GIB_BYTES, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=5 * GIB_BYTES, target_layerwise_resident_layers=(0,), - target_layerwise_pinned_layers=((0,),), - pinned_host_delta_bytes=30 * GIB_BYTES, + target_layerwise_pinned_layers=((),), + target_device_weight_bytes=GIB_BYTES, + current_placement=True, ) - rank1 = ResidencyTarget( + component_offload = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, target_resident_weight_bytes=0, - h2d_bytes_per_request=100 * GIB_BYTES, - target_layerwise_resident_layers=(0,), - target_layerwise_pinned_layers=((0,),), - pinned_host_delta_bytes=5 * GIB_BYTES, - ) - - plan = plan_auto_residency( - reports=[ - _report( - rank=0, - host_pin_capacity_gib=20, - candidates=[rank0], - ), - _report( - rank=1, - host_pin_capacity_gib=20, - candidates=[rank1], - ), - ] + h2d_bytes_per_request=11 * GIB_BYTES, + target_device_weight_bytes=0, ) - - assert [candidate.component_name for candidate in plan.changes] == [ - "transformer" - ] - assert plan.host_pin_target_bytes_by_rank == { - 0: 30 * GIB_BYTES, - 1: 5 * GIB_BYTES, - } - - def test_hostpin_capacity_is_scoped_per_node(self): - pin_more = ResidencyTarget( + less_layerwise = ResidencyTarget( component_name="transformer", residency_mode=LAYERWISE_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, target_resident_weight_bytes=0, - h2d_bytes_per_request=100 * GIB_BYTES, + h2d_bytes_per_request=7 * GIB_BYTES, target_layerwise_resident_layers=(0,), - target_layerwise_pinned_layers=((0,),), - pinned_host_delta_bytes=30 * GIB_BYTES, + target_layerwise_pinned_layers=((),), + target_device_weight_bytes=0, ) + candidates = [current, less_layerwise, component_offload] plan = plan_auto_residency( reports=[ _report( - rank=0, - node_rank=0, - host_pin_capacity_gib=20, - candidates=[pin_more], - ), - _report( - rank=1, - node_rank=1, - host_pin_capacity_gib=20, - candidates=[pin_more], - ), + budget_gib=12, + estimated_gib=2, + phase_peaks_gib={"static:transformer": 2}, + phase_components={"static:transformer": ("transformer",)}, + candidates=candidates, + require_feasible_placement=True, + estimated_request_duration_ns=1, + candidate_latency_savings_ns={ + candidate.option_key(): int(candidate.current_placement) + for candidate in candidates + }, + ) ] ) assert plan.changes == [] - assert plan.resource_budget_bytes["hostpin:node0"] == 20 * GIB_BYTES - assert plan.resource_budget_bytes["hostpin:node1"] == 20 * GIB_BYTES - def test_hostpin_repack_must_fit_transition_headroom(self): - pin_more = ResidencyTarget( - component_name="transformer", - residency_mode=LAYERWISE_OFFLOAD, + def test_pre_warmup_targets_defer_device_weight_promotion(self): + current = 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, - target_layerwise_resident_layers=(0,), - target_layerwise_pinned_layers=((0,),), - pinned_host_delta_bytes=4 * GIB_BYTES, - host_pin_scratch_bytes=6 * GIB_BYTES, + h2d_bytes_per_request=8 * GIB_BYTES, + target_device_weight_bytes=0, + current_placement=True, ) - plan = plan_auto_residency( - reports=[ - _report( - host_pin_capacity_gib=10, - host_transition_headroom_gib=5, - candidates=[pin_more], - ) - ] + resident = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=8 * GIB_BYTES, + h2d_bytes_per_request=0, + target_device_weight_bytes=8 * GIB_BYTES, ) - assert plan.changes == [] - assert plan.resource_budget_bytes["hostram:node0:pin"] == 5 * GIB_BYTES + assert pre_warmup_residency_targets([current, resident]) == [current] - def test_component_demotion_must_fit_host_materialization_headroom(self): - demote = ResidencyTarget( + def test_pre_warmup_targets_keep_active_working_set_demotion(self): + current = 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, + h2d_bytes_per_request=0, + target_device_weight_bytes=0, + current_placement=True, ) - plan = plan_auto_residency( - reports=[ - _report( - host_transition_headroom_gib=5, - candidates=[demote], - ) - ] + layerwise = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=0, + target_device_weight_bytes=GIB_BYTES, + active_device_delta_bytes=-16 * GIB_BYTES, + inactive_device_delta_bytes=GIB_BYTES, + device_transition_delta_bytes=GIB_BYTES, ) - assert plan.changes == [] - assert plan.resource_budget_bytes["hostram:node0:materialize"] == 5 * GIB_BYTES + assert pre_warmup_residency_targets([current, layerwise]) == [ + current, + layerwise, + ] - def test_layerwise_materialization_must_fit_transition_vram(self): - resident = ResidencyTarget( + def test_pre_warmup_report_uses_weight_model_below_legacy_threshold( + self, monkeypatch + ): + from sglang.multimodal_gen.runtime.managers import gpu_worker + + current = 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, + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=0, + h2d_bytes_per_request=8 * GIB_BYTES, + target_device_weight_bytes=0, + current_placement=True, ) - plan = plan_auto_residency( - reports=[ - _report( - budget_gib=100, - estimated_gib=50, - device_transition_allocated_gib=90, - target_workload_measured=True, - candidates=[resident], - ) - ] + modules = {"transformer": object()} + collect_calls = [] + + def collect_targets(_workload, *, allow_host_pin_reallocation=True): + collect_calls.append(allow_host_pin_reallocation) + return [current] + + worker = SimpleNamespace( + rank=0, + pipeline=SimpleNamespace(preload_residency_excluded_components=frozenset()), + server_args=SimpleNamespace( + num_gpus=1, + nnodes=1, + node_rank=0, + host_pin_budget=lambda: None, + ), + _auto_residency_budget_bytes=lambda: 30 * GIB_BYTES, + _collect_auto_residency_targets=collect_targets, + _auto_residency_modules=lambda: modules, + ) + monkeypatch.setattr( + gpu_worker, + "component_runtime_weight_bytes", + lambda _modules: {"transformer": 8 * GIB_BYTES}, + ) + monkeypatch.setattr( + gpu_worker, + "component_current_device_weight_bytes", + lambda _modules: {"transformer": 0}, + ) + monkeypatch.setattr( + gpu_worker, "layerwise_pinned_host_bytes", lambda *a, **k: 0 ) - - assert plan.changes == [] - assert ( - plan.resource_budget_bytes["gpu:rank0:placement-transition"] - == 5 * GIB_BYTES + monkeypatch.setattr( + gpu_worker, "layerwise_host_pin_capacity_bytes", lambda *a, **k: 0 ) - - 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, + monkeypatch.setattr( + gpu_worker, "host_memory_available_bytes", lambda: 100 * 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, + monkeypatch.setattr( + gpu_worker.torch, + "get_device_module", + lambda: SimpleNamespace(memory_allocated=lambda: 2 * 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], - ) - ] + + report = gpu_worker.GPUWorker._build_pre_warmup_auto_residency_report( + worker, + workload=DefaultWorkload( + width=832, + height=480, + num_frames=24, + num_inference_steps=50, + ), ) - assert {candidate.component_name for candidate in plan.changes} == { - "cold_transformer", - "hot_transformer", - } + assert report.budget_bytes == 30 * GIB_BYTES + assert report.estimated_peak_bytes == 10 * GIB_BYTES + assert report.candidates == [current] + assert report.skip_reason is None + assert report.estimated_request_duration_ns == 1 + assert report.candidate_latency_savings_ns == {current.option_key(): 1} + assert collect_calls == [False] - 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( + def test_pre_warmup_report_keeps_excluded_components_fixed(self, monkeypatch): + from sglang.multimodal_gen.runtime.managers import gpu_worker + + transformer = 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, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=20 * GIB_BYTES, + h2d_bytes_per_request=20 * GIB_BYTES, + permanent_residency=True, ) - encoder_resident = ResidencyTarget( - component_name="text_encoder", + vae = ResidencyTarget( + component_name="vae", 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, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=GIB_BYTES, + permanent_residency=True, + target_device_weight_bytes=GIB_BYTES, + current_placement=True, ) - 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, - ) - ] + modules = {"transformer": object(), "vae": object()} + worker = SimpleNamespace( + rank=0, + pipeline=SimpleNamespace( + preload_residency_excluded_components=frozenset(("transformer",)) + ), + server_args=SimpleNamespace( + num_gpus=1, + nnodes=1, + node_rank=0, + host_pin_budget=lambda: None, + ), + _auto_residency_budget_bytes=lambda: 80 * GIB_BYTES, + _collect_auto_residency_targets=lambda _workload, **_kwargs: [ + transformer, + vae, + ], + _auto_residency_modules=lambda: modules, + ) + monkeypatch.setattr( + gpu_worker, + "component_runtime_weight_bytes", + lambda _modules: {"transformer": 20 * GIB_BYTES, "vae": GIB_BYTES}, + ) + monkeypatch.setattr( + gpu_worker, + "component_current_device_weight_bytes", + lambda _modules: {"transformer": 0, "vae": 0}, + ) + monkeypatch.setattr( + gpu_worker, "layerwise_pinned_host_bytes", lambda *a, **k: 0 + ) + monkeypatch.setattr( + gpu_worker, "layerwise_host_pin_capacity_bytes", lambda *a, **k: 0 + ) + monkeypatch.setattr( + gpu_worker, "host_memory_available_bytes", lambda: 100 * GIB_BYTES + ) + monkeypatch.setattr( + gpu_worker.torch, + "get_device_module", + lambda: SimpleNamespace(memory_allocated=lambda: 2 * GIB_BYTES), ) - assert {candidate.option_key() for candidate in plan.changes} == { - "transformer:stage:layers=0:pins=-", - "text_encoder:resident", - } + report = gpu_worker.GPUWorker._build_pre_warmup_auto_residency_report( + worker, + workload=DefaultWorkload( + width=1920, + height=1088, + num_frames=24, + num_inference_steps=15, + ), + ) + + assert [candidate.component_name for candidate in report.candidates] == ["vae"] + assert report.estimated_peak_bytes_by_phase["static:transformer"] == ( + 22 * GIB_BYTES + ) + assert report.used_components_by_phase["static:transformer"] == ("transformer",) class _FakeLayerwiseManager: @@ -3189,70 +3139,13 @@ def test_measured_resident_frontier_streams_one_shot_and_unused_groups(self): {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_measured_resident_frontier_keeps_current_one_shot_layout(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)), - current_resident_layers=(2, 1), - ) - - assert targets == [(0, 0), (2, 1)] - - def test_large_measured_resident_frontier_is_bounded_and_keeps_anchors(self): - managers = [ - _FakeLayerwiseManager( - {f"first.{index}.w": torch.zeros(16) for index in range(70)} - ), - _FakeLayerwiseManager( - {f"second.{index}.w": torch.zeros(16) for index in range(70)} - ), - ] - - targets = _layerwise_resident_targets( - managers, - layer_uses=(tuple([8] * 70), tuple([8] * 70)), - current_resident_layers=(13, 17), - ) - - assert len(targets) <= MAX_LAYERWISE_RESIDENT_TARGETS - assert (0, 0) in targets - assert (70, 70) in targets - assert (70, 0) in targets - assert (0, 70) in targets - assert (13, 17) in targets - - def test_large_policy_frontier_is_bounded_and_keeps_extremes(self): - managers = [ - SimpleNamespace(num_layers=4, residency_policy=RESIDENCY_POLICY_LEADING) - for _ in range(80) - ] - - targets = _layerwise_policy_targets( - managers=managers, - resident_layers=tuple(2 for _ in managers), - tune_policy=True, + + targets = _layerwise_resident_targets( + managers, + layer_uses=((0, 0, 0), (1, 1)), ) - assert len(targets) == MAX_LAYERWISE_POLICY_TARGETS - assert tuple(RESIDENCY_POLICY_LEADING for _ in managers) in targets - assert tuple(RESIDENCY_POLICY_STRIDED for _ in managers) in targets + assert targets == [(0, 0)] def test_host_pin_frontier_includes_non_prefix_packings(self): manager = SimpleNamespace( @@ -3294,29 +3187,6 @@ def test_host_pin_frontier_coalesces_equal_transformer_layers(self): assert len(targets) == 41 - def test_large_host_pin_frontier_is_bounded_and_keeps_extremes(self): - layer_bytes = {index: 1 << index for index in range(24)} - manager = SimpleNamespace( - num_layers=24, - 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=1, - max_targets=16, - ) - - assert len(targets) <= 16 - assert ((),) in targets - assert (tuple(layer_bytes),) in targets - def test_nonbinding_host_resources_keep_only_maximum_pin_utility(self): manager = SimpleNamespace( num_layers=3, @@ -3364,31 +3234,6 @@ def test_collapsed_pin_frontier_retains_the_measured_state(self): assert len(current) == 1 assert current[0].target_layerwise_pinned_layers == ((0,),) - def test_component_latency_cap_keeps_the_full_pin_frontier(self): - manager = _FakeLayerwiseManager( - {f"layers.{index}.w": torch.zeros(16) for index in range(3)} - ) - layer_bytes = {index: GIB_BYTES for index in range(3)} - manager.layer_weight_bytes = lambda: layer_bytes - manager.layer_host_store_bytes = lambda: layer_bytes - 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=16 * GIB_BYTES, - host_pin_headroom_bytes=16 * GIB_BYTES, - request_duration_ns=1_000_000_000, - latency_upper_bound_ns_by_component={"transformer": 1}, - ) - - assert len(candidates) > 8 - def test_host_pin_frontier_drops_scratch_only_tradeoffs_when_nonbinding(self): manager = SimpleNamespace( num_layers=3, @@ -3442,62 +3287,334 @@ def test_filters_and_sizes(self): } ) candidates = collect_residency_targets( - 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, + 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, + ) + permanent_candidates = { + candidate.component_name: candidate + for candidate in candidates + if candidate.permanent_residency + } + # Unset resident components remain tunable. Explicit placement, + # custom strategy, and non-module objects stay outside the frontier. + assert set(permanent_candidates) == {"text_encoder", "transformer", "vae"} + assert permanent_candidates["vae"].current_placement + assert permanent_candidates["text_encoder"].h2d_bytes_per_request == ( + permanent_candidates["text_encoder"].target_resident_weight_bytes + ) + # The all-pageable layerwise DiT re-streams its layers every step. + assert permanent_candidates["transformer"].h2d_bytes_per_request == ( + permanent_candidates["transformer"].target_resident_weight_bytes + * 50 + * PAGEABLE_H2D_COST_MULTIPLIER + ) + 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): + # A module the compile window manages layerwise while its configured + # mode says component-offload: adjustment would be a silent no-op. + mismatch = _FakeLayerwiseDit( + [_FakeLayerwiseManager({"layers.0.w": torch.zeros(16)})] + ) + candidates = collect_residency_targets( + 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 == [] + + 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_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, + ) + + 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.target_resident_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( + { + "layers.0.w": torch.zeros(16), + "layers.1.w": torch.zeros(16), + } + ) + manager._pinned_layers = (0, 1) + 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=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_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, + 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 + + 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_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=10, + mixed_dtype_components={"vae"}, + ) + + assert candidates + assert {candidate.target_mode() for candidate in candidates} == { + LAYERWISE_OFFLOAD + } + + 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 + ) + 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_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) + + 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"), ) - permanent_candidates = { - candidate.component_name: candidate + + component = next( + candidate for candidate in candidates - if candidate.permanent_residency - } - # Unset resident components remain tunable. Explicit placement, - # custom strategy, and non-module objects stay outside the frontier. - assert set(permanent_candidates) == {"text_encoder", "transformer", "vae"} - assert permanent_candidates["vae"].current_placement - assert permanent_candidates["text_encoder"].h2d_bytes_per_request == ( - permanent_candidates["text_encoder"].target_resident_weight_bytes + if candidate.target_mode() == COMPONENT_OFFLOAD ) - # The all-pageable layerwise DiT re-streams its layers every step. - assert permanent_candidates["transformer"].h2d_bytes_per_request == ( - permanent_candidates["transformer"].target_resident_weight_bytes - * 50 - * PAGEABLE_H2D_COST_MULTIPLIER + 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 all( + candidate.h2d_bytes_per_request <= component.h2d_bytes_per_request + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD ) assert any( - candidate.component_name == "transformer" - and candidate.target_layerwise_resident_layers == (1,) + candidate.target_layerwise_resident_layers == (2,) for candidate in candidates ) + resident = next( + candidate for candidate in candidates if candidate.target_mode() == RESIDENT + ) + assert ( + resident.device_transition_delta_bytes + == resident.target_device_weight_bytes + ) + assert ( + resident.concurrent_prefetch_delta() == resident.active_device_delta_bytes + ) - def test_mechanism_mismatch_is_not_a_candidate(self): - # A module the compile window manages layerwise while its configured - # mode says component-offload: adjustment would be a silent no-op. - mismatch = _FakeLayerwiseDit( - [_FakeLayerwiseManager({"layers.0.w": torch.zeros(16)})] + 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 + assert ( + streamed.concurrent_prefetch_delta() == streamed.active_device_delta_bytes ) + + 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={"mismatch": mismatch}, - residency_mode_of=self._modes({"mismatch": COMPONENT_OFFLOAD}), + 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"), ) - assert candidates == [] - def test_layerwise_dit_can_release_existing_resident_layers(self): + 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_layerwise_frontier_accounts_for_immediate_materialization(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]) @@ -3508,55 +3625,50 @@ def test_layerwise_dit_can_release_existing_resident_layers(self): custom_strategy_names=(), num_inference_steps=10, ) - - target = next( + resident = next( + candidate for candidate in candidates if candidate.permanent_residency + ) + component = next( candidate for candidate in candidates - if not candidate.permanent_residency - and candidate.target_layerwise_resident_layers == (1,) - and candidate.target_layerwise_pinned_layers == ((),) + if candidate.target_mode() == COMPONENT_OFFLOAD ) - assert target.target_resident_weight_bytes > 0 - assert target.active_device_delta_bytes < 0 - assert target.h2d_bytes_per_request > 0 - - def test_auto_dit_frontier_tunes_policy_unless_explicit(self): - manager = _FakeLayerwiseManager( - {f"layers.{index}.w": torch.zeros(16) for index in range(4)}, - resident_layers=2, + assert ( + component.concurrent_prefetch_delta() == component.active_device_delta_bytes ) - module = _FakeLayerwiseDit([manager]) - common = dict( + assert ( + resident.device_transition_delta_bytes + == resident.target_device_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": LAYERWISE_OFFLOAD}), + 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, ) - - candidates = collect_residency_targets( - **common, - layerwise_policy_is_explicit=lambda _name, _dit_group: False, - ) - policies = { - candidate.target_layerwise_residency_policies + streamed = next( + candidate for candidate in candidates if candidate.target_mode() == LAYERWISE_OFFLOAD - and candidate.target_layerwise_resident_layers == (2,) - and candidate.target_layerwise_pinned_layers == ((),) - } - assert policies == { - (RESIDENCY_POLICY_LEADING,), - (RESIDENCY_POLICY_STRIDED,), - } - - explicit_candidates = collect_residency_targets( - **common, - layerwise_policy_is_explicit=lambda _name, _dit_group: True, ) - assert all( - candidate.target_layerwise_residency_policies is None - for candidate in explicit_candidates + assert ( + streamed.device_transition_delta_bytes == -manager.offloaded_weight_bytes() + ) + component = next( + candidate + for candidate in candidates + if candidate.target_mode() == COMPONENT_OFFLOAD + ) + assert ( + component.device_transition_delta_bytes + == -resident.target_device_weight_bytes ) def test_auto_dit_frontier_tunes_each_layer_group_policy(self): @@ -3592,23 +3704,24 @@ def test_auto_dit_frontier_tunes_each_layer_group_policy(self): (RESIDENCY_POLICY_STRIDED, RESIDENCY_POLICY_STRIDED), } - def test_repeated_non_dit_group_can_tune_policy(self): + def test_auto_dit_frontier_tunes_policy_unless_explicit(self): manager = _FakeLayerwiseManager( {f"layers.{index}.w": torch.zeros(16) for index in range(4)}, resident_layers=2, ) module = _FakeLayerwiseDit([manager]) - - candidates = collect_residency_targets( - modules={"custom_refiner": module}, - residency_mode_of=self._modes({"custom_refiner": LAYERWISE_OFFLOAD}), + common = dict( + 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, - layerwise_policy_is_explicit=lambda _name, _dit_group: False, - layerwise_layer_uses={"custom_refiner": {"layers": (4, 4, 4, 4)}}, ) + candidates = collect_residency_targets( + **common, + layerwise_policy_is_explicit=lambda _name, _dit_group: False, + ) policies = { candidate.target_layerwise_residency_policies for candidate in candidates @@ -3621,166 +3734,91 @@ def test_repeated_non_dit_group_can_tune_policy(self): (RESIDENCY_POLICY_STRIDED,), } - 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_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=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 + explicit_candidates = collect_residency_targets( + **common, + layerwise_policy_is_explicit=lambda _name, _dit_group: True, + ) + assert all( + candidate.target_layerwise_residency_policies is None + for candidate in explicit_candidates ) - def test_multi_worker_candidate_keeps_existing_hostpin_placement(self): + def test_component_latency_cap_keeps_the_full_pin_frontier(self): manager = _FakeLayerwiseManager( - { - "layers.0.w": torch.zeros(16), - "layers.1.w": torch.zeros(16), - } + {f"layers.{index}.w": torch.zeros(16) for index in range(3)} ) - manager._pinned_layers = (0, 1) + layer_bytes = {index: GIB_BYTES for index in range(3)} + manager.layer_weight_bytes = lambda: layer_bytes + manager.layer_host_store_bytes = lambda: layer_bytes module = _FakeLayerwiseDit([manager]) candidates = collect_residency_targets( modules={"transformer": module}, - residency_mode_of=self._modes({"transformer": LAYERWISE_OFFLOAD}), + residency_mode_of=lambda _name: 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 + num_inference_steps=3, + used_components={"transformer"}, + layerwise_layer_uses={"transformer": {"layers": (3, 3, 3)}}, + host_transition_headroom_bytes=16 * GIB_BYTES, + host_pin_headroom_bytes=16 * GIB_BYTES, + request_duration_ns=1_000_000_000, + latency_upper_bound_ns_by_component={"transformer": 1}, ) - 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]) + assert len(candidates) > 8 + def test_default_resident_component_keeps_offload_as_a_target(self): + module = nn.Linear(4, 4) candidates = collect_residency_targets( - modules={"vae": module}, - residency_mode_of=self._modes({"vae": LAYERWISE_OFFLOAD}), + modules={"text_encoder": module}, + residency_mode_of=self._modes({"text_encoder": RESIDENT}), explicit_residency_mode_of=lambda _name: None, custom_strategy_names=(), - num_inference_steps=10, - mixed_dtype_components={"vae"}, + num_inference_steps=1, ) - assert candidates - assert {candidate.target_mode() for candidate in candidates} == { - LAYERWISE_OFFLOAD - } - - def test_mixed_dtype_coarse_component_does_not_add_virtual_layerwise(self): - module = _FakeLazyLayerwiseDit() + 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 + assert by_mode[COMPONENT_OFFLOAD].device_transition_delta_bytes < 0 + assert by_mode[COMPONENT_OFFLOAD].host_materialize_scratch_bytes > 0 + def test_default_resident_dit_exposes_complete_layerwise_frontier(self): + module = _FakeLazyLayerwiseDit(num_layers=3) candidates = collect_residency_targets( - modules={"vae": module}, - residency_mode_of=self._modes({"vae": COMPONENT_OFFLOAD}), + modules={"transformer_2": module}, + residency_mode_of=self._modes({"transformer_2": RESIDENT}), explicit_residency_mode_of=lambda _name: None, custom_strategy_names=(), - num_inference_steps=10, - mixed_dtype_components={"vae"}, - layerwise_tuning_of=lambda _name, _dit_group: ( - 1.0, - 0.0, - "leading", - ), + num_inference_steps=4, + layerwise_tuning_of=lambda _name, _dit_group: (1.0, 0.0, "leading"), ) assert {candidate.target_mode() for candidate in candidates} == { COMPONENT_OFFLOAD, + LAYERWISE_OFFLOAD, RESIDENT, } - - 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 + assert { + candidate.target_layerwise_resident_layers for candidate in candidates - ) - component = next( - candidate + if candidate.target_mode() == LAYERWISE_OFFLOAD + } == {(0,), (1,), (2,), (3,)} + assert all( + candidate.residency_mode == RESIDENT 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_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", + if candidate.target_mode() == LAYERWISE_OFFLOAD ) - 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), - } - }, + resident = next( + candidate for candidate in candidates if candidate.target_mode() == RESIDENT ) - layerwise = [ - candidate + assert { + candidate.host_materialize_scratch_bytes 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 - ) + } == {resident.target_device_weight_bytes} @pytest.mark.parametrize( "component_name", @@ -3846,127 +3884,136 @@ def test_every_serving_component_group_gets_a_complete_frontier( for candidate in candidates ) - 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"), + def test_large_host_pin_frontier_is_bounded_and_keeps_extremes(self): + layer_bytes = {index: 1 << index for index in range(24)} + manager = SimpleNamespace( + num_layers=24, + 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), ) - 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 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 + targets = _layerwise_pin_targets( + managers=[manager], + resident_layers=(0,), + current_pinned_layers=((),), + uses_per_streamed_layer=1, + max_targets=16, ) - resident = next( - candidate for candidate in candidates if candidate.target_mode() == RESIDENT + + assert len(targets) <= 16 + assert ((),) in targets + assert (tuple(layer_bytes),) in targets + + def test_large_measured_resident_frontier_is_bounded_and_keeps_anchors(self): + managers = [ + _FakeLayerwiseManager( + {f"first.{index}.w": torch.zeros(16) for index in range(70)} + ), + _FakeLayerwiseManager( + {f"second.{index}.w": torch.zeros(16) for index in range(70)} + ), + ] + + targets = _layerwise_resident_targets( + managers, + layer_uses=(tuple([8] * 70), tuple([8] * 70)), + current_resident_layers=(13, 17), ) - assert ( - resident.device_transition_delta_bytes - == resident.target_device_weight_bytes + + assert len(targets) <= MAX_LAYERWISE_RESIDENT_TARGETS + assert (0, 0) in targets + assert (70, 70) in targets + assert (70, 0) in targets + assert (0, 70) in targets + assert (13, 17) in targets + + def test_large_policy_frontier_is_bounded_and_keeps_extremes(self): + managers = [ + SimpleNamespace(num_layers=4, residency_policy=RESIDENCY_POLICY_LEADING) + for _ in range(80) + ] + + targets = _layerwise_policy_targets( + managers=managers, + resident_layers=tuple(2 for _ in managers), + tune_policy=True, ) - assert ( - resident.concurrent_prefetch_delta() == resident.active_device_delta_bytes + + assert len(targets) == MAX_LAYERWISE_POLICY_TARGETS + assert tuple(RESIDENCY_POLICY_LEADING for _ in managers) in targets + assert tuple(RESIDENCY_POLICY_STRIDED for _ in managers) in targets + + def test_measured_resident_frontier_keeps_current_one_shot_layout(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)), + current_resident_layers=(2, 1), ) - def test_virtual_layerwise_prefetch_is_runtime_not_transition_memory(self): - module = _FakeLazyLayerwiseDit(num_layers=3) + assert targets == [(0, 0), (2, 1)] + + def test_mixed_dtype_coarse_component_does_not_add_virtual_layerwise(self): + module = _FakeLazyLayerwiseDit() + candidates = collect_residency_targets( - modules={"transformer": module}, - residency_mode_of=self._modes({"transformer": COMPONENT_OFFLOAD}), + modules={"vae": module}, + residency_mode_of=self._modes({"vae": 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"), + num_inference_steps=10, + mixed_dtype_components={"vae"}, + 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 - assert ( - streamed.concurrent_prefetch_delta() == streamed.active_device_delta_bytes - ) + assert {candidate.target_mode() for candidate in candidates} == { + COMPONENT_OFFLOAD, + RESIDENT, + } - 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)) + def test_repeated_non_dit_group_can_tune_policy(self): + manager = _FakeLayerwiseManager( + {f"layers.{index}.w": torch.zeros(16) for index in range(4)}, + resident_layers=2, + ) + module = _FakeLayerwiseDit([manager]) candidates = collect_residency_targets( - modules={"transformer": module}, - residency_mode_of=self._modes({"transformer": COMPONENT_OFFLOAD}), + modules={"custom_refiner": module}, + residency_mode_of=self._modes({"custom_refiner": LAYERWISE_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"), + layerwise_policy_is_explicit=lambda _name, _dit_group: False, + layerwise_layer_uses={"custom_refiner": {"layers": (4, 4, 4, 4)}}, ) - full_stage = next( - candidate + policies = { + candidate.target_layerwise_residency_policies 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() + and candidate.target_layerwise_pinned_layers == ((),) + } + assert policies == { + (RESIDENCY_POLICY_LEADING,), + (RESIDENCY_POLICY_STRIDED,), } - - def test_default_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}), - explicit_residency_mode_of=lambda _name: None, - custom_strategy_names=(), - num_inference_steps=1, - ) - - 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 - assert by_mode[COMPONENT_OFFLOAD].device_transition_delta_bytes < 0 - assert by_mode[COMPONENT_OFFLOAD].host_materialize_scratch_bytes > 0 def test_required_resident_component_has_no_tunable_frontier(self): candidates = collect_residency_targets( @@ -3980,41 +4027,6 @@ def test_required_resident_component_has_no_tunable_frontier(self): assert candidates == [] - def test_default_resident_dit_exposes_complete_layerwise_frontier(self): - module = _FakeLazyLayerwiseDit(num_layers=3) - candidates = collect_residency_targets( - modules={"transformer_2": module}, - residency_mode_of=self._modes({"transformer_2": RESIDENT}), - 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"), - ) - - assert {candidate.target_mode() for candidate in candidates} == { - COMPONENT_OFFLOAD, - LAYERWISE_OFFLOAD, - RESIDENT, - } - assert { - candidate.target_layerwise_resident_layers - for candidate in candidates - if candidate.target_mode() == LAYERWISE_OFFLOAD - } == {(0,), (1,), (2,), (3,)} - assert all( - candidate.residency_mode == RESIDENT - for candidate in candidates - if candidate.target_mode() == LAYERWISE_OFFLOAD - ) - resident = next( - candidate for candidate in candidates if candidate.target_mode() == RESIDENT - ) - assert { - candidate.host_materialize_scratch_bytes - for candidate in candidates - if candidate.target_mode() == LAYERWISE_OFFLOAD - } == {resident.target_device_weight_bytes} - def test_resident_seed_preserves_unconfigured_layerwise_frontier(self): module = _FakeLazyLayerwiseDit(num_layers=3) candidates = collect_residency_targets( @@ -4037,34 +4049,6 @@ def test_resident_seed_preserves_unconfigured_layerwise_frontier(self): candidate for candidate in candidates if candidate.target_mode() == RESIDENT ).current_placement - def test_virtual_layerwise_reuses_component_offload_host_weights(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 - resident = next( - candidate for candidate in candidates if candidate.target_mode() == RESIDENT - ) - assert ( - 0 - < streamed.host_materialize_scratch_bytes - < resident.target_device_weight_bytes - ) - def test_virtual_layerwise_frontier_includes_host_pin_prefixes(self): module = _FakeLazyLayerwiseDit(num_layers=3) candidates = collect_residency_targets( @@ -4098,137 +4082,36 @@ def test_virtual_layerwise_frontier_includes_host_pin_prefixes(self): fully_pinned.host_pin_scratch_bytes == fully_pinned.pinned_host_delta_bytes ) - 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 - ) - component = next( - candidate - for candidate in candidates - if candidate.target_mode() == COMPONENT_OFFLOAD - ) - assert ( - component.concurrent_prefetch_delta() == component.active_device_delta_bytes - ) - assert ( - resident.device_transition_delta_bytes - == resident.target_device_weight_bytes - ) - module.disable_offload() - module.register_parameter( - "materialized", - nn.Parameter(torch.zeros(manager.offloaded_weight_bytes() // 4)), - ) + def test_virtual_layerwise_reuses_component_offload_host_weights(self): + module = _FakeLazyLayerwiseDit(num_layers=3) candidates = collect_residency_targets( modules={"transformer": module}, - residency_mode_of=self._modes({"transformer": RESIDENT}), - baseline_residency_mode_of=self._modes({"transformer": LAYERWISE_OFFLOAD}), + residency_mode_of=self._modes({"transformer": COMPONENT_OFFLOAD}), explicit_residency_mode_of=lambda _name: None, - custom_strategy_names=(), - num_inference_steps=10, - ) - streamed = next( - candidate - for candidate in candidates - if candidate.target_mode() == LAYERWISE_OFFLOAD - ) - assert ( - streamed.device_transition_delta_bytes == -manager.offloaded_weight_bytes() - ) - component = next( - candidate - for candidate in candidates - if candidate.target_mode() == COMPONENT_OFFLOAD - ) - assert ( - component.device_transition_delta_bytes - == -resident.target_device_weight_bytes - ) - - -class TestApplyAndRollback: - def test_apply_and_rollback_restore_rank_hostpin_quota(self): - budget = HostPinBudget( - available_bytes=40 * GIB_BYTES, - reserve_bytes=2 * GIB_BYTES, - ) - args = _StubResidencyArgs(host_pin_budget=budget) - target = ResidencyTarget( - component_name="text_encoder", - residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=RESIDENT, - target_resident_weight_bytes=GIB_BYTES, - h2d_bytes_per_request=GIB_BYTES, - ) - plan = AutoResidencyPlan( - host_pin_target_bytes_by_rank={0: 12 * GIB_BYTES}, - changes=[target], - ) - - applied = apply_residency_changes( - plan=plan, - modules={"text_encoder": nn.Linear(2, 2)}, - server_args=args, - rank=0, - ) - assert budget.available_bytes == 12 * GIB_BYTES - assert budget.reserve_bytes == 0 - - rollback_residency_changes( - applied=applied, - modules={"text_encoder": nn.Linear(2, 2)}, - server_args=args, + custom_strategy_names=(), + num_inference_steps=4, + layerwise_tuning_of=lambda _name, _dit_group: (1.0, 0.0, "leading"), ) - assert budget.available_bytes == 40 * GIB_BYTES - assert budget.reserve_bytes == 2 * GIB_BYTES - assert budget.planning_capacity_bytes == 38 * GIB_BYTES - def test_apply_failure_restores_rank_hostpin_quota(self): - budget = HostPinBudget( - available_bytes=40 * GIB_BYTES, - reserve_bytes=2 * GIB_BYTES, + streamed = next( + candidate + for candidate in candidates + if candidate.target_mode() == LAYERWISE_OFFLOAD + and candidate.target_layerwise_resident_layers == (0,) ) - args = _StubResidencyArgs(host_pin_budget=budget) - plan = AutoResidencyPlan( - host_pin_target_bytes_by_rank={0: 12 * GIB_BYTES}, - changes=[ - ResidencyTarget( - component_name="missing", - residency_mode=COMPONENT_OFFLOAD, - target_residency_mode=RESIDENT, - target_resident_weight_bytes=GIB_BYTES, - h2d_bytes_per_request=GIB_BYTES, - ) - ], + assert streamed.device_transition_delta_bytes == 0 + assert streamed.active_device_delta_bytes < 0 + resident = next( + candidate for candidate in candidates if candidate.target_mode() == RESIDENT + ) + assert ( + 0 + < streamed.host_materialize_scratch_bytes + < resident.target_device_weight_bytes ) - with pytest.raises(RuntimeError, match="is missing"): - apply_residency_changes( - plan=plan, - modules={}, - server_args=args, - rank=0, - ) - - assert budget.available_bytes == 40 * GIB_BYTES - assert budget.reserve_bytes == 2 * GIB_BYTES - assert budget.planning_capacity_bytes == 38 * GIB_BYTES +class TestApplyAndRollback: def test_lazy_layerwise_configuration_rolls_back_to_component_offload(self): module = _FakeLazyLayerwiseDit(num_layers=4) args = _StubResidencyArgs() @@ -4463,38 +4346,6 @@ def test_partial_layerwise_promotion_restores_exact_group_counts(self): ) assert [manager.resident_layers for manager in managers] == [1, 0] - def test_partial_layerwise_policy_change_rolls_back_exact_layout(self): - manager = _FakeLayerwiseManager( - {f"layers.{index}.w": torch.zeros(16) for index in range(4)}, - resident_layers=2, - ) - module = _FakeLayerwiseDit([manager]) - args = _StubResidencyArgs() - candidate = ResidencyTarget( - component_name="transformer", - residency_mode=LAYERWISE_OFFLOAD, - target_resident_weight_bytes=manager.resident_weight_bytes(2), - h2d_bytes_per_request=GIB_BYTES, - target_layerwise_resident_layers=(2,), - target_layerwise_residency_policies=(RESIDENCY_POLICY_STRIDED,), - target_layerwise_pinned_layers=((),), - ) - - applied = apply_residency_changes( - plan=_plan_for([candidate]), - modules={"transformer": module}, - server_args=args, - ) - assert manager.residency_policy == RESIDENCY_POLICY_STRIDED - - rollback_residency_changes( - applied=applied, - modules={"transformer": module}, - server_args=args, - ) - assert manager.resident_layers == 2 - assert manager.residency_policy == RESIDENCY_POLICY_LEADING - def test_permanent_layerwise_can_be_demoted_and_rollback_restores_it(self): manager = _FakeLayerwiseManager( { @@ -4837,6 +4688,105 @@ def clear_auto_residency_mode(self, component_name): # the earlier-applied adjustment was still undone assert released == ["text_encoder"] + def test_apply_and_rollback_restore_rank_hostpin_quota(self): + budget = HostPinBudget( + available_bytes=40 * GIB_BYTES, + reserve_bytes=2 * GIB_BYTES, + ) + args = _StubResidencyArgs(host_pin_budget=budget) + target = ResidencyTarget( + component_name="text_encoder", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=GIB_BYTES, + ) + plan = AutoResidencyPlan( + host_pin_target_bytes_by_rank={0: 12 * GIB_BYTES}, + changes=[target], + ) + + applied = apply_residency_changes( + plan=plan, + modules={"text_encoder": nn.Linear(2, 2)}, + server_args=args, + rank=0, + ) + assert budget.available_bytes == 12 * GIB_BYTES + assert budget.reserve_bytes == 0 + + rollback_residency_changes( + applied=applied, + modules={"text_encoder": nn.Linear(2, 2)}, + server_args=args, + ) + assert budget.available_bytes == 40 * GIB_BYTES + assert budget.reserve_bytes == 2 * GIB_BYTES + assert budget.planning_capacity_bytes == 38 * GIB_BYTES + + def test_apply_failure_restores_rank_hostpin_quota(self): + budget = HostPinBudget( + available_bytes=40 * GIB_BYTES, + reserve_bytes=2 * GIB_BYTES, + ) + args = _StubResidencyArgs(host_pin_budget=budget) + plan = AutoResidencyPlan( + host_pin_target_bytes_by_rank={0: 12 * GIB_BYTES}, + changes=[ + ResidencyTarget( + component_name="missing", + residency_mode=COMPONENT_OFFLOAD, + target_residency_mode=RESIDENT, + target_resident_weight_bytes=GIB_BYTES, + h2d_bytes_per_request=GIB_BYTES, + ) + ], + ) + + with pytest.raises(RuntimeError, match="is missing"): + apply_residency_changes( + plan=plan, + modules={}, + server_args=args, + rank=0, + ) + + assert budget.available_bytes == 40 * GIB_BYTES + assert budget.reserve_bytes == 2 * GIB_BYTES + assert budget.planning_capacity_bytes == 38 * GIB_BYTES + + def test_partial_layerwise_policy_change_rolls_back_exact_layout(self): + manager = _FakeLayerwiseManager( + {f"layers.{index}.w": torch.zeros(16) for index in range(4)}, + resident_layers=2, + ) + module = _FakeLayerwiseDit([manager]) + args = _StubResidencyArgs() + candidate = ResidencyTarget( + component_name="transformer", + residency_mode=LAYERWISE_OFFLOAD, + target_resident_weight_bytes=manager.resident_weight_bytes(2), + h2d_bytes_per_request=GIB_BYTES, + target_layerwise_resident_layers=(2,), + target_layerwise_residency_policies=(RESIDENCY_POLICY_STRIDED,), + target_layerwise_pinned_layers=((),), + ) + + applied = apply_residency_changes( + plan=_plan_for([candidate]), + modules={"transformer": module}, + server_args=args, + ) + assert manager.residency_policy == RESIDENCY_POLICY_STRIDED + + rollback_residency_changes( + applied=applied, + modules={"transformer": module}, + server_args=args, + ) + assert manager.resident_layers == 2 + assert manager.residency_policy == RESIDENCY_POLICY_LEADING + class TestCandidateRanking: def test_post_request_hint_order_matches_promotion_order(self): @@ -4864,6 +4814,8 @@ def test_post_request_hint_order_matches_promotion_order(self): class TestAppliedChangesLog: + pass + def test_reports_runtime_transition_and_kill_switch(self): plan = _plan_for( [ @@ -5074,19 +5026,6 @@ def _skip_reason(self, args): return auto_residency_skip_reason(args) - def _static_skip_reason(self, args): - from sglang.multimodal_gen.runtime.server_args.auto_tune import ( - auto_residency_static_skip_reason, - ) - - return auto_residency_static_skip_reason(args) - - def test_static_planning_does_not_require_warmup(self, monkeypatch): - monkeypatch.delenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", raising=False) - monkeypatch.delenv("SGLANG_CACHE_DIT_ENABLED", raising=False) - assert self._static_skip_reason(self._base_args(warmup_mode="off")) is None - assert "server warmup" in self._skip_reason(self._base_args(warmup_mode="off")) - def test_env_kill_switch(self, monkeypatch): monkeypatch.setenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", "1") reason = self._skip_reason(self._base_args()) @@ -5218,3 +5157,66 @@ def test_transformer_weight_override_reaches_platform_gate(self, monkeypatch): self._base_args(transformer_weights_path="/x.safetensors") ) assert reason is None or reason == "requires CUDA" + + def _static_skip_reason(self, args): + from sglang.multimodal_gen.runtime.server_args.auto_tune import ( + auto_residency_static_skip_reason, + ) + + return auto_residency_static_skip_reason(args) + + def test_static_planning_does_not_require_warmup(self, monkeypatch): + monkeypatch.delenv("SGLANG_DIFFUSION_DISABLE_AUTO_RESIDENCY", raising=False) + monkeypatch.delenv("SGLANG_CACHE_DIT_ENABLED", raising=False) + assert self._static_skip_reason(self._base_args(warmup_mode="off")) is None + assert "server warmup" in self._skip_reason(self._base_args(warmup_mode="off")) + + +class TestEstimateAllocatorHeadroom: + def test_uses_covering_workload_allocator_gap(self): + small = _record( + num_frames=9, + peak_gib=10, + peak_reserved_gib=11, + ) + target = _record( + num_frames=17, + peak_gib=20, + peak_reserved_gib=26, + ) + + assert ( + estimate_allocator_headroom_bytes( + records=[small, target], target_units=target.workload_units() + ) + == 6 * GIB_BYTES + ) + + def test_keeps_largest_observed_gap_without_covering_workload(self): + first = _record(peak_gib=10, peak_reserved_gib=12) + second = _record(peak_gib=15, peak_reserved_gib=19) + + assert ( + estimate_allocator_headroom_bytes( + records=[first, second], target_units=second.workload_units() * 2 + ) + == 4 * GIB_BYTES + ) + + +def test_qwen_image_2511_excludes_transformer_only_during_preload(): + from sglang.multimodal_gen.runtime.pipelines.qwen_image import ( + QwenImageEditPlusPipeline, + ) + + pipeline = QwenImageEditPlusPipeline.__new__(QwenImageEditPlusPipeline) + + pipeline.server_args = SimpleNamespace( + pipeline_config=QwenImageEditPlus_2511_PipelineConfig() + ) + assert pipeline.preload_residency_excluded_components == frozenset(("transformer",)) + + pipeline.server_args = SimpleNamespace( + pipeline_config=QwenImageEditPlusPipelineConfig() + ) + assert pipeline.preload_residency_excluded_components == frozenset() 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 71d97cde8ef4..542cedc94e1b 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 @@ -78,214 +78,6 @@ def test_rewarm_step_limit_keeps_full_shape_but_shortens_denoising(self): self.assertEqual(req.extra["warmup_total"], 1) self.assertTrue(req.extra["server_warmup_rewarm"]) - def test_residency_hint_uses_effective_device_budget(self): - worker = GPUWorker.__new__(GPUWorker) - worker._auto_residency_budget_bytes = mock.Mock(return_value=12 * GIB_BYTES) - worker.get_can_stay_resident_components = mock.Mock(return_value=[]) - output = SimpleNamespace(metrics=None) - snapshot = SimpleNamespace( - peak_reserved_mb=11 * 1024, - peak_allocated_mb=10 * 1024, - ) - - with mock.patch.object( - gpu_worker_module, "capture_memory_snapshot", return_value=snapshot - ): - worker.do_mem_analysis(output) - - worker.get_can_stay_resident_components.assert_called_once_with(1.0) - self.assertEqual(output.peak_memory_mb, 11 * 1024) - - def test_sync_startup_uses_shared_auto_residency_sequence(self): - server_args = SimpleNamespace() - forward = mock.Mock() - - with ( - mock.patch.object( - server_warmup, "maybe_apply_pre_warmup_auto_residency" - ) as apply_static, - mock.patch.object(server_warmup, "run_async_client_warmup") as warmup, - mock.patch.object(server_warmup, "maybe_apply_auto_residency") as refine, - mock.patch.object( - server_warmup, - "should_run_synthetic_server_warmup", - return_value=True, - ), - ): - server_warmup.run_sync_startup_warmup(server_args, forward) - - apply_static.assert_awaited_once() - warmup.assert_awaited_once() - refine.assert_awaited_once() - self.assertIs(apply_static.await_args.args[0], server_args) - self.assertIs(warmup.await_args.args[0], server_args) - self.assertIs(refine.await_args.args[0], server_args) - self.assertTrue(warmup.await_args.kwargs["fail_open"]) - self.assertIs(apply_static.await_args.args[1], warmup.await_args.args[1]) - self.assertIs(apply_static.await_args.args[1], refine.await_args.args[1]) - - def test_sync_startup_static_planning_does_not_require_warmup(self): - with ( - mock.patch.object( - server_warmup, "maybe_apply_pre_warmup_auto_residency" - ) as apply_static, - mock.patch.object( - server_warmup, - "should_run_synthetic_server_warmup", - return_value=False, - ), - mock.patch.object(server_warmup, "run_async_client_warmup") as warmup, - mock.patch.object(server_warmup, "maybe_apply_auto_residency") as refine, - ): - server_warmup.run_sync_startup_warmup(SimpleNamespace(), mock.Mock()) - - apply_static.assert_awaited_once() - warmup.assert_not_awaited() - refine.assert_not_awaited() - - def test_pre_warmup_planner_uses_static_action(self): - actions = [] - - async def forward(req): - actions.append(req.action) - return OutputBatch(output={"status": PLACEMENT_STATUS_SKIPPED}) - - with mock.patch.object( - server_warmup, - "should_apply_pre_warmup_auto_residency", - return_value=True, - ): - asyncio.run( - server_warmup.maybe_apply_pre_warmup_auto_residency( - SimpleNamespace(), forward - ) - ) - - self.assertEqual(actions, ["apply_static"]) - - def test_pre_warmup_worker_commits_seed_without_pending_round(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 = [] - worker._auto_residency_applied = [] - worker._auto_residency_round_sizes = [] - worker._auto_residency_last_applied_plan = None - report = RankResidencyReport( - rank=0, - budget_bytes=10 * GIB_BYTES, - estimated_peak_bytes=5 * GIB_BYTES, - ) - worker._build_pre_warmup_auto_residency_report = mock.Mock(return_value=report) - worker._auto_residency_all_gather = mock.Mock(side_effect=lambda value: [value]) - worker._auto_residency_modules = mock.Mock(return_value={}) - worker._invalidate_component_strategies = mock.Mock() - candidate = ResidencyTarget( - component_name="transformer", - residency_mode=RESIDENT, - target_resident_weight_bytes=0, - h2d_bytes_per_request=1, - target_residency_mode=COMPONENT_OFFLOAD, - ) - plan = AutoResidencyPlan(changes=[candidate]) - applied = AppliedResidencyChange("transformer", RESIDENT) - workload = DefaultWorkload(512, 512, 1, 4) - - 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 - ), - mock.patch.object( - gpu_worker_module, "apply_residency_changes", return_value=[applied] - ), - mock.patch.object(gpu_worker_module, "commit_residency_changes") as commit, - ): - response = worker.apply_auto_residency(pre_warmup=True) - - self.assertEqual(response.output["status"], PLACEMENT_STATUS_ADJUSTED) - commit.assert_called_once_with( - applied=[applied], modules={}, server_args=worker.server_args - ) - self.assertEqual(worker._auto_residency_applied, []) - self.assertEqual(worker._auto_residency_round_sizes, []) - self.assertIsNone(worker._auto_residency_last_applied_plan) - - def test_auto_residency_uses_first_oom_without_degrading_probe(self): - req = SimpleNamespace() - forward = mock.AsyncMock(return_value=OutputBatch(error="CUDA out of memory")) - - with ( - mock.patch.object( - server_warmup, "auto_residency_skip_reason", return_value=None - ), - mock.patch.object( - server_warmup, - "should_include_warmup_image", - return_value=False, - ), - mock.patch.object( - server_warmup, "build_client_warmup_reqs", return_value=[req] - ), - mock.patch.object(server_warmup, "_degrade_after_oom") as degrade, - ): - asyncio.run( - server_warmup.run_async_client_warmup( - SimpleNamespace(), forward, fail_open=True - ) - ) - - forward.assert_awaited_once_with(req) - degrade.assert_not_called() - - def test_non_auto_warmup_keeps_oom_probe_degradation(self): - req = SimpleNamespace() - lighter = SimpleNamespace() - forward = mock.AsyncMock( - side_effect=[ - OutputBatch(error="CUDA out of memory"), - OutputBatch(), - ] - ) - - with ( - mock.patch.object( - server_warmup, - "auto_residency_skip_reason", - return_value="not auto", - ), - mock.patch.object( - server_warmup, - "should_include_warmup_image", - return_value=False, - ), - mock.patch.object( - server_warmup, "build_client_warmup_reqs", return_value=[req] - ), - mock.patch.object( - server_warmup, "_degrade_after_oom", return_value=lighter - ) as degrade, - ): - asyncio.run( - server_warmup.run_async_client_warmup(SimpleNamespace(), forward) - ) - - self.assertEqual( - [call.args[0] for call in forward.await_args_list], - [req, lighter], - ) - degrade.assert_called_once_with(mock.ANY, req) - def test_worker_rolls_back_a_materially_slower_calibrated_round(self): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 @@ -334,30 +126,35 @@ def test_worker_rolls_back_a_materially_slower_calibrated_round(self): self.assertFalse(kwargs["already_failed"]) self.assertTrue(kwargs["latest_round_only"]) - def test_worker_keeps_resident_only_round_despite_noisy_duration(self): + 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: RESIDENT) + worker.pipeline = SimpleNamespace(modules={}) + worker.server_args = SimpleNamespace( + residency_mode=lambda _name: COMPONENT_OFFLOAD + ) worker._auto_residency_warmup_records = [object()] - worker._auto_residency_applied = [ - AppliedResidencyChange("text_encoder", COMPONENT_OFFLOAD), - AppliedResidencyChange("vae", COMPONENT_OFFLOAD), - ] - worker._auto_residency_repeated_components = set() - worker._auto_residency_round_sizes = [2] - worker._auto_residency_last_applied_plan = AutoResidencyPlan() + worker._auto_residency_applied = [] + worker._auto_residency_round_sizes = [1] report = RankResidencyReport( rank=0, budget_bytes=10 * GIB_BYTES, estimated_peak_bytes=1, estimated_request_duration_ns=10_000_000_000, - measured_request_duration_ns=20_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(side_effect=lambda value: [value]) - worker._rollback_everywhere = mock.Mock() - self.assertTrue(worker._latest_auto_residency_round_supports_short_validation()) + candidate = ResidencyTarget( + component_name="vae", + residency_mode=COMPONENT_OFFLOAD, + target_resident_weight_bytes=1, + h2d_bytes_per_request=1, + ) + worker._auto_residency_last_applied_plan = AutoResidencyPlan( + changes=[candidate] + ) with ( mock.patch.object( @@ -370,122 +167,25 @@ def test_worker_keeps_resident_only_round_despite_noisy_duration(self): "resolve_measured_default_workload", return_value=SimpleNamespace(), ), + mock.patch.object( + gpu_worker_module, + "plan_auto_residency", + ) 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) - worker._rollback_everywhere.assert_not_called() - self.assertEqual(worker._auto_residency_round_sizes, []) + 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_validation_rolls_back_when_selected_placement_spends_the_reserve(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_applied = [] - worker._auto_residency_round_sizes = [1] - worker._auto_residency_last_applied_plan = AutoResidencyPlan( - resource_budget_bytes={"gpu:rank0:denoise": 8 * GIB_BYTES}, - resource_delta_bytes={"gpu:rank0:denoise": 7 * GIB_BYTES}, - ) - worker._build_auto_residency_report = mock.Mock( - return_value=RankResidencyReport( - rank=0, - budget_bytes=30 * GIB_BYTES, - estimated_peak_bytes=28 * GIB_BYTES, - target_workload_measured=True, - ) - ) - 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) - - 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") as plan, - ): - response = worker.apply_auto_residency(validate_only=True) - - self.assertIs(response, expected) - worker._rollback_everywhere.assert_called_once_with( - cause="VRAM reserve exceeded by 1.0 GiB", - already_failed=False, - latest_round_only=True, - ) - plan.assert_not_called() - - 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, - budget_bytes=10 * GIB_BYTES, - 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(side_effect=lambda value: [value]) - candidate = ResidencyTarget( - component_name="vae", - residency_mode=COMPONENT_OFFLOAD, - target_resident_weight_bytes=1, - h2d_bytes_per_request=1, - ) - worker._auto_residency_last_applied_plan = AutoResidencyPlan( - changes=[candidate] - ) - - 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", - ) 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_resident_promotions_use_short_validation(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] @@ -579,88 +279,6 @@ async def forward(req): step_limit=1, ) - def test_unrecoverable_warmup_oom_aborts_startup(self): - async def forward(_req): - return OutputBatch( - output={ - "status": PLACEMENT_STATUS_SKIPPED, - "recovering_from_oom": True, - } - ) - - server_args = SimpleNamespace( - performance_mode="auto", - warmup_resolutions=None, - ) - with ( - mock.patch.object( - server_warmup, "auto_residency_skip_reason", return_value=None - ), - self.assertRaisesRegex(RuntimeError, "found no feasible placement"), - ): - asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) - - def test_non_resident_adjustment_keeps_timing_validation_steps(self): - async def forward(req): - status = ( - PLACEMENT_STATUS_ADJUSTED - if req.action == "apply" - else PLACEMENT_STATUS_VALIDATED - ) - return OutputBatch(output={"status": status}) - - 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)) - - rewarm.assert_awaited_once_with( - server_args, - forward, - fail_open=False, - rewarm=True, - ) - - def test_regressed_validation_rolls_back_and_rewarms_original_placement(self): - responses = iter( - [ - OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), - OutputBatch( - error="request duration regressed", - output={"status": PLACEMENT_STATUS_ROLLED_BACK}, - ), - ] - ) - actions = [] - - async def forward(req): - 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", "validate"]) - self.assertEqual(rewarm.await_count, 2) - def test_failed_first_calibration_rolls_back_and_rewarms(self): actions = [] @@ -714,35 +332,6 @@ async def forward(req): ): asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) - def test_successful_calibration_does_not_start_another_apply_round(self): - actions = [] - - async def forward(req): - self.assertIsInstance(req, AutoResidencyReq) - actions.append(req.action) - status = ( - PLACEMENT_STATUS_ADJUSTED - if req.action == "apply" - else PLACEMENT_STATUS_VALIDATED - ) - return OutputBatch(output={"status": status}) - - 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", "validate"]) - self.assertEqual(rewarm.await_count, 1) - def test_failed_later_round_keeps_earlier_calibrated_promotions(self): first_round = AppliedResidencyChange("transformer", COMPONENT_OFFLOAD) latest_round = [ @@ -792,3 +381,414 @@ def test_failed_apply_round_does_not_rollback_prior_round(self): rollback.assert_not_called() self.assertEqual(worker._auto_residency_applied, [first_round]) self.assertEqual(worker._auto_residency_round_sizes, [1]) + + def test_auto_residency_uses_first_oom_without_degrading_probe(self): + req = SimpleNamespace() + forward = mock.AsyncMock(return_value=OutputBatch(error="CUDA out of memory")) + + with ( + mock.patch.object( + server_warmup, "auto_residency_skip_reason", return_value=None + ), + mock.patch.object( + server_warmup, + "should_include_warmup_image", + return_value=False, + ), + mock.patch.object( + server_warmup, "build_client_warmup_reqs", return_value=[req] + ), + mock.patch.object(server_warmup, "_degrade_after_oom") as degrade, + ): + asyncio.run( + server_warmup.run_async_client_warmup( + SimpleNamespace(), forward, fail_open=True + ) + ) + + forward.assert_awaited_once_with(req) + degrade.assert_not_called() + + def test_non_auto_warmup_keeps_oom_probe_degradation(self): + req = SimpleNamespace() + lighter = SimpleNamespace() + forward = mock.AsyncMock( + side_effect=[ + OutputBatch(error="CUDA out of memory"), + OutputBatch(), + ] + ) + + with ( + mock.patch.object( + server_warmup, + "auto_residency_skip_reason", + return_value="not auto", + ), + mock.patch.object( + server_warmup, + "should_include_warmup_image", + return_value=False, + ), + mock.patch.object( + server_warmup, "build_client_warmup_reqs", return_value=[req] + ), + mock.patch.object( + server_warmup, "_degrade_after_oom", return_value=lighter + ) as degrade, + ): + asyncio.run( + server_warmup.run_async_client_warmup(SimpleNamespace(), forward) + ) + + self.assertEqual( + [call.args[0] for call in forward.await_args_list], + [req, lighter], + ) + degrade.assert_called_once_with(mock.ANY, req) + + def test_non_resident_adjustment_keeps_timing_validation_steps(self): + async def forward(req): + status = ( + PLACEMENT_STATUS_ADJUSTED + if req.action == "apply" + else PLACEMENT_STATUS_VALIDATED + ) + return OutputBatch(output={"status": status}) + + 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)) + + rewarm.assert_awaited_once_with( + server_args, + forward, + fail_open=False, + rewarm=True, + ) + + def test_regressed_validation_rolls_back_and_rewarms_original_placement(self): + responses = iter( + [ + OutputBatch(output={"status": PLACEMENT_STATUS_ADJUSTED}), + OutputBatch( + error="request duration regressed", + output={"status": PLACEMENT_STATUS_ROLLED_BACK}, + ), + ] + ) + actions = [] + + async def forward(req): + 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", "validate"]) + self.assertEqual(rewarm.await_count, 2) + + def test_residency_hint_uses_effective_device_budget(self): + worker = GPUWorker.__new__(GPUWorker) + worker._auto_residency_budget_bytes = mock.Mock(return_value=12 * GIB_BYTES) + worker.get_can_stay_resident_components = mock.Mock(return_value=[]) + output = SimpleNamespace(metrics=None) + snapshot = SimpleNamespace( + peak_reserved_mb=11 * 1024, + peak_allocated_mb=10 * 1024, + ) + + with mock.patch.object( + gpu_worker_module, "capture_memory_snapshot", return_value=snapshot + ): + worker.do_mem_analysis(output) + + worker.get_can_stay_resident_components.assert_called_once_with(1.0) + self.assertEqual(output.peak_memory_mb, 11 * 1024) + + def test_successful_calibration_does_not_start_another_apply_round(self): + actions = [] + + async def forward(req): + self.assertIsInstance(req, AutoResidencyReq) + actions.append(req.action) + status = ( + PLACEMENT_STATUS_ADJUSTED + if req.action == "apply" + else PLACEMENT_STATUS_VALIDATED + ) + return OutputBatch(output={"status": status}) + + 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", "validate"]) + self.assertEqual(rewarm.await_count, 1) + + def test_unrecoverable_warmup_oom_aborts_startup(self): + async def forward(_req): + return OutputBatch( + output={ + "status": PLACEMENT_STATUS_SKIPPED, + "recovering_from_oom": True, + } + ) + + server_args = SimpleNamespace( + performance_mode="auto", + warmup_resolutions=None, + ) + with ( + mock.patch.object( + server_warmup, "auto_residency_skip_reason", return_value=None + ), + self.assertRaisesRegex(RuntimeError, "found no feasible placement"), + ): + asyncio.run(server_warmup.maybe_apply_auto_residency(server_args, forward)) + + def test_validation_rolls_back_when_selected_placement_spends_the_reserve(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_applied = [] + worker._auto_residency_round_sizes = [1] + worker._auto_residency_last_applied_plan = AutoResidencyPlan( + resource_budget_bytes={"gpu:rank0:denoise": 8 * GIB_BYTES}, + resource_delta_bytes={"gpu:rank0:denoise": 7 * GIB_BYTES}, + ) + worker._build_auto_residency_report = mock.Mock( + return_value=RankResidencyReport( + rank=0, + budget_bytes=30 * GIB_BYTES, + estimated_peak_bytes=28 * GIB_BYTES, + target_workload_measured=True, + ) + ) + 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) + + 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") as plan, + ): + response = worker.apply_auto_residency(validate_only=True) + + self.assertIs(response, expected) + worker._rollback_everywhere.assert_called_once_with( + cause="VRAM reserve exceeded by 1.0 GiB", + already_failed=False, + latest_round_only=True, + ) + plan.assert_not_called() + + def test_worker_keeps_resident_only_round_despite_noisy_duration(self): + worker = GPUWorker.__new__(GPUWorker) + worker.rank = 0 + worker.is_output_rank = False + worker.server_args = SimpleNamespace(residency_mode=lambda _name: RESIDENT) + worker._auto_residency_warmup_records = [object()] + worker._auto_residency_applied = [ + AppliedResidencyChange("text_encoder", COMPONENT_OFFLOAD), + AppliedResidencyChange("vae", COMPONENT_OFFLOAD), + ] + worker._auto_residency_repeated_components = set() + worker._auto_residency_round_sizes = [2] + worker._auto_residency_last_applied_plan = AutoResidencyPlan() + report = RankResidencyReport( + rank=0, + budget_bytes=10 * GIB_BYTES, + 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(side_effect=lambda value: [value]) + worker._rollback_everywhere = mock.Mock() + self.assertTrue(worker._latest_auto_residency_round_supports_short_validation()) + + 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(validate_only=True) + + self.assertEqual(response.output["status"], PLACEMENT_STATUS_VALIDATED) + worker._rollback_everywhere.assert_not_called() + self.assertEqual(worker._auto_residency_round_sizes, []) + + def test_sync_startup_uses_shared_auto_residency_sequence(self): + server_args = SimpleNamespace() + forward = mock.Mock() + + with ( + mock.patch.object( + server_warmup, "maybe_apply_pre_warmup_auto_residency" + ) as apply_static, + mock.patch.object(server_warmup, "run_async_client_warmup") as warmup, + mock.patch.object(server_warmup, "maybe_apply_auto_residency") as refine, + mock.patch.object( + server_warmup, + "should_run_synthetic_server_warmup", + return_value=True, + ), + ): + server_warmup.run_sync_startup_warmup(server_args, forward) + + apply_static.assert_awaited_once() + warmup.assert_awaited_once() + refine.assert_awaited_once() + self.assertIs(apply_static.await_args.args[0], server_args) + self.assertIs(warmup.await_args.args[0], server_args) + self.assertIs(refine.await_args.args[0], server_args) + self.assertTrue(warmup.await_args.kwargs["fail_open"]) + self.assertIs(apply_static.await_args.args[1], warmup.await_args.args[1]) + self.assertIs(apply_static.await_args.args[1], refine.await_args.args[1]) + + def test_sync_startup_static_planning_does_not_require_warmup(self): + with ( + mock.patch.object( + server_warmup, "maybe_apply_pre_warmup_auto_residency" + ) as apply_static, + mock.patch.object( + server_warmup, + "should_run_synthetic_server_warmup", + return_value=False, + ), + mock.patch.object(server_warmup, "run_async_client_warmup") as warmup, + mock.patch.object(server_warmup, "maybe_apply_auto_residency") as refine, + ): + server_warmup.run_sync_startup_warmup(SimpleNamespace(), mock.Mock()) + + apply_static.assert_awaited_once() + warmup.assert_not_awaited() + refine.assert_not_awaited() + + def test_pre_warmup_planner_uses_static_action(self): + actions = [] + + async def forward(req): + actions.append(req.action) + return OutputBatch(output={"status": PLACEMENT_STATUS_SKIPPED}) + + with mock.patch.object( + server_warmup, + "should_apply_pre_warmup_auto_residency", + return_value=True, + ): + asyncio.run( + server_warmup.maybe_apply_pre_warmup_auto_residency( + SimpleNamespace(), forward + ) + ) + + self.assertEqual(actions, ["apply_static"]) + + def test_pre_warmup_worker_commits_seed_without_pending_round(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 = [] + worker._auto_residency_applied = [] + worker._auto_residency_round_sizes = [] + worker._auto_residency_last_applied_plan = None + report = RankResidencyReport( + rank=0, + budget_bytes=10 * GIB_BYTES, + estimated_peak_bytes=5 * GIB_BYTES, + ) + worker._build_pre_warmup_auto_residency_report = mock.Mock(return_value=report) + worker._auto_residency_all_gather = mock.Mock(side_effect=lambda value: [value]) + worker._auto_residency_modules = mock.Mock(return_value={}) + worker._invalidate_component_strategies = mock.Mock() + candidate = ResidencyTarget( + component_name="transformer", + residency_mode=RESIDENT, + target_resident_weight_bytes=0, + h2d_bytes_per_request=1, + target_residency_mode=COMPONENT_OFFLOAD, + ) + plan = AutoResidencyPlan(changes=[candidate]) + applied = AppliedResidencyChange("transformer", RESIDENT) + workload = DefaultWorkload(512, 512, 1, 4) + + 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 + ), + mock.patch.object( + gpu_worker_module, "apply_residency_changes", return_value=[applied] + ), + mock.patch.object(gpu_worker_module, "commit_residency_changes") as commit, + ): + response = worker.apply_auto_residency(pre_warmup=True) + + self.assertEqual(response.output["status"], PLACEMENT_STATUS_ADJUSTED) + commit.assert_called_once_with( + applied=[applied], modules={}, server_args=worker.server_args + ) + self.assertEqual(worker._auto_residency_applied, []) + self.assertEqual(worker._auto_residency_round_sizes, []) + self.assertIsNone(worker._auto_residency_last_applied_plan) 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 8a54f1d17034..546d3991936b 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 @@ -16,7 +16,11 @@ import unittest from collections import deque from types import SimpleNamespace -from unittest.mock import ANY, MagicMock, patch +from unittest.mock import ( + ANY, + MagicMock, + patch, +) import torch @@ -346,30 +350,6 @@ def test_diff_generator_runs_explicit_warmup_through_scheduler_client(self): self.assertTrue(req.extra["server_based_warmup"]) self.assertEqual(req.extra["warmup_total"], 1) - def test_diff_generator_runs_server_style_startup_warmup(self): - generator = object.__new__(DiffGenerator) - generator.server_args = server_args = MagicMock() - - with ( - patch( - "sglang.multimodal_gen.runtime.entrypoints.diffusion_generator.should_run_synthetic_server_warmup", - return_value=True, - ), - patch( - "sglang.multimodal_gen.runtime.entrypoints.diffusion_generator.run_sync_startup_warmup" - ) as startup_warmup, - patch( - "sglang.multimodal_gen.runtime.entrypoints.diffusion_generator.run_sync_client_warmup" - ) as request_warmup, - ): - generator._run_client_warmup_if_needed() - - startup_warmup.assert_called_once_with( - server_args, - ANY, - ) - request_warmup.assert_not_called() - def test_server_based_warmup_uses_model_default_negative_prompt(self): server_args = MagicMock() server_args.warmup_steps = 1 @@ -984,6 +964,30 @@ def test_server_based_warmup_keeps_ti2v_image_input(self): self.assertEqual(reqs[0].image_path, ["/tmp/warmup.png"]) + def test_diff_generator_runs_server_style_startup_warmup(self): + generator = object.__new__(DiffGenerator) + generator.server_args = server_args = MagicMock() + + with ( + patch( + "sglang.multimodal_gen.runtime.entrypoints.diffusion_generator.should_run_synthetic_server_warmup", + return_value=True, + ), + patch( + "sglang.multimodal_gen.runtime.entrypoints.diffusion_generator.run_sync_startup_warmup" + ) as startup_warmup, + patch( + "sglang.multimodal_gen.runtime.entrypoints.diffusion_generator.run_sync_client_warmup" + ) as request_warmup, + ): + generator._run_client_warmup_if_needed() + + startup_warmup.assert_called_once_with( + server_args, + ANY, + ) + request_warmup.assert_not_called() + class TestFlux2FinetunedVaeEncodePreprocess(unittest.TestCase): def test_single_frame_custom_vae_encode_input_is_4d(self): 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 4cc9c4687cd4..03267679d051 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_residency.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_residency.py @@ -1,5 +1,8 @@ from types import SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import ( + Mock, + patch, +) import pytest import torch @@ -75,27 +78,6 @@ def test_component_offload_keeps_preferred_component_after_warmup(): strategy.finish_use.assert_not_called() -def test_component_offload_restores_cpu_after_failed_device_transfer(): - strategy = ComponentOffloadStrategy() - module = Mock() - use = ComponentUse("DenoiseStage", "transformer") - error = RuntimeError("CUDA out of memory") - - with patch.object( - component_residency_strategies, - "_module_to_local_device", - side_effect=error, - ): - try: - strategy.prepare_for_use(module, use, ResidencyState()) - except RuntimeError as caught: - assert caught is error - else: - raise AssertionError("expected the simulated device transfer to fail") - - module.to.assert_called_once_with("cpu", non_blocking=False) - - def test_request_tail_uses_dynamic_component_instance(): pipeline = SimpleNamespace( modules={}, @@ -120,30 +102,6 @@ def test_request_tail_uses_dynamic_component_instance(): ) -def test_placement_modules_include_lazy_request_components(): - pipeline_module = torch.nn.Linear(2, 2) - stale_seen_module = torch.nn.Linear(2, 2) - lazy_module = torch.nn.Linear(2, 2) - pipeline = SimpleNamespace( - modules={"transformer": pipeline_module}, - _stage_name_mapping={}, - component_residency_strategies={}, - ) - manager = ComponentResidencyManager( - pipeline, - SimpleNamespace(enable_layerwise_nvtx_marker=False), - ) - manager._modules_seen = { - "transformer": stale_seen_module, - "condition_image_encoder": lazy_module, - } - - assert manager.placement_modules() == { - "transformer": pipeline_module, - "condition_image_encoder": lazy_module, - } - - def test_strategy_cache_replaces_stale_component_instance(): server_args = SimpleNamespace( enable_layerwise_nvtx_marker=False, @@ -250,24 +208,6 @@ def test_reports_components_with_mixed_use_dtypes(): assert manager.components_with_mixed_use_dtypes() == {"vae"} -def test_reports_mixed_use_dtypes_before_first_request(): - encode = _Stage(ComponentUse("encode", "vae", target_dtype=torch.float32)) - decode = _Stage(ComponentUse("decode", "vae", target_dtype=torch.bfloat16)) - pipeline = SimpleNamespace( - modules={}, - _stage_name_mapping={"encode": encode, "decode": decode}, - component_residency_strategies={}, - ) - server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) - manager = ComponentResidencyManager(pipeline, server_args) - manager.refresh_pipeline(pipeline) - - assert manager.components_with_mixed_use_dtypes([encode, decode], server_args) == { - "vae" - } - assert manager._ordered_uses == () - - def test_warmup_records_use_and_transition_peaks(monkeypatch): device_module = SimpleNamespace( is_available=lambda: True, @@ -317,41 +257,6 @@ def test_warmup_records_use_and_transition_peaks(monkeypatch): ) -def test_failed_warmup_phase_survives_cleanup(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) - - use = ComponentUse("denoise", "transformer") - stage = _Stage(use) - pipeline = SimpleNamespace( - modules={"transformer": torch.nn.Linear(2, 2)}, - _stage_name_mapping={"denoise": stage}, - component_residency_strategies={}, - ) - server_args = _server_args() - manager = ComponentResidencyManager(pipeline, server_args) - 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_warmup_prefetch(use) - - manager.capture_failed_warmup_phase() - manager._warmup_phase_peaks = {} - manager._completed_warmup_phase_peaks = {"request:cleanup": WarmupPhasePeak((), 2)} - - peaks = manager.take_warmup_phase_peaks() - assert peaks["0:denoise:prefetch:transformer"] == WarmupPhasePeak( - ("transformer",), 7, prefetched_components=("transformer",) - ) - assert peaks["request:cleanup"] == WarmupPhasePeak((), 2) - - def test_warmup_skips_memory_tracking_for_unsupported_pipeline(monkeypatch): device_module = SimpleNamespace( is_available=lambda: True, @@ -548,46 +453,6 @@ def test_warmup_splits_sequential_component_transition(monkeypatch): ) -def test_warmup_preserves_repeated_phase_layouts(monkeypatch): - device_module = SimpleNamespace( - is_available=lambda: True, - reset_peak_memory_stats=Mock(), - max_memory_allocated=lambda: 7, - ) - 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() - manager = ComponentResidencyManager(pipeline, server_args) - manager.refresh_pipeline(pipeline) - manager.begin_request([stage], SimpleNamespace(is_warmup=True), server_args) - - manager._begin_warmup_phase( - key="0:stage:between", - components=("text_encoder",), - used_components=(), - ) - manager._begin_warmup_phase( - key="0:stage:between", - components=("transformer",), - used_components=(), - ) - manager._record_warmup_phase_peak() - - assert manager._warmup_phase_peaks["0:stage:between"] == WarmupPhasePeak( - ("text_encoder",), 7 - ) - assert manager._warmup_phase_peaks[ - "0:stage:between:occurrence:1" - ] == WarmupPhasePeak(("transformer",), 7) - - def _manager_for_stage(stage, modules): pipeline = SimpleNamespace( modules=modules, @@ -899,3 +764,141 @@ def test_component_is_not_kept_across_another_component_use(): manager.end_use(text_use) strategy.finish_use.assert_called_once_with(module, text_use, manager.state) + + +def test_component_offload_restores_cpu_after_failed_device_transfer(): + strategy = ComponentOffloadStrategy() + module = Mock() + use = ComponentUse("DenoiseStage", "transformer") + error = RuntimeError("CUDA out of memory") + + with patch.object( + component_residency_strategies, + "_module_to_local_device", + side_effect=error, + ): + try: + strategy.prepare_for_use(module, use, ResidencyState()) + except RuntimeError as caught: + assert caught is error + else: + raise AssertionError("expected the simulated device transfer to fail") + + module.to.assert_called_once_with("cpu", non_blocking=False) + + +def test_placement_modules_include_lazy_request_components(): + pipeline_module = torch.nn.Linear(2, 2) + stale_seen_module = torch.nn.Linear(2, 2) + lazy_module = torch.nn.Linear(2, 2) + pipeline = SimpleNamespace( + modules={"transformer": pipeline_module}, + _stage_name_mapping={}, + component_residency_strategies={}, + ) + manager = ComponentResidencyManager( + pipeline, + SimpleNamespace(enable_layerwise_nvtx_marker=False), + ) + manager._modules_seen = { + "transformer": stale_seen_module, + "condition_image_encoder": lazy_module, + } + + assert manager.placement_modules() == { + "transformer": pipeline_module, + "condition_image_encoder": lazy_module, + } + + +def test_reports_mixed_use_dtypes_before_first_request(): + encode = _Stage(ComponentUse("encode", "vae", target_dtype=torch.float32)) + decode = _Stage(ComponentUse("decode", "vae", target_dtype=torch.bfloat16)) + pipeline = SimpleNamespace( + modules={}, + _stage_name_mapping={"encode": encode, "decode": decode}, + component_residency_strategies={}, + ) + server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False) + manager = ComponentResidencyManager(pipeline, server_args) + manager.refresh_pipeline(pipeline) + + assert manager.components_with_mixed_use_dtypes([encode, decode], server_args) == { + "vae" + } + assert manager._ordered_uses == () + + +def test_warmup_preserves_repeated_phase_layouts(monkeypatch): + device_module = SimpleNamespace( + is_available=lambda: True, + reset_peak_memory_stats=Mock(), + max_memory_allocated=lambda: 7, + ) + 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() + manager = ComponentResidencyManager(pipeline, server_args) + manager.refresh_pipeline(pipeline) + manager.begin_request([stage], SimpleNamespace(is_warmup=True), server_args) + + manager._begin_warmup_phase( + key="0:stage:between", + components=("text_encoder",), + used_components=(), + ) + manager._begin_warmup_phase( + key="0:stage:between", + components=("transformer",), + used_components=(), + ) + manager._record_warmup_phase_peak() + + assert manager._warmup_phase_peaks["0:stage:between"] == WarmupPhasePeak( + ("text_encoder",), 7 + ) + assert manager._warmup_phase_peaks[ + "0:stage:between:occurrence:1" + ] == WarmupPhasePeak(("transformer",), 7) + + +def test_failed_warmup_phase_survives_cleanup(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) + + use = ComponentUse("denoise", "transformer") + stage = _Stage(use) + pipeline = SimpleNamespace( + modules={"transformer": torch.nn.Linear(2, 2)}, + _stage_name_mapping={"denoise": stage}, + component_residency_strategies={}, + ) + server_args = _server_args() + manager = ComponentResidencyManager(pipeline, server_args) + 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_warmup_prefetch(use) + + manager.capture_failed_warmup_phase() + manager._warmup_phase_peaks = {} + manager._completed_warmup_phase_peaks = {"request:cleanup": WarmupPhasePeak((), 2)} + + peaks = manager.take_warmup_phase_peaks() + assert peaks["0:denoise:prefetch:transformer"] == WarmupPhasePeak( + ("transformer",), 7, prefetched_components=("transformer",) + ) + assert peaks["request:cleanup"] == WarmupPhasePeak((), 2) 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 7603f5bc2add..46e7ed13684a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -138,37 +138,6 @@ def stream(_stream): return nullcontext() -def test_dtensor_layer_store_remains_rank_local(monkeypatch): - local_weight = torch.arange(4) - replacement = torch.arange(4) + 10 - wraps = [] - - class FakeDTensor: - def __init__(self, local_tensor): - self.local_tensor = local_tensor - self.device_mesh = "mesh" - self.placements = ("shard",) - - def to_local(self): - return self.local_tensor - - @classmethod - def from_local(cls, local_tensor, device_mesh, placements, **kwargs): - wraps.append((local_tensor, device_mesh, placements, kwargs)) - return cls(local_tensor) - - monkeypatch.setattr(layerwise_offload_mod, "DTensor", FakeDTensor) - target = FakeDTensor(local_weight) - - assert LayerwiseOffloadManager._to_local_tensor(target) is local_weight - wrapped = LayerwiseOffloadManager._wrap_for_target( - object.__new__(LayerwiseOffloadManager), target, replacement - ) - - assert wrapped.local_tensor is replacement - assert wraps == [(replacement, "mesh", ("shard",), {"run_check": False})] - - class _DummyBlock(torch.nn.Module): def __init__(self) -> None: super().__init__() @@ -723,23 +692,6 @@ def test_auto_layerwise_skips_unsupported_component(monkeypatch): assert server_args.residency_mode("text_encoder") == COMPONENT_OFFLOAD -def test_initial_auto_resident_seed_skips_configured_layerwise_setup(): - module = _NestedEncoderDummyModel() - server_args = _server_args( - layerwise_offload_components=["text_encoder"], - _auto_residency_modes={"text_encoder": RESIDENT}, - ) - - configured = configure_layerwise_offload_modules( - {"text_encoder": module}, - server_args, - component_names=["text_encoder"], - ) - - assert configured == [] - assert not is_layerwise_offloaded_module(module) - - def test_canonical_selector_does_not_make_auto_layerwise_selection_strict( monkeypatch, ): @@ -1125,44 +1077,6 @@ def test_configure_resolves_residency_policy(monkeypatch): ) -def test_explicit_residency_policy_is_scoped_to_the_selected_component(): - args = _server_args( - _explicit_arg_names={"dit_layerwise_residency_policy"}, - dit_layerwise_residency_policy=RESIDENCY_POLICY_STRIDED, - layerwise_residency_policy={ - "text_encoder": RESIDENCY_POLICY_STRIDED, - }, - ) - - assert args.is_layerwise_residency_policy_explicit("transformer", dit_group=True) - assert args.is_layerwise_residency_policy_explicit("text_encoder", dit_group=False) - assert not args.is_layerwise_residency_policy_explicit("vae", dit_group=False) - - -def test_residency_layout_switches_policy_without_rebuilding_host_stores(monkeypatch): - _patch_fake_device(monkeypatch) - manager = _resident_manager( - _MultiBlockModel(8), - num_layers=8, - prefetch_size=1, - resident_layers=2, - residency_policy=RESIDENCY_POLICY_LEADING, - ) - host_stores = manager._consolidated_cpu_weights - - previous = manager.set_residency_layout(3, RESIDENCY_POLICY_STRIDED) - - assert previous == (2, RESIDENCY_POLICY_LEADING) - assert manager.resident_layers == 3 - assert manager.residency_policy == RESIDENCY_POLICY_STRIDED - assert manager._streamed_order == compute_streamed_layers( - num_layers=8, - resident_layers=3, - policy=RESIDENCY_POLICY_STRIDED, - ) - assert manager._consolidated_cpu_weights is host_stores - - def test_configure_logs_component_start_and_completion(monkeypatch): _patch_fake_device(monkeypatch) logs = [] 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 c511ebd522c8..5e9a47b07351 100644 --- a/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py +++ b/python/sglang/multimodal_gen/test/unit/test_performance_metrics.py @@ -230,64 +230,6 @@ def test_server_warmup_does_not_treat_allocator_cache_as_untracked_live_memory() assert "request:untracked" not in record.phase_peak_allocated_bytes -def test_server_warmup_normalizes_profile_stage_names_for_residency_timing(): - worker = GPUWorker.__new__(GPUWorker) - worker._auto_residency_warmup_records = [] - stage = Mock() - stage._active_profile_stage_name.return_value = "RefinerStage" - stage._component_stage_name.return_value = "refiner" - worker.pipeline = SimpleNamespace(stages=[stage]) - residency_manager = Mock() - residency_manager.take_warmup_phase_peaks.return_value = { - "0:refiner:use:transformer_2": WarmupPhasePeak( - ("transformer_2",), - 8, - used_components=("transformer_2",), - ) - } - residency_manager.current_device_components.return_value = ("transformer_2",) - device_module = Mock() - device_module.max_memory_allocated.return_value = 8 - device_module.max_memory_reserved.return_value = 8 - req = SimpleNamespace( - metrics=SimpleNamespace( - total_duration_ms=12.0, - stages={"RefinerStage": 10.0}, - steps=[], - steps_by_stage={}, - stage_iterations={}, - ) - ) - - 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, - workload=(64, 64, 1, 1), - baseline_allocated_bytes=3, - baseline_reserved_bytes=3, - succeeded=True, - ) - - record = worker._auto_residency_warmup_records[0] - assert record.stage_duration_ms == {"refiner": 10.0} - _, stage_duration_ns, component_stages = ( - gpu_worker_module.estimate_default_workload_timing( - records=[record], - target_units=record.workload_units(), - target_num_inference_steps=1, - ) - ) - assert stage_duration_ns == {"refiner": 10_000_000} - assert component_stages == {"transformer_2": ("refiner",)} - - def test_loaded_prequantized_checkpoint_can_use_auto_residency(): worker = GPUWorker.__new__(GPUWorker) worker.rank = 0 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 7793a4aad53e..63a5a24f3e81 100644 --- a/python/sglang/multimodal_gen/test/unit/test_placement_budget.py +++ b/python/sglang/multimodal_gen/test/unit/test_placement_budget.py @@ -117,53 +117,6 @@ def test_joint_plan_beats_independent_greedy_allocations(): assert plan.estimated_latency_savings == 109 -def test_one_selection_vector_obeys_load_runtime_and_hostpin_budgets(): - plan = optimize_placement( - [ - PlacementOption( - group_key="text_encoder", - option_key="text-encoder-resident", - resource_delta_bytes={"gpu:load": 8, "gpu:runtime": 2}, - estimated_latency_savings=60, - ), - PlacementOption( - group_key="text_encoder", - option_key="text-encoder-pinned", - resource_delta_bytes={"gpu:load": 1, "hostpin": 6}, - estimated_latency_savings=50, - ), - PlacementOption( - group_key="transformer", - option_key="transformer-resident", - resource_delta_bytes={"gpu:load": 2, "gpu:runtime": 8}, - estimated_latency_savings=100, - ), - PlacementOption( - group_key="transformer", - option_key="transformer-pinned", - resource_delta_bytes={"gpu:runtime": 1, "hostpin": 6}, - estimated_latency_savings=20, - ), - ], - resource_budget_bytes={ - "gpu:load": 9, - "gpu:runtime": 8, - "hostpin": 6, - }, - require_selection_from_every_group=True, - ) - - assert [option.option_key for option in plan.selections] == [ - "text-encoder-pinned", - "transformer-resident", - ] - assert plan.resource_delta_bytes == { - "gpu:load": 3, - "gpu:runtime": 8, - "hostpin": 6, - } - - def test_selects_at_most_one_option_per_group(): plan = optimize_placement( [ @@ -313,32 +266,6 @@ def test_latency_tolerance_is_global_not_per_option(): assert plan.estimated_latency_savings == 60 -def test_lexicographic_preference_prunes_irrelevant_later_dimension(): - plan = optimize_placement( - [ - PlacementOption( - group_key="dit", - option_key="partial-pin", - resource_delta_bytes={"gpu": 1}, - estimated_latency_savings=90, - preference_cost=(0, -90, 5), - ), - PlacementOption( - group_key="dit", - option_key="full-pin", - resource_delta_bytes={"gpu": 1}, - estimated_latency_savings=100, - preference_cost=(0, -100, 10), - ), - ], - resource_budget_bytes={"gpu": 1}, - estimated_latency_tolerance=20, - require_selection_from_every_group=True, - ) - - assert [option.option_key for option in plan.selections] == ["full-pin"] - - def test_one_static_placement_satisfies_all_observed_phase_constraints(): plan = optimize_placement( [ @@ -681,6 +608,79 @@ def test_branch_and_bound_matches_exhaustive_multiresource_search(): assert plan.preference_cost == expected_cost +def test_one_selection_vector_obeys_load_runtime_and_hostpin_budgets(): + plan = optimize_placement( + [ + PlacementOption( + group_key="text_encoder", + option_key="text-encoder-resident", + resource_delta_bytes={"gpu:load": 8, "gpu:runtime": 2}, + estimated_latency_savings=60, + ), + PlacementOption( + group_key="text_encoder", + option_key="text-encoder-pinned", + resource_delta_bytes={"gpu:load": 1, "hostpin": 6}, + estimated_latency_savings=50, + ), + PlacementOption( + group_key="transformer", + option_key="transformer-resident", + resource_delta_bytes={"gpu:load": 2, "gpu:runtime": 8}, + estimated_latency_savings=100, + ), + PlacementOption( + group_key="transformer", + option_key="transformer-pinned", + resource_delta_bytes={"gpu:runtime": 1, "hostpin": 6}, + estimated_latency_savings=20, + ), + ], + resource_budget_bytes={ + "gpu:load": 9, + "gpu:runtime": 8, + "hostpin": 6, + }, + require_selection_from_every_group=True, + ) + + assert [option.option_key for option in plan.selections] == [ + "text-encoder-pinned", + "transformer-resident", + ] + assert plan.resource_delta_bytes == { + "gpu:load": 3, + "gpu:runtime": 8, + "hostpin": 6, + } + + +def test_lexicographic_preference_prunes_irrelevant_later_dimension(): + plan = optimize_placement( + [ + PlacementOption( + group_key="dit", + option_key="partial-pin", + resource_delta_bytes={"gpu": 1}, + estimated_latency_savings=90, + preference_cost=(0, -90, 5), + ), + PlacementOption( + group_key="dit", + option_key="full-pin", + resource_delta_bytes={"gpu": 1}, + estimated_latency_savings=100, + preference_cost=(0, -100, 10), + ), + ], + resource_budget_bytes={"gpu": 1}, + estimated_latency_tolerance=20, + require_selection_from_every_group=True, + ) + + assert [option.option_key for option in plan.selections] == ["full-pin"] + + def test_identical_resource_columns_keep_only_the_strictest_budget(): options = [ PlacementOption( From cf99a715648aeb8c83f475633f656eb3eea8c638 Mon Sep 17 00:00:00 2001 From: Mick Qian Date: Mon, 7 Sep 2026 22:06:21 +0800 Subject: [PATCH 68/70] fix: keep the threshold residency defaults under the static planner The pre-load planner only promotes DiTs, so starting from main's memory thresholds loses nothing and avoids serving below them when a promotion is declined (see the apply layer's fix for the measured regressions). Co-Authored-By: Claude Fable 5.1 --- .../multimodal_gen/runtime/server_args/auto_tune.py | 9 --------- 1 file changed, 9 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 9a09cfb77bae..806151f68932 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -225,12 +225,6 @@ 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_auto_residency_planner(): - # Keep the load-safe placement until the post-load static planner - # replaces model/card thresholds with component sizes and the - # default workload. A server warmup may refine that estimate. - return - # Explicit placement is component-scoped; unmatched components still # receive automatic defaults. @@ -646,9 +640,6 @@ 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_auto_residency_planner(): - return components - deployment_config = self._deployment_config() threshold_gb = self._resolve_keep_resident_min_available_gb(deployment_config) if threshold_gb is None: From a330811a770c8b69df43e322dbf356092b8b8645 Mon Sep 17 00:00:00 2001 From: Mick Qian Date: Mon, 7 Sep 2026 22:23:11 +0800 Subject: [PATCH 69/70] tests: restore the threshold placement tests after dropping the load-safe start The 21 'defers residency' variants asserted that auto mode starts every component offloaded so the static planner can place it later. Placement now starts from the memory thresholds main applies and the planner only promotes, so the apply layer's threshold tests are the right expectations again; the FSDP-versus-planner test stays u1's own. Co-Authored-By: Claude Fable 5.1 --- .../test/unit/test_server_args.py | 772 +++++++++--------- 1 file changed, 385 insertions(+), 387 deletions(-) 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 6892937bccda..4c96bcc04867 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -1509,7 +1509,7 @@ def test_cpu_offload_component_selector_keeps_unmatched_auto_defaults(self): ) self.assertEqual(args.residency_mode("vae"), COMPONENT_OFFLOAD) - self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD) + self.assertEqual(args.residency_mode("transformer"), RESIDENT) self.assertEqual(args.residency_mode("text_encoder"), LAYERWISE_OFFLOAD) def test_explicit_dit_layerwise_component_wins_over_auto_residency(self): @@ -1537,16 +1537,6 @@ def test_server_warmup_keeps_high_memory_residency_defaults(self): self.assertEqual(args.residency_mode("transformer"), RESIDENT) self.assertEqual(args.residency_mode("vae"), RESIDENT) - def test_no_warmup_keeps_load_safe_static_planner_seed(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"), COMPONENT_OFFLOAD) - self.assertEqual(args.residency_mode("vae"), LAYERWISE_OFFLOAD) - def test_sana_wm_and_helios_declare_model_residency_hints(self): sana_wm_deployment = SanaWMPipelineConfig().get_model_deployment_config() self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60) @@ -1672,25 +1662,6 @@ def test_pipeline_configs_declare_auto_tune_hints(self): self.assertEqual(qwen_deployment.keep_resident_components, ("vae",)) self.assertIsNone(qwen_deployment.keep_resident_min_available_gb) - def test_longlive_planner_candidates_do_not_use_a_coarse_vram_threshold(self): - high_memory_args = self._from_dict_with_pipeline_config( - LongLive2T2VConfig(), - memory_gb=80, - kwargs={"performance_mode": "auto"}, - ) - high_memory_offload = high_memory_args.layerwise_offload_components or [] - self.assertIn("text_encoder", high_memory_offload) - self.assertIn("vae", high_memory_offload) - - constrained_args = self._from_dict_with_pipeline_config( - LongLive2T2VConfig(), - memory_gb=50, - 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_sana_wm_residency_scales_with_available_memory(self): high_memory_args = self._from_dict_with_pipeline_config( SanaWMPipelineConfig(), @@ -1710,26 +1681,6 @@ def test_sana_wm_residency_scales_with_available_memory(self): self.assertIn("text_encoder", constrained_offload) self.assertIn("vae", constrained_offload) - def test_longcat_ar_planner_candidates_do_not_use_a_coarse_vram_threshold(self): - high_memory_args = self._from_dict_with_pipeline_config( - LongCatImagePipelineConfig(), - memory_gb=80, - kwargs={"performance_mode": "auto"}, - ) - self.assertIn( - "text_encoder", high_memory_args.layerwise_offload_components or [] - ) - self.assertFalse(high_memory_args.text_encoder_cpu_offload) - - constrained_args = self._from_dict_with_pipeline_config( - LongCatImagePipelineConfig(), - memory_gb=60, - kwargs={"performance_mode": "auto"}, - ) - self.assertIn( - "text_encoder", constrained_args.layerwise_offload_components or [] - ) - def test_auto_multi_gpu_sana_wm_prefers_fsdp_and_cfg_parallel(self): args = self._from_dict_with_pipeline_config( SanaWMPipelineConfig(), @@ -1833,21 +1784,6 @@ def test_manual_mode_preserves_unset_performance_args(self): self.assertIsNone(args.image_encoder_cpu_offload) self.assertFalse(args.enable_cfg_parallel) - def test_default_auto_exposes_all_implicit_image_placement_candidates(self): - args = self._from_dict_with_pipeline_config( - QwenImagePipelineConfig(), - kwargs={"model_path": "Qwen/Qwen-Image"}, - ) - - self.assertEqual(args.performance_mode, "auto") - self.assertFalse(args.use_fsdp_inference) - self.assertTrue(args.dit_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - self.assertFalse(args.vae_cpu_offload) - def test_auto_image_offloads_aux_below_resident_threshold(self): # 40gb < image threshold (45gb): aux incl. vae still offloaded to save vram args = self._from_dict_with_pipeline_config( @@ -1863,23 +1799,6 @@ def test_auto_image_offloads_aux_below_resident_threshold(self): ["text_encoder", "image_encoder", "vae"], ) - def test_auto_zimage_defers_residency_until_weights_are_known(self): - args = self._from_dict_with_pipeline_config( - ZImagePipelineConfig(), - memory_gb=32, - available_memory_gb=31, - kwargs={ - "model_path": "Tongyi-MAI/Z-Image-Turbo", - "performance_mode": "auto", - }, - ) - - self.assertTrue(args.dit_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - def test_auto_lingbot_keeps_dit_resident_on_h100(self): args = self._from_dict_with_pipeline_config( LingBotWorldCausalDMDConfig(), @@ -1910,45 +1829,6 @@ def test_auto_image_preserves_explicit_dit_cpu_offload(self): self.assertTrue(args.dit_cpu_offload) - def test_auto_ltx_original_defers_component_residency( - self, - ): - args = self._from_dict_with_pipeline_config( - LTX2PipelineConfig(), - available_memory_gb=76, - kwargs={ - "model_path": "Lightricks/LTX-2.3", - "pipeline_class_name": "LTX2TwoStageHQPipeline", - "performance_mode": "auto", - }, - ) - - self.assertEqual(args.ltx2_two_stage_device_mode, "original") - self.assertTrue(args.dit_cpu_offload) - self.assertTrue(args.layerwise_offload_components) - self.assertFalse(args.text_encoder_cpu_offload) - self.assertFalse(args.image_encoder_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - - def test_auto_wan_defers_single_dit_residency(self): - args = self._from_dict_with_pipeline_config( - WanT2V480PConfig(), - kwargs={"performance_mode": "auto"}, - ) - - self.assertTrue(args.layerwise_offload_components) - self.assertFalse(args.use_fsdp_inference) - self.assertTrue(args.dit_cpu_offload) - self.assertFalse(args.text_encoder_cpu_offload) - self.assertFalse(args.image_encoder_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - def test_auto_wan_offloads_single_dit_below_resident_threshold(self): args = self._from_dict_with_pipeline_config( WanT2V480PConfig(), @@ -1988,29 +1868,6 @@ def test_auto_wan2_2_a14b_layerwise_offload_adds_dit(self): ["dit", "text_encoder", "image_encoder", "vae"], ) - def test_auto_wan2_1_14b_defers_dit_residency(self): - for pipeline_config, model_path in ( - (WanT2V720PConfig(), "Wan-AI/Wan2.1-T2V-14B-Diffusers"), - (WanI2V480PConfig(), "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"), - (WanI2V720PConfig(), "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"), - ): - with self.subTest(pipeline_config=pipeline_config.__class__.__name__): - args = self._from_dict_with_pipeline_config( - pipeline_config, - kwargs={ - "model_path": model_path, - "performance_mode": "auto", - }, - ) - - self.assertTrue(args.layerwise_offload_components) - self.assertTrue(args.dit_cpu_offload) - self.assertEqual(args.dit_offload_prefetch_size, 0.0) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - def test_auto_wan2_1_14b_offloads_dit_below_resident_threshold(self): args = self._from_dict_with_pipeline_config( WanI2V720PConfig(), @@ -2104,39 +1961,6 @@ def test_auto_mova_layerwise_offload_adds_dit_below_memory_threshold(self): ["dit", "text_encoder", "image_encoder", "vae"], ) - def test_auto_mova_keeps_dit_layerwise_candidate_at_high_memory(self): - args = self._from_dict_with_pipeline_config( - MOVAPipelineConfig(), - memory_gb=140, - kwargs={ - "model_path": "OpenMOSS-Team/MOVA-360p", - "performance_mode": "auto", - }, - ) - - self.assertTrue(args.dit_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["dit", "text_encoder", "image_encoder", "vae"], - ) - - def test_auto_cosmos3_defers_high_memory_residency(self): - args = self._from_dict_with_pipeline_config( - Cosmos3Config(model_path="nvidia/Cosmos3-Nano"), - available_memory_gb=95, - kwargs={ - "model_path": "nvidia/Cosmos3-Nano", - "performance_mode": "auto", - }, - ) - - self.assertTrue(args.dit_cpu_offload) - self.assertFalse(args.vae_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - def test_auto_cosmos3_offloads_dit_below_resident_threshold(self): args = self._from_dict_with_pipeline_config( Cosmos3Config(model_path="nvidia/Cosmos3-Nano"), @@ -2150,19 +1974,6 @@ def test_auto_cosmos3_offloads_dit_below_resident_threshold(self): self.assertTrue(args.dit_cpu_offload) self.assertFalse(args.vae_cpu_offload) - def test_auto_cosmos3_super_defers_high_memory_residency(self): - args = self._from_dict_with_pipeline_config( - Cosmos3Config(model_path="nvidia/Cosmos3-Super"), - available_memory_gb=139, - kwargs={ - "model_path": "nvidia/Cosmos3-Super", - "performance_mode": "auto", - }, - ) - - self.assertTrue(args.dit_cpu_offload) - self.assertFalse(args.vae_cpu_offload) - def test_auto_cosmos3_super_offloads_dit_below_resident_threshold(self): args = self._from_dict_with_pipeline_config( Cosmos3Config(model_path="nvidia/Cosmos3-Super"), @@ -2190,22 +2001,6 @@ def test_memory_sana_wm_layerwise_offload_adds_dit(self): ["dit", "text_encoder", "image_encoder", "vae"], ) - def test_auto_fastwan_defers_dit_residency(self): - args = self._from_dict_with_pipeline_config( - FastWan2_2_TI2V_5B_Config(), - available_memory_gb=72, - kwargs={ - "model_path": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", - "performance_mode": "auto", - }, - ) - - self.assertTrue(args.dit_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - def test_auto_fastwan_offloads_dit_below_resident_threshold(self): args = self._from_dict_with_pipeline_config( FastWan2_2_TI2V_5B_Config(), @@ -2218,19 +2013,6 @@ def test_auto_fastwan_offloads_dit_below_resident_threshold(self): self.assertTrue(args.dit_cpu_offload) - def test_auto_fast_hunyuan_defers_residency(self): - args = self._from_dict_with_pipeline_config( - FastHunyuanConfig(), - available_memory_gb=72, - kwargs={ - "model_path": "FastVideo/FastHunyuan-diffusers", - "performance_mode": "auto", - }, - ) - - self.assertTrue(args.dit_cpu_offload) - self.assertFalse(args.vae_cpu_offload) - def test_auto_fast_hunyuan_offloads_dit_below_resident_threshold(self): args = self._from_dict_with_pipeline_config( FastHunyuanConfig(), @@ -2243,21 +2025,6 @@ def test_auto_fast_hunyuan_offloads_dit_below_resident_threshold(self): self.assertTrue(args.dit_cpu_offload) - def test_auto_turbo_wan_defers_dit_residency(self): - args = self._from_dict_with_pipeline_config( - TurboWanT2V480PConfig(), - kwargs={ - "model_path": "IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers", - "performance_mode": "auto", - }, - ) - - self.assertTrue(args.dit_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - def test_explicit_fastwan_dit_layerwise_still_selects_dit_group(self): args = self._from_dict_with_pipeline_config( FastWan2_2_TI2V_5B_Config(), @@ -2270,28 +2037,6 @@ def test_explicit_fastwan_dit_layerwise_still_selects_dit_group(self): self.assertTrue(args.dit_cpu_offload) self.assertEqual(args.layerwise_offload_components, ["dit"]) - def test_auto_multi_gpu_wan_defers_residency_without_cfg(self): - with patch.object(ServerArgs, "_model_default_uses_cfg", return_value=False): - args = self._from_dict_with_pipeline_config( - WanT2V480PConfig(), - kwargs={ - "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "num_gpus": 2, - "performance_mode": "auto", - }, - ) - - self.assertFalse(args.use_fsdp_inference) - self.assertFalse(args.enable_cfg_parallel) - self.assertTrue(args.dit_cpu_offload) - self.assertTrue(args.layerwise_offload_components) - self.assertFalse(args.text_encoder_cpu_offload) - self.assertFalse(args.image_encoder_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - def test_explicit_multi_gpu_dit_layerwise_only_selects_dit_group(self): args = self._from_dict_with_pipeline_config( MOVAPipelineConfig(), @@ -2309,62 +2054,16 @@ def test_explicit_multi_gpu_dit_layerwise_only_selects_dit_group(self): self.assertTrue(args.image_encoder_cpu_offload) self.assertEqual(args.layerwise_offload_components, ["dit"]) - def test_auto_multi_gpu_ltx_defers_component_residency(self): + def test_auto_high_memory_ltx23_original_keeps_default_layerwise_components(self): args = self._from_dict_with_pipeline_config( LTX2PipelineConfig(), - available_memory_gb=76, + memory_gb=140, + available_memory_gb=134, kwargs={ - "model_path": "Lightricks/LTX-2", + "model_path": "Lightricks/LTX-2.3", "num_gpus": 2, "pipeline_class_name": "LTX2TwoStagePipeline", - }, - ) - - self.assertFalse(args.use_fsdp_inference) - self.assertTrue(args.dit_cpu_offload) - self.assertTrue(args.layerwise_offload_components) - self.assertFalse(args.text_encoder_cpu_offload) - self.assertFalse(args.image_encoder_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - - def test_auto_high_memory_ltx23_resident_defers_aux_placement(self): - args = self._from_dict_with_pipeline_config( - LTX2PipelineConfig(), - memory_gb=140, - available_memory_gb=134, - kwargs={ - "model_path": "Lightricks/LTX-2.3", - "num_gpus": 2, - "pipeline_class_name": "LTX2TwoStagePipeline", - }, - ) - - self.assertEqual(args.ltx2_two_stage_device_mode, "resident") - self.assertFalse(args.use_fsdp_inference) - self.assertFalse(args.dit_cpu_offload) - self.assertEqual(args.residency_mode("transformer"), RESIDENT) - self.assertEqual(args.residency_mode("transformer_2"), RESIDENT) - self.assertFalse(args.text_encoder_cpu_offload) - self.assertFalse(args.image_encoder_cpu_offload) - self.assertFalse(args.vae_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - - def test_auto_high_memory_ltx23_original_keeps_default_layerwise_components(self): - args = self._from_dict_with_pipeline_config( - LTX2PipelineConfig(), - memory_gb=140, - available_memory_gb=134, - kwargs={ - "model_path": "Lightricks/LTX-2.3", - "num_gpus": 2, - "pipeline_class_name": "LTX2TwoStagePipeline", - "ltx2_two_stage_device_mode": "original", + "ltx2_two_stage_device_mode": "original", }, ) @@ -2403,25 +2102,6 @@ def test_explicit_layerwise_components_preserved_in_ltx23_resident(self): self.assertEqual(args.ltx2_two_stage_device_mode, "resident") self.assertEqual(args.layerwise_offload_components, ["text_encoder"]) - def test_auto_multi_gpu_qwen_defers_residency_with_cfg(self): - args = self._from_dict_with_pipeline_config( - QwenImagePipelineConfig(), - kwargs={ - "model_path": "Qwen/Qwen-Image", - "num_gpus": 2, - "performance_mode": "auto", - }, - ) - - self.assertFalse(args.use_fsdp_inference) - self.assertTrue(args.enable_cfg_parallel) - self.assertTrue(args.dit_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - self.assertFalse(args.vae_cpu_offload) - def test_auto_multi_gpu_zimage_defers_fsdp_to_weight_inventory(self): args = self._from_dict_with_pipeline_config( ZImagePipelineConfig(), @@ -2435,21 +2115,6 @@ def test_auto_multi_gpu_zimage_defers_fsdp_to_weight_inventory(self): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) - def test_server_warmup_keeps_zimage_dit_available_to_residency_planner(self): - args = self._from_dict_with_pipeline_config( - ZImagePipelineConfig(), - kwargs={ - "model_path": "Tongyi-MAI/Z-Image", - "num_gpus": 2, - "performance_mode": "auto", - "warmup_mode": "server", - }, - ) - - self.assertFalse(args.use_fsdp_inference) - self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD) - self.assertTrue(args.enable_cfg_parallel) - def test_auto_multi_gpu_zimage_turbo_skips_fsdp(self): args = self._from_dict_with_pipeline_config( ZImagePipelineConfig(), @@ -2476,13 +2141,15 @@ def test_auto_multi_gpu_qwen_preserves_explicit_fsdp_false(self): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) - # Explicit FSDP selection must not freeze unrelated, implicit - # residency decisions before component weights are known. - self.assertTrue(args.dit_cpu_offload) + # Explicit FSDP selection must not freeze unrelated, implicit DiT + # residency decisions on a high-memory GPU. + self.assertFalse(args.dit_cpu_offload) self.assertFalse(args.vae_cpu_offload) + # The layerwise filter still drops VAE (kept resident); encoders stay + # offloaded. self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], + ["text_encoder", "image_encoder"], ) def test_auto_multi_gpu_qwen_skips_fsdp_when_available_memory_is_low(self): @@ -2498,10 +2165,12 @@ def test_auto_multi_gpu_qwen_skips_fsdp_when_available_memory_is_low(self): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) - self.assertTrue(args.dit_cpu_offload) + # 50gb still > image threshold (45gb): vae and dit stay resident, while + # the encoders remain offloaded; qwen does not opt into auto fsdp. + self.assertFalse(args.dit_cpu_offload) self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], + ["text_encoder", "image_encoder"], ) self.assertFalse(args.vae_cpu_offload) @@ -2520,46 +2189,6 @@ def test_auto_multi_gpu_qwen_uses_selected_gpu_min_available_memory(self): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) - def test_auto_multi_gpu_qwen_defers_residency_with_headroom(self): - args = self._from_dict_with_pipeline_config( - QwenImagePipelineConfig(), - available_memory_gb={1: 72, 2: 80}, - kwargs={ - "model_path": "Qwen/Qwen-Image", - "base_gpu_id": 1, - "num_gpus": 2, - "performance_mode": "auto", - }, - ) - - self.assertFalse(args.use_fsdp_inference) - self.assertTrue(args.enable_cfg_parallel) - self.assertTrue(args.dit_cpu_offload) - self.assertEqual( - args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], - ) - self.assertFalse(args.vae_cpu_offload) - - def test_auto_minimax_h3_defers_large_component_residency(self): - args = self._from_dict_with_pipeline_config( - MiniMaxH3PipelineConfig(), - memory_gb=141, - available_memory_gb=130, - kwargs={ - "model_path": "MiniMaxAI/MiniMax-H3", - "num_gpus": 8, - "ulysses_degree": 8, - "performance_mode": "auto", - }, - ) - - self.assertTrue(args.dit_cpu_offload) - self.assertFalse(args.text_encoder_cpu_offload) - self.assertFalse(args.vae_cpu_offload) - self.assertIn("text_encoder", args.layerwise_offload_components or []) - self.assertIn("vae", args.layerwise_offload_components or []) - def test_auto_minimax_h3_keeps_memory_policy_below_residency_threshold(self): args = self._from_dict_with_pipeline_config( MiniMaxH3PipelineConfig(), @@ -2814,6 +2443,375 @@ def test_ltx23_snapshot_device_mode_cli_is_rejected(self): with self.assertRaises(SystemExit): parser.parse_known_args(argv) + 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_longlive_residency_scales_with_available_memory(self): + high_memory_args = self._from_dict_with_pipeline_config( + LongLive2T2VConfig(), + 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.assertNotIn("vae", high_memory_offload) + + constrained_args = self._from_dict_with_pipeline_config( + LongLive2T2VConfig(), + memory_gb=50, + 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(), + LongCatImagePipelineConfig(), + ) + + for pipeline_config in pipeline_configs: + high_memory_args = self._from_dict_with_pipeline_config( + pipeline_config, + memory_gb=80, + kwargs={"performance_mode": "auto"}, + ) + self.assertNotIn( + "text_encoder", high_memory_args.layerwise_offload_components or [] + ) + self.assertFalse(high_memory_args.text_encoder_cpu_offload) + + constrained_args = self._from_dict_with_pipeline_config( + pipeline_config, + memory_gb=60, + kwargs={"performance_mode": "auto"}, + ) + self.assertIn( + "text_encoder", constrained_args.layerwise_offload_components or [] + ) + + def test_default_auto_keeps_image_vae_resident_when_memory_allows(self): + args = self._from_dict_with_pipeline_config( + QwenImagePipelineConfig(), + kwargs={"model_path": "Qwen/Qwen-Image"}, + ) + + self.assertEqual(args.performance_mode, "auto") + self.assertFalse(args.use_fsdp_inference) + # 80gb > image threshold (45gb): vae and dit stay resident, while the + # large encoders use layerwise offload. + self.assertFalse(args.dit_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder"], + ) + self.assertFalse(args.vae_cpu_offload) + + def test_auto_zimage_keeps_dit_resident_on_5090(self): + args = self._from_dict_with_pipeline_config( + ZImagePipelineConfig(), + memory_gb=32, + available_memory_gb=31, + kwargs={ + "model_path": "Tongyi-MAI/Z-Image-Turbo", + "performance_mode": "auto", + }, + ) + + self.assertFalse(args.dit_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder"], + ) + + def test_auto_ltx_original_replaces_component_cpu_offload( + self, + ): + args = self._from_dict_with_pipeline_config( + LTX2PipelineConfig(), + available_memory_gb=76, + kwargs={ + "model_path": "Lightricks/LTX-2.3", + "pipeline_class_name": "LTX2TwoStageHQPipeline", + "performance_mode": "auto", + }, + ) + + self.assertEqual(args.ltx2_two_stage_device_mode, "original") + self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.layerwise_offload_components) + self.assertFalse(args.text_encoder_cpu_offload) + self.assertFalse(args.image_encoder_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder", "vae"], + ) + + def test_auto_wan_keeps_single_dit_resident_on_h100(self): + args = self._from_dict_with_pipeline_config( + WanT2V480PConfig(), + kwargs={"performance_mode": "auto"}, + ) + + self.assertTrue(args.layerwise_offload_components) + self.assertFalse(args.use_fsdp_inference) + self.assertFalse(args.dit_cpu_offload) + self.assertFalse(args.text_encoder_cpu_offload) + self.assertFalse(args.image_encoder_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder", "vae"], + ) + + def test_auto_wan2_1_14b_keeps_dit_resident_on_h100(self): + for pipeline_config, model_path in ( + (WanT2V720PConfig(), "Wan-AI/Wan2.1-T2V-14B-Diffusers"), + (WanI2V480PConfig(), "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"), + (WanI2V720PConfig(), "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"), + ): + with self.subTest(pipeline_config=pipeline_config.__class__.__name__): + args = self._from_dict_with_pipeline_config( + pipeline_config, + kwargs={ + "model_path": model_path, + "performance_mode": "auto", + }, + ) + + self.assertTrue(args.layerwise_offload_components) + self.assertFalse(args.dit_cpu_offload) + self.assertEqual(args.dit_offload_prefetch_size, 0.0) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder", "vae"], + ) + + def test_auto_mova_keeps_dit_resident_at_memory_threshold(self): + args = self._from_dict_with_pipeline_config( + MOVAPipelineConfig(), + memory_gb=140, + kwargs={ + "model_path": "OpenMOSS-Team/MOVA-360p", + "performance_mode": "auto", + }, + ) + + self.assertFalse(args.dit_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder"], + ) + + def test_auto_cosmos3_keeps_dit_resident_on_high_memory_gpu(self): + args = self._from_dict_with_pipeline_config( + Cosmos3Config(model_path="nvidia/Cosmos3-Nano"), + available_memory_gb=95, + kwargs={ + "model_path": "nvidia/Cosmos3-Nano", + "performance_mode": "auto", + }, + ) + + self.assertFalse(args.dit_cpu_offload) + self.assertFalse(args.vae_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder"], + ) + + def test_auto_cosmos3_super_keeps_dit_resident_on_high_memory_gpu(self): + # Super is a single-DiT pipeline like Nano, so above the threshold the + # component-offload round trip is pure per-request copy cost. + args = self._from_dict_with_pipeline_config( + Cosmos3Config(model_path="nvidia/Cosmos3-Super"), + available_memory_gb=139, + kwargs={ + "model_path": "nvidia/Cosmos3-Super", + "performance_mode": "auto", + }, + ) + + self.assertFalse(args.dit_cpu_offload) + self.assertFalse(args.vae_cpu_offload) + + def test_auto_fastwan_keeps_dit_resident_on_h100(self): + args = self._from_dict_with_pipeline_config( + FastWan2_2_TI2V_5B_Config(), + available_memory_gb=72, + kwargs={ + "model_path": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", + "performance_mode": "auto", + }, + ) + + self.assertFalse(args.dit_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder", "vae"], + ) + + def test_auto_fast_hunyuan_keeps_dit_resident_on_h100(self): + args = self._from_dict_with_pipeline_config( + FastHunyuanConfig(), + available_memory_gb=72, + kwargs={ + "model_path": "FastVideo/FastHunyuan-diffusers", + "performance_mode": "auto", + }, + ) + + self.assertFalse(args.dit_cpu_offload) + self.assertFalse(args.vae_cpu_offload) + + def test_auto_turbo_wan_keeps_dit_resident_on_h100(self): + args = self._from_dict_with_pipeline_config( + TurboWanT2V480PConfig(), + kwargs={ + "model_path": "IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers", + "performance_mode": "auto", + }, + ) + + self.assertFalse(args.dit_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder", "vae"], + ) + + def test_auto_multi_gpu_wan_uses_layerwise_offload_without_cfg(self): + with patch.object(ServerArgs, "_model_default_uses_cfg", return_value=False): + args = self._from_dict_with_pipeline_config( + WanT2V480PConfig(), + kwargs={ + "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "num_gpus": 2, + "performance_mode": "auto", + }, + ) + + self.assertFalse(args.use_fsdp_inference) + self.assertFalse(args.enable_cfg_parallel) + self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.layerwise_offload_components) + self.assertFalse(args.text_encoder_cpu_offload) + self.assertFalse(args.image_encoder_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder", "vae"], + ) + + def test_auto_multi_gpu_ltx_replaces_component_cpu_offload_with_resident_dit(self): + args = self._from_dict_with_pipeline_config( + LTX2PipelineConfig(), + available_memory_gb=76, + kwargs={ + "model_path": "Lightricks/LTX-2", + "num_gpus": 2, + "pipeline_class_name": "LTX2TwoStagePipeline", + }, + ) + + self.assertFalse(args.use_fsdp_inference) + self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.layerwise_offload_components) + self.assertFalse(args.text_encoder_cpu_offload) + self.assertFalse(args.image_encoder_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder", "vae"], + ) + + def test_auto_high_memory_ltx23_resident_keeps_aux_components_resident(self): + args = self._from_dict_with_pipeline_config( + LTX2PipelineConfig(), + memory_gb=140, + available_memory_gb=134, + kwargs={ + "model_path": "Lightricks/LTX-2.3", + "num_gpus": 2, + "pipeline_class_name": "LTX2TwoStagePipeline", + }, + ) + + self.assertEqual(args.ltx2_two_stage_device_mode, "resident") + self.assertFalse(args.use_fsdp_inference) + self.assertFalse(args.dit_cpu_offload) + self.assertFalse(args.text_encoder_cpu_offload) + self.assertFalse(args.image_encoder_cpu_offload) + self.assertFalse(args.vae_cpu_offload) + self.assertIsNone(args.layerwise_offload_components) + + def test_auto_multi_gpu_qwen_keeps_vae_resident_with_cfg(self): + args = self._from_dict_with_pipeline_config( + QwenImagePipelineConfig(), + kwargs={ + "model_path": "Qwen/Qwen-Image", + "num_gpus": 2, + "performance_mode": "auto", + }, + ) + + self.assertFalse(args.use_fsdp_inference) + self.assertTrue(args.enable_cfg_parallel) + # 80gb > image threshold (45gb): vae and dit stay resident, while the + # large encoders use layerwise offload. + self.assertFalse(args.dit_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder"], + ) + self.assertFalse(args.vae_cpu_offload) + + def test_auto_multi_gpu_qwen_keeps_vae_resident_with_headroom(self): + args = self._from_dict_with_pipeline_config( + QwenImagePipelineConfig(), + available_memory_gb={1: 72, 2: 80}, + kwargs={ + "model_path": "Qwen/Qwen-Image", + "base_gpu_id": 1, + "num_gpus": 2, + "performance_mode": "auto", + }, + ) + + self.assertFalse(args.use_fsdp_inference) + self.assertTrue(args.enable_cfg_parallel) + # min available across selected gpus is 72gb > image threshold (45gb): + # vae and dit stay resident, while the encoders remain offloaded. + self.assertFalse(args.dit_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder"], + ) + self.assertFalse(args.vae_cpu_offload) + + def test_auto_minimax_h3_keeps_large_components_resident_with_headroom(self): + args = self._from_dict_with_pipeline_config( + MiniMaxH3PipelineConfig(), + memory_gb=141, + available_memory_gb=130, + kwargs={ + "model_path": "MiniMaxAI/MiniMax-H3", + "num_gpus": 8, + "ulysses_degree": 8, + "performance_mode": "auto", + }, + ) + + self.assertFalse(args.dit_cpu_offload) + self.assertFalse(args.text_encoder_cpu_offload) + self.assertFalse(args.vae_cpu_offload) + self.assertNotIn("text_encoder", args.layerwise_offload_components or []) + self.assertNotIn("vae", args.layerwise_offload_components or []) + class TestKVGatherDegree(unittest.TestCase): def test_sp2_defaults_to_kv_gather(self): From 3fbc61f2096215a228a278fbff82be0f24fbd7c3 Mon Sep 17 00:00:00 2001 From: Mick Qian Date: Tue, 8 Sep 2026 01:41:15 +0800 Subject: [PATCH 70/70] ci: baseline the load peaks of the pre-warmup seed The seed keeps native DiTs resident from load, so the load peaks of the Wan2.2 A14B 2-GPU cases and the 5090 Wan/Turbo-Wan/FLUX.2-klein cases rise accordingly (run 34131148154 at a330811); runtime peaks follow the apply layer's refresh. Co-Authored-By: Claude Fable 5.1 --- .../test/server/perf_baselines/5090.json | 24 ++++++++-------- .../test/server/perf_baselines/h100.json | 28 +++++++++---------- 2 files changed, 26 insertions(+), 26 deletions(-) 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 76ba33833203..fa9d1d3059b7 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/5090.json @@ -60,10 +60,10 @@ "expected_e2e_ms": 18255.39, "expected_avg_denoise_ms": 338.78, "expected_median_denoise_ms": 345.95, - "load_peak_vram_mb": 1044.0, - "runtime_peak_vram_mb": 20048.0, - "load_peak_allocated_mb": 1019.0, - "runtime_peak_allocated_mb": 17724.0, + "load_peak_vram_mb": 8522.0, + "runtime_peak_vram_mb": 20056.0, + "load_peak_allocated_mb": 8412.0, + "runtime_peak_allocated_mb": 17723.0, "estimated_full_test_time_s": 94.0 }, "wan2_1_t2v_1.3b": { @@ -86,10 +86,10 @@ "expected_e2e_ms": 23418.73, "expected_avg_denoise_ms": 426.59, "expected_median_denoise_ms": 427.94, - "load_peak_vram_mb": 928.0, - "runtime_peak_vram_mb": 14074.0, - "load_peak_allocated_mb": 903.0, - "runtime_peak_allocated_mb": 9752.0, + "load_peak_vram_mb": 3808.0, + "runtime_peak_vram_mb": 14120.0, + "load_peak_allocated_mb": 3641.0, + "runtime_peak_allocated_mb": 9778.0, "estimated_full_test_time_s": 160.9 }, "turbo_wan2_1_t2v_1.3b": { @@ -111,10 +111,10 @@ "expected_e2e_ms": 2155.59, "expected_avg_denoise_ms": 143.62, "expected_median_denoise_ms": 160.64, - "load_peak_vram_mb": 928.0, - "runtime_peak_vram_mb": 13660.0, - "load_peak_allocated_mb": 903.0, - "runtime_peak_allocated_mb": 9599.0, + "load_peak_vram_mb": 3810.0, + "runtime_peak_vram_mb": 13962.0, + "load_peak_allocated_mb": 3643.0, + "runtime_peak_allocated_mb": 9625.0, "load_peak_host_anon_mb": 6759.33, "runtime_peak_host_anon_mb": 6759.33, "estimated_full_test_time_s": 200.3 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 70ee4338dca3..399803b44bb3 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -1628,10 +1628,10 @@ "expected_e2e_ms": 73811.61, "expected_avg_denoise_ms": 1689.73, "expected_median_denoise_ms": 1714.58, - "load_peak_vram_mb": 6440.0, - "runtime_peak_vram_mb": 26322.0, - "load_peak_allocated_mb": 6177.0, - "runtime_peak_allocated_mb": 19978.0, + "load_peak_vram_mb": 60038.0, + "runtime_peak_vram_mb": 75308.0, + "load_peak_allocated_mb": 57269.0, + "runtime_peak_allocated_mb": 71063.0, "estimated_full_test_time_s": 264.1 }, "wan2_1_i2v_14b_480P_2gpu": { @@ -1701,10 +1701,10 @@ "expected_avg_denoise_ms": 653.89, "expected_median_denoise_ms": 656.68, "load_peak_vram_mb": 35344.0, - "runtime_peak_vram_mb": 52260.0, + "runtime_peak_vram_mb": 52478.0, "warmup_peak_vram_mb": 45782.0, "load_peak_allocated_mb": 34033.0, - "runtime_peak_allocated_mb": 50073.0, + "runtime_peak_allocated_mb": 50074.0, "estimated_full_test_time_s": 243.0 }, "wan2_1_i2v_14b_720P_2gpu": { @@ -1833,11 +1833,11 @@ "expected_e2e_ms": 83297.78, "expected_avg_denoise_ms": 2029.88, "expected_median_denoise_ms": 2020.71, - "load_peak_vram_mb": 6440.0, - "runtime_peak_vram_mb": 21524.0, + "load_peak_vram_mb": 60038.0, + "runtime_peak_vram_mb": 73984.0, "warmup_peak_vram_mb": 46020.0, - "load_peak_allocated_mb": 6175.0, - "runtime_peak_allocated_mb": 15472.0, + "load_peak_allocated_mb": 57269.0, + "runtime_peak_allocated_mb": 68578.0, "estimated_full_test_time_s": 204.3 }, "wan2_1_t2v_14b_2gpu": { @@ -1964,10 +1964,10 @@ "expected_e2e_ms": 62134.19, "expected_avg_denoise_ms": 1437.82, "expected_median_denoise_ms": 1416.85, - "load_peak_vram_mb": 6784.0, - "runtime_peak_vram_mb": 21450.0, - "load_peak_allocated_mb": 6499.0, - "runtime_peak_allocated_mb": 14775.0, + "load_peak_vram_mb": 60358.0, + "runtime_peak_vram_mb": 75188.0, + "load_peak_allocated_mb": 57561.0, + "runtime_peak_allocated_mb": 68559.0, "estimated_full_test_time_s": 181.8 }, "wan2_1_t2v_1_3b_lora_1gpu": {