deepseek-ai/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00
Code Issues 162 Projects Packages Wiki Activity

[GH-ISSUE #15] Both layouts of Grouped GEMMs need to be aligned to the GEMM M block size ? #14

New issue
Closed
opened 2026-10-26 21:10:30 +00:00 by bsakaguchi · 6 comments
bsakaguchi commented 2026-10-26 21:10:30 +00:00
Owner
Copy link

Originally created by @CtfGo on GitHub (Feb 26, 2025).
Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/15

  1. The m_grouped_gemm_fp8_fp8_bf16_nt_contiguous function documentation says that

"On the M axis, inputs are grouped into several batches, of which batch sizes aligned to get_m_alignment_for_contiguous_layout() (128).",

which means the M of each group must be multiples of 128？

  1. I don't see such declaration on the m_grouped_gemm_fp8_fp8_bf16_nt_masked, but I test several M (8, 16, 64) that don't meet the requirement on test/test_core.py by test_m_grouped_gemm_masked and found they can't pass the correctness validation, so does the masked layout grouped gemm also has the limitation implicitly?
Originally created by @CtfGo on GitHub (Feb 26, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/15 1. The `m_grouped_gemm_fp8_fp8_bf16_nt_contiguous` function documentation says that > "On the M axis, inputs are grouped into several batches, of which batch sizes aligned to `get_m_alignment_for_contiguous_layout()` (128).", which means the M of each group must be multiples of 128？ 2. I don't see such declaration on the `m_grouped_gemm_fp8_fp8_bf16_nt_masked`, but I test several M `(8, 16, 64)` that don't meet the requirement on `test/test_core.py` by `test_m_grouped_gemm_masked` and found they can't pass the correctness validation, so does the masked layout grouped gemm also has the limitation implicitly?
bsakaguchi commented 2026-10-26 21:10:31 +00:00
Author
Owner
Copy link

@LyricZhao commented on GitHub (Feb 27, 2025):

which means the M of each group must be multiples of 128？

Yes.

but I test several M (8, 16, 64) that don't meet the requirement on test/test_core.py by test_m_grouped_gemm_masked and found they can't pass the correctness validation, so does the masked layout grouped gemm also has the limitation implicitly

No limitation for the masked kernel, could you please share you modified test script?

<!-- gh-comment-id:2686735320 --> @LyricZhao commented on GitHub (Feb 27, 2025): > which means the M of each group must be multiples of 128？ Yes. > but I test several M (8, 16, 64) that don't meet the requirement on test/test_core.py by test_m_grouped_gemm_masked and found they can't pass the correctness validation, so does the masked layout grouped gemm also has the limitation implicitly No limitation for the masked kernel, could you please share you modified test script?
bsakaguchi commented 2026-10-26 21:10:32 +00:00
Author
Owner
Copy link

@CtfGo commented on GitHub (Feb 27, 2025):

@LyricZhao Hi, thank you for your reply and appreciate your awesome work.

No limitation for the masked kernel, could you please share you modified test script?

I just modify the two lines on test/test_core.py test_m_grouped_gemm_masked to test other num_groups and m, the detail diff is:

Image

then run python test/test_core.py, and will see AssertionError of the diff check:

Image

test machine environment: Nvidia H800、cuda 12.4、torch 2.6.0+cu124

<!-- gh-comment-id:2686823697 --> @CtfGo commented on GitHub (Feb 27, 2025): @LyricZhao Hi, thank you for your reply and appreciate your awesome work. > No limitation for the masked kernel, could you please share you modified test script? I just modify the two lines on `test/test_core.py test_m_grouped_gemm_masked` to test other `num_groups` and `m`, the detail diff is: <img width="1262" alt="Image" src="https://github.com/user-attachments/assets/1ac1facd-9796-4ece-bb9f-340778eab3b0" /> then run `python test/test_core.py`, and will see AssertionError of the diff check: <img width="1032" alt="Image" src="https://github.com/user-attachments/assets/a43e765f-6cd4-45fe-b21e-e2116e7d196b" /> test machine environment: Nvidia H800、cuda 12.4、torch 2.6.0+cu124
bsakaguchi commented 2026-10-26 21:10:32 +00:00
Author
Owner
Copy link

@LyricZhao commented on GitHub (Feb 27, 2025):

Fixed in main, commit: github.com/deepseek-ai/DeepGEMM@ca13ce0fab and github.com/deepseek-ai/DeepGEMM@6da94d2d36

The reason is wrong TMA store block size when shape_m is smaller than BLOCK_M. So currently we restrict shape_m <= BLOCK_M or shape_m % BLOCK_M == 0 if for masked GEMM with more than one group.

Thanks for feedback!

<!-- gh-comment-id:2687457214 --> @LyricZhao commented on GitHub (Feb 27, 2025): Fixed in main, commit: https://github.com/deepseek-ai/DeepGEMM/commit/ca13ce0fab704c5bf9a4fc19547402b2d10e2f18 and https://github.com/deepseek-ai/DeepGEMM/commit/6da94d2d366b022d5251cee28b4422ecf4d95834 The reason is wrong TMA store block size when `shape_m` is smaller than `BLOCK_M`. So currently we restrict `shape_m <= BLOCK_M or shape_m % BLOCK_M == 0` if for masked GEMM with more than one group. Thanks for feedback!
bsakaguchi commented 2026-10-26 21:10:33 +00:00
Author
Owner
Copy link

@CtfGo commented on GitHub (Feb 28, 2025):

It looks well now.

Are these limitations necessary to achieve good performance？can we remove them to support any m ? for example, by passing extra params that can indicate batch size or data address of each group, like exclusive_sum.

<!-- gh-comment-id:2689816519 --> @CtfGo commented on GitHub (Feb 28, 2025): It looks well now. Are these limitations necessary to achieve good performance？can we remove them to support any `m` ? for example, by passing extra params that can indicate batch size or data address of each group, like exclusive_sum.
bsakaguchi commented 2026-10-26 21:10:33 +00:00
Author
Owner
Copy link

@LyricZhao commented on GitHub (Feb 28, 2025):

Not necessary, they are TMA limitations. If we are using 2D [num_groups * m, k] TMA store, I guess there's no way to bypass the limitation (one group may overwrite to the next group).

A possible solution is to use 3D TMA [num_groups, m, k], but I don't think adding such support for some rare cases is good.

<!-- gh-comment-id:2689828829 --> @LyricZhao commented on GitHub (Feb 28, 2025): Not necessary, they are TMA limitations. If we are using 2D `[num_groups * m, k]` TMA store, I guess there's no way to bypass the limitation (one group may overwrite to the next group). A possible solution is to use 3D TMA `[num_groups, m, k]`, but I don't think adding such support for some rare cases is good.
bsakaguchi commented 2026-10-26 21:10:34 +00:00
Author
Owner
Copy link

@CtfGo commented on GitHub (Mar 1, 2025):

Got it, thanks for your explanation.

<!-- gh-comment-id:2691872424 --> @CtfGo commented on GitHub (Mar 1, 2025): Got it, thanks for your explanation.
Sign in to join this conversation.
No Branch/Tag specified
Branches Tags
main
public-sync-20261014
0930-release
nv_dev
fix-mega-moe-fence
public-release-2609
public-release-2607
public-release-2606
mega-update
public-release-260416
ko3n1g/ci/build-wheels
nv_dev_f8e8fb5
nv_dev_950d31a
nv_dev_67fc648
nv_dev_c491439
nv_dev_4ff3f54
v2.1.1.post3
v2.1.1.post2
v2.1.1.post1
v2.1.1
v2.1.0
Labels
Clear labels   
pull-request
Mirrored from GitHub Pull Request

No labels
pull-request
Milestone
Clear milestone
No items
No milestone
Projects
Clear projects
No items
No project
Assignees
Clear assignees
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepGEMM#14
No description provided.