执行摘要
- 一句话:Weight Daemon 引入可插拔传输后端,默认路径保持 torch_ipc
- 推荐动作:值得 weight_cache 与模型加载相关开发者精读。核心看点:可插拔后端抽象的接口切分(导出/响应/接收/导入四段)、daemon 与 client 通过协议字段协商后端的设计、以及占位后端 fail loud 的降级策略(宁抛错不静默回退)。一般读者重点阅读 transport.py 的接口设计与 daemon._export_state 的接入方式即可;同时留意 _send_fd/_recv_fd 尚未启用,启用 VMM 后端前应补齐 GPU E2E 与 fd 故障测试。
功能与动机
PR body 是空模板,未填写实质内容,动机需要从代码意图推断。daemon 原先将张量导出/导入硬编码为 torch multiprocessing 序列化(MultiprocessingSerializer 输出 base64 字符串 handle),无法支撑更高效的 CUDA VMM + fd 传递方案。VmmFdTransportBackend 的 docstring 明确说明这是“Placeholder for the CUDA VMM + fd-passing transport. The backend is not wired up yet”;ipc_loader.py 模块 docstring 也写明 “Backends are negotiated per daemon response (torch IPC by default, VMM FD when available)”,即目标是让 daemon 与 client 通过响应协商传输方式,为后续跨进程共享同一物理 GPU 内存(零拷贝)铺路。
实现拆解
- 新增传输后端抽象层(python/sglang/srt/weight_cache/transport.py):定义 WeightCacheTransportBackend ABC,包含 prepare_export、send_fetch_state_response、recv_fetch_state_response、import_tensor 四个方法,把 daemon 导出与 client 导入彻底解耦。同文件实现 TorchIpcTransportBackend(默认)与 VmmFdTransportBackend(占位),并提供 choose_daemon_transport_backend 与 get_client_transport_backend 两个后端选择入口。
- daemon 侧迁移(python/sglang/srt/weight_cache/daemon.py):_export_state 不再直接调用 MultiprocessingSerializer.serialize,而是先收集 state_tensors(参数、持久 buffer、非持久 buffer 统一为 (tensor, is_param) 映射),再由 choose_daemon_transport_backend 选择后端并调用 prepare_export;_handle_connection 中成功响应改由后端 send_fetch_state_response 发出,并在响应中注入 transport_backend 字段;同时移除对 MultiprocessingSerializer 的导入。
- client 侧迁移(python/sglang/srt/weight_cache/ipc_loader.py):_fetch_from_cache 读取响应中的 transport_backend 字段(缺省视为 torch_ipc 以兼容旧 daemon),通过 get_client_transport_backend 构造后端并调用 recv_fetch_state_response;_load_zero_copy_mode 中 MultiprocessingSerializer.deserialize 替换为 backend.import_tensor(entry);后端实例挂在 model._weight_cache_transport_backend 上防止被 GC(为未来 VMM 后端持有 VA 映射做准备)。
- fd 传递原语与兜底逻辑:transport.py 内置 _send_fd / _recv_fd(基于 SCM_RIGHTS 的 fd 传递,含截断校验与异常时 fd 清理),当前仅供未来 VMM 后端使用;VmmFdTransportBackend 构造即抛 NotImplementedError、can_export_state 恒返回 False,保证 daemon 不会静默选到未就绪后端;client 端遇到未知后端名直接抛 RuntimeError,fail loud。
- 测试配套(test/registered/unit/model_loader/test_weight_cache_protocol.py):新增 TestTransportBackend,三个用例覆盖 get_client_transport_backend(None) 默认返回 torch_ipc、未知名称抛 RuntimeError、以及 TorchIpcTransportBackend 在真实 socketpair 上的导出→发送→接收→导入全链路 round-trip(用 torch.equal 校验张量内容)。
关键文件:
python/sglang/srt/weight_cache/transport.py(模块 传输后端;类别 source;类型 core-logic;符号 _send_fd, _recv_fd, WeightCacheTransportBackend, prepare_export): 本 PR 的核心新增文件,定义可插拔传输后端抽象与 torch_ipc/vmm_fd 两个实现,是 daemon 与 client 解耦的桥梁。
python/sglang/srt/weight_cache/daemon.py(模块 权重缓存;类别 source;类型 dependency-wiring;符号 _export_state, serve, _handle_connection): daemon 导出与响应发送路径迁移到传输后端,是后端抽象在服务端的落地点。
python/sglang/srt/weight_cache/ipc_loader.py(模块 模型加载;类别 source;类型 dependency-wiring;符号 _fetch_from_cache, _load_zero_copy_mode, load_model, init): client 侧按 daemon 响应协商后端并改用 import_tensor 导入,是后端抽象在客户端的关键接入点。
test/registered/unit/model_loader/test_weight_cache_protocol.py(模块 单元测试;类别 test;类型 test-coverage;符号 TestTransportBackend, test_default_backend_is_torch_ipc, test_unknown_backend_raises, test_torch_ipc_backend_round_trip): 新增后端抽象相关的 CPU 单测,覆盖默认选择、未知后端报错与 torch_ipc 全链路 round-trip。
关键符号:choose_daemon_transport_backend, get_client_transport_backend, prepare_export, send_fetch_state_response, recv_fetch_state_response, import_tensor, _send_fd, _recv_fd, _export_state, _fetch_from_cache, _load_zero_copy_mode
关键源码片段
python/sglang/srt/weight_cache/transport.py
本 PR 的核心新增文件,定义可插拔传输后端抽象与 torch_ipc/vmm_fd 两个实现,是 daemon 与 client 解耦的桥梁。
# SPDX-License-Identifier: Apache-2.0
"""可插拔的张量传输后端:daemon 导出与 client 导入的解耦点。"""
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
TORCH_IPC_BACKEND = "torch_ipc"
VMM_FD_BACKEND = "vmm_fd"
class WeightCacheTransportBackend(ABC):
"""传输后端抽象:daemon 负责 prepare_export / send_fetch_state_response,
client 负责 recv_fetch_state_response / import_tensor,
两侧通过响应里的 transport_backend 字段对齐实现。"""
name: str
@abstractmethod
def prepare_export(
self, state_tensors: Mapping[str, Tuple[torch.Tensor, bool]]
) -> Dict[str, Dict[str, Any]]:
"""daemon 侧:把全部张量转成可跨进程传输的 entry 元数据。"""
@abstractmethod
def send_fetch_state_response(self, conn, *, config, entries, pid) -> None:
"""daemon 侧:发送成功的 fetch_state 响应。"""
@abstractmethod
def recv_fetch_state_response(self, sock, result) -> Dict[str, Any]:
"""client 侧:收到响应后的后续处理(例如读取 fd 附属数据)。"""
@abstractmethod
def import_tensor(self, entry: Dict[str, Any]) -> torch.Tensor:
"""client 侧:从单个 entry 恢复出张量。"""
class TorchIpcTransportBackend(WeightCacheTransportBackend):
"""默认实现:沿用 torch multiprocessing 序列化(base64 字符串 handle),
行为与抽象引入前完全一致,保证零回归。"""
name = TORCH_IPC_BACKEND
def prepare_export(self, state_tensors):
entries = {}
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, *, config, entries, pid):
send_msg(conn, {
"status": "ok",
"config": config,
"entries": entries,
"pid": pid,
"transport_backend": self.name, # 新增字段,client 据此协商后端
})
def recv_fetch_state_response(self, sock, result):
return result # torch_ipc 无需额外处理
def import_tensor(self, entry):
return MultiprocessingSerializer.deserialize(entry["handle"])
class VmmFdTransportBackend(WeightCacheTransportBackend):
"""CUDA VMM + fd 传递后端占位。当前未接线:can_export_state 恒返回 False,
因此 daemon 始终选 torch_ipc;其他入口一律 fail loud,拒绝静默返回 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 implemented in this build"
)
@classmethod
def can_export_state(cls, state_tensors) -> bool:
return False # 未就绪,恒 False
def choose_daemon_transport_backend(state_tensors):
"""daemon 侧选择器:未来 VMM 具备导出能力时自动切换,当前固定 torch_ipc。"""
if VmmFdTransportBackend.can_export_state(state_tensors):
return VmmFdTransportBackend()
return TorchIpcTransportBackend()
def get_client_transport_backend(name):
"""client 侧按名称构造后端;未知名称直接抛错,避免静默降级。"""
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}")
python/sglang/srt/weight_cache/daemon.py
daemon 导出与响应发送路径迁移到传输后端,是后端抽象在服务端的落地点。
def _export_state(self):
"""把模型全部参数/缓冲交给所选传输后端导出,daemon 不再关心序列化格式。"""
self.state_entries.clear()
# remove_duplicate=False:tied weight 在每个名字下都被识别为参数,
# 避免 client 侧把重复键误注册成 buffer
param_names = set(
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]] = {}
# 参数 + 持久 buffer 统一收集;非持久 buffer(如 RoPE cos_sin_cache)
# 不在 state_dict 里,单独补上,保证 client 能完整重建模型状态
for name, tensor in self.model.state_dict().items():
state_tensors[name] = (tensor.data, name in param_names)
non_persistent_count = 0
for name, buf in self.model.named_buffers():
if name not in state_dict_names:
state_tensors[name] = (buf.data, False)
non_persistent_count += 1
# 选择后端并导出;这里是新增的抽象接入点
self.transport_backend = choose_daemon_transport_backend(state_tensors)
self.state_entries = self.transport_backend.prepare_export(state_tensors)
# 日志只统计 handle 本体(str/bytes)的长度,避免 stringify 整个 entry
# 造成大对象拷贝(review 中提出的性能问题)
total_bytes = sum(
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"transport={self.transport_backend.name}, "
f"metadata size ~{total_bytes / 1024 / 1024:.1f} MB"
)
评论区精华
两条 review 评论均已解决:
- alexnails 在 daemon.py 的 _export_state 差异上提问“
len(str(entry).encode("utf-8")) is expensive?”,担心 stringify 整个 entry 会带来不必要的拷贝开销。作者回复 fixed,最终实现改为只统计 handle 本体(str/bytes/bytearray)的长度,避免对完整 entry 做字符串化,日志口径也改为 metadata size。
-
alexnails 在 transport.py 的 VmmFdTransportBackend._raise_not_implemented 上建议采用带后端名的 f-string 错误消息(something like ...),作者采纳并回复 fixed。
-
daemon 导出日志统计方式的性能疑问 (performance): 作者回复 fixed,最终实现改为仅统计 handle(str/bytes/bytearray)长度,绕过对完整 entry 的字符串化,并把日志措辞改为 metadata size。
- VmmFdTransportBackend 未实现错误消息格式 (style): 作者回复 fixed,最终实现为 f"weight cache transport backend {VMM_FD_BACKEND!r} is not implemented in this build"。
风险与影响
- 风险:
- 加载主路径被重构(daemon._export_state 与 IpcModelLoader._load_zero_copy_mode),虽然默认后端行为等价,但 GPU 端 E2E 测试(test_weight_cache_daemon.py)不在本 PR 修改范围,CPU 单测无法覆盖 CUDA IPC 映射细节,存在回归未被单测捕获的风险。
- 协议新增 transport_backend 字段:client 用 result.get("transport_backend", TORCH_IPC_BACKEND) 缺省兜底,对旧 daemon 兼容;但若未来 daemon 返回 vmm_fd 而 client 侧该后端未实现(构造即抛 NotImplementedError),会硬失败——这是有意的 fail loud 策略,但跨版本混用时需要提前规划。
- _send_fd / _recv_fd 已随本 PR 提交但未启用:SCM_RIGHTS 的 CMSG 边界处理、fd 数量校验、部分发送失败等边界条件缺乏真实场景验证,未来启用 VMM 后端前需要补 fd 层面的故障注入测试。
- 日志统计口径变化:total_bytes 只统计 str/bytes/bytearray 类型的 handle,若未来后端改用整型 fd 或对象句柄,日志会低估元数据大小,可能误导排查。
- 影响:对用户与运行时:默认路径无行为变化,仍是 torch_ipc 后端,协议响应新增字段对旧客户端透明。对系统架构:weight_cache 模块完成导出/导入与具体序列化方式的解耦,后续 CUDA VMM 后端只需实现 WeightCacheTransportBackend 并在两个选择函数中注册即可。对团队协作:为 weight_cache 的多后端路线(跨进程零拷贝、降低序列化开销)划定了扩展点,风险集中在版本兼容与 GPU E2E 验证上。
- 风险标记:核心加载路径变更, 协议字段新增, fd 传递安全面, GPU E2E 未覆盖
关联脉络
参与讨论