执行摘要

本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内核优化）也显示团队在核心路径上的改进趋势。