@triton.jit
def mlstm_matmul_kernel(Q, K, V, F, I, M, B, H, 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_q = batch_id * NH * S * D + head_id * S * D
    batch_offset_f = batch_id * NH * S + head_id * S
    offset_q = tl.arange(0, SB) + sb_id * SB
    offset_k = tl.arange(0, SB) + sb_id * SB
    d_range = tl.arange(0, D)
    q_range = batch_offset_q + offset_q[:, None] * D + d_range[None, :]
    q_mask = (offset_q[:, None] < S) & (d_range[None, :] < D)
    q = tl.load(Q + q_range, q_mask)
    f = tl.load(F + batch_offset_f + offset_q, offset_q < S)
    f = tl.cumsum(tl.log(tl.sigmoid(f)))
    c_acc = tl.zeros((SB, D), dtype=tl.float32)
    b_acc = tl.zeros((SB,), dtype=tl.float32)
    m_acc = tl.zeros((SB,), dtype=tl.float32) - float('inf')
    for j in range(sb_id, -1, -1):
        kv_range = batch_offset_q + offset_k[:, None] * D + d_range[None, :]
        kv_mask = (offset_k[:, None] < S) & (d_range[None, :] < D)
        k = tl.load(K + kv_range, kv_mask) / tl.sqrt(tl.full((1,), D, dtype
            =tl.float32))
        v = tl.load(V + kv_range, kv_mask)
        f_next = tl.load(F + batch_offset_f + offset_k, offset_k < S)
        i = tl.load(I + batch_offset_f + offset_k, offset_k < 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_q[:, None] >= offset_k[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, :]
        m = tl.maximum(tl.max(d, 1), m_acc)
        d = tl.exp(d - m[:, None])
        c = matrix_mult(q, tl.trans(k), SB) * d
        b_acc = b_acc * tl.exp(m_acc - m) + tl.sum(c, 1)
        c = matrix_mult(c, v, SB)
        c_acc = c_acc * tl.exp(m_acc - m)[:, None] + c
        m_acc = m
        offset_k -= SB
    n = tl.maximum(tl.abs(b_acc), tl.exp(-m_acc)) + 1e-06
    h = c_acc / n[:, None]
    tl.store(H + q_range, h, q_mask)
    tl.store(B + batch_offset_f + offset_q, b_acc, offset_q < S)
    tl.store(M + batch_offset_f + offset_q, m_acc, offset_q < S)
