执行摘要
此PR为SGLang的PD disaggregation功能中的异构TP GPU暂存缓冲添加了端到端CI测试和详细文档，覆盖MHA模型在不同TP配置下的验证，并修复了代码中的MLA模型检查。通过测试增强功能可靠性，文档提升用户指导，整体为有意义的改进，适合工程师关注测试设计和配置细节。
功能与动机
动机源于需要验证和文档化异构TP（如prefill TP=4, decode TP=1）场景下的GPU暂存缓冲功能，以提升KV缓存传输吞吐量。PR body明确指出：'Add e2e test for disaggregation with staging buffer enabled ... Covers MLA and MHA models with both prefill-larger and decode-larger TP configurations.' 目标是确保该功能在复杂配置下正常工作，并提供清晰的使用指南。
实现拆解
实现分为三个主要部分：
- 文档更新：在
docs/advanced_features/pd_disaggregation.md中添加暂存缓冲章节，包括功能描述、环境变量列表和使用示例；在docs/references/environment_variables.md中补充相关环境变量。
- 代码修改：在
decode.py和prefill.py的__init__方法中添加MLA模型检查，抛出RuntimeError防止误用；在staging_buffer.py和staging_handler.py中修改日志消息，避免'NVLink incompatible!'等误导性警告。
- 测试添加：在
test/registered/distributed/test_disaggregation_different_tp.py中新增测试类TestDisaggregationStagingPrefillLargerTP和TestDisaggregationStagingDecodeLargerTP，使用环境变量SGLANG_DISAGG_STAGING_BUFFER=1进行端到端测试。
关键代码片段（来自decode.py）：
if self.enable_staging and self.is_mla_backend:
raise RuntimeError(
"SGLANG_DISAGG_STAGING_BUFFER is designed for non-MLA models "
"(e.g. GQA, MHA). MLA models should not set this flag."
)
评论区精华
review讨论中聚焦于两点：
- 文档准确性：gemini-code-assist[bot]指出示例中decode TP设置应改为TP=1以正确展示异构TP，作者回应已修正。
- 测试代码质量：ShangmingCai建议将测试合并到现有文件以减少重复，作者提交'Merge staging buffer tests into existing different-TP test file'解决。
引用讨论要点：
gemini-code-assist[bot]: "The usage example for the decode server uses --tp 4, which matches the prefill server's --tp 4. ... To correctly demonstrate the heterogeneous TP staging buffer feature, the decode server should use a different TP size."
ShangmingCai: "Is it possible that we put these tests in the different tp test file, instead of creating another file? I mean, they are testing one same feature after all."
风险与影响
风险：
- 文档示例错误可能误导用户配置不当，导致功能无效或性能下降。
- MLA模型检查可能影响现有用户，需调整环境变量以避免运行时错误。
- 测试仅覆盖MHA模型，未涵盖所有潜在场景，可能存在遗漏的bug。
- 日志消息修改可能隐藏硬件兼容性问题，在特定环境下引发性能隐患。
影响：
- 用户受益于清晰的文档，能更高效地启用暂存缓冲功能，提升系统吞吐量。
- 系统通过测试增强稳定性，代码变更确保MLA模型安全，减少生产事故风险。
- 团队获得标准化测试用例，便于后续功能扩展和维护，提升开发效率。
关联脉络
从近期历史PR看，此PR与测试和文档相关PR（如#22176修复测试导入）有相似模式，均聚焦于CI流程和代码一致性。暂存缓冲功能可能在前序PR中引入，但此PR专门针对测试验证和用户指导，体现了SGLang项目对异构TP优化的持续投入。结合讨论，团队倾向于将测试整合到现有文件以降低维护成本，这反映了代码重构和测试策略的演进趋势。
参与讨论