执行摘要
此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场景的性能，减少运行时开销。
实现拆解
关键改动模块：
- model_runner.py：
- 添加持久
self.intermediate_tensors缓冲，为非首PP rank预分配内存，确保图形重放时地址稳定。
- 更新
capture_model函数，传入中间张量以支持捕获。
- 在
execute_model中，实现从接收张量到缓冲的复制逻辑，示例代码片段：
python
n = input_batch.num_tokens_after_padding
for k, v in intermediate_tensors.tensors.items():
self.intermediate_tensors[k][:n].copy_(v[:n])
- cudagraph_utils.py：
- 扩展
capture函数，添加PP rank状态判断和中间张量参数。
- 引入
num_reqs调整逻辑，确保num_tokens可被整除，以避免TRTLLM decode断言错误，代码片段：
python
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讨论中聚焦于两个关键点：
- 中间张量键一致性：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."
- 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支持。
参与讨论