// hopper/flash_fwd_kernel_sm90.h:27-28
template <class CollectiveMainloop_, class CollectiveEpilogue_, class TileScheduler_>
class FlashAttnFwdSm90 {

// hopper/flash_fwd_kernel_sm90.h:74-78
static constexpr uint32_t NumLoadWarpGroups = 1;
static constexpr uint32_t NumMmaWarpGroups = size(TiledMmaPV{}) / NumThreadsPerWarpGroup;
static constexpr uint32_t MaxThreadsPerBlock = size(TiledMmaPV{}) + NumLoadWarpGroups * NumThreadsPerWarpGroup;
// NumMmaWarpGroups 可以是 1, 2, 或 3

// hopper/flash_fwd_kernel_sm90.h:82-83
static constexpr uint32_t LoadRegisterRequirement = NumMmaWarpGroups == 2 ? 24 : 32;
static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 2 ? 240 : 160;

// hopper/flash_fwd_kernel_sm90.h:94-117
struct SharedStorage {
    struct TensorStorage {
        union {                     // Mainloop 和 Epilogue 共享内存重叠
            struct {
                padding_;           // 对齐填充
                mainloop;           // smem_q, smem_k, smem_v
            };
            epilogue;               // smem_o（与 smem_v 重叠）
        };
    } tensors;
    struct PipelineStorage {
        barrier_Q;                  // Q 加载同步
        barrier_O;                  // O 输出同步
        pipeline_k;                 // K 加载流水线状态
        pipeline_v;                 // V 加载流水线状态
        pipeline_vt;                // V 转置流水线状态 (FP8)
        smem_scheduler;             // 调度器共享内存
    } pipelines;
};

// hopper/flash_fwd_kernel_sm90.h:308-452
if (warp_group_idx == 0) {  // Producer
    cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
    // ... TMA 加载逻辑
} else {  // Consumer
    cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
    // ... MMA 计算逻辑
}

// hopper/flash_fwd_kernel_sm90.h:328-359（简化）
for (auto work_tile_info = scheduler.get_initial_work<IsProducerWarp>(params.scheduler);
     work_tile_info.is_valid(params.scheduler);
     work_tile_info = scheduler.get_next_work<IsProducerWarp>(params.scheduler, work_tile_info)) {

    // 1. 构造序列长度信息
    SeqlenInfo_t seqlen_info{bidb, shape_Q, shape_K, ...};

    // 2. 如果启用 AppendKV，先加载新的 KV
    if constexpr (AppendKV) {
        mainloop.load_kv_new(params.mainloop, pipeline_k_new, pipeline_v_new, ...);
    }

    // 3. 预取下一个 tile 的工作信息
    auto scheduler_prefetch = [&]() {
        scheduler.prefetch_next_work(params.scheduler, work_tile_info);
    };

    // 4. 加载 Q, K, V 到共享内存
    mainloop.load(params.mainloop, pipeline_k, pipeline_v, pipeline_vt,
                  smem_pipe_write, shared_storage, scheduler_prefetch, seqlen_info, block_coord, work_idx);
}
// 5. 通知所有流水线加载结束
mainloop.load_tail(pipeline_k, pipeline_v, pipeline_vt, smem_pipe_write, ...);

// hopper/flash_fwd_kernel_sm90.h:360-451（简化）
cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
TiledMmaPV tiled_mma_pv;

for (auto work_tile_info = scheduler.get_initial_work(params.scheduler);
     work_tile_info.is_valid(params.scheduler);
     /* work_tile_info 在 epilogue 前更新 */) {

    // 1. FP8 缩放因子调整
    float softmax_scale_log2 = params.mainloop.softmax_scale_log2;
    if constexpr (Is_FP8) {
        softmax_scale_log2 *= q_descale * k_descale;
    }

    // 2. 初始化 Softmax 和输出累加器
    Softmax softmax(softmax_scale_log2);
    Tensor tOrO = partition_fragment_C(tiled_mma_pv, ...);  // 输出寄存器

    // 3. 执行主循环 MMA
    bool tile_valid = mainloop.mma(
        params.mainloop, pipeline_k, pipeline_v, smem_pipe_read,
        tOrO, softmax, threadIdx.x, work_idx, seqlen_info, block_coord, shared_storage);

    // 4. 预取下一个 tile（在 epilogue 之前！）
    work_tile_info = scheduler.get_next_work(params.scheduler, work_tile_info);

    // 5. 写回结果
    if (tile_valid) {
        epilogue.store(params.epilogue, tOrO, softmax.row_sum, ...);
    } else {
        epilogue.store_zero(params.epilogue, ...);  // 全零 O, -inf LSE
    }
}

for n_block = n_block_max - 1 downto n_block_min

// hopper/flash_fwd_kernel_sm90.h:442-449
if (tile_valid) {
    epilogue.store(params.epilogue, tOrO, softmax.row_sum, shared_storage,
                   tiled_mma_pv, threadIdx.x, block_coord);
} else {
    epilogue.store_zero(params.epilogue, threadIdx.x, block_coord);
}