@triton.jit
def mlstm_matmul_kernel_backward(dH, dB, Q, K, V, dQ, dK, dV, F, dF, I, dI,
    M, B, NH: tl.constexpr, S: tl.constexpr, D: tl.constexpr, SB: tl.constexpr
    ):
    bh_id = tl.program_id(0)
    sb_id = tl.program_id(1)
    batch_id = bh_id // NH
    head_id = bh_id % NH
    batch_offset_dh = batch_id * NH * S * D + head_id * S * D
    batch_offset_f = batch_id * NH * S + head_id * S
    offset_dh = tl.arange(0, SB) + sb_id * SB
    offset_vk = tl.arange(0, SB) + sb_id * SB
    d_range = tl.arange(0, D)
    dh_range = batch_offset_dh + offset_dh[:, None] * D + d_range[None, :]
    dh_mask = (offset_dh[:, None] < S) & (d_range[None, :] < D)
    dh = tl.load(dH + dh_range, dh_mask)
    m = tl.load(M + batch_offset_f + offset_dh, offset_dh < S)
    b = tl.load(B + batch_offset_f + offset_dh, offset_dh < S)
    f = tl.load(F + batch_offset_f + offset_dh, offset_dh < S)
    db = tl.load(dB + batch_offset_f + offset_dh, offset_dh < S)
    q = tl.load(Q + dh_range, dh_mask)
    scale = tl.sqrt(tl.full((1,), D, dtype=tl.float32))
    n = tl.maximum(tl.abs(b), tl.exp(-m)) + 1e-06
    f = tl.cumsum(tl.log(tl.sigmoid(f)))
    f_low = f
    df_acc = tl.zeros((SB,), dtype=tl.float32)
    dq_acc = tl.zeros((SB, D), dtype=tl.float32)
    for j in range(sb_id, -1, -1):
        vk_range = batch_offset_dh + offset_vk[:, None] * D + d_range[None, :]
        vk_mask = (offset_vk[:, None] < S) & (d_range[None, :] < D)
        f_next = tl.load(F + batch_offset_f + offset_vk, offset_vk < S)
        i = tl.load(I + batch_offset_f + offset_vk, offset_vk < S)
        f_next = tl.log(tl.sigmoid(f_next))
        if j == sb_id:
            f_next = tl.cumsum(f_next)
            d = f[:, None] - f_next[None, :] + i[None, :]
            mask = offset_dh[:, None] >= offset_vk[None, :]
            d = tl.where(mask, d, -float('inf'))
        else:
            f += tl.sum(f_next)
            f_next = tl.cumsum(f_next)
            d = f[:, None] - f_next[None, :] + i[None, :]
        d = tl.exp(d - m[:, None])
        v = tl.load(V + vk_range, vk_mask)
        dc_tilde = matrix_mult(dh, tl.trans(v), SB) * (1 / n)[:, None] + db[
            :, None]
        k = tl.load(K + vk_range, vk_mask) / scale
        dq_acc += matrix_mult(dc_tilde * d, k, SB)
        c_tilde = matrix_mult(q, tl.trans(k), SB) * d
        df_acc += tl.sum(c_tilde * dc_tilde, 1)
        offset_vk -= SB
    tl.store(dQ + dh_range, dq_acc, dh_mask)
    offset_q = tl.arange(0, SB) + sb_id * SB
    f = tl.zeros((1,), dtype=tl.float32)
    v = tl.load(V + dh_range, dh_mask)
    k = tl.load(K + dh_range, dh_mask)
    i = tl.load(I + batch_offset_f + offset_dh, offset_dh < S)
    dk_acc = tl.zeros((SB, D), dtype=tl.float32)
    dv_acc = tl.zeros((SB, D), dtype=tl.float32)
    di_acc = tl.zeros((SB,), dtype=tl.float32)
    for j in range(sb_id, tl.cdiv(S, SB)):
        q_range = batch_offset_dh + offset_q[:, None] * D + d_range[None, :]
        q_mask = (offset_q[:, None] < S) & (d_range[None, :] < D)
        f_next = tl.load(F + batch_offset_f + offset_q, offset_q < S)
        f_next = tl.log(tl.sigmoid(f_next))
        f_next_sum = tl.sum(f_next)
        f_next = f + tl.cumsum(f_next)
        d = f_next[None, :] - f_low[:, None] + i[:, None]
        f += f_next_sum
        if j == sb_id:
            mask = offset_dh[:, None] <= offset_q[None, :]
            d = tl.where(mask, d, -float('inf'))
        dh = tl.load(dH + q_range, q_mask)
        m = tl.load(M + batch_offset_f + offset_q, offset_q < S)
        b = tl.load(B + batch_offset_f + offset_q, offset_q < S)
        db = tl.load(dB + batch_offset_f + offset_q, offset_q < S)
        d = tl.exp(d - m[None, :])
        n = tl.maximum(tl.abs(b), tl.exp(-m)) + 1e-06
        dc_tilde_T = matrix_mult(v, tl.trans(dh), SB) * (1 / n)[None, :] + db[
            None, :]
        q = tl.load(Q + q_range, q_mask) / scale
        dk_acc += matrix_mult(dc_tilde_T * d, q, SB)
        c_tilde_T = matrix_mult(k, tl.trans(q), SB) * d
        dv_acc += matrix_mult(c_tilde_T / n[None, :], dh, SB)
        di_acc += tl.sum(c_tilde_T * dc_tilde_T, 1)
        offset_q += SB
    tl.store(dK + dh_range, dk_acc, dh_mask)
    tl.store(dV + dh_range, dv_acc, dh_mask)
    tl.store(dI + batch_offset_f + offset_dh, di_acc, offset_dh < S)
    tl.store(dF + batch_offset_f + offset_dh + 1, di_acc - df_acc, 
        offset_dh + 1 < S)
