File size: 4,190 Bytes
1a42548 34bf80d 1a42548 61cf87b 1a42548 34bf80d 1a42548 61cf87b 4f35b3d 34bf80d 4f35b3d 1a42548 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | ---
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.
|