rnnt-loss

The RNN-Transducer forward-backward on the GPU, loadable through kernels. The reference baselines are torchaudio.functional.rnnt_loss (the warp-transducer lineage), an independently written fp64 lattice, and torch autograd through a pure-torch logsumexp lattice. Companion to dpx-decode, the max-plus side of the same lattice family.

Training an RNN-T, the architecture behind most streaming speech recognizers, needs the log-sum over every monotone alignment of U labels to T frames and its gradient: a forward-backward pass over a [T, U+1] lattice whose rows and columns are both sequential. This kernel parallelizes the one direction that is free, the anti-diagonal, and runs the whole lattice as one block per utterance, beating the torchaudio implementation by 1.4x to 1.9x with the margin growing on the long utterances speech training is made of.

The real lattice of a 100-frame, 40-label utterance: alpha filling one anti-diagonal per step, beta swept in reverse, and alpha + beta, the posterior over alignments that the gradient is made of, blazing along the alignment ridge. The loss matches an independent fp64 lattice to 1e-4 with forward+backward in 0.28 ms.

Usage

import torch
from kernels import get_kernel

rnnt = get_kernel("phanerozoic/rnnt-loss", version=1, trust_remote_code=True)

# joint is [B, T, U+1, V] from the joint network
loss = rnnt.rnnt_loss_from_logits(joint, labels, input_lengths, label_lengths,
                                  blank=0, reduction="mean")
loss.backward()

# or pass log-probabilities directly, keeping the log_softmax where you want it
loss = rnnt.rnnt_loss(torch.log_softmax(joint, -1), labels,
                      input_lengths, label_lengths, blank=0)

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Symbol	Purpose
`rnnt_loss(log_probs, labels, input_lengths, label_lengths, blank, reduction)`	transducer loss from log-normalized joint output
`rnnt_loss_from_logits(logits, ...)`	wrapper that applies `log_softmax` first
`RNNTLoss(blank, reduction)`	`nn.Module` form, stateless
`ops.rnnt_forward(...)`	`(loss [B], alpha, beta)`, the state reused by backward
`ops.rnnt_backward(...)`	gradient with respect to `log_probs`

log_probs is [B, T, U+1, V] fp32, normalized over V; labels is [B, U] int64; lengths are [B].

Method

alpha[t, u] depends on alpha[t-1, u] and alpha[t, u-1], so rows and columns are both sequential; the anti-diagonal is not: every cell on t + u = k depends only on cells with t + u = k-1. One block per utterance sweeps T + U diagonals with a barrier between them, and beta sweeps the same diagonals in reverse. With both variables in hand the gradient is pointwise: only two entries per lattice cell are nonzero (the blank and the emitted label), so the gradient kernel writes 2 * T * (U+1) values and is embarrassingly parallel. Inputs are log-probabilities so the joint's log_softmax stays in autograd where it fuses with the rest of the joint.

Measured

fp32, full forward and backward through logits, against torchaudio.functional.rnnt_loss on the same host:

B	T	U	V	torchaudio	this	speedup
8	100	20	128	0.406 ms	0.259 ms	1.57x
16	200	40	256	3.388 ms	2.345 ms	1.44x
32	300	50	512	24.069 ms	14.943 ms	1.61x
16	500	80	256	19.856 ms	10.470 ms	1.90x

Work per barrier is min(T, U+1) cells, so long-utterance shapes keep more of the block busy per synchronization; the advantage grows with T.

Correctness

torchaudio: loss agrees to 4.5e-7 relative worst case across batch 2 to 16, T to 200, U to 40, V to 128, padded and ragged lengths; gradients through log_softmax to 1.3e-5 relative.
An fp64 lattice: the forward recurrence agrees with a scalar float64 Python lattice, written independently, to 6.6e-6 absolute.
Autograd: gradients agree to 1.2e-6 relative against autograd through a pure-torch logsumexp lattice.
Gradient structure: summed over the vocabulary the gradient vanishes at every reachable cell to 1e-3, the posterior-mass identity.
A single frame with an empty transcript returns exactly -log p(blank); repeated forward passes are bitwise identical.

Requirements and limits

NVIDIA GPU with compute capability 8.0+.
fp32 lattice arithmetic; no bf16 path, since the accumulation would lose the small-probability tail the loss is made of.
alpha and beta are each [B, T, U+1] fp32, held between forward and backward; the [B, T, U+1, V] joint output dominates memory regardless.
One block per utterance: batch size is what fills the device.