[PR #415] Fix SM90 k-grouped GEMM hang and silent corruption when TMA multicast clusters straddle group boundaries #428

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/415
Author: @shoemoney
Created: 8/24/2026
Status: 🔄 Open

Base: main ← Head: fix-sm90-kgrouped-multicast-group-straddle

📝 Commits (1)

7b270b5 Reject SM90 k-grouped multicast configs whose per-group grid is not divisible by the cluster size

📊 Changes

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

View changed files

📝 csrc/jit_kernels/heuristics/sm90.hpp (+8 -0)
📝 tests/generators.py (+3 -1)

📄 Description

SM90 k_grouped_fp8_gemm_nt_contiguous (and the BF16 k-grouped path, which shares the same candidate enumerator) can hang or silently produce wrong output when the heuristic enables TMA multicast with an odd per-group block grid.

Bug

The k-grouped scheduler lays the groups' block grids back-to-back. SM90 clusters are consecutive block-index pairs. With num_groups=2, m=192, n=3072, ks=[128, 256] on 132 SMs, csrc/jit_kernels/heuristics/sm90.hpp selects BLOCK_M=64, BLOCK_N=72, cluster size 2. The per-group grid is ceil(192/64) * ceil(3072/72) = 3 * 43 = 129 blocks, so group 1 starts at odd global block 129 and cluster {128, 129} spans both groups.

The two CTAs then disagree on is_tma_multicast_valid(). Rank 0 is remapped to num_blocks_in_group == 1 by the unaligned-multicast fix at deep_gemm/include/deep_gemm/scheduler/gemm.cuh:133-140; rank 1 is not. Only cluster rank 0 issues SM90_TMA_LOAD_MULTICAST_2D (deep_gemm/include/deep_gemm/common/tma_copy.cuh:47), so rank 1's full barrier never receives its B transaction bytes and the kernel hangs. Every later pair inside group 1 is parity-shifted: both CTAs consider multicast valid while needing different B tiles, which gives silently wrong output. Unequal ks also desync the empty-barrier phases across the straddling pair.

Fix

Reject k-grouped multicast candidates whose per-group grid is not divisible by the cluster size. The check sits next to the existing masked-layout legality check (sm90.hpp:87-91). This is the same invariant the SM100 heuristic already enforces at csrc/jit_kernels/heuristics/sm100.hpp:103 and :115.

The runtime remap in gemm.cuh:133-140 could instead disable multicast for a straddling pair, but SM90 bakes the multicast arrival count into barrier initialization at launch. The host-side filter keeps the invariant the kernel already assumes, as SM100 does.

Verification

I ported the exact SM90 candidate enumeration and cost model to a script (132 SMs). Without the filter, 269 (num_groups<=4, m, n) combos in a small scan select a group-straddling multicast config, including the shape above. With the filter, every one falls back to a legal even-grid or cluster-1 config; the example shape picks BLOCK_M=64, BLOCK_N=152, cluster 1, 63-block grid.

I added that shape to the k-grouped test generator, mirroring the existing SM90 tensor-map stress shape. The shipped shapes all have even per-group grids, which is why the suite never hit this.

I do not have Hopper hardware. The kernel-level failure is derived from the analysis and the model port, not an on-device run.

_{🔄 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/415 **Author:** [@shoemoney](https://github.com/shoemoney) **Created:** 8/24/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix-sm90-kgrouped-multicast-group-straddle` --- ### 📝 Commits (1) - [`7b270b5`](https://github.com/deepseek-ai/DeepGEMM/commit/7b270b53112c6ac8efa527da3ba2e3ca6a0e5ef6) Reject SM90 k-grouped multicast configs whose per-group grid is not divisible by the cluster size ### 📊 Changes **2 files changed** (+11 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/heuristics/sm90.hpp` (+8 -0) 📝 `tests/generators.py` (+3 -1) </details> ### 📄 Description SM90 `k_grouped_fp8_gemm_nt_contiguous` (and the BF16 k-grouped path, which shares the same candidate enumerator) can hang or silently produce wrong output when the heuristic enables TMA multicast with an odd per-group block grid. ## Bug The k-grouped scheduler lays the groups' block grids back-to-back. SM90 clusters are consecutive block-index pairs. With `num_groups=2, m=192, n=3072, ks=[128, 256]` on 132 SMs, `csrc/jit_kernels/heuristics/sm90.hpp` selects `BLOCK_M=64, BLOCK_N=72`, cluster size 2. The per-group grid is `ceil(192/64) * ceil(3072/72) = 3 * 43 = 129` blocks, so group 1 starts at odd global block 129 and cluster `{128, 129}` spans both groups. The two CTAs then disagree on `is_tma_multicast_valid()`. Rank 0 is remapped to `num_blocks_in_group == 1` by the unaligned-multicast fix at `deep_gemm/include/deep_gemm/scheduler/gemm.cuh:133-140`; rank 1 is not. Only cluster rank 0 issues `SM90_TMA_LOAD_MULTICAST_2D` (`deep_gemm/include/deep_gemm/common/tma_copy.cuh:47`), so rank 1's full barrier never receives its B transaction bytes and the kernel hangs. Every later pair inside group 1 is parity-shifted: both CTAs consider multicast valid while needing different B tiles, which gives silently wrong output. Unequal `ks` also desync the empty-barrier phases across the straddling pair. ## Fix Reject k-grouped multicast candidates whose per-group grid is not divisible by the cluster size. The check sits next to the existing masked-layout legality check (`sm90.hpp:87-91`). This is the same invariant the SM100 heuristic already enforces at `csrc/jit_kernels/heuristics/sm100.hpp:103` and `:115`. The runtime remap in `gemm.cuh:133-140` could instead disable multicast for a straddling pair, but SM90 bakes the multicast arrival count into barrier initialization at launch. The host-side filter keeps the invariant the kernel already assumes, as SM100 does. ## Verification I ported the exact SM90 candidate enumeration and cost model to a script (132 SMs). Without the filter, 269 `(num_groups<=4, m, n)` combos in a small scan select a group-straddling multicast config, including the shape above. With the filter, every one falls back to a legal even-grid or cluster-1 config; the example shape picks `BLOCK_M=64, BLOCK_N=152`, cluster 1, 63-block grid. I added that shape to the k-grouped test generator, mirroring the existing SM90 tensor-map stress shape. The shipped shapes all have even per-group grids, which is why the suite never hit this. I do not have Hopper hardware. The kernel-level failure is derived from the analysis and the model port, not an on-device run. --- <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:17:41 +00:00

bsakaguchi referenced this issue

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

[PR #428] [CLOSED] fix: reject non-positive block size multiples #439