@triton.jit
def fused_recurrent_fwd_kernel(q, k, k_l2, q_reflected, k_reflected,
    dk_l2_partial, T, D: tl.constexpr, BK: tl.constexpr):
    i_b, i_h = tl.program_id(0), tl.program_id(1)
    p_q = q + i_b * T * D + i_h * BK + tl.arange(0, BK)
    p_k = k + i_b * T * D + i_h * BK + tl.arange(0, BK)
    p_k_l2 = k_l2 + i_b * T * D + i_h * BK + tl.arange(0, BK)
    p_q_reflected = q_reflected + i_b * T * D + i_h * BK + tl.arange(0, BK)
    p_k_reflected = k_reflected + i_b * T * D + i_h * BK + tl.arange(0, BK)
    p_dk_l2_partial = dk_l2_partial + i_b * T * D + i_h * BK + tl.arange(0, BK)
    h = tl.zeros([BK, BK], dtype=tl.float32) + (tl.arange(0, BK)[:, None] ==
        tl.arange(0, BK)[None, :])
    for _ in range(0, T):
        b_k_l2 = tl.load(p_k_l2).to(tl.float32)
        b_q = tl.load(p_q).to(tl.float32)
        b_k = tl.load(p_k).to(tl.float32)
        tmp = tl.sum(h * b_k_l2[None, :], axis=1)
        h -= 2 * b_k_l2[None, :] * tmp[:, None]
        b_q = tl.sum(h * b_q[None, :], axis=1)
        b_k = tl.sum(h * b_k[None, :], axis=1)
        tl.store(p_q_reflected, b_q.to(p_q_reflected.dtype.element_ty))
        tl.store(p_k_reflected, b_k.to(p_k_reflected.dtype.element_ty))
        tl.store(p_dk_l2_partial, tmp.to(p_dk_l2_partial.dtype.element_ty))
        p_q += D
        p_k += D
        p_k_l2 += D
        p_q_reflected += D
        p_k_reflected += D
        p_dk_l2_partial += D
