Add cuctc decoder (#3096)

Summary:
This PR implements a CUDA based ctc prefix beam search decoder.

Attach serveral benchmark results using V100 below:
|decoder type| model |datasets       | decoding time (secs)| beam size | batch size | model unit | subsampling times | vocab size |
|--------------|---------|------|-----------------|------------|-------------|------------|-----------------------|------------|
| cuctc |  conformer nemo    |dev clean        |7.68s | 8           |  32       | bpe         |    4  | 1000|
| cuctc |  conformer nemo   |dev clean  (sort by length)      |1.6s | 8           |  32       | bpe         |    4  | 1000|
| cuctc |  wav2vec2.0 torchaudio |dev clean                                |22s | 10           |  1       | char         |    2  | 29|
| cuctc |   conformer espnet   |aishell1 test                             | 5s | 10           |  24       | char         |    4  | 4233|

Note:
1.  The design is to parallel computation through batch and vocab axis, for loop the frames axis. So it's more friendly with smaller sequence lengths, larger vocab size comparing with CPU implementations.
2. WER is the same as CPU implementations. However, it can't decode with LM now.

Resolves: https://github.com/pytorch/audio/issues/2957.

Pull Request resolved: https://github.com/pytorch/audio/pull/3096

Reviewed By: nateanl

Differential Revision: D44709397

Pulled By: mthrok

fbshipit-source-id: 3078c54a2b44dc00eb4a81b4c657487eeff8c155

CMakeLists.txt

@@ -56,6 +56,7 @@ option(BUILD_SOX "Build libsox statically" ON)
 option(BUILD_KALDI "Build kaldi statically" ON)
 option(BUILD_RIR "Enable RIR simulation" ON)
 option(BUILD_RNNT "Enable RNN transducer" ON)
+option(BUILD_CUDA_CTC_DECODER "Build CUCTC decoder" OFF)
 option(BUILD_TORCHAUDIO_PYTHON_EXTENSION "Build Python extension" OFF)
 option(USE_FFMPEG "Enable ffmpeg-based features" OFF)
 option(USE_CUDA "Enable CUDA support" OFF)
@@ -155,3 +156,9 @@ endif()
 if (USE_FFMPEG)
   add_subdirectory(torchaudio/csrc/ffmpeg)
 endif()
+if (BUILD_CUDA_CTC_DECODER)
+  if (NOT USE_CUDA)
+    message(FATAL "BUILD_CUDA_CTC_DECODER=1 but USE_CUDA=0.")
+  endif()
+  add_subdirectory(torchaudio/csrc/cuctc)
+endif()

CONTRIBUTING.md

@@ -77,6 +77,7 @@ Some environmnet variables that change the build behavior
 - `BUILD_KALDI`: Determines whether build Kaldi extension. This is required for `kaldi_pitch` function. Default value is 1 on Linux/macOS and 0 on Windows.
 - `BUILD_RNNT`: Determines whether build RNN-T loss function. Default value is 1.
 - `BUILD_CTC_DECODER`: Determines whether build decoder features based on KenLM and FlashLight CTC decoder. Default value is 1.
+- `BUILD_CUDA_CTC_DECODER`: Determines whether build decoder features based on CUDA CTC decoder. Default value is 1. (`USE_CUDA` has to be 1.)
 
 Please check the [./tools/setup_helpers/extension.py](./tools/setup_helpers/extension.py) for the up-to-date detail.
 

docs/source/_templates/autosummary/cuda_ctc_decoder_class.rst

+..
+  autogenerated from source/_templates/autosummary/cuda_ctc_decoder_class.rst
+
+
+{#
+    ################################################################################
+    # autosummary template for CUCTCDecoder
+    # Since the class has multiple methods and support structure.
+    # we want to have them show up in the table of contents.
+    # The default class template does not do this, so we use custom one here.
+    ################################################################################
+#}
+
+{{ name | underline }}
+
+{%- if name != "CUCTCDecoder" %}
+
+.. autofunction:: {{fullname}}
+
+{%- else %}
+
+.. autoclass:: {{ fullname }}()
+
+Methods
+=======
+
+{%- for item in members %}
+{%- if not item.startswith('_') or item == "__call__" %}
+
+{{ item | underline("-") }}
+
+.. container:: py attribute
+
+   .. automethod:: {{[fullname, item] | join('.')}}
+
+{%- endif %}
+{%- endfor %}
+
+Support Structures
+==================
+
+{%- for item in ["CUCTCHypothesis"] %}
+
+{{ item | underline("-") }}
+
+.. autoclass:: torchaudio.models.decoder.{{item}}
+   :members:
+
+{%- endfor %}
+
+{%- endif %}

docs/source/index.rst

@@ -167,7 +167,7 @@ Tutorials
    :image: https://download.pytorch.org/torchaudio/tutorial-assets/thumbnails/streamwriter_basic_tutorial.gif
    :link: tutorials/streamwriter_basic_tutorial.html
    :tags: I/O,StreamWriter
-   
+
 .. customcarditem::
    :header: Playing media with StreamWriter
    :card_description: Learn how to play audio/video with <code>torchaudio.io.StreamWriter</code>.
@@ -265,7 +265,7 @@ Tutorials
    :image: https://download.pytorch.org/torchaudio/tutorial-assets/thumbnails/online_asr_tutorial.gif
    :link: tutorials/online_asr_tutorial.html
    :tags: Pipelines,ASR,RNNT,StreamReader
-   
+
 .. customcarditem::
    :header: Real-time microphone ASR with Emformer RNN-T
    :card_description: Learn how to transcribe speech fomr microphone with Emformer RNN-T (<code>torchaudio.pipelines.RNNTBundle</code>) and <code>torchaudio.io.StreamReader</code>.
@@ -286,7 +286,7 @@ Tutorials
    :image: https://download.pytorch.org/torchaudio/tutorial-assets/thumbnails/tacotron2_pipeline_tutorial.png
    :link: tutorials/tacotron2_pipeline_tutorial.html
    :tags: Pipelines,TTS-(Text-to-Speech)
-   
+
 .. customcarditem::
    :header: Speech Enhancement with MVDR Beamforming
    :card_description: Learn how to improve speech quality with MVDR Beamforming.

docs/source/models.decoder.rst

@@ -20,3 +20,19 @@ CTC Decoder
 .. rubric:: Tutorials using CTC Decoder
 
 .. minigallery:: torchaudio.models.decoder.CTCDecoder
+
+CUDA CTC Decoder
+----------------
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+   :template: autosummary/cuda_ctc_decoder_class.rst
+
+   CUCTCDecoder
+   cuda_ctc_decoder
+
+
+.. rubric:: Tutorials using CUDA CTC Decoder
+
+.. minigallery:: torchaudio.models.decoder.CUCTCDecoder

examples/asr/librispeech_cuda_ctc_decoder/README.md

+# Speech Recognition Inference with CUDA CTC Beam Search Decoder
+
+This is an example inference script for running decoding on the LibriSpeech dataset and [zipformer](https://github.com/k2-fsa/icefall/tree/master/egs/librispeech/ASR/pruned_transducer_stateless7_ctc) models, using a CUDA-based CTC beam search decoder that supports parallel decoding through batch and vocabulary axises.
+
+## Usage
+Additional command line parameters and information can is available with the `--help` option.
+
+Sample command
+
+```
+pip install sentencepiece
+
+# download pretrained files
+wget -nc https://huggingface.co/Zengwei/icefall-asr-librispeech-pruned-transducer-stateless7-ctc-2022-12-01/resolve/main/data/lang_bpe_500/bpe.model
+wget -nc https://huggingface.co/Zengwei/icefall-asr-librispeech-pruned-transducer-stateless7-ctc-2022-12-01/resolve/main/exp/cpu_jit.pt
+
+python inference.py \
+    --librispeech_path ./librispeech/ \
+    --split test-other \
+    --model ./cpu_jit.pt \
+    --bp-model ./bpe.model \
+    --beam-size 10 \
+    --blank-skip-threshold 0.95
+```
+
+## Results
+The table below contains throughput and WER benchmark results on librispeech test_other set between cuda ctc decoder and flashlight cpu decoder.
+
+(Note: batch_size=4, beam_size=10, nbest=10, vocab_size=500, no LM, Intel(R) Xeon(R) CPU E5-2698 v4 @ 2.20GHz, V100 GPU)
+
+| Decoder | Setting | WER (%) | N-Best Oracle WER (%) | Decoder Cost Time (seconds) |
+|:-----------|-----------:|-----------:|-----------:|-----------:|
+|CUDA decoder|blank_skip_threshold=0.95| 5.81 | 4.11 | 2.57 |
+|CUDA decoder|blank_skip_threshold=1.0 (no frame-skip)| 5.81 | 4.09 | 6.24 |
+|flashlight decoder|beam_size_token=10| 5.86 | 4.30 | 28.61 |
+|flashlight decoder|beam_size_token=vocab_size| 5.86 | 4.30 | 791.80 |

examples/asr/librispeech_cuda_ctc_decoder/inference.py

+import argparse
+import logging
+import time
+
+import sentencepiece as spm
+import torch
+import torchaudio
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader
+from torchaudio.models.decoder import ctc_decoder, cuda_ctc_decoder
+
+logger = logging.getLogger(__name__)
+
+
+def collate_wrapper(batch):
+    speeches, labels = [], []
+    for (speech, _, label, _, _, _) in batch:
+        speeches.append(speech)
+        labels.append(label.strip().lower().strip())
+    return speeches, labels
+
+
+def run_inference(args):
+    device = torch.device("cuda", 0)
+    model = torch.jit.load(args.model)
+    model.to(device)
+    model.eval()
+
+    bpe_model = spm.SentencePieceProcessor()
+    bpe_model.load(args.bpe_model)
+    vocabs = [bpe_model.id_to_piece(id) for id in range(bpe_model.get_piece_size())]
+    if args.using_cpu_decoder:
+        cpu_decoder = ctc_decoder(
+            lexicon=None,
+            tokens=vocabs,
+            lm=None,
+            nbest=args.nbest,
+            beam_size=args.beam_size,
+            beam_size_token=args.beam_size_token,
+            beam_threshold=args.beam_threshold,
+            blank_token="<blk>",
+            sil_token="<blk>",
+        )
+    else:
+        assert vocabs[0] == "<blk>", "idx of blank token has to be zero"
+        blank_frame_skip_threshold = float(torch.log(torch.tensor(args.blank_skip_threshold)))
+        cuda_decoder = cuda_ctc_decoder(
+            vocabs, nbest=args.nbest, beam_size=args.beam_size, blank_skip_threshold=blank_frame_skip_threshold
+        )
+
+    dataset = torchaudio.datasets.LIBRISPEECH(args.librispeech_path, url=args.split, download=True)
+
+    total_edit_distance, oracle_edit_distance, total_length = 0, 0, 0
+
+    data_loader = DataLoader(
+        dataset, batch_size=args.batch_size, num_workers=4, pin_memory=True, collate_fn=collate_wrapper
+    )
+
+    decoding_duration = 0
+    for idx, batch in enumerate(data_loader):
+        waveforms, transcripts = batch
+        waveforms = [wave.to(device) for wave in waveforms]
+        features = [torchaudio.compliance.kaldi.fbank(wave, num_mel_bins=80, snip_edges=False) for wave in waveforms]
+        feature_lengths = [f.size(0) for f in features]
+
+        features = pad_sequence(features, batch_first=True, padding_value=torch.log(torch.tensor(1e-10)))
+        feature_lengths = torch.tensor(feature_lengths, device=device)
+
+        encoder_out, encoder_out_lens = model.encoder(
+            x=features,
+            x_lens=feature_lengths,
+        )
+        nnet_output = model.ctc_output(encoder_out)
+        log_prob = torch.nn.functional.log_softmax(nnet_output, -1)
+
+        decoding_start = time.perf_counter()
+        preds = []
+        if args.using_cpu_decoder:
+            results = cpu_decoder(log_prob.cpu())
+            duration = time.perf_counter() - decoding_start
+            for i in range(len(results)):
+                ith_preds = bpe_model.decode([results[i][j].tokens.tolist() for j in range(len(results[i]))])
+                ith_preds = [pred.lower().split() for pred in ith_preds]
+                preds.append(ith_preds)
+        else:
+            results = cuda_decoder(log_prob, encoder_out_lens.to(torch.int32))
+            duration = time.perf_counter() - decoding_start
+            for i in range(len(results)):
+                ith_preds = bpe_model.decode([results[i][j].tokens for j in range(len(results[i]))])
+                ith_preds = [pred.lower().split() for pred in ith_preds]
+                preds.append(ith_preds)
+        decoding_duration += duration
+
+        for transcript, nbest_pred in zip(transcripts, preds):
+            total_edit_distance += torchaudio.functional.edit_distance(transcript.split(), nbest_pred[0])
+            oracle_edit_distance += min(
+                [torchaudio.functional.edit_distance(transcript.split(), nbest_pred[i]) for i in range(len(nbest_pred))]
+            )
+            total_length += len(transcript.split())
+
+        if idx % 10 == 0:
+            logger.info(
+                f"Processed elem {idx}; "
+                f"WER: {total_edit_distance / total_length}, "
+                f"Oracle WER: {oracle_edit_distance / total_length}, ",
+                f"decoding time for batch size {args.batch_size}: {duration}",
+            )
+
+    logger.info(
+        f"Final WER: {total_edit_distance / total_length}, ",
+        f"Oracle WER: {oracle_edit_distance / total_length}, ",
+        f"time for decoding {decoding_duration} [sec].",
+    )
+
+
+def _parse_args():
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawTextHelpFormatter,
+    )
+    parser.add_argument(
+        "--librispeech_path",
+        type=str,
+        help="folder where LibriSpeech is stored",
+        default="./librispeech",
+    )
+    parser.add_argument(
+        "--split",
+        type=str,
+        help="LibriSpeech dataset split",
+        choices=["dev-clean", "dev-other", "test-clean", "test-other"],
+        default="test-other",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="./cpu_jit.pt",
+        help="pretrained ASR model using CTC loss",
+    )
+    parser.add_argument(
+        "--bpe-model",
+        type=str,
+        default="./bpe.model",
+        help="bpe file for pretrained ASR model",
+    )
+    parser.add_argument(
+        "--nbest",
+        type=int,
+        default=10,
+        help="number of best hypotheses to return",
+    )
+    parser.add_argument(
+        "--beam-size",
+        type=int,
+        default=10,
+        help="beam size for determining number of hypotheses to store",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=4,
+        help="batch size for decoding",
+    )
+    parser.add_argument(
+        "--blank-skip-threshold",
+        type=float,
+        default=0.95,
+        help="skip frames where prob_blank > 0.95, https://ieeexplore.ieee.org/document/7736093",
+    )
+    parser.add_argument("--debug", action="store_true", help="whether to use debug level for logging")
+    # cpu decoder specific parameters
+    parser.add_argument("--using-cpu-decoder", action="store_true", help="whether to use flashlight cpu ctc decoder")
+    parser.add_argument("--beam-threshold", type=int, default=50, help="beam threshold for pruning hypotheses")
+    parser.add_argument(
+        "--beam-size-token",
+        type=int,
+        default=None,
+        help="number of tokens to consider at each beam search step",
+    )
+    return parser.parse_args()
+
+
+def _init_logger(debug):
+    fmt = "%(asctime)s %(message)s" if debug else "%(message)s"
+    level = logging.DEBUG if debug else logging.INFO
+    logging.basicConfig(format=fmt, level=level, datefmt="%Y-%m-%d %H:%M:%S")
+
+
+def _main():
+    args = _parse_args()
+    _init_logger(args.debug)
+    run_inference(args)
+
+
+if __name__ == "__main__":
+    _main()

test/torchaudio_unittest/common_utils/__init__.py

@@ -7,6 +7,7 @@ from .case_utils import (
     skipIfCudaSmallMemory,
     skipIfNoAudioDevice,
     skipIfNoCtcDecoder,
+    skipIfNoCuCtcDecoder,
     skipIfNoCuda,
     skipIfNoExec,
     skipIfNoFFmpeg,
@@ -44,6 +45,7 @@ __all__ = [
     "is_ffmpeg_available",
     "skipIfNoAudioDevice",
     "skipIfNoCtcDecoder",
+    "skipIfNoCuCtcDecoder",
     "skipIfNoCuda",
     "skipIfCudaSmallMemory",
     "skipIfNoExec",

test/torchaudio_unittest/common_utils/case_utils.py

@@ -116,6 +116,7 @@ def is_ffmpeg_available():
 
 
 _IS_CTC_DECODER_AVAILABLE = None
+_IS_CUDA_CTC_DECODER_AVAILABLE = None
 
 
 def is_ctc_decoder_available():
@@ -130,6 +131,18 @@ def is_ctc_decoder_available():
     return _IS_CTC_DECODER_AVAILABLE
 
 
+def is_cuda_ctc_decoder_available():
+    global _IS_CUDA_CTC_DECODER_AVAILABLE
+    if _IS_CUDA_CTC_DECODER_AVAILABLE is None:
+        try:
+            from torchaudio.models.decoder import CUCTCDecoder  # noqa: F401
+
+            _IS_CUDA_CTC_DECODER_AVAILABLE = True
+        except Exception:
+            _IS_CUDA_CTC_DECODER_AVAILABLE = False
+    return _IS_CUDA_CTC_DECODER_AVAILABLE
+
+
 def _eval_env(var, default):
     if var not in os.environ:
         return default
@@ -236,6 +249,11 @@ skipIfNoCtcDecoder = _skipIf(
     reason="CTC decoder not available.",
     key="NO_CTC_DECODER",
 )
+skipIfNoCuCtcDecoder = _skipIf(
+    not is_cuda_ctc_decoder_available(),
+    reason="CUCTC decoder not available.",
+    key="NO_CUCTC_DECODER",
+)
 skipIfRocm = _skipIf(
     _eval_env("TORCHAUDIO_TEST_WITH_ROCM", default=False),
     reason="The test doesn't currently work on the ROCm stack.",

test/torchaudio_unittest/models/decoder/cuda_ctc_decoder_test.py

+import torch
+from torchaudio_unittest.common_utils import (
+    get_asset_path,
+    skipIfNoCuCtcDecoder,
+    skipIfNoCuda,
+    TempDirMixin,
+    TorchaudioTestCase,
+)
+
+NUM_TOKENS = 7
+
+
+@skipIfNoCuda
+@skipIfNoCuCtcDecoder
+class CUCTCDecoderTest(TempDirMixin, TorchaudioTestCase):
+    def _get_decoder(self, tokens=None, **kwargs):
+        from torchaudio.models.decoder import cuda_ctc_decoder
+
+        if tokens is None:
+            tokens = get_asset_path("decoder/tokens.txt")
+
+        return cuda_ctc_decoder(
+            tokens=tokens,
+            beam_size=5,
+            **kwargs,
+        )
+
+    def _get_emissions(self):
+        B, T, N = 4, 15, NUM_TOKENS
+
+        emissions = torch.rand(B, T, N).cuda()
+        emissions = torch.nn.functional.log_softmax(emissions, -1)
+
+        return emissions
+
+    def test_construct_basic_decoder_path(self):
+        tokens_path = get_asset_path("decoder/tokens.txt")
+        self._get_decoder(tokens=tokens_path)
+
+    def test_construct_basic_decoder_tokens(self):
+        tokens = ["-", "|", "f", "o", "b", "a", "r"]
+        self._get_decoder(tokens=tokens)
+
+    def test_shape(self):
+        log_probs = self._get_emissions()
+        encoder_out_lens = torch.tensor([15, 14, 13, 12], dtype=torch.int32).cuda()
+        decoder = self._get_decoder()
+        results = decoder(log_probs, encoder_out_lens)
+        self.assertEqual(len(results), log_probs.shape[0])

third_party/LICENSES_BUNDLED.txt

+The Torchaudio repository and source distributions bundle several libraries that are 
+compatibly licensed.  We list some here.
+
+Name: cuctc
+License: BSD-2-Clause (Files without specific notes)
+BSD-3-Clause File:
+     torchaudio/csrc/cuctc/src/ctc_fast_divmod.cuh,
+Apache 2.0 Files:
+     torchaudio/csrc/cuctc/src/bitonic_topk
+  For details, see: cuctc/LICENSE,
+     torchaudio/csrc/cuctc/src/bitonic_topk/LICENSE

tools/setup_helpers/extension.py

@@ -40,6 +40,7 @@ _BUILD_RNNT = _get_build("BUILD_RNNT", True)
 _USE_FFMPEG = _get_build("USE_FFMPEG", False)
 _USE_ROCM = _get_build("USE_ROCM", torch.backends.cuda.is_built() and torch.version.hip is not None)
 _USE_CUDA = _get_build("USE_CUDA", torch.backends.cuda.is_built() and torch.version.hip is None)
+_BUILD_CUDA_CTC_DECODER = _get_build("BUILD_CUDA_CTC_DECODER", _USE_CUDA)
 _USE_OPENMP = _get_build("USE_OPENMP", True) and "ATen parallel backend: OpenMP" in torch.__config__.parallel_info()
 _TORCH_CUDA_ARCH_LIST = os.environ.get("TORCH_CUDA_ARCH_LIST", None)
 
@@ -56,6 +57,13 @@ def get_ext_modules():
                 Extension(name="torchaudio.lib._torchaudio_sox", sources=[]),
             ]
         )
+    if _BUILD_CUDA_CTC_DECODER:
+        modules.extend(
+            [
+                Extension(name="torchaudio.lib.libctc_prefix_decoder", sources=[]),
+                Extension(name="torchaudio.lib.pybind11_prefixctc", sources=[]),
+            ]
+        )
     if _USE_FFMPEG:
         modules.extend(
             [
@@ -110,6 +118,7 @@ class CMakeBuild(build_ext):
             f"-DBUILD_KALDI:BOOL={'ON' if _BUILD_KALDI else 'OFF'}",
             f"-DBUILD_RIR:BOOL={'ON' if _BUILD_RIR else 'OFF'}",
             f"-DBUILD_RNNT:BOOL={'ON' if _BUILD_RNNT else 'OFF'}",
+            f"-DBUILD_CUDA_CTC_DECODER:BOOL={'ON' if _BUILD_CUDA_CTC_DECODER else 'OFF'}",
             "-DBUILD_TORCHAUDIO_PYTHON_EXTENSION:BOOL=ON",
             f"-DUSE_ROCM:BOOL={'ON' if _USE_ROCM else 'OFF'}",
             f"-DUSE_CUDA:BOOL={'ON' if _USE_CUDA else 'OFF'}",

torchaudio/csrc/cuctc/CMakeLists.txt

+# Custom CMakeLists for building cuda ctc decoder
+
+set(CMAKE_CXX_VISIBILITY_PRESET default)
+
+# the following line is added in order to export symbols when building on Windows
+# this approach has some limitations as documented in https://github.com/pytorch/pytorch/pull/3650
+if (MSVC)
+  set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS ON)
+endif()
+
+set(
+  libctc_prefix_decoder_src
+  src/ctc_prefix_decoder.cpp
+  src/ctc_prefix_decoder_kernel_v2.cu
+  )
+
+set(
+  additional_libs
+  )
+
+list(
+  APPEND
+  additional_libs
+  cuda_deps
+  )
+
+torchaudio_library(
+  libctc_prefix_decoder
+  "${libctc_prefix_decoder_src}"
+  "${CMAKE_CURRENT_SOURCE_DIR}"
+  "${additional_libs}"
+  ""
+  )
+
+if (BUILD_TORCHAUDIO_PYTHON_EXTENSION)
+  torchaudio_extension(
+    pybind11_prefixctc
+    src/python_binding.cpp
+    "${CMAKE_CURRENT_SOURCE_DIR}"
+    "libctc_prefix_decoder;${additional_libs}"
+    ""
+    )
+endif()

torchaudio/csrc/cuctc/LICENSE

+BSD 2-Clause License
+
+Copyright (c) 2023 Nvidia 
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

torchaudio/csrc/cuctc/include/ctc_prefix_decoder.h

+// Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions
+// are met:
+//  * Redistributions of source code must retain the above copyright
+//    notice, this list of conditions and the following disclaimer.
+//  * Redistributions in binary form must reproduce the above copyright
+//    notice, this list of conditions and the following disclaimer in the
+//    documentation and/or other materials provided with the distribution.
+//  * Neither the name of NVIDIA CORPORATION nor the names of its
+//    contributors may be used to endorse or promote products derived
+//    from this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+// PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#ifndef __ctc_prefix_decoder_h_
+#define __ctc_prefix_decoder_h_
+
+#include <cuda_runtime.h>
+#include <cstdint>
+#include <tuple>
+#include <vector>
+namespace cu_ctc {
+
+struct InternalData;
+std::uintptr_t prefixCTC_alloc(std::uintptr_t stream_ptr);
+void prefixCTC_free(std::uintptr_t inter_data_ptr);
+
+std::tuple<size_t, int> calculate_require_buff_and_init_internal_data(
+    InternalData* inter_data,
+    int batch_size,
+    int seq_len,
+    int vocab_size,
+    int beam,
+    std::uintptr_t buff_ptr,
+    size_t buff_size,
+    float* log_prob_data_ptr,
+    int* original_lens,
+    const std::vector<int>& prob_sizes,
+    const std::vector<int>& prob_strides,
+    int blid,
+    float threshold);
+int ctc_beam_search_decoder_batch_gpu(
+    InternalData* inter_data,
+    float* pp,
+    int blid,
+    int spid,
+    int* clist,
+    int* clen,
+    float* score);
+
+} // namespace cu_ctc
+
+#endif

torchaudio/csrc/cuctc/include/ctc_prefix_decoder_host.h

+// Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions
+// are met:
+//  * Redistributions of source code must retain the above copyright
+//    notice, this list of conditions and the following disclaimer.
+//  * Redistributions in binary form must reproduce the above copyright
+//    notice, this list of conditions and the following disclaimer in the
+//    documentation and/or other materials provided with the distribution.
+//  * Neither the name of NVIDIA CORPORATION nor the names of its
+//    contributors may be used to endorse or promote products derived
+//    from this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+// PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#ifndef __ctc_prefix_decoder_host_h_
+#define __ctc_prefix_decoder_host_h_
+
+#include <cuda_runtime.h>
+
+#define CUDA_CHECK(X)                                   \
+  do {                                                  \
+    auto result = X;                                    \
+    if (result != cudaSuccess) {                        \
+      const char* p_err_str = cudaGetErrorName(result); \
+      fprintf(                                          \
+          stderr,                                       \
+          "File %s Line %d %s returned %s.\n",          \
+          __FILE__,                                     \
+          __LINE__,                                     \
+          #X,                                           \
+          p_err_str);                                   \
+      abort();                                          \
+    }                                                   \
+  } while (0)
+
+#define CHECK(X, ERROR_INFO)                        \
+  do {                                              \
+    auto result = (X);                              \
+    if (!result) {                                  \
+      fprintf(                                      \
+          stderr,                                   \
+          " File %s Line %d %s ERROR_INFO: %s .\n", \
+          __FILE__,                                 \
+          __LINE__,                                 \
+          #X,                                       \
+          ERROR_INFO);                              \
+      abort();                                      \
+    }                                               \
+  } while (0)
+
+namespace cu_ctc {
+
+struct LogProb;
+int init_log_prob_and_cal_max_select_seq_len(
+    LogProb* log_prob_struct,
+    int blid,
+    float threshold,
+    cudaStream_t stream);
+
+int CTC_prob_matrix_V2(
+    LogProb* log_prob_struct,
+    int step,
+    float2* pprev,
+    float* ptable,
+    float* ptablen,
+    int* clast,
+    int lc,
+    int ldc,
+    int beam,
+    int ldbeam,
+    int bs,
+    int blid,
+    int spid,
+    cudaStream_t stream);
+int CTC_prob_merge_V2(
+    LogProb* log_prob_struct,
+    int step,
+    float* ptable,
+    float* ptablen,
+    int* ptid,
+    int* clast,
+    int* clist,
+    int* clen,
+    int lc,
+    int ldc,
+    int beam,
+    int ldbeam,
+    int ldseq_len,
+    int bs,
+    cudaStream_t stream,
+    int blid);
+
+int CTC_prob_first_step_V2(
+    LogProb* log_prob_struct,
+    int step,
+    float2* pprev,
+    int* ptid,
+    int* clast,
+    int* clen,
+    int* clist,
+    int beam,
+    int ldbeam,
+    int ldseq_len,
+    int bs,
+    float* score,
+    cudaStream_t stream,
+    int blid);
+int CTC_prob_topK_V2(
+    LogProb* log_prob_struct,
+    int step,
+    float2* pprev,
+    float* ptable,
+    float* ptablen,
+    int* ptid,
+    int* clast,
+    int* clen,
+    int* clen2,
+    int* clist,
+    int* clist2,
+    int lc,
+    int ldc,
+    int beam,
+    int ldbeam,
+    int ldseq_len,
+    int blid,
+    int bs,
+    float* score,
+    float* topk_key_buff,
+    int* topk_value_buff,
+    cudaStream_t stream,
+    bool is_last_step);
+int CTC_copy_list_len_for_differnet_parity(
+    LogProb* log_prob_struct,
+    int step,
+    int max_select_seq_len,
+    int* clen,
+    int* clen2,
+    int* clist,
+    int* clist2,
+    int bs,
+    int beam,
+    int ldbeam,
+    int ldseq_len,
+    cudaStream_t stream);
+
+} // namespace cu_ctc
+
+#endif

torchaudio/csrc/cuctc/src/bitonic_topk/LICENSE

+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "{}"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2020 NVIDIA Corporation
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

torchaudio/csrc/cuctc/src/bitonic_topk/bitonic_sort.cuh

+/**
+ *  Modified from Rapidsai/raft(https://github.com/rapidsai/raft)
+ *
+ */
+
+#pragma once
+#include <cstdint>
+
+namespace cu_ctc {
+
+namespace topk {
+static constexpr int WarpSize = 32;
+template <typename IntType>
+constexpr inline __host__ __device__ bool isPo2(IntType num) {
+  return (num && !(num & (num - 1)));
+}
+
+inline __device__ int laneId() {
+  int id;
+  asm("mov.s32 %0, %%laneid;" : "=r"(id));
+  return id;
+}
+/**
+ * @brief Shuffle the data inside a warp
+ * @tparam T the data type (currently assumed to be 4B)
+ * @param val value to be shuffled
+ * @param laneMask mask to be applied in order to perform xor shuffle
+ * @param width lane width
+ * @param mask mask of participating threads (Volta+)
+ * @return the shuffled data
+ */
+template <typename T>
+inline __device__ T shfl_xor(
+    T val,
+    int laneMask,
+    int width = WarpSize,
+    uint32_t mask = 0xffffffffu) {
+#if CUDART_VERSION >= 9000
+  return __shfl_xor_sync(mask, val, laneMask, width);
+#else
+  return __shfl_xor(val, laneMask, width);
+#endif
+}
+
+/**
+ * @brief Shuffle the data inside a warp
+ * @tparam T the data type (currently assumed to be 4B)
+ * @param val value to be shuffled
+ * @param srcLane lane from where to shuffle
+ * @param width lane width
+ * @param mask mask of participating threads (Volta+)
+ * @return the shuffled data
+ */
+template <typename T>
+inline __device__ T
+shfl(T val, int srcLane, int width = WarpSize, uint32_t mask = 0xffffffffu) {
+#if CUDART_VERSION >= 9000
+  return __shfl_sync(mask, val, srcLane, width);
+#else
+  return __shfl(val, srcLane, width);
+#endif
+}
+
+/** warp-wide any boolean aggregator */
+inline __device__ bool any(bool inFlag, uint32_t mask = 0xffffffffu) {
+#if CUDART_VERSION >= 9000
+  inFlag = __any_sync(mask, inFlag);
+#else
+  inFlag = __any(inFlag);
+#endif
+  return inFlag;
+}
+
+template <typename T>
+constexpr T lower_bound() {
+  if constexpr (
+      std::numeric_limits<T>::has_infinity &&
+      std::numeric_limits<T>::is_signed) {
+    return -std::numeric_limits<T>::infinity();
+  }
+  return std::numeric_limits<T>::lowest();
+}
+
+template <typename T>
+constexpr T upper_bound() {
+  if constexpr (std::numeric_limits<T>::has_infinity) {
+    return std::numeric_limits<T>::infinity();
+  }
+  return std::numeric_limits<T>::max();
+}
+
+namespace helpers {
+
+template <typename T>
+__device__ __forceinline__ void swap(T& x, T& y) {
+  T t = x;
+  x = y;
+  y = t;
+}
+
+template <typename T>
+__device__ __forceinline__ void conditional_assign(bool cond, T& ptr, T x) {
+  if (cond) {
+    ptr = x;
+  }
+}
+
+} // namespace helpers
+
+/**
+ * Warp-wide bitonic merge and sort.
+ * The data is strided among `warp_width` threads,
+ * e.g. calling `bitonic<4>(ascending=true).sort(arr)` takes a unique 4-element
+ * array as input of each thread in a warp and sorts them, such that for a fixed
+ * i, arr[i] are sorted within the threads in a warp, and for any i < j, arr[j]
+ * in any thread is not smaller than arr[i] in any other thread. When
+ * `warp_width < WarpSize`, the data is sorted within all subwarps of the warp
+ * independently.
+ *
+ * As an example, assuming `Size = 4`, `warp_width = 16`, and `WarpSize = 32`,
+ * sorting a permutation of numbers 0-63 in each subwarp yield the following
+ * result:
+ * `
+ *  arr_i \ laneId()
+ *       0   1   2   3   4   5   6   7   8   9  10  11  12  13  14  15    16  17
+ * 18 ... subwarp_1 subwarp_2 0   0   1   2   3   4   5   6   7   8   9  10  11
+ * 12  13  14  15     0   1   2 ... 1  16  17  18  19  20  21  22  23  24  25 26
+ * 27  28  29  30  31    16  17  18 ... 2  32  33  34  35  36  37  38  39  40 41
+ * 42  43  44  45  46  47    32  33  34 ... 3  48  49  50  51  52  53  54  55 56
+ * 57  58  59  60  61  62  63    48  49  50 ...
+ * `
+ *
+ * @tparam Size
+ *   number of elements processed in each thread;
+ *   i.e. the total data size is `Size * warp_width`.
+ *   Must be power-of-two.
+ *
+ */
+template <int Size = 1>
+class bitonic {
+  static_assert(
+      isPo2(Size),
+      "class bitonic<Size> , size should be power of 2 \n");
+
+ public:
+  /**
+   * Initialize bitonic sort config.
+   *
+   * @param ascending
+   *   the resulting order (true: ascending, false: descending).
+   * @param warp_width
+   *   the number of threads participating in the warp-level primitives;
+   *   the total size of the sorted data is `Size * warp_width`.
+   *   Must be power-of-two, not larger than the WarpSize.
+   */
+  __device__ __forceinline__ explicit bitonic(
+      bool ascending,
+      int warp_width = WarpSize)
+      : ascending_(ascending), warp_width_(warp_width) {}
+
+  bitonic(bitonic const&) = delete;
+  bitonic(bitonic&&) = delete;
+  auto operator=(bitonic const&) -> bitonic& = delete;
+  auto operator=(bitonic&&) -> bitonic& = delete;
+
+  /**
+   * You can think of this function in two ways:
+   *
+   *   1) Sort any bitonic sequence.
+   *   2) Merge two halfs of the input data assuming they're already sorted, and
+   * their order is opposite (i.e. either ascending, descending or vice-versa).
+   *
+   * The input pointers are unique per-thread.
+   * See the class description for the description of the data layout.
+   *
+   * @param keys
+   *   is a device pointer to a contiguous array of keys, unique per thread;
+   * must be at least `Size` elements long.
+   * @param payloads
+   *   are zero or more associated arrays of the same size as keys, which are
+   * sorted together with the keys; must be at least `Size` elements long.
+   */
+  template <typename KeyT, typename... PayloadTs>
+  __device__ __forceinline__ void merge(
+      KeyT* __restrict__ keys,
+      PayloadTs* __restrict__... payloads) const {
+    return bitonic<Size>::merge_(ascending_, warp_width_, keys, payloads...);
+  }
+
+  /**
+   * Sort the data.
+   * The input pointers are unique per-thread.
+   * See the class description for the description of the data layout.
+   *
+   * @param keys
+   *   is a device pointer to a contiguous array of keys, unique per thread;
+   * must be at least `Size` elements long.
+   * @param payloads
+   *   are zero or more associated arrays of the same size as keys, which are
+   * sorted together with the keys; must be at least `Size` elements long.
+   */
+  template <typename KeyT, typename... PayloadTs>
+  __device__ __forceinline__ void sort(
+      KeyT* __restrict__ keys,
+      PayloadTs* __restrict__... payloads) const {
+    return bitonic<Size>::sort_(ascending_, warp_width_, keys, payloads...);
+  }
+
+  /**
+   * @brief `merge` variant for the case of one element per thread
+   *        (pass input by a reference instead of a pointer).
+   *
+   * @param key
+   * @param payload
+   */
+  template <typename KeyT, typename... PayloadTs, int S = Size>
+  __device__ __forceinline__ auto merge(
+      KeyT& __restrict__ key,
+      PayloadTs& __restrict__... payload) const
+      -> std::enable_if_t<S == 1, void> // SFINAE to enable this for Size == 1
+                                        // only
+  {
+    static_assert(S == Size);
+    return merge(&key, &payload...);
+  }
+
+  /**
+   * @brief `sort` variant for the case of one element per thread
+   *        (pass input by a reference instead of a pointer).
+   *
+   * @param key
+   * @param payload
+   */
+  template <typename KeyT, typename... PayloadTs, int S = Size>
+  __device__ __forceinline__ auto sort(
+      KeyT& __restrict__ key,
+      PayloadTs& __restrict__... payload) const
+      -> std::enable_if_t<S == 1, void> // SFINAE to enable this for Size == 1
+                                        // only
+  {
+    static_assert(S == Size);
+    return sort(&key, &payload...);
+  }
+
+ private:
+  const int warp_width_;
+  const bool ascending_;
+
+  template <int AnotherSize>
+  friend class bitonic;
+
+  template <typename KeyT, typename... PayloadTs>
+  static __device__ __forceinline__ void merge_(
+      bool ascending,
+      int warp_width,
+      KeyT* __restrict__ keys,
+      PayloadTs* __restrict__... payloads) {
+#pragma unroll
+    for (int size = Size; size > 1; size >>= 1) {
+      const int stride = size >> 1;
+#pragma unroll
+      for (int offset = 0; offset < Size; offset += size) {
+#pragma unroll
+        for (int i = offset + stride - 1; i >= offset; i--) {
+          const int other_i = i + stride;
+          KeyT& key = keys[i];
+          KeyT& other = keys[other_i];
+          if (ascending ? key > other : key < other) {
+            helpers::swap(key, other);
+            (helpers::swap(payloads[i], payloads[other_i]), ...);
+          }
+        }
+      }
+    }
+    const int lane = laneId();
+#pragma unroll
+    for (int i = 0; i < Size; i++) {
+      KeyT& key = keys[i];
+      for (int stride = (warp_width >> 1); stride > 0; stride >>= 1) {
+        const bool is_second = lane & stride;
+        const KeyT other = shfl_xor(key, stride, warp_width);
+        const bool do_assign =
+            (ascending != is_second) ? key > other : key < other;
+
+        helpers::conditional_assign(do_assign, key, other);
+        // NB: don't put shfl_xor in a conditional; it must be called by all
+        // threads in a warp.
+        (helpers::conditional_assign(
+             do_assign, payloads[i], shfl_xor(payloads[i], stride, warp_width)),
+         ...);
+      }
+    }
+  }
+
+  template <typename KeyT, typename... PayloadTs>
+  static __device__ __forceinline__ void sort_(
+      bool ascending,
+      int warp_width,
+      KeyT* __restrict__ keys,
+      PayloadTs* __restrict__... payloads) {
+    if constexpr (Size == 1) {
+      const int lane = laneId();
+      for (int width = 2; width < warp_width; width <<= 1) {
+        bitonic<1>::merge_(lane & width, width, keys, payloads...);
+      }
+    } else {
+      constexpr int kSize2 = Size / 2;
+      bitonic<kSize2>::sort_(false, warp_width, keys, payloads...);
+      bitonic<kSize2>::sort_(
+          true, warp_width, keys + kSize2, (payloads + kSize2)...);
+    }
+    bitonic<Size>::merge_(ascending, warp_width, keys, payloads...);
+  }
+};
+} // namespace topk
+} // namespace cu_ctc

torchaudio/csrc/cuctc/src/bitonic_topk/pow2_utils.cuh

+/**
+ *  Modified from Rapidsai/raft(https://github.com/rapidsai/raft)
+ *
+ */
+
+#pragma once
+#include <type_traits>
+
+namespace cu_ctc {
+/**
+ * @brief Give logarithm of the number to base-2
+ * @tparam IntType data type (checked only for integers)
+ */
+template <typename IntType>
+constexpr __device__ IntType log2(IntType num, IntType ret = IntType(0)) {
+  return num <= IntType(1) ? ret : log2(num >> IntType(1), ++ret);
+}
+
+/**
+ * @brief Fast arithmetics and alignment checks for power-of-two values known at
+ * compile time.
+ *
+ * @tparam Value_ a compile-time value representable as a power-of-two.
+ */
+template <auto Value_>
+struct Pow2 {
+  typedef decltype(Value_) Type;
+  static constexpr Type Value = Value_;
+  static constexpr Type Log2 = log2(Value);
+  static constexpr Type Mask = Value - 1;
+
+  static_assert(std::is_integral<Type>::value, "Value must be integral.");
+  static_assert(Value && !(Value & Mask), "Value must be power of two.");
+
+#define Pow2_FUNC_QUALIFIER static constexpr __host__ __device__ __forceinline__
+#define Pow2_WHEN_INTEGRAL(I) std::enable_if_t<Pow2_IS_REPRESENTABLE_AS(I), I>
+#define Pow2_IS_REPRESENTABLE_AS(I) \
+  (std::is_integral<I>::value && Type(I(Value)) == Value)
+
+  /**
+   * Integer division by Value truncated toward zero
+   * (same as `x / Value` in C++).
+   *
+   *  Invariant: `x = Value * quot(x) + rem(x)`
+   */
+  template <typename I>
+  Pow2_FUNC_QUALIFIER Pow2_WHEN_INTEGRAL(I) quot(I x) noexcept {
+    if constexpr (std::is_signed<I>::value)
+      return (x >> I(Log2)) + (x < 0 && (x & I(Mask)));
+    if constexpr (std::is_unsigned<I>::value)
+      return x >> I(Log2);
+  }
+
+  /**
+   *  Remainder of integer division by Value truncated toward zero
+   *  (same as `x % Value` in C++).
+   *
+   *  Invariant: `x = Value * quot(x) + rem(x)`.
+   */
+  template <typename I>
+  Pow2_FUNC_QUALIFIER Pow2_WHEN_INTEGRAL(I) rem(I x) noexcept {
+    if constexpr (std::is_signed<I>::value)
+      return x < 0 ? -((-x) & I(Mask)) : (x & I(Mask));
+    if constexpr (std::is_unsigned<I>::value)
+      return x & I(Mask);
+  }
+
+  /**
+   * Integer division by Value truncated toward negative infinity
+   * (same as `x // Value` in Python).
+   *
+   * Invariant: `x = Value * div(x) + mod(x)`.
+   *
+   * Note, `div` and `mod` for negative values are slightly faster
+   * than `quot` and `rem`, but behave slightly different
+   * compared to normal C++ operators `/` and `%`.
+   */
+  template <typename I>
+  Pow2_FUNC_QUALIFIER Pow2_WHEN_INTEGRAL(I) div(I x) noexcept {
+    return x >> I(Log2);
+  }
+
+  /**
+   * x modulo Value operation (remainder of the `div(x)`)
+   * (same as `x % Value` in Python).
+   *
+   * Invariant: `mod(x) >= 0`
+   * Invariant: `x = Value * div(x) + mod(x)`.
+   *
+   * Note, `div` and `mod` for negative values are slightly faster
+   * than `quot` and `rem`, but behave slightly different
+   * compared to normal C++ operators `/` and `%`.
+   */
+  template <typename I>
+  Pow2_FUNC_QUALIFIER Pow2_WHEN_INTEGRAL(I) mod(I x) noexcept {
+    return x & I(Mask);
+  }
+
+#define Pow2_CHECK_TYPE(T)                                     \
+  static_assert(                                               \
+      std::is_pointer<T>::value || std::is_integral<T>::value, \
+      "Only pointer or integral types make sense here")
+
+  /**
+   * Tell whether the pointer or integral is Value-aligned.
+   * NB: for pointers, the alignment is checked in bytes, not in elements.
+   */
+  template <typename PtrT>
+  Pow2_FUNC_QUALIFIER bool isAligned(PtrT p) noexcept {
+    Pow2_CHECK_TYPE(PtrT);
+    if constexpr (Pow2_IS_REPRESENTABLE_AS(PtrT))
+      return mod(p) == 0;
+    if constexpr (!Pow2_IS_REPRESENTABLE_AS(PtrT))
+      return mod(reinterpret_cast<Type>(p)) == 0;
+  }
+
+  /** Tell whether two pointers have the same address modulo Value. */
+  template <typename PtrT, typename PtrS>
+  Pow2_FUNC_QUALIFIER bool areSameAlignOffsets(PtrT a, PtrS b) noexcept {
+    Pow2_CHECK_TYPE(PtrT);
+    Pow2_CHECK_TYPE(PtrS);
+    Type x, y;
+    if constexpr (Pow2_IS_REPRESENTABLE_AS(PtrT))
+      x = Type(mod(a));
+    else
+      x = mod(reinterpret_cast<Type>(a));
+    if constexpr (Pow2_IS_REPRESENTABLE_AS(PtrS))
+      y = Type(mod(b));
+    else
+      y = mod(reinterpret_cast<Type>(b));
+    return x == y;
+  }
+
+  /** Get this or next Value-aligned address (in bytes) or integral. */
+  template <typename PtrT>
+  Pow2_FUNC_QUALIFIER PtrT roundUp(PtrT p) noexcept {
+    Pow2_CHECK_TYPE(PtrT);
+    if constexpr (Pow2_IS_REPRESENTABLE_AS(PtrT))
+      return (p + PtrT(Mask)) & PtrT(~Mask);
+    if constexpr (!Pow2_IS_REPRESENTABLE_AS(PtrT)) {
+      auto x = reinterpret_cast<Type>(p);
+      return reinterpret_cast<PtrT>((x + Mask) & (~Mask));
+    }
+  }
+
+  /** Get this or previous Value-aligned address (in bytes) or integral. */
+  template <typename PtrT>
+  Pow2_FUNC_QUALIFIER PtrT roundDown(PtrT p) noexcept {
+    Pow2_CHECK_TYPE(PtrT);
+    if constexpr (Pow2_IS_REPRESENTABLE_AS(PtrT))
+      return p & PtrT(~Mask);
+    if constexpr (!Pow2_IS_REPRESENTABLE_AS(PtrT)) {
+      auto x = reinterpret_cast<Type>(p);
+      return reinterpret_cast<PtrT>(x & (~Mask));
+    }
+  }
+#undef Pow2_CHECK_TYPE
+#undef Pow2_IS_REPRESENTABLE_AS
+#undef Pow2_FUNC_QUALIFIER
+#undef Pow2_WHEN_INTEGRAL
+};
+
+}; // namespace cu_ctc

torchaudio/csrc/cuctc/src/bitonic_topk/warpsort_topk.cuh

+/**
+ *  Modified from Rapidsai/raft(https://github.com/rapidsai/raft)
+ *
+ */
+
+#pragma once
+
+#include <algorithm>
+#include <functional>
+#include <type_traits>
+#include "bitonic_sort.cuh"
+#include "pow2_utils.cuh"
+
+namespace cu_ctc {
+
+/*
+  Three APIs of different scopes are provided:
+    1. host function: warp_sort_topk()
+    2. block-wide API: class block_sort
+    3. warp-wide API: class warp_sort_filtered and class warp_sort_immediate
+
+
+  1. warp_sort_topk()
+    (see the docstring)
+
+  2. class block_sort
+    It can be regarded as a fixed size priority queue for a thread block,
+    although the API is not typical.
+    class warp_sort_filtered and warp_sort_immediate can be used to instantiate
+  block_sort.
+
+    It uses dynamic shared memory as an intermediate buffer.
+    So the required shared memory size should be calculated using
+    calc_smem_size_for_block_wide() and passed as the 3rd kernel launch
+  parameter.
+
+    To add elements to the queue, use add(T val, IdxT idx) with unique values
+  per-thread. Use WarpSortClass<...>::kDummy constant for the threads outside of
+  input bounds.
+
+    After adding is finished, function done() should be called. And finally,
+  store() is used to get the top-k result.
+
+    Example:
+      __global__ void kernel() {
+        block_sort<warp_sort_immediate, ...> queue(...);
+
+        for (IdxT i = threadIdx.x; i < len, i += blockDim.x) {
+          queue.add(in[i], in_idx[i]);
+        }
+
+        queue.done();
+        queue.store(out, out_idx);
+     }
+
+     int smem_size = calc_smem_size_for_block_wide<T>(...);
+     kernel<<<grid_dim, block_dim, smem_size>>>();
+
+
+  3. class warp_sort_filtered and class warp_sort_immediate
+    These two classes can be regarded as fixed size priority queue for a warp.
+    Usage is similar to class block_sort. No shared memory is needed.
+
+    The host function (warp_sort_topk) uses a heuristic to choose between these
+  two classes for sorting, warp_sort_immediate being chosen when the number of
+  inputs per warp is somewhat small (see the usage of
+  LaunchThreshold<warp_sort_immediate>::len_factor_for_choosing).
+
+    Example:
+      __global__ void kernel() {
+        warp_sort_immediate<...> queue(...);
+        int warp_id = threadIdx.x / WarpSize;
+        int lane_id = threadIdx.x % WarpSize;
+
+        for (IdxT i = lane_id; i < len, i += WarpSize) {
+          queue.add(in[i], idx[i]);
+        }
+
+        queue.done();
+        // each warp outputs to a different offset
+        queue.store(out + warp_id * k, out_idx + warp_id * k);
+      }
+ */
+
+namespace topk {
+static constexpr int kMaxCapacity = 256;
+/** Whether 'left` should indeed be on the left w.r.t. `right`. */
+template <bool Ascending, typename T>
+__device__ __forceinline__ auto is_ordered(T left, T right) -> bool {
+  if constexpr (Ascending) {
+    return left < right;
+  }
+  if constexpr (!Ascending) {
+    return left > right;
+  }
+}
+
+constexpr inline auto calc_capacity(int k) -> int {
+  int capacity = isPo2(k) ? k : (1 << (log2(k) + 1));
+  return capacity;
+}
+
+/**
+ * A fixed-size warp-level priority queue.
+ * By feeding the data through this queue, you get the `k <= Capacity`
+ * smallest/greatest values in the data.
+ *
+ * @tparam Capacity
+ *   maximum number of elements in the queue.
+ * @tparam Ascending
+ *   which comparison to use: `true` means `<`, collect the smallest elements,
+ *   `false` means `>`, collect the greatest elements.
+ * @tparam T
+ *   the type of keys (what is being compared)
+ * @tparam IdxT
+ *   the type of payload (normally, indices of elements), i.e.
+ *   the content sorted alongside the keys.
+ */
+template <int Capacity, bool Ascending, typename T, typename IdxT>
+class warp_sort {
+  static_assert(isPo2(Capacity));
+
+ public:
+  /**
+   *  The `empty` value for the choosen binary operation,
+   *  i.e. `Ascending ? upper_bound<T>() : lower_bound<T>()`.
+   */
+  static constexpr T kDummy = Ascending ? upper_bound<T>() : lower_bound<T>();
+  /** Width of the subwarp. */
+  static constexpr int kWarpWidth = std::min<int>(Capacity, WarpSize);
+  /** The number of elements to select. */
+  const int k;
+
+  /**
+   * Construct the warp_sort empty queue.
+   *
+   * @param k
+   *   number of elements to select.
+   */
+  __device__ warp_sort(int k) : k(k) {
+#pragma unroll
+    for (int i = 0; i < kMaxArrLen; i++) {
+      val_arr_[i] = kDummy;
+    }
+  }
+
+  /**
+   * Load k values from the pointers at the given position, and merge them in
+   * the storage.
+   *
+   * When it actually loads the values, it always performs some collective warp
+   * operations in the end, thus enforcing warp sync. This means, it's safe to
+   * call `store` with the same arguments after `load_sorted` without extra
+   * sync. Note, however, that this is not neccesarily true for the reverse
+   * order, because the access patterns of `store` and `load_sorted` are
+   * different.
+   *
+   * @param[in] in
+   *    a device pointer to a contiguous array, unique per-subwarp
+   *    (length: k <= kWarpWidth * kMaxArrLen).
+   * @param[in] in_idx
+   *    a device pointer to a contiguous array, unique per-subwarp
+   *    (length: k <= kWarpWidth * kMaxArrLen).
+   * @param[in] do_merge
+   *    must be the same for all threads within a subwarp of size `kWarpWidth`.
+   *    It serves as a conditional; when `false` the function does nothing.
+   *    We need it to ensure threads within a full warp don't diverge calling
+   * `bitonic::merge()`.
+   */
+  __device__ void load_sorted(
+      const T* in,
+      const IdxT* in_idx,
+      bool do_merge = true) {
+    if (do_merge) {
+      int idx = Pow2<kWarpWidth>::mod(laneId()) ^ Pow2<kWarpWidth>::Mask;
+#pragma unroll
+      for (int i = kMaxArrLen - 1; i >= 0; --i, idx += kWarpWidth) {
+        if (idx < k) {
+          T t = in[idx];
+          if (is_ordered<Ascending>(t, val_arr_[i])) {
+            val_arr_[i] = t;
+            idx_arr_[i] = in_idx[idx];
+          }
+        }
+      }
+    }
+    if (kWarpWidth < WarpSize || do_merge) {
+      topk::bitonic<kMaxArrLen>(Ascending, kWarpWidth)
+          .merge(val_arr_, idx_arr_);
+    }
+  }
+
+  /**
+   *  Save the content by the pointer location.
+   *
+   * @param[out] out
+   *   device pointer to a contiguous array, unique per-subwarp of size
+   * `kWarpWidth` (length: k <= kWarpWidth * kMaxArrLen).
+   * @param[out] out_idx
+   *   device pointer to a contiguous array, unique per-subwarp of size
+   * `kWarpWidth` (length: k <= kWarpWidth * kMaxArrLen).
+   */
+  __device__ void store(T* out, IdxT* out_idx) const {
+    int idx = Pow2<kWarpWidth>::mod(laneId());
+#pragma unroll kMaxArrLen
+    for (int i = 0; i < kMaxArrLen && idx < k; i++, idx += kWarpWidth) {
+      out[idx] = val_arr_[i];
+      out_idx[idx] = idx_arr_[i];
+    }
+  }
+
+ protected:
+  static constexpr int kMaxArrLen = Capacity / kWarpWidth;
+
+  T val_arr_[kMaxArrLen];
+  IdxT idx_arr_[kMaxArrLen];
+
+  /**
+   * Merge another array (sorted in the opposite direction) in the queue.
+   * Thanks to the other array being sorted in the opposite direction,
+   * it's enough to call bitonic.merge once to maintain the valid state
+   * of the queue.
+   *
+   * @tparam PerThreadSizeIn
+   *   the size of the other array per-thread (compared to `kMaxArrLen`).
+   *
+   * @param keys_in
+   *   the values to be merged in. Pointers are unique per-thread. The values
+   *   must already be sorted in the opposite direction.
+   *   The layout of `keys_in` must be the same as the layout of `val_arr_`.
+   * @param ids_in
+   *   the associated indices of the elements in the same format as `keys_in`.
+   */
+  template <int PerThreadSizeIn>
+  __device__ __forceinline__ void merge_in(
+      const T* __restrict__ keys_in,
+      const IdxT* __restrict__ ids_in) {
+#pragma unroll
+    for (int i = std::min(kMaxArrLen, PerThreadSizeIn); i > 0; i--) {
+      T& key = val_arr_[kMaxArrLen - i];
+      T other = keys_in[PerThreadSizeIn - i];
+      if (is_ordered<Ascending>(other, key)) {
+        key = other;
+        idx_arr_[kMaxArrLen - i] = ids_in[PerThreadSizeIn - i];
+      }
+    }
+    topk::bitonic<kMaxArrLen>(Ascending, kWarpWidth).merge(val_arr_, idx_arr_);
+  }
+};
+
+/**
+ * This version of warp_sort compares each input element against the current
+ * estimate of k-th value before adding it to the intermediate sorting buffer.
+ * This makes the algorithm do less sorting steps for long input sequences
+ * at the cost of extra checks on each step.
+ *
+ * This implementation is preferred for large len values.
+ */
+template <int Capacity, bool Ascending, typename T, typename IdxT>
+class warp_sort_filtered : public warp_sort<Capacity, Ascending, T, IdxT> {
+ public:
+  using warp_sort<Capacity, Ascending, T, IdxT>::kDummy;
+  using warp_sort<Capacity, Ascending, T, IdxT>::kWarpWidth;
+  using warp_sort<Capacity, Ascending, T, IdxT>::k;
+
+  __device__ warp_sort_filtered(int k)
+      : warp_sort<Capacity, Ascending, T, IdxT>(k), buf_len_(0), k_th_(kDummy) {
+#pragma unroll
+    for (int i = 0; i < kMaxBufLen; i++) {
+      val_buf_[i] = kDummy;
+    }
+  }
+
+  __device__ void add(T val, IdxT idx) {
+    // comparing for k_th should reduce the total amount of updates:
+    // `false` means the input value is surely not in the top-k values.
+    bool do_add = is_ordered<Ascending>(val, k_th_);
+    // merge the buf if it's full and we cannot add an element anymore.
+    if (any(buf_len_ + do_add > kMaxBufLen)) {
+      // still, add an element before merging if possible for this thread
+      if (do_add && buf_len_ < kMaxBufLen) {
+        add_to_buf_(val, idx);
+        do_add = false;
+      }
+      merge_buf_();
+    }
+    // add an element if necessary and haven't already.
+    if (do_add) {
+      add_to_buf_(val, idx);
+    }
+  }
+
+  __device__ void done() {
+    if (any(buf_len_ != 0)) {
+      merge_buf_();
+    }
+  }
+
+ private:
+  __device__ __forceinline__ void set_k_th_() {
+    // NB on using srcLane: it's ok if it is outside the warp size / width;
+    //                      the modulo op will be done inside the __shfl_sync.
+    // const int id = (k - 1) / kWarpWidth;
+    const int id = Pow2<kWarpWidth>::div(k - 1);
+#pragma unroll
+    for (int i = 0; i < kMaxArrLen; i++) {
+      if (i == id) {
+        k_th_ = shfl(val_arr_[i], k - 1, kWarpWidth);
+      }
+    }
+    // k_th_ = shfl(val_arr_[kMaxArrLen - 1], k - 1, kWarpWidth);
+  }
+
+  __device__ __forceinline__ void merge_buf_() {
+    topk::bitonic<kMaxBufLen>(!Ascending, kWarpWidth).sort(val_buf_, idx_buf_);
+    this->merge_in<kMaxBufLen>(val_buf_, idx_buf_);
+    buf_len_ = 0;
+    set_k_th_(); // contains warp sync
+#pragma unroll
+    for (int i = 0; i < kMaxBufLen; i++) {
+      val_buf_[i] = kDummy;
+    }
+  }
+
+  __device__ __forceinline__ void add_to_buf_(T val, IdxT idx) {
+    // NB: the loop is used here to ensure the constant indexing,
+    //     to not force the buffers spill into the local memory.
+#pragma unroll
+    for (int i = 0; i < kMaxBufLen; i++) {
+      if (i == buf_len_) {
+        val_buf_[i] = val;
+        idx_buf_[i] = idx;
+      }
+    }
+    buf_len_++;
+  }
+
+  using warp_sort<Capacity, Ascending, T, IdxT>::kMaxArrLen;
+  using warp_sort<Capacity, Ascending, T, IdxT>::val_arr_;
+  using warp_sort<Capacity, Ascending, T, IdxT>::idx_arr_;
+
+  static constexpr int kMaxBufLen = (Capacity <= 64) ? 2 : 4;
+
+  T val_buf_[kMaxBufLen];
+  IdxT idx_buf_[kMaxBufLen];
+  int buf_len_;
+
+  T k_th_;
+};
+
+/**
+ * This version of warp_sort adds every input element into the intermediate
+ * sorting buffer, and thus does the sorting step every `Capacity` input
+ * elements.
+ *
+ * This implementation is preferred for very small len values.
+ */
+template <int Capacity, bool Ascending, typename T, typename IdxT>
+class warp_sort_immediate : public warp_sort<Capacity, Ascending, T, IdxT> {
+ public:
+  using warp_sort<Capacity, Ascending, T, IdxT>::kDummy;
+  using warp_sort<Capacity, Ascending, T, IdxT>::kWarpWidth;
+  using warp_sort<Capacity, Ascending, T, IdxT>::k;
+
+  __device__ warp_sort_immediate(int k)
+      : warp_sort<Capacity, Ascending, T, IdxT>(k), buf_len_(0) {
+#pragma unroll
+    for (int i = 0; i < kMaxArrLen; i++) {
+      val_buf_[i] = kDummy;
+    }
+  }
+
+  __device__ void add(T val, IdxT idx) {
+    // NB: the loop is used here to ensure the constant indexing,
+    //     to not force the buffers spill into the local memory.
+#pragma unroll
+    for (int i = 0; i < kMaxArrLen; ++i) {
+      if (i == buf_len_) {
+        val_buf_[i] = val;
+        idx_buf_[i] = idx;
+      }
+    }
+
+    ++buf_len_;
+    if (buf_len_ == kMaxArrLen) {
+      topk::bitonic<kMaxArrLen>(!Ascending, kWarpWidth)
+          .sort(val_buf_, idx_buf_);
+      this->merge_in<kMaxArrLen>(val_buf_, idx_buf_);
+#pragma unroll
+      for (int i = 0; i < kMaxArrLen; i++) {
+        val_buf_[i] = kDummy;
+      }
+      buf_len_ = 0;
+    }
+  }
+
+  __device__ void done() {
+    if (buf_len_ != 0) {
+      topk::bitonic<kMaxArrLen>(!Ascending, kWarpWidth)
+          .sort(val_buf_, idx_buf_);
+      this->merge_in<kMaxArrLen>(val_buf_, idx_buf_);
+    }
+  }
+
+ private:
+  using warp_sort<Capacity, Ascending, T, IdxT>::kMaxArrLen;
+  using warp_sort<Capacity, Ascending, T, IdxT>::val_arr_;
+  using warp_sort<Capacity, Ascending, T, IdxT>::idx_arr_;
+
+  T val_buf_[kMaxArrLen];
+  IdxT idx_buf_[kMaxArrLen];
+  int buf_len_;
+};
+
+/**
+ * @brief Provide a ceiling division operation ie. ceil(a / b)
+ * @tparam IntType supposed to be only integers for now!
+ */
+template <typename IntType>
+constexpr inline __host__ __device__ IntType ceildiv(IntType a, IntType b) {
+  return (a + b - 1) / b;
+}
+template <typename IntType>
+constexpr inline __device__ IntType roundUp256(IntType num) {
+  // return (num + 255) / 256 * 256;
+  constexpr int MASK = 255;
+  return (num + MASK) & (~MASK);
+}
+
+template <typename T, typename IdxT>
+auto calc_smem_size_for_block_wide(int num_of_subwarp, int k) -> int {
+  return roundUp256(ceildiv(num_of_subwarp, 2) * sizeof(T) * k) +
+      ceildiv(num_of_subwarp, 2) * sizeof(IdxT) * k;
+}
+
+template <
+    template <int, bool, typename, typename>
+    class WarpSortWarpWide,
+    int Capacity,
+    bool Ascending,
+    typename T,
+    typename IdxT>
+class block_sort {
+  using queue_t = WarpSortWarpWide<Capacity, Ascending, T, IdxT>;
+
+ public:
+  __device__ block_sort(int k, uint8_t* smem_buf) : queue_(k) {
+    val_smem_ = reinterpret_cast<T*>(smem_buf);
+    const int num_of_warp = subwarp_align::div(blockDim.x);
+    idx_smem_ = reinterpret_cast<IdxT*>(
+        smem_buf + roundUp256(ceildiv(num_of_warp, 2) * sizeof(T) * k));
+  }
+
+  __device__ void add(T val, IdxT idx) {
+    queue_.add(val, idx);
+  }
+
+  /**
+   * At the point of calling this function, the warp-level queues consumed all
+   * input independently. The remaining work to be done is to merge them
+   * together.
+   *
+   * Here we tree-merge the results using the shared memory and block sync.
+   */
+  __device__ void done() {
+    queue_.done();
+
+    const int warp_id = subwarp_align::div(threadIdx.x);
+    // NB: there is no need for the second __synchthreads between .load_sorted
+    // and .store:
+    //     we shift the pointers every iteration, such that individual warps
+    //     either access the same locations or do not overlap with any of the
+    //     other warps. The access patterns within warps are different for the
+    //     two functions, but .load_sorted implies warp sync at the end, so
+    //     there is no need for __syncwarp either.
+    for (int shift_mask = ~0,
+             nwarps = subwarp_align::div(blockDim.x),
+             split = (nwarps + 1) >> 1;
+         nwarps > 1;
+         nwarps = split, split = (nwarps + 1) >> 1) {
+      if (warp_id < nwarps && warp_id >= split) {
+        int dst_warp_shift = (warp_id - (split & shift_mask)) * queue_.k;
+        queue_.store(val_smem_ + dst_warp_shift, idx_smem_ + dst_warp_shift);
+      }
+      __syncthreads();
+
+      shift_mask = ~shift_mask; // invert the mask
+      {
+        int src_warp_shift = (warp_id + (split & shift_mask)) * queue_.k;
+        // The last argument serves as a condition for loading
+        //  -- to make sure threads within a full warp do not diverge on
+        //  `bitonic::merge()`
+        queue_.load_sorted(
+            val_smem_ + src_warp_shift,
+            idx_smem_ + src_warp_shift,
+            warp_id < nwarps - split);
+      }
+    }
+  }
+
+  /** Save the content by the pointer location. */
+  __device__ void store(T* out, IdxT* out_idx) const {
+    if (threadIdx.x < subwarp_align::Value) {
+      queue_.store(out, out_idx);
+    }
+  }
+
+ private:
+  using subwarp_align = Pow2<queue_t::kWarpWidth>;
+  queue_t queue_;
+  T* val_smem_;
+  IdxT* idx_smem_;
+};
+
+} // namespace topk
+
+} // namespace cu_ctc