[PR #193] [MERGED] Fix multicast bug and optimize masked GEMM #285

New issue

Closed

opened 2026-10-26 21:15:46 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/193
Author: @yukuai26
Created: 9/12/2025
Status: ✅ Merged
Merged: 9/12/2025
Merged by: @LyricZhao

Base: main ← Head: multicast-fixed

📝 Commits (2)

a43a750 Fix multicast bug and profile masked GEMM
821ad1e Updates and lint

📊 Changes

3 files changed (+10 additions, -5 deletions)

View changed files

📝 csrc/jit_kernels/heuristics/common.hpp (+6 -3)
📝 csrc/jit_kernels/heuristics/sm100.hpp (+1 -1)
📝 csrc/jit_kernels/heuristics/sm90.hpp (+3 -1)

📄 Description

In the previous code, the variable used for determining whether to use multicast was incorrectly applied, the variable for judging A multicast capability was used to determine B multicast capability. Fortunately, this error does not affect correctness, and it also does not impact performance in the case of non-masked GEMM.

Additionally, I added B multicast logic for masked GEMM and modified the parameters. Masked GEMM achieved a 10%-20% speedup in H800.

Before:

Testing m-grouped masked GEMM:
 > Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D):   79 us |  778 TFLOPS |  578 GB/s
 > Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D):   45 us |  727 TFLOPS |  732 GB/s
 > Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D):   92 us |  744 TFLOPS |  839 GB/s
 > Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D):   44 us |  660 TFLOPS | 1041 GB/s
 > Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D):   95 us |  688 TFLOPS | 1425 GB/s
 > Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D):   46 us |  573 TFLOPS | 1585 GB/s

After:

Testing m-grouped masked GEMM:
 > Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D):   67 us |  920 TFLOPS |  683 GB/s
 > Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D):   35 us |  931 TFLOPS |  937 GB/s
 > Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D):   77 us |  879 TFLOPS |  992 GB/s
 > Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D):   38 us |  751 TFLOPS | 1184 GB/s
 > Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D):   82 us |  798 TFLOPS | 1652 GB/s
 > Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D):   43 us |  613 TFLOPS | 1698 GB/s

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/193 **Author:** [@yukuai26](https://github.com/yukuai26) **Created:** 9/12/2025 **Status:** ✅ Merged **Merged:** 9/12/2025 **Merged by:** [@LyricZhao](https://github.com/LyricZhao) **Base:** `main` ← **Head:** `multicast-fixed` --- ### 📝 Commits (2) - [`a43a750`](https://github.com/deepseek-ai/DeepGEMM/commit/a43a7500b5bbd8e0c179cfcf525c41efa2cabbaf) Fix multicast bug and profile masked GEMM - [`821ad1e`](https://github.com/deepseek-ai/DeepGEMM/commit/821ad1e49cc4fd6677f44ed2ec51d526d8e435a3) Updates and lint ### 📊 Changes **3 files changed** (+10 additions, -5 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/heuristics/common.hpp` (+6 -3) 📝 `csrc/jit_kernels/heuristics/sm100.hpp` (+1 -1) 📝 `csrc/jit_kernels/heuristics/sm90.hpp` (+3 -1) </details> ### 📄 Description In the previous code, the variable used for determining whether to use multicast was incorrectly applied, the variable for judging A multicast capability was used to determine B multicast capability. Fortunately, this error does not affect correctness, and it also does not impact performance in the case of non-masked GEMM. Additionally, I added B multicast logic for masked GEMM and modified the parameters. Masked GEMM achieved a 10%-20% speedup in H800. Before: ``` Testing m-grouped masked GEMM: > Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D): 79 us | 778 TFLOPS | 578 GB/s > Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D): 45 us | 727 TFLOPS | 732 GB/s > Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D): 92 us | 744 TFLOPS | 839 GB/s > Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D): 44 us | 660 TFLOPS | 1041 GB/s > Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D): 95 us | 688 TFLOPS | 1425 GB/s > Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D): 46 us | 573 TFLOPS | 1585 GB/s ``` After: ``` Testing m-grouped masked GEMM: > Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D): 67 us | 920 TFLOPS | 683 GB/s > Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D): 35 us | 931 TFLOPS | 937 GB/s > Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D): 77 us | 879 TFLOPS | 992 GB/s > Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D): 38 us | 751 TFLOPS | 1184 GB/s > Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D): 82 us | 798 TFLOPS | 1652 GB/s > Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D): 43 us | 613 TFLOPS | 1698 GB/s ``` --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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