[PR #47] [CLOSED] Add kernel‘s config cache and improve TMA alignment #218

New issue

Closed

opened 2026-10-26 21:14:59 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:14:59 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/47
Author: @DXHPC
Created: 3/7/2025
Status: ❌ Closed

Base: main ← Head: dev_config_cache

📝 Commits (3)

ba142b8 config cache and tensor alignment fix
0b5d353 tensor alignment fix
de1603d fix typo

📊 Changes

4 files changed (+122 additions, -92 deletions)

View changed files

➕ deep_gemm/jit_kernels/config.py (+99 -0)
📝 deep_gemm/jit_kernels/gemm.py (+2 -75)
📝 deep_gemm/jit_kernels/m_grouped_gemm.py (+4 -3)
📝 deep_gemm/jit_kernels/utils.py (+17 -14)

📄 Description

When testing the function deep_gemm.gemm_fp8_fp8_bf16_nt, we observed that for certain examples, especially matrices with small shapes, the overhead outside of the GEMM kernel's execution was non-negligible. By using PyTorch's profiler, we found that the execution of the functions get_col_major_tma_aligned_tensor and get_best_configs introduces significant overhead.
截屏2025-03-07 11 53 19

We tried to perform memory alignment in advance by calling the function get_col_major_tma_aligned_tensor before invoking the function deep_gemm.gemm_fp8_fp8_bf16_nt. But we still noticed that the alignment-related operations were triggered during the execution of deep_gemm.gemm_fp8_fp8_bf16_nt. Additionally, the function get_best_configs redundantly executes for the same inputs, yielding the same results, which introduces unnecessary overhead.
First, we fixed the improper handling of 2D matrices in function get_col_major_tma_aligned_tensor. Additionally, we added a cache to obtain the optimal configuration before the GEMM kernel launches, which reduces overhead for inputs of the same size.

_{🔄 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/47 **Author:** [@DXHPC](https://github.com/DXHPC) **Created:** 3/7/2025 **Status:** ❌ Closed **Base:** `main` ← **Head:** `dev_config_cache` --- ### 📝 Commits (3) - [`ba142b8`](https://github.com/deepseek-ai/DeepGEMM/commit/ba142b891ab10e4cb757434af06add7691822dcc) config cache and tensor alignment fix - [`0b5d353`](https://github.com/deepseek-ai/DeepGEMM/commit/0b5d353dba2ab84e02170c6afe6aeee95cfc3efb) tensor alignment fix - [`de1603d`](https://github.com/deepseek-ai/DeepGEMM/commit/de1603d5e4aee83b05cec3048ce0fda4f3efeb66) fix typo ### 📊 Changes **4 files changed** (+122 additions, -92 deletions) <details> <summary>View changed files</summary> ➕ `deep_gemm/jit_kernels/config.py` (+99 -0) 📝 `deep_gemm/jit_kernels/gemm.py` (+2 -75) 📝 `deep_gemm/jit_kernels/m_grouped_gemm.py` (+4 -3) 📝 `deep_gemm/jit_kernels/utils.py` (+17 -14) </details> ### 📄 Description When testing the function **deep_gemm.gemm_fp8_fp8_bf16_nt**, we observed that for certain examples, especially matrices with small shapes, the overhead outside of the GEMM kernel's execution was non-negligible. By using PyTorch's profiler, we found that the execution of the functions **get_col_major_tma_aligned_tensor** and **get_best_configs** introduces significant overhead. <img width="1994" alt="截屏2025-03-07 11 53 19" src="https://github.com/user-attachments/assets/a8924d21-8504-40a6-a809-16da326dd8bb" /> We tried to perform memory alignment in advance by calling the function **get_col_major_tma_aligned_tensor** before invoking the function **deep_gemm.gemm_fp8_fp8_bf16_nt**. But we still noticed that the alignment-related operations were triggered during the execution of **deep_gemm.gemm_fp8_fp8_bf16_nt**. Additionally, the function **get_best_configs** redundantly executes for the same inputs, yielding the same results, which introduces unnecessary overhead. First, we fixed the improper handling of 2D matrices in function **get_col_major_tma_aligned_tensor**. Additionally, we added a cache to obtain the optimal configuration before the GEMM kernel launches, which reduces overhead for inputs of the same size. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 21:14:59 +00:00