执行摘要
- 一句话：修复CUDA Graph捕获时num_token_non_padded计算逻辑，确保捕获与重放行为一致。
- 推荐动作：建议CUDA Graph和attention TP相关开发者精读此PR，理解捕获与重放路径对齐的设计决策。关注条件判断逻辑和compute_local_num_token_non_padded函数的实现，确保在不同配置下行为正确。
功能与动机
根据PR body中的描述，修复动机是"Correctly compute the num-non-padded-tokens during graph capture, which align with the behavior in replay"，并引用了代码行https://github.com/sgl-project/sglang/blob/1b7c33a5b751dac6187367d798a7b80bd12ccaaf/python/sglang/srt/model_executor/cuda_graph_runner.py#L326。这表明需要确保CUDA Graph捕获时的计算逻辑与重放阶段保持一致，避免因不一致导致的问题。
实现拆解
该PR仅修改了一个文件python/sglang/srt/model_executor/cuda_graph_runner.py。在capture_one_batch_size函数中，原本只是简单地将buffers.num_token_non_padded[...]设置为num_tokens。现在增加了条件判断逻辑：当启用enable_num_token_non_padded且需要gathered buffer且未启用NSA预填充检查点时，调用compute_local_num_token_non_padded函数计算本地token数量，并复制到buffers.num_token_non_padded中。这确保了捕获路径与重放路径中的populate_from_forward_batch函数行为一致。
关键文件：
python/sglang/srt/model_executor/cuda_graph_runner.py（模块 model_executor）: 这是唯一修改的文件，包含了CUDA Graph捕获的核心逻辑修复，直接影响图执行的一致性。
关键符号：capture_one_batch_size, compute_local_num_token_non_padded
评论区精华
由于review评论为空，仅有的交互是ispobock的批准和触发CI的评论。从上下文看，这个修复相对直接，没有引发深入的技术讨论。ispobock作为合并者，通过/tag-and-rerun-ci触发了CI测试，表明需要验证修复的正确性。
- CUDA Graph捕获与重放一致性修复 (correctness): 通过添加条件判断和调用compute_local_num_token_non_padded函数来修复。
风险与影响
- 风险：风险较低但需关注：1. 核心路径变更：修改了CUDA Graph捕获的关键逻辑，如果条件判断有误或compute_local_num_token_non_padded函数行为不符合预期，可能导致图捕获失败或重放时行为异常。2. 条件复杂性：新增的条件涉及多个标志（enable_num_token_non_padded、require_gathered_buffer、nsa_enable_prefill_cp），需要确保这些标志在捕获和重放时状态一致。3. 缺少测试覆盖：从PR body看，作者未提供准确性测试或性能测试结果，依赖CI测试验证。
- 影响：影响范围有限但重要：1. 对系统：修复CUDA Graph执行的一致性问题，确保捕获的图能正确重放，提升推理稳定性。2. 对用户：透明修复，不会改变API或可见行为，但能避免潜在的图执行错误。3. 对团队：涉及CUDA Graph和attention TP的交互，需要相关开发者关注此修复的逻辑。
- 风险标记：核心路径变更, 条件复杂性, 缺少测试覆盖
关联脉络
- PR #21104 perf: precompute FA3 scheduler_metadata to eliminate per-layer prepare_varlen_num_blocks: 同样涉及CUDA Graph和性能优化，可能共享类似的图捕获逻辑。
- PR #22051 [MUSA][9/N] Add FA3 attention backend support through MATE (MUSA AI Tensor Engine): 涉及attention后端支持，与本PR的attention TP调整相关。
参与讨论