dpx-decode

Dynamic-programming decoders for speech and alignment pipelines on NVIDIA GPUs, loadable through kernels: Viterbi decoding, monotonic DTW, and CTC forced alignment, batched, deterministic, and bitwise reproducible. The reference baselines are exhaustive fixed-point and float64 DP references, matched exactly. The integer paths run on the DPX fused min/max instructions, hardware on sm_90+ and compiler-emulated bit-identically on Ampere and Ada, so results are architecture-independent.

After the neural encoder of an ASR, diarization, or TTS pipeline has been accelerated, the decoders that follow it, max-plus and min-plus wavefronts over lattices, commonly still run on CPU one utterance at a time. This kernel runs them as batched GPU wavefronts: hundreds of utterances decode in a fraction of a millisecond, with tie-breaking documented and results reproducible to the bit.

A 48-state, 120-frame Viterbi decode: the max-plus wavefront fills the trellis, and the backtrace recovers the single best path through 48^120 state sequences, equal state for state to an exhaustive DP reference, with the int32 DPX path agreeing on every frame; a batch of 256 utterances decodes in 0.2 ms.

Usage

import torch
from kernels import get_kernel

dpx = get_kernel("phanerozoic/dpx-decode", version=1, trust_remote_code=True)

# CTC forced alignment (float path)
frames, score = dpx.ctc_forced_align(log_probs, targets, blank=0)

# Viterbi on the DPX fixed-point path
path, score = dpx.viterbi(dpx.quantize(emissions), dpx.quantize(transitions))

# word-timestamp style DTW
path, n, D = dpx.dtw(cost)

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

API

Each op accepts float32 or int32 inputs; quantize(x, scale) produces the fixed-point representation (default 256 quanta per unit).

Method

On the int path, Viterbi packs (score << 16 | state) so the DPX 3-way max resolves the argmax within the instruction; per-timestep renormalization bounds packed scores, and clamping at the int16 floor cannot change the argmax. DTW steps diagonal/down/up as in word-timestamp alignment with backtrace preferring diagonal then vertical. CTC alignment uses the standard blank-interleaved trellis with the repeated-label constraint, targets right-padded with -1. Tie-breaking is documented per op and deterministic on both paths. __vimin3_s32/__vimax3_s32 are single hardware instructions on sm_90 and newer and compiler-emulated bit-identically on Ampere and Ada.

Measured

Verified on L4 (sm_89, compiler-emulated DPX) and H200 (sm_90, hardware DPX): the int32 Viterbi path is bit-exact against a fixed-point reference replicating packing, clamping, and renormalization; the float Viterbi path reproduces a float64 reference's paths exactly; DTW satisfies endpoint, monotonicity, and accumulated-cost checks against exhaustive references; CTC alignment collapses to the transcript with path score equal to the summed per-frame log-probabilities; repeated runs are bitwise identical. A batch of 256 utterances at 48 states and 120 frames decodes in 0.17 ms.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+.
• Viterbi: S <= 32768; backpointer workspace is [B, T, S] int32.
• All log-domain inputs must be finite.

References

Viterbi (1967); Sakoe and Chiba (DTW, 1978); Graves et al., "Connectionist Temporal Classification" (2006); NVIDIA DPX instructions (Hopper, 2022).

License

Apache-2.0.