deepseek-ai / deepseek-ai/DeepGEMM
[nv_dev][SM120] k_grouped_fp8_gemm_nt_contiguous returns non-deterministic wrong results — sm120 kernel is missing the tensormap drain fixed for SM90 in #343
- Dominant language
- Cuda
- Stars
- 7.8k
- Forks
- 1.3k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 3
Description
The SM120 1D1D kernel was copied from `sm90_fp8_gemm_1d1d.cuh` (merged into nv_dev on 6/24 in #324), **twelve days before #343** fixed the in-flight tensormap race on the SM90 side. Since the copy lives at a different file path (`sm120_fp8_fp4_gemm_1d1d.cuh`), the #343 fix cannot propagate to it through merges — the SM120 group-switch still publishes GMEM tensormaps without draining in-flight TMA loads.
Notably, the regression shape that #343 added to `test_k_grouped_gemm_contiguous` (8 groups, m=768, n=2048, k≈128/group) does its job on SM120: it fails consistently on our RTX 5090 D, which is how we found this.
## Environment
RTX 5090 D (SM120), CUDA 13.0, PyTorch 2.11.0+cu130. Built from the nv_dev tarball @ `2642b32` with pinned submodules (CUTLASS @ f3fde58, fmt @553ec11).
## Repro
Run `test_k_grouped_gemm_contiguous` (tests/test_fp8_fp4.py). The first shape fails with diff (the test's relative-error metric) of 0.0015–0.006, threshold 0.001. Observed trigger rate: **13/14 runs** (10/10 with warm JIT cache, 3/3 with cold cache, plus 1 fail / 1 pass in early manual runs — the single pass itself matches race behavior). The same shape passes via the TN (MN-major) path — isolated to the NT kernel.
## It is a race
Identical inputs (fixed seeds, data regenerated each run), 3 fresh runs:
```
diff = 0.006319 / 0.005205 / 0.003889
```
Errors spread across all 8 groups (max abs error ~122 on k=128 groups) — every group switch is a hazard.
## Root cause
`sm120_fp8_fp4_gemm_1d1d.cuh`, KGroupedContiguous group-switch (~L236-282):
```cpp
ptx::tensor_map_replace_global_addr_in_smem(smem_tm_a, a_base + a_offset);
ptx::tensor_map_replace_global_addr_in_smem(smem_tm_b, b_base + b_offset);
...
// ← the #343 drain (commit_group + wait_group) is missing here
*gmem_tm_a = *smem_tm_a; // published while TMA loads may still be in flight
*gmem_tm_b = *smem_tm_b;
ptx::tensor_map_release_gpu();
ptx::tensor_map_acquire_gpu(gmem_tm_a);
ptx::tensor_map_acquire_gpu(gmem_tm_b);
```
Same publish-without-drain pattern that #365 root-caused on SM90; the fixed SM90 path has `cute::tma_desc_commit_group() + cute::tma_desc_wait_group()` before the GMEM publish (sm90_fp8_gemm_1d1d.cuh:202-207).
## Fix validated on SM120
Porting the #343 drain (same 4 lines, inserted before the GMEM publish):
```cpp
cute::tma_desc_commit_group();
cute::tma_desc_wait_group();
__syncwarp(1u << lane_idx);
```
- diff drops to **0.000105, bit-identical across runs** (vs 0.006 non-deterministic before)
- full k_grouped + m_grouped suites pass (56/56 configs); 100× soak clean
- Perf A/B (same seeds/data): sentinel shape 87.0 → 87.8 µs (**+0.9%**, worst case for group-switch density); typical EP shape 3106 → 3114 µs (**+0.26%**) — the same trade-off SM90 accepted in #343
## Scope / timing
`main` is not affected (no SM120 support there yet). But when nv_dev next syncs to main, this would ship with the first official SM120 support. Recent SM120 issues (#392, #405, #417, #425) were all found on real hardware — there appears to be no SM120 CI runner, so release testing is unlikely to catch this.
## Note on the bf16 kernel
`sm120_bf16_gemm.cuh` has the same un-drained group-switch pattern. We could not make it fail on our hardware (full suite + 50 targeted stress runs pass), but per the commit-group ordering analysis in #365 it is the same UB class — happy to include a defensive fix for it too if maintainers prefer.
---
Repro script and raw logs: https://gist.github.com/aganhui/3b374ffb3d6d13b52d5688780fafac28
PR incoming (base nv_dev). Happy to run any additional validation on this 5090 D (soak, sanitizer, benchmarks) if useful, and to adjust to whatever structure maintainers prefer — including a 3-phase restructure à la #375.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with sm120_fp8_fp4_gemm_1d1d.cuh at the KGroupedContiguous group-switch and compare it with the fixed section in sm90_fp8_gemm_1d1d.cuh. Run test_k_grouped_gemm_contiguous in tests/test_fp8_fp4.py on SM120 and verify that repeated runs are deterministic. Done means the grouped suites pass with the reported regression shape and the race-related error is gone.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 52/100