[GH-ISSUE #125] [Question] Do you have any plans to improve grouped wgrad gemm? #74

New issue

Open

opened 2026-10-26 21:11:56 +00:00 by bsakaguchi · 1 comment

bsakaguchi commented

2026-10-26 21:11:56 +00:00

Owner

Originally created by @shcho1118 on GitHub (Jul 9, 2025).
Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/125

First of all, thank you for sharing your great implementation.
We are experimenting with using some non-gemm kernels (e.g. quantize, transpose, permute, 128-align per group) with DeepGEMM to train our moe models.
In practice, we observed significant speedups over TransformerEngine GroupedLinear for fprop and dgrad, but negligible performance gains for wgrad.
It seems that there is something that can be optimized, since it launches a non-grouped version of wgrad for each group separately.
TransformerEngine has a similar implementation, but at least it utilizes multi-stream.
Do you have any plans or hints to improve these?

Originally created by @shcho1118 on GitHub (Jul 9, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/125 First of all, thank you for sharing your great implementation. We are experimenting with using some non-gemm kernels (e.g. quantize, transpose, permute, 128-align per group) with DeepGEMM to train our moe models. In practice, we observed significant speedups over [TransformerEngine GroupedLinear](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html#transformer_engine.pytorch.GroupedLinear) for fprop and dgrad, but negligible performance gains for wgrad. It seems that there is something that can be optimized, since it launches a non-grouped version of wgrad for each group separately. TransformerEngine has a similar implementation, but at least it utilizes [multi-stream](https://github.com/NVIDIA/TransformerEngine/blob/f05f12c974b37c5bd8dfca3d2d294be53b66abfa/transformer_engine/common/gemm/cublaslt_gemm.cu#L629). Do you have any plans or hints to improve these?

bsakaguchi commented

2026-10-26 21:11:57 +00:00

Author

Owner

@LyricZhao commented on GitHub (Jul 21, 2025):

Sorry for my late reply.

2 weeks ago, the wgrad kernel utilizes 2 streams. And we believe the performance is good enough (memory bound) and has similar perf with cuBLAS.

For now, we currently remove the SM90 impl on the latest commit (but we have a fast impl for SM100 by NVIDIA). Later, @zheanxu will follow this up and push a single-stream single-kernel (even for grouped GEMMs) SM90 impl for wgrad.

@LyricZhao commented on GitHub (Jul 21, 2025): Sorry for my late reply. 2 weeks ago, the wgrad kernel utilizes 2 streams. And we believe the performance is good enough (memory bound) and has similar perf with cuBLAS. For now, we currently remove the SM90 impl on the latest commit (but we have a fast impl for SM100 by NVIDIA). Later, @zheanxu will follow this up and push a single-stream single-kernel (even for grouped GEMMs) SM90 impl for wgrad.

bsakaguchi referenced this issue

2026-10-26 21:15:07 +00:00

[PR #74] [MERGED] Performance: Larger BlockTile optimizations enable 1470+ TF FP8 on the "H800"-SXM #231