Fix slow tests

[NouamaneTazi]

Fixes slow tests after the PR #371

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[Birch-san]

@NouamaneTazi @patrickvonplaten, do you remember any details about why this needed to be reverted?

by my measurements (admittedly on CompVis): reintroducing baddbmm would make inference 19% faster on Mac/MPS:
https://twitter.com/Birchlabs/status/1588703520433541120

was it the case that tests/test_pipelines.py::PipelineTesterMixin::test_stable_diffusion_memory_chunking went slower? did it do so on all platforms, or just on certain hardware (e.g. CUDA vs MPS, high vs low VRAM)? or this is only a problem when sliced attention is used, and we should prefer to use baddbmm for unsliced attention?

[patrickvonplaten]

Hey @Birch-san,

Good question! Now that we have strong unet testing we could maybe re-open this PR to see if it improves overall performance.
However we need to be sure that it doesn't slow down any other use case such as ONNX, FP16 Torch, ...

What are your thoughts here @anton-l @pcuenca ?

[patrickvonplaten]

And @NouamaneTazi ?

[Birch-san]

Ah, lemme correct myself. the 19% speed increase was the outcome of 2 changes:

• changing the (q,k) matmul to baddbmm (as shown in this PR)
• changing the (attn, v) matmul below to a bmm(). not featured in this PR. the batched matmul was way faster than regular matmul on MPS.

[patil-suraj]

Think we could open a PR and do through tests for other cases as well. @NouamaneTazi would you like to tackle this ?

[pcuenca]

I agree we should open a new PR and explore this again :)

[Birch-san]

@patrickvonplaten @pcuenca @patil-suraj @NouamaneTazi
I've restored the original baddbmm, added a new bmm, implemented for sliced attention, added fast-path to Decoder:
#1203