rescuing flash attention #386

Green-Sky · 2024-09-01T13:08:15Z

I ported over the flash attention code changes form #221 .
~~This does not yet fix the existing code behind the define, which broke when gg removed the code from ggml assuming noone was using it :)~~
~~The old broken code behind the define also only gets invoked by the VAE.~~
The new code gets used by UNET, DiT and clip/t5. However I have not come across any clip or t5 that just works without extra work, but there are sizable memory reductions for ~~sd1,~~ sd2, sdxl and flux compute buffers, as well as a speed increase on cuda.
I switched VAE over to using the new attention code, but flash attention makes no difference, so it is disabled there.
There is alot more left on the table if we employ padding, so flash attention can be applied to more ops.

Flash attention for diffusion models is available via a runtime flag --diffusion-fa.

flux1-schnell performance numbers

all tests done on a debug build, so the real number are likely a bit better.
4 steps, cfg-scale of 1

CPU

512x512:

compute buffer size: 397MB -> 247MB
speed q3_k: 80.00s/it -> 83.38s/it

768x768:

compute buffer size: 1103MB -> 503MB
speed q3_k: 197.11s/it -> 209.16s/it

CUDA

512x512:

compute buffer size: 398MB -> 248MB
speed q3_k: 1.98s/it -> 1.69s/it
speed q4_k: 1.84s/it -> 1.54s/it

768x768:

compute buffer size: 1105MB -> 505MB
speed q3_k: 4.58s/it -> 3.58s/it
speed q4_k: OOM -> 3.29s/it

direct comparison of the cuda images:

nofa	fa

difference, darker/colorful is worse:

error	error^2

SD2 turbo

8 steps, cfg-scale of 1

CPU

512x512:

compute buffer size: 367MB -> 130MB
speed q8_0: 6.40s/it -> 6.65s/it

768x768:

compute buffer size: 1718MB -> 294MB
speed q8_0: 20.59s/it -> 22.41s/it

CUDA

512x512:

compute buffer size: 367MB -> 130MB
speed q8_0: 6.24it/s -> 8.17it/s

768x768:

compute buffer size: 1718MB -> 295MB
speed q8_0: 1.84it/s -> 3.17it/s

direct comparison of the cuda images:

nofa	fa

difference, darker/colorful is worse:

error	error^2

SDXL realvisxl_v50Lightning

6 steps, cfg-scale of 1.8
dpm++2mv2 karras

CPU

512x512:

compute buffer size: 131MB -> 131MB
speed q8_0: 15.04s/it -> 15.21s/it

768x768:

compute buffer size: 330MB -> 280MB
speed q8_0: 37.36s/it -> 39.30s/it

CUDA

512x512:

compute buffer size: 132MB -> 132MB
speed q8_0: 2.01it/s -> 2.36it/s

768x768:

compute buffer size: 331MB -> 280MB
speed q8_0: 1.23s/it -> 1.15s/it

direct comparison of the cuda images:

nofa	fa

difference, darker/colorful is worse:

error	error^2

There still is the opportunity to pad some tensors to make them fit.

TODO

Please test this code.

props to @FSSRepo for having the code laying around

fixes: #297

images:
sd2_turbo.zip
flux1-schnell-q4_k.zip
flux1-schnell-q3_k.zip

update: added sd2 and sdxl numbers
update2: see rocm tests down in the thread. it works but behaves similar to cpu
udpate3: added flash_attn param and expose via --diffusion-fa runtime flag for supported models.

theaerotoad · 2024-09-03T22:58:26Z

@Green-Sky -- would this enable flash attention for Vulkan builds as well?

Green-Sky · 2024-09-04T07:32:58Z

@Green-Sky -- would this enable flash attention for Vulkan builds as well?

No, sadly not. Vulkan does not implement GGML_OP_FLASH_ATTN_EXT.
However, it looks like cuda, cuda built as rocm (and musa?) and metal all support it.

Would be cool if someone could try rocm and metal builds.

Green-Sky · 2024-09-06T08:54:46Z

I just tested SD3 (2B), and it appears the kv dimensions are not multiples of 256, so flash attention wont work on the mmdit without padding.

MineGame159 · 2024-09-06T13:53:53Z

I wanted to test the PR with the ROCM backend but I couldn't get it to build. Did you somehow add back the missing ggml_flash_attn function that is used in ggml_extend.hpp? I am not familiar with ggml so it would be nice if you could explain how you got the CUDA backend working.

Green-Sky · 2024-09-06T13:58:07Z

I wanted to test the PR with the ROCM backend but I couldn't get it to build. Did you somehow add back the missing ggml_flash_attn function that is used in ggml_extend.hpp? I am not familiar with ggml so it would be nice if you could explain how you got the CUDA backend working.

Oh no, sorry for the confusion, dont use the old define or cmake option, I did not add that back in.
Just build it as-is and it has flash attention (for diffusion models only) enabled.
The old code that gets enabled with the define is what would get used by VAE, but I have not touched that part yet.

edit: I might push a change soon, where you can enable flash attention for diffusion models via a command line option.

MineGame159 · 2024-09-06T14:06:34Z

Oh interesting. Because I already tried a comparison between the current master and your PR applied without the cmake flag. And the results didn't look like it was enabled. Here are the results (no Flash Attention is how I marked that I didn't enable the flag). I used the flux1-schnell model using q8_0.

Master - no Flash Attention

total params memory size = 21471.11MB (VRAM 12152.28MB, RAM 9318.83MB)
sampling completed, taking 149.66s - 18.69s/it

PR - no Flash Attention

total params memory size = 21471.11MB (VRAM 12152.28MB, RAM 9318.83MB)
sampling completed, taking 173.13s - 21.62s/it

Green-Sky · 2024-09-06T14:11:05Z

@MineGame159 I see, so speed went down on rocm...
Also, run it with -v and look for

[DEBUG] ggml_extend.hpp:739  - attention_ext L_q:2304 L_k:2304 n_head:24 C:3072 d_head:128 N:1
[DEBUG] ggml_extend.hpp:763  - using flash attention

If an attention_ext line with tensor sizes is follows by using flash attention, then that specific attention was converted to a flash attention.

Also look for flux compute buffer size: 456.75 MB(VRAM) (in the spam with -v) to see how the compute buffer is affected. Total params size does not change with this patch. 😃

MineGame159 · 2024-09-06T14:23:58Z

Oh sorry 😅, I am pretty new to sd.cpp and AI as a whole. And yes, I can see that it is using flash attention.

[DEBUG] ggml_extend.hpp:738  - attention_ext L_k:1280 n_head:24 C:3072 d_head:128
[DEBUG] ggml_extend.hpp:755  - using flash attention

And the compute buffer size:
Master - flux compute buffer size: 398.50 MB(VRAM)
PR - flux compute buffer size: 248.50 MB(VRAM)

Green-Sky · 2024-09-06T14:26:01Z

And the resulting image look close to identical?
The numbers look good, the compute buffer size reduction is the same as on cuda. Try larger images :)

MineGame159 · 2024-09-06T15:18:07Z

Tested with a 768x768 image and I can't see any difference. The size of the compute buffer is basically the same as on CUDA for you. Only thing is the decreased speed.

1105.07 MB(VRAM) -> 505.07 MB(VRAM)
41.66s/it -> 57.59s/it

Green-Sky · 2024-09-06T19:19:41Z

I diffed your images.

error	error^2

The error is even less then what I have with cuda.

Tested with a 768x768 image and I can't see any difference. The size of the compute buffer is basically the same as on CUDA for you. Only thing is the decreased speed.

Yea, good to know, thanks for testing.
In summery: use it if you use cuda and/or need the extra vram savings.

Green-Sky · 2024-09-07T10:09:40Z

edit: I might push a change soon, where you can enable flash attention for diffusion models via a command line option.

I did.
I am now looking into vae.

update: flash attention does not seem to improve VAE in anyway + it does not work on cuda, since cuda supports d_head up to 256, but VAE needs 512. So I think I will not enable VAE flash attention, but still remove the old crusty code.

Green-Sky · 2024-09-07T11:39:39Z

Would be nice if someone could test metal.

Green-Sky · 2024-09-08T07:48:12Z

Accelerated memory-efficient CPU inference
- Only requires ~2.3GB when using txt2img with fp16 precision to generate a 512x512 image, enabling Flash Attention just requires ~1.8GB.

imma leave @leejet to update this example, not sure where those numbers came from.

Green-Sky · 2024-09-08T16:57:00Z

If someone wants to play with kv-padding, to enable more cases where flash attention can be used, here is a patch:

diff --git a/ggml_extend.hpp b/ggml_extend.hpp
index 8452a0b..518eb6f 100644
--- a/ggml_extend.hpp
+++ b/ggml_extend.hpp
@@ -710,6 +710,7 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention_ext(struct ggml_context*

     float scale = (1.0f / sqrt((float)d_head));

+    int kv_pad = 0;
     //if (flash_attn) {
     //    LOG_DEBUG("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
     //}
@@ -717,7 +718,14 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention_ext(struct ggml_context*
     GGML_ASSERT(((L_k % 256 == 0) && L_q == L_k) || !(L_k % 256 == 0));

     bool can_use_flash_attn = true;
-    can_use_flash_attn = can_use_flash_attn && L_k % 256 == 0;
+    if (can_use_flash_attn && L_k % 256 != 0) {
+        if (L_k == 77) {
+            kv_pad = 256 - (L_k % 256);
+        } else {
+            can_use_flash_attn = false;
+        }
+    }
+    //can_use_flash_attn = can_use_flash_attn && L_k % 256 == 0;
     can_use_flash_attn = can_use_flash_attn && d_head % 64 == 0; // double check

     // cuda max d_head seems to be 256, cpu does seem to work with 512
@@ -734,11 +742,18 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention_ext(struct ggml_context*
     ggml_tensor* kqv = nullptr;
     //GGML_ASSERT((flash_attn && can_use_flash_attn) || !flash_attn);
     if (can_use_flash_attn && flash_attn) {
-        //LOG_DEBUG("using flash attention");
+        LOG_DEBUG("using flash attention");
+        if (kv_pad != 0) {
+            LOG_DEBUG("padding kv by %d", kv_pad);
+            k = ggml_pad(ctx, k, 0, kv_pad, 0, 0);
+        }
         k = ggml_cast(ctx, k, GGML_TYPE_F16);

         v = ggml_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3));  // [N, n_head, L_k, d_head]
         v = ggml_reshape_3d(ctx, v, d_head, L_k, n_head * N);  // [N * n_head, L_k, d_head]
+        if (kv_pad != 0) {
+            v = ggml_pad(ctx, v, 0, kv_pad, 0, 0);
+        }
         v = ggml_cast(ctx, v, GGML_TYPE_F16);

         kqv = ggml_flash_attn_ext(ctx, q, k, v, mask, scale, 0, 0);

I wont be committing this here since the memory saving are minuscule, or it uses even more memory because of the padding.
Plus, the images look way off, so maybe there is a mistake somewhere or it's just the f16 cast in the wrong situation...

Green-Sky · 2024-10-25T14:21:59Z

rebased on master, did not yet look at sd3.5, it is disabled, same as sd3 for now.
I still consider this ready for merge as-is, I will probably look at sd3.5 in the coming weeks.

this does not fix the currently broken fa behind the define, which is only used by VAE Co-authored-by: FSSRepo <[email protected]>

no support for sd3 or video

leejet · 2024-11-23T04:38:06Z

Thank you for your contribution

Green-Sky changed the title ~~rescue flash attention~~ rescueing flash attention Sep 1, 2024

Green-Sky changed the title ~~rescueing flash attention~~ rescuing flash attention Sep 1, 2024

Green-Sky force-pushed the rescue_flash_attn branch from 2634bea to f88143f Compare September 3, 2024 07:41

Green-Sky force-pushed the rescue_flash_attn branch from b0aecf4 to e904b86 Compare September 7, 2024 11:04

Green-Sky marked this pull request as ready for review September 7, 2024 11:28

Green-Sky force-pushed the rescue_flash_attn branch from f69fbfd to cc7efa2 Compare September 7, 2024 11:42

Green-Sky force-pushed the rescue_flash_attn branch from c943406 to 90d420a Compare September 8, 2024 07:49

This was referenced Sep 12, 2024

SD_FLASH_ATTN for Mac, Linux, and non-CUBLAS Windows builds? #407

Closed

metal-flash-attention support #22

Open

Green-Sky force-pushed the rescue_flash_attn branch from 90d420a to 947340b Compare October 25, 2024 14:14

Green-Sky force-pushed the rescue_flash_attn branch from 947340b to a06ec06 Compare October 25, 2024 14:24

Green-Sky and others added 3 commits November 12, 2024 13:03

repair flash attention in _ext

652f58f

this does not fix the currently broken fa behind the define, which is only used by VAE Co-authored-by: FSSRepo <[email protected]>

make flash attention in the diffusion model a runtime flag

cc65d7e

no support for sd3 or video

remove old flash attention option and switch vae over to attn_ext

78aeee8

update docs

cbf0489

Green-Sky force-pushed the rescue_flash_attn branch from a06ec06 to cbf0489 Compare November 12, 2024 12:03

rmatif mentioned this pull request Nov 17, 2024

Flash Attention CPU-only build fails in cross-compilation for Android, succeeds with Vulkan backend #463

Open

Green-Sky mentioned this pull request Nov 21, 2024

Feat: add Flux 1 Lite 8B (Freepik) support #474

Merged

leejet added 2 commits November 23, 2024 12:04

Merge branch 'master' into rescue_flash_attn

6600b89

format code

5ef1821

leejet merged commit 1c168d9 into leejet:master Nov 23, 2024
9 checks passed

Green-Sky deleted the rescue_flash_attn branch November 23, 2024 12:46

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

rescuing flash attention #386

rescuing flash attention #386

Green-Sky commented Sep 1, 2024 •

edited

Loading

theaerotoad commented Sep 3, 2024 •

edited

Loading

Green-Sky commented Sep 4, 2024 •

edited

Loading

Green-Sky commented Sep 6, 2024

MineGame159 commented Sep 6, 2024

Green-Sky commented Sep 6, 2024 •

edited

Loading

MineGame159 commented Sep 6, 2024

Green-Sky commented Sep 6, 2024

MineGame159 commented Sep 6, 2024

Green-Sky commented Sep 6, 2024

MineGame159 commented Sep 6, 2024

Green-Sky commented Sep 6, 2024

Green-Sky commented Sep 7, 2024 •

edited

Loading

Green-Sky commented Sep 7, 2024

Green-Sky commented Sep 8, 2024

Green-Sky commented Sep 8, 2024

Green-Sky commented Oct 25, 2024

leejet commented Nov 23, 2024

rescuing flash attention #386

rescuing flash attention #386

Conversation

Green-Sky commented Sep 1, 2024 • edited Loading

flux1-schnell performance numbers

CPU

512x512:

768x768:

CUDA

512x512:

768x768:

SD2 turbo

CPU

512x512:

768x768:

CUDA

512x512:

768x768:

SDXL realvisxl_v50Lightning

CPU

512x512:

768x768:

CUDA

512x512:

768x768:

TODO

theaerotoad commented Sep 3, 2024 • edited Loading

Green-Sky commented Sep 4, 2024 • edited Loading

Green-Sky commented Sep 6, 2024

MineGame159 commented Sep 6, 2024

Green-Sky commented Sep 6, 2024 • edited Loading

MineGame159 commented Sep 6, 2024

Master - no Flash Attention

PR - no Flash Attention

Green-Sky commented Sep 6, 2024

MineGame159 commented Sep 6, 2024

Green-Sky commented Sep 6, 2024

MineGame159 commented Sep 6, 2024

Green-Sky commented Sep 6, 2024

Green-Sky commented Sep 7, 2024 • edited Loading

Green-Sky commented Sep 7, 2024

Green-Sky commented Sep 8, 2024

Green-Sky commented Sep 8, 2024

Green-Sky commented Oct 25, 2024

leejet commented Nov 23, 2024

Green-Sky commented Sep 1, 2024 •

edited

Loading

theaerotoad commented Sep 3, 2024 •

edited

Loading

Green-Sky commented Sep 4, 2024 •

edited

Loading

Green-Sky commented Sep 6, 2024 •

edited

Loading

Green-Sky commented Sep 7, 2024 •

edited

Loading