// mainloop_fwd_sm90_tma_gmma_ws.hpp::load()（简化结构）
template <typename SharedStorage>
CUTLASS_DEVICE void load(
    Params const& params,
    MainloopPipelineK pipeline_k,
    MainloopPipelineV pipeline_v,
    MainloopPipelineVt pipeline_vt,
    PipelineState& smem_pipe_write,
    SharedStorage& shared_storage,
    auto scheduler_prefetch,
    SeqlenInfo_t const& seqlen_info,
    cute::tuple<int32_t, int32_t, int32_t, int32_t> block_coord,
    int work_idx) {
    // 1. 加载 Q → smem_q
    // 2. 计算有效 N 块范围
    // 3. 循环加载 K/V → smem_k[stage], smem_v[stage]
    // 4. 调度器预取
}

// mainloop_fwd_sm90_tma_gmma_ws.hpp:~812-846（简化）
// TMA 模式
if constexpr (Use_TMA_Q) {
    if (warp_idx_in_warpgroup == 0) {
        // 等待 Consumer 消费完上一个 Q
        cutlass::arch::NamedBarrier::sync(NumMmaThreadsQK + cutlass::NumThreadsPerWarp,
            static_cast<uint32_t>(FwdNamedBarriers::QueryEmpty));

        if (cute::elect_one_sync()) {
            // 通知 TMA 即将写入的字节数
            shared_storage.pipelines.barrier_Q.arrive_and_expect_tx(TmaTransactionBytesQ);
            // 发起 TMA 加载（仅 1 个线程）
            copy(params.tma_load_Q.with(
                     reinterpret_cast<BarrierQ::ValueType&>(shared_storage.pipelines.barrier_Q),
                     0 /*mcast_mask*/),
                 tQgQ(_, m_block, bidh_kv, bidb_kv),
                 tQsQ);
        }
    }
}

// mainloop_fwd_sm90_tma_gmma_ws.hpp:~750-770（简化）
auto load_K = [&](int const n_block, auto const& smem_pipe_write, auto need_seqlenk_masking_type) {
    pipeline_k.producer_acquire(smem_pipe_write);     // 等待 stage 可用

    if constexpr (!PagedKVNonTMA) {
        // TMA 模式
        copy(params.tma_load_K.with(
                 *pipeline_k.producer_get_barrier(smem_pipe_write),
                 mcast_mask_kv,
                 TMA::CacheHintSm90::EVICT_LAST),     // L2 缓存提示
             tKgK_TMA(_, n_block_idx, bidb_kv_idx),
             tKsK_TMA(_, smem_pipe_write.index()));
    } else {
        // cp.async 模式（PagedKV 非 TMA）
        paged_kv_manager.template load_K<Seqlenk_mask>(n_block, sK_pi(_, _, smem_pipe_write.index()));
        pipeline_k.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive);
    }
};

auto load_V = [&](int const n_block, auto const& smem_pipe_write, auto need_seqlenk_masking_type) {
    auto pipeline_v_load = cute::conditional_return<!Transpose_V>(pipeline_v, pipeline_vt);
    pipeline_v_load.producer_acquire(smem_pipe_write);

    copy(params.tma_load_V.with(
             *pipeline_v_load.producer_get_barrier(smem_pipe_write),
             mcast_mask_kv,
             TMA::CacheHintSm90::EVICT_LAST),
         tVgVt_TMA(_, n_block_idx, bidb_kv_idx),
         tVsVt_TMA(_, smem_pipe_write.index()));
};

// mainloop_fwd_sm90_tma_gmma_ws.hpp:~860-888（简化）
int n_block_prev = n_block_max - 1;

// 首先加载第一个 K 块
load_K(n_block_max - 1, smem_pipe_write, ...);

// 内层循环：从后向前遍历 N 块
#pragma unroll (!Transpose_V && Use_TMA_KV ? 2 : 1)
for (; n_block >= n_block_min; --n_block) {
    PipelineState smem_pipe_write_v = smem_pipe_write;  // V 使用当前 stage
    ++smem_pipe_write;                                   // K 推进到下一个 stage

    load_K(n_block, smem_pipe_write, ...);               // 加载下一个 K

    if constexpr (IntraWGOverlap) {
        load_V(n_block_prev, smem_pipe_write_v, ...);    // 加载上一个 V（重叠）
    } else {
        load_V(n_block, smem_pipe_write_v, ...);         // 加载当前 V
    }

    n_block_prev = n_block;
}

// mainloop_fwd_sm90_tma_gmma_ws.hpp（简化）
CUTLASS_DEVICE void load_tail(
    MainloopPipelineK pipeline_k, MainloopPipelineV pipeline_v,
    MainloopPipelineVt pipeline_vt, PipelineState& smem_pipe_write, ...) {
    // 在最后一个 stage 发送"空"提交，通知 Consumer 数据流结束
    pipeline_k.producer_tail(smem_pipe_write);
    pipeline_v.producer_tail(smem_pipe_write);
}

// mainloop_fwd_sm90_tma_gmma_ws.hpp::mma()（简化）
CUTLASS_DEVICE bool mma(
    Params const& params,
    MainloopPipelineK pipeline_k,
    MainloopPipelineV pipeline_v,
    PipelineState& smem_pipe_read,
    Tensor<Engine, Layout>& tOrO,       // 输出累加器（寄存器）
    Softmax& softmax,
    int thread_idx, int work_idx,
    SeqlenInfo_t const& seqlen_info,
    cute::tuple<...> block_coord,
    SharedStorage& shared_storage) {

    // 共享内存张量绑定
    Tensor sQ = make_tensor(make_smem_ptr(shared_storage.tensors.mainloop.smem_q.data()), SmemLayoutQ{});
    Tensor sK = make_tensor(make_smem_ptr(shared_storage.tensors.mainloop.smem_k.data()), SmemLayoutK{});
    Tensor sV = make_tensor(make_smem_ptr(shared_storage.tensors.mainloop.smem_v.data()), SmemLayoutVtMma{});

    // MMA 对象
    TiledMmaQK tiled_mma_qk;    // Q × K → Score
    TiledMmaPV tiled_mma_pv;    // P × V → Output

    // 线程级张量分区
    Tensor tSrQ = tiled_mma_qk.partition_fragment_A(sQ);              // Q 在寄存器中的片段
    Tensor tSrK = tiled_mma_qk.partition_fragment_B(sK);              // K 的共享内存视图
    Tensor tOrV = tiled_mma_pv.partition_fragment_B(sV);              // V 的共享内存视图

    // 等待 Q 加载完成
    shared_storage.pipelines.barrier_Q.wait(work_idx % 2);

    // 计算有效的 N 块范围
    auto [n_block_min, n_block_max] = BlockMN::get_n_block_min_max(
        seqlen_info, m_block, bidb, split_idx, params.num_splits,
        kBlockM, kBlockN, params.window_size_left, params.window_size_right);

    if (n_block_max <= n_block_min) {
        // 通知 Producer 可以加载新 Q
        if (warp_group_thread_idx == 0) {
            NamedBarrier::arrive(NumMmaThreadsQK + NumThreadsPerWarp,
                static_cast<uint32_t>(FwdNamedBarriers::QueryEmpty));
        }
        return false;  // 无有效块，跳过
    }

    // mainloop_fwd_sm90_tma_gmma_ws.hpp:~1254-1349（简化）
    clear(tOrO);  // 输出归零

    auto fwd_step = [&](int n_block, auto mask_fn, auto check_inf_type) {
        // === Step 1: 等待 K 就绪 ===
        consumer_wait(pipeline_k, smem_pipe_read);

        // === Step 2: Score GEMM — S = Q @ K^T ===
        Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_MNK{}));
        flash::gemm</*zero_init=*/true, /*wg_wait=*/-1>(
            tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read.index()), tSrS);
        warpgroup_wait<0>();                    // 等待 GMMA 完成

        // === Step 3: Softcap（可选）===
        if constexpr (Has_softcap) {
            flash::apply_softcap(tSrS, params.softcap_val);
        }

        // === Step 4: Masking ===
        mask_fn(tSrS, n_block);

        // === Step 5: Online Softmax ===
        Tensor scores_scale = softmax.template max_get_scale<Is_first, Check_inf>(tSrS);
        softmax.template online_softmax<Is_first, Check_inf>(tSrS);

        // === Step 6: Rescale 已有输出 ===
        softmax.rescale_o(tOrO, scores_scale);

        // === Step 7: 将 P 写入 smem 或保持在寄存器 ===
        if constexpr (!MmaPV_is_RS) {
            // 类型转换：FP32 → FP16/BF16
            Tensor tOrP = make_tensor_like<Element>(tSrS);
            flash::convert_type_out(tSrS, tOrP);
            // 写入共享内存
            cute::copy(smem_tiled_copy_P, tOrP, tPsP);
        }

        // === Step 8: 等待 V 就绪 ===
        consumer_wait(pipeline_v, smem_pipe_read);

        // === Step 9: Output GEMM — O += P @ V ===
        flash::gemm</*zero_init=*/false, /*wg_wait=*/-1>(
            tiled_mma_pv,
            cute::conditional_return<MmaPV_is_RS>(tOrP_reg, tOsP),
            tOrV(_, _, _, smem_pipe_read.index()),
            tOrO);
        warpgroup_wait<0>();

        // === Step 10: 释放 K/V Pipeline stage ===
        pipeline_k.consumer_release(smem_pipe_read);
        ++smem_pipe_read;
    };

    // 遍历所有 N 块（从后向前）
    for (int n_block = n_block_max - 1; n_block >= n_block_min; --n_block) {
        fwd_step(n_block, mask_fn, ...);
    }

    // === 最终归一化 ===
    Tensor final_scale = softmax.finalize();
    softmax.rescale_o(tOrO, final_scale);

    return true;

// mainloop_fwd_sm90_tma_gmma_ws.hpp:~1170-1207（简化）
auto fwd_step = [&](int n_block, auto mask_fn, auto check_inf_type) {
    // V 使用上一个 stage，K 使用当前 stage
    PipelineState smem_pipe_read_v(smem_pipe_read.index(), smem_pipe_read.phase(), ...);
    ++smem_pipe_read;  // K 推进到下一个 stage

    // === 等待 K 就绪并发射 QK GEMM ===
    consumer_wait(pipeline_k, smem_pipe_read);
    warp_scheduler_barrier_sync();

    Tensor tSrS = partition_fragment_C(tiled_mma_qk, ...);
    flash::gemm</*zero_init=*/true, /*wg_wait=*/-1>(
        tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read.index()), tSrS);

    // === 在 QK GMMA 异步执行的同时，做上一轮的 rescale ===
    if constexpr (RescaleOBeforeGemm) {
        softmax.rescale_o(tOrO, scores_scale);  // 使用上一轮的 scale
    }

    // === 发射 PV GEMM（使用上一轮的 P 和上一个 stage 的 V） ===
    flash::gemm</*zero_init=*/false, /*wg_wait=*/-1>(
        tiled_mma_pv,
        cute::conditional_return<MmaPV_is_RS>(tOrP, tOsP),
        tOrV(_, _, _, smem_pipe_read_v.index()),  // 上一个 V stage
        tOrO);

    warp_scheduler_barrier_arrive();

    // === 等待 QK GEMMA 完成（PV 可能仍在后台执行） ===
    warpgroup_wait<1>();     // 等待到最多 1 个 GMMA 未完成（即 PV）
    pipeline_k.consumer_release(smem_pipe_read);

    // === 对当前 Score 做 Softcap + Mask + Softmax ===
    scoremod_premask_fn(tSrS);
    mask_fn(tSrS, n_block);

    cute::copy(softmax.template max_get_scale<false, Check_inf>(tSrS), scores_scale);
    softmax.template online_softmax<false, Check_inf>(tSrS);

    // === 将 P 转换为 FP16 并写入 smem ===
    if constexpr (!MmaPV_is_RS) {
        write_P_to_smem(tOrP);
    }

    if constexpr (!RescaleOBeforeGemm) {
        softmax.rescale_o(tOrO, scores_scale);
    }
};

auto [n_block_min, n_block_max] = BlockMN::get_n_block_min_max(...);
// n_block_min ~ n_block_max 之外的块完全跳过

// hopper/mask.h（简化）
template <bool Seqlenk_mask, bool Causal_mask, bool Local_mask>
CUTLASS_DEVICE void apply(Tensor& tSrS, int m_block, int n_block) {
    #pragma unroll
    for (int mi = 0; mi < size<0>(tSrS); ++mi) {
        int const row_idx = m_block * kBlockM + get<0>(idx_row(mi));
        #pragma unroll
        for (int ni = 0; ni < size<1>(tSrS); ++ni) {
            int const col_idx = n_block * kBlockN + get<1>(idx_col(ni));
            if constexpr (Causal_mask) {
                if (col_idx > row_idx + seqlen_k - seqlen_q) {
                    tSrS(mi, ni) = -INFINITY;  // 因果遮蔽
                }
            }
            if constexpr (Local_mask) {
                if (col_idx > row_idx + window_size_right || col_idx < row_idx - window_size_left) {
                    tSrS(mi, ni) = -INFINITY;  // 窗口遮蔽
                }
            }
        }
    }
}

// 类型 1: 需要 Causal/Local masking 的边界块
for (n_block = n_block_max - 1; n_block >= n_block_min_causal_local_mask; --n_block) {
    fwd_step(n_block, causal_local_mask_fn, /*Check_inf=*/true);
}

// 类型 2: 完全有效的块（无需 masking）
for (; n_block >= n_block_min_before_local_mask; --n_block) {
    fwd_step(n_block, no_mask_fn, /*Check_inf=*/false);
}

// 类型 3: 需要 Local masking 的边界块（窗口左边界）
if constexpr (Is_local) {
    for (; n_block >= n_block_min; --n_block) {
        fwd_step(n_block, local_mask_fn, /*Check_inf=*/true);
    }
}

// epilogue_fwd.hpp:~251-265（简化）
if constexpr (Use_smem) {
    // 类型转换
    Tensor tOrO_out = make_tensor_like<Element>(tOrO);
    flash::convert_type_out(tOrO, tOrO_out);  // FP32 → FP16/BF16

    // 写入 smem_o
    auto smem_tiled_copy_O = make_tiled_copy_C(SmemCopyAtomO{}, tiled_mma);
    auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(thread_idx);
    Tensor taccOrO = smem_thr_copy_O.retile_S(tOrO_out);
    Tensor taccOsO = smem_thr_copy_O.partition_D(sO);
    cute::copy(smem_tiled_copy_O, taccOrO, taccOsO);

    // 如果使用 TMA Store，通知 TMA 数据已就绪
    if constexpr (Use_TMA_O) {
        cutlass::arch::fence_view_async_shared();
        cutlass::arch::NamedBarrier::arrive(NumEpilogueThreads + cutlass::NumThreadsPerWarp,
            cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
    }
}

// epilogue_fwd.hpp:~311-330（简化）
if constexpr (Use_TMA_O) {
    // 构造 TMA 张量
    Tensor mO = params.tma_store_O.get_tma_tensor(params.shape_O)(_, _, bidh, bidb, split_idx);
    Tensor gO = local_tile(mO, select<0, 1>(TileShape_MNK_PV{}), make_coord(m_block, _0{}));

    auto block_tma_O = params.tma_store_O.get_slice(_0{});
    Tensor tOgO = block_tma_O.partition_D(gO);
    Tensor tOsO = block_tma_O.partition_S(sO);

    // 最后一个 warp 执行 TMA Store
    int warp_idx_sync = __shfl_sync(0xffffffff, thread_idx / NumThreadsPerWarp, 0);
    if (warp_idx_sync == NumEpilogueThreads / NumThreadsPerWarp - 1) {
        // 等待所有线程完成 smem 写入
        cutlass::arch::NamedBarrier::sync(NumEpilogueThreads + cutlass::NumThreadsPerWarp,
            cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
        if (cute::elect_one_sync()) {
            cute::copy(params.tma_store_O, tOsO, tOgO);    // TMA 异步写回
            tma_store_arrive();
            tma_store_wait<0>();                              // 等待写回完成
            // 通知 Cluster 内其他 block
            for (uint32_t cta_id = 0; cta_id < size(ClusterShape{}); ++cta_id) {
                shared_storage.pipelines.barrier_O.arrive(cta_id);
            }
        }
    }
}

// epilogue_fwd.hpp:~331-401（简化）
if constexpr (!Use_TMA_O) {
    Tensor mO = make_tensor(make_gmem_ptr(params.ptr_O + offset_o * stride_O), ...);
    Tensor gO = local_tile(mO, TileShape_PV{}, make_coord(m_block, _0{}));

    // 使用向量化全局内存写入
    GmemTiledCopyO gmem_tiled_copy_O;
    auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(thread_idx);
    Tensor tOrO = make_fragment_like(tOsO);
    cute::copy(gmem_tiled_copy_O, tOsO, tOrO);     // smem → register

    // 带 predicate 的全局内存写入
    Tensor tOpO = make_tensor<bool>(make_shape(size<2>(tOsO)));
    for (int k = 0; k < size(tOpO); ++k) {
        tOpO(k) = get<1>(tOcO(_0{}, _0{}, k)) < get<1>(params.shape_O);  // 列边界检查
    }
    flash::copy</*Is_even_MN=*/false, /*Is_even_K=*/false>(
        gmem_tiled_copy_O, tOrO, tOgO, tOcO, tOpO, seqlen_o - m_block * kBlockM);
}

// epilogue_fwd.hpp:~281-308（简化）
Tensor mLSE = make_tensor(make_gmem_ptr(params.ptr_LSE + offset * stride_LSE), ...)(_, bidh, bidb);

for (int mi = 0; mi < size(lse); ++mi) {
    int const row = m_block * kBlockM + get<0>(taccOcO_row(mi));
    if (get<1>(taccOcO_row(_0{})) == 0 && row < seqlen_o) {
        mLSE(row) = lse(mi);   // 直接全局内存写入
    }
}

// 非 Split 模式：直接写入 gO, gLSE
// Split 模式：写入 gO_partial, gLSE_partial
if constexpr (Split) {
    Tensor mOpartial = make_tensor(make_gmem_ptr(params.ptr_O_partial + ...),
        params.shape_O_packed, ...)(_, _, bidh, bidb, split_idx);
    // 写入临时缓冲区
}

// Producer 端的 V 转置
auto copy_Vt_to_V = [&](auto const& smem_pipe_write) {
    pipeline_vt.consumer_wait(smem_pipe_write);     // 等待 TMA 加载完成
    pipeline_v.producer_acquire(smem_pipe_write);   // 获取 V pipeline stage

    // LDSM: 读取 smem_vt 到寄存器
    cute::copy(LDSM_copy, tVsVt(_, _, _, smem_pipe_write.index()), tVrV);
    // STSM: 转置写入寄存器到 smem_v
    cute::copy(STSM_copy, tVrV, tVsV(_, _, _, smem_pipe_write.index()));

    pipeline_vt.consumer_release(smem_pipe_write);
    cute::arch::fence_view_async_shared();
    pipeline_v.producer_commit(smem_pipe_write);
};

// hopper/flash_fwd_kernel_sm90.h:408-415
if constexpr (Is_FP8 && !Has_softcap) {
    int const bidh_kv = !PackGQA ? params.mainloop.qhead_per_khead_divmod.divide(bidh) : bidh;
    float const q_descale = params.mainloop.ptr_q_descale == nullptr
        ? 1.0f : params.mainloop.ptr_q_descale[bidb * stride + bidh_kv * stride];
    float const k_descale = params.mainloop.ptr_k_descale == nullptr
        ? 1.0f : params.mainloop.ptr_k_descale[bidb * stride + bidh_kv * stride];
    softmax_scale_log2 *= q_descale * k_descale;  // 将 descale 融入 softmax scale
}

// hopper/flash_fwd_kernel_sm90.h:416
flash::Softmax<..., /*Max_offset=*/!Is_FP8 ? 0 : 8> softmax(softmax_scale_log2);