跳转至

Reasoning Compiler: LLM-Guided Optimizations for Efficient Model Serving

会议: NeurIPS 2025
arXiv: 2506.01374
代码: https://github.com/he-actlab/REASONING_COMPILER
领域: 编译器优化 / 模型服务
关键词: LLM引导编译, MCTS, 程序优化, 神经编译器, 采样效率

一句话总结

提出 Reasoning Compiler，将编译器优化建模为序列决策过程，用 LLM 作为上下文感知提案引擎 + MCTS 平衡探索/利用，在 5 个代表性 benchmark 和 5 个硬件平台上实现平均 5.0× 加速且采样效率比 TVM 进化搜索提升 10.8×。

研究背景与动机

  1. 领域现状：大规模模型推理成本高昂，编译器优化（tiling, fusion, vectorization 等）是关键加速手段。现有神经编译器（TVM, Ansor 等）使用进化搜索或模拟退火探索变换空间。

  2. 现有痛点：

  3. 规则编译器依赖手工启发式，对特定负载或硬件过拟合
  4. 随机搜索方法采样效率低——不利用变换间的上下文依赖关系

  5. 变换组合空间指数级大且高度相互依赖（如 loop tiling 的收益取决于之前是否做了 fusion）

  6. 核心矛盾：编译器变换间存在复杂的非局部交互，盲目搜索无法有效捕获这些上下文依赖。

  7. 切入角度：LLM 天然擅长上下文推理——给它当前代码状态、变换历史、性能反馈，让它推理下一步应该做什么变换。

  8. 核心idea一句话：LLM 提供上下文感知的变换提案 + MCTS 提供结构化搜索 = 高采样效率的编译器优化。

方法详解

整体框架

将程序优化建模为有限视野 MDP：\(\mathcal{M} = \langle \mathcal{S}, \mathcal{A}, \mathcal{P}, \mathcal{R} \rangle\)。状态是程序变体，动作是编译器变换，转移确定性，奖励由硬件代价模型评估。MCTS 做结构化搜索，LLM 在每个扩展节点提出变换提案。

关键设计

  1. LLM 引导的上下文推理:
  2. 做什么：在 MCTS 每次扩展时，让 LLM 基于完整上下文提出下一步变换
  3. 核心思路：Prompt 包含当前程序 \(p_i\) 及其父/祖父程序的源码、性能预测和变换历史，以及所有可用变换集 \(O\)。指示 LLM 进行 CoT 推理——分析变换间的协同/拮抗效应

  4. 设计动机：变换组合存在复杂交互（如 tiling 后做 vectorization 可能更优），LLM 能推理这些关系

  5. MCTS 结构化搜索:

  6. 做什么：平衡探索和利用，复用公共变换前缀
  7. 核心思路：UCT 选择 \(\text{UCT}(p_i) = \frac{W(p_i)}{N(p_i)} + c\sqrt{\frac{\ln N(p_{i-1})}{N(p_i)}}\)，LLM 提案做节点扩展，随机 rollout + 代价模型做评估，反向传播更新统计量

  8. 设计动机：MCTS 的树结构天然支持变换序列的回溯和复用，解决了纯随机搜索的低效问题

  9. 硬件感知代价模型:

  10. 做什么：替代真实硬件运行，提供快速性能评估
  11. 核心思路：学习的代理模型 \(\hat{f}\) 预测延迟，用于 rollout 奖励计算
  12. 设计动机：真实编译+运行太慢（分钟级），代理模型评估毫秒级

训练策略

  • LLM 零样本使用（GPT-4o mini），无需微调
  • MCTS 超参：探索系数 \(c = \sqrt{2}\)，分支因子 \(B = 2\)
  • 变换验证：LLM 输出经解析和过滤，无效变换回退至随机采样

实验关键数据

主实验（5 个 kernel × 5 个硬件平台）

方法 平均加速比 平均采样数 采样效率提升
TVM 进化搜索 基线 5.8× 更多 1.0×
MCTS (无LLM) 中等 中等 中等
Reasoning Compiler 5.0× 5.8× 更少 10.8×

端到端 Llama-3-8B

方法 加速比 采样数减少 采样效率提升
TVM 基线 基线 基线
Reasoning Compiler 4.0× 3.9× 更少 5.6×

关键发现

  • 在 FLUX attention 上，Reasoning Compiler 36 样本达 2× 加速，TVM 需 600+ 样本（16× 减少）
  • 在 Llama-4 MLP 上，20 样本达 12.7× 加速，TVM 3000 样本仍未达到
  • 低预算场景优势最大——前几十个样本就能发现高质量优化序列

亮点与洞察

  • LLM 做提案而非做决策：LLM 的角色被精确定位为"上下文感知的变换提案器"，正确性和搜索结构由 MCTS 保证。这种 LLM-in-the-loop 设计范式值得学习
  • 编译优化的 MDP 建模：将变换序列建模为 MDP，让 MCTS 的理论保证自动适用
  • 跨硬件泛化：同一框架在 ARM、x86、Apple Silicon 上均有效，不需要硬件特定调整

局限性 / 可改进方向

  • 依赖 LLM API 调用，对于大规模部署成本可能较高
  • GPT-4o mini 的代码理解能力有限，更强的 LLM 可能进一步提升
  • rollout 中的随机变换序列可能不够高效
  • 未与其他 LLM for code 的编译优化工作比较

相关工作与启发

  • vs TVM/Ansor/MetaSchedule：这些用进化搜索/模拟退火，采样效率低；Reasoning Compiler 通过 LLM 推理大幅提升
  • vs STOKE：STOKE 用 MCMC 做超级优化但也不利用上下文；本文的 LLM+MCTS 组合更高效
  • 展示了 LLM 在系统级优化中的巨大潜力

评分

  • 新颖性: ⭐⭐⭐⭐⭐ LLM+MCTS 的编译器优化框架是首创，将 LLM 推理与结构化搜索优雅结合
  • 实验充分度: ⭐⭐⭐⭐⭐ 5 个 benchmark × 5 个硬件平台 + 端到端评估，非常全面
  • 写作质量: ⭐⭐⭐⭐⭐ 形式化清晰，动机阐述到位
  • 价值: ⭐⭐⭐⭐⭐ 对编译器和模型服务社区有重大影响力