[PR #298] feat: support bf16 output and plain TMA writes in k_grouped_gemm on SM90; #350

New issue

Open

opened 2026-10-26 21:16:34 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:16:34 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/298
Author: @fedorovgv
Created: 3/26/2026
Status: 🔄 Open

Base: main ← Head: feat/k_grouped_sm90_bfp16

📝 Commits (1)

89ec9f4 add support of bfp16 out and plain tma write in k_grouped_gemm sm90;

📊 Changes

4 files changed (+41 additions, -25 deletions)

View changed files

📝 csrc/apis/gemm.hpp (+2 -2)
📝 csrc/jit_kernels/heuristics/sm90.hpp (+7 -7)
📝 csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp (+6 -4)
📝 deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh (+26 -12)

📄 Description

Add two features to SM90 FP8 1D1D k-grouped GEMM:

Support plain TMA store as an alternative to atomic accumulation, controlled by the presence of the c tensor
Support BF16 output dtype, casting WGMMA FP32 accumulators to BF16 before the TMA store

_{🔄 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/298 **Author:** [@fedorovgv](https://github.com/fedorovgv) **Created:** 3/26/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `feat/k_grouped_sm90_bfp16` --- ### 📝 Commits (1) - [`89ec9f4`](https://github.com/deepseek-ai/DeepGEMM/commit/89ec9f42ea7af8ed191c86e4a20cc425b5c74de6) add support of bfp16 out and plain tma write in k_grouped_gemm sm90; ### 📊 Changes **4 files changed** (+41 additions, -25 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/gemm.hpp` (+2 -2) 📝 `csrc/jit_kernels/heuristics/sm90.hpp` (+7 -7) 📝 `csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp` (+6 -4) 📝 `deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh` (+26 -12) </details> ### 📄 Description Add two features to SM90 FP8 1D1D k-grouped GEMM: - Support plain TMA store as an alternative to atomic accumulation, controlled by the presence of the c tensor - Support BF16 output dtype, casting WGMMA FP32 accumulators to BF16 before the TMA store --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:16:34 +00:00

bsakaguchi referenced this issue

2026-10-26 21:17:06 +00:00

[PR #350] Add shape-aware recommended alignment for SM90 small-M grouped GEMM #383