执行摘要
本PR修复了Piecewise CUDA Graph (PCG)模式下注意力计算中填充令牌被错误处理的问题，通过切片张量排除填充令牌，避免FlashInfer后端产生未定义行为（如NaN、输出损坏）。修复后，先前禁用的测试test_qwen3_next_models_pcg重新启用并通过，提升了PCG的鲁棒性和推理可靠性。
功能与动机
为什么做：当PCG启用时，注意力元数据使用实际令牌数（real_num_tokens）初始化，但输入张量仍包含填充令牌，导致FlashInfer等注意力后端无法正确处理，可能引发NaN值、损坏的输出（如重复“!!!!!”）或异常输出长度。PR body中明确指出“To fix this, exclude the padded tokens and make PCG more robust。”
实现拆解
改动模块：
- 注意力计算层：
radix_attention.py和radix_linear_attention.py中，unified_attention_with_output和unified_linear_attention_with_output函数新增切片逻辑：
- 使用
forward_batch.num_token_non_padded_cpu获取实际令牌数。
- 将query、key、value等张量切片到该长度，排除填充令牌。
- 动态修改和恢复
out_cache_loc，确保缓存写入正确位置。
python
real_num_tokens = forward_batch.num_token_non_padded_cpu
query = query[:real_num_tokens]
forward_batch.out_cache_loc = original_out_cache_loc[:real_num_tokens]
- 后端清理：
flashinfer_backend.py移除PCG填充相关代码（如extra_kv和pad_tokens逻辑），简化实现。
- PCG运行器：
piecewise_cuda_graph_runner.py添加num_token_non_padded_cpu参数传递，确保实际令牌数在重放时可用。
- 上下文管理：
piecewise_context_manager.py删除num_tokens字段，减少冗余。
评论区精华
讨论焦点：
- 字段命名与注释：
gemini-code-assist[bot]建议更新real_num_tokens注释，ch-wan指出其等同于num_token_non_padded，作者采纳并重命名为num_token_non_padded_cpu，体现代码清晰度优化。
- 输出缓冲区初始化：
Oasis-Git和hebiao064讨论是否用zeros代替empty，结论是移除零初始化以保持设计简洁，避免不必要更改。
- sinks形状处理：
ispobock提醒sinks可能无令牌维度，作者更新代码保持sinks不变，确保模型兼容性。
风险与影响
技术风险：
- 回归风险：核心注意力路径变更可能影响所有PCG模式下的模型推理，需依赖CI测试覆盖。
- 性能开销：添加切片操作可能引入微小延迟，但相比未定义行为修复，利大于弊。
- 测试覆盖不足：边缘情况（如不同填充场景）可能未充分测试，建议补充单元测试。
影响评估：
- 用户：修复输出损坏问题，提升推理可靠性和确定性，尤其对Qwen3-next等模型用户有益。
- 系统：PCG路径更稳健，减少未定义行为，可能间接改善性能。
- 团队：代码简化便于维护，移除冗余逻辑降低未来错误概率。
关联脉络
历史PR关联：
- PR #21452：被本PR回滚，原PR可能引入了填充处理逻辑，但本PR提供了更彻底的修复。
- PR #17404：评论中提及修复了Mamba缓存问题，与本PR共同提升PCG稳定性。
- 近期PR趋势：仓库近期多个PR涉及PCG、填充和模型特定修复（如PR #22739），显示团队持续优化推理路径的稳健性和性能。
演进方向：本PR是PCG优化线的一部分，旨在通过简化逻辑和排除填充令牌，提升大规模模型推理的确定性和效率。
参与讨论