Tempo: Compiled Dynamic Deep Learning with Symbolic Dependence Graphs (SOSP 2025)

一句话总结：用 recurrent tensors + symbolic dependence graph 把动态时序依赖编译化，Llama-3.2-3B decode 7× 于 JAX，三种 RL 算法端到端 54× 加速且峰值显存 16× 更低。

问题与动机

许多 DL 算法（Attention、RL、diffusion）用 recurrence 表达跨 timestep 依赖，但 eager 系统（PyTorch）无法 whole-program 优化；graph 系统（JAX/TensorFlow）要求静态 shape，迫使用 padding/mask、loop unroll 或拆成多个 static subgraph。

后果：LLM decode 的 KV-Cache 策略需为每种 attention 类型手工工程；RL 的 actor-learner 拆分导致 重复前向、串行 acting/learning、高峰显存（Fig. 4 三大缺陷）。

关键观察 / 隐含假设

观察 1：动态依赖的结构可用 timestep 符号表达式紧凑编码，从而支持代数化简、vectorization、tiling 等 whole-program 优化。
- 依赖假设：依赖表达式在编译期可符号操作（invert、slice range 等）；T 上界已知或可参数化。
- 可能失效场景：数据依赖的控制流（依赖运行时 tensor 值的分支）超出 RT 表达能力。
- 证据强度：强——RL 与 LLM decode 两类 workload 均显著加速。
观察 2：RL actor-learner 拆分的根本原因是现有系统无法在同一程序里表达 acting 与 learning 的细粒度未来依赖。
- 依赖假设：experience 数据可通过 RT 自动 memory management 保留/释放，无需 replay buffer 语义。
- 可能失效场景：极大 T 导致 schedule/search 爆炸；需显式 checkpoint 的极长 horizon。
- 证据强度：中——三种经典 RL 算法验证，但未覆盖 distributed RL。
假设 1：polyhedral schedule + 自动 memory management 可在单 GPU 上优于手工调度的 JAX/Torch 实现。
- 证据强度：强——7×/54× 量级差距，ablation 支持 tiling 与 memory 贡献。

核心方法

Recurrent tensors (RT)：显式 temporal dimension，用符号表达式索引过去/未来 timestep；支持 domain unioning、符号自动微分（dependence invert）。
Symbolic dependence graph (SDG)：顶点带 temporal domain，边用符号表达跨 timestep 依赖；在此上做代数化简、vectorization、tiling（动态依赖 tile 成 static block 复用现有 codegen）。
Polyhedral scheduling：SDG 转 ILP 约束求全局执行顺序。
Automated memory management：在 SDG 上插入 dealloc/donate/swap 操作，与 schedule 联合求解。

后端可接 Torch/JAX codegen，按 schedule AST 运行时解析符号依赖。

设计取舍

取舍 1：声明式 RT 编程模型学习曲线高于 imperative PyTorch，换 compile-time 优化空间。
取舍 2：单 GPU 原型，分布式 dynamic DL 未涉及。
边界条件：依赖结构在编译期可符号化的算法收益最大；完全动态 shape 且无 recurrence 结构的模型收益有限。

实验与结果

Llama-3.2-3B decode：最高 7× vs JAX/Torch，序列长度扩展 4×
REINFORCE/A2C/PPO：54× 端到端 vs 五套 RL baseline，峰值显存 16× 更低
内存：tiling + memory planning 是 LLM 长序列的关键

Critical Analysis

论证链条

「动态依赖有结构 → 符号图可优化」是核心 insight，从 observation 到 SDG/polyhedral 设计链条闭合。RL 结果尤其能证明 actor-learner 拆分非必然，具有 counterintuitive 价值。

假设压力测试

编译时间随 T、模型深度增长如何 scale？论文对 compilation time 讨论有限。
多 GPU / pipeline parallel decode 未覆盖，7× 结论限于 single-GPU。
与 FlashAttention、PagedAttention 等手工 KV 优化系统的正面对比深度不足。

实验可信度

Baseline 包含多种 RL 框架，覆盖面好。LLM 实验用 3B 模型，对 70B+ production serving 的外推需谨慎。

系统性缺陷

论文未讨论：编译失败/debug 体验；与 torch.compile/XLA 动态 shape 进展的横向对比；生产部署的 warmup 与 AOT 缓存策略。

局限与 Future Work

局限 1：单 GPU，无分布式 dynamic compilation。
局限 2：用户需适应 RT 编程模型。
Future work 1：与 KV-Cache 专用 runtime（vLLM 类）的 hybrid——编译期 schedule + runtime paging。

Awesome System Papers Wiki

探索

Tempo-SOSP25