deepseek-ai / deepseek-ai/FlashMLA
dual gemm
- Dominant language
- C++
- Stars
- 12.9k
- Forks
- 1.2k
- Avg merge
- 4h 20m
- Merged PRs (30d)
- 2
Description
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>{})
);
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>{})
);
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?
Contributor guide
No contributing guide indexed for this repository
Research direction
Start from the sparse decode head64 implementation containing the V32 dual-gemm code, then trace SmemLayoutKTiles_DualGemm_SW64 and SmemLayoutKTiles_DualGemm_SW128 together with ku::utcmma_ts. Follow how tP is consumed by the Scale & Exp warpgroup; done means explaining the K-dimension layout split and whether dual gemm requires a reduction and why.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend, performance
- Issue type
- Documentation
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100