Implement Flash Attention for 50% speedup

[n00mkrad]

Check out this reddit post:

https://www.reddit.com/r/StableDiffusion/comments/xmr3ic/speed_up_stable_diffusion_by_50_using_flash/

Claims to achieve a 50% speedup by using Flash Attention.

Is it possible to implement this into this repo?

[n00mkrad]

See here for a reference implementation:

https://github.com/TheLastBen/fast-stable-diffusion

[Any-Winter-4079]

Yep, I had just read that post. I'll take a look.

[Any-Winter-4079]

Repo for Flash Attention https://github.com/HazyResearch/flash-attention

[krummrey]

The pace at which SD is maturing is unbelievable. The true power of going Open Source.

[tildebyte]

The only issue I see, at least in the "official" paper-accompaniment repo, is that there's a CUDA-side c++ binary library (I guess? I don't see how it's being built) - building binaries via conda/pip is a non-starter in this repo. We have enough trouble with pure-Python installs 😂

If they are going to be releasing binaries for the Big Three platforms, then OK, otherwise...

[nguyenkm]

There's am active thread on how to get it to work for non-Linux systems. huggingface/diffusers#532

[Any-Winter-4079]

There's am active thread on how to get it to work for non-Linux systems. huggingface/diffusers#532

That looks interesting. I couldn't get it to work on MPS, but curious to see if it works on CUDA.
NotImplementedError: No operator found for this attention: AttentionOpDispatch(dtype=torch.float32, device=device(type='mps', index=0), k=40, has_dropout=False, attn_bias_type=<class 'NoneType'>, kv_len=4096, q_len=4096)

[tildebyte]

@Any-Winter-4079;

I'm guessing that's because, as I said, "there's a CUDA-side c++ binary library" which has to be built per-platform; I doubt that it has been built for macOS/M* systems yet.

[Any-Winter-4079]

@tildebyte I was testing https://github.com/huggingface/diffusers/pull/532/files, with
pip install git+https://github.com/facebookresearch/xformers@51dd119#egg=xformers
and
export USE_MEMORY_EFFICIENT_ATTENTION=1
plus updating attention.py but got the error from above.

That seemed to be another repo different from the C++ one initially mentioned (this one is using xformers), although I might be tripping and it might internally be using the C++ lib, cause it's very late here, and I'm half awake/half asleep already :)

[tildebyte]

There's a "Fast Attention" as well, different thing...

[blefaudeux]

There's a "Fast Attention" as well, different thing...

No it’s the same idea actually, the xformers repo dispatches to the flash attention kernels when applicable, else uses it’s own kernels

[i3oc9i]

@Any-Winter-4079
Here is an implemetation of xformer from facebook research
https://github.com/facebookresearch/xformers

looking at the README, it can be compiled on MacOs. Read build from source section.
MACOSX_DEPLOYMENT_TARGET=10.9 CC=clang CXX=clang++ pip install -e .

[Any-Winter-4079]

I installed w/ MACOSX_DEPLOYMENT_TARGET=10.9 CC=clang CXX=clang++ pip install -e git+https://github.com/facebookresearch/xformers.git#egg=xformers without doing any of their environment steps. Hope it works!
Note I already had the requirements from requirements.txt
Now let's try to use it...

---

Update: I could run the benchmark up to 100% with
cd src/xformers/xformers
python3 benchmarks/benchmark_encoder.py --activations relu --plot -emb 256 -bs 32 -heads 16
after running pip install -r requirements-benchmark.txt and commenting out triton == 2.0.0.dev20220701 from requirements-test.txt and running pip install seaborn

But it says cpu so that's not good most likely
32 1024 256 True cpu pooling
32 1024 256 True cpu lambda

The error with triton was

ERROR: Could not find a version that satisfies the requirement triton==2.0.0.dev20220701 (from versions: 0.1, 0.1.1, 0.1.2, 0.1.3, 0.2.0, 0.2.1, 0.2.2, 0.2.3, 0.3.0)
ERROR: No matching distribution found for triton==2.0.0.dev20220701

and pip install triton==0.3.0 gives: RuntimeError: Could not find llvm-config. Please install llvm-{8, 9, 10}-dev, which I can't seem to be able to install (or know how to install) correctly.

And omitting triton installation and trying to run adding MemoryEfficientCrossAttention to attention.py from https://github.com/huggingface/diffusers/pull/532/files:
NotImplementedError: No operator found for this attention: AttentionOpDispatch(dtype=torch.float32, device=device(type='mps', index=0), k=40, has_dropout=False, attn_bias_type=<class 'NoneType'>, kv_len=4096, q_len=4096, kv=40, batch_size=16, num_heads=1)

@i3oc9i Did you manage to use it?

[Any-Winter-4079]

Update to say I tried ignoring the triton part, and searching a bit where the AttentionOpDispatch error comes from, it seems to come from import xformers.ops and ops.py file. There is a bunch of cuda stuff there so not sure how we would run it with mps.

Commenting out / replacing some cuda stuff, the furthest I could come is
NotImplementedError: Could not run 'xformers::efficient_attention_forward_cutlass' with arguments from the 'CPU' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'xformers::efficient_attention_forward_cutlass' is only available for these backends: [Metal, UNKNOWN_TENSOR_TYPE_ID, QuantizedXPU, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, SparseCPU, SparseCUDA, SparseHIP, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, SparseVE, UNKNOWN_TENSOR_TYPE_ID, NestedTensorCUDA, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID, UNKNOWN_TENSOR_TYPE_ID].

A somewhat similar problem is reported here facebookresearch/xformers#489 (but for cuda)

@mauwii or someone with m1, could you take a look?