@triton.jit
def _associative_rnn_scan_fwd_kernel(x_ptr, a_ptr, cum_a_ptr, out_ptr,
    stride_x_batch, stride_x_len, stride_x_dim, stride_a_batch,
    stride_a_len, stride_a_dim, stride_out_batch, stride_out_len,
    stride_out_dim, stride_cum_a_batch, stride_cum_a_len, stride_cum_a_dim,
    BLOCK_SIZE_LEN: tl.constexpr, BLOCK_SIZE_DIM: tl.constexpr):
    pid_batch = tl.program_id(0)
    pid_len = tl.program_id(1)
    pid_dim = tl.program_id(2)
    x_ptr += pid_batch * stride_x_batch
    a_ptr += pid_batch * stride_a_batch
    if cum_a_ptr is not None:
        cum_a_ptr += pid_batch * stride_cum_a_batch
    out_ptr += pid_batch * stride_out_batch
    offsets_dim = pid_dim * BLOCK_SIZE_DIM + tl.arange(0, BLOCK_SIZE_DIM)
    offsets_len = pid_len * BLOCK_SIZE_LEN + tl.arange(0, BLOCK_SIZE_LEN)
    x_ptrs = x_ptr + offsets_dim[None, :] * stride_x_dim + offsets_len[:, None
        ] * stride_x_len
    a_ptrs = a_ptr + offsets_dim[None, :] * stride_a_dim + offsets_len[:, None
        ] * stride_a_len
    out_ptrs = out_ptr + offsets_dim[None, :] * stride_out_dim + offsets_len[
        :, None] * stride_out_len
    if cum_a_ptr is not None:
        cum_a_ptrs = cum_a_ptr + offsets_dim[None, :
            ] * stride_cum_a_dim + offsets_len[:, None] * stride_cum_a_len
    x = tl.load(x_ptrs).to(tl.float32)
    a = tl.load(a_ptrs).to(tl.float32)
    cum_a, all_hiddens = tl.associative_scan(input=(a, x), axis=0,
        combine_fn=_associative_scan_op)
    mask = (offsets_len == (pid_len + 1) * BLOCK_SIZE_LEN - 1)[:, None]
    if cum_a_ptr is not None:
        tl.store(cum_a_ptrs, cum_a.to(cum_a_ptr.dtype.element_ty), mask=mask)
    tl.store(out_ptrs, all_hiddens.to(out_ptr.dtype.element_ty), mask=mask)
