SIGN IN SIGN UP

[BugFix][ROCm] Reject sub-wavefront block sizes instead of crashing (#2918)

A T.gemm in a block narrower than one wavefront aborted the whole process with
SIGFPE during compilation. On CDNA the wavefront is 64 wide, so

    num_warps = block_size / TargetRocmGetWarpSize(target)

floors to zero for a 32-thread block, and ComputeDefaultWarpPartition then
evaluates `M % (m_warp * kMPerWarp)` with m_warp == 0.

Easy to hit by accident: a thread count tuned for NVIDIA's 32-lane warp is valid
on CUDA and lands exactly here. All four examples/deepseek_nsa kernels hardcode
`threads = 32`, and because the failure killed the process rather than raising,
it took down the pytest worker and produced collateral failures in unrelated
tests under xdist.

Adds ICHECK_GE(block_size, warp_size) before the division, plus a defensive
ICHECK_GT(num_warps, 0) for callers that bypass ComputeWarpPartition.

ROCm-only. The CUDA path structures the same computation as a loop that does not
execute when num_warps == 0, so it never had this failure mode.

Before: examples/deepseek_nsa exits 136 (SIGFPE), no diagnostic.
After:  exits 1 with a message naming the block size and the wavefront size.

Verified on gfx942 (MI300X, ROCm 7.2): 1932 passed, 1223 skipped, 0 failed.
A
andyluo7 committed
5b6fcac526d5baf5ab1a69d9b7ec14acca1bb612
Parent: 7fc4706
Committed by GitHub <noreply@github.com> on 8/8/2026, 4:38:51 PM