Flash RNN-T: Memory-Efficient Joint Network and Transducer Loss Computation

Mahmoud Ashraf

cs.LGOct 16, 2026

Flash RNN-T: Memory-Efficient Joint Network and Transducer Loss Computation

Authors: Mahmoud Ashraf

Organizations: Tarteel AI

Abstract

The recurrent neural network transducer (RNN-T) remains a practical choice for streaming speech recognition, but training requires large intermediate tensors. We present Flash RNN-T, a joint and loss implementation combining specialized kernels with optional tiled recomputation over the full alignment space and vocabulary. Tiled and untiled Flash form the entire measured throughput-memory frontier on medium and long workloads. On H100, untiled Flash achieves 24-39% higher throughput with 26-56% less memory than the fastest baseline on these workloads wherever it fits. Tiling reduces memory further through recomputation. On medium, tiled Flash uses ninefold less memory than that baseline at 16% lower throughput when both use one backward call per batch. On long, it fits where all tested alternatives with that backward schedule run out of memory with 80 GB. These comparisons use each implementation's lowest-memory configuration reaching at least 95% of its peak throughput.

Figures & tables

Panel of Figure 1 — Figure 1: Packing two utterances into tiles with an illustrative limit of $R=6$ alignment positions.

Panel of Figure 2 — Figure 2: Dashed enclosure marks intermediates discarded after each tile’s forward computation and recomputed during backward. Outputs from all tiles feed the loss.

Workload	Mean duration (s)	Encoder length $T$	Transcript length $U$	Valid alignment positions
Short	4.92	47–69	29–50	282,687
Medium	18.72	227–244	40–96	2,300,131
Long	69.62	851–885	325–458	41,122,726

Table 1: Representative workload batches.

Panel of Figure 3 — Figure 3: Throughput–memory tradeoffs under retain and release scheduling. Points correspond to tile limits for tiled Flash and sub-batch sizes for the other implementations.

Components	Short	Medium	Long
Untiled
Compiled	22743 / 2.80	3177 / 21.97	OOM
Specialized joint	21741 / 2.58	3205 / 20.44	OOM
Specialized norm.	26564 / 1.65	3878 / 11.80	OOM
Both specialized	25333 / 1.44	3852 / 10.27	OOM
Tiled

Table 2: Throughput (samples/s) / peak allocated memory (GiB). Tiled variants use a tile limit of 200k alignment positions.

Memory-Efficient Training

Memory-efficient training reduces the memory needed to optimize neural networks by compressing, recomputing, or relocating training state.

See all 89 papers

Chunk-wise Attention Transducers for Fast and Accurate Streaming Speech-to-Text
Feb 27, 2026Hainan Xu, Vladimir Bataev, Travis M. Bartley +1Streaming ASR Attention Mechanisms
NVIDIA Corporation.
The Recurrent Transformer: Greater Effective Depth and Efficient Decoding
Apr 23, 2026Costin-Andrei Oncescu, Depen Morwani, Samy Jelassi +3Autoregressive Language Modeling Self-Attention
Harvard University
Memory Caching: RNNs with Growing Memory
Feb 27, 2026Ali Behrouz, Zeman Li, Yuan Deng +3Recurrent Neural Networks Memory-Augmented Language Models

Feb 27, 2026J/K move · Enter open · S save

Chunk-wise Attention Transducers for Fast and Accurate Streaming Speech-to-Text

Hainan Xu, Vladimir Bataev, Travis M. Bartley, Jagadeesh Balam

NVIDIA Corporation.

Streaming ASR Attention Mechanisms Speech Translation

Open paper AskPDF arXiv

We propose Chunk-wise Attention Transducer (CHAT), a novel extension to RNN-T models that processes audio in fixed-size chunks while employing cross-attention within each chunk. This hybrid approach maintains RNN-T's streaming capability while introducing controlled flexibility for local alignment modeling. CHAT significantly reduces the temporal dimension that RNN-T must handle, yielding substantial efficiency improvements: up to 46.2% reduction in peak training memory, up to 1.36X faster training, and up to 1.69X faster inference. Alongside these efficiency gains, CHAT achieves consistent accuracy improvements over RNN-T across multiple languages and tasks -- up to 6.3% relative WER reduction for speech recognition and up to 18.0% BLEU improvement for speech translation. The method proves particularly effective for speech translation, where RNN-T's strict monotonic alignment hurts performance. Our results demonstrate that the CHAT model offers a practical solution for deploying more capable streaming speech models without sacrificing real-time constraints.

Abstract

Memory-Efficient Training

Explore similar work

Chunk-wise Attention Transducers for Fast and Accurate Streaming Speech-to-Text

The Recurrent Transformer: Greater Effective Depth and Efficient Decoding

Memory Caching: RNNs with Growing Memory