PolyGraph:面向自适应投机解码的 plan-buffer CUDA Graph
原标题:PolyGraph: Plan-Buffered CUDA Graphs for Adaptive Speculative Decoding
AI 摘要
论文提出 PolyGraph,一种面向自定义 Triton 内核的 plan-buffer 机制,用于解决自适应投机解码中 CUDA Graph 只能重放固定形状的问题。它只捕获一个最大网格的图,并在每次重放前用 28 字节异步拷贝把当前活跃形状写入 GPU 内存,内核据此跳过超出范围的 tile。在单张 NVIDIA H800 上,以 Qwen2.5-3B 维度做 trace 驱动原型,用一张图替代 64 张形状专用图,P50 延迟与预热基线持平(17.7 vs 17.8 ms),预服务图构建时间从 3915 ms 降至 93 ms。作者强调这只是机制原型,未集成到 vLLM 等生产引擎,且仅限单 GPU 合成 trace。
正文节选
PolyGraph: Plan-Buffered CUDA Graphs for Adaptive Speculative Decoding Abstract Adaptive speculative decoding changes the active batch size and draft length from one decode step to the next, whereas a CUDA Graph replays a fixed launch configuration. The resulting verification width forces serving systems to pre-capture many shapes, capture new ones on the serving path, or pad every step to the largest draft. We present PolyGraph, a plan-buffer mechanism for custom Triton kernels that avoids this