Prhub

#27139 [FEAT] Support fast engine recovery through weight cache

原始 PR 作者 QiuMike 合并时间 2026-07-25 14:31 文件变更 18 提交数 78 评论 115 代码增减 +3404 / -2

执行摘要

引入权重缓存守护进程 , 引擎重启提速至 <1 秒

SGLang engine restarts are expensive. When an engine crashes, rolls out, or needs to restart for any reason, it must reload model weights from storage into GPU memory, re-initialize the distributed backend, recapture CUDA graphs, and re-warm JIT kernels — a process that takes 3–6+ minutes for large models (e.g., Qwen3-235B FP8 takes ~6.5 minutes on 4×GPU). Weight loading from disk dominates at ~79% of total startup time (~306/327 seconds). Eliminating weight loading is the single highest-leverage optimization. 此外,权重缓存守护进程还解锁了多实例权重共享、优先级协同服务(priority co-serving)、主动-备用故障切换(active-standby failover)等生产场景,这些在传统加载模式下因重启代价过高而难以实现。

里程碑级 PR,值得所有生产运维团队和架构师精读。重点关注的决策包括:(1) CUDA IPC 零拷贝 vs 进程级 checkpoint (CRIU+cuda-checkpoint) 的设计权衡;(2) 量化白名单机制确保数值正确性;(3) CacheConfig 指纹匹配策略(环境戳 + revision + 量化哈希);(4) 进程生命周期管理(PR_SET_PDEATHSIG + force 清理);(5) 与现有模块(权重更新、CPU 备份、内存保存器)的交互边界。该 PR 为后续快速重启路线图奠定了坚实基础,代码质量高、讨论充分,是值得学习的系统设计案例。

讨论亮点
  1. IPC 量化方法兼容性(alexnails & Claude):_post_load_weights 跳过后,某些量化方法(如 per-tensor FP8)的 Python 属性(如 weight_scale_inv.format_ue8m0)不会跨 IPC 传递,导致静默输出错误。作者引入 IPC_QUANT_ALLOWLIST 白名单解决,仅允许 unquantized 和 block-wise FP8 通过。

  2. 守护进程生命周期与脏文件清理(Kangyan-Zhou, alexnails & Claude):引擎崩溃后守护进程成为孤儿,残留 .ready/.sock 文件阻止下次启动。作者通过 PR_SET_PDEATHSIGshutdown() 中先 SIGTERM 再 SIGKILL、cleanup_stale_daemon_files 添加 --force 标志解决。

  3. 连接超时与 socket 泄漏(Kangyan-Zhou & alexnails):守护进程 accept 后未设置超时,一个挂起客户端会阻塞所有其他引擎等级。作者添加 conn.settimeout(30s)。socket 连接失败时未关闭导致泄漏,已修复。

  4. 扩展段内存分配器不兼容(alexnails & Claude):PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 使 _share_cuda_ 失败。作者添加 _assert_ipc_compatible_allocator 提前报错。

  5. 权重在线更新保护(alexnails & Claude):IPC 零拷贝模式下 update_weights_from_disk 等操作会损坏共享内存。作者在 weight_updater.py 的四个入口处插入 _assert_weight_cache_inactive 硬性报错。

  6. 测试覆盖计划(Kangyan-Zhou & alexnails):初始版本仅手动测试。作者添加 CPU 单元测试(无 GPU)和 GPU E2E 测试,TP=1 冒烟测试始终运行在 base-b 阶段,TP=2 测试作为 extra-a 阶段。

实现拆解

  1. 协议层设计:在 protocol.py 中定义 CacheConfig 类(基于 msgspec.Struct),包含模型路径、架构、TP/PP/DP/EP 规模、量化方法、量化配置哈希、dtyperevision、环境戳(CUDA 算力 + torch 版本)。实现长度前缀的 socket 消息帧协议 (send_msg/recv_msg) 和 IPC 量化方法白名单机制 (IPC_QUANT_ALLOWLIST),仅允许 unquantized 和 block-wise FP8 通过 IPC 加载。

  2. 守护进程实现daemon.py 中的 WeightCacheDaemon 类是一个持久化进程,负责:初始化分布式环境 (_init_distributed),使用 DefaultModelLoader 从磁盘加载完整模型权重并执行后量化处理,然后通过 _export_state 遍历所有参数和缓冲区,使用 torch.Tensor._share_cuda_ 导出为 CUDA IPC 句柄(base64 编码)。最后启动 Unix socket 服务循环 (serve),在 CLIENT_CONNECTION_TIMEOUT (30s) 超时内接受引擎连接,验证 CacheConfig 匹配后发送所有句柄。

  3. IPC 加载器ipc_loader.py 中的 IpcModelLoader 类(继承 BaseModelLoader)是引擎端的加载器。通过 check_ipc_quant_support 硬性检查量化方法是否在白名单内。连接守护进程的 Unix socket,发送引擎的 CacheConfig 进行匹配验证。成功后接收 IPC 句柄列表,使用 meta device 初始化模型(_initialize_modeldevice_config='meta'),然后遍历句柄通过 torch.Tensor._new_shared_cuda 映射为 GPU 张量并赋值给对应参数/缓冲区。如果守护进程不可用但 socket 文件不存在,允许回退到磁盘加载;所有其他失败(connection refused、配置不匹配、协议错误)均硬性报错。

  4. 引擎集成:在 engine.py 中新增 _launch_weight_cache_daemons_terminate_weight_cache_daemons 类方法。前者根据 TP/PP 划分启动相应数量的守护进程子进程(使用 subprocess.Popen 避免父进程 CUDA 初始化干扰),并等待所有 .ready 文件就绪。shutdown() 负责优雅终止守护进程(先 SIGTERM 释放文件再 SIGKILL 剩余进程),确保 socket/ready 文件被清理。新增 --weight-cache-mode 参数(off/daemon/client)、--weight-cache-socket 等配置。

  5. 安全与兼容性机制:实现 IPC_QUANT_ALLOWLIST 在守护进程和客户端两侧都硬性检查。检测 PYTORCH_CUDA_ALLOC_CONF 中的 expandable_segments:True 并提前失败。使用 PR_SET_PDEATHSIG 确保守护进程在父进程崩溃时自动终止。在权重更新路径(update_weights_from_disk_distributed_tensor_ipc)中插入 _assert_weight_cache_inactive 防止零拷贝模式下的数据损坏,release_memory_occupation 中对应权重的分支也被保护。

  6. 测试配套:新增 CPU 单元测试 test_weight_cache_protocol.py,覆盖协议帧协议、CacheConfig 匹配、量化哈希、白名单检查等,运行快速且无需 GPU。新增 GPU E2E 测试 test_weight_cache_daemon.py,包含 TP=2 的完整守护进程+客户端生成测试和 TP=1 的快速冒烟测试,通过日志断言(如 "Fetched ... IPC handles from daemon")验证 IPC 路径被实际使用。

文件 模块 状态 重要度
python/sglang/srt/weight_cache/daemon.py 权重守护 added 9.25
python/sglang/srt/weight_cache/ipc_loader.py IPC 加载 added 9.25
python/sglang/srt/weight_cache/protocol.py 缓存协议 added 9.25
python/sglang/srt/entrypoints/engine.py 引擎入口 modified 8.58
python/sglang/srt/server_args.py 服务配置 modified 6.99
python/sglang/srt/model_executor/model_runner_components/load_model_utils.py 模型加载 modified 7.54
python/sglang/srt/model_loader/loader.py 加载器 modified 6.81
python/sglang/srt/model_executor/model_runner_components/weight_updater.py 权重更新 modified 7.07
python/sglang/srt/managers/scheduler_components/weight_updater.py 调度权重 modified 7.05
test/registered/unit/model_loader/test_weight_cache_protocol.py 协议测试 added 8.42
test/registered/model_loading/test_weight_cache_daemon.py 守护测试 added 8.42

关键符号

WeightCacheDaemon.__init__ WeightCacheDaemon._init_distributed WeightCacheDaemon.load WeightCacheDaemon._assert_ipc_compatible_allocator WeightCacheDaemon._export_state WeightCacheDaemon.serve WeightCacheDaemon._signal_handler IpcModelLoader.__init__ IpcModelLoader.load_model IpcModelLoader._start_daemon_liveness_watchdog IpcModelLoader._daemon_alive IpcModelLoader._watch IpcModelLoader._resolve_engine_quant IpcModelLoader._rebuild_stale_views CacheConfig.matches CacheConfig.to_dict CacheConfig.from_dict hash_quant_config get_quant_method_name check_ipc_quant_support UnsupportedQuantForIPCError _get_quant_field engine._launch_weight_cache_daemons engine._terminate_weight_cache_daemons maybe_enable_ipc_weight_cache weight_updater._assert_weight_cache_inactive

关键源码片段

python/sglang/srt/weight_cache/protocol.py data-contract

协议定义和工具函数,包含 CacheConfig 指纹结构、IPC 量化白名单、环境戳计算、socket 消息帧协议、脏文件清理等。是整个框架的契约基础。

# SPDX-License-Identifier: Apache-2.0
"""权重缓存守护进程的协议定义。
定义 CacheConfig 指纹、socket 消息帧协议、IPC 量化白名单等。
"""class CacheConfig(msgspec.Struct):
    """缓存权重的指纹,用于验证守护进程与引擎之间的兼容性。
    任何不匹配都会触发磁盘 fallback 或硬性报错。
    """
    model_path: str # 模型路径
    model_arch: str # 模型架构(如 LlamaForCausalLM)
    tp_size: int # 张量并行度
    tp_rank: int # 当前 TP rank
    pp_size: int # 流水线并行度
    pp_rank: int # 当前 PP rank
    dp_size: int # 数据并行度
    ep_size: int # 专家并行度
    quant_method: str # 量化方法,如 "fp8"、"gptq_marlin"、"" 无量化
    quant_config_hash: str # 量化配置的 SHA-256 哈希
    dtype: str # 例如 "torch.float16"
    revision: str # 模型版本("" 表示未设置)
    # 环境戳:不同 GPU 算力或 torch 版本可能产生不兼容的权重
    device_capability: str # 本地计算能力,如 "8.0"
    torch_version: str # torch.__version__
​
    def matches(self, other: "CacheConfig") -> bool:
        """检查两个配置是否兼容。"""
        return self == other
​
    def to_dict(self) -> Dict[str, Any]:
        return {f: getattr(self, f) for f in self.__struct_fields__}
​
    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "CacheConfig":
        return cls(**d)
​
​
# ---------------------------------------------------------------------------
# IPC 量化方法白名单
# ---------------------------------------------------------------------------
# CUDA IPC 零拷贝共享只导出原始张量数据,因此仅在
# process_weights_after_loading 的完整效果被张量数据捕获时才正确。
# 需要 Python 元数据(如 block-FP8 的 format_ue8m0)或重排 / 转置的方法
# (per-tensor FP8、Marlin、AWQ/GPTQ)会提供静默错误。
# 只有验证可 round-trip 通过纯张量导出的方法才被列入白名单。
IPC_QUANT_ALLOWLIST: Dict[str, List[str]] = {
    # "" 表示未量化的模型
    "": ["*"],
    # block-wise FP8:仅当 weight_block_size 已设置(非默认)
    "fp8": ["*"],
}
​
​
def check_ipc_quant_support(quant_method: str, quant_config: Any, where: str):
    """检查量化方法是否支持 IPC 零拷贝。不支持时硬性报错。"""
    if quant_method not in IPC_QUANT_ALLOWLIST:
        raise UnsupportedQuantForIPCError(
            f"Quantization method '{quant_method}' is not supported for IPC "
            f"weight cache (checked in {where}). Supported methods: "
            f"{list(IPC_QUANT_ALLOWLIST.keys())}. "
            f"Use `--quantization fp8` with `--weight-block-size` or load "
            f"without quantization for IPC mode."
        )
​
​
def hash_quant_config(quant_config: Any) -> str:
    """计算量化配置的稳定哈希。
    避免使用 str()/repr() 因为嵌入内存地址会导致不同进程产生不同哈希。
    """
    if quant_config is None:
        return ""
    try:
        if hasattr(quant_config, "to_dict"):
            config_str = json.dumps(quant_config.to_dict(), sort_keys=True)
        elif isinstance(quant_config, dict):
            config_str = json.dumps(quant_config, sort_keys=True)
        elif hasattr(quant_config, "__dict__"):
            config_str = type(quant_config).__name__ + ":" + json.dumps(
                {k: v for k, v in sorted(quant_config.__dict__.items())
                 if not k.startswith("_")
                 and isinstance(v, (str, int, float, bool, type(None), list, dict))},
                sort_keys=True)
        else:
            config_str = type(quant_config).__name__
        return hashlib.sha256(config_str.encode()).hexdigest()
    except Exception:
        config_str = type(quant_config).__name__
        return hashlib.sha256(config_str.encode()).hexdigest()

评论区精华

IPC 量化方法兼容性 正确性

跳过 `_post_load_weights` 后某些量化方法(如 per-tensor FP8)的 Python 属性(如 `weight_scale_inv.format_ue8m0`)不会跨 IPC 传递,导致静默输出错误。

结论:引入 `IPC_QUANT_ALLOWLIST` 白名单,仅允许 unquantized 和 block-wise FP8(weight_block_size 已设置)通过 IPC 加载,其他方法硬性报错。 · 已解决

守护进程孤儿进程与脏文件清理 安全

引擎崩溃后守护进程成为孤儿,残留 `.ready`/`.sock` 文件阻止下次启动。engine-spawned daemon 无 pdeathsig,父进程 SIGKILL 后 daemon 成为孤儿且不清理文件。

结论:添加 `PR_SET_PDEATHSIG`(`kill_itself_when_parent_died`);`shutdown()` 先 SIGTERM 后 SIGKILL;`cleanup_stale_daemon_files` 添加 `--force` 标志手动接管。 · 已解决

连接超时与 socket 泄漏 正确性

守护进程 accept 后未设置超时,一个挂起客户端会阻塞所有其他引擎等级;socket 连接失败时未关闭导致泄漏。

结论:添加 `conn.settimeout(30s)`;修复 socket 泄漏(使用 `with` 或 `try/finally`)。 · 已解决

扩展段内存分配器不兼容 性能

`PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` 使 `_share_cuda_` 失败,导致 daemon 导出过程 mid-way 崩溃。

结论:添加 `_assert_ipc_compatible_allocator()` 在 load 开始时检测并提前报错。 · 已解决

权重在线更新保护 正确性

IPC 零拷贝模式下 `update_weights_from_disk` 等操作会 in-place 修改 `param.data`,从而损坏守护进程的 master 副本和所有 co-attached 引擎。

结论:在 `weight_updater.py` 的四个入口(disk/distributed/tensor/ipc)添加 `_assert_weight_cache_inactive` 硬性报错;`scheduler_components/weight_updater.py` 的 `release_memory_occupation` 的权重分支也添加保护。 · 已解决

PP/EP/DP 参数传递遗漏 正确性

引擎侧 `_launch_weight_cache_daemons` 未传递 `pp_size`/`dp_size`/`ep_size` 给守护进程,导致配置指纹不匹配。

结论:已修复为从 `server_args` 传递所有维度参数给 daemon 构造器。 · 已解决

测试覆盖计划 测试

初始版本仅手动测试(`test/manual/test_weight_cache_e2e.py`),缺少 CI 覆盖。

结论:添加 CPU 单元测试(`test_weight_cache_protocol.py`,无 GPU)和 GPU E2E 测试(`test_weight_cache_daemon.py`,TP=1 冒烟测试始终运行在 base-b 阶段)。 · 已解决

全局参数 vs 实际并行配置不一致 正确性

Claude 指出 daemon 初始化时使用全局 `ServerArgs` 但 `pp_size`/`dp_size`/`nnodes` 未被正确传递,可能导致 `get_server_args().pp_size` 看到错误的值。

结论:daemon 构造器接受独立的 `pp_size`/`dp_size`/`ep_size` 参数,由调用方从实际配置传入,不再依赖全局 ServerArgs。 · 已解决

风险与影响

  • 兼容性风险:IPC 零拷贝仅支持 unquantized 和 block-wise FP8(weight_block_size 已设置),其他量化方法(per-tensor FP8、Marlin、AWQ/GPTQ 等)在守护进程和客户端两侧都会硬性报错,不存在静默错误。环境戳(CUDA 算力、torch 版本)不匹配也会导致拒绝服务,但这是保护机制而非风险。
  • 进程生命周期风险:尽管添加了 PR_SET_PDEATHSIG 保护,在极端场景(如父进程被 SIGKILL 后子进程被 reparent 到 init)仍可能产生 orphan 守护进程,残留 GPU 内存。--force 标志提供了手动恢复手段。
  • 内存安全风险:零拷贝模式下引擎和守护进程共享 GPU 内存,已通过 _assert_weight_cache_inactive 禁止权重更新操作,并强制 requires_grad=False 防止 autograd 意外修改。但若第三方代码直接操作 param.data 仍可能破坏共享数据。
  • 性能风险--weight-cache-mode daemon 模式下首次启动比直接加载更慢(因为多一次磁盘加载到守护进程)。只有在守护进程预启动(standalone launcher)并使用 --weight-cache-mode client 时才能体现加速优势,需在文档中明确说明。
  • 安全风险:Unix socket 文件位于 /tmp/,同一节点上其他进程可尝试连接。建议限制 socket 文件权限(当前实现未显示设置)。
  • 分布式初始化风险:多节点场景下若未设置 --dist-init-addr,守护进程无法形成联合 NCCL 组,引擎会陷入 readiness 超时(默认 1800s)。已在 engine.py 中添加对应验证,提前报错。
  • 用户影响:主要面向生产环境部署用户。通过 --weight-cache-mode client 配合预启动守护进程可大幅缩短引擎重启时间(从分钟级到秒级),需增加守护进程的部署和管理。默认 off,不影响现有用户。
  • 系统影响:每个 GPU 上多运行一个守护进程(额外 GPU 内存占用:零拷贝模式下引擎和守护进程共享同一份权重,不增加整体内存使用,但守护进程本身需常驻 GPU 上下文)。CPU 和网络开销极小。
  • 团队影响:新增约 2000 行核心代码和 1000 行测试,需要维护与现有特性(推测解码、睡眠模式、多节点、弹性 EP 等)的兼容性。后续 Roadmap(CUDA 图序列化、内核预热缓存等)依赖此基础结构。
  • 兼容性影响:未对现有 API 和协议做破坏性变更,参数均为新增可选参数。需确保与 model_runner.pyweight_updater.py 等模块的交互正确。
量化白名单限制 守护进程孤儿风险 扩展段不兼容 共享内存更新保护 首次启动延迟 分布式初始化超时

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论