[PR #343] [MERGED] Fix a race condition in contiguous k-grouped GEMM where in-flight tensormaps are updated in-place #377

New issue

Closed

opened 2026-10-26 21:17:01 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/343
Author: @dfyz
Created: 5/29/2026
Status: ✅ Merged
Merged: 7/6/2026
Merged by: @zheanxu

Base: main ← Head: main

📝 Commits (1)

a6a9339 Fix a race condition where in-flight tensormaps are updated in-place

📊 Changes

2 files changed (+10 additions, -1 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh (+9 -0)
📝 tests/generators.py (+1 -1)

📄 Description

The weight gradient kernel for SM90 modifies the A/B tensormaps stored in GMEM when the scheduler decides to switch groups, but it doesn't make sure that previous TMA loads have already finished reading the tensormaps at that point. A TMA load reading from a tensormap being modified concurrently corrupts the output for obvious reasons.

A reliable way to reproduce this corruption (at least on an H200) is to run the test_k_grouped_gemm_contiguous test a couple of times with the shape added in this PR. Eventually the corruption will be large enough to make the test fail:

[...]
Testing k-grouped contiguous GEMM:
 > Perf (num_groups= 4, m= 4096, n= 7168, k=34304, gran_k=128): 1966 us | 1025 TFLOPS |  681 GB/s
 > Perf (num_groups= 4, m= 7168, n= 2048, k=32896, gran_k=128):  970 us |  996 TFLOPS |  807 GB/s
Traceback (most recent call last):
  File "/persistent/DeepGEMM/tests/test_fp8_fp4.py", line 222, in <module>
    test_k_grouped_gemm_contiguous()
  File "/persistent/DeepGEMM/tests/test_fp8_fp4.py", line 194, in test_k_grouped_gemm_contiguous
    assert diff < 0.001, f'{m=}, {n=}, {k=}, {ks=}, {diff:.5f}'
           ^^^^^^^^^^^^
AssertionError: m=768, n=2048, k=1408, ks=[128, 128, 256, 256, 128, 256, 256, 256], 0.00103

Even when the test doesn't fail, the corruption still happens. You can see that by printing the torch.norm() of the result here and seeing that it differs between the test runs:

[...]
tensor(124499.3047, device='cuda:0')
tensor(122834.4531, device='cuda:0')
 > Perf (num_groups= 8, m=  768, n= 2048, k= 1664, gran_k=128):   44 us |  118 TFLOPS | 2372 GB/s
[...]
tensor(124495.9609, device='cuda:0')
tensor(122833.1406, device='cuda:0')
 > Perf (num_groups= 8, m=  768, n= 2048, k= 1664, gran_k=128):   45 us |  117 TFLOPS | 2356 GB/s
[...]
tensor(124501.3516, device='cuda:0')
tensor(122833.2109, device='cuda:0')
 > Perf (num_groups= 8, m=  768, n= 2048, k= 1664, gran_k=128):   45 us |  117 TFLOPS | 2367 GB/s

The fix for this is conceptually straightforward: section 9.7.9.26.5.2 from the PTX docs says that bulk async-group based completion mechanism can be used "for the completion of reading of the tensormap object" in (otherwise mbarrier-based) TMA loads. So running cp.async.bulk.commit_group + cp.async.bulk.wait_group.read 0 before modifying GMEM ensures all in-flight TMA loads have finished reading their tensormaps, which fixes the race condition.

The only catch is that writing to a TMA descriptor from a single thread generates multiple STG.E.128 instructions, and ptxas decides to move some of them before the DEPBAR.LE SB0, 0x0 generated by cp.async.bulk.wait_group.read 0, which re-introduces the race condition. CUTLASS does a __syncwarp() before updating tensormaps in GMEM, and inserting a warp sync indeed prevents ptxas from reordering instructions, but this seems more like a coincidence, since we only have one thread from the warp anyway and I don't see a good way to explain this in term of the PTX memory model. In other words, this is an ugly hack, so if you have better solutions, I will gladly implement them. :)

_{🔄 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/343 **Author:** [@dfyz](https://github.com/dfyz) **Created:** 5/29/2026 **Status:** ✅ Merged **Merged:** 7/6/2026 **Merged by:** [@zheanxu](https://github.com/zheanxu) **Base:** `main` ← **Head:** `main` --- ### 📝 Commits (1) - [`a6a9339`](https://github.com/deepseek-ai/DeepGEMM/commit/a6a9339a1140538c151e0e4a5d0660528fb9a1b9) Fix a race condition where in-flight tensormaps are updated in-place ### 📊 Changes **2 files changed** (+10 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh` (+9 -0) 📝 `tests/generators.py` (+1 -1) </details> ### 📄 Description The weight gradient kernel for SM90 modifies the A/B tensormaps stored in GMEM when the scheduler decides to switch groups, but it doesn't make sure that previous TMA loads have already finished reading the tensormaps at that point. A TMA load reading from a tensormap being modified concurrently corrupts the output for obvious reasons. A reliable way to reproduce this corruption (at least on an H200) is to run the `test_k_grouped_gemm_contiguous` test a couple of times with the shape added in this PR. Eventually the corruption will be large enough to make the test fail: ``` [...] Testing k-grouped contiguous GEMM: > Perf (num_groups= 4, m= 4096, n= 7168, k=34304, gran_k=128): 1966 us | 1025 TFLOPS | 681 GB/s > Perf (num_groups= 4, m= 7168, n= 2048, k=32896, gran_k=128): 970 us | 996 TFLOPS | 807 GB/s Traceback (most recent call last): File "/persistent/DeepGEMM/tests/test_fp8_fp4.py", line 222, in <module> test_k_grouped_gemm_contiguous() File "/persistent/DeepGEMM/tests/test_fp8_fp4.py", line 194, in test_k_grouped_gemm_contiguous assert diff < 0.001, f'{m=}, {n=}, {k=}, {ks=}, {diff:.5f}' ^^^^^^^^^^^^ AssertionError: m=768, n=2048, k=1408, ks=[128, 128, 256, 256, 128, 256, 256, 256], 0.00103 ``` Even when the test doesn't fail, the corruption still happens. You can see that by printing the `torch.norm()` of the result [here](https://github.com/deepseek-ai/DeepGEMM/blob/714dd1a4a980f7937a74343d19a8eba4fe321480/tests/test_fp8_fp4.py#L194) and seeing that it differs between the test runs: ``` [...] tensor(124499.3047, device='cuda:0') tensor(122834.4531, device='cuda:0') > Perf (num_groups= 8, m= 768, n= 2048, k= 1664, gran_k=128): 44 us | 118 TFLOPS | 2372 GB/s [...] tensor(124495.9609, device='cuda:0') tensor(122833.1406, device='cuda:0') > Perf (num_groups= 8, m= 768, n= 2048, k= 1664, gran_k=128): 45 us | 117 TFLOPS | 2356 GB/s [...] tensor(124501.3516, device='cuda:0') tensor(122833.2109, device='cuda:0') > Perf (num_groups= 8, m= 768, n= 2048, k= 1664, gran_k=128): 45 us | 117 TFLOPS | 2367 GB/s ``` The fix for this is conceptually straightforward: section [9.7.9.26.5.2](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-tensor) from the PTX docs says that bulk async-group based completion mechanism can be used "for the completion of reading of the tensormap object" in (otherwise mbarrier-based) TMA loads. So running `cp.async.bulk.commit_group` + `cp.async.bulk.wait_group.read 0` before modifying GMEM ensures all in-flight TMA loads have finished reading their tensormaps, which fixes the race condition. The only catch is that writing to a TMA descriptor from a single thread generates multiple `STG.E.128` instructions, and `ptxas` decides to move some of them before the `DEPBAR.LE SB0, 0x0` generated by `cp.async.bulk.wait_group.read 0`, which re-introduces the race condition. CUTLASS [does](https://github.com/NVIDIA/cutlass/blob/1732ed7da3b81d9f28b0130370ba70755a7e6dda/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp#L664) a `__syncwarp()` before updating tensormaps in GMEM, and inserting a warp sync indeed prevents `ptxas` from reordering instructions, but this seems more like a coincidence, since we only have one thread from the warp anyway and I don't see a good way to explain this in term of the PTX memory model. In other words, this is an ugly hack, so if you have better solutions, I will gladly implement them. :) --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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