@triton.jit
def chunk_mlstm_fwd_kernel_C(k, v, C, n, m, i, f, initial_C, initial_n,
    initial_m, final_C, final_n, final_m, s_qk_h, s_qk_t, s_qk_d, s_vh_h,
    s_vh_t, s_vh_d, s_C_h, s_C_t, s_n_h, H: tl.constexpr, T: tl.constexpr,
    K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr,
    BV: tl.constexpr, NT: tl.constexpr, USE_INITIAL_STATE: tl.constexpr,
    STORE_FINAL_STATE: tl.constexpr):
    i_k, i_v, i_bC = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    if USE_INITIAL_STATE:
        p_C0 = tl.make_block_ptr(initial_C + i_bC * K * V, (K, V), (V, 1),
            (i_k * BK, i_v * BV), (BK, BV), (1, 0))
        p_n0 = tl.make_block_ptr(initial_n + i_bC * K, (K,), (1,), (i_k *
            BK,), (BK,), (0,))
        p_m0 = initial_m
        b_C = tl.load(p_C0, boundary_check=(0, 1)).to(tl.float32)
        b_n = tl.load(p_n0, boundary_check=(0,)).to(tl.float32)
        b_m = tl.load(p_m0).to(tl.float32)
    else:
        b_C = tl.zeros([BK, BV], dtype=tl.float32)
        b_n = tl.zeros([BK], dtype=tl.float32)
        b_m = 0.0
    b_m_next = 0.0
    for i_t in range(NT):
        p_k = tl.make_block_ptr(k + i_bC * s_qk_h, (K, T), (s_qk_d, s_qk_t),
            (i_k * BK, i_t * BT), (BK, BT), (0, 1))
        p_v = tl.make_block_ptr(v + i_bC * s_vh_h, (T, V), (s_vh_t, s_vh_d),
            (i_t * BT, i_v * BV), (BT, BV), (1, 0))
        p_C = tl.make_block_ptr(C + i_bC * s_C_h + i_t * K * V, (K, V), (
            s_C_t, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
        p_n = tl.make_block_ptr(n + i_bC * s_n_h + i_t * K, (K,), (1,), (
            i_k * BK,), (BK,), (0,))
        tl.store(p_C, b_C.to(p_C.dtype.element_ty), boundary_check=(0, 1))
        tl.store(p_n, b_n.to(p_n.dtype.element_ty), boundary_check=(0,))
        tl.store(m + i_bC * (NT + 1) + i_t, b_m)
        b_k = tl.load(p_k, boundary_check=(0, 1))
        b_v = tl.load(p_v, boundary_check=(0, 1))
        b_f_last = tl.load(f + i_bC * T + i_t * BT + BT - 1)
        b_f = tl.load(f + i_bC * T + i_t * BT + tl.arange(0, BT))
        b_i = tl.load(i + i_bC * T + i_t * BT + tl.arange(0, BT))
        b_g = b_i + b_f_last - b_f
        b_m_next, _ = tl.max(b_g)
        b_m_next = tl.maximum(b_f_last + b_m, b_m_next)
        b_C *= tl.math.exp2(b_f_last - b_m_next + b_m)
        b_n *= tl.math.exp2(b_f_last - b_m_next + b_m)
        b_C += tl.dot(b_k, (b_v * tl.math.exp2(b_g - b_m_next)[:, None]).to
            (b_k.dtype), allow_tf32=False)
        b_n += tl.sum(b_k * tl.math.exp2(b_g - b_m_next), axis=1)
        b_m = b_m_next
    if STORE_FINAL_STATE:
        p_Ct = tl.make_block_ptr(final_C + i_bC * K * V, (K, V), (V, 1), (
            i_k * BK, i_v * BV), (BK, BV), (1, 0))
        p_n = tl.make_block_ptr(final_n + i_bC * K, (K,), (1,), (i_k * BK,),
            (BK,), (0,))
        tl.store(p_Ct, b_C.to(p_Ct.dtype.element_ty), boundary_check=(0, 1))
        tl.store(p_n, b_n.to(p_n.dtype.element_ty), boundary_check=(0,))
        tl.store(final_m + i_bC, b_m)
        tl.store(m + i_bC * (NT + 1) + NT, b_m)
