Skip to content

Commit b69f164

Browse files
JohannesGaesslerJohannesGaessler
authored
CUDA: skip fully masked-out KV in FA vec kernel (#13584)
* CUDA: skip fully masked-out KV in FA vec kernel
1 parent 759e37b commit b69f164

File tree

2 files changed

+95
-8
lines changed

2 files changed

+95
-8
lines changed

‎ggml/src/ggml-cuda/fattn-vec-f16.cuh

Lines changed: 48 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22
#include "fattn-common.cuh"
33

44
template<int D, int ncols, ggml_type type_K, ggml_type type_V, bool use_logit_softcap> // D == head size
5-
#if !(defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__))
5+
#ifndef GGML_USE_HIP
66
__launch_bounds__(D, 1)
7-
#endif // !(defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__))
7+
#endif // GGML_USE_HIP
88
static __global__ void flash_attn_vec_ext_f16(
99
const char * __restrict__ Q,
1010
const char * __restrict__ K,
@@ -48,6 +48,12 @@ static __global__ void flash_attn_vec_ext_f16(
4848
NO_DEVICE_CODE;
4949
return;
5050
}
51+
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
52+
if (ncols > 1) {
53+
NO_DEVICE_CODE;
54+
return;
55+
}
56+
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
5157

5258
//In this kernel Q, K, V are matrices while i, j, k are matrix indices.
5359

@@ -91,6 +97,13 @@ static __global__ void flash_attn_vec_ext_f16(
9197
kqsum_shared[j][threadIdx.x] = 0.0f;
9298
}
9399
}
100+
101+
__shared__ half maskh_shared[ncols*D];
102+
#pragma unroll
103+
for (int j = 0; j < ncols; ++j) {
104+
maskh_shared[j*D + tid] = 0.0f;
105+
}
106+
94107
__syncthreads();
95108

96109
// Convert Q to half2 (f16 K) or q8_1 (quantized K) and store in registers:
@@ -175,6 +188,35 @@ static __global__ void flash_attn_vec_ext_f16(
175188
for (int k_VKQ_0 = blockIdx.y*D; k_VKQ_0 < ne11; k_VKQ_0 += gridDim.y*D) {
176189
// Calculate KQ tile and keep track of new maximum KQ values:
177190

191+
if (mask) {
192+
#pragma unroll
193+
for (int j = 0; j < ncols; ++j) {
194+
maskh_shared[j*D + tid] = slopeh*maskh[j*ne11 + k_VKQ_0 + tid];
195+
}
196+
197+
__syncthreads();
198+
199+
// When using multiple parallel sequences in llama.cpp, some KV slices can be fully masked out.
200+
// In such cases, skip the KV slice.
201+
// On AMD __all_sync would not work correctly because it assumes a warp size of 64.
202+
#ifndef GGML_USE_HIP
203+
bool skip = true;
204+
#pragma unroll
205+
for (int j = 0; j < ncols; ++j) {
206+
#pragma unroll
207+
for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) {
208+
const int i = i0 + threadIdx.x;
209+
210+
const float2 tmp = __half22float2(((const half2 *) maskh_shared)[j*(D/2) + i]);
211+
skip = skip && isinf(tmp.x) && isinf(tmp.y);
212+
}
213+
}
214+
if (__all_sync(0xFFFFFFFF, skip)) {
215+
continue;
216+
}
217+
#endif // GGML_USE_HIP
218+
}
219+
178220
// For unknown reasons using a half array of size 1 for kqmax_new causes a performance regression,
179221
// see https://github.com/ggerganov/llama.cpp/pull/7061 .
180222
// Therefore this variable is defined twice but only used once (so that the compiler can optimize out the unused variable).
@@ -202,7 +244,7 @@ static __global__ void flash_attn_vec_ext_f16(
202244
sum = logit_softcap*tanhf(sum);
203245
}
204246

205-
sum += mask ? slopeh*maskh[j*ne11 + k_VKQ_0 + i_KQ] : __float2half(0.0f);
247+
sum += maskh_shared[j*D + i_KQ];
206248

207249
if (ncols == 1) {
208250
kqmax_new = ggml_cuda_hmax(kqmax_new, sum);
@@ -335,7 +377,9 @@ void ggml_cuda_flash_attn_ext_vec_f16_case(ggml_backend_cuda_context & ctx, ggml
335377
float logit_softcap;
336378
memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float));
337379

338-
if (Q->ne[1] == 1) {
380+
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
381+
382+
if (Q->ne[1] == 1 || GGML_CUDA_CC_IS_NVIDIA(cc)) {
339383
constexpr int cols_per_block = 1;
340384
if (logit_softcap == 0.0f) {
341385
constexpr bool use_logit_softcap = false;

‎ggml/src/ggml-cuda/fattn-vec-f32.cuh

Lines changed: 47 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22
#include "fattn-common.cuh"
33

44
template<int D, int ncols, ggml_type type_K, ggml_type type_V, bool use_logit_softcap> // D == head size
5-
#if !(defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__))
5+
#ifndef GGML_USE_HIP
66
__launch_bounds__(D, 1)
7-
#endif // !(defined(GGML_USE_HIP) && defined(__HIP_PLATFORM_AMD__))
7+
#endif // GGML_USE_HIP
88
static __global__ void flash_attn_vec_ext_f32(
99
const char * __restrict__ Q,
1010
const char * __restrict__ K,
@@ -60,6 +60,12 @@ static __global__ void flash_attn_vec_ext_f32(
6060
NO_DEVICE_CODE;
6161
return;
6262
}
63+
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
64+
if (ncols > 1) {
65+
NO_DEVICE_CODE;
66+
return;
67+
}
68+
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
6369

6470
//In this kernel Q, K, V are matrices while i, j, k are matrix indices.
6571

@@ -104,6 +110,13 @@ static __global__ void flash_attn_vec_ext_f32(
104110
kqsum_shared[j][threadIdx.x] = 0.0f;
105111
}
106112
}
113+
114+
__shared__ float maskf_shared[ncols*D];
115+
#pragma unroll
116+
for (int j = 0; j < ncols; ++j) {
117+
maskf_shared[j*D + tid] = 0.0f;
118+
}
119+
107120
__syncthreads();
108121

109122
// Convert Q to float2 (f16 K) or q8_1 (quantized K) and store in registers:
@@ -181,6 +194,34 @@ static __global__ void flash_attn_vec_ext_f32(
181194
for (int k_VKQ_0 = blockIdx.y*D; k_VKQ_0 < ne11; k_VKQ_0 += gridDim.y*D) {
182195
// Calculate KQ tile and keep track of new maximum KQ values:
183196

197+
if (mask) {
198+
#pragma unroll
199+
for (int j = 0; j < ncols; ++j) {
200+
maskf_shared[j*D + tid] = slope*__half2float(maskh[j*ne11 + k_VKQ_0 + tid]);
201+
}
202+
203+
__syncthreads();
204+
205+
// When using multiple parallel sequences in llama.cpp, some KV slices can be fully masked out.
206+
// In such cases, skip the KV slice.
207+
// On AMD __all_sync would not work correctly because it assumes a warp size of 64.
208+
#ifndef GGML_USE_HIP
209+
bool skip = true;
210+
#pragma unroll
211+
for (int j = 0; j < ncols; ++j) {
212+
#pragma unroll
213+
for (int i0 = 0; i0 < D; i0 += WARP_SIZE) {
214+
const int i = i0 + threadIdx.x;
215+
216+
skip = skip && isinf(maskf_shared[j*D + i]);
217+
}
218+
}
219+
if (__all_sync(0xFFFFFFFF, skip)) {
220+
continue;
221+
}
222+
#endif // GGML_USE_HIP
223+
}
224+
184225
float kqmax_new_arr[ncols];
185226
#pragma unroll
186227
for (int j = 0; j < ncols; ++j) {
@@ -204,7 +245,7 @@ static __global__ void flash_attn_vec_ext_f32(
204245
sum = logit_softcap*tanhf(sum);
205246
}
206247

207-
sum += mask ? slope*__half2float(maskh[j*ne11 + k_VKQ_0 + i_KQ]) : 0.0f;
248+
sum += maskf_shared[j*D + i_KQ];
208249

209250
kqmax_new_arr[j] = fmaxf(kqmax_new_arr[j], sum);
210251

@@ -326,7 +367,9 @@ void ggml_cuda_flash_attn_ext_vec_f32_case(ggml_backend_cuda_context & ctx, ggml
326367
float logit_softcap;
327368
memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float));
328369

329-
if (Q->ne[1] == 1) {
370+
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
371+
372+
if (Q->ne[1] == 1 || GGML_CUDA_CC_IS_NVIDIA(cc)) {
330373
constexpr int cols_per_block = 1;
331374
if (logit_softcap == 0.0f) {
332375
constexpr bool use_logit_softcap = false;

0 commit comments

Comments
 (0)