Deepgram 如何用 AWS Trainium 优化大规模语音 AI

As voice becomes a first-class modality for AI, Deepgram’s models span speech-to-text, text-to-speech, and voice agents. They now process quadrillions of audio tokens a year for more than 200,000 developers and 1,400 organizations. Customers use them through the Deepgram API or run them in their own environments, including on Amazon SageMaker AI.

For Deepgram, serving efficiency determines how much audio the system can process for a given infrastructure budget. High-volume batch workloads put heavy pressure on that budget, so the more efficiently each accelerator hour is used, the more audio Deepgram can serve on the same fleet. That made inference efficiency a natural target and Deepgram turned to AWS AI chips to improve it. They worked with AWS to bring its Nova-3 speech-to-text model to AWS Trainium2 for offline batch workloads such as contact-center analytics and medical transcription.

The engagement began with Deepgram delivering Nova-3 as standard PyTorch code. Over several months, Deepgram and the AWS Forward Deployed Engineering (FDE) team ported Nova-3 to AWS Trainium2 in PyTorch reconstruction and tuned it for the new hardware. The final configuration runs eight model replicas on a single Amazon Elastic Compute Cloud (Amazon EC2) trn2.3xlarge instance. On Deepgram’s offline, throughput-oriented batch workloads, that setup processed 11–70 percent more audio per dollar than the compute options Deepgram had been running. We measured this using instance pricing at the time of the evaluation. These gains help Deepgram scale efficiently, keep shipping better models, and bring them to customers at lower prices. These benefits can extend to customers who self-host Deepgram models in their own virtual private cloud (VPC) on Amazon SageMaker AI. There, Deepgram and AWS continue to improve ease of setup, reliability, and observability.

In this post, we walk through the full inference optimization effort on Trainium: porting Nova-3 to the AWS Neuron software stack, the kernel-level tuning that followed, the data-parallel Rust runtime that lets Deepgram’s own model server call Trainium, and the throughput and cost results we measured on batch transcription.

Optimizing Nova-3 to Trainium2 with native PyTorch

Starting with standard PyTorch code, we used TorchNeuron, the native PyTorch integration of the Neuron Compiler, which supports both eager execution and torch.compile. We ran Nova-3 in eager mode on Trainium2 first and checked which operations fell back to the CPU, then compiled with torch.compile. The compiled path ran every model operation on Trainium2 with no CPU fallbacks detected.

We compiled the encoder and decoder as separate static-shape graphs using torch.compile. Throughout, we checked numerical accuracy at two levels: layer-level output consistency against CPU, and the end-to-end Word Error Rate (WER), which had to stay in line with Deepgram’s reference values and target.

From that working baseline, we kicked off the optimization work by replacing several layers with fused kernels from the public NKI Library (NKILib), a collection of pre-built reference kernels written with the Neuron Kernel Interface (NKI). The kernels covered operations such as fused QKV projections, fused feed-forward network (FFN) layers, and depthwise 1D convolution.

We optimized for throughput and, by extension, cost efficiency, targeting Deepgram’s offline inference workloads. Throughput is reported as a relative gain over a baseline configuration rather than as absolute audio processed per second, and we swept batch sizes 1–128 to find the throughput-maximizing batch size.

Kernel-level tuning with Neuron Explorer and NAD

Enabling kernels is only the first step of an optimization flywheel: enable an NKI kernel (custom or from NKILib), profile the workload, then either tune the existing kernel or replace it with a new one based on the profiling insights. This cycle repeats until the performance target is met. At each step, a different part of the Neuron stack accelerates the work: NKILib provides drop-in replacement kernels, Neuron Explorer is the go-to tool for profiling Neuron workloads, and the Neuron Agentic Development (NAD) toolkit speeds up profiling analysis and kernel tuning through its profiling and kernel-writing skills.

In our case, the NKILib kernels got us most of the way, and profiling showed where a hand-written kernel would do better. Each time, the NAD profiling skills took us from a Neuron Explorer profile to a root cause, and the NAD kernel-writing skills helped produce the replacement kernels.

Layer	What we found and changed	Gain
FFN	The reference kernel’s tiling schedule reloaded the full set of projection weights from off-chip memory for every token tile, so memory bandwidth limited throughput as the token count grew. Our hand-tuned kernel loads each weight tile once, reuses it across multiple token tiles, and pipelines weight reads with on-chip compute.	2.0x lower latency compared to the reference kernel at batch size 96
Depthwise conv1d	The experimental NKILib kernel framed the operation as an implicit general matrix multiply (GEMM), which left Tensor Engine utilization low. We rewrote it as a broadcast multiply-accumulate loop on the Vector Engine, reaching full engine utilization, enough to more than offset the Vector Engine’s lower peak throughput.	1.5x speedup compared to the reference kernel at batch size 96
Conv2d	The model’s conv2d layers hit a compilation bottleneck: beyond batch size 32, these layers prevented the model from compiling as a single graph. Since NKILib had no conv2d kernel, we wrote one from scratch in NKI. The kernel went through several profiler-informed iterations following the flywheel described earlier, unlocking batch sizes above 32.	Up to 50x on the convolution module at batch sizes ≤32

Data-parallel runtime in Rust

Each Trainium2 device features eight physical NeuronCores, exposed by default as four logical NeuronCores (LNC2). Because each Nova-3 model instance fits comfortably on a single physical core, we reduced the number of physical cores per logical core to one (LNC1). This configuration lets us host eight model replicas (DP8) per Trainium2 device. Inputs are processed concurrently while maximizing per-device throughput. The data-parallel runtime is driven by a work-stealing queue with one thread per replica.

Deepgram’s production serving stack is written in Rust, so porting Nova-3 to Trainium2 required more than compiling the model: it also had to run without Python or PyTorch in the serving path. We therefore built a lightweight runtime in which Rust coordinates model replicas through a shared work-stealing queue. A C layer loads and executes ahead-of-time-compiled Neuron Executable File Format (NEFFs) directly through the Neuron Runtime API (libnrt). The runtime integrates directly with Deepgram’s serving architecture and reproduced the reference token sequences exactly at every tested batch size.

We benchmarked it against multithreaded (MT) and multiprocessing (MP) Python data-parallel runtimes. Compared with the MP baseline, the Rust runtime removes the overhead of the Python interpreter entirely. Compared with the MT baseline, it further alleviates global interpreter lock (GIL) serialization, so dispatch calls from concurrent worker threads are no longer contended within a single process. These gains matter because Nova-3 performs a sequence of small decoder calls for each audio chunk, making host-side dispatch the bottleneck at low batch sizes. The throughput improvement is therefore largest there, and remains meaningful even at the throughput-maximizing batch size. There, the Rust runtime delivered a 7.5% throughput gain over the multiprocessing Python baseline.

Figure 1: Deepgram Nova-3 on AWS Trainium2, compiled ahead of time and served from Rust

Cost saving impact

We measured price-performance as audio processed per dollar. On Deepgram’s offline, throughput-oriented batch workloads, the single-device Trainium2 setup (trn2.3xlarge) processed 11–70 percent more audio per dollar than the compute options Deepgram had been running.

Conclusion

In this post, we demonstrated how Deepgram and the AWS FDE ported the Nova-3 speech-to-text PyTorch model to Trainium hardware using TorchNeuron, which provides native PyTorch framework integration for AWS Trainium. We improved performance by replacing layers directly with pre-built kernels from NKILib. We then assessed performance using Neuron Explorer and used the Neuron Agentic Development (NAD) profiling skills to query profiling data and accelerate performance troubleshooting. Building on these insights, we adapted kernel source code using the NAD kernel-writing skills. Finally, we built a lightweight Rust data-parallel runtime for AWS Trainium, calling the Neuron Runtime API from Rust. Our setup demonstrated cost savings in offline, throughput-oriented scenarios, validating AWS Trainium as a viable and efficient option for Deepgram’s inference workloads.

To try this approach on your own models, start with the AWS Neuron documentation, which covers the compiler, the runtime, and the supported frameworks. The Neuron Kernel Interface (NKI) guide explains how to write the kernels described in this post, and the NKI samples repository holds the reference kernels we started from. The Neuron developer tools documentation covers the profiling and agentic tooling we used here, including Neuron Explorer and Neuron Agentic Development. For instance sizes and pricing, see Amazon EC2 Trn2 Instances. For help scoping a similar port, contact the AWS Forward Deployed Engineering.

About the authors

Pierre Lienhart

Pierre is a Senior Machine Learning Engineer with the AWS Forward Deployed Engineering team in London. He partners with model providers and AI companies to unlock the full potential of their inference workloads across a range of compute types, helping them serve models faster and at lower cost in production.

Victor Wang

Victor is a Staff Software Engineer at Deepgram and Technical Advisor to the VP of Engineering. Before joining Deepgram, Victor held multiple roles at AWS, including Sr. Solutions Architect, Technical Program Manager, Proserve Consultant, and Software Developer. His passion is learning new technologies and traveling the world, and he has flown over 2 million miles.

Andre Gomes

Andre is a Solutions Architect at Annapurna Labs ML, AWS, with experience helping startups run large-scale training and real-time inference. A former startup co-founder and CEO, he works with founders to scale their workloads on Trainium and Inferentia. Andre earned his PhD at UFMG, including research at the University of Cambridge, where he applied large language models and neural topic modeling to large-scale data analysis and decision making.

Elyse Zhang

Elyse is a Senior GenAI Strategist in AWS, where she helps enterprise and startup customers with their model customization journey. She focuses on translating business requirements into effective customization strategies across industries.

Haotian An

Haotian is a Machine Learning Engineer with the AWS Forward Deployed Engineering team, where he specializes in customizing foundation models and optimizing inference at scale. He works closely with customers to adapt generative AI to their specific use cases, helping them unlock new capabilities and drive measurable business outcomes.

Niraj Chaudhari

Niraj is a Deep Learning Architect with the AWS Forward Deployed Engineering team. He works with enterprise customers to customize foundation models through supervised and reinforcement fine-tuning, build agentic systems, and optimize AI workloads. He also helps customers take generative AI applications from proof of concept to production on Amazon SageMaker AI and Amazon Bedrock.

Rohit Thekkanal

Rohit serves as Senior Manager, ML Engineering for Model Customization at the AWS Forward Deployed Engineering team, where he leads the development of scalable generative AI applications focused on model optimization. With nearly a decade at Amazon, he has contributed to machine learning initiatives that significantly impact Amazon’s retail catalog. Rohit holds an MBA from The University of Chicago Booth School of Business and a Master’s degree from Carnegie Mellon University.