dual gemm

[WalrusWTQ]

Hi, I'm studying the source code of flashmla. I noticed that in sparse decode head64 implementaion, it uses "dual gemm" to compute P=QK^T.
I have a few questions about this design.

if constexpr (MODEL_TYPE == ModelType::V32) {
    // V3.2: RoPE behaves like an extra block with size 64, so we can do RoPE first
    // QK RoPE
    plan.bar_rope_ready[rs.buf_idx].wait(rs.bar_phase);
    ku::tcgen05_after_thread_sync();
    Tensor tQ_rope = tiled_mma_P.get_slice(_0{}).make_fragment_A(
        partition_shape_A(tiled_mma_P, Shape<Int<B_H>, Int<D_ROPE/2>>{})
    );
    tQ_rope.data().get() = tmem_cols::Q_Tail;
    Tensor sK_rope = make_tensor(make_smem_ptr(plan.u.kv.dequant[rs.buf_idx].rope.data()), SmemLayoutKTiles_DualGemm_SW64<2/2>{});
    ku::utcmma_ts(tiled_mma_P, tQ_rope, sK_rope, tP, true);

    // QK NoPE
    plan.bar_nope_ready[rs.buf_idx].wait(rs.bar_phase);
    ku::tcgen05_after_thread_sync();
    Tensor tQ_nope = tiled_mma_P.get_slice(_0{}).make_fragment_A(
        partition_shape_A(tiled_mma_P, Shape<Int<B_H>, Int<D_NOPE/2>>{})
    );
    tQ_nope.data().get() = tmem_cols::Q;
    Tensor sK_nope = make_tensor(make_smem_ptr(plan.u.kv.dequant[rs.buf_idx].nope.data()), SmemLayoutKTiles_DualGemm_SW128<512/64/2>{});
    ku::utcmma_ts(tiled_mma_P, tQ_nope, sK_nope, tP, false);

1. In the layout, the K dimension is divided by 2, but there is no corresponding loop iterating twice for the computation. Why is that?
2. Why use dual gemm? Doesn't it introduce reduction overhead afterwards in the Scale & Exp warpgroup?

[zchenCN]

I also do not understand the "dual gemm" trick. Does the memory calculation of tA_frag in gemm.cuh matters.

[imisszxq]

When processing MMA computations with M=64 using the tcgen05.mma.ws instruction, a 2×2 duplicated data pattern allows executing 2N computations with a single MMA instruction. In this case, the computation from N to 2N requires fetching data from the tmem corresponding to warpgroup2-3. Dual GEMM leverages this feature by partitioning the 2K dimension to simulate the 2N computation, enabling warpgroup0-1 to obtain MMA results for rows 0–63 and columns 0–k-1, while warpgroup2-3 obtains results for rows 0–63 and columns k–2k. The final result is then obtained through the retrieve_mask_and_reduce_p function.