NaNs from CUTLASS GQA op

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.

To Reproduce

import torch
import math
import xformers.ops as xops
from xformers.ops.fmha import cutlass

torch.manual_seed(0)

device = torch.device("cuda:4")
dtype = torch.float16  # float16 and bfloat16 break

batch = 1
seqlen = 129  # > 128 breaks
nheads_k = 2
ngroups = 8  # 4D MHA still works
headdim = 512  # >= 256 breaks
q = torch.randn(batch, seqlen, nheads_k, ngroups, headdim, device = device, dtype = dtype)
k = torch.randn(batch, seqlen, nheads_k, ngroups, headdim, device = device, dtype = dtype)
v = torch.randn(batch, seqlen, nheads_k, ngroups, headdim, device = device, dtype = dtype)

nan_count = 0
total_runs = 1000
for i in range(total_runs):

    o = xops.memory_efficient_attention(
        q, k, v,
        attn_bias=None,
        scale=1.0 / math.sqrt(headdim),
        op=(cutlass.FwOp, None),
    )

    if torch.isnan(o).any():
        nan_count += 1
        print(f"Run {i}: NaN in output, "
              f"o has {torch.isnan(o).sum().item()} NaNs out of {o.numel()})")

print(f"\n{nan_count}/{total_runs} runs produced NaN output")
print(f"xformers version: {xops.__version__ if hasattr(xops, '__version__') else 'unknown'}")
print(f"torch version: {torch.__version__}")
print(f"GPU: {torch.cuda.get_device_name(device)}")

Example output:

...
Run 980: NaN in output, o has 1 NaNs out of 1056768)
Run 996: NaN in output, o has 3125 NaNs out of 1056768)
Run 998: NaN in output, o has 2887 NaNs out of 1056768)

363/1000 runs produced NaN output
xformers version: unknown
torch version: 2.10.0+cu130
GPU: NVIDIA GeForce RTX 3090

Expected behavior

No NaNs.

Environment

• PyTorch Version (e.g., 1.0): 2.10.0+cu130
• OS (e.g., Linux): Manjaro Linux (x86_64)
• How you installed PyTorch (conda, pip, source): pip
• Build command you used (if compiling from source):
• Python version: 3.13.11
• CUDA/cuDNN version:
• GPU models and configuration:
  • GPU 0: NVIDIA GeForce RTX 4090
  • GPU 1: NVIDIA RTX PRO 6000 Blackwell Workstation Edition
  • GPU 2: NVIDIA GeForce RTX 5090
  • GPU 3: NVIDIA GeForce RTX 4090
  • GPU 4: NVIDIA GeForce RTX 3090
• Any other relevant information:

Additional context

I understand GQA is marked as experimental and this may not be a priority. But there aren't a lot of options for large head sizes currently. :|