Together AI 发布 Together Compiler，为 AI 推理提供持续编译优化

Together Compiler：面向 AI 推理的持续编译 #4 on TrendingalphaXiv 独家 Together AI Reyna Abhyankar DL Dylan Lim DW Drew Wadsworth VH Venkata Hasith Vattikuti YJ YuTong Julia Zhang JH Jung Hwan Heo TK Taeksang Kim KR Kesavan Ramakrishnan ES Evangelene Stanley XF Xinyu Fang DY Daniel Y. Fu SA Simran Arora +4 更多 收起 摘要 如今的 AI 工作负载正在快速变化：新模型发布得越来越频繁，模型架构日益多样化，流量分布也经常变动。我们在 Together AI 面临的挑战，是针对这些不断变化的工作负载进行持续优化。将给定部署编译为优化的前向计算过程，正是一个非常适合用 agent 爬山法解决的问题。现有系统在持续编译上面临困难，原因在于测量吞吐量与保真度之间的权衡。许多系统使用 agent 借助离线基准来优化推理前向计算中的单个 kernel，但由于混合专家（MoE）等原语存在数据依赖行为，计时结果的保真度很低。另一方面，也有系统使用 agent 直接优化端到端引擎，但受困于引擎启动时间过长。为应对这些挑战，我们构建了 Together Compiler（TKCC），这是首个真正为 AI 推理优化而设计的持续 agent 循环。为解决吞吐量与保真度之间的权衡，我们提出：（1）分层持久化引擎，将引擎启动开销与前向计算解耦；（2）一个成本模型，用于选择为优化过程提供初始种子的并行配置。我们还对 agent 上下文管理进行了大量消融实验，并分享了我们关于在生产环境中部署 TKCC 的见解。在 Llama 1B 主力工作负载上，TKCC 相比经过高度优化的人工编写 Megakernel 最高提速 1.7 倍；在 SemiAnalysis 的 InferenceX 基准上，跨硬件平台针对最新前沿模型（Kimi K3、GLM 5.3、MiniMax M3.1 Flash、DeepSeek v4.1 Flash），相比最佳开源引擎最高可达到 2.6 倍的每秒 token 数（TPS）。TKCC 还在低延迟、batch size 1 的推理上创下了最先进性能：在 MiniMax M3.1 Flash 上生成 1111 TPS，在 DeepSeek v4.1 Flash 上生成 1205 TPS。最后，TKCC 具备泛化能力，其加速效果不仅限于前沿语言模型推理。TKCC 在我们的生产集群上服务了数万亿 token。 引用 AI 概览 复制 编辑 推理部署会逐渐过时。某个实验室的模型在两周内其典型上下文从 10,000 增长到 18,000 token，同时缓存命中率上升了八个百分点。这个例子说明了为什么一次性调优的配置会随着请求变化而偏离最佳选择；它并不是下文结果所对应的基准测试条件。 Together Compiler（TKCC）正是围绕这一问题构建的。Agent 提出替换用的前向计算代码，这些代码会在一个已经加载的引擎内部接受测试。快速筛选会淘汰较差的候选，而有前景的候选则会接受更慢、更具代表性的检查。一个成本模型负责挑选跨 GPU 工作的初始排布。在一个 batch-size-one 的 MiniMax 示例中，该流程相比已发布的配置带来了更高的单用户交互性，不过 GPU 并行配置也发生了变化，因此这一比较并非仅针对代码。 让引擎保持运行 为每个候选都重启推理引擎太慢了。TKCC 改为只启动一次持久化引擎，并只加载一次模型权重。当 agent 提交候选时，服务器在热重载时替换模型的前向函数指针，并重新捕获 CUDA/HIP 图；拆除时则恢复基线。多个 agent 并行工作并相互通信，因此当一个 agent 在思考时，GPU 时间不太可能闲置。 测量分为两个层级。快速层只服务较少的层，并以较低成本拒绝缓慢、无效或挂起的候选。慢速层服务更多的层，并为数据依赖行为提供更具代表性的信号。在 B200 上优化 GLM 5.3 时，作者将这种两层设置与单一慢速层进行了对比：它评估了 1,354 个候选，而非 168 个，约为后者的 8.1 倍，并找到了快 7.2% 的解决方案。这检验的是测量设计，而非任何特定的候选生成策略。 作者还与 Apple 的 RSI-Inference Bench 进行了对比，后者是一个端到端 agent 优化测试框架，其更高保真度的测试会因每次测量都必须启动引擎而牺牲实验吞吐量。他们的对比涵盖了两层设置、单一慢速层、无分层以及 Apple 的测试框架。Apple 测试框架有 10 小时的预算和 8 小时的验证器超时，并使用与 TKCC 相同的 GLM 5.3 服务镜像、不含量化的前向计算编辑面和 AgentX 工作负载，测量每秒输出 token 数。 慢速层之所以重要，是因为混合专家（MoE）模型会将每个 token 只路由到部分专家，而不是对每个 token 应用全部权重。因此，输入分布会改变实际执行的工作。作者观察到，随机初始化的数据比真实数据多激活了 13–38% 的专家，这导致引擎端到端慢了 1.3 倍。这正是论文不信任孤立的、随机输入 kernel 测量的动机。GEAK 等系统使用 agent 调优 kernel，但论文并未与它们进行正面对比。 挑选有前景的 GPU 布局 这里的并行指的是模型的工作和权重如何在 GPU 之间划分。TKCC 的成本模型利用分析估计以及实测的计算、内存和通信性能来缩小候选配置范围，并考虑生产工作负载特征，例如专家不均衡和请求长度不均。其目标是既避免穷举搜索，也避免仅依赖硬件规格。

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）