[BUG] Fix sm90 grouped multicast deadlock

[baptiste-bt]

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

[baptiste-bt]

Had missed this but this issue looks like a duplicate of #415

[ds-review-bot]

🔵 suggestion: 代码风格与周边不一致：本文件中其他单语句 if ... continue; 均不带花括号，且新增块与下一段注释（"The block sizes cannot be too large..."）之间缺少空行。建议去掉花括号并补一个空行，保持与上下文一致。

_{🤖 v5}

[ds-review-bot]

🔵 suggestion: 该检查比严格必要条件更保守：死锁只在相邻 group 的 K-loop 长度不同（即 cluster 跨 group）时发生。对于 num_groups == 1 的退化情况，普通 GEMM 已能处理末尾不完整 cluster，此处会不必要地禁用多播。如在意该场景性能，可加上 desc.num_groups > 1 条件；若认为该场景罕见，保持现状也可，但建议在注释中说明这是保守限制。

_{🤖 v5}

[ds-review-bot]

🔵 suggestion: 建议后续核对 SM100 启发式（csrc/jit_kernels/heuristics/sm100.hpp）的 K-grouped 调度是否存在同类跨组 cluster 的 barrier 风险；若确认无风险，最好在此处或彼处注释中说明原因，避免读者疑惑为何限制仅加在 SM90。

_{🤖 v5}

[ds-review-bot]

🔵 suggestion: 建议把注释 "SM90 multicast boundary regression" 与具体问题关联（例如引用 issue/commit 号），方便后续读者理解为何存在这个非典型形状；另注意该用例在 SM100 上也会随更多 k_alignment 选项执行，会略微增加测试时长，如需可加架构条件限制。

_{🤖 v5}

[ds-review-bot]

🔵 suggestion: Style suggestion only: the new guard uses braces and spans three lines, while the surrounding legality checks (e.g. the masked-layout check just above) omit braces and break after the condition. Consider matching the local style for readability, e.g. 'if (desc.gemm_type == GemmType::KGroupedContiguous and (ceil_div(desc.m, block_m) * ceil_div(desc.n, block_n)) % cluster_size != 0) continue;'. Cosmetic; no functional impact.

_{🤖 v4f}

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

变更正确限制了 SM90 K-grouped multicast 布局，确保每组 tile 数与 CTA cluster 对齐，并补充了对应回归场景。未发现会破坏现有功能的缺陷。

v5

本 MR 修复了 SM90 上 k_grouped_bf16_gemm 的多播（multicast）死锁：当每组输出 tile 数不能被 CTA cluster size 整除时，一个 cluster 的两个 CTA 可能跨越两个 K 长度不同的 group，导致分布式 barrier 的 arrive/wait 计数不匹配而永久挂起。修复方式是在 SM90 启发式（csrc/jit_kernels/heuristics/sm90.hpp）中拒绝 KGroupedContiguous 且 (ceil_div(m, block_m) * ceil_div(n, block_n)) % cluster_size != 0 的候选布局，并顺带把重复的 cluster_m * cluster_n 提取为局部变量 cluster_size。tests/generators.py 新增了回归形状 (2, 1408, 1408, 256)：m=n=1408 在 block 128 下恰好产生 11*11=121 个 tile（奇数），配合两组随机且按 128 对齐后大概率不同的 K，可确定性复现修复前的死锁。审查结论：修复正确且保守安全——KGroupedContiguous 下所有 group 共享同一 m×n 输出形状，用 desc.m/desc.n 计算每组 tile 数是正确的；cluster_size==1 的候选恒满足该条件，因此 DG_HOST_ASSERT(not candidates.empty()) 不会因新过滤条件而触发；该检查作用于所有走 get_best_config<SM90ArchSpec> 的 K-grouped 路径（含 FP8 1D1D），覆盖面合理。整体建议合入，下方仅为少量风格与可选改进意见。

v4f

The MR correctly fixes the SM90 grouped-multicast deadlock reported for k_grouped_bf16_gemm_tn_contiguous (m=n=1408, ks=[384,256]) on H100. Root cause: with a persistent scheduler, each CTA pair in a 2-CTA cluster always processes two consecutive global tile indices (next_block_idx = iter * num_sms + blockIdx.x), and with num_sms divisible by the cluster size the pair starts are always even. A K-group boundary therefore only falls inside a cluster pair when the tile count per group (ceil_div(m, block_m) * ceil_div(n, block_n)) is odd (e.g. 11x11=121 tiles for 128x128 blocks). The two CTAs then enter different groups whose K-loop lengths differ (384/64=6 vs 256/64=4 blocks), so the TMA producers and math consumers wait on the distributed full/empty cluster barriers with phases that can never match, hanging the kernel. The fix rejects every multicast layout candidate whose per-group tile count is not divisible by the cluster size, so group boundaries always align to cluster boundaries. For the repro the heuristic now falls back to Layout(block_m=64, block_n=256, block_k=64, cluster_m=2, cluster_n=1) (132 tiles/group, even), which completes. I verified the fix against the scheduler/barrier arithmetic in deep_gemm/include/deep_gemm/scheduler/gemm.cuh and deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh: the check correctly uses desc.m/desc.n (the M/N shape shared by every group in a K-grouped GEMM), the K dim is intentionally not part of the check since the differing K-loop length across groups is exactly what the alignment prevents from ever being co-resident in one cluster. Because cluster_size==1 always passes the modulo test, non-multicast candidates survive and DG_HOST_ASSERT(not candidates.empty()) cannot fire; num_groups > 4 already disables multicast. Placing the guard in SM90ArchSpec::get_layout_candidates() also covers all SM90 K-grouped kernels (bf16 and fp8 1d1d/1d2d) that share this heuristic. The deadlock mechanism is SM90-specific (SM100 k-grouped uses a different 2-CTA scheduling scheme), so leaving SM100ArchSpec untouched is appropriate. The accompanying change to tests/generators.py adds a regression shape (2, 1408, 1408, 256) that reproduces the odd-tile (121 for 128x128) cluster-2 boundary and exercises a small num_groups=2 case not covered by the existing EP/16-32-64 cases; the cluster_size refactor in sm90.hpp is behavior-neutral cleanup. I found no correctness or safety issues; only the minor style suggestion below. Recommendation: approve.

Files reviewed: 2
Issues found: 🔵 5 suggestion
Inline comments posted: 5