Prhub

#19890 [Disagg] GPU staging buffer with dynamic ring allocator for heterogeneous TP KV transfer

原始 PR 作者 YAMY1234 合并时间 2026-04-02 05:09 文件变更 9 提交数 56 评论 62 代码增减 +2117 / -8

feature performance scheduling

执行摘要

为异构 TP KV 传输引入 GPU 暂存缓冲区和动态环形分配器，提升高并发下传输吞吐量。

PR body指出，'When prefill and decode use different TP layouts in disaggregated serving (e.g., TP4 prefill + DEP4 decode), KV cache transfer requires head slice redistribution. The original per-token RDMA approach (send_kvcache_slice) generates ~30,000 small RDMA requests per request, severely limiting transfer throughput at high concurrency.' 因此，需要staging buffer来批量传输，提升吞吐量。

该PR值得精读，展示了高性能异构传输的关键设计决策，如GPU kernels和动态分配器。关注staging_handler的设计以学习生命周期管理，以及review中的设计权衡讨论，如字段移动和代码清理。

讨论亮点

review中核心讨论包括：ShangmingCai建议将staging字段移至CommonKVReceiver中以支持多prefill实例（设计优化）；代码清理方面，建议使用直接env变量检查而非getattr，提升可读性（YAMY1234采纳并调整）；性能讨论中，YAMY1234解释了多queue设计以避免单线程瓶颈；命名改进，如重命名为PrefillStagingContext以更好描述作用域。

实现拆解

实现分为几个关键部分：1) 新增staging_buffer.py，包含Triton kernels（_fused_gather_to_staging_kernel和_fused_scatter_from_staging_kernel）用于GPU gather/scatter操作；2) 新增staging_handler.py，定义DecodeStagingHandler和PrefillStagingContext，管理staging生命周期；3) 修改mooncake/conn.py，集成send_kvcache_staged方法处理staging传输路径；4) 修改decode.py和prefill.py，添加staging相关逻辑，如_poll_with_staging和maybe_prefetch_staging_for_batch；5) 在environ.py中添加环境变量SGLANG_DISAGG_STAGING_BUFFER控制功能开关；6) 在server_args.py中验证staging仅与mooncake后端兼容。

文件 模块 状态 重要度
python/sglang/srt/disaggregation/common/staging_buffer.py disaggregation added 10.0
python/sglang/srt/disaggregation/common/staging_handler.py disaggregation added 9.0
python/sglang/srt/disaggregation/mooncake/conn.py disaggregation modified 8.0
python/sglang/srt/disaggregation/decode.py disaggregation modified 7.0
python/sglang/srt/disaggregation/prefill.py disaggregation modified 6.0

分析完成后，这里会展示 LLM 生成的相对完整源码片段和详细注释。

关键符号

_fused_gather_to_staging_kernel _fused_scatter_from_staging_kernel DecodeStagingHandler.advance_scatter send_kvcache_staged maybe_prefetch_staging_for_batch

评论区精华

将 staging 字段移至 kv_receiver 以支持多 prefill 实例 设计

ShangmingCai 建议在 CommonKVReceiver 中添加 require_staging 字段，而不是在 DecodeRequest 中，以灵活处理不同 TP 大小

结论：采纳，移除 prefill_attn_tp_size 字段，使用 prefill_info.attn_tp_size 计算 require_staging · 已解决

使用 enable_staging 标志清理代码，避免 getattr hack 设计

ShangmingCai 建议直接使用 env 变量 SGLANG_DISAGG_STAGING_BUFFER 检查，提升可读性和维护性

结论：采纳，在 server_args.py 添加验证，确保仅 mooncake 后端支持 staging · 已解决

queue 大小对性能的影响，讨论单队列与多队列设计 性能

ShangmingCai 询问固定 queue 大小为 1 是否影响性能，YAMY1234 解释多 queue 可避免单线程调度瓶颈

结论：保持多 queue 设计，不改变 SGLANG_DISAGGREGATION_QUEUE_SIZE · 已解决

风险与影响

技术风险具体包括：1) 兼容性风险：staging仅支持mooncake传输后端，其他后端会报错（server_args.py中验证）；2) 内存风险：动态环形分配器在高并发下可能OOM，但通过分块传输缓解；3) 正确性风险：长上下文场景buffer不足，PR通过per-chunk传输解决；4) 性能风险：多线程调度可能引入延迟，需优化（如queue大小讨论）。

影响范围：对用户，异构TP传输吞吐量提升2-3倍，接近同构性能（基准测试显示）；对系统，新增GPU缓冲区和环境变量，增加内存使用但可控（默认64MB/4096MB）；对团队，代码库增加新模块，需维护staging相关逻辑，但设计隔离良好（staging_handler.py）。

仅限 mooncake 后端 内存使用增加 长上下文需分块 多线程复杂性

关联 Issue

未识别关联 Issue

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

完整报告

Markdown PDF

执行摘要

本PR为sglang的disaggregated serving引入GPU staging buffer机制，针对异构TP布局下的KV传输优化。通过将分散的head slices收集到连续内存进行批量RDMA传输，减少请求数量约1000倍，提升高并发下吞吐量2-3倍。实现包括动态环形分配器和专用kernels，设计隔离良好，但仅支持mooncake后端。

功能与动机

在异构TP场景（如TP4 prefill + DEP4 decode），原有每令牌RDMA方法生成约30,000小请求，限制吞吐量。如PR body所述：'severely limiting transfer throughput at high concurrency'。新机制旨在通过批量传输解决此瓶颈，基准测试显示在128并发下吞吐量从2,466 tok/s提升至4,914 tok/s。

实现拆解

  • staging_buffer.py: 新增Triton kernels _fused_gather_to_staging_kernel 和 _fused_scatter_from_staging_kernel，执行GPU gather/scatter，内核延迟低至微秒级。
  • staging_handler.py: 定义 DecodeStagingHandler 和 PrefillStagingContext，管理staging生命周期，隔离于decode.py和conn.py。
  • mooncake/conn.py: 集成 send_kvcache_staged 方法，处理staging传输路径，更新bootstrap协议支持动态偏移。
  • decode.py: 添加 _poll_with_staging 函数（后移至utils.py），处理staging感知的轮询，异步scatter在专用CUDA流执行。
  • prefill.py: 添加 maybe_prefetch_staging_for_batch 预取逻辑，优化decode侧内存分配。
  • environ.py: 新增环境变量 SGLANG_DISAGG_STAGING_BUFFER 控制功能开关，默认禁用。

评论区精华

  • 设计优化: ShangmingCai建议：'Can we add an attr in PrefillServerInfo, let us call it require_staging'，最终将逻辑移至 CommonKVReceiver，提升多实例支持。YAMY1234回复：'Fully agree. I’ve updated the code to support cases where prefill uses different TP sizes'。
  • 代码可读性: 讨论中强调使用 enable_staging 直接检查，避免 getattr hack，如ShangmingCai评论：'I think we should let prefill/decode both have this attr'。YAMY1234采纳并调整代码。
  • 性能权衡: 关于queue大小，ShangmingCai询问：'Have you checked whether fixing the queue size to 1 affects the performance?' YAMY1234解释：'With a single queue there's only one dispatcher... can become a single-threaded bottleneck'，决定保持多queue设计。

风险与影响

  • 风险:
    1. 兼容性：staging仅支持mooncake传输后端，其他后端会报错（server_args.py验证）。
    2. 内存：动态环形分配器可能在高并发下OOM，但通过分块传输和默认4GB池缓解。
    3. 正确性：长上下文场景buffer不足，PR通过per-chunk传输解决，已验证128k ISL场景。
    4. 性能：多线程调度增加复杂性，但基准测试显示显著提升。
  • 影响:
    • 用户：异构TP传输吞吐量提升2-3倍，接近同构性能（基准测试数据）。
    • 系统：新增GPU缓冲区和环境变量，增加内存使用但可控（环境变量可调）。
    • 团队：代码库增加新模块，需维护staging相关逻辑，但设计隔离良好（staging_handler.py）。

关联脉络

与历史PR #18158 和 #18163 相关，后者曾尝试类似方法但未解决长上下文问题，如issue评论中Ch3ngY1提到：'in long-context scenarios, the buffer can hardly accommodate the entire context's kv cache'。本PR通过分块传输和动态分配器克服限制，体现了disaggregation性能优化的持续演进。近期PR如#21233（清理冗余代码）和#21834（JIT内核优化）也显示团队在核心路径上的改进趋势。

参与讨论