[Paper Review] A GPU-based WFST Decoder with Exact Lattice Generation

Zhehuai Chen, Justin Luitjens+4 more|arXiv (Cornell University)|Apr 9, 2018

Speech Recognition and Synthesis28 references3 citations

TL;DR

This paper presents a GPU-accelerated WFST decoder for Kaldi that enables exact lattice generation and achieves 3–15× speedup over CPU decoding by leveraging atomic operations for token recombination, dynamic load balancing, and optimized lattice processing. The system supports general acoustic and language models, runs on diverse GPU architectures, and delivers identical recognition and confidence measure results with up to 46× speedup using sequence parallelism and multi-process service (MPS).

ABSTRACT

We describe initial work on an extension of the Kaldi toolkit that supports weighted finite-state transducer (WFST) decoding on Graphics Processing Units (GPUs). We implement token recombination as an atomic GPU operation in order to fully parallelize the Viterbi beam search, and propose a dynamic load balancing strategy for more efficient token passing scheduling among GPU threads. We also redesign the exact lattice generation and lattice pruning algorithms for better utilization of the GPUs. Experiments on the Switchboard corpus show that the proposed method achieves identical 1-best results and lattice quality in recognition and confidence measure tasks, while running 3 to 15 times faster than the single process Kaldi decoder. The above results are reported on different GPU architectures. Additionally we obtain a 46-fold speedup with sequence parallelism and multi-process service (MPS) in GPU.

Motivation & Objective

• To accelerate WFST decoding in automatic speech recognition (ASR) using GPU parallelism without requiring specialized acoustic or language models.
• To address the challenge of high memory usage and poor scalability of large language models (LMs) in GPU-based decoding by enabling 2-pass decoding with exact lattice generation.
• To design a general-purpose, portable GPU decoder that supports older GPU architectures (e.g., Kepler) and integrates seamlessly with existing Kaldi workflows.
• To improve decoding efficiency through optimized token recombination, dynamic load balancing, and parallel lattice processing with reduced synchronization and memory transfer overhead.

Proposed method

• Implement token recombination as an atomic GPU operation using CUDA's atomicMin to eliminate serialization and enable full parallelism in Viterbi beam search.
• Design a dynamic load balancing strategy to distribute token passing work across GPU threads, minimizing idle time and improving thread utilization.
• Redesign exact lattice generation and pruning algorithms for GPU execution, reducing data movement and kernel launch overhead.
• Use a two-pass decoding framework: first pass generates an exact lattice for post-processing, second pass performs rescoring with large language models.
• Leverage sequence parallelism and multi-process service (MPS) to reduce context switching and achieve up to 46× speedup on modern GPUs.
• Integrate the system into the Kaldi toolkit as an open-source extension, ensuring backward compatibility with all released Kaldi recipes.

Experimental results

Research questions

• RQ1 Can GPU-based WFST decoding achieve significant speedup over CPU decoding while maintaining identical 1-best and lattice quality?
• RQ2 How can token recombination in Viterbi beam search be efficiently parallelized on GPUs without precision loss or synchronization bottlenecks?
• RQ3 What dynamic load balancing strategy can effectively balance work across GPU threads during token passing in WFST decoding?
• RQ4 Can exact lattice generation and pruning be efficiently parallelized on GPUs to reduce memory transfer and kernel launch overhead?
• RQ5 To what extent can sequence parallelism and multi-process service (MPS) further accelerate GPU-based WFST decoding in a real-world ASR setup?

Key findings

• The GPU decoder achieves 3–15× speedup over the single-process CPU decoder on various GPU architectures, with identical 1-best recognition results and lattice quality.
• The system achieves a 46-fold speedup when using sequence parallelism and multi-process service (MPS), compared to only 1.8× on CPU under similar conditions.
• The use of atomic operations for token recombination eliminates critical sections, enabling compatibility with older GPU architectures such as Kepler.
• Dynamic load balancing improves performance over static load balancing and graph partitioning methods, achieving 15× speedup for 1-best decoding and 9.7× for lattice decoding.
• Lattice pruning contributes significantly to performance gains by reducing data transfer and kernel launch costs, with the GPU version outperforming CPU in both speed and scalability.
• The decoder maintains consistent performance across different language model sizes (13MB to 258MB), with memory usage capped at 11GB on a 12GB GPU, demonstrating effective memory optimization.