Memory-efficient attention (without xformers)

[Birch-san]

I implemented sub-quadratic attention (as described in https://arxiv.org/abs/2112.05682v2):
https://twitter.com/Birchlabs/status/1607503573906063362
Birch-san#1
Birch-san/diffusers-play@a573e3d

is this worth upstreaming? it enables creation of images larger than can be achieved with attention slicing.

[patrickvonplaten]

Hey @Birch-san,

Thanks a lot for the issue! @patil-suraj what do you think?

[patil-suraj]

Very cool @Birch-san , is this more efficient than xformers ? Also, xformers installation situation is getting better now cf https://pypi.org/project/xformers/#history, so not sure if we need another efficient attention for PT. This could be a good addition to flax.

[Birch-san]

this implements the same paper as xformers memory efficient attention.
it's unlikely to be more efficient than xformers, since they have the advantage of custom CUDA kernels.

xformers is CUDA-only, I presume? no support for MPS or ROCm or Metal or CPU backends?

there are Mac users trying to run stable-diffusion on Mac Minis with 8GB of unified memory. IIRC they couldn't even fit 512x512 images. sliced attention helped with that, but this goes further: you can chunk up attention far finer, arbitrarily so.

if my calculations are correct: a 2048x2048 image's self-attention can require 80GB VRAM ordinarily. setting sliced attention to its most aggressive setting can get this down to 40GB slices (MPS backend refuses to allocate this). but the chunked attention implementation I've provided here can get it down to anything you like, e.g. 80MB chunks.

[Lime-Cakes]

This memory efficient attention might not be faster than xformer implementation, but since it doesn't rely on custom CUDA kernel, this means better support for non-CUDA device, which would be good for mac and other non-cuda accelerator on pytorch.

[patrickvonplaten]

@pcuenca could we maybe run some tests for MPS for this?

xformers will soon be natively supported in PyTorch so I'm wondering how important this for PyTorch-only. I definitely see an important use case for MPS though.

Also with the new attention processor system it should be relatively easy to add either way.

[patil-suraj]

Good point regarding MPS @Birch-san ! In that case, it would be cool to have this. cc @pcuenca

[keturn]

There's another implementation here: https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/ldm/modules/sub_quadratic_attention.py#L153

comfy reports:

I tweaked the sub-quadratic optimization when I implemented it in my own UI and it gave me a nice speed boost.
but I had to tweak it first cause by default it didn't give better performance than the split optimization on my 6800XT

[pcuenca]

Hi @Birch-san, sorry for being so slow to react here! Is there any chance you could submit a PR that applies these optimizations to mps? That way it would be easier for us to test and discuss. Perhaps you can use the new attention processor mechanism so people can opt in (or make it default maybe) if they are running on mps. If you can't then I'll test your branch :)

Another question I have is: will this become obsolete with the upcoming memory-efficient attention integrated in PyTorch 2.x? Or does that not work for mps at all?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[keturn]

Still interested, Stalebot!

Related:

• [Flax] added memory efficient attention for U-net  #2231
• efficient memory attention #2461

[Birch-san]

yup, planning to submit PRs for memory-efficient attention and cross-attention masks soon.

[patrickvonplaten]

PyTorch 2.0 will provide memory efficient attention out of the box - is this PR still relevant then?

[Birch-san]

will PyTorch 2.0 provide memory-efficient attention for Mac?

[Birch-san]

@patrickvonplaten @pcuenca

https://pytorch.org/docs/2.0/generated/torch.nn.functional.scaled_dot_product_attention.html

the docs for memory-efficient attention link to xformers, which to my knowledge does not support MPS (it's focused on CUDA and triton).

the pytorch context manager for activating memory-efficient attention is torch.backends.cuda.enable_mem_efficient_sdp(), which again is CUDA-specific.

so I think there's still a case for implementing this: for Mac, iOS, ROCm, CPU. it also means that if you trace the torchscript operations: you'd get a memory-efficient torchscript model, which you could convert to CoreML†.

† if you wanted the CoreML model to be optimal for Neural Engine: you'd need to reshape the algorithm a bit; Neural Engine prefers B,C,1,S tensors over B,S,C tensors.