diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 82f9f7ff9..dfe597381 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -972,6 +972,10 @@ class Scheduler( and self.tp_worker.model_runner.token_to_kv_pool_allocator is not None ): return + preloaded_weights_bytes = self.tp_worker.preloaded_weights_bytes + if self.draft_worker is not None: + preloaded_weights_bytes += self.draft_worker.preloaded_weights_bytes + self.tp_worker.model_runner.account_preloaded_weights(preloaded_weights_bytes) self.tp_worker.alloc_memory_pool() def init_memory_pools(self): diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index fbfa399ca..52404e2dc 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -127,6 +127,11 @@ class BaseTpWorker(ABC): runners = self.model_runner_list or [self.model_runner] return sum(runner.weight_load_time for runner in runners) + @property + def preloaded_weights_bytes(self) -> int: + runners = self.model_runner_list or [self.model_runner] + return sum(runner.preloaded_weights_bytes for runner in runners) + def get_pad_input_ids_func(self): return getattr(self.model_runner.model, "pad_input_ids", None) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 6738521f4..45c584d59 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -844,6 +844,24 @@ class ModelRunner: max_rows = max(max_rows, max(capture_bs) * num_tokens_per_req) return max_rows + @property + def preloaded_weights_bytes(self) -> int: + value = self.loader.preloaded_weights_bytes + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError( + "ModelLoader.preloaded_weights_bytes must be a non-negative int, " + f"got {value!r}" + ) + return value + + def account_preloaded_weights(self, preloaded_weights_bytes: int) -> None: + # Dist-init sampled B after the daemon already held weights, so slack + # (B * (1 - mem_fraction_static)) is too small. Add those bytes back + # onto the existing MIN'd baseline. Skip when nothing was preloaded. + if preloaded_weights_bytes == 0: + return + self.pre_model_load_memory += preloaded_weights_bytes / (1 << 30) + def alloc_memory_pool(self, memory_pool_config: Optional[MemoryPoolConfig] = None): """Allocate KV cache memory pools only (no backends or cuda graphs).""" if memory_pool_config is not None: diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index dc68cef8a..8bd1d3714 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -310,6 +310,10 @@ def _post_load_weights(model: nn.Module) -> None: class BaseModelLoader(ABC): """Base class for model loaders.""" + # Rank-local weight memory already resident when ModelRunner sampled its + # pre-load baseline. Shared allocations must be reported by only one loader. + preloaded_weights_bytes: int = 0 + def __init__(self, load_config: LoadConfig): self.load_config = load_config diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py index ba0ffd616..7f6edb336 100644 --- a/python/sglang/srt/speculative/base_spec_worker.py +++ b/python/sglang/srt/speculative/base_spec_worker.py @@ -95,6 +95,10 @@ class EagleDraftWorkerBase(ABC): def weight_load_time(self) -> float: return sum(runner.weight_load_time for runner in self.draft_runners) + @property + def preloaded_weights_bytes(self) -> int: + return sum(runner.preloaded_weights_bytes for runner in self.draft_runners) + def alloc_memory_pool(self, **kwargs): pass @@ -211,6 +215,12 @@ class BaseSpecWorker(ABC): return 0.0 return self.draft_worker.weight_load_time + @property + def preloaded_weights_bytes(self) -> int: + if self.draft_worker is None: + return 0 + return self.draft_worker.preloaded_weights_bytes + @property def last_shared_read_runner(self): # The runner that runs the step's LAST shared-buffer-reading phase -- diff --git a/python/sglang/srt/weight_cache/daemon.py b/python/sglang/srt/weight_cache/daemon.py index 30240cd00..8bb80ae33 100644 --- a/python/sglang/srt/weight_cache/daemon.py +++ b/python/sglang/srt/weight_cache/daemon.py @@ -184,6 +184,7 @@ class WeightCacheDaemon: self.config: Optional[CacheConfig] = None # name -> transport-specific tensor entry metadata (shape/dtype/is_param + payload metadata) self.state_entries: Dict[str, Dict[str, Any]] = {} + self.preloaded_weights_bytes = 0 self.transport_backend = None def _init_distributed(self, server_args, model_config): @@ -347,6 +348,9 @@ class WeightCacheDaemon: **compute_env_stamp(), ) + current_platform.empty_cache() + memory_before_load = torch.cuda.memory_reserved(self.gpu_id) + # Build load config load_config = LoadConfig( load_format=self.load_format, @@ -377,6 +381,10 @@ class WeightCacheDaemon: # memory: clients map these tensors read-only via IPC and would otherwise # risk observing half-written weights. current_platform.synchronize() + current_platform.empty_cache() + self.preloaded_weights_bytes = max( + 0, torch.cuda.memory_reserved(self.gpu_id) - memory_before_load + ) # Export all parameters and buffers as IPC handles self._export_state() @@ -582,6 +590,7 @@ class WeightCacheDaemon: # process dies while clients hold IPC mappings, their # param.data (and any CUDA-graph-captured addresses) dangle. pid=os.getpid(), + preloaded_weights_bytes=self.preloaded_weights_bytes, ) elif req.get("type") == "ping": diff --git a/python/sglang/srt/weight_cache/ipc_loader.py b/python/sglang/srt/weight_cache/ipc_loader.py index ddc54bed7..fa0d5c72c 100644 --- a/python/sglang/srt/weight_cache/ipc_loader.py +++ b/python/sglang/srt/weight_cache/ipc_loader.py @@ -76,6 +76,7 @@ class IpcModelLoader(BaseModelLoader): self.weight_cache_mode = weight_cache_mode self._fallback_loader_cls = fallback_loader_cls self._fallback_load_format = fallback_load_format + self.preloaded_weights_bytes = 0 self._transport_backend = get_client_transport_backend(TORCH_IPC_BACKEND) def load_model( @@ -90,6 +91,7 @@ class IpcModelLoader(BaseModelLoader): (fallback to disk loading would cause OOM on shared GPUs). In client mode, falls back to DefaultModelLoader. """ + self.preloaded_weights_bytes = 0 tic = time.perf_counter() # Hard-gate unsupported quant methods before touching the daemon, so an @@ -119,6 +121,19 @@ class IpcModelLoader(BaseModelLoader): return self._fallback_load(model_config, device_config) entries = cache_data["entries"] + # Older daemons omit this field; missing metadata means no correction. + preloaded_weights_bytes = cache_data.get("preloaded_weights_bytes", 0) + if preloaded_weights_bytes is None: + preloaded_weights_bytes = 0 + if ( + isinstance(preloaded_weights_bytes, bool) + or not isinstance(preloaded_weights_bytes, int) + or preloaded_weights_bytes < 0 + ): + raise RuntimeError( + "[IpcModelLoader] Daemon returned invalid weight-memory metadata: " + f"{preloaded_weights_bytes=}" + ) logger.info( f"[IpcModelLoader] Fetched {len(entries)} tensors from daemon " f"(transport={self._transport_backend.name}) " @@ -137,6 +152,7 @@ class IpcModelLoader(BaseModelLoader): entries, quant_config, ) + self.preloaded_weights_bytes = preloaded_weights_bytes # Skip _post_load_weights: the daemon already ran # process_weights_after_loading on the weights before exporting diff --git a/python/sglang/srt/weight_cache/transport.py b/python/sglang/srt/weight_cache/transport.py index 7318db824..f770fbb33 100644 --- a/python/sglang/srt/weight_cache/transport.py +++ b/python/sglang/srt/weight_cache/transport.py @@ -73,6 +73,7 @@ class WeightCacheTransportBackend(ABC): config: Dict[str, Any], entries: Dict[str, Dict[str, Any]], pid: int, + preloaded_weights_bytes: int = 0, ) -> None: """Send a successful fetch_state response.""" @@ -112,6 +113,7 @@ class TorchIpcTransportBackend(WeightCacheTransportBackend): config: Dict[str, Any], entries: Dict[str, Dict[str, Any]], pid: int, + preloaded_weights_bytes: int = 0, ) -> None: send_msg( conn, @@ -121,6 +123,7 @@ class TorchIpcTransportBackend(WeightCacheTransportBackend): "entries": entries, "pid": pid, "transport_backend": self.name, + "preloaded_weights_bytes": preloaded_weights_bytes, }, ) @@ -171,6 +174,7 @@ class VmmFdTransportBackend(WeightCacheTransportBackend): config: Dict[str, Any], entries: Dict[str, Dict[str, Any]], pid: int, + preloaded_weights_bytes: int = 0, ) -> None: self._raise_not_implemented() diff --git a/test/registered/unit/model_loader/test_weight_cache_protocol.py b/test/registered/unit/model_loader/test_weight_cache_protocol.py index bacbfa415..e1aa82da0 100644 --- a/test/registered/unit/model_loader/test_weight_cache_protocol.py +++ b/test/registered/unit/model_loader/test_weight_cache_protocol.py @@ -149,12 +149,14 @@ class TestTransportBackend(CustomTestCase): config={"k": "v"}, entries=entries, pid=123, + preloaded_weights_bytes=65536, ) resp = recv_msg(b) resp = backend.recv_fetch_state_response(b, resp) imported = backend.import_tensor(resp["entries"]["x"]) self.assertTrue(torch.equal(imported.cpu(), state_tensors["x"][0])) self.assertEqual(resp["transport_backend"], TORCH_IPC_BACKEND) + self.assertEqual(resp["preloaded_weights_bytes"], 65536) finally: a.close() b.close()