dpx-decode

Dynamic-programming decoders for speech and alignment pipelines on NVIDIA GPUs: Viterbi decoding, monotonic DTW, and CTC forced alignment, batched, deterministic, and bitwise reproducible. The integer paths run on the DPX fused min/max instructions introduced with Hopper (__vimin3_s32 / __vimax3_s32), single hardware instructions on sm_90 and newer and compiler-emulated bit-identically on Ampere and Ada, so results are architecture-independent.

These recurrences (max-plus and min-plus wavefronts) are the parts of ASR, diarization, and TTS dataset pipelines that commonly still run on CPU after the neural encoder has been accelerated.

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

Op	Recurrence	Outputs
`viterbi(emissions [B,T,S], transitions [S,S], priors [S])`	max-plus over predecessors	path [B,T], score [B]
`dtw(cost [B,N,M])`	3-way min-plus, monotonic steps	path, path length, accumulated cost matrix
`ctc_forced_align(log_probs [B,T,C], targets [B,L], blank)`	CTC-trellis Viterbi (2L+1 states)	per-frame labels [B,T], score [B]

Each op accepts float32 or int32 inputs; quantize(x, scale) produces the fixed-point representation (default 256 quanta per unit). 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. Tie-breaking is documented per op and deterministic on both paths.

Semantics and limits

Viterbi: S <= 32768; backpointer workspace is [B, T, S] int32.
DTW: steps diagonal/down/up as in word-timestamp alignment; backtrace prefers diagonal, then vertical.
CTC alignment: targets right-padded with -1; standard blank-interleaved trellis with the repeated-label constraint.
All log-domain inputs must be finite.

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.