@triton.jit
def chunk_mlstm_fwd_kernel_h(q, k, v, C, n, m, m_total, i, f, h, norm,
    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,
    scale, H: tl.constexpr, T: tl.constexpr, K: tl.constexpr, V: tl.
    constexpr, BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, NT: tl
    .constexpr):
    i_v, i_t, i_bC = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    h_i = tl.arange(0, BT)
    m_s = h_i[:, None] >= h_i[None, :]
    b_h = tl.zeros([BT, BV], dtype=tl.float32)
    b_s = tl.zeros([BT, BT], dtype=tl.float32)
    b_norm = tl.zeros([BT, BV], dtype=tl.float32)
    for i_k in range(tl.cdiv(K, BK)):
        p_q = tl.make_block_ptr(q + i_bC * s_qk_h, (T, K), (s_qk_t, s_qk_d),
            (i_t * BT, i_k * BK), (BT, BK), (1, 0))
        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_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, BV), (1, 0),
            (i_k * BK, 0), (BK, BV), (0, 1))
        b_q = tl.load(p_q, boundary_check=(0, 1))
        b_k = tl.load(p_k, boundary_check=(0, 1))
        b_C = tl.load(p_C, boundary_check=(0, 1))
        b_n = tl.load(p_n, boundary_check=(0,))
        b_h += tl.dot(b_q, b_C, allow_tf32=False)
        b_s += tl.dot(b_q, b_k, allow_tf32=False)
        b_n2 = tl.dot(b_q, b_n, allow_tf32=False)
        b_norm += b_n2
    p_f = f + i_bC * T + i_t * BT + tl.arange(0, BT)
    b_f = tl.load(p_f)
    p_i = i + i_bC * T + i_t * BT + tl.arange(0, BT)
    b_i = tl.load(p_i)
    b_m = tl.load(m + i_bC * (NT + 1) + i_t)
    b_logD = b_i[None, :] + b_f[:, None] - b_f[None, :]
    b_logD = tl.where(m_s, b_logD, -float('inf'))
    b_mlogD = tl.max(b_logD, axis=1)
    b_m_total = tl.maximum(b_f + b_m, b_mlogD)
    p_m_total = tl.make_block_ptr(m_total + T * i_bC, (T,), (1,), (i_t * BT
        ,), (BT,), (0,))
    tl.store(p_m_total, b_m_total.to(p_m_total.dtype.element_ty),
        boundary_check=(0,))
    b_D = tl.math.exp2(b_logD - b_m_total[:, None])
    b_h = b_h * tl.math.exp2(b_f + b_m - b_m_total)[:, None] * scale
    b_s = b_s * b_D * scale
    b_norm = b_norm * tl.math.exp2(b_f + b_m - b_m_total)[:, None] * scale
    b_s = tl.where(m_s, b_s, 0)
    b_norm += tl.sum(b_s, axis=1)[:, None]
    b_norm = tl.abs(b_norm)
    b_norm = tl.maximum(b_norm, tl.math.exp2(-b_m_total)[:, None])
    tl.store(norm + i_bC * T + i_t * BT + tl.arange(0, BT), tl.max(b_norm,
        axis=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))
    b_v = tl.load(p_v, boundary_check=(0, 1))
    b_h = (b_h + tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False)) / b_norm
    p_h = tl.make_block_ptr(h + i_bC * s_vh_h, (T, V), (s_vh_t, s_vh_d), (
        i_t * BT, i_v * BV), (BT, BV), (1, 0))
    tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1))
