[GH-ISSUE #85] Why is get_m_alignment_for_contiguous_layout fixed at 128? Can it be reduced to 64? #49

Closed
opened 2026-10-26 21:11:14 +00:00 by bsakaguchi · 9 comments
Owner

Originally created by @wangxiyu191 on GitHub (Apr 11, 2025).
Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/85

The function get_m_alignment_for_contiguous_layout appears to enforce an alignment of 128 for the M dimension. But both normal GEMM and group GEMM (masked layout) support BLOCK_M=64.
In the case of grouped_contiguous, would it be valid to reduce both BLOCK_M and get_m_alignment_for_contiguous_layout to 64 in order to more efficiently support smaller values of m_per_group? When m_per_group is only 64, using BLOCK_M=64(a single math warp group) appears to deliver better performance.

Originally created by @wangxiyu191 on GitHub (Apr 11, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/85 The function get_m_alignment_for_contiguous_layout appears to enforce an alignment of 128 for the M dimension. But both normal GEMM and group GEMM (masked layout) support BLOCK_M=64. In the case of grouped_contiguous, would it be valid to reduce both BLOCK_M and get_m_alignment_for_contiguous_layout to 64 in order to more efficiently support smaller values of m_per_group? When m_per_group is only 64, using BLOCK_M=64(a single math warp group) appears to deliver better performance.
Author
Owner

@LyricZhao commented on GitHub (Apr 14, 2025):

Normally for training/prefill inference engine, only contiguous layout is used, for such cases, the shape of M is normal ~4K, so larger BLOCK_M delivers better performance, and BLOCK_M=64 rarely used in real cases.

<!-- gh-comment-id:2800759624 --> @LyricZhao commented on GitHub (Apr 14, 2025): Normally for training/prefill inference engine, only contiguous layout is used, for such cases, the shape of M is normal ~4K, so larger `BLOCK_M` delivers better performance, and `BLOCK_M=64` rarely used in real cases.
Author
Owner

@wangxiyu191 commented on GitHub (Apr 14, 2025):

Thank you very much for the explanation; this reason is quite reasonable. So, can I understand it this way: Using BLOCK_M=64 in continuous layout is valid, but such usage scenarios are rare.
The background is that when decoding, there might be cases where expert selection within a single batch is heavily skewed (e.g., all prompts in the batch are identical). In such cases, it may be necessary to configure a large m_max (e.g., m_max=global_batch_size) for masked layout. If there are also many experts on a single GPU, this will result in a large lhs matrix. Therefore, I am trying to use the contiguous layout to reduce the size of the lhs matrix. Additionally, I would like to ask if there are any suggestions for this case? Thank you again.

<!-- gh-comment-id:2800954127 --> @wangxiyu191 commented on GitHub (Apr 14, 2025): Thank you very much for the explanation; this reason is quite reasonable. So, can I understand it this way: Using BLOCK_M=64 in continuous layout is valid, but such usage scenarios are rare. The background is that when decoding, there might be cases where expert selection within a single batch is heavily skewed (e.g., all prompts in the batch are identical). In such cases, it may be necessary to configure a large m_max (e.g., m_max=global_batch_size) for masked layout. If there are also many experts on a single GPU, this will result in a large lhs matrix. Therefore, I am trying to use the contiguous layout to reduce the size of the lhs matrix. Additionally, I would like to ask if there are any suggestions for this case? Thank you again.
Author
Owner

@LyricZhao commented on GitHub (Apr 14, 2025):

So, can I understand it this way: Using BLOCK_M=64 in continuous layout is valid, but such usage scenarios are rare.

Yes.

If there are also many experts on a single GPU, this will result in a large lhs matrix.

You mean, for example, expert 0 has 4K tokens, and the other experts have like 0-64 tokens? In that case, I guess you should blame the expert balance algorithm or gate?

<!-- gh-comment-id:2801138761 --> @LyricZhao commented on GitHub (Apr 14, 2025): > So, can I understand it this way: Using BLOCK_M=64 in continuous layout is valid, but such usage scenarios are rare. Yes. > If there are also many experts on a single GPU, this will result in a large lhs matrix. You mean, for example, expert 0 has 4K tokens, and the other experts have like 0-64 tokens? In that case, I guess you should blame the expert balance algorithm or gate?
Author
Owner

@LyricZhao commented on GitHub (Apr 14, 2025):

I think such cases should never happen, if bsz is small, use masked layout, else use contiguous layout and ensure experts are balanced.

<!-- gh-comment-id:2801141203 --> @LyricZhao commented on GitHub (Apr 14, 2025): I think such cases should never happen, if bsz is small, use masked layout, else use contiguous layout and ensure experts are balanced.
Author
Owner

@LyricZhao commented on GitHub (Apr 14, 2025):

expert 0 has 4K tokens, and the other experts have like 0-64 tokens

Even for such cases, "other experts" are memory-bound, I guess the BLOCK_M doesn't really matter.

<!-- gh-comment-id:2801143442 --> @LyricZhao commented on GitHub (Apr 14, 2025): > expert 0 has 4K tokens, and the other experts have like 0-64 tokens Even for such cases, "other experts" are memory-bound, I guess the `BLOCK_M` doesn't really matter.
Author
Owner

@wangxiyu191 commented on GitHub (Apr 15, 2025):

You mean, for example, expert 0 has 4K tokens, and the other experts have like 0-64 tokens? In that case, I guess you should blame the expert balance algorithm or gate?

Yes. Maybe I should blame the benchmark. I agree this case is really rare.

Even for such cases, "other experts" are memory-bound, I guess the BLOCK_M doesn't really matter.

It is true when using H800. But when using H20(~300 FP8 FLOPS and ~4TB/s memory bandwidth), experts with 64 tokens are still compute-bound:

Testing grouped contiguous GEMM with BLOCK_M=64:
> Performance (num_groups=16, m_per_group=  64, n=4096, k=7168):  234 us | throughput:  257 TFLOPS, 2074 GB/s
Testing grouped contiguous GEMM with BLOCK_M=128:
> Performance (num_groups=16, m_per_group= 128, n=4096, k=7168):  452 us | throughput:  266 TFLOPS, 1109 GB/s
<!-- gh-comment-id:2803525276 --> @wangxiyu191 commented on GitHub (Apr 15, 2025): > You mean, for example, expert 0 has 4K tokens, and the other experts have like 0-64 tokens? In that case, I guess you should blame the expert balance algorithm or gate? Yes. Maybe I should blame the benchmark. I agree this case is really rare. > Even for such cases, "other experts" are memory-bound, I guess the BLOCK_M doesn't really matter. It is true when using H800. But when using H20(~300 FP8 FLOPS and ~4TB/s memory bandwidth), experts with 64 tokens are still compute-bound: ``` Testing grouped contiguous GEMM with BLOCK_M=64: > Performance (num_groups=16, m_per_group= 64, n=4096, k=7168): 234 us | throughput: 257 TFLOPS, 2074 GB/s Testing grouped contiguous GEMM with BLOCK_M=128: > Performance (num_groups=16, m_per_group= 128, n=4096, k=7168): 452 us | throughput: 266 TFLOPS, 1109 GB/s ```
Author
Owner

@LyricZhao commented on GitHub (Apr 15, 2025):

But when using H20(~300 FP8 FLOPS and ~4TB/s memory bandwidth), experts with 64 tokens are still compute-bound:

Oh, for H20, the answer is yes. But maybe a workaround is to skip the second warpgroup computation (see the TODO at https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/include/deep_gemm/fp8_gemm.cuh#L297), rather than changing the alignment, changing the alignment will affect other kernels in the training/inference frameworks (at least at deepseek).

Will do the TODO later.

<!-- gh-comment-id:2803534652 --> @LyricZhao commented on GitHub (Apr 15, 2025): > But when using H20(~300 FP8 FLOPS and ~4TB/s memory bandwidth), experts with 64 tokens are still compute-bound: Oh, for H20, the answer is yes. But maybe a workaround is to skip the second warpgroup computation (see the TODO at https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/include/deep_gemm/fp8_gemm.cuh#L297), rather than changing the alignment, changing the alignment will affect other kernels in the training/inference frameworks (at least at deepseek). Will do the TODO later.
Author
Owner

@wangxiyu191 commented on GitHub (Apr 15, 2025):

You are right. Keeping BLOCK_M=128 and skipping the second warp group when necessary is a better approach.
My original idea was to support adaptive m_alignment/BLOCK_M (including 64 and 128) for contiguous layout. However, this would require changes to the signature of get_m_alignment_for_contiguous_layout function (add expected_m as an argument).
Thank you again. Looking forward to this optimization.

<!-- gh-comment-id:2803945215 --> @wangxiyu191 commented on GitHub (Apr 15, 2025): You are right. Keeping BLOCK_M=128 and skipping the second warp group when necessary is a better approach. My original idea was to support adaptive m_alignment/BLOCK_M (including 64 and 128) for contiguous layout. However, this would require changes to the signature of get_m_alignment_for_contiguous_layout function (add expected_m as an argument). Thank you again. Looking forward to this optimization.
Author
Owner

@LyricZhao commented on GitHub (Apr 15, 2025):

cc @yukuai26, please do it later.

<!-- gh-comment-id:2804099151 --> @LyricZhao commented on GitHub (Apr 15, 2025): cc @yukuai26, please do it later.
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
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#49
No description provided.