@triton.jit
def mlstm_triton_scan(Q, K, V, F, I, O, F_REDUCED, C, N, M, H, NH: tl.
    constexpr, S: tl.constexpr, D: tl.constexpr, SB: tl.constexpr, VB: tl.
    constexpr):
    bh_id = tl.program_id(0)
    sb_id = tl.program_id(1)
    vb_id = tl.program_id(2)
    num_sequence_blocks = tl.num_programs(1)
    batch_id = bh_id // NH
    head_id = bh_id % NH
    v_range = tl.arange(0, VB) + vb_id * VB
    v_range_3d = v_range[None, :, None]
    k_range = tl.arange(0, VB)
    sb_range = tl.arange(0, SB)
    sb_range_2d = tl.arange(0, SB)[:, None]
    sb_range_3d = tl.arange(0, SB)[:, None, None]
    sb_range_offset = tl.arange(0, SB) + sb_id * SB
    sb_range_offset_2d = sb_range_offset[:, None]
    batch_offset_fi = batch_id * NH * S + head_id * S
    batch_offset_f_reduced = (batch_id * NH * num_sequence_blocks + head_id *
        num_sequence_blocks)
    batch_offset_qkv = batch_id * NH * S * D + head_id * S * D
    batch_offset_n = (batch_id * NH * num_sequence_blocks * D + head_id *
        num_sequence_blocks * D)
    batch_offset_c = (batch_id * NH * num_sequence_blocks * D * D + head_id *
        num_sequence_blocks * D * D)
    f = tl.load(F + batch_offset_fi + sb_range_offset, sb_range_offset < S)
    i = tl.load(I + batch_offset_fi + sb_range_offset, sb_range_offset < S)
    f_reduced_range = batch_offset_f_reduced + sb_range * 0 + sb_id - 1
    f_reduced_mask = (sb_range == 0) & (sb_id != 0)
    f_reduced = tl.load(F_REDUCED + f_reduced_range, f_reduced_mask, other=1.0)
    vo_range = batch_offset_qkv + sb_range_offset_2d * D + v_range[None, :]
    vo_mask = (sb_range_offset_2d < S) & (v_range[None, :] < D)
    v = tl.load(V + vo_range, vo_mask)
    normalizer = tl.zeros((SB,), dtype=tl.float32)
    h = tl.zeros((SB, VB), dtype=tl.float32)
    k_scale_factor = tl.sqrt(tl.full((1,), D, dtype=tl.float32))
    for j in tl.range(tl.cdiv(D, VB)):
        k_range_ = batch_offset_qkv + sb_range_offset_2d * D + k_range[None, :]
        k_mask = (sb_range_offset_2d < S) & (k_range[None, :] < D)
        k = tl.load(K + k_range_, k_mask) / k_scale_factor
        q = tl.load(Q + k_range_, k_mask)
        c_range = batch_offset_c + sb_range_3d + (sb_id - 1
            ) * D * D + v_range_3d * D + k_range[None, None, :]
        c_mask = (sb_range_3d == 0) & (v_range_3d < D) & (k_range[None,
            None, :] < D) & (sb_id != 0)
        c_reduced = tl.load(C + c_range, c_mask, other=0)
        vk = v[:, :, None] * k[:, None, :] * i[:, None, None]
        f_tmp, vk_tmp = scan_op(f_reduced[:, None, None], c_reduced, f[:,
            None, None], vk)
        f_tmp = tl.broadcast_to(f_tmp, (SB, VB, VB))
        _, c = tl.associative_scan((f_tmp, vk_tmp), 0, scan_op)
        h += tl.sum(c * q[:, None, :], -1)
        n_range = batch_offset_n + sb_range_2d + (sb_id - 1) * D + k_range[
            None, :]
        n_mask = (sb_range_2d == 0) & (k_range[None, :] < D) & (sb_id != 0)
        n = tl.load(N + n_range, n_mask, other=0.0)
        f_tmp, n_tmp = scan_op(f_reduced[:, None], n, f[:, None], k * i[:,
            None])
        _, n = tl.associative_scan((tl.broadcast_to(f_tmp, (SB, VB)), n_tmp
            ), 0, scan_op)
        normalizer += tl.sum(n * q, -1)
        k_range += VB
    m = tl.load(M + batch_offset_fi + sb_range_offset, sb_range_offset < S)
    o = tl.load(O + vo_range, vo_mask)
    normalizer = tl.maximum(tl.abs(normalizer), tl.exp(-m))[:, None] + 1e-06
    h = h / normalizer * silu(o)
    tl.store(H + vo_range, h, vo_mask)
