[FEAT] Weight Daemon abstraction (#33279)

Co-authored-by: liusy58 <liusy58@smail.nju.edu.cn>

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

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:

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("<Q")
+
+
+def _send_fd(sock: socket.socket, fd: int, index: int) -> 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}")

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