deepseek-ai/DeepGEMM

Fork 0

mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00

[PR #480] Use the allocated TMEM base in the SM100 HC prenorm GEMM #482

New issue

Open

opened 2026-10-26 21:18:22 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:18:22 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/480
Author: @XFDG
Created: 10/21/2026
Status: 🔄 Open

Base: main ← Head: fix/479-sm100-hc-prenorm-tmem-base

📝 Commits (1)

6563a81 Use the allocated TMEM base in the SM100 HC prenorm GEMM

📊 Changes

1 file changed (+8 additions, -4 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm100_tf32_hc_prenorm_gemm.cuh (+8 -4)

📄 Description

Fixes #479.

sm100_tf32_hc_prenorm_gemm_impl allocates 256 tensor memory columns, but it addresses its cast A stages and its accumulator from column 0 and frees them with free(0, kNumTmemCols). The allocator can return a non-zero base when another CTA on the same SM already holds tensor memory. This kernel uses at most about 200 KB of shared memory, so a small CTA of another kernel, for example one on another stream, can run next to it. The GEMM then writes into that CTA's columns, and tcgen05.dealloc on columns it does not own hits the tcgen05 guardrail trap.

This change reads the base that the allocator wrote after __syncthreads() and adds it to every tensor memory address and to free, as sm100_mega_mhc already does.

Validation on B200 (CUDA 13.1). A test kernel is launched first on one stream with one CTA per SM (148 CTAs). Each CTA allocates 32 or 256 columns (all at base 0), fills them with a pattern and holds them until the GEMM on a second stream has finished, then checks them. In every concurrent run that completed, the GEMM finished while all 148 CTAs were still resident.

case	`main` (`5081f85`)	only `free` fixed (diagnostic)	this PR
test kernel alone, 32 or 256 columns	pass	pass	pass
test kernel holds 32 columns, GEMM m=13 with 16 splits or m=4096 without splits; holds 256 columns, GEMM m=13 with 16 splits	3/3 `unspecified launch failure`	held columns overwritten on every SM that ran a GEMM CTA	pass, GEMM output bitwise equal to a run alone
test kernel keeps rewriting its 32 columns, GEMM m=4096 with 16 splits or m=137, k=28672 without splits	2/2 `unspecified launch failure`	wrong GEMM output (`calc_diff` is NaN)	pass, GEMM output bitwise equal to a run alone

Under cuda-gdb, main stops in sm100_tf32_hc_prenorm_gemm_impl<24u, 7168u, 64u, 32u, 64u, 16u, 128u, 12u, 128u, 128u> with the error from the issue:

CUDA Exception: Warp Tensor Memory Access Check
Tensor Memory column 3152 being dealloced by instruction tcgen05.dealloc is not allocated. Columns allocated are: 192-223.

tests/test_hyperconnection.py passes with and without this change. Kernel times are unchanged within noise: across the 24 shapes of that test, median per-kernel times differ by -8% to +6% in both directions (geometric mean 0.996). Another process shared the GPU during timing, so these numbers are not exact. The SASS of the two configs I compared has 24 and 32 fewer instructions, with the same or one fewer register.

The other SM100 kernels that free from column 0 (sm100_bf16_gemm, sm100_fp8_fp4_gemm_1d1d, sm100_bmk_bnk_mn, sm100_sparse_mqa_logits, sm100_mqa_logits, sm100_bf16_mega_gate, and the two Mega MoE kernels) are not changed here. I can extend this PR if you want them handled the same way. The test kernel and script are about 340 lines; I can share them if that helps.

🤖 Generated with Claude Code

_{🔄 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/480 **Author:** [@XFDG](https://github.com/XFDG) **Created:** 10/21/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/479-sm100-hc-prenorm-tmem-base` --- ### 📝 Commits (1) - [`6563a81`](https://github.com/deepseek-ai/DeepGEMM/commit/6563a81fa6ad4f247199543d4d0ca40ff0e7f9a6) Use the allocated TMEM base in the SM100 HC prenorm GEMM ### 📊 Changes **1 file changed** (+8 additions, -4 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm100_tf32_hc_prenorm_gemm.cuh` (+8 -4) </details> ### 📄 Description Fixes #479. `sm100_tf32_hc_prenorm_gemm_impl` allocates 256 tensor memory columns, but it addresses its cast A stages and its accumulator from column 0 and frees them with `free(0, kNumTmemCols)`. The allocator can return a non-zero base when another CTA on the same SM already holds tensor memory. This kernel uses at most about 200 KB of shared memory, so a small CTA of another kernel, for example one on another stream, can run next to it. The GEMM then writes into that CTA's columns, and `tcgen05.dealloc` on columns it does not own hits the tcgen05 guardrail trap. This change reads the base that the allocator wrote after `__syncthreads()` and adds it to every tensor memory address and to `free`, as `sm100_mega_mhc` already does. Validation on B200 (CUDA 13.1). A test kernel is launched first on one stream with one CTA per SM (148 CTAs). Each CTA allocates 32 or 256 columns (all at base 0), fills them with a pattern and holds them until the GEMM on a second stream has finished, then checks them. In every concurrent run that completed, the GEMM finished while all 148 CTAs were still resident. | case | `main` (5081f85) | only `free` fixed (diagnostic) | this PR | |---|---|---|---| | test kernel alone, 32 or 256 columns | pass | pass | pass | | test kernel holds 32 columns, GEMM m=13 with 16 splits or m=4096 without splits; holds 256 columns, GEMM m=13 with 16 splits | 3/3 `unspecified launch failure` | held columns overwritten on every SM that ran a GEMM CTA | pass, GEMM output bitwise equal to a run alone | | test kernel keeps rewriting its 32 columns, GEMM m=4096 with 16 splits or m=137, k=28672 without splits | 2/2 `unspecified launch failure` | wrong GEMM output (`calc_diff` is NaN) | pass, GEMM output bitwise equal to a run alone | Under cuda-gdb, `main` stops in `sm100_tf32_hc_prenorm_gemm_impl<24u, 7168u, 64u, 32u, 64u, 16u, 128u, 12u, 128u, 128u>` with the error from the issue: ```text CUDA Exception: Warp Tensor Memory Access Check Tensor Memory column 3152 being dealloced by instruction tcgen05.dealloc is not allocated. Columns allocated are: 192-223. ``` `tests/test_hyperconnection.py` passes with and without this change. Kernel times are unchanged within noise: across the 24 shapes of that test, median per-kernel times differ by -8% to +6% in both directions (geometric mean 0.996). Another process shared the GPU during timing, so these numbers are not exact. The SASS of the two configs I compared has 24 and 32 fewer instructions, with the same or one fewer register. The other SM100 kernels that free from column 0 (`sm100_bf16_gemm`, `sm100_fp8_fp4_gemm_1d1d`, `sm100_bmk_bnk_mn`, `sm100_sparse_mqa_logits`, `sm100_mqa_logits`, `sm100_bf16_mega_gate`, and the two Mega MoE kernels) are not changed here. I can extend this PR if you want them handled the same way. The test kernel and script are about 340 lines; I can share them if that helps. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- <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:18:22 +00:00

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#482

No description provided.

Rows
Columns