@triton.jit
def bwd_kernel_dq(Q, K, V, sm_scale, Out, DO, DQ, L, D, stride_qz,
    stride_qh, stride_qm, stride_qk, stride_kz, stride_kh, stride_kn,
    stride_kk, stride_vz, stride_vh, stride_vk, stride_vn, seqlen_q,
    seqlen_k, dropout_p, philox_seed, philox_offset_base, BLOCK_M: tl.
    constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr, CAUSAL:
    tl.constexpr, ENABLE_DROPOUT: tl.constexpr):
    start_m = tl.program_id(0) * BLOCK_M
    off_h = tl.program_id(1)
    off_z = tl.program_id(2)
    num_h = tl.num_programs(1)
    num_z = tl.num_programs(2)
    offs_m = start_m + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)
    q_offset = off_h * stride_qh + off_z * stride_qz
    Q_block_ptr = tl.make_block_ptr(base=Q + q_offset, shape=(seqlen_q,
        BLOCK_DMODEL), strides=(stride_qm, stride_qk), offsets=(start_m, 0),
        block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0))
    k_offset = off_h * stride_kh + off_z * stride_kz
    K_block_ptr = tl.make_block_ptr(base=K + k_offset, shape=(BLOCK_DMODEL,
        seqlen_k), strides=(stride_kk, stride_kn), offsets=(0, 0),
        block_shape=(BLOCK_DMODEL, BLOCK_N), order=(0, 1))
    v_offset = off_h * stride_vh + off_z * stride_vz
    V_block_ptr = tl.make_block_ptr(base=V + v_offset, shape=(BLOCK_DMODEL,
        seqlen_k), strides=(stride_vn, stride_vk), offsets=(0, 0),
        block_shape=(BLOCK_DMODEL, BLOCK_N), order=(0, 1))
    DO_block_ptr = tl.make_block_ptr(base=DO + q_offset, shape=(seqlen_q,
        BLOCK_DMODEL), strides=(stride_qm, stride_qk), offsets=(start_m, 0),
        block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0))
    off_zh = off_z * num_h + off_h * 1
    D_ptrs = D + off_zh * seqlen_q
    l_ptrs = L + off_zh * seqlen_q
    qk_scale = sm_scale * 1.44269504
    q = tl.load(Q_block_ptr)
    q = (q * qk_scale).to(Q_block_ptr.type.element_ty)
    do = tl.load(DO_block_ptr)
    Di = tl.load(D_ptrs + offs_m)
    l_i = tl.load(l_ptrs + offs_m)
    dq = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
    lo = 0
    hi = min(start_m + BLOCK_M, seqlen_k) if CAUSAL else seqlen_k
    batch_philox_offset = philox_offset_base + off_zh * seqlen_q * seqlen_k
    """
           K1   K2      (d)V      dO
    Q1    qk11 qk12     (d)v1     dO1
    Q2    qk21 qk22     (d)v2     dO2

    QK: (seqlen_q, seqlen_k)
    dO: (seqlen_q, hdim)
    dV: (seqlen_k, hdim)
    """
    for start_n in range(lo, hi, BLOCK_N):
        kt = tl.load(K_block_ptr)
        vt = tl.load(V_block_ptr)
        qk = dot(BLOCK_M, BLOCK_DMODEL, BLOCK_DMODEL, q, kt)
        if CAUSAL:
            qk = tl.where(offs_m[:, None] >= offs_n[None, :] + start_n, qk,
                float('-inf'))
        p = tl.math.exp2(qk - l_i[:, None])
        dp = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        dp += dot(BLOCK_M, BLOCK_DMODEL, BLOCK_DMODEL, do, vt)
        if ENABLE_DROPOUT:
            philox_offset = batch_philox_offset + start_m * seqlen_k + start_n
            keep = dropout_mask(philox_seed, philox_offset, dropout_p,
                BLOCK_M, BLOCK_N, seqlen_k)
            dp = tl.where(keep, dp / (1 - dropout_p), 0)
        ds = p * (dp - Di[:, None])
        if BLOCK_M == 1:
            dq += tl.view(kt, [BLOCK_DMODEL]) * ds.to(Q.type.element_ty)
        else:
            dq += tl.dot(ds.to(Q.type.element_ty), tl.trans(kt))
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
        V_block_ptr = tl.advance(V_block_ptr, (0, BLOCK_N))
    DQ_block_ptr = tl.make_block_ptr(base=DQ + q_offset, shape=(seqlen_q,
        BLOCK_DMODEL), strides=(stride_qm, stride_qk), offsets=(start_m, 0),
        block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0))
    tl.store(DQ_block_ptr, (dq * sm_scale).to(DQ_block_ptr.type.element_ty))
