@triton.jit
def precompute_mlstm_triton_scan(K, V, F, I, F_REDUCED, C, N, 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)
    batch_id = bh_id // NH
    head_id = bh_id % NH
    num_sequence_blocks = tl.num_programs(1)
    v_range = tl.arange(0, VB) + vb_id * VB
    v_range_2d = v_range[None, :]
    v_range_3d = v_range[None, :, None]
    k_range = tl.arange(0, VB)
    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_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 + sb_range_offset + batch_offset_fi, sb_range_offset < S)
    i = tl.load(I + sb_range_offset + batch_offset_fi, sb_range_offset < S)
    v_range_ = batch_offset_qkv + sb_range_offset_2d * D + v_range_2d
    v_mask = (sb_range_offset_2d < S) & (v_range_2d < D)
    v = tl.load(V + v_range_, v_mask)
    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
        vk = v[:, :, None] * k[:, None, :] * i[:, None, None]
        _, c = tl.associative_scan((tl.broadcast_to(f[:, None, None], (SB,
            VB, VB)), vk), 0, scan_op)
        c_range = (batch_offset_c + sb_range_3d * 0 + sb_id * D * D + 
            v_range_3d * D + k_range[None, None, :])
        c_mask = (sb_range_3d == SB - 1) & (v_range_3d < D) & (k_range[None,
            None, :] < D)
        tl.store(C + c_range, c, c_mask)
        f_reduced, n = tl.associative_scan((tl.broadcast_to(f[:, None], (SB,
            VB)), i[:, None] * k), 0, scan_op)
        n_range = batch_offset_n + sb_range_2d * 0 + sb_id * D + k_range[
            None, :]
        n_mask = (sb_range_2d == SB - 1) & (k_range[None, :] < D)
        tl.store(N + n_range, n, n_mask)
        if j == 0:
            f_range = batch_offset_fi + sb_range_2d * 0 + sb_id + tl.arange(
                0, VB)[None, :]
            f_mask = (sb_range_2d == SB - 1) & (tl.arange(0, VB)[None, :] == 0)
            tl.store(F_REDUCED + f_range, f_reduced, f_mask)
        k_range += VB
