// hopper/flash_bwd_kernel_sm90.h:211-240
if (warp_group_idx == 0) {  // Producer
    cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
    int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0);

    if (warp_idx_in_warpgroup == 0) {
        // Warp 0: TMA 加载 Q, dO, LSE, dPsum → 共享内存
        for (auto work_tile_info = scheduler.get_initial_work<true>(...); ...) {
            mainloop.load(params, pipeline_q, pipeline_do, smem_pipe_write,
                          smem_pipe_write_do, shared_storage, scheduler_prefetch, block_coord);
        }
        mainloop.load_tail(pipeline_q, pipeline_do, smem_pipe_write, smem_pipe_write_do);

    } else if (warp_idx_in_warpgroup == 1) {
        // Warp 1: TMA 存储 dQ ← 共享内存 → HBM
        for (auto work_tile_info = scheduler.get_initial_work<false>(...); ...) {
            mainloop.store_dq(params, shared_storage, block_coord);
        }
    }
}

// hopper/flash_bwd_kernel_sm90.h:82-83
alignas(16) typename CollectiveMainloop::MainloopPipeline::SharedStorage pipeline_q;
alignas(16) typename CollectiveMainloop::MainloopPipeline_dO::SharedStorage pipeline_do;

// hopper/flash_bwd_kernel_sm90.h:178-179
pipeline_params.transaction_bytes =
    CollectiveMainloop::TmaTransactionBytesQ + CollectiveMainloop::TmaTransactionBytesLSE;

// mainloop_bwd_sm90_tma_gmma_ws.hpp::load()（简化）
// 在 Q-inner 循环开始前，一次性加载 K, V
if (cute::elect_one_sync()) {
    shared_storage.pipelines.barrier_KV.arrive_and_expect_tx(
        TmaTransactionBytesK + TmaTransactionBytesV);
    copy(params.tma_load_K.with(...), tKgK, tKsK);   // TMA Load K
    copy(params.tma_load_V.with(...), tVgV, tVsV);   // TMA Load V
}

// hopper/flash_bwd_kernel_sm90.h:81
alignas(16) cutlass::arch::ClusterTransactionBarrier barrier_KV;

// mainloop_bwd_sm90_tma_gmma_ws.hpp::mma()（简化结构）
CUTLASS_DEVICE bool mma(
    Params const& params,
    MainloopPipeline pipeline_q,
    MainloopPipeline_dO pipeline_do,
    PipelineState& smem_pipe_read,
    PipelineState_dO& smem_pipe_read_do,
    FrgTensordKV& tdKrdK,      // dK 累加器（寄存器，跨 Q-block 累加）
    FrgTensordKV& tdVrdV,      // dV 累加器（寄存器，跨 Q-block 累加）
    int thread_idx, int& work_idx,
    cute::tuple<...> block_coord,
    SharedStorage& shared_storage) {

    // 1. 初始化
    clear(tdKrdK);  // dK = 0
    clear(tdVrdV);  // dV = 0

    // 2. 等待 K/V 加载完成
    shared_storage.pipelines.barrier_KV.wait(...);

    // 3. 计算有效 Q-block 范围
    auto [m_block_min, m_block_max] = get_m_block_min_max(...);

    // 4. Q-inner 循环
    for (int m_block = m_block_min; m_block < m_block_max; ++m_block) {
        bwd_step(m_block, mask_fn);
    }

    // 5. dK 乘以 softmax_scale
    for (int i = 0; i < size(tdKrdK); ++i) {
        tdKrdK(i) *= params.softmax_scale;
    }

    return true;
}

// mainloop_bwd_sm90_tma_gmma_ws.hpp:~1003-1033（简化）
// 阶段 1: 需要 Causal Masking 的 Q-block
if constexpr ((Is_causal || Is_local) && SeparateMaskingIterations) {
    int const m_block_masking_max = ((n_block + 1) * kBlockN - 1 + seqlen_q - seqlen_k
                                      - params.window_size_right) / kBlockM + 1;
    for (; m_block < std::min(m_block_max, m_block_masking_max); ++m_block) {
        bwd_step(m_block, causal_mask_fn);
    }
}

// 阶段 2: 完全有效的 Q-block
for (; m_block < m_block_max_before_local_mask; ++m_block) {
    bwd_step(m_block, no_mask_fn);
}

// 阶段 3: 需要 Local Masking 的 Q-block
if constexpr (Is_local && SeparateMaskingIterations) {
    for (; m_block < m_block_max; ++m_block) {
        bwd_step(m_block, local_mask_fn);
    }
}

// mainloop_bwd_sm90_tma_gmma_ws.hpp:~835-851（简化）
// GEMM 1: S = Q @ K^T
Tensor tSrS = partition_fragment_C(tiled_mma_SdP, ...);
consumer_wait(pipeline_q, smem_pipe_read);          // 等待 Q 就绪
flash::gemm</*zero_init=*/true, /*wg_wait=*/-1, /*SwapAB=*/SdP_swapAB>(
    tiled_mma_SdP, tSrQ(_, _, _, smem_pipe_read.index()), tSrK, tSrS);

// GEMM 2: dP = dO @ V^T（与 GEMM 1 异步重叠）
Tensor tdPrdP = partition_fragment_C(tiled_mma_SdP, ...);
consumer_wait(pipeline_do, smem_pipe_read_do_cur);  // 等待 dO 就绪
flash::gemm</*zero_init=*/true, /*wg_wait=*/-1, /*SwapAB=*/SdP_swapAB>(
    tiled_mma_SdP, tdPrdO(_, _, _, smem_pipe_read_do_cur.index()), tdPrV, tdPrdP);

// mainloop_bwd_sm90_tma_gmma_ws.hpp:~852-896（简化）
warpgroup_wait<1>();  // 等待 GEMM 1 (S) 完成，GEMM 2 (dP) 可仍在后台

// 应用 softcap（可选）
if constexpr (Has_softcap) {
    flash::apply_softcap(tSrS, params.softcap_val);
    auto dtanh = flash::calculate_dtanh(scores);     // 预计算 tanh 导数
}

// 应用 mask
mask_fn(tSrS, m_block);

// 恢复 P = exp2(S * scale * log2e - LSE)
for (int mi = 0; mi < size<0>(scores); ++mi) {
    float lse_scaled = tLSErLSE(mi);
    for (int ni = 0; ni < size<1>(scores); ++ni) {
        scores(mi, ni) = exp2f(scores(mi, ni) * params.softmax_scale_log2 - lse_scaled);
    }
}

warpgroup_wait<0>();  // 等待 GEMM 2 (dP) 完成

// 计算 dS = P * (dP - dPsum)
for (int mi = 0; mi < size<0>(dS); ++mi) {
    float dP_sum_cur = tLSErdPsum(mi);
    for (int ni = 0; ni < size<1>(dS); ++ni) {
        dS(mi, ni) = scores(mi, ni) * (dS(mi, ni) - dP_sum_cur);
        if constexpr (Has_softcap) { dS(mi, ni) *= dtanh(mi, ni); }
    }
}

// mainloop_bwd_sm90_tma_gmma_ws.hpp:~925-935（简化）
if constexpr (Mma_dKV_is_RS) {
    // RS 模式：P 直接从寄存器
    Tensor tdVrP = make_tensor(rP.data(), convert_layout_acc_Aregs<TiledMmadKV>(...));
    flash::gemm</*zero_init=*/false, /*wg_wait=*/-1>(
        tiled_mma_dKV, tdVrP, tdVrdO(_, _, _, smem_pipe_read_do_cur.index()), tdVrdV);
} else {
    // 标准模式：P 从共享内存
    flash::gemm</*zero_init=*/false, /*wg_wait=*/-1, /*SwapAB=*/dKV_swapAB>(
        tiled_mma_dKV, tdVrP_cur, tdVrdO(_, _, _, smem_pipe_read_do_cur.index()), tdVrdV);
}

// mainloop_bwd_sm90_tma_gmma_ws.hpp:~938-950（简化）
// GEMM 4: dQ = dS @ K（每个 Q-block 重新计算）
flash::gemm</*zero_init=*/true, /*wg_wait=*/1, /*SwapAB=*/dQ_swapAB>(
    tiled_mma_dQ, tdQrdS_cur, tdQrK, tdQrdQ);

pipeline_do.consumer_release(smem_pipe_read_do_cur);  // 释放 dO smem

// GEMM 5: dK += dS^T @ Q（累加）
if constexpr (Mma_dKV_is_RS) {
    flash::gemm</*zero_init=*/false, /*wg_wait=*/1>(
        tiled_mma_dKV, tdKrdS, tdKrQ(_, _, _, smem_pipe_read.index()), tdKrdK);
} else {
    flash::gemm</*zero_init=*/false, /*wg_wait=*/1, /*SwapAB=*/dKV_swapAB>(
        tiled_mma_dKV, tdKrdS_cur, tdKrQ(_, _, _, smem_pipe_read.index()), tdKrdK);
}

// mainloop_bwd_sm90_tma_gmma_ws.hpp:~951-958（简化）
if constexpr (dQacc_use_TMA) {
    int warp_group_idx = flash::canonical_warp_group_idx_nosync() - 1;

    // 等待 smem_dq 空闲
    NamedBarrier::sync(NumThreadsPerWarpGroup + NumThreadsPerWarp,
        static_cast<uint32_t>(BwdNamedBarriers::dQEmptyWG1) + warp_group_idx);

    // R2S: dQ 从寄存器写入共享内存
    cute::copy(r2s_tiled_copy_dQaccum, taccdQrdQ, tdQsdQaccum);
    fence_view_async_shared();

    // 通知 Warp 1 可以执行 TMA Store
    NamedBarrier::arrive(NumThreadsPerWarpGroup + NumThreadsPerWarp,
        static_cast<uint32_t>(BwdNamedBarriers::dQFullWG1) + warp_group_idx);
}

// mainloop_bwd_sm90_tma_gmma_ws.hpp::store_dq()（简化）
CUTLASS_DEVICE void store_dq(Params const& params, SharedStorage& shared_storage, ...) {
    for (int m_block = m_block_min; m_block < m_block_max; ++m_block) {
        // 等待 Consumer 写入 dQ
        for (int wg_idx = 0; wg_idx < NumMmaWarpGroups; ++wg_idx) {
            NamedBarrier::sync(NumThreadsPerWarpGroup + NumThreadsPerWarp,
                static_cast<uint32_t>(BwdNamedBarriers::dQFullWG1) + wg_idx);
        }

        // TMA Store: smem_dq → dQaccum (HBM)
        if (cute::elect_one_sync()) {
            cute::copy(params.tma_store_dQaccum, tdQsdQaccum, tdQgdQaccum(_, m_block));
            tma_store_arrive();
        }
        tma_store_wait<0>();

        // 通知 Consumer smem_dq 已释放
        for (int wg_idx = 0; wg_idx < NumMmaWarpGroups; ++wg_idx) {
            NamedBarrier::arrive(NumThreadsPerWarpGroup + NumThreadsPerWarp,
                static_cast<uint32_t>(BwdNamedBarriers::dQEmptyWG1) + wg_idx);
        }
    }
}

// mainloop_bwd_sm90_tma_gmma_ws.hpp:~960-966
if constexpr (!dQacc_use_TMA) {
    Tensor tdQrdQ_atomic = recast<float4>(r2s_thr_copy_dQaccum.retile_S(tdQrdQ));
    Tensor tdQgdQaccum_atomic = recast<float4>(tdQgdQaccum(_, _, _, m_block));

    #pragma unroll
    for (int i = 0; i < size(tdQrdQ_atomic); ++i) {
        atomicAdd(&tdQgdQaccum_atomic(i), tdQrdQ_atomic(i));
    }
}

# flash_attn_interface.py（简化）
dq_accum = torch.empty(..., dtype=torch.float32)
# 反向内核写入 dq_accum
dq = dq_accum.to(q.dtype)  # FP32 → FP16/BF16

// hopper/flash_bwd_kernel_sm90.h:267-273
if (tile_valid) {
    epilogue.store(params.epilogue, tdKrdK, tdVrdV, shared_storage,
                   tiled_mma_dKV, threadIdx.x - NumCopyThreads, block_coord);
} else {
    epilogue.store_zero(params.epilogue, threadIdx.x - NumCopyThreads, block_coord);
}

// epilogue_bwd.hpp::store()（简化）
// Step 1: FP32 → FP16/BF16
Tensor tdVrdV_out = make_tensor_like<Element>(tdVrdV);
flash::convert_type_out(tdVrdV, tdVrdV_out);
Tensor tdKrdK_out = make_tensor_like<Element>(tdKrdK);
flash::convert_type_out(tdKrdK, tdKrdK_out);

// Step 2: R2S（寄存器 → 共享内存）
named_barrier_sync(NumEpilogueThreads, EpilogueBarrier);
cute::copy(smem_tiled_copy_dKV, taccdVrdV, taccdVsdV);  // dV → smem_dv
cute::copy(smem_tiled_copy_dKV, taccdKrdK, taccdKsdK);  // dK → smem_dk

// Step 3: S2G（共享内存 → HBM）
if constexpr (Use_TMA) {
    // TMA Store 路径
    fence_view_async_shared();
    NamedBarrier::arrive(NumEpilogueThreads + NumThreadsPerWarp, EpilogueBarrier);

    // 最后一个 warp 执行 TMA Store
    if (warp_idx_sync == last_warp && cute::elect_one_sync()) {
        NamedBarrier::sync(NumEpilogueThreads + NumThreadsPerWarp, EpilogueBarrier);
        cute::copy(params.tma_store_dV, tdVsdV, tdVgdV);
        cute::copy(params.tma_store_dK, tdKsdK, tdKgdK);
        tma_store_arrive();
        tma_store_wait<0>();
    }
} else {
    // Direct Store 路径（带 predicate 的全局内存写入）
    flash::copy<Is_even_MN, Is_even_K>(gmem_tiled_copy_dKV, tdKVrdV, tdKVgdV, ...);
}

// epilogue_bwd.hpp::store()（GQA 路径，简化）
int bidh_kv = params.qhead_per_khead_divmod.divmod(bidh_idx_in_group, bidh);

// Deterministic 模式：信号量串行化
if constexpr (Deterministic) {
    Barrier::wait_eq(lock_ptr, thread_idx,
        n_block * num_batch * num_head_kv, bidh_idx_in_group);
}

// 原子累加 dV 到 dVaccum (FP32)
for (int i = 0; i < size(tdVrdV_atomic); ++i) {
    atomicAdd(&tdVgdV_atomic(i), tdVrdV_atomic(i));
}

if constexpr (Deterministic) {
    Barrier::arrive_inc(lock_ptr, thread_idx, n_block * num_batch * num_head_kv);
}

// 同样处理 dK

// mainloop_bwd_sm90_tma_gmma_ws.hpp:~838-845
if constexpr (ShuffleLSE) {
    // 每 8 个线程共享一份 LSE，通过 warp shuffle 广播
    for (int i = 0; i < kStatsPerThread; ++i) {
        tLSErLSE(i) = tLSEsLSE((thread_idx % 32) / 4 + i * 8, smem_pipe_read.index());
    }
    // 使用时：
    float lse = __shfl_sync(0xffffffff, tLSErLSE(mi / 8), (mi % 8) * 4 + (thread_idx % 4));
}

// mainloop_bwd_sm90_tma_gmma_ws.hpp:~967-995（简化）
if constexpr (Slice_dQKV_Mma) {
    // 第一半 dV
    flash::gemm<false, -1, dKV_swapAB, /*M_slice=*/0>(
        tiled_mma_dKV, tdVrP_cur, tdVrdO, tdVrdV);

    // 同步 + 第一半 dQ
    NamedBarrier::sync(NumMmaThreads, BwdNamedBarriers::PdS);
    flash::gemm<true, -1, dQ_swapAB, /*M_slice=*/0>(
        tiled_mma_dQ, tdQrdS_cur, tdQrK, tdQrdQ);

    // 第二半 dV（与第一半 dQ 的 atomicAdd 重叠）
    flash::gemm<false, 1, dKV_swapAB, /*M_slice=*/1>(
        tiled_mma_dKV, tdVrP_cur, tdVrdO, tdVrdV);

    // 第一半 dQ 的 atomicAdd（在第二半 dV GEMM 的同时执行）
    for (int i = 0; i < size(tdQrdQ_atomic) / 2; ++i) {
        atomicAdd(&tdQgdQaccum_atomic(i), tdQrdQ_atomic(i));
    }

    // 继续 dK 的两半 + 剩余 dQ atomicAdd ...
}

// 控制 P/dS 在共享内存中的布局
static constexpr bool dKV_swapAB = CollectiveMainloop::dKV_swapAB;