执行摘要

此PR为vllm-project/vllm仓库的V2模型runner添加了流水线并行下的piecewise CUDA graph支持，解决了此前PP场景只能使用eager模式导致的性能瓶颈。通过引入持久中间张量缓冲和PP-aware捕获逻辑，性能吞吐量提升约66%，TTFT和TPOT显著改善，使V2模型runner在PP下性能对齐V1基线。PR包含关键设计权衡，如num_reqs调整workaround，值得工程师精读以理解分布式CUDA graph优化。

功能与动机

为什么做：V2模型runner在流水线并行启用时无法使用CUDA graph捕获，只能回退到eager模式，限制了推理性能。PR作者在body中明确表示："V2 model runner did not support CUDA graph capture with PP, falling back to eager mode. This PR adds piecewise CUDA graph capture for PP." 关联Issue #33960可能提供了更多背景。目标是通过启用CUDA graph捕获来提升PP场景的性能，减少运行时开销。

实现拆解

关键改动模块：

1. model_runner.py：
   • 添加持久self.intermediate_tensors缓冲，为非首PP rank预分配内存，确保图形重放时地址稳定。
   • 更新capture_model函数，传入中间张量以支持捕获。
   • 在execute_model中，实现从接收张量到缓冲的复制逻辑，示例代码片段：
     

     n = input_batch.num_tokens_after_padding
     for k, v in intermediate_tensors.tensors.items():
         self.intermediate_tensors[k][:n].copy_(v[:n])
     

2. cudagraph_utils.py：
   • 扩展capture函数，添加PP rank状态判断和中间张量参数。
   • 引入num_reqs调整逻辑，确保num_tokens可被整除，以避免TRTLLM decode断言错误，代码片段：
     

     if num_reqs > 0 and num_tokens > num_reqs and num_tokens % num_reqs != 0:
         tokens_per_req = cdiv(num_tokens, num_reqs)
         num_reqs = num_tokens // tokens_per_req
         if num_tokens % num_reqs != 0:
             num_reqs = 1
     

评论区精华

Review讨论中聚焦于两个关键点：

1. 中间张量键一致性：gemini-code-assist[bot]建议添加assertion以确保键匹配，但代码未实现，留下潜在风险。

   "The code assumes that the keys in intermediate_tensors.tensors received from the previous pipeline stage are identical to the keys in the persistent self.intermediate_tensors.tensors... To improve robustness... it's safer to assert that the sets of keys are identical."

2. num_reqs调整设计：yewentao256询问调整原因，ZhanqiuHu解释为应对TRTLLM decode断言，但不确定是否为正确方法。

   "With CUDA graph capture, I ran into AssertionError: TRTLLM decode requires uniform query lengths per request... So I added this workaround... but not sure if this is the right approach."
   WoosukKwon批准PR并提及将跟进FULL graph支持，表明此PR是阶段性改进。

风险与影响

技术风险：

• 中间张量键不匹配可能导致运行时错误或静默数据损坏。
• num_reqs调整是workaround，可能在某些场景下影响性能或引入非预期行为。
• PP下CUDA graph捕获首次启用，需警惕内存对齐、跨rank同步等边缘情况。

影响分析：

• 性能提升：根据测试结果，吞吐量从13.89 req/s提升至23.07 req/s，TTFT从231ms降至167ms，TPOT从17.5ms降至10.4ms，显著改善用户体验和系统效率。
• 系统演进：为V2模型runner添加核心功能，促进向新架构迁移，团队可借鉴中间张量管理设计。

关联脉络

与历史PR的关系：

• PR #34903被yewentao256在issue评论中提及，同为V2模型runner的CUDA graph支持PR，可能涉及full graph功能，可作为后续参考。
• 从近期历史PR看，仓库持续优化性能（如FP8 kernel、ROCm改进），此PR延续了性能提升趋势，聚焦于CUDA graph在分布式场景的应用。
  整体上，此PR是V2模型runner演进中的关键一步，为流水线并行提供了高效的图形捕获方案，后续可能扩展至full graph支持。