diff --git a/python/sglang/srt/weight_cache/daemon.py b/python/sglang/srt/weight_cache/daemon.py index 1c638ddb7..5c6c0e1f5 100644 --- a/python/sglang/srt/weight_cache/daemon.py +++ b/python/sglang/srt/weight_cache/daemon.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 """Weight Cache Daemon — a persistent process that holds post-quantized, -TP-sharded model weights in GPU memory and serves them via CUDA IPC handles. +TP-sharded model weights in GPU memory and serves them via pluggable transport backends. Each GPU runs one daemon process for its TP rank. The daemon: 1. Loads model weights from disk (full pipeline: disk → TP shard → quantize) 2. Exports every parameter/buffer as a CUDA IPC handle -3. Serves handles over a Unix socket to requesting engine processes +3. Serves transport entries over a Unix socket to requesting engine processes 4. Validates CacheConfig compatibility before serving Usage: @@ -39,7 +39,7 @@ import os import signal import socket import time -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, Tuple import torch import torch.distributed as dist @@ -47,7 +47,6 @@ import torch.distributed as dist from sglang.srt.configs.load_config import LoadConfig from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import publish -from sglang.srt.utils import MultiprocessingSerializer from .protocol import ( CacheConfig, @@ -63,6 +62,7 @@ from .protocol import ( recv_msg, send_msg, ) +from .transport import choose_daemon_transport_backend logger = logging.getLogger(__name__) @@ -121,8 +121,9 @@ class WeightCacheDaemon: self.model = None self.config: Optional[CacheConfig] = None - # name -> {"handle": base64_str, "shape": list, "dtype": str, "is_param": bool} + # name -> transport-specific tensor entry metadata (shape/dtype/is_param + payload metadata) self.state_entries: Dict[str, Dict[str, Any]] = {} + self.transport_backend = None def _init_distributed(self, server_args, model_config): """Initialize the distributed backend required for model loading. @@ -344,12 +345,7 @@ class WeightCacheDaemon: ) def _export_state(self): - """Export model parameters and buffers as CUDA IPC handles. - - This includes both persistent buffers (in state_dict) and non-persistent - buffers (e.g. rotary embedding cos_sin_cache) so the engine can fully - reconstruct the model state via zero-copy IPC. - """ + """Export model state entries through the selected transport backend.""" self.state_entries.clear() # remove_duplicate=False so tied weights are recognized as parameters @@ -360,45 +356,37 @@ class WeightCacheDaemon: name for name, _ in self.model.named_parameters(remove_duplicate=False) ) state_dict_names = set(self.model.state_dict().keys()) + state_tensors: Dict[str, Tuple[torch.Tensor, bool]] = {} # Export all items from state_dict (parameters + persistent buffers) for name, tensor in self.model.state_dict().items(): - ipc_handle = MultiprocessingSerializer.serialize( - tensor.data, output_str=True - ) - self.state_entries[name] = { - "handle": ipc_handle, - "shape": list(tensor.shape), - "dtype": str(tensor.dtype).replace("torch.", ""), - "is_param": name in param_names, - } + state_tensors[name] = (tensor.data, name in param_names) # Also export non-persistent buffers (not in state_dict but needed # for inference, e.g. rotary embedding cos_sin_cache) non_persistent_count = 0 for name, buf in self.model.named_buffers(): if name not in state_dict_names: - ipc_handle = MultiprocessingSerializer.serialize( - buf.data, output_str=True - ) - self.state_entries[name] = { - "handle": ipc_handle, - "shape": list(buf.shape), - "dtype": str(buf.dtype).replace("torch.", ""), - "is_param": False, - } + state_tensors[name] = (buf.data, False) non_persistent_count += 1 - # Log total size + self.transport_backend = choose_daemon_transport_backend(state_tensors) + self.state_entries = self.transport_backend.prepare_export(state_tensors) + + # Log approximate serialized metadata size (not payload-backed bytes). + # Only the handle blob carries real weight, so measure it directly: + # stringifying every entry would allocate a copy of all handles. total_bytes = sum( - entry["handle"].__len__() if hasattr(entry["handle"], "__len__") else 0 - for entry in self.state_entries.values() + len(handle) + for handle in (entry.get("handle") for entry in self.state_entries.values()) + if isinstance(handle, (str, bytes, bytearray)) ) logger.info( f"[WeightCacheDaemon gpu={self.gpu_id}] " f"Exported {len(self.state_entries)} tensors " f"({non_persistent_count} non-persistent buffers), " - f"serialized handle size ~{total_bytes / 1024 / 1024:.1f} MB" + f"transport={self.transport_backend.name}, " + f"metadata size ~{total_bytes / 1024 / 1024:.1f} MB" ) def serve(self): @@ -497,19 +485,17 @@ class WeightCacheDaemon: logger.info( f"[WeightCacheDaemon gpu={self.gpu_id}] " - f"Serving {len(self.state_entries)} IPC handles to engine" + f"Serving {len(self.state_entries)} tensors via " + f"{self.transport_backend.name} transport" ) - send_msg( + self.transport_backend.send_fetch_state_response( conn, - { - "status": "ok", - "config": self.config.to_dict(), - "entries": self.state_entries, - # PID so the client can watch daemon liveness: if this - # process dies while clients hold IPC mappings, their - # param.data (and any CUDA-graph-captured addresses) dangle. - "pid": os.getpid(), - }, + config=self.config.to_dict(), + entries=self.state_entries, + # PID so the client can watch daemon liveness: if this + # process dies while clients hold IPC mappings, their + # param.data (and any CUDA-graph-captured addresses) dangle. + pid=os.getpid(), ) 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 a761df838..a649b7fdd 100644 --- a/python/sglang/srt/weight_cache/ipc_loader.py +++ b/python/sglang/srt/weight_cache/ipc_loader.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -"""IPC Model Loader — loads model weights from a Weight Cache Daemon via CUDA IPC. +"""IPC Model Loader — loads model weights from a Weight Cache Daemon. -Zero-copy mode: param.data points directly to IPC-mapped GPU memory. Only 1x GPU -memory needed — engine and daemon share the same physical GPU memory via CUDA IPC. -Engine depends on daemon staying alive. +Zero-copy mode: param.data points directly to transport-mapped GPU memory. +Backends are negotiated per daemon response (torch IPC by default, VMM FD when +available). Engine depends on daemon staying alive. """ import logging @@ -22,7 +22,6 @@ from sglang.srt.model_loader.loader import ( BaseModelLoader, _initialize_model, ) -from sglang.srt.utils import MultiprocessingSerializer from .protocol import ( CacheConfig, @@ -33,6 +32,7 @@ from .protocol import ( recv_msg, send_msg, ) +from .transport import TORCH_IPC_BACKEND, get_client_transport_backend logger = logging.getLogger(__name__) @@ -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._transport_backend = get_client_transport_backend(TORCH_IPC_BACKEND) def load_model( self, @@ -119,7 +120,8 @@ class IpcModelLoader(BaseModelLoader): entries = cache_data["entries"] logger.info( - f"[IpcModelLoader] Fetched {len(entries)} IPC handles from daemon " + f"[IpcModelLoader] Fetched {len(entries)} tensors from daemon " + f"(transport={self._transport_backend.name}) " f"in {time.perf_counter() - tic:.2f}s" ) @@ -337,7 +339,7 @@ class IpcModelLoader(BaseModelLoader): # This ensures post-quantization parameters (weight_scale, etc.) # that were created by process_weights_after_loading are also mapped. for name, entry in entries.items(): - imported_tensor = MultiprocessingSerializer.deserialize(entry["handle"]) + imported_tensor = self._transport_backend.import_tensor(entry) is_param = entry.get("is_param", True) if name in existing_names: @@ -416,6 +418,9 @@ class IpcModelLoader(BaseModelLoader): # Stash IPC refs on the model to prevent GC (which would unmap the memory) if imported_refs: model._ipc_imported_tensors = imported_refs + # Keep transport backend alive for the model lifetime (VMM backend owns + # VA mappings that must stay mapped while tensors are in use). + model._weight_cache_transport_backend = self._transport_backend logger.info( f"[IpcModelLoader] Zero-copy: mapped {imported_count} tensors " @@ -529,6 +534,9 @@ class IpcModelLoader(BaseModelLoader): f" Daemon config: {daemon_config}" ) + backend_name = result.get("transport_backend", TORCH_IPC_BACKEND) + self._transport_backend = get_client_transport_backend(backend_name) + result = self._transport_backend.recv_fetch_state_response(sock, result) return result except RuntimeError: diff --git a/python/sglang/srt/weight_cache/transport.py b/python/sglang/srt/weight_cache/transport.py new file mode 100644 index 000000000..7318db824 --- /dev/null +++ b/python/sglang/srt/weight_cache/transport.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pluggable tensor transport backends for weight_cache.""" + +from __future__ import annotations + +import array +import logging +import os +import socket +import struct +from abc import ABC, abstractmethod +from typing import Any, Dict, Mapping, NoReturn, Optional, Tuple + +import torch + +from sglang.srt.utils import MultiprocessingSerializer + +from .protocol import send_msg + +logger = logging.getLogger(__name__) + +TORCH_IPC_BACKEND = "torch_ipc" +VMM_FD_BACKEND = "vmm_fd" + +_FD_INDEX_STRUCT = struct.Struct(" None: + payload = _FD_INDEX_STRUCT.pack(index) + fds = array.array("i", [int(fd)]) + sent = sock.sendmsg( + [payload], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds.tobytes())] + ) + if sent != len(payload): + raise RuntimeError(f"sendmsg sent {sent} bytes, expected {len(payload)}") + + +def _recv_fd(sock: socket.socket) -> Tuple[int, int]: + fd_item_size = array.array("i").itemsize + data, ancdata, _, _ = sock.recvmsg( + _FD_INDEX_STRUCT.size, socket.CMSG_SPACE(fd_item_size) + ) + if len(data) != _FD_INDEX_STRUCT.size: + raise RuntimeError( + f"received truncated fd header: {len(data)} < {_FD_INDEX_STRUCT.size}" + ) + index = _FD_INDEX_STRUCT.unpack(data)[0] + fds = array.array("i") + for level, cmsg_type, cmsg_data in ancdata: + if level == socket.SOL_SOCKET and cmsg_type == socket.SCM_RIGHTS: + fds.frombytes(cmsg_data[: len(cmsg_data) - (len(cmsg_data) % fd_item_size)]) + if len(fds) != 1: + for fd in fds: + os.close(fd) + raise RuntimeError(f"expected one fd, got {len(fds)}") + return int(index), int(fds[0]) + + +class WeightCacheTransportBackend(ABC): + name: str + + @abstractmethod + def prepare_export( + self, state_tensors: Mapping[str, Tuple[torch.Tensor, bool]] + ) -> Dict[str, Dict[str, Any]]: + """Prepare daemon-side entries for all tensors.""" + + @abstractmethod + def send_fetch_state_response( + self, + conn: socket.socket, + *, + config: Dict[str, Any], + entries: Dict[str, Dict[str, Any]], + pid: int, + ) -> None: + """Send a successful fetch_state response.""" + + @abstractmethod + def recv_fetch_state_response( + self, sock: socket.socket, result: Dict[str, Any] + ) -> Dict[str, Any]: + """Client-side receive hook after recv_msg.""" + + @abstractmethod + def import_tensor(self, entry: Dict[str, Any]) -> torch.Tensor: + """Import a single tensor from one entry.""" + + +class TorchIpcTransportBackend(WeightCacheTransportBackend): + name = TORCH_IPC_BACKEND + + def prepare_export( + self, state_tensors: Mapping[str, Tuple[torch.Tensor, bool]] + ) -> Dict[str, Dict[str, Any]]: + entries: Dict[str, Dict[str, Any]] = {} + for name, (tensor, is_param) in state_tensors.items(): + entries[name] = { + "handle": MultiprocessingSerializer.serialize( + tensor.data, output_str=True + ), + "shape": list(tensor.shape), + "dtype": str(tensor.dtype).replace("torch.", ""), + "is_param": is_param, + } + return entries + + def send_fetch_state_response( + self, + conn: socket.socket, + *, + config: Dict[str, Any], + entries: Dict[str, Dict[str, Any]], + pid: int, + ) -> None: + send_msg( + conn, + { + "status": "ok", + "config": config, + "entries": entries, + "pid": pid, + "transport_backend": self.name, + }, + ) + + def recv_fetch_state_response( + self, sock: socket.socket, result: Dict[str, Any] + ) -> Dict[str, Any]: + return result + + def import_tensor(self, entry: Dict[str, Any]) -> torch.Tensor: + return MultiprocessingSerializer.deserialize(entry["handle"]) + + +class VmmFdTransportBackend(WeightCacheTransportBackend): + """Placeholder for the CUDA VMM + fd-passing transport. + + The backend is not wired up yet: can_export_state reports False so the + daemon keeps selecting torch_ipc, and every other entry point fails loudly + instead of silently returning None. + """ + + name = VMM_FD_BACKEND + + def __init__(self): + self._raise_not_implemented() + + @staticmethod + def _raise_not_implemented() -> NoReturn: + raise NotImplementedError( + f"weight cache transport backend {VMM_FD_BACKEND!r} is not " + f"implemented in this build" + ) + + @classmethod + def can_export_state( + cls, state_tensors: Mapping[str, Tuple[torch.Tensor, bool]] + ) -> bool: + return False + + def prepare_export( + self, state_tensors: Mapping[str, Tuple[torch.Tensor, bool]] + ) -> Dict[str, Dict[str, Any]]: + self._raise_not_implemented() + + def send_fetch_state_response( + self, + conn: socket.socket, + *, + config: Dict[str, Any], + entries: Dict[str, Dict[str, Any]], + pid: int, + ) -> None: + self._raise_not_implemented() + + def recv_fetch_state_response( + self, sock: socket.socket, result: Dict[str, Any] + ) -> Dict[str, Any]: + self._raise_not_implemented() + + def import_tensor(self, entry: Dict[str, Any]) -> torch.Tensor: + self._raise_not_implemented() + + +def choose_daemon_transport_backend( + state_tensors: Mapping[str, Tuple[torch.Tensor, bool]], +) -> WeightCacheTransportBackend: + if VmmFdTransportBackend.can_export_state(state_tensors): + logger.info("[weight_cache] Using transport backend: %s", VMM_FD_BACKEND) + return VmmFdTransportBackend() + logger.info("[weight_cache] Using transport backend: %s", TORCH_IPC_BACKEND) + return TorchIpcTransportBackend() + + +def get_client_transport_backend(name: Optional[str]) -> WeightCacheTransportBackend: + if name in (None, "", TORCH_IPC_BACKEND): + return TorchIpcTransportBackend() + if name == VMM_FD_BACKEND: + return VmmFdTransportBackend() + raise RuntimeError(f"Unknown weight cache transport backend {name!r}") 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 4788897fd..ae4b4354c 100644 --- a/test/registered/unit/model_loader/test_weight_cache_protocol.py +++ b/test/registered/unit/model_loader/test_weight_cache_protocol.py @@ -22,6 +22,8 @@ import socket import struct import unittest +import torch + from sglang.srt.weight_cache.protocol import ( IPC_QUANT_ALLOWLIST, CacheConfig, @@ -38,6 +40,11 @@ from sglang.srt.weight_cache.protocol import ( recv_msg, send_msg, ) +from sglang.srt.weight_cache.transport import ( + TORCH_IPC_BACKEND, + TorchIpcTransportBackend, + get_client_transport_backend, +) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -112,6 +119,38 @@ class TestProtocolFraming(CustomTestCase): b.close() +class TestTransportBackend(CustomTestCase): + def test_default_backend_is_torch_ipc(self): + backend = get_client_transport_backend(None) + self.assertEqual(backend.name, TORCH_IPC_BACKEND) + + def test_unknown_backend_raises(self): + with self.assertRaises(RuntimeError): + get_client_transport_backend("does_not_exist") + + def test_torch_ipc_backend_round_trip(self): + backend = TorchIpcTransportBackend() + state_tensors = {"x": (torch.arange(8, dtype=torch.float32), True)} + entries = backend.prepare_export(state_tensors) + + a, b = socket.socketpair() + try: + backend.send_fetch_state_response( + a, + config={"k": "v"}, + entries=entries, + pid=123, + ) + 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) + finally: + a.close() + b.close() + + class TestCacheConfig(CustomTestCase): def test_identical_configs_match(self): self.assertTrue(_make_cache_config().matches(_make_cache_config()))