[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