已实现

v07_cuda_graph

本版本引入什么

本版本使用 CUDA Graph 捕获静态解码路径,并在后续解码步骤中重放。

引擎仍以普通方式执行动态预填充。

为什么引入

解码通常会为每个生成词元启动许多小型 GPU 操作,Python 和 CUDA 启动开销可能占据延迟的很大部分。CUDA Graph 通过捕获固定 GPU 工作序列并重放来降低这项开销。

核心原则

CUDA Graph 捕获针对静态解码张量的操作。后续步骤原地更新这些张量,再重放已捕获的图。

它的约束很严格:

  • 张量地址相同。
  • 形状相同。
  • 捕获区域中的控制流相同。
  • 捕获路径内没有动态分配。

建议对比的文件

  • engine.py:预热、捕获和重放。
  • scheduler.py:原地修改静态解码输入。
  • kvcache.py:固定槽位缓存访问。
  • layer/gqa.py:在静态槽位形状输入上执行解码注意力。

保留的权衡

本版本在分页 KV 缓存出现前展示 CUDA Graph,从而单独呈现图约束,但缓存布局仍不是最终面向服务的内存形态。