问题域/PD-28

CUDA Graph 与编译优化

CUDA Graph & Torch Compile Optimization

通过 CUDA Graph 捕获和 torch.compile 消除 GPU kernel launch 开销。model_runner.py 在 decode 阶段预捕获不同 batch size 的 CUDA Graph,运行时直接 replay 避免重复 launch。sampler/rotary_embedding/layernorm/activation 均用 @torch.compile 装饰,attention 层使用 Triton 自定义 kernel (store_kvcache_kernel) 实现高效 KV Cache 写入。

子问题

1.CUDA Graph 捕获与 replay(decode 阶段)

2.graph 变量预分配与填充(padding 到捕获尺寸)

3.torch.compile 应用于计算密集算子

4.Triton JIT kernel 实现 KV Cache store

各项目的解法0 solutions

Signals

最佳实践

1.仅对 decode 阶段捕获 CUDA Graph(batch size 可预测)

2.使用 graph pool 共享显存,减少碎片