来自 CUTLASS GQA 操作的 NaNs

Bug The CUTLASS op seems to produce NaN values from certain tensor shapes when using GQA. It's nondeterministic, but I've observed it on both Ampere and Ada GPUs. I also tried patching xformers.ops.fmha.FwOp.CUDA_MAXIMUM_COMPUTE_CAPABILITY to (12, 0) to confirm the same behavior on Blackwell. Sequence lengths <= 128 don't cause the behavior at all, longer sequences do. 129 seems to be particularly bad but it's hard to say for sure. Head dim is a factor, too. Curiously the dtype also makes a difference, with float32 not showing any issues, although it doesn't appear to be an overflow condition since bfloat16 still breaks. I've seen failure rates as high as 50% sometimes, but it's not entirely consistent. Maybe depends on the state of uninitialized memory? CUDA_LAUNCH_BLOCKING=1 prevents the error from happening, though, so presumably it's a race condition. Repro script below.