[PR #453] Fix SM120 BF16 GEMM shared-memory reuse ordering #458

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/453
Author: @Sunt-ing
Created: 9/21/2026
Status: 🔄 Open

Base: nv_dev ← Head: fix/sm120-bf16-stage-reuse

📝 Commits (2)

6092535 Fix SM120 BF16 GEMM shared-memory reuse ordering
7516924 Clarify the BF16 repeatability test limitation

📊 Changes

2 files changed (+23 additions, -0 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm120_bf16_gemm.cuh (+2 -0)
📝 tests/test_bf16.py (+21 -0)

📄 Description

Summary

Fixes intermittent silent data corruption in SM120 BF16 GEMM caused by a generic-to-async proxy WAR (Write-After-Read) hazard.

Problem

Consumer warps read smem tiles via the generic proxy and signal the empty barrier without an async proxy fence. Producer TMA loads (async proxy) can subsequently overwrite the smem stage before generic reads complete across all consumer threads.

Fix

Add cutlass::arch::fence_view_async_shared() before the empty-barrier arrival. Placed outside the lane-0 check so all consumer threads issue the fence for their own smem reads.

Targets nv_dev (based on 572557e7ae9ad5331b81a1c250f141fba2c57962).

Reproduction and validation

Environment: RTX 5090 (SM120), Driver 590.48.01, PyTorch 2.13.0+cu130.

Added test_bf16_repeatability in tests/test_bf16.py ([4096, 1024] × [1024, 896], seed 233, 6,000 iterations against initial output).

To run the standalone regression:

cd tests
python -c "from test_bf16 import test_bf16_repeatability; test_bf16_repeatability()"

Note: The corruption is an intermittent race condition; a single baseline run passing does not guarantee correctness.

Repeatability (A/B/A Test)

Original: 2 / 6,000 mismatches
Patched: 0 / 6,000 mismatches
Reverted: 2 / 6,000 mismatches

Performance Overhead

CUDA graph timing on [4096, 1024] × [1024, 896] (median per-call time across 9 batches of 200 calls):

Version	Median Latency
Original	50.67 μs
Patched	52.01 μs
Reverted	50.68 μs

Adds ~2.6% latency overhead on this shape from the fence.

Supplemental End-to-End Inference Validation

Internal TP2 inference test (weights/prompts omitted; public verification relies on the standalone GEMM test above).

Setup: 32,640-token prompt, 89 decode steps -> re-prefill -> 16 decode steps vs. continuous generation. Unified prefill/decode attention was held constant; CUDA graphs were disabled.
Original / Reverted: 6 / 12 runs with divergent logprobs.
Patched: 0 / 12 runs diverged (token IDs and logprobs match strictly).

#389: Addressed the same class of generic-to-async proxy hazard in SM100 MegaMoE combine.
#447: The consumer release path (reviewed at 1f78b85a1fb974cc69ea63e864d817c603688cab) does not yet include this fence.

_{🔄 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/453 **Author:** [@Sunt-ing](https://github.com/Sunt-ing) **Created:** 9/21/2026 **Status:** 🔄 Open **Base:** `nv_dev` ← **Head:** `fix/sm120-bf16-stage-reuse` --- ### 📝 Commits (2) - [`6092535`](https://github.com/deepseek-ai/DeepGEMM/commit/6092535b1aeae7e4db627bdd7ec052423a86366b) Fix SM120 BF16 GEMM shared-memory reuse ordering - [`7516924`](https://github.com/deepseek-ai/DeepGEMM/commit/751692402bff8db3ae892adaabd8b410d50ebddf) Clarify the BF16 repeatability test limitation ### 📊 Changes **2 files changed** (+23 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm120_bf16_gemm.cuh` (+2 -0) 📝 `tests/test_bf16.py` (+21 -0) </details> ### 📄 Description ## Summary Fixes intermittent silent data corruption in SM120 BF16 GEMM caused by a generic-to-async proxy WAR (Write-After-Read) hazard. ### Problem Consumer warps read smem tiles via the generic proxy and signal the empty barrier without an async proxy fence. Producer TMA loads (async proxy) can subsequently overwrite the smem stage before generic reads complete across all consumer threads. ### Fix Add `cutlass::arch::fence_view_async_shared()` before the empty-barrier arrival. Placed outside the lane-0 check so all consumer threads issue the fence for their own smem reads. Targets `nv_dev` (based on `572557e7ae9ad5331b81a1c250f141fba2c57962`). ## Reproduction and validation Environment: RTX 5090 (SM120), Driver 590.48.01, PyTorch 2.13.0+cu130. Added `test_bf16_repeatability` in `tests/test_bf16.py` (`[4096, 1024] × [1024, 896]`, seed 233, 6,000 iterations against initial output). To run the standalone regression: ```bash cd tests python -c "from test_bf16 import test_bf16_repeatability; test_bf16_repeatability()" ``` *Note: The corruption is an intermittent race condition; a single baseline run passing does not guarantee correctness.* ### Repeatability (A/B/A Test) - **Original:** 2 / 6,000 mismatches - **Patched:** 0 / 6,000 mismatches - **Reverted:** 2 / 6,000 mismatches ### Performance Overhead CUDA graph timing on `[4096, 1024] × [1024, 896]` (median per-call time across 9 batches of 200 calls): | Version | Median Latency | | :--- | ---: | | Original | 50.67 μs | | Patched | 52.01 μs | | Reverted | 50.68 μs | Adds ~2.6% latency overhead on this shape from the fence. ### Supplemental End-to-End Inference Validation *Internal TP2 inference test (weights/prompts omitted; public verification relies on the standalone GEMM test above).* - **Setup:** 32,640-token prompt, 89 decode steps -> re-prefill -> 16 decode steps vs. continuous generation. Unified prefill/decode attention was held constant; CUDA graphs were disabled. - **Original / Reverted:** 6 / 12 runs with divergent logprobs. - **Patched:** 0 / 12 runs diverged (token IDs and logprobs match strictly). ## Related work - #389: Addressed the same class of generic-to-async proxy hazard in SM100 MegaMoE combine. - #447: The consumer release path (reviewed at `1f78b85a1fb974cc69ea63e864d817c603688cab`) does not yet include this fence. --- <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:05 +00:00

bsakaguchi referenced this issue

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

[PR #459] [MERGED] Sparse MQA logits follow-ups for nv_dev: 16-byte paged KV stride, workspace reuse under graph capture #463

bsakaguchi referenced this issue

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

[PR #458] [CLOSED] Merge main (Public Release 26/09) into nv_dev, with sparse MQA-logits integration fixes #465