| --- |
| library_name: kernels |
| license: apache-2.0 |
| --- |
| |
| # 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 |
| |
| ```python |
| 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). |
| |
| ## 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. |
| |
