[PR #429] [BUG] Fix sm90 grouped multicast deadlock #438

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/429
Author: @baptiste-bt
Created: 9/8/2026
Status: 🔄 Open

Base: main ← Head: fix/sm90-k-grouped-multicast-deadlock

📝 Commits (1)

acc504d Fix sm90 grouped multicast deadlock

📊 Changes

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

View changed files

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

📄 Description

I noticed a bug on sm90, k_grouped_bf16_gemm can deadlock when the number of output tiles per group is not divisible by the CTA cluster size.

Repro:

import os

os.environ["DG_JIT_DEBUG"] = "1"
os.environ["DG_PRINT_CONFIGS"] = "1"
os.environ["DG_JIT_WITH_LINEINFO"] = "1"

import deep_gemm
import torch

ks = [384, 256]
m = n = 1408
sum_k = sum(ks)

a = torch.randn(sum_k, m, device="cuda", dtype=torch.bfloat16)
b = torch.randn(sum_k, n, device="cuda", dtype=torch.bfloat16)
d = torch.zeros(len(ks), m, n, device="cuda", dtype=torch.float32)
grouped_layout = torch.tensor(ks, device="cuda", dtype=torch.int32)

deep_gemm.k_grouped_bf16_gemm_tn_contiguous(a, b, d, ks, grouped_layout, d)
torch.cuda.synchronize()  # Hangs without the fix.
print("Ok")

Setup: H100 with DeepGEMM 2.6.1 (559d79f), pytorch 2.11 + CUDA 13.2

With debug output enabled, the unpatched heuristic selects Layout(swap_ab=0, block_m=128, block_n=128, block_k=64, cluster_m=1, cluster_n=2) and hangs after:

Launch kernel with {132, 1} x 384, shared memory: 230144 bytes, cluster: 2, pdl: 0, stream: 0

This gives 11 * 11 = 121 tiles per group for a cluster size of 2. One CTA pair can therefore straddle two groups with different K-loop lengths and deadlock on the distributed barriers.

With -G -O0 appended to the JIT NVCC command, CUDA-GDB shows the two CTAs spanning the group boundary:

$ cuda-gdb --args python repro.py
(cuda-gdb) run
^C
Thread 1 "python" received signal SIGINT, Interrupt.

(cuda-gdb) info cuda clusters
  ClusterIdx  BlockIdx To BlockIdx Count
Kernel 0
ClusterDim (2,1,1)
    (60,0,0) (120,0,0)   (121,0,0)     2

(cuda-gdb) info cuda blocks
   BlockIdx To BlockIdx Count   State
Kernel 0
  (120,0,0)   (121,0,0)     2 running

(cuda-gdb) cuda block (120,0,0)
(cuda-gdb) cuda thread (320,0,0)
(cuda-gdb) bt full
#0  cutlass::arch::ClusterBarrier::wait(smem_ptr=0x...38428, phase=0)
    at cutlass/arch/barrier.h:424
#2  deep_gemm::sm90_bf16_gemm_impl<...>()
    at deep_gemm/impls/sm90_bf16_gemm.cuh:169
        num_total_k_blocks = 6
        scheduler = {... current_group_idx = 0, current_shape_k = 384 ...}
        phase = 1

(cuda-gdb) cuda block (121,0,0)
(cuda-gdb) cuda thread (320,0,0)
(cuda-gdb) bt full
#0  cutlass::arch::ClusterBarrier::wait(smem_ptr=0x...38428, phase=0)
    at cutlass/arch/barrier.h:424
#2  deep_gemm::sm90_bf16_gemm_impl<...>()
    at deep_gemm/impls/sm90_bf16_gemm.cuh:205
        i = 1
        scheduler = {... current_group_idx = 2, current_shape_k = 256 ...}

(cuda-gdb) cuda thread (0,0,0)
(cuda-gdb) bt full
#0  cutlass::arch::ClusterBarrier::wait(smem_ptr=0x...38400, phase=0)
    at cutlass/arch/barrier.h:424
#2  deep_gemm::sm90_bf16_gemm_impl<...>()
    at deep_gemm/impls/sm90_bf16_gemm.cuh:249
        k_block_idx = 0
        num_total_k_blocks = 4
        scheduler = {... current_group_idx = 1, current_shape_k = 256 ...}
        phase = 0

CTA 120's TMA producer is still processing group 0 at empty_barriers[stage_idx]->wait(phase ^ 1) on line 169. CTA 121's producer has reached the final cleanup wait on line 205, while its consumers are waiting on full_barriers[stage_idx]->wait(phase) on line 249 for group 1. The matching peer arrival never happens.

This change rejects K-grouped layouts unless the tile count per group is divisible by the cluster size. For this case, the heuristic selects Layout(swap_ab=0, block_m=64, block_n=256, block_k=64, cluster_m=2, cluster_n=1) and completes:

Launch kernel with {132, 1} x 256, shared memory: 230144 bytes, cluster: 2, pdl: 0, stream: 0
Ok

_{🔄 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/429 **Author:** [@baptiste-bt](https://github.com/baptiste-bt) **Created:** 9/8/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/sm90-k-grouped-multicast-deadlock` --- ### 📝 Commits (1) - [`acc504d`](https://github.com/deepseek-ai/DeepGEMM/commit/acc504d4d9f878b1bb42dbb4401ace16d6a08295) Fix sm90 grouped multicast deadlock ### 📊 Changes **2 files changed** (+10 additions, -4 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/heuristics/sm90.hpp` (+8 -3) 📝 `tests/generators.py` (+2 -1) </details> ### 📄 Description I noticed a bug on sm90, `k_grouped_bf16_gemm` can deadlock when the number of output tiles per group is not divisible by the CTA cluster size. Repro: ```python import os os.environ["DG_JIT_DEBUG"] = "1" os.environ["DG_PRINT_CONFIGS"] = "1" os.environ["DG_JIT_WITH_LINEINFO"] = "1" import deep_gemm import torch ks = [384, 256] m = n = 1408 sum_k = sum(ks) a = torch.randn(sum_k, m, device="cuda", dtype=torch.bfloat16) b = torch.randn(sum_k, n, device="cuda", dtype=torch.bfloat16) d = torch.zeros(len(ks), m, n, device="cuda", dtype=torch.float32) grouped_layout = torch.tensor(ks, device="cuda", dtype=torch.int32) deep_gemm.k_grouped_bf16_gemm_tn_contiguous(a, b, d, ks, grouped_layout, d) torch.cuda.synchronize() # Hangs without the fix. print("Ok") ``` Setup: H100 with DeepGEMM 2.6.1 (`559d79f`), pytorch `2.11` + CUDA `13.2` With debug output enabled, the unpatched heuristic selects `Layout(swap_ab=0, block_m=128, block_n=128, block_k=64, cluster_m=1, cluster_n=2)` and hangs after: ```text Launch kernel with {132, 1} x 384, shared memory: 230144 bytes, cluster: 2, pdl: 0, stream: 0 ``` This gives `11 * 11 = 121` tiles per group for a cluster size of 2. One CTA pair can therefore straddle two groups with different K-loop lengths and deadlock on the distributed barriers. With `-G -O0` appended to the JIT NVCC command, CUDA-GDB shows the two CTAs spanning the group boundary: ```text $ cuda-gdb --args python repro.py (cuda-gdb) run ^C Thread 1 "python" received signal SIGINT, Interrupt. (cuda-gdb) info cuda clusters ClusterIdx BlockIdx To BlockIdx Count Kernel 0 ClusterDim (2,1,1) (60,0,0) (120,0,0) (121,0,0) 2 (cuda-gdb) info cuda blocks BlockIdx To BlockIdx Count State Kernel 0 (120,0,0) (121,0,0) 2 running (cuda-gdb) cuda block (120,0,0) (cuda-gdb) cuda thread (320,0,0) (cuda-gdb) bt full #0 cutlass::arch::ClusterBarrier::wait(smem_ptr=0x...38428, phase=0) at cutlass/arch/barrier.h:424 #2 deep_gemm::sm90_bf16_gemm_impl<...>() at deep_gemm/impls/sm90_bf16_gemm.cuh:169 num_total_k_blocks = 6 scheduler = {... current_group_idx = 0, current_shape_k = 384 ...} phase = 1 (cuda-gdb) cuda block (121,0,0) (cuda-gdb) cuda thread (320,0,0) (cuda-gdb) bt full #0 cutlass::arch::ClusterBarrier::wait(smem_ptr=0x...38428, phase=0) at cutlass/arch/barrier.h:424 #2 deep_gemm::sm90_bf16_gemm_impl<...>() at deep_gemm/impls/sm90_bf16_gemm.cuh:205 i = 1 scheduler = {... current_group_idx = 2, current_shape_k = 256 ...} (cuda-gdb) cuda thread (0,0,0) (cuda-gdb) bt full #0 cutlass::arch::ClusterBarrier::wait(smem_ptr=0x...38400, phase=0) at cutlass/arch/barrier.h:424 #2 deep_gemm::sm90_bf16_gemm_impl<...>() at deep_gemm/impls/sm90_bf16_gemm.cuh:249 k_block_idx = 0 num_total_k_blocks = 4 scheduler = {... current_group_idx = 1, current_shape_k = 256 ...} phase = 0 ``` CTA 120's TMA producer is still processing group 0 at `empty_barriers[stage_idx]->wait(phase ^ 1)` on line 169. CTA 121's producer has reached the final cleanup wait on line 205, while its consumers are waiting on `full_barriers[stage_idx]->wait(phase)` on line 249 for group 1. The matching peer arrival never happens. This change rejects K-grouped layouts unless the tile count per group is divisible by the cluster size. For this case, the heuristic selects `Layout(swap_ab=0, block_m=64, block_n=256, block_k=64, cluster_m=2, cluster_n=1)` and completes: ```text Launch kernel with {132, 1} x 256, shared memory: 230144 bytes, cluster: 2, pdl: 0, stream: 0 Ok ``` --- <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:47 +00:00

bsakaguchi referenced this issue

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

[PR #439] [CLOSED] [Bugfix][Layout] Add cuda_fp8.h include to fix __nv_fp8_e4m3 undefined #447