通过 AWS Trainium 和 vLLM 上的推测性解码加速解码繁重的 LLM 推理

在本文中,您将了解推测性解码的工作原理以及它为何有助于降低 AWS Trainium2 上每个生成令牌的成本。

来源:亚马逊云科技 _机器学习

实用基准测试显示,使用 vLLM、Kubernetes 和 AWS AI 芯片部署 Qwen3 模型时,代币间延迟更快。

AWS Trainium 上的推测性解码可以将解码繁重的工作负载的令牌生成速度提高多达 3 倍,有助于降低每个输出令牌的成本并提高吞吐量,而无需牺牲输出质量。如果您构建人工智能写作助手、编码代理或其他生成式人工智能应用程序,您的工作负载可能会产生比它们消耗的更多的令牌,从而使解码阶段成为推理的主要成本。在自回归解码期间,令牌是按顺序生成的,从而使硬件加速器受到内存带宽限制且未得到充分利用。这会增加每个生成令牌的成本。推测性解码通过让小型草稿模型一次提出多个标记来解决此瓶颈,目标模型在单个前向传递中验证这些标记。更少的串行解码步骤意味着更低的延迟和更高的硬件利用率,有助于降低推理成本。

在这篇文章中,您将学到:

  • 推测性解码的工作原理以及它为何有助于降低 AWS Trainium2 上每个生成令牌的成本
  • 如何在 Trainium 上使用 vLLM 启用推测解码
  • 我们用于评估性能的基准测试方法
  • 如何调整草稿模型选择和工作负载的推测令牌窗口大小
  • 使用 Qwen3 重现结果的分步说明
  • 什么是推测解码?

    推测解码通过使用两个模型加速自回归生成:

  • 草稿模型快速提出 n 个候选标记。
  • 目标模型在一次前向传递中验证它们。
  • 将其设置得太高会增加早期拒绝的可能性,浪费草稿计算并提高目标模型验证成本。您可以根据观察到的接受率,通过平衡草稿计算与验证成本来调整此值。

    图 1 推测性解码配置权衡

    NeuronX 分布式推理 (NxD Inference) 支持什么

    图2 系统架构

    结果