[Fix] Account resident weight memory in KV sizing (#34053)

Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>

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):

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)
 

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:

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
 

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

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":

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

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()
 

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()