@staticmethod
def backward(ctx, dout, *args):
    # 步骤 1: 恢复前向保存的张量
    q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
    # 注意：q, k, v 是 padding 后的版本

    # 步骤 2: 预分配梯度张量
    dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)

    # 步骤 3: dout padding
    head_size_og = dout.size(3)
    dout_padded = dout
    if head_size_og % 8 != 0:
        dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])

    # 步骤 4: 调用 CUDA 反向
    _wrapped_flash_attn_backward(
        dout_padded, q, k, v, out, softmax_lse,
        dq, dk, dv,                    # 原地写入
        ctx.dropout_p, ctx.softmax_scale,
        ctx.causal, ctx.window_size[0], ctx.window_size[1],
        ctx.softcap, ctx.alibi_slopes, ctx.deterministic,
        rng_state=rng_state,
    )

    # 步骤 5: 裁剪 padding
    dq = dq[..., :head_size_og]
    dk = dk[..., :head_size_og]
    dv = dv[..., :head_size_og]

    return dq, dk, dv, None, None, None, None, None, None, None, None, None

# flash_attn/flash_attn_interface.py:241-290
@_torch_custom_op_wrapper(
    "flash_attn::_flash_attn_backward",
    mutates_args=("dq", "dk", "dv"),  # 声明原地修改
    device_types="cuda"
)
def _flash_attn_backward(
    dout, q, k, v, out, softmax_lse,
    dq, dk, dv,                       # 这些会被原地修改
    dropout_p, softmax_scale, causal,
    window_size_left, window_size_right, softcap,
    alibi_slopes, deterministic,
    rng_state=None,
):
    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
    dq, dk, dv, softmax_d = flash_attn_gpu.bwd(
        dout, q, k, v, out, softmax_lse,
        dq, dk, dv,
        alibi_slopes,
        dropout_p, softmax_scale, causal,
        window_size_left, window_size_right, softcap,
        deterministic, None, rng_state,
    )
    return softmax_d  # 返回行归一化因子（调试用）

std::vector<at::Tensor> mha_bwd(
    const at::Tensor &dout,   // (B, S_q, H, D)
    const at::Tensor &q,      // (B, S_q, H, D)
    const at::Tensor &k,      // (B, S_k, Hk, D)
    const at::Tensor &v,      // (B, S_k, Hk, D)
    const at::Tensor &out,    // (B, S_q, H, D)
    const at::Tensor &softmax_lse,  // (B, H, S_q)
    c10::optional<at::Tensor> &dq_,  // 预分配梯度（可选）
    c10::optional<at::Tensor> &dk_,
    c10::optional<at::Tensor> &dv_,
    c10::optional<at::Tensor> &alibi_slopes_,
    const float p_dropout,
    const float softmax_scale,
    const bool is_causal,
    int window_size_left, int window_size_right,
    const float softcap,
    const bool deterministic,
    c10::optional<at::Generator> gen_,
    c10::optional<at::Tensor> &rng_state,
) {

// csrc/flash_attn/flash_api.cpp:849-904

// 梯度张量分配（如果 Python 层没有传入）
at::Tensor dq, dk, dv;
if (dq_.has_value()) { dq = dq_.value(); }
else { dq = torch::empty_like(q); }

// GQA 时 dk, dv 需要扩展维度
if (num_heads_k != num_heads) {
    dk = torch::empty({batch_size, seqlen_k, num_heads, head_size}, ...);
    dv = torch::empty({batch_size, seqlen_k, num_heads, head_size}, ...);
    // 后面会归约: (B, S, H, D) → (B, S, Hk, D)
}

// softmax_d: 反向中间结果
auto softmax_d = torch::empty(
    {batch_size, num_heads, seqlen_q_rounded},
    opts.dtype(at::kFloat)
);

// dq_accum: dQ 累加缓冲（FP32）
auto dq_accum = torch::empty(
    {batch_size, seqlen_q_rounded, num_heads, head_size_rounded},
    opts.dtype(at::kFloat)
);

void set_params_dgrad(Flash_bwd_params &params, ...) {
    // 复用前向参数设置
    set_params_fprop(params, batch_size, seqlen_q, seqlen_k,
                     seqlen_q_rounded, seqlen_k_rounded,
                     num_heads, num_heads_k, head_size, head_size_rounded,
                     q, k, v, out, ...);

    // 反向特有指针
    params.do_ptr = dout.data_ptr();       // 上游梯度
    params.dq_ptr = dq.data_ptr();         // Q 梯度
    params.dk_ptr = dk.data_ptr();         // K 梯度
    params.dv_ptr = dv.data_ptr();         // V 梯度

    // 步长设置
    params.do_row_stride = dout.stride(-3);
    params.dq_row_stride = dq.stride(-3);
    params.dk_row_stride = dk.stride(-3);
    params.dv_row_stride = dv.stride(-3);

    // 累加缓冲
    params.dq_accum_ptr = dq_accum.data_ptr();
    params.dsoftmax_sum = softmax_d.data_ptr();  // Σ(dout · out) 行和

    // 确定性标志
    params.deterministic = deterministic;
}

// csrc/flash_attn/flash_api.cpp:936-951
if (rng_state.has_value()) {
    // 使用前向保存的 RNG 状态
    params.rng_state = reinterpret_cast<uint64_t*>(rng_state.value().data_ptr());
}

// csrc/flash_attn/flash_api.cpp:956
if (seqlen_q > 0) {
    auto stream = at::cuda::getCurrentCUDAStream().stream();
    run_mha_bwd(params, stream);
}

// hopper/flash_api.cpp:1246-1257
void run_mha_bwd(Flash_bwd_params &params, cudaStream_t stream) {
    ARCH_SWITCH(params.arch, Arch, [&] {
        SOFTCAP_SWITCH(params.softcap > 0.f, Has_softcap, [&] {
            run_mha_bwd_constexpr<Arch, Has_softcap>(params, stream);
        });
    });
}

// hopper/flash_api.cpp:1206-1244
template<int Arch, bool Has_softcap>
void run_mha_bwd_constexpr(Flash_bwd_params &params, cudaStream_t stream) {
    if (!params.is_bf16) {
        if (params.d_rounded == 64)
            return run_mha_bwd_<Arch, cutlass::half_t, 64, Has_softcap>(...);
        if (params.d_rounded == 96)
            return run_mha_bwd_<Arch, cutlass::half_t, 96, Has_softcap>(...);
        if (params.d_rounded == 128)
            return run_mha_bwd_<Arch, cutlass::half_t, 128, Has_softcap>(...);
        // ... 192, 256
    } else {
        // BF16 路径
    }
}

CUTLASS_DEVICE void operator()(Params const& params, char* smem_buf) {
    if (warp_group_idx == 0) {
        // Producer 0: 加载 Q, dO, LSE
        collective_mainloop.load_Q_dO(params, ...);
    }
    else if (warp_group_idx == NumMmaWarpGroups) {
        // Producer 1: 通过 TMA 存储 dQ
        collective_mainloop.store_dQ(params, ...);
    }
    else {
        // Consumer: 计算 dQ, dK, dV
        collective_mainloop.mma(params, ...);  // 主循环
        collective_epilogue.store(params, ...); // 写出 dK, dV
    }
}

// 每个 bwd_step 后通过原子操作写回 dQ
atomicAdd(dQ_global[m_block], dQ_local);

// Producer 1 通过 TMA 将 dQ 写到 dq_accum 缓冲区
// 后续通过 reduction kernel 归约

// csrc/flash_attn/flash_api.cpp:965-968
if (num_heads_k != num_heads) {
    // dk_expanded: (B, S, H, D) → reshape → (B, S, Hk, H/Hk, D)
    at::sum_out(dk,
        at::reshape(dk_expanded, {batch_size, seqlen_k, num_heads_k,
                                   num_heads / num_heads_k, head_size}),
        {3}  // 在 group 维度求和
    );
    at::sum_out(dv,
        at::reshape(dv_expanded, {batch_size, seqlen_k, num_heads_k,
                                   num_heads / num_heads_k, head_size}),
        {3}
    );
}

// 确定性模式下的梯度归约在 epilogue_bwd.hpp 中
// 通过控制 thread block 的执行顺序避免不确定性