Fix flash-attn for AMD #6773
Closed
Fix flash-attn for AMD #6773
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
8 tasks
|
📈 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 554 iterations"
y-axis "llamacpp:prompt_tokens_seconds"
x-axis "llamacpp:prompt_tokens_seconds" 1713551797 --> 1713552421
line [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 370.76, 370.76, 370.76, 370.76, 370.76, 690.32, 690.32, 690.32, 690.32, 690.32, 685.99, 685.99, 685.99, 685.99, 685.99, 727.61, 727.61, 727.61, 727.61, 727.61, 788.87, 788.87, 788.87, 788.87, 788.87, 785.93, 785.93, 785.93, 785.93, 785.93, 783.55, 783.55, 783.55, 783.55, 783.55, 806.51, 806.51, 806.51, 806.51, 806.51, 811.08, 811.08, 811.08, 811.08, 811.08, 819.73, 819.73, 819.73, 819.73, 819.73, 819.37, 819.37, 819.37, 819.37, 819.37, 845.45, 845.45, 845.45, 845.45, 845.45, 853.93, 853.93, 853.93, 853.93, 853.93, 866.95, 866.95, 866.95, 866.95, 866.95, 814.63, 814.63, 814.63, 814.63, 814.63, 797.98, 797.98, 797.98, 797.98, 797.98, 799.75, 799.75, 799.75, 799.75, 799.75, 797.54, 797.54, 797.54, 797.54, 797.54, 808.86, 808.86, 808.86, 808.86, 808.86, 808.81, 808.81, 808.81, 808.81, 808.81, 808.43, 808.43, 808.43, 808.43, 808.43, 814.93, 814.93, 814.93, 814.93, 814.93, 819.42, 819.42, 819.42, 819.42, 819.42, 837.72, 837.72, 837.72, 837.72, 837.72, 836.82, 836.82, 836.82, 836.82, 836.82, 830.97, 830.97, 830.97, 830.97, 830.97, 844.77, 844.77, 844.77, 844.77, 844.77, 841.46, 841.46, 841.46, 841.46, 841.46, 840.42, 840.42, 840.42, 840.42, 840.42, 840.76, 840.76, 840.76, 840.76, 840.76, 845.35, 845.35, 845.35, 845.35, 845.35, 844.62, 844.62, 844.62, 844.62, 844.62, 845.7, 845.7, 845.7, 845.7, 845.7, 851.0, 851.0, 851.0, 851.0, 851.0, 865.77, 865.77, 865.77, 865.77, 865.77, 877.85, 877.85, 877.85, 877.85, 877.85, 873.99, 873.99, 873.99, 873.99, 873.99, 872.27, 872.27, 872.27, 872.27, 872.27, 874.52, 874.52, 874.52, 874.52, 874.52, 875.99, 875.99, 875.99, 875.99, 875.99, 881.04, 881.04, 881.04, 881.04, 881.04, 872.55, 872.55, 872.55, 872.55, 872.55, 860.13, 860.13, 860.13, 860.13, 860.13, 858.55, 858.55, 858.55, 858.55, 858.55, 855.49, 855.49, 855.49, 855.49, 855.49, 860.4, 860.4, 860.4, 860.4, 860.4, 862.58, 862.58, 862.58, 862.58, 862.58, 861.13, 861.13, 861.13, 861.13, 861.13, 862.37, 862.37, 862.37, 862.37, 862.37, 865.12, 865.12, 865.12, 865.12, 865.12, 867.84, 867.84, 867.84, 867.84, 867.84, 867.15, 867.15, 867.15, 867.15, 867.15, 873.08, 873.08, 873.08, 873.08, 873.08, 869.69, 869.69, 869.69, 869.69, 869.69, 870.15, 870.15, 870.15, 870.15, 870.15, 869.09, 869.09, 869.09, 869.09, 869.09, 870.33, 870.33, 870.33, 870.33, 870.33, 870.85, 870.85, 870.85, 870.85, 870.85, 870.99, 870.99, 870.99, 870.99, 870.99, 871.96, 871.96, 871.96, 871.96, 871.96, 872.21, 872.21, 872.21]
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 554 iterations"
y-axis "llamacpp:predicted_tokens_seconds"
x-axis "llamacpp:predicted_tokens_seconds" 1713551797 --> 1713552421
line [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 42.01, 42.01, 42.01, 42.01, 42.01, 41.2, 41.2, 41.2, 41.2, 41.2, 31.2, 31.2, 31.2, 31.2, 31.2, 33.77, 33.77, 33.77, 33.77, 33.77, 33.84, 33.84, 33.84, 33.84, 33.84, 33.74, 33.74, 33.74, 33.74, 33.74, 34.35, 34.35, 34.35, 34.35, 34.35, 35.23, 35.23, 35.23, 35.23, 35.23, 35.19, 35.19, 35.19, 35.19, 35.19, 34.76, 34.76, 34.76, 34.76, 34.76, 33.99, 33.99, 33.99, 33.99, 33.99, 33.98, 33.98, 33.98, 33.98, 33.98, 34.04, 34.04, 34.04, 34.04, 34.04, 33.16, 33.16, 33.16, 33.16, 33.16, 33.17, 33.17, 33.17, 33.17, 33.17, 32.31, 32.31, 32.31, 32.31, 32.31, 32.41, 32.41, 32.41, 32.41, 32.41, 32.56, 32.56, 32.56, 32.56, 32.56, 32.47, 32.47, 32.47, 32.47, 32.47, 32.1, 32.1, 32.1, 32.1, 32.1, 31.92, 31.92, 31.92, 31.92, 31.92, 31.8, 31.8, 31.8, 31.8, 31.8, 32.02, 32.02, 32.02, 32.02, 32.02, 32.08, 32.08, 32.08, 32.08, 32.08, 31.81, 31.81, 31.81, 31.81, 31.81, 31.91, 31.91, 31.91, 31.91, 31.91, 31.96, 31.96, 31.96, 31.96, 31.96, 31.75, 31.75, 31.75, 31.75, 31.75, 31.61, 31.61, 31.61, 31.61, 31.61, 31.69, 31.69, 31.69, 31.69, 31.69, 31.88, 31.88, 31.88, 31.88, 31.88, 31.92, 31.92, 31.92, 31.92, 31.92, 32.21, 32.21, 32.21, 32.21, 32.21, 32.23, 32.23, 32.23, 32.23, 32.23, 32.08, 32.08, 32.08, 32.08, 32.08, 31.97, 31.97, 31.97, 31.97, 31.97, 31.8, 31.8, 31.8, 31.8, 31.8, 31.87, 31.87, 31.87, 31.87, 31.87, 32.0, 32.0, 32.0, 32.0, 32.0, 32.14, 32.14, 32.14, 32.14, 32.14, 32.25, 32.25, 32.25, 32.25, 32.25, 31.92, 31.92, 31.92, 31.92, 31.92, 31.9, 31.9, 31.9, 31.9, 31.9, 31.8, 31.8, 31.8, 31.8, 31.8, 30.59, 30.59, 30.59, 30.59, 30.59, 30.55, 30.55, 30.55, 30.55, 30.55, 30.52, 30.52, 30.52, 30.52, 30.52, 30.64, 30.64, 30.64, 30.64, 30.64, 30.63, 30.63, 30.63, 30.63, 30.63, 30.84, 30.84, 30.84, 30.84, 30.84, 30.79, 30.79, 30.79, 30.79, 30.79, 30.75, 30.75, 30.75, 30.75, 30.75, 30.78, 30.78, 30.78, 30.78, 30.78, 30.77, 30.77, 30.77, 30.77, 30.77, 30.86, 30.86, 30.86, 30.86, 30.86, 30.99, 30.99, 30.99, 30.99, 30.99, 31.02, 31.02, 31.02, 31.02, 31.02, 31.14, 31.14, 31.14, 31.14, 31.14, 31.16, 31.16, 31.16, 31.16, 31.16, 31.14, 31.14, 31.14, 31.14, 31.14, 31.14, 31.14, 31.14]
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 554 iterations"
y-axis "llamacpp:kv_cache_usage_ratio"
x-axis "llamacpp:kv_cache_usage_ratio" 1713551797 --> 1713552421
line [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.35, 0.35, 0.35, 0.35, 0.35, 0.08, 0.08, 0.08, 0.08, 0.08, 0.16, 0.16, 0.16, 0.16, 0.16, 0.18, 0.18, 0.18, 0.18, 0.18, 0.16, 0.16, 0.16, 0.16, 0.16, 0.08, 0.08, 0.08, 0.08, 0.08, 0.17, 0.17, 0.17, 0.17, 0.17, 0.2, 0.2, 0.2, 0.2, 0.2, 0.25, 0.25, 0.25, 0.25, 0.25, 0.11, 0.11, 0.11, 0.11, 0.11, 0.16, 0.16, 0.16, 0.16, 0.16, 0.3, 0.3, 0.3, 0.3, 0.3, 0.18, 0.18, 0.18, 0.18, 0.18, 0.33, 0.33, 0.33, 0.33, 0.33, 0.2, 0.2, 0.2, 0.2, 0.2, 0.19, 0.19, 0.19, 0.19, 0.19, 0.12, 0.12, 0.12, 0.12, 0.12, 0.28, 0.28, 0.28, 0.28, 0.28, 0.3, 0.3, 0.3, 0.3, 0.3, 0.11, 0.11, 0.11, 0.11, 0.11, 0.16, 0.16, 0.16, 0.16, 0.16, 0.12, 0.12, 0.12, 0.12, 0.12, 0.25, 0.25, 0.25, 0.25, 0.25, 0.12, 0.12, 0.12, 0.12, 0.12, 0.13, 0.13, 0.13, 0.13, 0.13, 0.32, 0.32, 0.32, 0.32, 0.32, 0.22, 0.22, 0.22, 0.22, 0.22, 0.15, 0.15, 0.15, 0.15, 0.15, 0.12, 0.12, 0.12, 0.12, 0.12, 0.09, 0.09, 0.09, 0.09, 0.09, 0.09, 0.09, 0.09, 0.09, 0.09, 0.14, 0.14, 0.14, 0.14, 0.14, 0.13, 0.13, 0.13, 0.13, 0.13, 0.18, 0.18, 0.18, 0.18, 0.18, 0.28, 0.28, 0.28, 0.28, 0.28, 0.16, 0.16, 0.16, 0.16, 0.16, 0.17, 0.17, 0.17, 0.17, 0.17, 0.13, 0.13, 0.13, 0.13, 0.13, 0.12, 0.12, 0.12, 0.12, 0.12, 0.24, 0.24, 0.24, 0.24, 0.24, 0.54, 0.54, 0.54, 0.54, 0.54, 0.46, 0.46, 0.46, 0.46, 0.46, 0.49, 0.49, 0.49, 0.49, 0.49, 0.27, 0.27, 0.27, 0.27, 0.27, 0.13, 0.13, 0.13, 0.13, 0.13, 0.15, 0.15, 0.15, 0.15, 0.15, 0.14, 0.14, 0.14, 0.14, 0.14, 0.11, 0.11, 0.11, 0.11, 0.11, 0.17, 0.17, 0.17, 0.17, 0.17, 0.26, 0.26, 0.26, 0.26, 0.26, 0.17, 0.17, 0.17, 0.17, 0.17, 0.19, 0.19, 0.19, 0.19, 0.19, 0.13, 0.13, 0.13, 0.13, 0.13, 0.14, 0.14, 0.14, 0.14, 0.14, 0.07, 0.07, 0.07, 0.07, 0.07, 0.16, 0.16, 0.16, 0.16, 0.16, 0.16, 0.16, 0.16, 0.16, 0.16, 0.19, 0.19, 0.19, 0.19, 0.19, 0.18, 0.18, 0.18, 0.18, 0.18, 0.16, 0.16, 0.16]
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 554 iterations"
y-axis "llamacpp:requests_processing"
x-axis "llamacpp:requests_processing" 1713551797 --> 1713552421
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, 8.0, 8.0, 8.0, 8.0, 8.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, 6.0, 6.0, 6.0, 6.0, 6.0, 4.0, 4.0, 4.0, 4.0, 4.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 4.0, 4.0, 4.0, 4.0, 4.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 3.0, 3.0, 3.0, 3.0, 3.0, 8.0, 8.0, 8.0, 8.0, 8.0, 3.0, 3.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 4.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 8.0, 8.0, 8.0, 8.0, 8.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 7.0, 7.0, 7.0, 7.0, 7.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, 7.0, 7.0, 7.0, 7.0, 7.0, 4.0, 4.0, 4.0, 4.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0, 8.0, 8.0, 8.0, 8.0, 8.0, 5.0, 5.0, 5.0, 5.0, 5.0, 8.0, 8.0, 8.0, 8.0, 8.0, 4.0, 4.0, 4.0, 4.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 3.0, 3.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 4.0, 4.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, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 6.0, 6.0, 6.0, 6.0, 6.0, 6.0, 6.0, 6.0, 6.0, 6.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 8.0, 8.0, 8.0, 8.0, 8.0, 4.0, 4.0, 4.0, 4.0, 4.0, 3.0, 3.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 4.0, 4.0, 7.0, 7.0, 7.0, 7.0, 7.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0, 8.0, 8.0, 8.0, 8.0, 8.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 7.0, 7.0, 7.0, 7.0, 7.0, 4.0, 4.0, 4.0]
|
I don't think it work for gfx1100 at least not yet |
|
Sorry, I forgot: you need to install rocWMMA https://github.com/ROCm/rocWMMA . |
|
I just tried this PR (ac6ae5d), but could not compile. I installed rocwmma from source and I'm using ROCm 6.0.2. I tried to complie with:
|
|
Well, this looks like it would be non-trivial to fix. I was hoping it would be possible to just use rocWMMA as a drop-in replacement. But as I said, I don't have an AMD GPU with tensor cores to debug this with. And quite honestly I don't want to invest a lot of time into this either way because as far as I am concerned none of the current AMD GPUs are worth buying anyways. So unless another dev wants to take over the current llama.cpp FlashAttention implementation using tensor cores will be NVIDIA only. |
|
I just tried to compile with C++17, It could be something with my setup, as there should be a fill_fragment function in rocwmma. |
|
nvm, there are still rocwmma compilation issue with c++17: Something is not right with the types. @JohannesGaessler, indeed does not seem to look like an easy fix. |
|
I managed to compile the branch.
According to https://rocm.docs.amd.com/projects/HIPIFY/en/amd-staging/tables/CUDA_Device_API_supported_by_HIP.html, I know next to nothing about HIP/Cuda, so this is probably wrong.
For this error, here is my workaround: --- a/ggml-cuda/fattn.cu (revision ac6ae5daca029e554af08281a3fd839169725c8c)
+++ b/ggml-cuda/fattn.cu (revision 21b0bf477a56122d8302b218579956b034deaa36)
@@ -339,7 +351,7 @@
frag_c_KQ KQ_c[ncols/frag_n];
#pragma unroll
for (int j = 0; j < ncols/frag_n; ++j) {
- wmma::fill_fragment(KQ_c[j], 0.0f);
+ wmma::fill_fragment(KQ_c[j], KQ_acc_t{0.0f});
}
#pragma unroll
for (int k_KQ_0 = 0; k_KQ_0 < D; k_KQ_0 += 16) {
@@ -470,7 +482,7 @@
for (int i_VKQ_0 = 0; i_VKQ_0 < D; i_VKQ_0 += VKQ_stride) {
#pragma unroll
for (int j = 0; j < ncols/frag_n; ++j) {
- wmma::fill_fragment(VKQ_c[i_VKQ_0/VKQ_stride][j], 0.0f);
+ wmma::fill_fragment(VKQ_c[i_VKQ_0/VKQ_stride][j], __half{0.0f});
}
#pragma unrollHowever, the resulting binary crashes, and it doesn't seem directly related to flash-attn: |
|
@jdecourval just tried your suggestions (I also know nothing about HIP/CUDA), but ChatGPT has a similar proposal for hmax2. And for me it compiles and works. I tried: with the output: |
|
Thanks @tbocek, it was indeed something on my side, probably a leftover from another test. Having the code compile is a nice improvement, but it still doesn't work here if I add the It comes from the fact that However, simply removing the check does make it work! --- a/ggml-cuda/common.cuh (revision 21b0bf477a56122d8302b218579956b034deaa36)
+++ b/ggml-cuda/common.cuh (date 1713632170170)
@@ -364,16 +364,11 @@
}
static __device__ __forceinline__ half2 warp_reduce_sum(half2 a) {
-#if !(defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__)) && __CUDA_ARCH__ >= CC_PASCAL
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1) {
a = __hadd2(a, __shfl_xor_sync(0xffffffff, a, mask, 32));
}
return a;
-#else
- GGML_UNUSED(a);
- NO_DEVICE_CODE;
-#endif // !(defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__)) && __CUDA_ARCH__ >= CC_PASCAL
}
static __device__ __forceinline__ float warp_reduce_max(float x) {@JohannesGaessler do you know why that check is there? |
|
That check is there because that is the old check for FlashAttention and I forgot to change it. |
|
Here are some results. Token generation is about 2-3% faster. I pushed all my fixes to this branch: https://github.com/jdecourval/llama.cpp/tree/fixflashattn |
|
If it's only token generation that is faster then this PR is pretty much pointless because the FlashAttention kernel for batch size 1 does not use tensor cores at all (except for Phi-2). Also can you check llama-bench with |
|
@jdecourval I applied your patches. Here are my results with -p 4096: |
And here are mine
Updated with more models |
ggerganov
force-pushed
the
gg/flash-attn
branch
3 times, most recently
from
5ac4690 to
f725ca9
Compare
ggerganov
approved these changes
Apr 24, 2024
There was a problem hiding this comment.
@JohannesGaessler What do you think given the numbers so far - is it worth it to merge AMD support?
Maybe in the future there can be more progress. I don't know how difficult it would be for people to install the rocWMMA dependency. We can also put it behind a compile flag?
|
My current stance is that I don't think that a speedup of 2% is large enough to justify adding a dependency, especially when there is no dev with the ability to test and support the implementation themselves. I think right now it makes more sense to just disable FlashAttention for AMD. One of my next goals is to write kernels that don't use tensor cores at all. It may turn out that those are faster on AMD than rocWMMA anyways. |
|
Alright. Thank you very much for the help. I will update the target branch to disable flash attention when HIP is enabled for now |
I agree that a 2% performance improvement is not much, but the 500-600MB VRAM reduction may be significant: #5021 (comment) |
|
On Arch Linux at least, I had to manually install the rocWMMA dependency from github, as there is no current library (also not in AUR). |
There's this one: https://aur.archlinux.org/packages/rocwmma |
ggerganov
force-pushed
the
gg/flash-attn
branch
from
82b282c to
ce281b9
Compare
|
Obsolete now that #5021 has been merged. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR attempts to make the llama.cpp FlashAttention code run on AMD via HIP. I do not have an AMD GPU with tensor cores so I cannot test the code myself.
Edit: this PR needs rocWMMA https://github.com/ROCm/rocWMMA to be installed.
This PR also changes the order of some definitions in
common.cuhbecause otherwiseNO_DEVICE_CODEis not functional on AMD.