@triton.heuristics({'USE_INITIAL_STATE': lambda args: args['h0'] is not
    None, 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None,
    'USE_OFFSETS': lambda args: args['offsets'] is not None})
@triton.jit
def fused_recurrent_retention_bwd_kernel(q, k, v, h0, do, dq, dk, dv, dh0,
    dht, offsets, scale, B: tl.constexpr, H: tl.constexpr, T: tl.constexpr,
    K: tl.constexpr, V: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr,
    REVERSE: tl.constexpr, USE_INITIAL_STATE: tl.constexpr,
    USE_FINAL_STATE_GRADIENT: tl.constexpr, USE_OFFSETS: tl.constexpr,
    HEAD_FIRST: tl.constexpr):
    i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    i_n, i_h = i_nh // H, i_nh % H
    if USE_OFFSETS:
        bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets +
            i_n + 1).to(tl.int64)
        all = T
        T = eos - bos
    else:
        bos, eos = i_n * T, i_n * T + T
        all = B * T
    b_b = 1 - tl.math.exp2(-5 - i_h * 1.0)
    if HEAD_FIRST:
        p_k = k + i_nh * T * K + ((T - 1) * K if REVERSE else 0
            ) + i_k * BK + tl.arange(0, BK)
        p_v = v + i_nh * T * V + ((T - 1) * V if REVERSE else 0
            ) + i_v * BV + tl.arange(0, BV)
        p_do = do + i_nh * T * V + ((T - 1) * V if REVERSE else 0
            ) + i_v * BV + tl.arange(0, BV)
        p_dq = dq + (i_v * B * H + i_nh) * T * K + ((T - 1) * K if REVERSE else
            0) + i_k * BK + tl.arange(0, BK)
    else:
        p_k = k + (bos + (T - 1 if REVERSE else 0)
            ) * H * K + i_h * K + i_k * BK + tl.arange(0, BK)
        p_v = v + (bos + (T - 1 if REVERSE else 0)
            ) * H * V + i_h * V + i_v * BV + tl.arange(0, BV)
        p_do = do + (bos + (T - 1 if REVERSE else 0)
            ) * H * V + i_h * V + i_v * BV + tl.arange(0, BV)
        p_dq = dq + (i_v * all + bos + (T - 1 if REVERSE else 0)
            ) * H * K + i_h * K + i_k * BK + tl.arange(0, BK)
    mask_k = i_k * BK + tl.arange(0, BK) < K
    mask_v = i_v * BV + tl.arange(0, BV) < V
    mask_h = mask_k[:, None] & mask_v[None, :]
    b_h = tl.zeros([BK, BV], dtype=tl.float32)
    if USE_INITIAL_STATE:
        p_h0 = h0 + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]
            ) * V + (i_v * BV + tl.arange(0, BV)[None, :])
        b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
    for _ in range(0, T):
        b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
        b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)
        b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32)
        b_h = b_b * b_h + b_k[:, None] * b_v[None, :]
        b_dq = tl.sum(b_h * b_do[None, :], axis=1) * scale
        tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=mask_k)
        p_k += (-1 if REVERSE else 1) * (1 if HEAD_FIRST else H) * K
        p_v += (-1 if REVERSE else 1) * (1 if HEAD_FIRST else H) * V
        p_do += (-1 if REVERSE else 1) * (1 if HEAD_FIRST else H) * V
        p_dq += (-1 if REVERSE else 1) * (1 if HEAD_FIRST else H) * K
    tl.debug_barrier()
    if HEAD_FIRST:
        p_q = q + i_nh * T * K + ((T - 1) * K if not REVERSE else 0
            ) + i_k * BK + tl.arange(0, BK)
        p_k = k + i_nh * T * K + ((T - 1) * K if not REVERSE else 0
            ) + i_k * BK + tl.arange(0, BK)
        p_v = v + i_nh * T * V + ((T - 1) * V if not REVERSE else 0
            ) + i_v * BV + tl.arange(0, BV)
        p_do = do + i_nh * T * V + ((T - 1) * V if not REVERSE else 0
            ) + i_v * BV + tl.arange(0, BV)
        p_dk = dk + (i_v * B * H + i_nh) * T * K + ((T - 1) * K if not
            REVERSE else 0) + i_k * BK + tl.arange(0, BK)
        p_dv = dv + (i_k * B * H + i_nh) * T * V + ((T - 1) * V if not
            REVERSE else 0) + i_v * BV + tl.arange(0, BV)
    else:
        p_q = q + (bos + (T - 1 if not REVERSE else 0)
            ) * H * K + i_h * K + i_k * BK + tl.arange(0, BK)
        p_k = k + (bos + (T - 1 if not REVERSE else 0)
            ) * H * K + i_h * K + i_k * BK + tl.arange(0, BK)
        p_v = v + (bos + (T - 1 if not REVERSE else 0)
            ) * H * V + i_h * V + i_v * BV + tl.arange(0, BV)
        p_do = do + (bos + (T - 1 if not REVERSE else 0)
            ) * H * V + i_h * V + i_v * BV + tl.arange(0, BV)
        p_dk = dk + (i_v * all + bos + (T - 1 if not REVERSE else 0)
            ) * H * K + i_h * K + i_k * BK + tl.arange(0, BK)
        p_dv = dv + (i_k * all + bos + (T - 1 if not REVERSE else 0)
            ) * H * V + i_h * V + i_v * BV + tl.arange(0, BV)
    b_dh = tl.zeros([BK, BV], dtype=tl.float32)
    if USE_FINAL_STATE_GRADIENT:
        p_ht = dht + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]
            ) * V + (i_v * BV + tl.arange(0, BV)[None, :])
        b_dh += tl.load(p_ht, mask=mask_h, other=0).to(tl.float32)
    for _ in range(T):
        b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale
        b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
        b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)
        b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32)
        b_dh += b_q[:, None] * b_do[None, :]
        b_dk = tl.sum(b_dh * b_v[None, :], axis=1)
        b_dv = tl.sum(b_dh * b_k[:, None], axis=0)
        b_dh *= b_b
        tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=mask_k)
        tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=mask_v)
        p_q += (1 if REVERSE else -1) * (1 if HEAD_FIRST else H) * K
        p_k += (1 if REVERSE else -1) * (1 if HEAD_FIRST else H) * K
        p_v += (1 if REVERSE else -1) * (1 if HEAD_FIRST else H) * V
        p_do += (1 if REVERSE else -1) * (1 if HEAD_FIRST else H) * V
        p_dk += (1 if REVERSE else -1) * (1 if HEAD_FIRST else H) * K
        p_dv += (1 if REVERSE else -1) * (1 if HEAD_FIRST else H) * V
    if USE_INITIAL_STATE:
        p_dh0 = dh0 + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]
            ) * V + (i_v * BV + tl.arange(0, BV)[None, :])
        tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=mask_h)
