Skip to content

Commit 4946c19

Browse files
am17anam17an
committed
CUDA: add attention sinks for tile and wmma
1 parent cd6983d commit 4946c19

4 files changed

Lines changed: 110 additions & 13 deletions

File tree

‎ggml/src/ggml-cuda/fattn-tile-f16.cu‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,8 @@ static __global__ void flash_attn_tile_ext_f16(
5252
const float2 * Q_f2 = (const float2 *) (Q + nb03* sequence + nb02* head + nb01*ic0);
5353
const half2 * K_h2 = (const half2 *) (K + nb13* sequence + nb12*(head / gqa_ratio));
5454
const half2 * V_h2 = (const half2 *) (V + nb13* sequence + nb12*(head / gqa_ratio)); // K and V have same shape
55-
const half * maskh = (const half *) (mask + nb33*(sequence % ne33) + nb31*ic0);
55+
const half * maskh = (const half *) (mask + nb33*(sequence % ne33) + nb31*ic0);
56+
const float * sinksf = (const float *) (sinks);
5657

5758
const int stride_KV2 = nb11 / sizeof(half2);
5859

@@ -242,6 +243,33 @@ static __global__ void flash_attn_tile_ext_f16(
242243
__syncthreads();
243244
}
244245

246+
//Attention sink: adjust running max and sum once per head
247+
if (sinksf && blockIdx.y == 0) {
248+
const half sink = __float2half(sinksf[head]);
249+
250+
#pragma unroll
251+
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
252+
half kqmax_new_j = fmaxf(kqmax[j0/nwarps], sink);
253+
kqmax_new_j = warp_reduce_max(kqmax_new_j);
254+
255+
const half2 KQ_max_scale = __half2half2(hexp(kqmax[j0/nwarps] - kqmax_new_j));
256+
kqmax[j0/nwarps] = kqmax_new_j;
257+
258+
const half val = hexp(sink - kqmax[j0/nwarps]);
259+
kqsum[j0/nwarps] = kqsum[j0/nwarps] * KQ_max_scale;
260+
if (threadIdx.x == 0) {
261+
kqsum[j0/nwarps].x = __hadd(kqsum[j0/nwarps].x, val);
262+
}
263+
264+
#pragma unroll
265+
for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) {
266+
VKQ[j0/nwarps][i0/WARP_SIZE] *= KQ_max_scale;
267+
}
268+
}
269+
270+
__syncthreads();
271+
}
272+
245273
float2 * dst2 = (float2 *) dst;
246274

247275
#pragma unroll

‎ggml/src/ggml-cuda/fattn-tile-f32.cu‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,8 @@ static __global__ void flash_attn_tile_ext_f32(
6363
const float2 * Q_f2 = (const float2 *) (Q + nb03* sequence + nb02* head + nb01*ic0);
6464
const half2 * K_h2 = (const half2 *) (K + nb13* sequence + nb12*(head / gqa_ratio));
6565
const half2 * V_h2 = (const half2 *) (V + nb13* sequence + nb12*(head / gqa_ratio)); // K and V have same shape
66-
const half * maskh = (const half *) (mask + nb33*(sequence % ne33) + nb31*ic0);
66+
const half * maskh = (const half *) (mask + nb33*(sequence % ne33) + nb31*ic0);
67+
const float * sinksf = (const float *) (sinks);
6768

6869
const int stride_KV2 = nb11 / sizeof(half2);
6970

@@ -252,6 +253,35 @@ static __global__ void flash_attn_tile_ext_f32(
252253
__syncthreads();
253254
}
254255

256+
257+
//Attention sink: adjust running max and sum once per head
258+
if (sinksf && blockIdx.y == 0) {
259+
const float sink = sinksf[head];
260+
261+
#pragma unroll
262+
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
263+
float kqmax_new_j = fmaxf(kqmax[j0/nwarps], sink);
264+
kqmax_new_j = warp_reduce_max(kqmax_new_j);
265+
266+
const float KQ_max_scale = expf(kqmax[j0/nwarps] - kqmax_new_j);
267+
kqmax[j0/nwarps] = kqmax_new_j;
268+
269+
const float val = expf(sink - kqmax[j0/nwarps]);
270+
kqsum[j0/nwarps] = kqsum[j0/nwarps] * KQ_max_scale;
271+
if (threadIdx.x == 0) {
272+
kqsum[j0/nwarps] += val;
273+
}
274+
275+
#pragma unroll
276+
for (int i0 = 0; i0 < D/2; i0 += WARP_SIZE) {
277+
VKQ[j0/nwarps][i0/WARP_SIZE].x *= KQ_max_scale;
278+
VKQ[j0/nwarps][i0/WARP_SIZE].y *= KQ_max_scale;
279+
}
280+
}
281+
282+
__syncthreads();
283+
}
284+
255285
float2 * dst2 = (float2 *) dst;
256286

257287
#pragma unroll

‎ggml/src/ggml-cuda/fattn-wmma-f16.cu‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ static __global__ void flash_attn_ext_f16(
8787
const half * V_h = (const half *) (V + nb13* sequence + nb12*(head / gqa_ratio)); // K and V have same shape
8888
const half * maskh = (const half *) (mask + nb33*(sequence % ne33) + nb31*ic0);
8989
const half2 * mask2 = (const half2 *) maskh;
90+
const float * sinksf = (const float *) sinks;
9091

9192
const int stride_Q = nb01 / sizeof(float);
9293
const int stride_KV = nb11 / sizeof(half);
@@ -381,6 +382,55 @@ static __global__ void flash_attn_ext_f16(
381382
__syncthreads();
382383
}
383384

385+
// Apply attention sinks
386+
if (sinksf && blockIdx.y == 0) {
387+
const float sinkf = sinksf[head];
388+
const half sinkh = __float2half(sinkf);
389+
390+
#pragma unroll
391+
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
392+
const int j = j0 + threadIdx.y;
393+
394+
if (std::is_same<KQ_acc_t, float>::value) {
395+
float kqmax_new = fmaxf(KQ_max_f[j0/nwarps], sinkf);
396+
kqmax_new = warp_reduce_max<warp_size>(kqmax_new);
397+
398+
const float KQ_max_scale = expf(KQ_max_f[j0/nwarps] - kqmax_new);
399+
KQ_max_f[j0/nwarps] = kqmax_new;
400+
401+
KQ_rowsum_f[j0/nwarps] = KQ_rowsum_f[j0/nwarps] * KQ_max_scale + expf(sinkf - KQ_max_f[j0/nwarps]);
402+
403+
const half2 scale_h2 = make_half2(KQ_max_scale, KQ_max_scale);
404+
#pragma unroll
405+
for (int i0 = 0; i0 < D/2; i0 += warp_size) {
406+
const int i = i0 + threadIdx.x;
407+
if (i0 + warp_size > D/2 && i >= D/2) break;
408+
VKQ2[j*(D_padded/2) + i] *= scale_h2;
409+
}
410+
} else {
411+
half kqmax_old = __low2half(KQ_max_h2[j0/nwarps]);
412+
half kqmax_new = fmaxf(kqmax_old, sinkh);
413+
kqmax_new = warp_reduce_max<warp_size>(kqmax_new);
414+
KQ_max_h2[j0/nwarps] = __half2half2(kqmax_new);
415+
416+
const half KQ_max_scale_h = hexp(kqmax_old - kqmax_new);
417+
const half2 KQ_max_scale = __half2half2(KQ_max_scale_h);
418+
419+
KQ_rowsum_h2[j0/nwarps] = KQ_rowsum_h2[j0/nwarps] * KQ_max_scale;
420+
const half val = hexp(sinkh - kqmax_new);
421+
KQ_rowsum_h2[j0/nwarps].x = __hadd(KQ_rowsum_h2[j0/nwarps].x, val);
422+
423+
#pragma unroll
424+
for (int i0 = 0; i0 < D/2; i0 += warp_size) {
425+
const int i = i0 + threadIdx.x;
426+
if (i0 + warp_size > D/2 && i >= D/2) break;
427+
VKQ2[j*(D_padded/2) + i] *= KQ_max_scale;
428+
}
429+
}
430+
}
431+
432+
__syncthreads();
433+
}
384434
#pragma unroll
385435
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
386436
const int j_VKQ = j0 + threadIdx.y;

‎ggml/src/ggml-cuda/fattn.cu‎

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -274,23 +274,12 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
274274
const ggml_tensor * K = dst->src[1];
275275
const ggml_tensor * V = dst->src[2];
276276
const ggml_tensor * mask = dst->src[3];
277-
const ggml_tensor * sinks = dst->src[4];
278277

279278
ggml_cuda_set_device(ctx.device);
280279
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
281280
const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
282281
const enum ggml_prec prec = ggml_flash_attn_ext_get_prec(KQV);
283282

284-
// TODO: currently only vec implementation for sinks is supported [TAG_ATTN_SINKS]
285-
if (sinks && !fp16_mma_available(cc)) {
286-
if (prec == GGML_PREC_DEFAULT && fast_fp16_available(cc)) {
287-
ggml_cuda_flash_attn_ext_vec_f16(ctx, dst);
288-
} else {
289-
ggml_cuda_flash_attn_ext_vec_f32(ctx, dst);
290-
}
291-
return;
292-
}
293-
294283
#if defined(GGML_HIP_ROCWMMA_FATTN)
295284
if (GGML_CUDA_CC_IS_AMD(cc) && fp16_mma_available(cc)) {
296285
ggml_cuda_flash_attn_ext_wmma_f16(ctx, dst);

0 commit comments

Comments
 (0)