CUDA: generalize FP16 fattn vec kernel #7061
Conversation
|
This Pr Main branch For reference |
|
There is no ROCm FlashAttention implementation on master. The only reason the test is passing is that |
|
📈 llama.cpp server for bench-server-baseline on Standard_NC4as_T4_v3 for Expand details for performance related PR only
More---
config:
xyChart:
titleFontSize: 12
width: 900
height: 600
themeVariables:
xyChart:
titleColor: "#000000"
---
xychart-beta
title "llama.cpp bench-server-baseline on Standard_NC4as_T4_v3
duration=10m 557 iterations"
y-axis "llamacpp:prompt_tokens_seconds"
x-axis "llamacpp:prompt_tokens_seconds" 1715256904 --> 1715257530
line [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 381.62, 381.62, 381.62, 381.62, 381.62, 422.36, 422.36, 422.36, 422.36, 422.36, 474.0, 474.0, 474.0, 474.0, 474.0, 536.01, 536.01, 536.01, 536.01, 536.01, 559.47, 559.47, 559.47, 559.47, 559.47, 567.91, 567.91, 567.91, 567.91, 567.91, 591.48, 591.48, 591.48, 591.48, 591.48, 612.77, 612.77, 612.77, 612.77, 612.77, 638.65, 638.65, 638.65, 638.65, 638.65, 643.97, 643.97, 643.97, 643.97, 643.97, 669.36, 669.36, 669.36, 669.36, 669.36, 689.16, 689.16, 689.16, 689.16, 689.16, 717.83, 717.83, 717.83, 717.83, 717.83, 677.78, 677.78, 677.78, 677.78, 677.78, 682.48, 682.48, 682.48, 682.48, 682.48, 694.24, 694.24, 694.24, 694.24, 694.24, 711.37, 711.37, 711.37, 711.37, 711.37, 718.22, 718.22, 718.22, 718.22, 718.22, 726.76, 726.76, 726.76, 726.76, 726.76, 727.9, 727.9, 727.9, 727.9, 727.9, 733.53, 733.53, 733.53, 733.53, 733.53, 740.01, 740.01, 740.01, 740.01, 740.01, 740.97, 740.97, 740.97, 740.97, 740.97, 743.4, 743.4, 743.4, 743.4, 743.4, 761.54, 761.54, 761.54, 761.54, 761.54, 760.26, 760.26, 760.26, 760.26, 760.26, 761.67, 761.67, 761.67, 761.67, 761.67, 763.48, 763.48, 763.48, 763.48, 763.48, 772.56, 772.56, 772.56, 772.56, 772.56, 771.82, 771.82, 771.82, 771.82, 771.82, 774.36, 774.36, 774.36, 774.36, 774.36, 782.99, 782.99, 782.99, 782.99, 782.99, 782.43, 782.43, 782.43, 782.43, 782.43, 797.79, 797.79, 797.79, 797.79, 797.79, 801.21, 801.21, 801.21, 801.21, 801.21, 799.82, 799.82, 799.82, 799.82, 799.82, 799.97, 799.97, 799.97, 799.97, 799.97, 803.67, 803.67, 803.67, 803.67, 803.67, 801.64, 801.64, 801.64, 801.64, 801.64, 807.95, 807.95, 807.95, 807.95, 807.95, 803.3, 803.3, 803.3, 803.3, 803.3, 803.7, 803.7, 803.7, 803.7, 803.7, 802.88, 802.88, 802.88, 802.88, 802.88, 802.13, 802.13, 802.13, 802.13, 802.13, 806.87, 806.87, 806.87, 806.87, 806.87, 810.68, 810.68, 810.68, 810.68, 810.68, 811.47, 811.47, 811.47, 811.47, 811.47, 816.64, 816.64, 816.64, 816.64, 816.64, 817.12, 817.12, 817.12, 817.12, 817.12, 822.36, 822.36, 822.36, 822.36, 822.36, 822.54, 822.54, 822.54, 822.54, 822.54, 822.34, 822.34, 822.34, 822.34, 822.34, 828.93, 828.93, 828.93, 828.93, 828.93, 830.1, 830.1, 830.1, 830.1, 830.1, 830.61, 830.61, 830.61, 830.61, 830.61, 831.91, 831.91, 831.91, 831.91, 831.91, 830.75, 830.75, 830.75, 830.75, 830.75, 831.98, 831.98, 831.98, 831.98, 831.98, 833.57, 833.57, 833.57, 833.57, 833.57, 833.79, 833.79, 833.79, 833.79, 833.79, 833.89, 833.89, 833.89, 833.89]
More---
config:
xyChart:
titleFontSize: 12
width: 900
height: 600
themeVariables:
xyChart:
titleColor: "#000000"
---
xychart-beta
title "llama.cpp bench-server-baseline on Standard_NC4as_T4_v3
duration=10m 557 iterations"
y-axis "llamacpp:predicted_tokens_seconds"
x-axis "llamacpp:predicted_tokens_seconds" 1715256904 --> 1715257530
line [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 42.2, 42.2, 42.2, 42.2, 42.2, 43.62, 43.62, 43.62, 43.62, 43.62, 33.3, 33.3, 33.3, 33.3, 33.3, 33.93, 33.93, 33.93, 33.93, 33.93, 34.15, 34.15, 34.15, 34.15, 34.15, 35.78, 35.78, 35.78, 35.78, 35.78, 36.37, 36.37, 36.37, 36.37, 36.37, 36.78, 36.78, 36.78, 36.78, 36.78, 37.02, 37.02, 37.02, 37.02, 37.02, 36.83, 36.83, 36.83, 36.83, 36.83, 36.7, 36.7, 36.7, 36.7, 36.7, 36.4, 36.4, 36.4, 36.4, 36.4, 35.02, 35.02, 35.02, 35.02, 35.02, 34.14, 34.14, 34.14, 34.14, 34.14, 34.23, 34.23, 34.23, 34.23, 34.23, 34.69, 34.69, 34.69, 34.69, 34.69, 33.98, 33.98, 33.98, 33.98, 33.98, 34.0, 34.0, 34.0, 34.0, 34.0, 33.78, 33.78, 33.78, 33.78, 33.78, 33.7, 33.7, 33.7, 33.7, 33.7, 33.55, 33.55, 33.55, 33.55, 33.55, 33.56, 33.56, 33.56, 33.56, 33.56, 33.48, 33.48, 33.48, 33.48, 33.48, 33.51, 33.51, 33.51, 33.51, 33.51, 33.49, 33.49, 33.49, 33.49, 33.49, 33.02, 33.02, 33.02, 33.02, 33.02, 32.99, 32.99, 32.99, 32.99, 32.99, 33.27, 33.27, 33.27, 33.27, 33.27, 33.44, 33.44, 33.44, 33.44, 33.44, 33.44, 33.44, 33.44, 33.44, 33.44, 33.55, 33.55, 33.55, 33.55, 33.55, 33.56, 33.56, 33.56, 33.56, 33.56, 33.47, 33.47, 33.47, 33.47, 33.47, 33.25, 33.25, 33.25, 33.25, 33.25, 32.75, 32.75, 32.75, 32.75, 32.75, 32.75, 32.75, 32.75, 32.75, 32.75, 32.92, 32.92, 32.92, 32.92, 32.92, 32.97, 32.97, 32.97, 32.97, 32.97, 33.01, 33.01, 33.01, 33.01, 33.01, 33.13, 33.13, 33.13, 33.13, 33.13, 32.91, 32.91, 32.91, 32.91, 32.91, 32.35, 32.35, 32.35, 32.35, 32.35, 32.23, 32.23, 32.23, 32.23, 32.23, 30.59, 30.59, 30.59, 30.59, 30.59, 30.52, 30.52, 30.52, 30.52, 30.52, 30.53, 30.53, 30.53, 30.53, 30.53, 30.62, 30.62, 30.62, 30.62, 30.62, 30.74, 30.74, 30.74, 30.74, 30.74, 30.88, 30.88, 30.88, 30.88, 30.88, 30.88, 30.88, 30.88, 30.88, 30.88, 30.86, 30.86, 30.86, 30.86, 30.86, 30.76, 30.76, 30.76, 30.76, 30.76, 30.68, 30.68, 30.68, 30.68, 30.68, 30.77, 30.77, 30.77, 30.77, 30.77, 30.93, 30.93, 30.93, 30.93, 30.93, 31.01, 31.01, 31.01, 31.01, 31.01, 31.09, 31.09, 31.09, 31.09, 31.09, 31.13, 31.13, 31.13, 31.13, 31.13, 31.19, 31.19, 31.19, 31.19, 31.19, 31.15, 31.15, 31.15, 31.15, 31.15, 31.17, 31.17, 31.17, 31.17]
Details
More---
config:
xyChart:
titleFontSize: 12
width: 900
height: 600
themeVariables:
xyChart:
titleColor: "#000000"
---
xychart-beta
title "llama.cpp bench-server-baseline on Standard_NC4as_T4_v3
duration=10m 557 iterations"
y-axis "llamacpp:kv_cache_usage_ratio"
x-axis "llamacpp:kv_cache_usage_ratio" 1715256904 --> 1715257530
line [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.11, 0.11, 0.11, 0.11, 0.11, 0.45, 0.45, 0.45, 0.45, 0.45, 0.14, 0.14, 0.14, 0.14, 0.14, 0.18, 0.18, 0.18, 0.18, 0.18, 0.13, 0.13, 0.13, 0.13, 0.13, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.12, 0.12, 0.12, 0.12, 0.12, 0.21, 0.21, 0.21, 0.21, 0.21, 0.16, 0.16, 0.16, 0.16, 0.16, 0.16, 0.16, 0.16, 0.16, 0.16, 0.27, 0.27, 0.27, 0.27, 0.27, 0.32, 0.32, 0.32, 0.32, 0.32, 0.14, 0.14, 0.14, 0.14, 0.14, 0.13, 0.13, 0.13, 0.13, 0.13, 0.17, 0.17, 0.17, 0.17, 0.17, 0.21, 0.21, 0.21, 0.21, 0.21, 0.18, 0.18, 0.18, 0.18, 0.18, 0.19, 0.19, 0.19, 0.19, 0.19, 0.16, 0.16, 0.16, 0.16, 0.16, 0.17, 0.17, 0.17, 0.17, 0.17, 0.32, 0.32, 0.32, 0.32, 0.32, 0.08, 0.08, 0.08, 0.08, 0.08, 0.16, 0.16, 0.16, 0.16, 0.16, 0.35, 0.35, 0.35, 0.35, 0.35, 0.15, 0.15, 0.15, 0.15, 0.15, 0.11, 0.11, 0.11, 0.11, 0.11, 0.08, 0.08, 0.08, 0.08, 0.08, 0.1, 0.1, 0.1, 0.1, 0.1, 0.16, 0.16, 0.16, 0.16, 0.16, 0.14, 0.14, 0.14, 0.14, 0.14, 0.16, 0.16, 0.16, 0.16, 0.16, 0.11, 0.11, 0.11, 0.11, 0.11, 0.28, 0.28, 0.28, 0.28, 0.28, 0.28, 0.28, 0.28, 0.28, 0.28, 0.18, 0.18, 0.18, 0.18, 0.18, 0.11, 0.11, 0.11, 0.11, 0.11, 0.09, 0.09, 0.09, 0.09, 0.09, 0.1, 0.1, 0.1, 0.1, 0.1, 0.35, 0.35, 0.35, 0.35, 0.35, 0.55, 0.55, 0.55, 0.55, 0.55, 0.63, 0.63, 0.63, 0.63, 0.63, 0.6, 0.6, 0.6, 0.6, 0.6, 0.22, 0.22, 0.22, 0.22, 0.22, 0.14, 0.14, 0.14, 0.14, 0.14, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.2, 0.2, 0.2, 0.2, 0.2, 0.13, 0.13, 0.13, 0.13, 0.13, 0.23, 0.23, 0.23, 0.23, 0.23, 0.19, 0.19, 0.19, 0.19, 0.19, 0.12, 0.12, 0.12, 0.12, 0.12, 0.09, 0.09, 0.09, 0.09, 0.09, 0.15, 0.15, 0.15, 0.15, 0.15, 0.08, 0.08, 0.08, 0.08, 0.08, 0.14, 0.14, 0.14, 0.14, 0.14, 0.12, 0.12, 0.12, 0.12, 0.12, 0.12, 0.12, 0.12, 0.12, 0.12, 0.17, 0.17, 0.17, 0.17, 0.17, 0.16, 0.16, 0.16, 0.16, 0.16, 0.19, 0.19, 0.19, 0.19]
More---
config:
xyChart:
titleFontSize: 12
width: 900
height: 600
themeVariables:
xyChart:
titleColor: "#000000"
---
xychart-beta
title "llama.cpp bench-server-baseline on Standard_NC4as_T4_v3
duration=10m 557 iterations"
y-axis "llamacpp:requests_processing"
x-axis "llamacpp:requests_processing" 1715256904 --> 1715257530
line [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.0, 3.0, 3.0, 3.0, 3.0, 7.0, 7.0, 7.0, 7.0, 7.0, 1.0, 1.0, 1.0, 1.0, 1.0, 7.0, 7.0, 7.0, 7.0, 7.0, 5.0, 5.0, 5.0, 5.0, 5.0, 4.0, 4.0, 4.0, 4.0, 4.0, 3.0, 3.0, 3.0, 3.0, 3.0, 5.0, 5.0, 5.0, 5.0, 5.0, 4.0, 4.0, 4.0, 4.0, 4.0, 2.0, 2.0, 2.0, 2.0, 2.0, 6.0, 6.0, 6.0, 6.0, 6.0, 8.0, 8.0, 8.0, 8.0, 8.0, 3.0, 3.0, 3.0, 3.0, 3.0, 5.0, 5.0, 5.0, 5.0, 5.0, 7.0, 7.0, 7.0, 7.0, 7.0, 6.0, 6.0, 6.0, 6.0, 6.0, 3.0, 3.0, 3.0, 3.0, 3.0, 7.0, 7.0, 7.0, 7.0, 7.0, 6.0, 6.0, 6.0, 6.0, 6.0, 8.0, 8.0, 8.0, 8.0, 8.0, 3.0, 3.0, 3.0, 3.0, 3.0, 5.0, 5.0, 5.0, 5.0, 5.0, 4.0, 4.0, 4.0, 4.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0, 6.0, 6.0, 6.0, 6.0, 6.0, 4.0, 4.0, 4.0, 4.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0, 2.0, 2.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 3.0, 3.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 5.0, 5.0, 5.0, 5.0, 5.0, 8.0, 8.0, 8.0, 8.0, 8.0, 7.0, 7.0, 7.0, 7.0, 7.0, 2.0, 2.0, 2.0, 2.0, 2.0, 4.0, 4.0, 4.0, 4.0, 4.0, 2.0, 2.0, 2.0, 2.0, 2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 5.0, 5.0, 5.0, 5.0, 5.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 2.0, 2.0, 2.0, 2.0, 2.0, 6.0, 6.0, 6.0, 6.0, 6.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 2.0, 2.0, 2.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 5.0, 5.0, 5.0, 5.0, 5.0, 6.0, 6.0, 6.0, 6.0, 6.0, 7.0, 7.0, 7.0, 7.0, 7.0, 3.0, 3.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 4.0, 4.0, 3.0, 3.0, 3.0, 3.0, 3.0, 7.0, 7.0, 7.0, 7.0, 7.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 7.0, 7.0, 7.0, 7.0, 7.0, 6.0, 6.0, 6.0, 6.0, 6.0, 3.0, 3.0, 3.0, 3.0]
|
|
I get an error if I try |
|
The issue was that the logic for determining when to compile the kernel was incorrect. Does it work now? |
Now it runs, but the output is broken when FlashAttention is enabled and the --batch-size 16 --batch-size 8 |
|
Thank you for the bug report, the issue was that I wrote back the data in the wrong order for batch sizes 2-8 and I didn't notice because those batch sizes were not being used when I tested for correctness. It should now be fixed. |
|
Can confirm this pr is quite a bit slower for my workload. I have another question. Would there be a workaround for GPU like P40 i.e a kernel for FP32 or it's just not worth it. |
This PR works perfect on my P40. |
but it's gonna slow right? P40 is pretty slow on FP16 as it only support int8/fp32. |
The current PR will produce correct results but for usable performance P40s will need a dedicated FP32 kernel which I will also add. |
At that time, I did not try to quantify the difference in quality because comparing the two versions side by side seemed to produce very similar results. If what you said is true for me, it should translate in a higher perplexity, right? Here what I get: My PR: Same result. This is with a random Q6_K model over wikitext. About performance, here is what I get. My PR:
This PR:
|
so with FA, it got better PPL consistently . |
JohannesGaessler
force-pushed
the
cuda-fa-no-tc-5
branch
from
e4643c9 to
57bde8c
Compare
|
@jdecourval thank you for the data. I'll still try to write a kernel without any tensor cores that is more optimized for large batch sizes first but given this data I think that rocWMMA is a viable option. @sorasoras I don't think that enabling/disabling FlashAttention should significantly affect results since the differences should just be due to floating point rounding error. |
RDNA3 is gonna slowdown, Any there any plan for this? |
Just don't use FlashAttention until it's faster? |
|
Is there something in this PR that could cause a slowdown in TG?
There are also some warnings: |
|
Thank you for pointing this out, I did not think to test this but the new kernel is for whatever reason massively slower than the one on master for a batch size of 1. According to NSight Compute the kernel runtime that I used for testing increased from 70 µs to 92 µs. I really don't understand why this is happening. The compiler should be able to unroll all loops over Should we just keep both kernel versions? |
|
Personally I am not interested about the hardware that benefits from this change, and I don't expect to make any changes to this code, so if you want to maintain another copy of the flash attn kernel that's entirely up to you. It would be nice to keep the same performance with tensor cores as in master, but IMO the difference in overall performance is small enough that it could be ignored. More results with the latest commit:
|
If you start reverting the |
|
Intuitively I would think the issue has to do with the data types (array vs. regular variables) rather than the loops. |
|
FWIW, I tried having a quick look at the performance on my 7900xtx, was able to bring it back to the level of your previous kernel, by removing most checks for AMD, which meant re-enabling rocwmma. Compare with above:
It looks like this (pile of hacks): jdecourval@b82649d |
|
I can confirm that the performance regression is caused by In any case, I think this is worth asking my contact at NVIDIA about. This could be a legitimate compiler bug. |
JohannesGaessler
force-pushed
the
cuda-fa-no-tc-5
branch
from
78ee06e to
fece1fe
Compare
|
From what I can tell the issue is caused specifically |
mofosyne
added
the
Review Complexity : High
|
I think so.
|
This PR adds a FlashAttention kernel that only uses regular FP16 arithmetic and not any tensor core operations by generalizing the kernel I wrote for batch size 1 to larger batch sizes. It works reasonably well for batch sizes <= 8. The target hardware is the NVIDIA P100 and AMD RX 5000/6000 GPUs although the code also makes it possible to run FlashAttention on AMD RX 7000 (with likely much worse performance) and on other Pascal GPUs (with effectively unusable performance). On my RX 6800 the performance changes as follows:
This PR also rearranges the order of some definitions in
ggml-cuda/common.cuh. This is because on masterNO_DEVICE_CODEis broken on AMD due to the definitions being in the wrong order. I also implemented some warp reduce functions that didn't work on AMD.@jdecourval when you worked on #6773 , did you check that the code actually produces correct results? I encountered an issue where due to missing implementations in
ggml-cuda/common.cuhthe compiler would optimize out part of the kernel which resulted in significantly faster but useless code.