PolyGraph：自适应推测解码的CUDA图优化方案

摘要：PolyGraph提出了一种基于计划缓冲机制的CUDA图优化方案，用于自适应推测解码。该方案通过捕获最大网格图并在每次重放前将活动形状写入GPU内存，显著减少了构建时间并避免了填充开销。在单个NVIDIA H800 GPU上的实验表明，PolyGraph将预服务图构建时间从3915毫秒降至93毫秒，同时保持了与基线相当的延迟表现（17.7 vs 17.8毫秒）。这项研究为自定义Triton内核提供了一种高效的工作形状管理机制，但目前仅限于单GPU的合成轨迹测试。

---

技术机制剖析

PolyGraph旨在解决自适应推测解码中CUDA图构建的效率问题。传统方法中，CUDA图需要为不同的活动形状构建多个图，这导致预服务图构建时间过长。PolyGraph通过以下机制进行优化：

1. 最大网格图捕获：PolyGraph捕获一个最大网格图，并在每次重放前将当前活动形状写入GPU内存。
2. 跳过非活动区域：内核在执行时跳过超出活动形状的tiles，而图拓扑、分配、布局、指针和启动网格保持固定。
3. 轻量级更新：每次更新仅传输28字节的数据，极大降低了内存开销。

工程权衡与实测表现

• 优点：

  • 显著减少预服务图构建时间，从3915毫秒降至93毫秒。
  • 保持了与基线相当的延迟表现（17.7 vs 17.8毫秒）。
  • 避免了填充开销，在100%填充浪费情况下，P99延迟减少了18.7%。

• 限制：

  • 仅限于单GPU的合成轨迹测试，未在生产环境中验证。
  • 依赖于Triton内核的特定实现，可能不适用于其他框架。

开发者落地与部署建议

PolyGraph为自定义Triton内核提供了一种高效的工作形状管理机制，特别适用于对延迟敏感且形状变化频繁的应用场景。开发者可以参考以下建议进行部署：

1. 评估应用场景：在应用场景中，形状变化频繁且对延迟敏感时，PolyGraph的优势最为明显。
2. 硬件兼容性：目前仅在NVIDIA H800 GPU上进行了测试，其他硬件平台可能需要进一步验证。
3. 性能优化：结合自适应批处理机制和内存优化技术，可以进一步提升整体性能。

结论

PolyGraph展示了在自适应推测解码中通过计划缓冲机制优化CUDA图的潜力。尽管目前仅限于单GPU的合成轨迹测试，但其创新性机制为未来在生产环境中的应用提供了新的思路。

---

消息来源：ArXiv Machine Learning (cs.LG) (2026-10-21)

—— 完 ——

主题标签： #PolyGraph 关注 #CUDA图 关注 #自适应推测解码 关注 #Triton内核 关注 #GPU优化 关注