dual gemm
Author: WalrusWTQCreated Mar 6, 2026Updated Jul 13, 2026
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);- In the layout, the K dimension is divided by 2, but there is no corresponding loop iterating twice for the computation. Why is that?
- Why use dual gemm? Doesn't it introduce reduction overhead afterwards in the Scale & Exp warpgroup?
Source: deepseek-ai/FlashMLA