@triton.jit
def _backward_one_row(seq_prog_id, seq_length, qk_scale, M_range, N_range,
    D_range, D_mask, cm, DO_head_seq_ptr, stride_dom, stride_dod: tl.
    constexpr, DR_head_seq_ptr, stride_drm, A_head_seq_ptr, stride_am: tl.
    constexpr, Q_head_seq_ptr, stride_qm, stride_qd: tl.constexpr,
    K_head_seq_ptr, stride_kn, stride_kd: tl.constexpr, V_head_seq_ptr,
    stride_vn, stride_vd: tl.constexpr, DQ_head_seq_ptr, stride_dqm,
    stride_dqd: tl.constexpr, DK_head_seq_ptr, stride_dkn, stride_dkd: tl.
    constexpr, DV_head_seq_ptr, stride_dvn, stride_dvd: tl.constexpr,
    KV_Lock_ptr, KV_Count_ptr, logit_scale, BLOCK_D: tl.constexpr,
    NO_D_MASK: tl.constexpr, NO_M_MASK: tl.constexpr, ALLOW_TF32: tl.
    constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, acc_dtype: tl.
    constexpr=tl.float32, is_compiling: tl.constexpr=False, attend_current:
    tl.constexpr=False):
    block_start_offset = BLOCK_M * seq_prog_id
    M_blk_idxs = block_start_offset + M_range
    M_mask = M_blk_idxs < seq_length
    NO_M_MASK = block_start_offset + BLOCK_M - 1 < seq_length
    N_blk_idxs_start = 0
    N_blk_idxs = N_blk_idxs_start + N_range
    DO_blk_ptrs = DO_head_seq_ptr + (stride_dom * M_blk_idxs[:, None] + 
        stride_dod * D_range[None, :])
    K_blk_ptrs = K_head_seq_ptr + (stride_kn * N_blk_idxs[:, None] + 
        stride_kd * D_range[None, :])
    Q_blk_ptrs = Q_head_seq_ptr + (stride_qm * M_blk_idxs[:, None] + 
        stride_qd * D_range[None, :])
    V_blk_ptrs = V_head_seq_ptr + (stride_vn * N_blk_idxs[:, None] + 
        stride_vd * D_range[None, :])
    A_blk_ptrs = A_head_seq_ptr + stride_am * M_blk_idxs
    DQ_blk_ptrs = DQ_head_seq_ptr + (stride_dqm * M_blk_idxs[:, None] + 
        stride_dqd * D_range[None, :])
    DK_blk_ptrs = DK_head_seq_ptr + (stride_dkn * N_blk_idxs[:, None] + 
        stride_dkd * D_range[None, :])
    DV_blk_ptrs = DV_head_seq_ptr + (stride_dvn * N_blk_idxs[:, None] + 
        stride_dvd * D_range[None, :])
    DR_blk_ptrs = DR_head_seq_ptr + stride_drm * M_blk_idxs
    if NO_D_MASK:
        if NO_M_MASK:
            q = tl.load(Q_blk_ptrs)
            do = tl.load(DO_blk_ptrs)
            dr = tl.load(DR_blk_ptrs)
            neg_log_acc = tl.load(A_blk_ptrs, mask=M_mask)
        else:
            q = tl.load(Q_blk_ptrs, mask=M_mask[:, None])
            do = tl.load(DO_blk_ptrs, mask=M_mask[:, None])
            dr = tl.load(DR_blk_ptrs, mask=M_mask)
            neg_log_acc = tl.load(A_blk_ptrs, mask=M_mask)
    else:
        MD_mask = M_mask[:, None] & D_mask[None, :]
        q = tl.load(Q_blk_ptrs, mask=MD_mask)
        do = tl.load(DO_blk_ptrs, mask=MD_mask)
        dr = tl.load(DR_blk_ptrs, mask=M_mask)
        neg_log_acc = tl.load(A_blk_ptrs, mask=M_mask)
    neg_log_acc = neg_log_acc.to(dtype=acc_dtype)
    grad_prev_acc = tl.zeros((BLOCK_M,), dtype=acc_dtype)
    dq = tl.zeros((BLOCK_M, BLOCK_D), dtype=acc_dtype)
    fwd_cm = tl.trans(cm)
    iters = (block_start_offset + BLOCK_M) // BLOCK_N
    for i in range(iters):
        on_band = iters - i - 1 < BLOCK_M // BLOCK_N
        N_mask = N_blk_idxs < seq_length
        NO_N_MASK = N_blk_idxs_start + BLOCK_N - 1 < seq_length
        k, v = load_kv(K_blk_ptrs, V_blk_ptrs, N_mask=N_mask, NO_N_MASK=
            N_blk_idxs_start + BLOCK_N - 1 < seq_length, D_mask=D_mask,
            NO_D_MASK=NO_D_MASK)
        p, log_om_beta, neg_log_acc = compute_block(q, k, qk_scale,
            neg_log_acc, M_blk_idxs, N_blk_idxs, cm, on_band, ALLOW_TF32,
            attend_current=attend_current, backward=True, is_compiling=
            is_compiling)
        if not NO_M_MASK:
            neg_log_acc = tl.where(M_mask, neg_log_acc, 0.0)
        att_dA = p * (tl.dot(do, tl.trans(v), allow_tf32=ALLOW_TF32) - dr[:,
            None])
        cumul_att_dA = tl.dot(att_dA.to(cm.dtype), fwd_cm, allow_tf32=
            ALLOW_TF32) + grad_prev_acc[:, None]
        grad_prev_acc += tl.sum(att_dA, axis=1)
        beta = 1 - tl.exp2(log_om_beta)
        dqk = att_dA - beta * cumul_att_dA
        dq = tl.dot(dqk.to(k.dtype), k, acc=dq, allow_tf32=ALLOW_TF32)
        block_dk = tl.dot(tl.trans(dqk).to(q.dtype), q, allow_tf32=ALLOW_TF32
            ) * logit_scale
        block_dv = tl.dot(tl.trans(p), do.to(p.dtype), allow_tf32=ALLOW_TF32)
        locked_add(KV_Lock_ptr + i, KV_Count_ptr + i, DK_blk_ptrs, block_dk,
            DV_blk_ptrs, block_dv, N_mask, NO_N_MASK, D_mask, NO_D_MASK)
        N_blk_idxs += BLOCK_N
        N_blk_idxs_start += BLOCK_N
        K_blk_ptrs += BLOCK_N * stride_kn
        V_blk_ptrs += BLOCK_N * stride_vn
        DK_blk_ptrs += BLOCK_N * stride_dkn
        DV_blk_ptrs += BLOCK_N * stride_dvn
    dq = (logit_scale * dq).to(DQ_head_seq_ptr.type.element_ty)
    if NO_D_MASK:
        tl.store(DQ_blk_ptrs, dq, mask=M_mask[:, None])
    else:
        tl.store(DQ_blk_ptrs, dq, mask=M_mask[:, None] & D_mask[None, :])
