deepseek-ai / deepseek-ai/FlashMLA

dual gemm

Open
#169 2 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.