[Model Loading] Overlap checkpoint staging with CUDA graph capture during startup (#32017)

Co-authored-by: Wenhui Zhu <wzhu59@asu.edu>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>

python/sglang/benchmark/one_batch.py

@@ -350,9 +350,13 @@ def load_model(server_args, port_args, gpu_id, tp_rank):
         model_runner = MlxModelRunnerStub(**runner_kwargs)
     else:
         model_runner = ModelRunner(**runner_kwargs)
+        if server_args.is_startup_weight_load_overlap:
+            model_runner.start_startup_weight_load()
         model_runner.alloc_memory_pool()
         model_runner.init_attention_backends()
         model_runner.init_cuda_graphs()
+        if server_args.is_startup_weight_load_overlap:
+            model_runner.finalize_startup_weight_load()
     rank_print(f"max_total_num_tokens={model_runner.max_total_num_tokens}")
     tokenizer = get_tokenizer(
         server_args.tokenizer_path,

python/sglang/srt/hardware_backend/mlx/model_runner_stub.py

@@ -107,7 +107,15 @@ class MlxModelRunnerStub(ModelRunner):
     # that path working instead of raising AttributeError.
     prefill_aware_swa = False
 
+    @staticmethod
+    def validate_startup_weight_load_mode(server_args) -> None:
+        if server_args.is_startup_weight_load_overlap:
+            raise ValueError(
+                "--startup-weight-load-mode=overlap is not supported: CUDA only"
+            )
+
     def __init__(self, *args, mlx_pool_size: int | None = None, **kwargs):
+        self.validate_startup_weight_load_mode(kwargs["server_args"])
         self._mlx_pool_size = mlx_pool_size
         super().__init__(*args, **kwargs)
 

python/sglang/srt/hardware_backend/mlx/tp_worker.py

@@ -78,11 +78,14 @@ class MlxTpModelWorker(TpModelWorker):
 
     def _init_model_runner(self):
         """Create MLX runner first (auto-sizes pool), then stub with matching size."""
-        from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
         from sglang.srt.hardware_backend.mlx.model_runner_stub import (
             MlxModelRunnerStub,
         )
 
+        MlxModelRunnerStub.validate_startup_weight_load_mode(self.server_args)
+
+        from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
+
         logger.info("Initializing MlxModelRunner for end-to-end MLX inference")
         init_kwargs = dict(
             model_path=get_model().model_path,

python/sglang/srt/managers/scheduler.py

@@ -986,6 +986,8 @@ class Scheduler(
     def init_model_worker(self):
         # Load model weights.
         self.init_tp_model_worker()
+        if self.server_args.is_startup_weight_load_overlap:
+            self.tp_worker.start_startup_weight_load()
         self.maybe_init_draft_worker()
 
         # Prepare KV cache pools for all workers
@@ -1002,6 +1004,9 @@ class Scheduler(
             model_runner.post_capture_resize_kv_pool()
             self.kv_cache_allocation_time += time.perf_counter() - tic
 
+        if self.server_args.is_startup_weight_load_overlap:
+            self.tp_worker.finalize_startup_weight_load()
+
         if (
             get_exec().moe.elastic_ep_backend is not None
             and get_exec().moe.ep_join_mode == "recover"

python/sglang/srt/managers/tp_worker.py

@@ -429,6 +429,18 @@ class TpModelWorker(BaseTpWorker):
         for mr in self.model_runner_list[1:]:
             mr.init_cuda_graphs(capture_decode_cuda_graph=capture_decode_cuda_graph)
 
+    def start_startup_weight_load(self) -> None:
+        """Start deferred checkpoint prefetching for all model runners."""
+        self.model_runner.start_startup_weight_load()
+        for mr in self.model_runner_list[1:]:
+            mr.start_startup_weight_load()
+
+    def finalize_startup_weight_load(self) -> None:
+        """Commit deferred startup weights for all model runners."""
+        self.model_runner.finalize_startup_weight_load()
+        for mr in self.model_runner_list[1:]:
+            mr.finalize_startup_weight_load()
+
     def _init_model_config(self):
         from sglang.srt.configs.model_config import ModelConfig
 

python/sglang/srt/model_executor/model_runner.py

@@ -431,6 +431,11 @@ class ModelRunner:
         # For hisparse (must be set before initialize() so CUDA graph capture can see it)
         self.hisparse_coordinator = None
 
+        # The native overlap path replaces this during load_model(). Keep the
+        # no-pending-work invariant for lightweight backends that override the
+        # base initialization and weight-loading flow.
+        self.startup_weight_load = None
+
         # Load model weights and configure
         self.initialize()
         self.check_quantized_moe_compatibility()
@@ -1115,6 +1120,7 @@ class ModelRunner:
             )
         self.loader = loaded.loader
         self.model = loaded.model
+        self.startup_weight_load = loaded.startup_weight_load
         if loaded.remote_instance_weight_info is not None:
             self.remote_instance_weight_transporter.weight_info = (
                 loaded.remote_instance_weight_info
@@ -1158,14 +1164,15 @@ class ModelRunner:
         # This handles both config.json (standard) and hf_quant_config.json (ModelOpt)
         quant_str = self.model_config.get_quantization_config_log_str()
 
-        logger.info(
-            f"Load weight end. "
-            f"elapsed={self.weight_load_time:.2f} s, "
-            f"type={type(self.model).__name__}, "
-            f"{quant_str + ', ' if quant_str else ''}"
-            f"avail mem={after_avail_memory:.2f} GB, "
-            f"mem usage={self.weight_load_mem_usage:.2f} GB."
-        )
+        if self.startup_weight_load is None:
+            logger.info(
+                f"Load weight end. "
+                f"elapsed={self.weight_load_time:.2f} s, "
+                f"type={type(self.model).__name__}, "
+                f"{quant_str + ', ' if quant_str else ''}"
+                f"avail mem={after_avail_memory:.2f} GB, "
+                f"mem usage={self.weight_load_mem_usage:.2f} GB."
+            )
 
         report_online_quantization(model=self.model, server_args=self.server_args)
 
@@ -1191,11 +1198,36 @@ class ModelRunner:
             logger,
         )
 
+        if self.startup_weight_load is None:
+            dist_barrier_after_load(
+                elastic_ep_backend=get_exec().moe.elastic_ep_backend,
+                tp_rank=self.ps.tp_rank,
+                is_ep_joiner=self.server_args.is_ep_joiner,
+            )
+
+    def start_startup_weight_load(self) -> None:
+        assert self.startup_weight_load is not None
+        self.startup_weight_load.start_prefetch()
+
+    def finalize_startup_weight_load(self) -> None:
+        """Commit the real weights, then run the post-load barrier.
+
+        The barrier moves here because ``load_model`` returns with sentinel
+        values under overlap, so this is the first point at which "weights are
+        loaded" is true for this rank. It follows the commit and its validation
+        deliberately: a rank that fails to commit must not report readiness. A
+        commit failure is terminal for the process, so peer ranks observe it as
+        a barrier timeout rather than a clean collective abort, which matches
+        the existing startup contract for load failures.
+        """
+        assert self.startup_weight_load is not None
+        self.startup_weight_load.finalize()
         dist_barrier_after_load(
             elastic_ep_backend=get_exec().moe.elastic_ep_backend,
             tp_rank=self.ps.tp_rank,
             is_ep_joiner=is_ep_joiner(),
         )
+        self.startup_weight_load = None
 
     def maybe_precompile_model_kernels_after_loading(self) -> None:
         maybe_precompile_model_kernels_after_loading(self.model, self.device)

python/sglang/srt/model_executor/model_runner_components/load_model_utils.py

@@ -59,6 +59,7 @@ class LoadedModel(msgspec.Struct, frozen=True, kw_only=True):
     loader: Any
     model: Any
     remote_instance_weight_info: Optional[Any]
+    startup_weight_load: Optional[Any] = None
 
 
 def maybe_downgrade_dtype_for_legacy_gpu(
@@ -292,6 +293,7 @@ def load_model_with_memory_saver(
         enable_cpu_backup = False
 
     remote_instance_weight_info = None
+    startup_weight_load = None
     with memory_saver_adapter.region(
         GPU_MEMORY_TYPE_WEIGHTS,
         enable_cpu_backup=enable_cpu_backup,
@@ -300,10 +302,26 @@ def load_model_with_memory_saver(
             load_config=load_config,
             model_config=model_config,
         )
-        model = loader.load_model(
-            model_config=model_config,
-            device_config=DeviceConfig(device, gpu_id),
-        )
+        device_config = DeviceConfig(device, gpu_id)
+        if server_args.is_startup_weight_load_overlap:
+            from sglang.srt.model_executor.model_runner_components.startup_weight_load import (
+                StartupWeightLoadManager,
+            )
+
+            startup_weight_load = StartupWeightLoadManager.create_from_server_args(
+                loader=loader,
+                model_config=model_config,
+                load_config=load_config,
+                device_config=device_config,
+                server_args=server_args,
+                is_draft_worker=is_draft_worker,
+            )
+            model = startup_weight_load.prepare()
+        else:
+            model = loader.load_model(
+                model_config=model_config,
+                device_config=device_config,
+            )
         if hasattr(loader, "remote_instance_transfer_engine_weight_info"):
             remote_instance_weight_info = (
                 loader.remote_instance_transfer_engine_weight_info
@@ -318,6 +336,7 @@ def load_model_with_memory_saver(
         loader=loader,
         model=model,
         remote_instance_weight_info=remote_instance_weight_info,
+        startup_weight_load=startup_weight_load,
     )
 
 

python/sglang/srt/model_executor/model_runner_components/startup_weight_load.py

@@ -0,0 +1,591 @@
+from __future__ import annotations
+
+import dataclasses
+import enum
+import logging
+import time
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+from torch import nn
+
+from sglang.srt.configs.device_config import DeviceConfig
+from sglang.srt.configs.load_config import LoadConfig, LoadFormat
+from sglang.srt.distributed.parallel_state import monkey_patch_vllm_parallel_state
+from sglang.srt.model_executor.cuda_graph_config import Backend, Phase
+from sglang.srt.model_loader.loader import DefaultModelLoader
+from sglang.srt.model_loader.utils import get_model_architecture
+from sglang.srt.model_loader.weight_utils import (
+    CAPTURE_SAFE_WEIGHT_SENTINEL,
+    CheckpointFilePrefetchHandle,
+)
+from sglang.srt.platforms import current_platform
+
+if TYPE_CHECKING:
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+_SUPPORTED_ARCHITECTURES = frozenset(
+    {
+        "LlamaForCausalLM",
+        "Qwen2ForCausalLM",
+        "Qwen3ForCausalLM",
+    }
+)
+_SUPPORTED_DTYPES = frozenset({torch.float16, torch.bfloat16})
+
+
+def _get_canonical_model_class(architecture: str):
+    if architecture == "LlamaForCausalLM":
+        from sglang.srt.models.llama import LlamaForCausalLM
+
+        return LlamaForCausalLM
+    if architecture == "Qwen2ForCausalLM":
+        from sglang.srt.models.qwen2 import Qwen2ForCausalLM
+
+        return Qwen2ForCausalLM
+    if architecture == "Qwen3ForCausalLM":
+        from sglang.srt.models.qwen3 import Qwen3ForCausalLM
+
+        return Qwen3ForCausalLM
+    raise ValueError(f"Unsupported startup-overlap architecture: {architecture}")
+
+
+class StartupWeightLoadState(str, enum.Enum):
+    CREATED = "created"
+    PREPARING = "preparing"
+    CAPTURE_READY = "capture_ready"
+    PREFETCHING = "prefetching"
+    COMMITTING = "committing"
+    READY = "ready"
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class StartupWeightLoadOptions:
+    device: str
+    is_cuda_platform: bool
+    cuda_graph_enabled: bool
+    prefill_cuda_graph_backend: Backend
+    is_draft_worker: bool
+    speculative_algorithm: Optional[str]
+    tp_size: int
+    attn_cp_size: int
+    dcp_size: int
+    pp_size: int
+    dp_size: int
+    ep_size: int
+    cpu_offload_gb: int
+    offload_group_size: int
+    enable_memory_saver: bool
+    enable_weights_cpu_backup: bool
+    torchao_config: str
+    enable_lora: bool
+    has_lora_paths: bool
+    weight_loader_disable_mmap: bool
+    weight_loader_drop_cache_after_load: bool
+    has_custom_weight_loader: bool
+    enable_torch_compile: bool
+    prefetch_num_threads: int
+
+    @classmethod
+    def from_server_args(
+        cls,
+        *,
+        server_args: ServerArgs,
+        is_draft_worker: bool,
+    ) -> StartupWeightLoadOptions:
+        cuda_graph_config = server_args.cuda_graph_config
+        cuda_graph_enabled = any(
+            getattr(cuda_graph_config, phase).backend != Backend.DISABLED
+            for phase in Phase.ALL
+        )
+        return cls(
+            device=server_args.device,
+            is_cuda_platform=current_platform.is_cuda(),
+            cuda_graph_enabled=cuda_graph_enabled,
+            prefill_cuda_graph_backend=cuda_graph_config.prefill.backend,
+            is_draft_worker=is_draft_worker,
+            speculative_algorithm=server_args.speculative_algorithm,
+            tp_size=server_args.tp_size,
+            attn_cp_size=server_args.attn_cp_size,
+            dcp_size=server_args.dcp_size,
+            pp_size=server_args.pp_size,
+            dp_size=server_args.dp_size,
+            ep_size=server_args.ep_size,
+            cpu_offload_gb=server_args.cpu_offload_gb,
+            offload_group_size=server_args.offload_group_size,
+            enable_memory_saver=server_args.enable_memory_saver,
+            enable_weights_cpu_backup=server_args.enable_weights_cpu_backup,
+            torchao_config=server_args.torchao_config,
+            enable_lora=server_args.enable_lora,
+            has_lora_paths=bool(server_args.lora_paths),
+            weight_loader_disable_mmap=server_args.weight_loader_disable_mmap,
+            weight_loader_drop_cache_after_load=(
+                server_args.weight_loader_drop_cache_after_load
+            ),
+            has_custom_weight_loader=bool(server_args.custom_weight_loader),
+            enable_torch_compile=server_args.enable_torch_compile,
+            prefetch_num_threads=server_args.weight_loader_prefetch_num_threads,
+        )
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class TensorStorageMetadata:
+    tensor: torch.Tensor = dataclasses.field(repr=False, compare=False)
+    data_ptr: int
+    shape: Tuple[int, ...]
+    stride: Tuple[int, ...]
+    dtype: torch.dtype
+    device: torch.device
+    storage_offset: int
+
+    @classmethod
+    def from_tensor(cls, tensor: torch.Tensor) -> TensorStorageMetadata:
+        return cls(
+            tensor=tensor,
+            data_ptr=tensor.data_ptr(),
+            shape=tuple(tensor.shape),
+            stride=tuple(tensor.stride()),
+            dtype=tensor.dtype,
+            device=tensor.device,
+            storage_offset=tensor.storage_offset(),
+        )
+
+    def matches(self, other: TensorStorageMetadata) -> bool:
+        return self.tensor is other.tensor and (
+            self.data_ptr,
+            self.shape,
+            self.stride,
+            self.dtype,
+            self.device,
+            self.storage_offset,
+        ) == (
+            other.data_ptr,
+            other.shape,
+            other.stride,
+            other.dtype,
+            other.device,
+            other.storage_offset,
+        )
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class ModelStorageManifest:
+    tensors: Tuple[Tuple[str, TensorStorageMetadata], ...]
+
+    @classmethod
+    def capture(cls, model: nn.Module) -> ModelStorageManifest:
+        entries = []
+        for kind, tensors in (
+            ("parameter", model.named_parameters(remove_duplicate=False)),
+            ("buffer", model.named_buffers(remove_duplicate=False)),
+        ):
+            entries.extend(
+                (f"{kind}:{name}", TensorStorageMetadata.from_tensor(tensor))
+                for name, tensor in tensors
+            )
+        # Key explicitly by name because TensorStorageMetadata is not orderable,
+        # and stable name ordering keeps diagnostics deterministic for aliases.
+        return cls(tensors=tuple(sorted(entries, key=lambda entry: entry[0])))
+
+    def changed_names(self, model: nn.Module) -> Tuple[str, ...]:
+        before = dict(self.tensors)
+        after = dict(ModelStorageManifest.capture(model).tensors)
+        return tuple(
+            name
+            for name in sorted(before.keys() | after.keys())
+            if name not in before
+            or name not in after
+            or not before[name].matches(after[name])
+        )
+
+    def unchanged_parameter_names(self, value: float) -> Tuple[str, ...]:
+        """Return floating-point parameters still entirely equal to ``value``.
+
+        This is the capture-sentinel check, and it is deliberately strict: every
+        floating-point parameter must be rewritten by ``model.load_weights()``.
+        A model that keeps an ``__init__``-computed floating-point parameter with
+        no checkpoint entry will fail startup here rather than silently serve the
+        sentinel, so this doubles as the admission gate for widening
+        ``_SUPPORTED_ARCHITECTURES``. Buffers are excluded because
+        ``initialize_capture_safe_weights`` never overwrites them.
+        """
+        names = []
+        checks = []
+        seen_tensor_ids = set()
+        for name, metadata in self.tensors:
+            tensor = metadata.tensor
+            if (
+                not name.startswith("parameter:")
+                or not torch.is_floating_point(tensor)
+                or id(tensor) in seen_tensor_ids
+            ):
+                continue
+            seen_tensor_ids.add(id(tensor))
+            names.append(name)
+            checks.append(torch.all(tensor == value))
+
+        if not checks:
+            return ()
+        unchanged = torch.stack(checks).cpu().tolist()
+        return tuple(
+            name for name, is_unchanged in zip(names, unchanged) if is_unchanged
+        )
+
+
+class StartupWeightLoadManager:
+    """Coordinate native CPU staging with capture and post-capture commit."""
+
+    def __init__(
+        self,
+        *,
+        loader: DefaultModelLoader,
+        model_config: ModelConfig,
+        device_config: DeviceConfig,
+        options: StartupWeightLoadOptions,
+    ) -> None:
+        self._loader = loader
+        self._model_config = model_config
+        self._device_config = device_config
+        self._options = options
+        self._model: Optional[nn.Module] = None
+        self._resolved_sources: Tuple[DefaultModelLoader.ResolvedSource, ...] = ()
+        self._prefetch_handle: Optional[CheckpointFilePrefetchHandle] = None
+        self._state = StartupWeightLoadState.CREATED
+        self._created_at = time.perf_counter()
+        self._capture_ready_at: Optional[float] = None
+        self._prefetch_started_at: Optional[float] = None
+        self._prefetch_failure_reported = False
+
+    @classmethod
+    def create_from_server_args(
+        cls,
+        *,
+        loader,
+        model_config: ModelConfig,
+        load_config: LoadConfig,
+        device_config: DeviceConfig,
+        server_args: ServerArgs,
+        is_draft_worker: bool,
+    ) -> StartupWeightLoadManager:
+        """Build a manager straight from ``ServerArgs``.
+
+        Callers on the model-loading path only decide *whether* to overlap; the
+        knowledge of which server arguments matter, and every support rule,
+        stays in this module.
+        """
+        return cls.create(
+            loader=loader,
+            model_config=model_config,
+            load_config=load_config,
+            device_config=device_config,
+            options=StartupWeightLoadOptions.from_server_args(
+                server_args=server_args,
+                is_draft_worker=is_draft_worker,
+            ),
+        )
+
+    @classmethod
+    def create(
+        cls,
+        *,
+        loader,
+        model_config: ModelConfig,
+        load_config: LoadConfig,
+        device_config: DeviceConfig,
+        options: StartupWeightLoadOptions,
+    ) -> StartupWeightLoadManager:
+        unsupported_reason = cls._get_unsupported_reason(
+            loader=loader,
+            model_config=model_config,
+            load_config=load_config,
+            options=options,
+        )
+        if unsupported_reason is not None:
+            raise ValueError(
+                "--startup-weight-load-mode=overlap is not supported: "
+                f"{unsupported_reason}"
+            )
+        return cls(
+            loader=loader,
+            model_config=model_config,
+            device_config=device_config,
+            options=options,
+        )
+
+    @staticmethod
+    def _get_unsupported_reason(
+        *,
+        loader,
+        model_config: ModelConfig,
+        load_config: LoadConfig,
+        options: StartupWeightLoadOptions,
+    ) -> Optional[str]:
+        architectures = tuple(model_config.hf_config.architectures or ())
+        # NOTE(2026-08): The initial rollout supports only the configurations
+        # admitted below. Expand this matrix only with storage-stability,
+        # capture-sentinel, and startup-correctness coverage for the new case.
+        # Keep these checks here because they depend on resolved loader and model
+        # state; ServerArgs owns only the mode selection.
+        basic_rules = (
+            (not options.is_cuda_platform or options.device != "cuda", "CUDA only"),
+            (not options.cuda_graph_enabled, "CUDA graph capture is disabled"),
+            (
+                options.prefill_cuda_graph_backend == Backend.TC_PIECEWISE,
+                "tc_piecewise prefill CUDA graphs are not supported",
+            ),
+            (type(loader) is not DefaultModelLoader, "DefaultModelLoader only"),
+            (
+                load_config.load_format
+                not in (LoadFormat.AUTO, LoadFormat.SAFETENSORS),
+                "load format must be auto or safetensors",
+            ),
+            (options.is_draft_worker, "draft workers are not supported"),
+            (
+                load_config.draft_model_idx is not None,
+                "draft model loading is unsupported",
+            ),
+            (
+                options.speculative_algorithm is not None,
+                "speculative decoding is not supported",
+            ),
+            (options.tp_size not in (1, 2), "only TP1 and TP2 are supported"),
+            (
+                options.attn_cp_size != 1,
+                "attention context parallelism is not supported",
+            ),
+            (
+                options.dcp_size != 1,
+                "decode context parallelism is not supported",
+            ),
+            (options.pp_size != 1, "pipeline parallelism is not supported"),
+            (options.dp_size != 1, "data parallelism is not supported"),
+            (options.ep_size != 1, "expert parallelism is not supported"),
+            (options.cpu_offload_gb > 0, "CPU offload is not supported"),
+            (
+                options.offload_group_size > 0,
+                "layer-group offloading is not supported",
+            ),
+            (options.enable_memory_saver, "memory saver is not supported"),
+            (
+                options.enable_weights_cpu_backup,
+                "CPU weight backup is not supported",
+            ),
+            (bool(options.torchao_config), "TorchAO is not supported"),
+            (
+                options.enable_lora or options.has_lora_paths,
+                "LoRA is not supported",
+            ),
+            (
+                options.weight_loader_disable_mmap,
+                "safetensors mmap must be enabled",
+            ),
+            (
+                options.weight_loader_drop_cache_after_load,
+                "dropping the page cache during load is not supported",
+            ),
+            (
+                options.has_custom_weight_loader,
+                "custom weight loaders are not supported",
+            ),
+            (options.enable_torch_compile, "torch.compile is not supported"),
+        )
+        unsupported_reason = next(
+            (reason for unsupported, reason in basic_rules if unsupported),
+            None,
+        )
+        if unsupported_reason is not None:
+            return unsupported_reason
+
+        model_rules = (
+            (model_config.dtype not in _SUPPORTED_DTYPES, "FP16 or BF16 only"),
+            (model_config.quantization is not None, "quantization is not supported"),
+            (
+                bool(getattr(model_config, "modelopt_quant", False)),
+                "ModelOpt is not supported",
+            ),
+            (model_config.is_multimodal, "multimodal models are not supported"),
+            (not model_config.is_generation, "generation models only"),
+            (
+                len(architectures) != 1
+                or architectures[0] not in _SUPPORTED_ARCHITECTURES,
+                "model architecture is not in the startup-overlap allowlist",
+            ),
+        )
+        unsupported_reason = next(
+            (reason for unsupported, reason in model_rules if unsupported),
+            None,
+        )
+        if unsupported_reason is not None:
+            return unsupported_reason
+
+        architecture = architectures[0]
+        resolved_model_class, resolved_architecture = get_model_architecture(
+            model_config
+        )
+        if (
+            resolved_architecture != architecture
+            or resolved_model_class is not _get_canonical_model_class(architecture)
+        ):
+            return "the native SGLang model implementation is required"
+        return None
+
+    @property
+    def state(self) -> StartupWeightLoadState:
+        return self._state
+
+    def prepare(self) -> nn.Module:
+        if self._state != StartupWeightLoadState.CREATED:
+            raise RuntimeError(
+                f"Cannot prepare startup weights from state {self._state}"
+            )
+        self._state = StartupWeightLoadState.PREPARING
+        model = self._loader.initialize_model_for_startup(
+            model_config=self._model_config,
+            device_config=self._device_config,
+        )
+        resolved_sources = self._loader.resolve_model_weights(
+            self._model_config,
+            model,
+        )
+        if len(resolved_sources) != 1:
+            raise ValueError(
+                "Startup weight-loading overlap does not support secondary weights"
+            )
+        model = self._loader.prepare_model_for_capture(
+            model=model,
+            model_config=self._model_config,
+        )
+        self._model = model
+        self._resolved_sources = resolved_sources
+        self._capture_ready_at = time.perf_counter()
+        self._state = StartupWeightLoadState.CAPTURE_READY
+        logger.info(
+            "Prepared capture-safe model in %.2f s",
+            self._capture_ready_at - self._created_at,
+        )
+        return model
+
+    def start_prefetch(self) -> None:
+        if self._state != StartupWeightLoadState.CAPTURE_READY:
+            raise RuntimeError(
+                f"Cannot prefetch startup weights from state {self._state}"
+            )
+        assert self._capture_ready_at is not None
+        prefetch_started_at = time.perf_counter()
+        self._prefetch_handle = self._loader.start_checkpoint_prefetch(
+            self._resolved_sources,
+            num_threads=self._options.prefetch_num_threads,
+        )
+        self._prefetch_started_at = prefetch_started_at
+        self._state = StartupWeightLoadState.PREFETCHING
+        logger.info(
+            "Started checkpoint prefetching %.2f s after capture-safe model prep",
+            self._prefetch_started_at - self._capture_ready_at,
+        )
+
+    def finalize(self) -> None:
+        if self._state == StartupWeightLoadState.READY:
+            return
+        if self._state != StartupWeightLoadState.PREFETCHING:
+            raise RuntimeError(
+                f"Cannot finalize startup weights from state {self._state}"
+            )
+        assert self._model is not None
+        assert self._capture_ready_at is not None
+        assert self._prefetch_started_at is not None
+        self._state = StartupWeightLoadState.COMMITTING
+        manifest = ModelStorageManifest.capture(self._model)
+        startup_prefetch_active = self._prepare_prefetch_for_commit()
+        commit_started_at = time.perf_counter()
+        monkey_patch_vllm_parallel_state()
+        self._loader.commit_model_weights(
+            model=self._model,
+            model_config=self._model_config,
+            resolved_sources=self._resolved_sources,
+            target_device=torch.device(self._device_config.device),
+            startup_prefetch_active=startup_prefetch_active,
+        )
+        torch.cuda.synchronize()
+        changed_names = manifest.changed_names(self._model)
+        if changed_names:
+            preview = ", ".join(changed_names[:8])
+            raise RuntimeError(
+                "Startup weight commit changed graph-visible tensor storage: "
+                f"{preview}"
+            )
+        unchanged_names = manifest.unchanged_parameter_names(
+            CAPTURE_SAFE_WEIGHT_SENTINEL
+        )
+        if unchanged_names:
+            preview = ", ".join(unchanged_names[:8])
+            raise RuntimeError(
+                "Startup weight commit did not replace capture-safe dummy values: "
+                f"{preview}"
+            )
+        monkey_patch_vllm_parallel_state(reverse=True)
+        self._stop_prefetch()
+        self._state = StartupWeightLoadState.READY
+        logger.info(
+            "Load weight end. Committed real weights after CUDA graph capture in %.2f s "
+            "(capture overlap window %.2f s, startup overlap total %.2f s)",
+            time.perf_counter() - commit_started_at,
+            commit_started_at - self._prefetch_started_at,
+            time.perf_counter() - self._created_at,
+        )
+
+    def _prepare_prefetch_for_commit(self) -> bool:
+        assert self._prefetch_handle is not None
+        if not self._prefetch_handle.failed:
+            return not self._prefetch_handle.done
+
+        self._prefetch_handle.stop()
+        self._report_prefetch_failure(falling_back=True)
+        return False
+
+    def _stop_prefetch(self) -> None:
+        if self._prefetch_handle is None:
+            return
+        try:
+            if self._prefetch_handle.done:
+                self._prefetch_handle.wait()
+            else:
+                self._prefetch_handle.stop()
+        except TimeoutError:
+            # Only reached after the real weights are committed and validated,
+            # so a stager that outlives its stop timeout must not fail an
+            # otherwise-successful startup. The worker is a daemon thread and
+            # cannot keep the process alive.
+            logger.warning(
+                "Checkpoint prefetch did not stop within its timeout after the "
+                "weight commit; leaving the daemon stager to exit on its own."
+            )
+        self._report_prefetch_failure(falling_back=False)
+        self._prefetch_handle = None
+
+    def _report_prefetch_failure(self, *, falling_back: bool) -> None:
+        handle = self._prefetch_handle
+        if handle is None or not handle.failed or self._prefetch_failure_reported:
+            return
+
+        if handle.errors:
+            path, error = handle.errors[0]
+            failure_detail = (
+                f"{len(handle.errors)} recorded failure(s), first: {path!r}: {error}"
+            )
+        else:
+            failure_detail = "the background worker terminated before completion"
+        action = (
+            "falling back to normal weight loading"
+            if falling_back
+            else "real weight loading completed despite incomplete staging"
+        )
+        logger.warning(
+            "Checkpoint prefetch was incomplete because %s; %s",
+            failure_detail,
+            action,
+        )
+        self._prefetch_failure_reported = True

python/sglang/srt/model_loader/loader.py

@@ -103,6 +103,8 @@ DEFAULT_GPU_MEMORY_FRACTION_FOR_CALIBRATION = (
 )
 from sglang.srt.environ import envs
 from sglang.srt.model_loader.weight_utils import (
+    CheckpointFilePrefetchHandle,
+    _prefetch_all_checkpoints,
     buffered_multi_thread_safetensors_weights_iterator,
     download_safetensors_index_file_from_hf,
     download_weights_from_hf,
@@ -112,6 +114,7 @@ from sglang.srt.model_loader.weight_utils import (
     get_gguf_extra_tensor_names,
     get_quant_config,
     gguf_quant_weights_iterator,
+    initialize_capture_safe_weights,
     initialize_dummy_weights,
     maybe_add_mtp_safetensors,
     multi_thread_pt_weights_iterator,
@@ -410,6 +413,15 @@ class DefaultModelLoader(BaseModelLoader):
                 model_config=model_config,
             )
 
+    @dataclasses.dataclass(frozen=True)
+    class ResolvedSource:
+        """A weight source whose local checkpoint files are already resolved."""
+
+        source: DefaultModelLoader.Source
+        hf_folder: str
+        weight_files: Tuple[str, ...]
+        use_safetensors: bool
+
     counter_before_loading_weights: float = 0.0
     counter_after_loading_weights: float = 0.0
 
@@ -571,22 +583,31 @@ class DefaultModelLoader(BaseModelLoader):
         return hf_folder, hf_weights_files, use_safetensors
 
     def _get_weights_iterator(
-        self, source: Source
+        self,
+        source: Source,
+        *,
+        resolved_source: Optional[ResolvedSource] = None,
+        startup_prefetch_started: bool = False,
+        startup_prefetch_active: bool = False,
     ) -> Generator[Tuple[str, torch.Tensor], None, None]:
         """Get an iterator for the model weights based on the load format."""
         extra_config = self.load_config.model_loader_extra_config
         use_multithread = extra_config.get("enable_multithread_load", True)
-        hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
-            source.model_or_path, source.revision, source.fall_back_to_pt
-        )
-
-        if use_safetensors and source.model_config is not None:
-            hf_weights_files = maybe_add_mtp_safetensors(
-                hf_weights_files,
-                hf_folder,
-                "model.safetensors.index.json",
-                source.model_config.hf_config,
+        if resolved_source is None:
+            hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
+                source.model_or_path, source.revision, source.fall_back_to_pt
             )
+            if use_safetensors and source.model_config is not None:
+                hf_weights_files = maybe_add_mtp_safetensors(
+                    hf_weights_files,
+                    hf_folder,
+                    "model.safetensors.index.json",
+                    source.model_config.hf_config,
+                )
+        else:
+            hf_folder = resolved_source.hf_folder
+            hf_weights_files = list(resolved_source.weight_files)
+            use_safetensors = resolved_source.use_safetensors
 
         if self.load_config.load_format == LoadFormat.NPCACHE:
             # Currently np_cache only support *.bin checkpoints
@@ -599,7 +620,13 @@ class DefaultModelLoader(BaseModelLoader):
             )
         elif use_safetensors:
             weight_loader_disable_mmap = get_model().weight_loader_disable_mmap
-            weight_loader_prefetch = get_model().weight_loader_prefetch_checkpoints
+            configured_prefetch = get_model().weight_loader_prefetch_checkpoints
+            start_iterator_prefetch = (
+                configured_prefetch and not startup_prefetch_started
+            )
+            concurrent_prefetch_active = (
+                startup_prefetch_active or start_iterator_prefetch
+            )
             prefetch_num_threads = get_model().weight_loader_prefetch_num_threads
             weight_loader_drop_cache_after_load = (
                 get_model().weight_loader_drop_cache_after_load
@@ -616,7 +643,7 @@ class DefaultModelLoader(BaseModelLoader):
             # e.g. local NVMe, where prefetch is a no-op and multi-threading
             # helps.
             if (
-                weight_loader_prefetch
+                concurrent_prefetch_active
                 and not weight_loader_disable_mmap
                 and self.load_config.load_format != LoadFormat.FASTSAFETENSORS
                 and use_multithread
@@ -625,7 +652,7 @@ class DefaultModelLoader(BaseModelLoader):
                 )
             ):
                 logger.warning(
-                    "--weight-loader-prefetch-checkpoints is enabled; falling "
+                    "Checkpoint prefetching is active; falling "
                     "back to single-threaded weight loading to avoid I/O "
                     "oversubscription with the prefetch threads. Set "
                     "enable_multithread_load=true in --model-loader-extra-config "
@@ -647,7 +674,7 @@ class DefaultModelLoader(BaseModelLoader):
                         "num_threads", self.DEFAULT_NUM_THREADS
                     ),
                     disable_mmap=weight_loader_disable_mmap,
-                    prefetch=weight_loader_prefetch,
+                    prefetch=start_iterator_prefetch,
                     prefetch_num_threads=prefetch_num_threads,
                     drop_cache_after_load=weight_loader_drop_cache_after_load,
                 )
@@ -655,7 +682,7 @@ class DefaultModelLoader(BaseModelLoader):
                 weights_iterator = safetensors_weights_iterator(
                     hf_weights_files,
                     disable_mmap=weight_loader_disable_mmap,
-                    prefetch=weight_loader_prefetch,
+                    prefetch=start_iterator_prefetch,
                     prefetch_num_threads=prefetch_num_threads,
                     drop_cache_after_load=weight_loader_drop_cache_after_load,
                 )
@@ -716,6 +743,143 @@ class DefaultModelLoader(BaseModelLoader):
         for source in secondary_weights:
             yield from self._get_weights_iterator(source)
 
+    def resolve_model_weights(
+        self,
+        model_config: ModelConfig,
+        model: nn.Module,
+    ) -> Tuple[ResolvedSource, ...]:
+        """Resolve all checkpoint files before background startup prefetching."""
+        sources = [DefaultModelLoader.Source.init_new(model_config, model)]
+        sources.extend(
+            cast(
+                Iterable[DefaultModelLoader.Source],
+                getattr(model, "secondary_weights", ()),
+            )
+        )
+
+        resolved_sources = []
+        for source in sources:
+            hf_folder, weight_files, use_safetensors = self._prepare_weights(
+                source.model_or_path,
+                source.revision,
+                source.fall_back_to_pt,
+            )
+            if use_safetensors and source.model_config is not None:
+                weight_files = maybe_add_mtp_safetensors(
+                    weight_files,
+                    hf_folder,
+                    "model.safetensors.index.json",
+                    source.model_config.hf_config,
+                )
+            resolved_sources.append(
+                DefaultModelLoader.ResolvedSource(
+                    source=source,
+                    hf_folder=hf_folder,
+                    weight_files=tuple(weight_files),
+                    use_safetensors=use_safetensors,
+                )
+            )
+        return tuple(resolved_sources)
+
+    @staticmethod
+    def start_checkpoint_prefetch(
+        resolved_sources: Tuple[ResolvedSource, ...],
+        *,
+        num_threads: int,
+    ) -> CheckpointFilePrefetchHandle:
+        """Start CPU-only page-cache staging for already-resolved sources."""
+        if not all(source.use_safetensors for source in resolved_sources):
+            raise ValueError(
+                "Startup weight-loading overlap requires safetensors checkpoints"
+            )
+        weight_files = sorted(
+            {path for source in resolved_sources for path in source.weight_files}
+        )
+        return _prefetch_all_checkpoints(weight_files, num_threads=num_threads)
+
+    def initialize_model_for_startup(
+        self,
+        *,
+        model_config: ModelConfig,
+        device_config: DeviceConfig,
+    ) -> nn.Module:
+        """Build the final model structure and GPU parameter storage."""
+        target_device = torch.device(device_config.device)
+        quant_config = _get_quantization_config(model_config, self.load_config)
+        with set_default_torch_dtype(model_config.dtype):
+            with target_device:
+                model = _initialize_model(
+                    model_config,
+                    self.load_config,
+                    quant_config,
+                )
+        return model
+
+    def prepare_model_for_capture(
+        self,
+        *,
+        model: nn.Module,
+        model_config: ModelConfig,
+    ) -> nn.Module:
+        """Initialize final storage with values safe for graph warmup.
+
+        Mirrors the post-initialization sequence of ``DummyModelLoader``, except
+        that parameters are filled with a detectable sentinel instead of random
+        values so ``commit_model_weights`` can prove every one of them was
+        replaced.
+
+        Note that this runs ``process_weights_after_loading`` on the sentinel
+        values, and ``commit_model_weights`` runs it again on the real weights,
+        so overlap invokes it once more than the serial path. That is safe for
+        the currently supported matrix, where the CUDA unquantized path is a
+        no-op, and it is not covered by the storage manifest, which proves
+        tensor identity rather than idempotence. Any quantization method that
+        mutates weights in place therefore has to be evaluated here before its
+        configuration is added to the supported set.
+        """
+        with set_default_torch_dtype(model_config.dtype):
+            initialize_capture_safe_weights(model)
+            _post_load_weights(model)
+            for _, module in model.named_modules():
+                quant_method = getattr(module, "quant_method", None)
+                if quant_method is None:
+                    continue
+                if (
+                    hasattr(module, "is_weights_quantized")
+                    and module.is_weights_quantized()
+                ):
+                    continue
+                quant_method.process_weights_after_loading(module)
+        return model.eval()
+
+    def commit_model_weights(
+        self,
+        *,
+        model: nn.Module,
+        model_config: ModelConfig,
+        resolved_sources: Tuple[ResolvedSource, ...],
+        target_device: torch.device,
+        startup_prefetch_active: bool,
+    ) -> None:
+        """Load real checkpoint values into a capture-ready model."""
+
+        def weights_iterator():
+            for resolved_source in resolved_sources:
+                yield from self._get_weights_iterator(
+                    resolved_source.source,
+                    resolved_source=resolved_source,
+                    startup_prefetch_started=True,
+                    startup_prefetch_active=startup_prefetch_active,
+                )
+
+        with set_default_torch_dtype(model_config.dtype):
+            self.load_weights_and_postprocess(
+                model,
+                weights_iterator(),
+                target_device,
+            )
+        self.counter_after_loading_weights = time.perf_counter()
+
     def download_model(self, model_config: ModelConfig) -> None:
         self._prepare_weights(
             model_config.model_path, model_config.revision, fall_back_to_pt=True

python/sglang/srt/model_loader/weight_utils.py

@@ -16,6 +16,7 @@ import os
 import re
 import struct
 import tempfile
+import threading
 from collections import defaultdict
 from pathlib import Path
 from typing import (
@@ -131,6 +132,8 @@ def probe_routed_expert_weight_dtype(model_path: str) -> Optional[str]:
 
 # Block size for sequential checkpoint prefetch reads (page cache warming).
 _PREFETCH_BLOCK_SIZE = None
+_PREFETCH_STOP_TIMEOUT_SECONDS = 60.0
+CAPTURE_SAFE_WEIGHT_SENTINEL = 1e-3
 
 
 def _get_prefetch_block_size() -> int:
@@ -856,21 +859,72 @@ def np_cache_weights_iterator(
         yield name, torch.from_numpy(param)
 
 
-def _prefetch_checkpoint_file(file_path: str) -> None:
+def _prefetch_checkpoint_file(
+    file_path: str,
+    cancel_event: Optional[threading.Event] = None,
+) -> None:
     """Prefetch a checkpoint file into the OS page cache.
 
     Reads the file sequentially in 16 MB blocks so the kernel caches its pages
     before workers load the same file via mmap.
     """
     with open(file_path, "rb") as f:
-        while f.read(_get_prefetch_block_size()):
-            pass
+        while cancel_event is None or not cancel_event.is_set():
+            if not f.read(_get_prefetch_block_size()):
+                break
+
+
+class CheckpointFilePrefetchHandle:
+    """Lifecycle handle for background checkpoint page-cache prefetching."""
+
+    def __init__(
+        self,
+        *,
+        thread: threading.Thread,
+        cancel_event: threading.Event,
+        succeeded_event: threading.Event,
+        errors: List[Tuple[str, Exception]],
+    ) -> None:
+        self._thread = thread
+        self._cancel_event = cancel_event
+        self._succeeded_event = succeeded_event
+        self._errors = errors
+
+    def wait(self, timeout: Optional[float] = None) -> None:
+        self._thread.join(timeout)
+        if self._thread.is_alive():
+            raise TimeoutError("Timed out waiting for checkpoint prefetching")
+
+    def cancel(self) -> None:
+        """Stop scheduling shards and interrupt reads at the next block."""
+        self._cancel_event.set()
+
+    def stop(self, timeout: Optional[float] = _PREFETCH_STOP_TIMEOUT_SECONDS) -> None:
+        """Cancel prefetching and wait for the background worker to finish."""
+        self.cancel()
+        self.wait(timeout)
+
+    @property
+    def done(self) -> bool:
+        return not self._thread.is_alive()
+
+    @property
+    def failed(self) -> bool:
+        return bool(self._errors) or (self.done and not self._succeeded_event.is_set())
+
+    @property
+    def cancelled(self) -> bool:
+        return self._cancel_event.is_set()
+
+    @property
+    def errors(self) -> Tuple[Tuple[str, Exception], ...]:
+        return tuple(self._errors)
 
 
 def _prefetch_all_checkpoints(
     sorted_files: List[str],
     num_threads: int = 4,
-) -> None:
+) -> CheckpointFilePrefetchHandle:
     """Start prefetching checkpoint files into page cache in a background thread.
 
     When multiple ranks on the same node load the same checkpoint (e.g.
@@ -886,7 +940,6 @@ def _prefetch_all_checkpoints(
     naturally adapts to any RAM size — even if the full checkpoint does
     not fit in page cache, the prefetch thread stays ahead of the loader.
     """
-    import threading
     import time
 
     if num_threads < 1:
@@ -905,6 +958,9 @@ def _prefetch_all_checkpoints(
 
     my_files = sorted_files[local_rank::local_world_size]
     total_for_rank = len(my_files)
+    cancel_event = threading.Event()
+    succeeded_event = threading.Event()
+    errors: List[Tuple[str, Exception]] = []
 
     logger.info(
         "Rank %d: prefetching %d/%d checkpoint shards into page cache "
@@ -941,7 +997,11 @@ def _prefetch_all_checkpoints(
             pending: Dict[concurrent.futures.Future, str] = {}
 
             for path in itertools.islice(file_iter, num_threads):
-                pending[executor.submit(_prefetch_checkpoint_file, path)] = path
+                if cancel_event.is_set():
+                    break
+                pending[
+                    executor.submit(_prefetch_checkpoint_file, path, cancel_event)
+                ] = path
 
             while pending:
                 done, _ = concurrent.futures.wait(
@@ -950,35 +1010,46 @@ def _prefetch_all_checkpoints(
                 )
                 for future in done:
                     path = pending.pop(future)
-                    try:
-                        future.result()
-                    except Exception:
+                    exc = future.exception()
+                    if exc is not None:
+                        errors.append((path, exc))
                         logger.warning(
-                            "Failed to prefetch checkpoint file %r.",
+                            "Failed to prefetch checkpoint file %r: %s",
                             path,
-                            exc_info=True,
+                            exc,
                         )
-                    finally:
-                        record_complete()
+                    record_complete()
 
-                    next_path = next(file_iter, None)
+                    next_path = None if cancel_event.is_set() else next(file_iter, None)
                     if next_path is not None:
                         pending[
-                            executor.submit(_prefetch_checkpoint_file, next_path)
+                            executor.submit(
+                                _prefetch_checkpoint_file,
+                                next_path,
+                                cancel_event,
+                            )
                         ] = next_path
 
     def _run_prefetch() -> None:
         start = time.perf_counter()
         _prefetch_all()
-        elapsed = time.perf_counter() - start
+        succeeded_event.set()
         logger.info(
             "Rank %d: prefetching checkpoint files into page cache "
             "finished in %.2fs",
             local_rank,
-            elapsed,
+            time.perf_counter() - start,
         )
 
-    threading.Thread(target=_run_prefetch, daemon=True).start()
+    thread = threading.Thread(target=_run_prefetch, daemon=True)
+    handle = CheckpointFilePrefetchHandle(
+        thread=thread,
+        cancel_event=cancel_event,
+        succeeded_event=succeeded_event,
+        errors=errors,
+    )
+    thread.start()
+    return handle
 
 
 def _drop_file_cache_after_load(path: str) -> None:
@@ -1550,6 +1621,21 @@ def set_runai_streamer_env(load_config: LoadConfig):
         os.environ["RUNAI_STREAMER_S3_ENDPOINT"] = aws_endpoint_url
 
 
+@torch.no_grad()
+def initialize_capture_safe_weights(
+    model: torch.nn.Module,
+    value: float = CAPTURE_SAFE_WEIGHT_SENTINEL,
+) -> None:
+    """Fill floating-point parameters with finite values for graph warmup.
+
+    Persistent buffers are intentionally left intact: unlike parameters, they
+    are not guaranteed to be replaced by ``model.load_weights()``.
+    """
+    for param in model.parameters():
+        if torch.is_floating_point(param):
+            param.fill_(value)
+
+
 def initialize_dummy_weights(
     model: torch.nn.Module,
     low: float = -1e-3,

python/sglang/srt/server_args.py

@@ -3172,6 +3172,16 @@ class ServerArgs:
     # -------------------------------------------------------------------------
     # Model weight update and weight loading
     # -------------------------------------------------------------------------
+    startup_weight_load_mode: A[
+        Literal["serial", "overlap"],
+        (
+            "Control startup weight loading relative to CUDA graph capture. "
+            "'serial' preserves the existing startup order; 'overlap' stages "
+            "checkpoint files while CUDA graphs are captured and commits the "
+            "real weights afterward."
+        ),
+        NS("model"),
+    ] = "serial"
     custom_weight_loader: A[
         Optional[List[str]],
         Arg(
@@ -8811,6 +8821,10 @@ class ServerArgs:
     def is_ep_scale_joiner(self) -> bool:
         return self.ep_join_mode == "scale"
 
+    @property
+    def is_startup_weight_load_overlap(self) -> bool:
+        return self.startup_weight_load_mode == "overlap"
+
     def ssl_verify(self):
         """Return the value for the requests library's verify= parameter.
 

test/registered/model_loading/test_startup_weight_load.py

@@ -0,0 +1,69 @@
+"""End-to-end parity test for post-capture startup weight loading."""
+
+import unittest
+
+import sglang as sgl
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    CustomTestCase,
+)
+
+register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-small")
+
+
+class TestStartupWeightLoad(CustomTestCase):
+    @staticmethod
+    def _generate(startup_weight_load_mode=None):
+        kwargs = dict(
+            model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+            dtype="bfloat16",
+            random_seed=42,
+            cuda_graph_max_bs_decode=1,
+            max_total_tokens=256,
+        )
+        if startup_weight_load_mode is not None:
+            kwargs["startup_weight_load_mode"] = startup_weight_load_mode
+
+        with sgl.Engine(**kwargs) as engine:
+            return engine.generate(
+                "The capital of France is",
+                sampling_params={
+                    "temperature": 0,
+                    "max_new_tokens": 8,
+                    "ignore_eos": True,
+                },
+                return_logprob=True,
+                logprob_start_len=0,
+            )
+
+    def test_overlap_matches_default_serial_startup(self):
+        # Omitting the flag is intentional: it pins the merge-safe default path.
+        serial = self._generate()
+        overlap = self._generate("overlap")
+
+        self.assertEqual(serial["output_ids"], overlap["output_ids"])
+        self.assertEqual(serial["text"], overlap["text"])
+
+        serial_logprobs = serial["meta_info"]["output_token_logprobs"]
+        overlap_logprobs = overlap["meta_info"]["output_token_logprobs"]
+        self.assertEqual(len(serial_logprobs), len(overlap_logprobs))
+        self.assertGreater(len(serial_logprobs), 0)
+        for index, (serial_item, overlap_item) in enumerate(
+            zip(serial_logprobs, overlap_logprobs)
+        ):
+            self.assertEqual(
+                serial_item[1],
+                overlap_item[1],
+                f"token id differs at output position {index}",
+            )
+            self.assertAlmostEqual(
+                serial_item[0],
+                overlap_item[0],
+                delta=1e-5,
+                msg=f"logprob differs at output position {index}",
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py

@@ -220,6 +220,20 @@ class TestMlxExtendRouting(CustomTestCase):
         worker._mlx_pool_initialized = True
         return worker
 
+    def test_startup_weight_overlap_is_rejected_before_mlx_model_load(self):
+        from sglang.srt.hardware_backend.mlx.model_runner_stub import (
+            MlxModelRunnerStub,
+        )
+        from sglang.srt.hardware_backend.mlx.tp_worker import MlxTpModelWorker
+
+        worker = MlxTpModelWorker.__new__(MlxTpModelWorker)
+        worker.server_args = SimpleNamespace(is_startup_weight_load_overlap=True)
+
+        with self.assertRaisesRegex(ValueError, "CUDA only"):
+            MlxModelRunnerStub.validate_startup_weight_load_mode(worker.server_args)
+        with self.assertRaisesRegex(ValueError, "CUDA only"):
+            worker._init_model_runner()
+
     # ---------- the shared decision helper ----------
     # The helper takes no seq_len: length cannot distinguish a 1-token
     # continuation from a genuine decode -- request state does.

test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py

@@ -0,0 +1,732 @@
+"""Unit tests for the post-capture startup weight-loading component."""
+
+import dataclasses
+import re
+import unittest
+from types import SimpleNamespace
+from unittest.mock import call, patch
+
+import torch
+from torch import nn
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
+
+maybe_stub_sgl_kernel()
+
+from sglang.srt.configs.device_config import DeviceConfig
+from sglang.srt.configs.load_config import LoadConfig, LoadFormat
+from sglang.srt.configs.model_config import ModelImpl
+from sglang.srt.managers.tp_worker import TpModelWorker
+from sglang.srt.model_executor.cuda_graph_config import Backend
+from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.model_executor.model_runner_components.startup_weight_load import (
+    ModelStorageManifest,
+    StartupWeightLoadManager,
+    StartupWeightLoadOptions,
+    StartupWeightLoadState,
+)
+from sglang.srt.model_loader.loader import DefaultModelLoader
+from sglang.srt.model_loader.weight_utils import initialize_capture_safe_weights
+from sglang.srt.runtime_context import get_context
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+_STARTUP_MODULE = (
+    "sglang.srt.model_executor.model_runner_components.startup_weight_load"
+)
+
+
+class _CanonicalModel:
+    pass
+
+
+class _ExternalModel:
+    pass
+
+
+def _make_options(**overrides):
+    options = StartupWeightLoadOptions(
+        device="cuda",
+        is_cuda_platform=True,
+        cuda_graph_enabled=True,
+        prefill_cuda_graph_backend=Backend.FULL,
+        is_draft_worker=False,
+        speculative_algorithm=None,
+        tp_size=1,
+        attn_cp_size=1,
+        dcp_size=1,
+        pp_size=1,
+        dp_size=1,
+        ep_size=1,
+        cpu_offload_gb=0,
+        offload_group_size=-1,
+        enable_memory_saver=False,
+        enable_weights_cpu_backup=False,
+        torchao_config="",
+        enable_lora=False,
+        has_lora_paths=False,
+        weight_loader_disable_mmap=False,
+        weight_loader_drop_cache_after_load=False,
+        has_custom_weight_loader=False,
+        enable_torch_compile=False,
+        prefetch_num_threads=4,
+    )
+    return dataclasses.replace(options, **overrides)
+
+
+def _make_model_config(**overrides):
+    values = dict(
+        hf_config=SimpleNamespace(architectures=["LlamaForCausalLM"]),
+        dtype=torch.bfloat16,
+        quantization=None,
+        modelopt_quant=None,
+        is_multimodal=False,
+        is_generation=True,
+        model_impl=ModelImpl.SGLANG,
+        _resolved_model_impl=ModelImpl.SGLANG,
+    )
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+class _RecordingPrefetchHandle:
+    def __init__(self, trace, *, done=False, errors=()):
+        self._trace = trace
+        self.done = done
+        self.errors = errors
+
+    @property
+    def failed(self):
+        return bool(self.errors)
+
+    def wait(self, timeout=None):
+        self._trace.append("wait_prefetch")
+
+    def stop(self, timeout=None):
+        self._trace.append("stop_prefetch")
+        self.wait()
+        self.done = True
+
+
+class _RecordingLoader:
+    def __init__(self, model, trace):
+        self._model = model
+        self._trace = trace
+        self.prefetch_handle = _RecordingPrefetchHandle(trace)
+
+    def initialize_model_for_startup(self, *, model_config, device_config):
+        self._trace.append("initialize")
+        return self._model
+
+    def resolve_model_weights(self, model_config, model):
+        self._trace.append("resolve")
+        return (object(),)
+
+    def start_checkpoint_prefetch(self, resolved_sources, *, num_threads):
+        self._trace.append("start_prefetch")
+        return self.prefetch_handle
+
+    def prepare_model_for_capture(self, *, model, model_config):
+        self._trace.append("prepare_capture")
+        return model
+
+    def commit_model_weights(
+        self,
+        *,
+        model,
+        model_config,
+        resolved_sources,
+        target_device,
+        startup_prefetch_active,
+    ):
+        self._trace.append("commit")
+        self.startup_prefetch_active = startup_prefetch_active
+        with torch.no_grad():
+            for parameter in model.parameters():
+                parameter.fill_(3)
+
+
+class _TiedWeightModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(2, 2))
+        self.tied_weight = self.weight
+        self.register_buffer("scale", torch.ones(2))
+
+
+class TestStartupWeightLoadSelector(CustomTestCase):
+    def setUp(self):
+        self.load_config = LoadConfig(load_format=LoadFormat.SAFETENSORS)
+        self.loader = DefaultModelLoader(self.load_config)
+        self.device_config = DeviceConfig("cuda", 0)
+
+    def _create(
+        self,
+        *,
+        options=None,
+        model_config=None,
+        load_config=None,
+        loader=None,
+        resolved_model_class=None,
+    ):
+        model_config = _make_model_config() if model_config is None else model_config
+        architecture = model_config.hf_config.architectures[0]
+        with (
+            patch(
+                f"{_STARTUP_MODULE}.get_model_architecture",
+                return_value=(
+                    resolved_model_class or _CanonicalModel,
+                    architecture,
+                ),
+            ),
+            patch(
+                f"{_STARTUP_MODULE}._get_canonical_model_class",
+                return_value=_CanonicalModel,
+            ),
+        ):
+            return StartupWeightLoadManager.create(
+                loader=self.loader if loader is None else loader,
+                model_config=model_config,
+                load_config=self.load_config if load_config is None else load_config,
+                device_config=self.device_config,
+                options=_make_options() if options is None else options,
+            )
+
+    def test_supported_overlap_creates_a_manager(self):
+        self.assertIsInstance(self._create(), StartupWeightLoadManager)
+        self.assertIsInstance(
+            self._create(options=_make_options(tp_size=2)),
+            StartupWeightLoadManager,
+        )
+
+    def test_unsupported_overlap_is_rejected_instead_of_falling_back(self):
+        cases = (
+            (
+                "non_cuda",
+                dict(options=_make_options(device="cpu", is_cuda_platform=False)),
+                "CUDA only",
+            ),
+            (
+                "graphs_disabled",
+                dict(options=_make_options(cuda_graph_enabled=False)),
+                "CUDA graph capture is disabled",
+            ),
+            (
+                "tc_piecewise_prefill",
+                dict(
+                    options=_make_options(
+                        prefill_cuda_graph_backend=Backend.TC_PIECEWISE
+                    )
+                ),
+                "tc_piecewise prefill CUDA graphs are not supported",
+            ),
+            (
+                "pt_checkpoint",
+                dict(load_config=LoadConfig(load_format=LoadFormat.PT)),
+                "load format must be auto or safetensors",
+            ),
+            (
+                "draft_worker",
+                dict(options=_make_options(is_draft_worker=True)),
+                "draft workers are not supported",
+            ),
+            (
+                "draft_model_checkpoint",
+                dict(
+                    load_config=LoadConfig(
+                        load_format=LoadFormat.SAFETENSORS,
+                        draft_model_idx=0,
+                    )
+                ),
+                "draft model loading is unsupported",
+            ),
+            (
+                "speculative_decoding",
+                dict(options=_make_options(speculative_algorithm="EAGLE")),
+                "speculative decoding is not supported",
+            ),
+            (
+                "tp3",
+                dict(options=_make_options(tp_size=3)),
+                "only TP1 and TP2 are supported",
+            ),
+            (
+                "attention_context_parallel",
+                dict(options=_make_options(tp_size=2, attn_cp_size=2)),
+                "attention context parallelism is not supported",
+            ),
+            (
+                "decode_context_parallel",
+                dict(options=_make_options(tp_size=2, dcp_size=2)),
+                "decode context parallelism is not supported",
+            ),
+            (
+                "quantized_model",
+                dict(model_config=_make_model_config(quantization="fp8")),
+                "quantization is not supported",
+            ),
+            (
+                "layer_group_offload",
+                dict(options=_make_options(offload_group_size=1)),
+                "layer-group offloading is not supported",
+            ),
+            (
+                "torch_compile",
+                dict(options=_make_options(enable_torch_compile=True)),
+                "torch.compile is not supported",
+            ),
+            (
+                "transformers_model_impl",
+                dict(
+                    model_config=_make_model_config(
+                        model_impl=ModelImpl.TRANSFORMERS,
+                        _resolved_model_impl=ModelImpl.TRANSFORMERS,
+                    ),
+                    resolved_model_class=_ExternalModel,
+                ),
+                "the native SGLang model implementation is required",
+            ),
+            (
+                "external_model_implementation",
+                dict(resolved_model_class=_ExternalModel),
+                "the native SGLang model implementation is required",
+            ),
+            (
+                "unknown_architecture",
+                dict(
+                    model_config=_make_model_config(
+                        hf_config=SimpleNamespace(architectures=["OtherForCausalLM"])
+                    )
+                ),
+                "model architecture is not in the startup-overlap allowlist",
+            ),
+        )
+        for name, kwargs, reason in cases:
+            with self.subTest(name=name):
+                with self.assertRaisesRegex(ValueError, re.escape(reason)):
+                    self._create(**kwargs)
+
+
+class TestStartupWeightLoadManager(CustomTestCase):
+    def _manager(self, loader):
+        return StartupWeightLoadManager(
+            loader=loader,
+            model_config=_make_model_config(),
+            device_config=DeviceConfig("cpu", 0),
+            options=_make_options(),
+        )
+
+    def test_prepare_capture_finalize_state_and_order(self):
+        trace = []
+        model = _TiedWeightModel()
+        manager = self._manager(_RecordingLoader(model, trace))
+
+        self.assertEqual(manager.state, StartupWeightLoadState.CREATED)
+        self.assertIs(manager.prepare(), model)
+        self.assertEqual(manager.state, StartupWeightLoadState.CAPTURE_READY)
+        manager.start_prefetch()
+        self.assertEqual(manager.state, StartupWeightLoadState.PREFETCHING)
+
+        # CUDA graph capture is owned by Scheduler and occurs between these calls.
+        trace.append("capture")
+        with (
+            patch(
+                f"{_STARTUP_MODULE}.monkey_patch_vllm_parallel_state"
+            ) as parallel_state_patch,
+            patch(f"{_STARTUP_MODULE}.torch.cuda.synchronize"),
+            patch(f"{_STARTUP_MODULE}.logger.info") as log_info,
+        ):
+            manager.finalize()
+
+        self.assertEqual(manager.state, StartupWeightLoadState.READY)
+        self.assertEqual(
+            trace,
+            [
+                "initialize",
+                "resolve",
+                "prepare_capture",
+                "start_prefetch",
+                "capture",
+                "commit",
+                "stop_prefetch",
+                "wait_prefetch",
+            ],
+        )
+
+        # Finalization is idempotent after a successful commit.
+        manager.finalize()
+        self.assertEqual(trace.count("commit"), 1)
+        self.assertIs(model.weight, model.tied_weight)
+        torch.testing.assert_close(model.weight, torch.full_like(model.weight, 3))
+        self.assertTrue(log_info.call_args.args[0].startswith("Load weight end."))
+        self.assertTrue(manager._loader.startup_prefetch_active)
+        self.assertEqual(
+            parallel_state_patch.call_args_list,
+            [call(), call(reverse=True)],
+        )
+
+    def test_finalize_rejects_graph_visible_storage_rebind(self):
+        trace = []
+        model = _TiedWeightModel()
+        loader = _RecordingLoader(model, trace)
+
+        def rebind_tied_weight(**kwargs):
+            trace.append("commit")
+            model.tied_weight = nn.Parameter(model.tied_weight.detach().clone())
+
+        loader.commit_model_weights = rebind_tied_weight
+        manager = self._manager(loader)
+        manager.prepare()
+        manager.start_prefetch()
+
+        with (
+            patch(f"{_STARTUP_MODULE}.monkey_patch_vllm_parallel_state"),
+            patch(f"{_STARTUP_MODULE}.torch.cuda.synchronize"),
+            self.assertRaisesRegex(
+                RuntimeError,
+                "changed graph-visible tensor storage: parameter:tied_weight",
+            ),
+        ):
+            manager.finalize()
+
+    def test_finalize_rejects_parameter_left_at_capture_sentinel(self):
+        trace = []
+        model = _TiedWeightModel()
+        loader = _RecordingLoader(model, trace)
+
+        def skip_commit(**kwargs):
+            trace.append("commit")
+
+        loader.commit_model_weights = skip_commit
+        manager = self._manager(loader)
+        manager.prepare()
+        with torch.no_grad():
+            model.weight.fill_(1e-3)
+        manager.start_prefetch()
+
+        with (
+            patch(f"{_STARTUP_MODULE}.monkey_patch_vllm_parallel_state"),
+            patch(f"{_STARTUP_MODULE}.torch.cuda.synchronize"),
+            self.assertRaisesRegex(
+                RuntimeError,
+                "did not replace capture-safe dummy values: parameter:tied_weight",
+            ),
+        ):
+            manager.finalize()
+
+    def test_completed_prefetch_restores_normal_loader(self):
+        trace = []
+        model = _TiedWeightModel()
+        loader = _RecordingLoader(model, trace)
+        loader.prefetch_handle.done = True
+        manager = self._manager(loader)
+        manager.prepare()
+        manager.start_prefetch()
+
+        with (
+            patch(f"{_STARTUP_MODULE}.monkey_patch_vllm_parallel_state"),
+            patch(f"{_STARTUP_MODULE}.torch.cuda.synchronize"),
+        ):
+            manager.finalize()
+
+        self.assertFalse(loader.startup_prefetch_active)
+        self.assertIn("wait_prefetch", trace)
+        self.assertNotIn("stop_prefetch", trace)
+
+    def test_failed_prefetch_falls_back_and_logs_summary(self):
+        trace = []
+        model = _TiedWeightModel()
+        loader = _RecordingLoader(model, trace)
+        loader.prefetch_handle.errors = (("bad.safetensors", OSError("failed")),)
+        manager = self._manager(loader)
+        manager.prepare()
+        manager.start_prefetch()
+
+        with (
+            patch(f"{_STARTUP_MODULE}.monkey_patch_vllm_parallel_state"),
+            patch(f"{_STARTUP_MODULE}.torch.cuda.synchronize"),
+            patch(f"{_STARTUP_MODULE}.logger.warning") as warning,
+        ):
+            manager.finalize()
+
+        self.assertFalse(loader.startup_prefetch_active)
+        warning.assert_called_once()
+        self.assertIn("falling back", warning.call_args.args[2])
+
+    def test_stop_timeout_after_commit_does_not_fail_startup(self):
+        trace = []
+        model = _TiedWeightModel()
+        loader = _RecordingLoader(model, trace)
+
+        def _stop_times_out(timeout=None):
+            trace.append("stop_prefetch")
+            raise TimeoutError("Timed out waiting for checkpoint prefetching")
+
+        loader.prefetch_handle.stop = _stop_times_out
+        manager = self._manager(loader)
+        manager.prepare()
+        manager.start_prefetch()
+
+        with (
+            patch(f"{_STARTUP_MODULE}.monkey_patch_vllm_parallel_state"),
+            patch(f"{_STARTUP_MODULE}.torch.cuda.synchronize"),
+            patch(f"{_STARTUP_MODULE}.logger.warning") as warning,
+        ):
+            manager.finalize()
+
+        self.assertEqual(manager.state, StartupWeightLoadState.READY)
+        self.assertIn("stop_prefetch", trace)
+        warning.assert_called_once()
+        self.assertIn("did not stop within its timeout", warning.call_args.args[0])
+
+    def test_start_prefetch_requires_capture_ready_and_starts_once(self):
+        trace = []
+        manager = self._manager(_RecordingLoader(nn.Linear(2, 2), trace))
+
+        with self.assertRaisesRegex(RuntimeError, "from state"):
+            manager.start_prefetch()
+
+        manager.prepare()
+        manager.start_prefetch()
+        self.assertEqual(manager.state, StartupWeightLoadState.PREFETCHING)
+
+        with self.assertRaisesRegex(RuntimeError, "from state"):
+            manager.start_prefetch()
+        self.assertEqual(trace.count("start_prefetch"), 1)
+
+
+class TestModelStorageManifest(CustomTestCase):
+    def test_in_place_updates_preserve_the_manifest(self):
+        model = _TiedWeightModel()
+        manifest = ModelStorageManifest.capture(model)
+
+        with torch.no_grad():
+            model.weight.fill_(2)
+            model.scale.fill_(3)
+
+        self.assertEqual(manifest.changed_names(model), ())
+
+    def test_manifest_keeps_strong_tensor_references(self):
+        model = _TiedWeightModel()
+        manifest = ModelStorageManifest.capture(model)
+
+        metadata = dict(manifest.tensors)["parameter:weight"]
+        self.assertIs(metadata.tensor, model.weight)
+
+    def test_capture_sentinel_check_ignores_buffers(self):
+        model = _TiedWeightModel()
+        with torch.no_grad():
+            model.weight.fill_(1e-3)
+            model.scale.fill_(1e-3)
+        manifest = ModelStorageManifest.capture(model)
+
+        self.assertEqual(
+            manifest.unchanged_parameter_names(1e-3),
+            ("parameter:tied_weight",),
+        )
+
+    def test_parameter_rebind_and_alias_break_are_detected(self):
+        model = _TiedWeightModel()
+        manifest = ModelStorageManifest.capture(model)
+
+        model.tied_weight = nn.Parameter(model.tied_weight.detach().clone())
+
+        self.assertEqual(
+            manifest.changed_names(model),
+            ("parameter:tied_weight",),
+        )
+
+
+class TestCaptureSafeWeightInitialization(CustomTestCase):
+    def test_only_parameters_are_filled(self):
+        model = _TiedWeightModel()
+
+        initialize_capture_safe_weights(model, value=0.125)
+
+        torch.testing.assert_close(model.weight, torch.full_like(model.weight, 0.125))
+        torch.testing.assert_close(model.scale, torch.ones_like(model.scale))
+
+
+class _LifecycleRunner:
+    def __init__(self, name, trace):
+        self._name = name
+        self._trace = trace
+
+    def start_startup_weight_load(self):
+        self._trace.append(f"start:{self._name}")
+
+    def finalize_startup_weight_load(self):
+        self._trace.append(f"finalize:{self._name}")
+
+
+class TestStartupWeightLoadFanout(CustomTestCase):
+    def test_primary_and_multi_runner_extras_are_started_once(self):
+        trace = []
+        primary = _LifecycleRunner("primary", trace)
+        extra_1 = _LifecycleRunner("extra_1", trace)
+        extra_2 = _LifecycleRunner("extra_2", trace)
+        worker = TpModelWorker.__new__(TpModelWorker)
+        worker._model_runner = primary
+        worker.model_runner_list = [primary, extra_1, extra_2]
+
+        worker.start_startup_weight_load()
+
+        self.assertEqual(
+            trace,
+            ["start:primary", "start:extra_1", "start:extra_2"],
+        )
+
+    def test_primary_and_multi_runner_extras_are_finalized_once(self):
+        for multi_runner in (False, True):
+            with self.subTest(multi_runner=multi_runner):
+                trace = []
+                primary = _LifecycleRunner("primary", trace)
+                extra_1 = _LifecycleRunner("extra_1", trace)
+                extra_2 = _LifecycleRunner("extra_2", trace)
+                worker = TpModelWorker.__new__(TpModelWorker)
+                worker._model_runner = primary
+                worker.model_runner_list = (
+                    [primary, extra_1, extra_2] if multi_runner else []
+                )
+
+                worker.finalize_startup_weight_load()
+
+                self.assertEqual(
+                    trace,
+                    (
+                        ["finalize:primary", "finalize:extra_1", "finalize:extra_2"]
+                        if multi_runner
+                        else ["finalize:primary"]
+                    ),
+                )
+
+
+class _RunnerStartupManager:
+    def __init__(self, trace):
+        self._trace = trace
+
+    def start_prefetch(self):
+        self._trace.append("start_prefetch")
+
+    def finalize(self):
+        self._trace.append("finalize")
+
+
+class TestModelRunnerStartupWeightLoadOwnership(CustomTestCase):
+    @staticmethod
+    def _runner(manager):
+        runner = ModelRunner.__new__(ModelRunner)
+        runner.startup_weight_load = manager
+        runner.server_args = SimpleNamespace(
+            elastic_ep_backend=None,
+            is_ep_joiner=False,
+        )
+        runner.ps = SimpleNamespace(tp_rank=0)
+        return runner
+
+    def test_start_delegates_to_the_manager(self):
+        trace = []
+        runner = self._runner(_RunnerStartupManager(trace))
+
+        runner.start_startup_weight_load()
+
+        self.assertEqual(trace, ["start_prefetch"])
+
+    def test_success_releases_ownership_after_the_barrier(self):
+        trace = []
+        manager = _RunnerStartupManager(trace)
+        runner = self._runner(manager)
+
+        def barrier(**kwargs):
+            self.assertIs(runner.startup_weight_load, manager)
+            trace.append("barrier")
+
+        with (
+            patch(
+                "sglang.srt.model_executor.model_runner.dist_barrier_after_load",
+                side_effect=barrier,
+            ),
+            get_context().override_server_args(),
+        ):
+            runner.finalize_startup_weight_load()
+
+        self.assertEqual(trace, ["finalize", "barrier"])
+        self.assertIsNone(runner.startup_weight_load)
+
+
+class _SchedulerWorker:
+    def __init__(self, trace, *, post_capture_active=False):
+        self._trace = trace
+        self.model_runner = SimpleNamespace(
+            token_to_kv_pool=SimpleNamespace(post_capture_active=post_capture_active),
+            post_capture_resize_kv_pool=lambda: trace.append("resize"),
+        )
+
+    def start_startup_weight_load(self):
+        self._trace.append("start")
+
+    def finalize_startup_weight_load(self):
+        self._trace.append("finalize")
+
+
+class TestStartupWeightLoadSchedulerRouting(CustomTestCase):
+    @staticmethod
+    def _scheduler(worker, trace, *, mode):
+        from sglang.srt.managers.scheduler import Scheduler
+
+        scheduler = Scheduler.__new__(Scheduler)
+        scheduler.server_args = SimpleNamespace(
+            is_startup_weight_load_overlap=mode == "overlap"
+        )
+        scheduler.init_tp_model_worker = lambda: setattr(scheduler, "tp_worker", worker)
+        scheduler.maybe_init_draft_worker = lambda: setattr(
+            scheduler, "draft_worker", None
+        )
+        scheduler.init_memory_pools = lambda: trace.append("memory_pool")
+        scheduler.init_all_attention_backends = lambda: trace.append("attention")
+        scheduler.init_all_cuda_graphs = lambda: trace.append("capture")
+        return scheduler
+
+    def _run_startup(self, mode):
+        trace = []
+        worker = _SchedulerWorker(trace, post_capture_active=True)
+        scheduler = self._scheduler(worker, trace, mode=mode)
+
+        def stop_after_startup():
+            raise RuntimeError("stop after startup")
+
+        scheduler.spec_algorithm = SimpleNamespace(is_none=stop_after_startup)
+
+        with (
+            patch(
+                "sglang.srt.managers.scheduler.get_exec",
+                return_value=SimpleNamespace(
+                    moe=SimpleNamespace(
+                        elastic_ep_backend=None,
+                        ep_join_mode=None,
+                    )
+                ),
+            ),
+            self.assertRaisesRegex(RuntimeError, "stop after startup"),
+        ):
+            scheduler.init_model_worker()
+
+        return trace
+
+    def test_serial_path_skips_overlap_hooks(self):
+        self.assertEqual(
+            self._run_startup("serial"),
+            ["memory_pool", "attention", "capture", "resize"],
+        )
+
+    def test_overlap_starts_before_capture_and_finalizes_after(self):
+        self.assertEqual(
+            self._run_startup("overlap"),
+            ["start", "memory_pool", "attention", "capture", "resize", "finalize"],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/model_loader/test_prefetch_checkpoints.py

@@ -7,10 +7,11 @@ to weights loaded without prefetch.
 
 import os
 import tempfile
+import threading
 import unittest
 from concurrent.futures import Future
 from types import SimpleNamespace
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import safetensors.torch
 import torch
@@ -18,6 +19,7 @@ import torch
 from sglang.srt.configs.load_config import LoadConfig, LoadFormat
 from sglang.srt.model_loader.loader import DefaultModelLoader
 from sglang.srt.model_loader.weight_utils import (
+    CheckpointFilePrefetchHandle,
     _prefetch_all_checkpoints,
     buffered_multi_thread_safetensors_weights_iterator,
     fastsafetensors_weights_iterator,
@@ -37,6 +39,12 @@ class _InlineThread:
     def start(self):
         self.target()
 
+    def join(self, timeout=None):
+        pass
+
+    def is_alive(self):
+        return False
+
 
 class _InlineExecutor:
     def __init__(self, max_workers):
@@ -99,6 +107,45 @@ class TestPrefetchCheckpoints(CustomTestCase):
         with self.assertRaisesRegex(ValueError, "num_threads"):
             _prefetch_all_checkpoints(["dummy.safetensors"], num_threads=0)
 
+    @patch("torch.distributed.is_initialized", return_value=False)
+    def test_wait_returns_after_worker_thread_failure(self, _):
+        worker_errors = []
+        with (
+            patch(
+                "concurrent.futures.ThreadPoolExecutor",
+                side_effect=RuntimeError("worker failed"),
+            ),
+            patch(
+                "threading.excepthook",
+                side_effect=lambda args: worker_errors.append(args.exc_value),
+            ),
+        ):
+            handle = _prefetch_all_checkpoints(["dummy.safetensors"], num_threads=1)
+            handle.wait(timeout=5)
+
+        self.assertTrue(handle.done)
+        self.assertTrue(handle.failed)
+        self.assertEqual(handle.errors, ())
+        self.assertEqual(len(worker_errors), 1)
+        self.assertIsInstance(worker_errors[0], RuntimeError)
+
+    def test_prefetch_stop_has_a_bounded_default_wait(self):
+        thread = MagicMock()
+        thread.is_alive.return_value = True
+        cancel_event = threading.Event()
+        handle = CheckpointFilePrefetchHandle(
+            thread=thread,
+            cancel_event=cancel_event,
+            succeeded_event=threading.Event(),
+            errors=[],
+        )
+
+        with self.assertRaisesRegex(TimeoutError, "checkpoint prefetching"):
+            handle.stop()
+
+        self.assertTrue(cancel_event.is_set())
+        thread.join.assert_called_once_with(60.0)
+
     @patch("torch.distributed.is_initialized", return_value=False)
     def test_prefetch_keeps_bounded_pending_window(self, _):
         paths = [f"model-{i:05d}.safetensors" for i in range(20)]
@@ -106,9 +153,9 @@ class TestPrefetchCheckpoints(CustomTestCase):
         submitted_paths = []
 
         class RecordingExecutor(_InlineExecutor):
-            def submit(self, fn, path):
+            def submit(self, fn, path, *args):
                 submitted_paths.append(path)
-                return super().submit(fn, path)
+                return super().submit(fn, path, *args)
 
         def record_pending_size(fs, return_when):
             pending_sizes.append(len(fs))
@@ -129,7 +176,7 @@ class TestPrefetchCheckpoints(CustomTestCase):
     def test_prefetch_logs_failed_futures(self, _):
         paths = ["bad.safetensors"]
 
-        def fail_prefetch(path):
+        def fail_prefetch(path, cancel_event):
             raise OSError(f"failed {path}")
 
         with (
@@ -142,15 +189,18 @@ class TestPrefetchCheckpoints(CustomTestCase):
             ),
             patch("sglang.srt.model_loader.weight_utils.logger.warning") as warning,
         ):
-            _prefetch_all_checkpoints(paths, num_threads=1)
+            handle = _prefetch_all_checkpoints(paths, num_threads=1)
 
+        handle.wait()
+        self.assertEqual(handle.errors[0][0], paths[0])
+        self.assertIsInstance(handle.errors[0][1], OSError)
         warning.assert_called_once()
         self.assertEqual(
             warning.call_args.args[0],
-            "Failed to prefetch checkpoint file %r.",
+            "Failed to prefetch checkpoint file %r: %s",
         )
         self.assertEqual(warning.call_args.args[1], paths[0])
-        self.assertTrue(warning.call_args.kwargs["exc_info"])
+        self.assertIsInstance(warning.call_args.args[2], OSError)
 
     @patch("torch.distributed.is_initialized", return_value=False)
     def test_prefetch_progress_logs_all_crossed_buckets(self, _):
@@ -193,13 +243,40 @@ class TestPrefetchCheckpoints(CustomTestCase):
             ),
             patch(
                 "sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file",
-                side_effect=loaded_paths.append,
+                side_effect=lambda path, cancel_event: loaded_paths.append(path),
             ),
         ):
             _prefetch_all_checkpoints(paths, num_threads=2)
 
         self.assertEqual(sorted(loaded_paths), sorted(paths[1::3]))
 
+    @patch("torch.distributed.is_initialized", return_value=False)
+    def test_prefetch_handle_cancels_before_scheduling_next_shard(self, _):
+        paths = [f"model-{i:05d}.safetensors" for i in range(3)]
+        started = threading.Event()
+        release = threading.Event()
+        loaded_paths = []
+
+        def block_first_prefetch(path, cancel_event):
+            loaded_paths.append(path)
+            started.set()
+            self.assertTrue(release.wait(timeout=5))
+
+        with patch(
+            "sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file",
+            side_effect=block_first_prefetch,
+        ):
+            handle = _prefetch_all_checkpoints(paths, num_threads=1)
+            self.assertTrue(started.wait(timeout=5))
+            with self.assertRaisesRegex(TimeoutError, "checkpoint prefetching"):
+                handle.wait(timeout=0)
+            handle.cancel()
+            release.set()
+            handle.wait(timeout=5)
+
+        self.assertTrue(handle.cancelled)
+        self.assertEqual(loaded_paths, paths[:1])
+
     @patch("torch.distributed.is_initialized", return_value=False)
     def test_buffered_loader_drops_cache_after_each_loaded_shard(self, _):
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -319,11 +396,16 @@ class TestPrefetchDispatch(CustomTestCase):
             weight_loader_drop_cache_after_load=drop_cache,
         )
 
-    def _run(self, loader):
+    def _run(self, loader, **iterator_kwargs):
         # _get_weights_iterator returns a generator wrapping the chosen
         # iterator; consuming it forces the eager dispatch (the if/elif/else
         # that calls the iterator factory) to execute.
-        list(loader._get_weights_iterator(self._make_source()))
+        list(
+            loader._get_weights_iterator(
+                self._make_source(),
+                **iterator_kwargs,
+            )
+        )
 
     def _patch_dispatch(self, prefetch, disable_mmap=False, drop_cache=False):
         return (
@@ -332,14 +414,6 @@ class TestPrefetchDispatch(CustomTestCase):
                 "_prepare_weights",
                 return_value=("/dummy", ["f.safetensors"], True),
             ),
-            patch(
-                "sglang.srt.model_loader.loader.get_server_args",
-                return_value=self._server_args(
-                    prefetch,
-                    disable_mmap,
-                    drop_cache,
-                ),
-            ),
             patch(
                 "sglang.srt.model_loader.loader.get_model",
                 return_value=self._server_args(prefetch, disable_mmap, drop_cache),
@@ -360,12 +434,11 @@ class TestPrefetchDispatch(CustomTestCase):
         """Prefetch on + no explicit multithread config -> single-threaded,
         and the opt-out warning fires once."""
         loader = self._make_loader({})
-        p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
+        p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
             prefetch=True
         )
         with (
             p_prep,
-            p_args,
             p_model,
             p_buffered as mock_buffered,
             p_single as mock_single,
@@ -380,12 +453,11 @@ class TestPrefetchDispatch(CustomTestCase):
         """Explicit enable_multithread_load=true is the escape hatch; the
         override and its warning must not fire."""
         loader = self._make_loader({"enable_multithread_load": True})
-        p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
+        p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
             prefetch=True
         )
         with (
             p_prep,
-            p_args,
             p_model,
             p_buffered as mock_buffered,
             p_single as mock_single,
@@ -401,12 +473,11 @@ class TestPrefetchDispatch(CustomTestCase):
         default) also signals multi-thread intent, so the override must not
         fire and num_threads stays live."""
         loader = self._make_loader({"num_threads": 64})
-        p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
+        p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
             prefetch=True
         )
         with (
             p_prep,
-            p_args,
             p_model,
             p_buffered as mock_buffered,
             p_single as mock_single,
@@ -423,12 +494,11 @@ class TestPrefetchDispatch(CustomTestCase):
         """Prefetch off -> multi-threaded iterator is used (default), no
         override warning."""
         loader = self._make_loader({})
-        p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
+        p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
             prefetch=False
         )
         with (
             p_prep,
-            p_args,
             p_model,
             p_buffered as mock_buffered,
             p_single as mock_single,
@@ -439,16 +509,117 @@ class TestPrefetchDispatch(CustomTestCase):
         mock_single.assert_not_called()
         mock_warning.assert_not_called()
 
+    def test_startup_prefetch_reuses_existing_background_handle(self):
+        """Startup commit reuses resolved shards and the active prefetch handle."""
+        loader = self._make_loader({})
+        source = self._make_source()
+        resolved_source = DefaultModelLoader.ResolvedSource(
+            source=source,
+            hf_folder="/dummy",
+            weight_files=("f.safetensors",),
+            use_safetensors=True,
+        )
+        p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
+            prefetch=False
+        )
+        with (
+            p_prep as mock_prepare,
+            p_model,
+            p_buffered as mock_buffered,
+            p_single as mock_single,
+            p_warn as mock_warning,
+        ):
+            list(
+                loader._get_weights_iterator(
+                    source,
+                    resolved_source=resolved_source,
+                    startup_prefetch_started=True,
+                    startup_prefetch_active=True,
+                )
+            )
+
+        mock_prepare.assert_not_called()
+        mock_single.assert_called_once()
+        self.assertFalse(mock_single.call_args.kwargs["prefetch"])
+        mock_buffered.assert_not_called()
+        mock_warning.assert_called_once()
+
+    def test_completed_startup_prefetch_restores_multithread_loader(self):
+        loader = self._make_loader({})
+        source = self._make_source()
+        resolved_source = DefaultModelLoader.ResolvedSource(
+            source=source,
+            hf_folder="/dummy",
+            weight_files=("f.safetensors",),
+            use_safetensors=True,
+        )
+        p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
+            prefetch=False
+        )
+        with (
+            p_prep as mock_prepare,
+            p_model,
+            p_buffered as mock_buffered,
+            p_single as mock_single,
+            p_warn as mock_warning,
+        ):
+            list(
+                loader._get_weights_iterator(
+                    source,
+                    resolved_source=resolved_source,
+                    startup_prefetch_started=True,
+                    startup_prefetch_active=False,
+                )
+            )
+
+        mock_prepare.assert_not_called()
+        mock_buffered.assert_called_once()
+        self.assertFalse(mock_buffered.call_args.kwargs["prefetch"])
+        mock_single.assert_not_called()
+        mock_warning.assert_not_called()
+
+    def test_completed_startup_prefetch_is_not_started_twice(self):
+        loader = self._make_loader({})
+        source = self._make_source()
+        resolved_source = DefaultModelLoader.ResolvedSource(
+            source=source,
+            hf_folder="/dummy",
+            weight_files=("f.safetensors",),
+            use_safetensors=True,
+        )
+        p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
+            prefetch=True
+        )
+        with (
+            p_prep,
+            p_model,
+            p_buffered as mock_buffered,
+            p_single as mock_single,
+            p_warn as mock_warning,
+        ):
+            list(
+                loader._get_weights_iterator(
+                    source,
+                    resolved_source=resolved_source,
+                    startup_prefetch_started=True,
+                    startup_prefetch_active=False,
+                )
+            )
+
+        mock_buffered.assert_called_once()
+        self.assertFalse(mock_buffered.call_args.kwargs["prefetch"])
+        mock_single.assert_not_called()
+        mock_warning.assert_not_called()
+
     def test_prefetch_does_not_override_when_mmap_disabled(self):
         """Prefetch is a no-op without mmap, so the override and its warning
         must not fire."""
         loader = self._make_loader({})
-        p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
+        p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
             prefetch=True, disable_mmap=True
         )
         with (
             p_prep,
-            p_args,
             p_model,
             p_buffered as mock_buffered,
             p_single as mock_single,
@@ -463,7 +634,7 @@ class TestPrefetchDispatch(CustomTestCase):
         """FASTSAFETENSORS ignores both flags; override + warning must not
         fire."""
         loader = self._make_loader({}, load_format=LoadFormat.FASTSAFETENSORS)
-        p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
+        p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
             prefetch=True
         )
         with (
@@ -472,7 +643,6 @@ class TestPrefetchDispatch(CustomTestCase):
                 return_value=iter([]),
             ) as mock_fast,
             p_prep,
-            p_args,
             p_model,
             p_buffered as mock_buffered,
             p_single as mock_single,
@@ -492,7 +662,7 @@ class TestPrefetchDispatch(CustomTestCase):
         loader = self._make_loader(
             {"enable_gds": False}, load_format=LoadFormat.FASTSAFETENSORS
         )
-        p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
+        p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
             prefetch=False,
             drop_cache=True,
         )
@@ -502,7 +672,6 @@ class TestPrefetchDispatch(CustomTestCase):
                 return_value=iter([]),
             ) as mock_fast,
             p_prep,
-            p_args,
             p_model,
             p_buffered,
             p_single,

test/registered/unit/test_server_args_migration.py

@@ -93,6 +93,25 @@ class TestServerArgsAnnotatedCli(CustomTestCase):
                 sa = self._parse(["--image-processor-backend", backend])
                 self.assertEqual(sa.image_processor_backend, backend)
 
+    def test_startup_weight_load_mode(self):
+        """The startup loading mode keeps serial as the safe default."""
+        serial = self._parse([])
+        overlap = self._parse(["--startup-weight-load-mode", "overlap"])
+        self.assertEqual(serial.startup_weight_load_mode, "serial")
+        self.assertFalse(serial.is_startup_weight_load_overlap)
+        self.assertEqual(overlap.startup_weight_load_mode, "overlap")
+        self.assertTrue(overlap.is_startup_weight_load_overlap)
+
+        with self.assertRaises(SystemExit):
+            self.parser.parse_args(
+                [
+                    "--model",
+                    "dummy",
+                    "--startup-weight-load-mode",
+                    "unsupported",
+                ]
+            )
+
     def test_deprecated_flags_still_work(self):
         """Deprecated flags set the correct dest field."""
         sa = self._parse(["--stream-output"])