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.
