// hopper/flash_fwd_kernel_sm90.h:411
int const bidh_kv = !PackGQA ? params.mainloop.qhead_per_khead_divmod.divide(bidh) : bidh;

// hopper/flash_fwd_launch_template.h:147-151
int qhead_per_khead = !PackGQA ? 1 : cutlass::ceil_div(params.h, params.h_k);
int num_blocks_m = cutlass::ceil_div(params.seqlen_q * qhead_per_khead, get<0>(TileShape_MNK{}));
// Grid: (num_blocks_m, num_heads_kv, batch_size)  // 注意用 KV 头数

// hopper/pack_gqa.h:55-73
template <int NumThreadsPerRow, typename Engine, typename Layout, typename TensorC>
static auto compute_ptr(Tensor<Engine, Layout> &tensor, TensorC const &tRows,
                        cutlass::FastDivmod const &qhead_per_khead_divmod,
                        int const thread_idx, int const m_block) {
    Tensor tPrPtr = make_tensor<TensorType const*>(Shape<Int<NumPtrPerThread>>{});
    for (int i = 0; i < NumPtrPerThread; ++i) {
        int const row = i * NumThreads + get<0>(tRows(thread_idx % NumThreadsPerRow));
        int const idx = m_block * kBlockM + row;
        int m_idx, h_idx;
        m_idx = qhead_per_khead_divmod.divmod(h_idx, idx);  // 一次操作同时得到商和余数
        tPrPtr[i] = &tensor(make_coord(make_coord(h_idx, m_idx)));
    }
    return tPrPtr;
}

// hopper/pack_gqa.h:106-109
for (int m = 0; m < size<1>(tQsQ); ++m) {
    // 通过 __shfl_sync 获取指针，避免重复计算 divmod
    Element const* q_ptr = reinterpret_cast<Element const*>(
        __shfl_sync(0xffffffff,
                    reinterpret_cast<uint64_t>(tPrQPtr(m / kGmemThreadsPerRow)),
                    m % kGmemThreadsPerRow, kGmemThreadsPerRow));
    // ... 使用 q_ptr 加载数据
}

// hopper/pack_gqa.h:84-121
static void load_Q(TensormQ const &mQ,  // ((qhead_per_khead, seqlen_q), headdim)
                   TensorsQ &sQ,        // (kBlockM, kHeadDim)
                   ...) {
    // 使用 CpAsync 替代 TMA
    GmemTiledCopyQCpAsync gmem_tiled_copy_Q_cp_async;

    // 计算每行的指针
    Tensor tPrQPtr = compute_ptr(mQ_0, tQcQ_row, qhead_per_khead_divmod, ...);

    for (int m = 0; m < size<1>(tQsQ); ++m) {
        int idx = m_block * kBlockM + get<0>(tQcQ(_0{}, m, _0{}));
        // 获取共享指针
        Element const* q_ptr = /* __shfl_sync 获取 */;
        if (idx < seqlen_q * qhead_per_khead) {
            // 从 q_ptr 加载一行到 SMEM
            Tensor mQ_cur = make_tensor(make_gmem_ptr(q_ptr), Shape<Int<kHeadDim>>{});
            for (int k = 0; k < size<2>(tQsQ); ++k) {
                cute::copy(gmem_tiled_copy_Q_cp_async.with(tQpQ(k)),
                           mQ_cur_copy(_, ki), tQsQ(_, m, k));
            }
        }
    }
}

// hopper/pack_gqa.h:26-29
static constexpr int kBytePerRow = kHeadDim * sizeof(Element);
// 选择 128/64/32 字节中能整除 kBytePerRow 的最大值
static constexpr int kBlockKGmem = (kBytePerRow % 128 == 0 ? 128
                                  : (kBytePerRow % 64 == 0 ? 64 : 32)) / sizeof(Element);
static constexpr int kGmemThreadsPerRow = kBlockKGmem / kGmemElemsPerLoad;

static_assert(cutlass::NumThreadsPerWarp % kGmemThreadsPerRow == 0,
              "kGmemThreadsPerRow must divide NumThreadsPerWarp");

// hopper/pack_gqa.h:159-198  store_O
static void store_O(TensormO &mO, TensorrO const &tOrO,
                    cutlass::FastDivmod const &qhead_per_khead_divmod, ...) {
    // 与 load_Q 类似的指针计算和 __shfl_sync 广播
    Tensor tPrOPtr = compute_ptr(mO_0, tOcO_row, qhead_per_khead_divmod, ...);
    for (int m = 0; m < size<1>(tOrO); ++m) {
        Element* o_ptr = /* __shfl_sync 获取 */;
        if (idx < seqlen_o * qhead_per_khead) {
            // 写回到正确的 O 位置
            cute::copy(gmem_tiled_copy_O, tOrO(_, m, k), mO_cur_copy(_, ki));
        }
    }
}

// hopper/mask.h:79-86
if constexpr (PackGQA) {
    // 从全局索引中反推实际的序列位置
    mma_m_idx = qhead_per_khead_divmod.divide(
        m_block * kBlockM + get<Row>(tScS_rowcol(...))
    );
    // 使用 __shfl_sync 广播到同组线程
}

// hopper/block.h:24-34
if constexpr (Is_causal || Is_local) {
    int m_idx_max = (m_block + 1) * kBlockM;
    if (PackGQA) {
        // 从打包索引计算实际的最大序列位置
        m_idx_max = qhead_per_khead_divmod.divide(m_idx_max - 1) + 1;
    }
    int const n_idx = m_idx_max + seqlen_k - seqlen_q;
    n_block_max = std::min(n_block_max, cute::ceil_div(n_idx_right, kBlockN));
}

// hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp:1105-1119
int const qhead_per_khead = !PackGQA ? 1 : params.qhead_per_khead_divmod.divisor;
if (params.is_rotary_interleaved) {
    auto [tRrCos, tRrSin] = cute::conditional_return<!PackGQA>(
        rotary.template load_cos_sin<true>(m_block),
        rotary.template load_cos_sin_packgqa<true>(m_block, params.qhead_per_khead_divmod)
    );
    // ...
    rotary.apply_Q_interleaved(sQ_pi, tRrCos, tRrSin, m_block, qhead_per_khead);
}

// hopper/flash_bwd_launch_template.h:99
// 反向使用 SingleTileScheduler，PackGQA 固定为 false
flash::SingleTileScheduler<Varlen, false /*Split*/, false /*PackGQA*/, kBlockN>

# flash_attn/flash_attn_interface.py 中的反向传播
# dK 原始形状: (batch, seqlen_k, num_heads_q, headdim)
# 需要归约到:   (batch, seqlen_k, num_heads_k, headdim)
if num_heads_q != num_heads_k:
    dK = dK.reshape(batch, seqlen_k, num_heads_k, num_heads_q // num_heads_k, headdim)
    dK = dK.sum(dim=3)  # 沿 Q 头组维度求和
    # dV 同理

// 内核输出 dK 形状为 (batch, seqlen_k, num_heads_q, headdim)
// 如果 num_heads_q > num_heads_k，需要归约
// reshape 为 (batch, seqlen_k, num_heads_k, G, headdim)
// 对 G 维度求和

// hopper/heuristics.h:9-17
inline bool should_pack_gqa(bool varlen_q, int seqlen_q, int qhead_per_khead, int blockM) {
    // 如果是变长序列，始终启用 PackGQA
    if (varlen_q) return true;

    // 比较两种模式的 Tile 利用率
    float nopack_gqa_efficiency = float(seqlen_q) / float(round_up(seqlen_q, blockM));
    float pack_gqa_efficiency = float(seqlen_q * qhead_per_khead)
                               / float(round_up(seqlen_q * qhead_per_khead, blockM));

    // 只有当打包模式效率显著更高时才启用（90% 阈值）
    return nopack_gqa_efficiency < 0.9 * pack_gqa_efficiency;
}

// hopper/flash_api.cpp:426-439
inline bool get_pack_gqa(Flash_fwd_params const& params) {
    // SM8x、非TMA PagedKV、Split 模式下强制启用
    if (params.arch < 90 || (params.page_table && !params.pagedkv_tma) || params.num_splits > 1) {
        return true;
    }
    // MHA (num_heads == num_heads_k) 不需要 PackGQA
    if (params.h == params.h_k) { return false; }
    // 其他情况使用启发式决策
    return should_pack_gqa(...);
}

from flash_attn import flash_attn_func

# 自动决策（默认）
out = flash_attn_func(q, k, v, causal=True)

# 手动启用 PackGQA
out = flash_attn_func(q, k, v, causal=True, pack_gqa=True)

# 手动禁用 PackGQA
out = flash_attn_func(q, k, v, causal=True, pack_gqa=False)