Deepgram Voice AI Optimization: Porting Nova-3 to AWS Trainium2 in Practice

Voice AI company Deepgram, in collaboration with the AWS Forward Deployed Engineering team, ported its latest Nova-3 speech-to-text model to AWS Trainium2 chips for offline batch inference. This work involved not only migrating code from GPUs to specialized AI chips, but also kernel-level tuning and a custom data-parallel runtime, ultimately achieving significant gains in throughput and cost. For teams focused on AI infrastructure and deploying voice technology, this is a rare full-stack optimization case study.

Why Choose Trainium2 for Batch Inference

Deepgram's Nova-3 is its flagship speech recognition model, supporting multilingual, high-accuracy, low-latency real-time transcription. But as customer demand for batch transcription of historical audio grew, inference cost became a key bottleneck. AWS Trainium2 is an accelerator chip designed specifically for machine learning training and inference, with a lower cost per unit of compute than general-purpose GPUs, making it especially suitable for large-scale batch tasks that can tolerate some latency.

The goal for Deepgram and AWS was clear: increase Nova-3's batch inference throughput severalfold while reducing the compute cost per transcription, all without compromising model accuracy. This required solving problems at three levels: software stack adaptation, compute kernel optimization, and runtime scheduling.

Trainium2 is AWS's second-generation in-house machine learning accelerator chip, built on the NeuronCore-v3 architecture, offering higher high-bandwidth memory (HBM) capacity and higher FP8/BF16 compute density per chip. Unlike general-purpose GPUs, it has dedicated data path optimizations for large-scale matrix operations and supports multi-chip interconnect via NeuronLink, making it easy to build high-throughput inference clusters. Its cost advantage comes mainly from two aspects: first, the chip itself is designed for cost-effectiveness; second, it is offered as an AWS cloud service, eliminating the procurement and operational overhead of building your own GPU cluster. For batch inference, it allows trading slightly higher per-request latency for higher overall throughput—exactly the trade-off needed for offline transcription scenarios.

Porting to the AWS Neuron Software Stack

Trainium2 provides programming interfaces through the AWS Neuron SDK, supporting mainstream frameworks such as PyTorch and TensorFlow. Deepgram's first step was to migrate Nova-3's inference graph from its original GPU environment to the Neuron stack. This was not a simple recompile—the speech model contains many custom operators (such as convolutions, attention variants, and audio feature extraction layers), each of which needed to be verified for numerical correctness and performance on Neuron.

The team used Neuron's graph compiler to perform operator fusion and memory planning on the model, and adjusted tensor layouts to match Trainium2's architectural characteristics. Operators that could not be directly mapped were implemented via Neuron's custom operator interface as fallbacks, ensuring the model's functionality remained complete.

AWS Neuron SDK is the official software stack for accessing Trainium and Inferentia chips, including the compiler (neuronx-cc), runtime libraries, and performance analysis tools (Neuron Profiler). It provides adaptation layers for PyTorch (torch-neuronx) and TensorFlow, so developers only need to make minor substitutions in their existing framework code to hand the computation graph to the Neuron compiler for operator fusion, memory planning, and instruction scheduling. The compiler partitions the graph into subgraphs executable on NeuronCore, and unsupported operators fall back to CPU or are implemented via the custom operator interface. Understanding this "graph compilation + operator fallback" mechanism is key to judging porting effort and performance ceilings.

Kernel-Level Tuning: Squeezing Out Every Drop of Compute

Porting was only the starting point; the real performance gains came from kernel-level tuning. Deepgram engineers analyzed in depth the execution bottlenecks of Neuron kernels on the speech model and found that audio feature extraction and attention computation were the main hotspots.

For the audio frontend, the team rewrote part of the data preprocessing logic to better fit Trainium2's vectorized instructions; for the attention layers, they reduced memory bandwidth pressure by adjusting tile sizes and pipelining strategies. According to the original article, after multiple rounds of tuning, per-chip inference throughput improved substantially over the initial ported version. Specific numbers were not disclosed, but they were sufficient to support the cost targets of its batch transcription service.

In addition, the team used Neuron's performance analysis tools to identify several inefficient operators and removed them from the critical path through operator replacement and fusion.

Kernel-level tuning refers to further optimization on top of the low-level compute kernels generated by the compiler, rather than staying at the framework-level parameter tuning. Common techniques include: adjusting matrix multiplication tiling sizes to match the chip's on-chip cache capacity, reordering data layouts to enable vectorized instructions, fusing multiple small operators into one large kernel to reduce kernel launch and memory round-trip overhead, and using pipelining to hide memory access latency. The audio frontend of speech models involves many non-matrix operations (such as STFT and Mel filter banks), and such workloads are often difficult for general-purpose compilers to map efficiently, so kernels need to be hand-written or rewritten to fit Trainium2's vectorized and parallel execution units.

Data-Parallel Rust Runtime

Another challenge of batch inference is task scheduling. Deepgram did not rely on an off-the-shelf inference server; instead, it built a lightweight data-parallel runtime in Rust, responsible for splitting massive audio files, distributing them to multiple Trainium2 instances, and aggregating results.

Rust was chosen for both performance and memory safety: the runtime needs to handle highly concurrent file I/O and network communication while ensuring stability over long-running operation. The runtime supports dynamic batching, automatically adjusting batch size based on queue length to fully utilize chip compute. The data-parallel strategy lets multiple Trainium2 nodes work together, scaling throughput linearly.

Dynamic batching is a scheduling strategy for aggregating requests in inference services: the system maintains a request queue and, within a very short time window, packs multiple pending requests into a single batch sent to the chip, thereby amortizing the cost of each forward pass and improving chip utilization. Window length and maximum batch size require trade-offs—longer windows and larger batches yield higher throughput but also higher queuing latency. For offline batch transcription, latency tolerance is high, so larger batches and aggressive aggregation strategies can be used. Data parallelism means replicating the model across multiple chips or nodes, each processing a different data shard, thereby scaling total throughput approximately linearly; it complements model parallelism (splitting a single model).

Throughput and Cost Results

Although Deepgram and AWS did not disclose precise benchmark numbers, the official blog clearly states that after migrating to Trainium2, Nova-3's batch transcription throughput improved significantly and the per-unit audio processing cost dropped noticeably. This result enables Deepgram to offer large-scale offline transcription services at more competitive prices while reserving compute headroom for deploying larger models in the future.

From a technical perspective, this collaboration validated the feasibility of Trainium2 for voice AI batch inference scenarios and provided a reference path for other teams hoping to reduce GPU dependence and optimize inference costs: start with porting to the Neuron stack, go deep into kernel tuning, and then pair it with a custom scheduling runtime to gradually unlock the potential of specialized chips.

Implications for Voice AI Infrastructure

Deepgram's case shows that deploying specialized AI chips is not "plug-and-play"; it requires co-design of the model, compiler, and runtime. For scenarios with relatively fixed inference workloads and large batch sizes, investing engineering resources in chip-level optimization can yield substantial long-term cost benefits. As the Trainium ecosystem matures, batch inference in areas such as voice, recommendation, and vision is likely to see more similar optimization practices.

One-line Summary

Deepgram ported Nova-3 to Trainium2, significantly reducing batch transcription costs through kernel tuning and a custom runtime.

Paragraph Summary

Deepgram, in collaboration with the AWS Forward Deployed Engineering team, migrated its flagship speech-to-text model Nova-3 from GPUs to AWS Trainium2 chips for offline batch inference. The work spanned three levels: first, porting the inference graph via the AWS Neuron SDK, handling numerical correctness and fallbacks for custom operators; second, deep kernel-level tuning, rewriting preprocessing logic for audio feature extraction and attention computation and adjusting tiling and pipelining strategies to relieve memory bandwidth pressure; finally, building a custom data-parallel runtime in Rust supporting dynamic batching and linear cross-node scaling. Although official precise numbers were not disclosed, batch transcription throughput improved significantly and unit cost dropped noticeably, validating the feasibility of specialized AI chips for voice batch inference scenarios.