Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/fattn-wmma-f16.cu (CUDA)

Code before:
                }

                half2 VKQ_add = make_half2(0.0f, 0.0f);
#pragma unroll
                for (int l = 0; l < VKQ_ratio; ++l) {
                    VKQ_add += KQ2[l*(ncols*D_padded/2) + j*(D_padded/2) + i];
                }
                VKQ2[j*(D_padded/2) + i] = VKQ_scale*VKQ2[j*(D_padded/2) + i] + VKQ_add;
            }
        }

        __syncthreads();
    }

    // Apply attention sinks
    if (sinksf && blockIdx.y == 0) {
        const float sinkf = sinksf[head];
        const half  sinkh = __float2half(sinkf);

        #pragma unroll
        for (int j0 = 0; j0 < ncols; j0 += nwarps) {
            const int j = j0 + threadIdx.y;

            if (std::is_same<KQ_acc_t, float>::value) {
                float kqmax_new = fmaxf(KQ_max_f[j0/nwarps], sinkf);
                kqmax_new = warp_reduce_max<warp_size>(kqmax_new);

                const float KQ_max_scale = expf(KQ_max_f[j0/nwarps] - kqmax_new);
                KQ_max_f[j0/nwarps] = kqmax_new;

                KQ_rowsum_f[j0/nwarps] = KQ_rowsum_f[j0/nwarps] * KQ_max_scale + expf(sinkf - KQ_max_f[j0/nwarps]);

                const half2 scale_h2 = make_half2(KQ_max_scale, KQ_max_scale);
                #pragma unroll
                for (int i0 = 0; i0 < D/2; i0 += warp_size) {
                    const int i = i0 + threadIdx.x;
                    if (i0 + warp_size > D/2 && i >= D/2) break;
                    VKQ2[j*(D_padded/2) + i] *= scale_h2;
                }
            } else {
                half kqmax_old = __low2half(KQ_max_h2[j0/nwarps]);
                half kqmax_new = fmaxf(kqmax_old, sinkh);
                kqmax_new = warp_reduce_max<warp_size>(kqmax_new);
                KQ_max_h2[j0/nwarps] = __half2half2(kqmax_new);

                const half  KQ_max_scale_h = hexp(kqmax_old - kqmax_new);
                const half2 KQ_max_scale   = __half2half2(KQ_max_scale_h);

                KQ_rowsum_h2[j0/nwarps] = KQ_rowsum_h2[j0/nwarps] * KQ_max_scale;
                const half val = hexp(sinkh - kqmax_new);
                KQ_rowsum_h2[j0/nwarps].x = __hadd(KQ_rowsum_h2[j0/nwarps].x, val);

Review comment (suggestion):
```suggestion
```

I think this reduction can be removed since `KQ_max_f` should be synchronized at this point in the code and all threads in a warp have compared it against the sink.

Code after:
                }

                half2 VKQ_add = make_half2(0.0f, 0.0f);
#pragma unroll
                for (int l = 0; l < VKQ_ratio; ++l) {
                    VKQ_add += KQ2[l*(ncols*D_padded/2) + j*(D_padded/2) + i];
                }
                VKQ2[j*(D_padded/2) + i] = VKQ_scale*VKQ2[j*(D_padded/2) + i] + VKQ_add;
            }
        }

        __syncthreads();
    }

    // Apply attention sinks
    if (sinksf && blockIdx.y == 0) {
        const float sinkf = sinksf[head];
        const half  sinkh = __float2half(sinkf);

#pragma unroll
        for (int j0 = 0; j0 < ncols; j0 += nwarps) {
            const int j = j0 + threadIdx.y;

            if (std::is_same<KQ_acc_t, float>::value) {
                float kqmax_new = fmaxf(KQ_max_f[j0/nwarps], sinkf);

                const float KQ_max_scale = expf(KQ_max_f[j0/nwarps] - kqmax_new);
                KQ_max_f[j0/nwarps] = kqmax_new;

                KQ_rowsum_f[j0/nwarps] = KQ_rowsum_f[j0/nwarps] * KQ_max_scale + expf(sinkf - KQ_max_f[j0/nwarps]);

                const half2 scale_h2 = make_half2(KQ_max_scale, KQ_max_scale);
#pragma unroll
                for (int i0 = 0; i0 < D/2; i0 += warp_size) {
                    const int i = i0 + threadIdx.x;
                    if (i0 + warp_size > D/2 && i >= D/2) break;
                    VKQ2[j*(D_padded/2) + i] *= scale_h2;
                }
            } else {
                half kqmax_old = __low2half(KQ_max_h2[j0/nwarps]);
                half kqmax_new = fmaxf(kqmax_old, sinkh);
                KQ_max_h2[j0/nwarps] = __half2half2(kqmax_new);

                const half  KQ_max_scale_h = hexp(kqmax_old - kqmax_new);
                const half2 KQ_max_scale   = __half2half2(KQ_max_scale_h);

                KQ_rowsum_h2[j0/nwarps] = KQ_rowsum_h2[j0/nwarps] * KQ_max_scale;
                const half val = hexp(sinkh - kqmax_new);
                KQ_rowsum_h2[j0/nwarps].x = __hadd(KQ_rowsum_h2[j0/nwarps].x, val);

#pragma unroll