[PR #431] fix: port the #343 tensormap drain to the SM120 k_grouped NT contiguous kernel #441

Open
opened 2026-10-26 21:17:49 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/431
Author: @aganhui
Created: 9/9/2026
Status: 🔄 Open

Base: nv_dev ← Head: fix/sm120-kgrouped-tensormap-drain


📝 Commits (1)

  • f115a07 fix: drain in-flight TMA loads before publishing grouped tensormaps on SM120 (#430)

📊 Changes

1 file changed (+7 additions, -0 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh (+7 -0)

📄 Description

Fixes #430.

The SM120 1D1D kernel (sm120_fp8_fp4_gemm_1d1d.cuh, introduced in #324) was copied from sm90_fp8_gemm_1d1d.cuh twelve days before #343 fixed the in-flight tensormap race there. Since the copy lives at a different file path, the #343 fix never reached it. This ports the same drain to the SM120 group-switch path so that in-flight TMA loads complete before the grouped tensormaps are published to GMEM.

Validation on RTX 5090 D (SM120, CUDA 13.0, PyTorch 2.11.0+cu130, nv_dev @ 2642b32):

  • Before: test_k_grouped_gemm_contiguous fails at the first shape in 13/14 runs; diff 0.0015–0.006, non-deterministic on identical inputs
  • After: diff 0.000105, bit-identical across runs; full k_grouped + m_grouped suites pass; 100× soak clean
  • Perf A/B (same seeds/data): +0.9% on the sentinel shape (worst-case group-switch density), +0.26% on typical EP shapes — the same trade-off SM90 accepted in #343

Note on __syncwarp: same construct as the merged SM90 fix — for the prior discussion of its validity see #343.

sm120_bf16_gemm.cuh has the same un-drained pattern, but we could not reproduce failures there (full suite + 50 targeted stress runs pass), so this PR leaves it alone — details in the linked issue.

Evidence and raw logs: https://gist.github.com/aganhui/3b374ffb3d6d13b52d5688780fafac28


🔄 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/431 **Author:** [@aganhui](https://github.com/aganhui) **Created:** 9/9/2026 **Status:** 🔄 Open **Base:** `nv_dev` ← **Head:** `fix/sm120-kgrouped-tensormap-drain` --- ### 📝 Commits (1) - [`f115a07`](https://github.com/deepseek-ai/DeepGEMM/commit/f115a07e4a443b7403e9d702fc908be5dc49f0ac) fix: drain in-flight TMA loads before publishing grouped tensormaps on SM120 (#430) ### 📊 Changes **1 file changed** (+7 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh` (+7 -0) </details> ### 📄 Description Fixes #430. The SM120 1D1D kernel (`sm120_fp8_fp4_gemm_1d1d.cuh`, introduced in #324) was copied from `sm90_fp8_gemm_1d1d.cuh` twelve days before #343 fixed the in-flight tensormap race there. Since the copy lives at a different file path, the #343 fix never reached it. This ports the same drain to the SM120 group-switch path so that in-flight TMA loads complete before the grouped tensormaps are published to GMEM. Validation on RTX 5090 D (SM120, CUDA 13.0, PyTorch 2.11.0+cu130, nv_dev @ 2642b32): - Before: `test_k_grouped_gemm_contiguous` fails at the first shape in 13/14 runs; diff 0.0015–0.006, non-deterministic on identical inputs - After: diff **0.000105, bit-identical across runs**; full k_grouped + m_grouped suites pass; 100× soak clean - Perf A/B (same seeds/data): +0.9% on the sentinel shape (worst-case group-switch density), +0.26% on typical EP shapes — the same trade-off SM90 accepted in #343 Note on `__syncwarp`: same construct as the merged SM90 fix — for the prior discussion of its validity see #343. `sm120_bf16_gemm.cuh` has the same un-drained pattern, but we could not reproduce failures there (full suite + 50 targeted stress runs pass), so this PR leaves it alone — details in the linked issue. Evidence and raw logs: https://gist.github.com/aganhui/3b374ffb3d6d13b52d5688780fafac28 --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
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#441
No description provided.