code: def compute_alphas_kernel(
      acts: torch.Tensor,
      denom: torch.Tensor,
      alphas: torch.Tensor,
      llForward: torch.Tensor,
      xlen: torch.Tensor,
      ylen: torch.Tensor,
      mlabels: torch.Tensor,  # [B]
      minibatch: int,
      maxT: int,
      maxU: int,
      alphabet_size: int,
      blank_: int,
  ):
      """
      Compute alpha (forward variable) probabilities over the transduction step.
  
      Args:
          acts: Tensor of shape [B, T, U, V+1] flattened.
              Represents the logprobs activation tensor.
          denom: Tensor of shape [B, T, U] flattened. Represents the denominator of the
              logprobs activation tensor across entire vocabulary.
          alphas: Zero tensor of shape [B, T, U]. Will be updated inside the kernel
              with the forward variable probabilities.
          llForward: Zero tensor of shape [B]. Represents the log-likelihood of the
              forward pass. Returned as the forward pass loss that is reduced by
              the optimizer.
          xlen: Vector of length B which contains the actual acoustic sequence
              lengths in the padded activation tensor.
          ylen: Vector of length B which contains the actual target sequence
              lengths in the padded activation tensor.
          mlabels: Matrix of shape [B, U+1] (+1 here is due to <SOS> token
              - usually the RNNT blank). The matrix contains the padded target
              transcription that must be predicted.
          minibatch: Int representing the batch size.
          maxT: The maximum possible acoustic sequence length.
              Represents T in the logprobs tensor.
          maxU: The maximum possible target sequence length.
              Represents U in the logprobs tensor.
          alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
          blank_: Index of the RNNT blank token in the vocabulary.
              Generally the first or last token in the vocab.
  
      Updates:
          Kernel inplace updates the following inputs:
          -   alphas: forward variable scores.
          -   llForward: log-likelihood of forward variable.
      """
  
      # // launch B blocks, each block has U threads
      b = cuda.blockIdx.x  # // batch id
      u = cuda.threadIdx.x  # label id, u
      T = xlen[b]  # select AM length of current sample
      U = (
          ylen[b] + 1
      )  # select target length of current sample, +1 for the blank token
  
      labels: torch.Tensor = mlabels[
          b
      ]  # mb label start point, equivalent to mlabels + b * (maxU - 1)
      offset = b * maxT * maxU  # pointer indexing offset
  
      # alphas += offset # pointer offset, ignored since we explicitly add offset
  
      # Initilize alpha[b, t=0, u=0] for all b in B
      if u == 0:
          alphas[offset] = 0
  
      # sync until all alphas are initialized
      cuda.syncthreads()
  
      # Ordinary alpha calculations, broadcast across B=b and U=u
      # Look up forward variable calculation from rnnt_numpy.forward_pass()
      for n in range(1, T + U - 1):
          t = n - u
  
          if u == 0:
              # for t in range(1, T) step to initialize alphas[b, t, 0]
              if t > 0 and t < T:
                  alphas[offset + t * maxU + u] = alphas[
                      offset + (t - 1) * maxU + u
                  ] + logp(
                      denom,
                      acts,
                      maxT,
                      maxU,
                      alphabet_size,
                      b,
                      t - 1,
                      0,
                      blank_,
                  )
          elif u < U:
              # for u in range(1, U) step to initialize alphas[b, 0, u]
              if t == 0:
                  alphas[offset + u] = alphas[offset + u - 1] + logp(
                      denom,
                      acts,
                      maxT,
                      maxU,
                      alphabet_size,
                      b,
                      0,
                      u - 1,
                      labels[u - 1],
                  )
  
              # for t in range(1, T) for u in range(1, U) step to compute alphas[b, t, u]
              elif t > 0 and t < T:
                  no_emit = alphas[offset + (t - 1) * maxU + u] + logp(
                      denom,
                      acts,
                      maxT,
                      maxU,
                      alphabet_size,
                      b,
                      t - 1,
                      u,
                      blank_,
                  )
                  emit = alphas[offset + t * maxU + u - 1] + logp(
                      denom,
                      acts,
                      maxT,
                      maxU,
                      alphabet_size,
                      b,
                      t,
                      u - 1,
                      labels[u - 1],
                  )
  
                  alphas[
                      offset + t * maxU + u
                  ] = rnnt_helper.log_sum_exp(emit, no_emit)
  
          # sync across all B=b and U=u
          cuda.syncthreads()
  
      # After final sync, alphas[b, T-1, U - 1] + logprobs[b, T-1, U-1, blank]
      # + denom[b, T-1, U-1] gives log-likelihood of forward pass.
      if u == 0:
          loglike = alphas[offset + (T - 1) * maxU + U - 1] + logp(
              denom,
              acts,
              maxT,
              maxU,
              alphabet_size,
              b,
              T - 1,
              U - 1,
              blank_,
          )
          llForward[b] = loglike
  
file_path: espnet/espnet2/asr/transducer/rnnt_multi_blank/utils/cuda_utils/gpu_rnnt_kernel.py
imports: import math
  import torch
  from numba import cuda
  from espnet2.asr.transducer.rnnt_multi_blank.utils import rnnt_helper
label_desc: Use image node_id 2 for calling a global function with example usage: compute_alphas_kernel(acts, denom, alphas, llForward, xlen, ylen, mlabels, minibatch, maxT, maxU, alphabet_size, blank_) without return types