•2025年10月
ai-systems

Together AI 中使用内核融合的推测解码工程实践：实现 4x LLM 推理加速

探讨 Together AI 如何通过推测解码和内核融合优化 LLM 推理，焦点低延迟 token 生成和硬件参数配置，实现 4x 加速。

在大型语言模型（LLM）的推理阶段，延迟和吞吐量是关键瓶颈，尤其是在实时应用如聊天机器人或代码生成中。Together AI 通过整合推测解码（Speculative Decoding）和内核融合（Kernel Fusion）技术，实现了 LLM 推理速度的 4x 提升。这种方法的核心在于利用辅助模型（draft model）并行预测多个 token，同时优化 GPU 内核执行以减少开销，从而实现低延迟的 token 生成和高效的硬件利用。

推测解码的基本原理是使用一个小型、高效的 draft model 来生成候选 token 序列，然后由目标模型（target model）在单次前向传播中并行验证这些 token 的有效性。这种机制避免了传统自回归解码的逐 token 计算瓶颈，提高了整体吞吐量。在 Together AI 的实现中，draft model 基于 RedPajama 数据集训练，针对目标模型如 Llama 或 DeepSeek 进行对齐，确保 acceptance rate（接受率）保持在高水平。根据 Together AI 的 ATLAS 系统，这种自适应学习机制能在运行时从历史流量中学习，进一步优化 draft model 的预测准确性，从而动态提升 acceptance rate 至 2-3x 以上[1]。

证据显示，这种方法在实际部署中显著降低了延迟。例如，在 NVIDIA H100 GPU 上运行 Llama-3 8B 模型时，标准解码的 token 生成速率约为 150 tokens/s，而启用推测解码后可达 400 tokens/s，实现了 2.65x 的加速。更进一步，结合内核融合后，总加速比接近 4x。内核融合技术将 Transformer 层中的多个操作（如 LayerNorm、GeMM 和 Attention）合并为单个 CUDA 内核，减少了 kernel 启动开销和中间内存访问。Together AI 的自定义 FP8 内核比基线 PyTorch 快 75%，这得益于 subgraph fusion，将多头注意力块的操作融合，降低了数据传输量并提高了数学计算密度。

要落地这种优化，需要关注几个关键参数和配置。首先，draft model 的选择至关重要：推荐使用参数规模为目标模型 1/10 的小型 Transformer，如 7B 模型的 draft 使用 700M 参数的变体。draft length（γ）设置为 4-8，根据 acceptance rate 动态调整；如果 α > 0.7，则可增加 γ 以最大化并行验证收益。其次，硬件优化方面，优先使用 NVIDIA H100 或 H200 GPU，支持 FP8 量化以保持精度同时加速计算。batch size 建议为 1-16，针对低延迟场景保持小 batch 以最小化等待时间。

工程实践清单包括：

  • 模型对齐：使用知识蒸馏或联合训练，使 draft model 的 logit 分布接近目标模型，目标 KL 散度 < 0.1。
  • 内核融合实现：集成 FasterTransformer 或 DeepSpeed-Inference 库，将 attention 和 FFN 层融合；监控融合后 kernel 利用率，应 > 80%。
  • 监控与调优：实时追踪 acceptance rate 和 relative latency (c < 0.2)；如果 α 下降，使用 ATLAS-like 自适应机制，每 1000 请求更新 draft model。
  • 量化与精度：应用 QTIP 量化技术，确保 perplexity 损失 < 5%；测试 FP8 vs BF16 的权衡。
  • 回滚策略：如果推测解码失败率 > 20%，fallback 到标准解码；设置超时阈值 500ms 以防 draft model 延迟过高。

在硬件层面，Together AI 的 GPU 集群支持 InfiniBand 互联，确保多 GPU 场景下 KV cache 共享高效。对于长上下文（如 128K tokens），结合 PagedAttention 避免内存碎片。风险包括 draft model 在领域外数据上的低 α，导致退化；缓解方式是通过持续 fine-tune 和流量采样保持适应性。

总体而言，这种推测解码与内核融合的结合，不仅提升了 token 生成的低延迟性，还优化了 GPU 资源利用，提供可扩展的工程解决方案。在生产环境中，预期 ROI 为成本降低 50% 同时性能提升 4x，适用于高并发 AI 系统[2]。

（字数：1025）

[1] Together AI Blog: ATLAS System for Adaptive Speculative Decoding.

[2] Together AI Inference Engine Documentation.