{{def_kernel("Q", "K", "V", "OUT", "D_OUT", "LSE", "DK", "DV")}}
    from flash_attn.cute.interface import _flash_attn_bwd

    q_transposed = Q.transpose(1, 2)
    k_transposed = K.transpose(1, 2)
    v_transposed = V.transpose(1, 2)
    out_transposed = OUT.transpose(1, 2)
    d_out_transposed = D_OUT.transpose(1, 2)

    dq_transposed, dk_transposed, dv_transposed = _flash_attn_bwd(
        q_transposed,
        k_transposed,
        v_transposed,
        out_transposed,
        d_out_transposed,
        LSE,
        softmax_scale={{SM_SCALE}},
    )

    dq = dq_transposed.transpose(1, 2)
    dk = dk_transposed.transpose(1, 2)
    dv = dv_transposed.transpose(1, 2)

    dq_out = {{get_output()}}
    {#  TODO: add out support to flash #}
    dq_out.copy_(dq)
    DK.copy_(dk)
    DV.copy_(dv)
