[0.0.18] memory_efficient_attention NaNs when seqlen>32768

[comfyanonymous]

🐛 Bug

Command

To Reproduce

Steps to reproduce the behavior:

import xformers
import xformers.ops
import torch

q = torch.zeros(([1, 33728, 512])).cuda()
k = torch.zeros(([1, 33728, 512])).cuda()
v = torch.zeros(([1, 33728, 512])).cuda()
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)
print(out)
print(torch.isnan(out).any())

Expected behavior

out should not contain Nan.

Environment

This test was done on the free google colab with their T4 GPU using the 0.0.18 package on pip.

Additional context

It works fine on 0.0.17 but fails on 0.0.18.

People have been reporting that my ComfyUI returns black images during the VAE decoding phase when the resolution is higher than a certain amount and I have narrowed it down to this issue.

[leppie]

I hit this issue too yesterday after upgrading to 0.0.18.

[danthe3rd]

Hi,
Thanks for the minimal reproduction example. I can repro on A100 - I'll have a look.
Also pinning this issue as it's quite important

[danthe3rd]

It seems to happen due to a cast to int16 at some point in the code, so it happens when the sequence length is larger than 32768. Updating the title according.

[danthe3rd]

I have a tentative fix - hopefully we can land that soon and release the 0.0.19 early next week to address that

[danthe3rd]

It should be fixed as of 68dce69, and will be included in the next release (0.0.19). In the meantime, you can also use a development build >=0.0.19.dev516

[patrickvonplaten]

Thanks a mille for the fix @danthe3rd - very helpful thread here!