[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.