Deepgram 如何借助 AWS Trainium 大规模优化语音 AI 推理效率

Deepgram 的语音 AI 业务规模庞大，每年处理的音频 token 数量达到数千万亿级别。对于这类高吞吐量服务商而言，推理效率直接决定了在固定基础设施预算下能够处理多少音频。批量离线工作负载对算力成本尤为敏感，因此 Deepgram 将目光投向了 AWS 自研 AI 芯片，希望通过硬件层面的优化来提升每加速器小时的利用率。

合作的起点是 Deepgram 提供标准 PyTorch 格式的 Nova-3 模型代码。AWS 与 Deepgram 团队首先在 Trainium2 上以 eager 模式运行模型，逐一排查回退到 CPU 执行的算子，随后使用 torch.compile 将编码器和解码器分别编译为静态形状计算图。编译后的路径实现了全部算子在 Trainium2 上执行，未检测到任何 CPU 回退，同时在层级输出一致性和端到端词错误率两个维度上通过了精度验证。

在基础移植完成后，团队引入了 Neuron Kernel Interface 公共库中的预置融合内核，替换了 QKV 投影、前馈网络层和深度一维卷积等关键算子。随后通过 Neuron Explorer 进行性能剖析，发现参考内核在内存带宽和引擎利用率上存在明显瓶颈。针对前馈网络层，团队手写了一个新内核，通过权重分块复用和流水线调度将延迟降低了 2 倍；针对深度卷积层，将计算范式从隐式 GEMM 改写为向量引擎上的广播乘累加循环，速度提升 1.5 倍。

conv2d 层的优化尤为关键。原有实现在批量大小超过 32 时会触发编译瓶颈，导致模型无法作为单一计算图编译。由于 NKILib 中没有现成的 conv2d 内核，团队从零开始用 NKI 编写，经过多轮基于性能剖析的迭代，最终解锁了批量大小 32 以上的编译能力，在小批量场景下卷积模块的性能提升幅度最高达 50 倍。整个调优过程遵循「启用内核—剖析—调优或替换」的飞轮循环，Neuron Agentic Development 工具包在根因分析和内核编写两个环节均发挥了加速作用。

在运行时层面，Deepgram 的生产服务栈以 Rust 编写，因此需要一套无需 Python 或 PyTorch 的推理路径。团队构建了一个轻量级 Rust 运行时：通过将每个逻辑 NeuronCore 的物理核数从默认的 2 个降为 1 个，单台 Trainium2 设备可承载 8 个 Nova-3 模型副本；一个 C 层直接加载并执行预编译的 Neuron 可执行文件，Rust 层通过工作窃取队列协调各副本并发处理请求，实现了最大化的设备吞吐量。

最终测试结果显示，在离线批量转录场景下，该配置与 Deepgram 此前使用的计算方案相比，每美元可处理的音频量提升了 11% 至 70%。这一成果将直接转化为更低的服务价格，并惠及在 Amazon SageMaker AI 上通过自有 VPC 部署 Deepgram 模型的企业客户。双方目前仍在持续改进自托管场景下的部署便捷性、可靠性和可观测性。

要点

• Deepgram 与 AWS 联合将 Nova-3 语音转文字模型移植到 Trainium2，在离线批量推理场景下每美元处理音频量提升 11%–70%
• 通过手写 NKI 内核替换参考实现，前馈网络层延迟降低 2 倍，深度卷积层速度提升 1.5 倍，conv2d 模块在小批量下最高提升 50 倍
• 将每个逻辑 NeuronCore 的物理核数降为 1，使单台 trn2.3xlarge 实例可同时运行 8 个模型副本，最大化设备吞吐量
• Deepgram 构建了绕开 Python 运行时的 Rust 数据并行调度层，通过工作窃取队列直接调用预编译 Neuron 可执行文件
• 优化成果将同步惠及在 Amazon SageMaker AI 上自托管 Deepgram 模型的企业客户