@triton.jit
def reduce_mlstm_triton(F_REDUCED_IN, F_REDUCED_OUT, C, N, NH: tl.constexpr,
    D: tl.constexpr, NSB: tl.constexpr, BLOCK_SIZE: tl.constexpr):
    bh_id = tl.program_id(0)
    x_id = tl.program_id(1)
    y_id = tl.program_id(2)
    batch_id = bh_id // NH
    head_id = bh_id % NH
    nsb_range = tl.arange(0, NSB)
    nsb_range_2d = tl.arange(0, NSB)[:, None]
    nsb_range_3d = tl.arange(0, NSB)[:, None, None]
    block_range = tl.arange(0, BLOCK_SIZE)
    block_range_2d = block_range[None, :]
    x_range = block_range + x_id * BLOCK_SIZE
    x_range_3d = x_range[None, :, None]
    y_range = block_range + y_id * BLOCK_SIZE
    y_range_3d = y_range[None, None, :]
    batch_offset_f = batch_id * NH * NSB + head_id * NSB
    batch_offset_n = batch_id * NH * NSB * D + head_id * NSB * D
    batch_offset_c = batch_id * NH * NSB * D * D + head_id * NSB * D * D
    f = tl.load(F_REDUCED_IN + batch_offset_f + nsb_range)
    c_range = (batch_offset_c + nsb_range_3d * D * D + x_range_3d * D +
        y_range_3d)
    c_mask = (nsb_range_3d < NSB) & (x_range_3d < D) & (y_range_3d < D)
    c = tl.load(C + c_range, c_mask)
    f_reduced, c = tl.associative_scan((tl.broadcast_to(f[:, None, None], (
        NSB, BLOCK_SIZE, BLOCK_SIZE)), c), 0, scan_op)
    tl.store(C + c_range, c, c_mask)
    if x_id == 0 and y_id == 0:
        f_range = batch_offset_f + nsb_range_3d + block_range[None, :, None
            ] * 0 + block_range[None, None, :] * 0
        f_mask = (nsb_range_3d < NSB) & (block_range[None, :, None] == 0) & (
            block_range[None, None, :] == 0)
        tl.store(F_REDUCED_OUT + f_range, f_reduced, f_mask)
    if x_id == 0:
        n_range = (batch_offset_n + nsb_range_2d * D + block_range_2d + 
            y_id * BLOCK_SIZE)
        n_mask = (nsb_range_2d < NSB) & (block_range_2d + y_id * BLOCK_SIZE < D
            )
        n = tl.load(N + n_range, n_mask)
        _, n = tl.associative_scan((tl.broadcast_to(f[:, None], (NSB,
            BLOCK_SIZE)), n), 0, scan_op)
        tl.store(N + n_range, n, n_mask)
