char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
73a43d1f6
char/aliceai-t5-moe
1#include "common.cuh" 2#include "fattn-common.cuh" 3 4static int ggml_cuda_fattn_vec_get_nthreads_host(const int cc) { 5 return 128; 6 GGML_UNUSED(cc); 7} 8 9static constexpr __device__ int ggml_cuda_fattn_vec_get_nthreads_device() { 10 return 128; 11} 12 13// Currently llvm with the amdgcn target does not support unrolling loops 14// that contain a break that can not be resolved at compile time. 15#ifdef __clang__ 16#pragma clang diagnostic push 17#pragma clang diagnostic ignored "-Wpass-failed" 18#endif // __clang__ 19template<int D, int ncols, ggml_type type_K, ggml_type type_V, bool use_logit_softcap> // D == head size 20__launch_bounds__(ggml_cuda_fattn_vec_get_nthreads_device(), 1) 21static __global__ void flash_attn_ext_vec( 22 const char * Q_ptr, 23 const char * K_ptr, 24 const char * V_ptr, 25 const char * mask_ptr, 26 const char * sinks_ptr, 27 const int * KV_max_ptr, 28 float * dst_ptr, 29 float2 * dst_meta_ptr, 30 const float scale, 31 const float max_bias, 32 const float m0, 33 const float m1, 34 const uint32_t n_head_log2, 35 const float logit_softcap, 36 const int32_t ne00, const uint3 ne01, const int32_t ne02, const int32_t ne03, 37 const int32_t nb01, const int32_t nb02, const int32_t nb03, 38 const int32_t ne10, const int32_t ne11, const int32_t ne12, const int32_t ne13, 39 const int32_t nb11, const int32_t nb12, const int64_t nb13, 40 const int32_t nb21, const int32_t nb22, const int64_t nb23, 41 const int32_t ne31, const int32_t ne32, const int32_t ne33, 42 const int32_t nb31, const int32_t nb32, const int64_t nb33) { 43 ggml_cuda_pdl_lc(); 44#ifdef FLASH_ATTN_AVAILABLE 45 const char * GGML_CUDA_RESTRICT Q = Q_ptr; 46 const char * GGML_CUDA_RESTRICT K = K_ptr; 47 const char * GGML_CUDA_RESTRICT V = V_ptr; 48 const char * GGML_CUDA_RESTRICT mask = mask_ptr; 49 const char * GGML_CUDA_RESTRICT sinks = sinks_ptr; 50 const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr; 51 float * GGML_CUDA_RESTRICT dst = dst_ptr; 52 float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr; 53 54 // Skip unused kernel variants for faster compilation: 55 if (use_logit_softcap && !(D == 128 || D == 256)) { 56 GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale, 57 max_bias, m0, m1, n_head_log2, logit_softcap, 58 ne00, ne01, ne02, ne03, 59 nb01, nb02, nb03, 60 ne10, ne11, ne12, ne13, 61 nb11, nb12, nb13, 62 nb21, nb22, nb23, 63 ne31, ne32, ne33, 64 nb31, nb32, nb33); 65 NO_DEVICE_CODE; 66 return; 67 } 68 69 //In this kernel Q, K, V are matrices while i, j, k are matrix indices. 70 71 constexpr int cpy_nb = ggml_cuda_get_max_cpy_bytes(); 72 constexpr int cpy_ne = cpy_nb / 4; 73 74#ifdef GGML_USE_HIP 75#ifdef RDNA 76 constexpr int nthreads_KQ_q = 2; 77#else 78 constexpr int nthreads_KQ_q = 4; 79#endif // RDNA 80 constexpr int nthreads_V_q = (D/4 < 32 ? D/4 : 32); 81#else 82 constexpr int nthreads_KQ_q = (D/4 < 32 ? D/4 : 32); 83 constexpr int nthreads_V_q = (D/4 < 32 ? D/4 : 32); 84#endif // GGML_USE_HIP 85 86 constexpr int nthreads = ggml_cuda_fattn_vec_get_nthreads_device(); 87 constexpr int nthreads_KQ = (type_K == GGML_TYPE_F16 || type_K == GGML_TYPE_BF16) ? 128 / cpy_nb : nthreads_KQ_q; 88 constexpr int nthreads_V = (type_V == GGML_TYPE_F16 || type_V == GGML_TYPE_BF16) ? 128 / cpy_nb : nthreads_V_q; 89 90 static_assert(WARP_SIZE % nthreads_KQ == 0, "bad nthreads_K"); 91 static_assert(WARP_SIZE % nthreads_V == 0, "bad nthreads_V"); 92 93 constexpr int V_rows_per_thread = (type_V == GGML_TYPE_F16 || type_V == GGML_TYPE_BF16) ? 2*cpy_ne : 4; 94 constexpr int V_cols_per_iter = WARP_SIZE / nthreads_V; 95 96 constexpr vec_dot_KQ_t vec_dot_KQ = get_vec_dot_KQ<type_K, D, nthreads_KQ>(); 97 constexpr bool Q_q8_1 = type_K != GGML_TYPE_F16 && type_K != GGML_TYPE_BF16; 98#ifdef V_DOT2_F32_F16_AVAILABLE 99 constexpr dequantize_V_t dequantize_V = get_dequantize_V<type_V, half, V_rows_per_thread>(); 100#else 101 constexpr dequantize_V_t dequantize_V = get_dequantize_V<type_V, float, V_rows_per_thread>(); 102#endif // V_DOT2_F32_F16_AVAILABLE 103 104 const int ic0 = blockIdx.x * ncols; // Index of the Q/QKV column to work on. 105 106 const int sequence = blockIdx.z / ne02; 107 const int head = blockIdx.z - sequence*ne02; 108 const int gqa_ratio = ne02 / ne12; // With grouped query attention there are > 1 Q matrices per K, V matrix. 109 Q += nb03*sequence + nb02* head + nb01*ic0; 110 K += nb13*sequence + nb12*(head / gqa_ratio); 111 V += nb23*sequence + nb22*(head / gqa_ratio); 112 113 const half * maskh = (const half *) (mask + nb33*(sequence % ne33) + nb31*ic0); 114 115 const float slope = get_alibi_slope(max_bias, head, n_head_log2, m0, m1); 116 117 static_assert(D % (2*WARP_SIZE) == 0, "D not divisible by 2*WARP_SIZE == 64."); 118 constexpr int nwarps = nthreads / WARP_SIZE; 119 const int tid = WARP_SIZE*threadIdx.y + threadIdx.x; 120 __builtin_assume(tid < nthreads); 121 122 constexpr int ne_KQ = ncols*D; 123 constexpr int ne_combine = nwarps*V_cols_per_iter*D; 124#ifdef V_DOT2_F32_F16_AVAILABLE 125 half2 VKQ[ncols][(D/2)/nthreads_V] = {{{0.0f, 0.0f}}}; 126 __shared__ half KQ[ne_KQ > ne_combine ? ne_KQ : ne_combine]; 127#else 128 float2 VKQ[ncols][(D/2)/nthreads_V] = {{{0.0f, 0.0f}}}; 129 __shared__ float KQ[ne_KQ > ne_combine ? ne_KQ : ne_combine]; 130#endif // V_DOT2_F32_F16_AVAILABLE 131 132 float KQ_max[ncols]; 133 float KQ_sum[ncols]; 134#pragma unroll 135 for (int j = 0; j < ncols; ++j) { 136 KQ_max[j] = -FLT_MAX/2.0f; 137 KQ_sum[j] = 0.0f; 138 } 139 140 // Convert Q to float2 (f16 K) or q8_1 (quantized K) and store in registers: 141#ifdef V_DOT2_F32_F16_AVAILABLE 142 half2 Q_reg[ncols][(D/2)/nthreads_KQ]; // Will be initialized completely. 143#else 144 __align__(16) float2 Q_reg[ncols][(D/2)/nthreads_KQ] = {{{0.0f, 0.0f}}}; // May be only partially initialized. 145#endif // V_DOT2_F32_F16_AVAILABLE 146 int Q_i32[ncols][1 > D/(sizeof(int)*nthreads_KQ) ? 1 : D/(sizeof(int)*nthreads_KQ)]; 147 float2 Q_ds[ncols][1 > D/(sizeof(int)*nthreads_KQ) ? 1 : D/(sizeof(int)*nthreads_KQ)]; 148 149 ggml_cuda_pdl_sync(); 150 if constexpr (Q_q8_1) { 151#pragma unroll 152 for (int j0 = 0; j0 < ncols; j0 += nwarps) { 153 const int j = j0 + threadIdx.y; 154 155 if (j0 + nwarps > ncols && j >= ncols) { 156 break; 157 } 158 159 // Reuse KQ as temporary storage for converting Q to q8_1: 160 int * tmp_q_i32 = (int *) &KQ[j*D]; 161 float2 * tmp_q_ds = (float2 *) (tmp_q_i32 + D/sizeof(int)); 162 163 // Set memory to zero if out of bounds: 164 if (ncols > 1 && ic0 + j >= int(ne01.z)) { 165#pragma unroll 166 for (int i0 = 0; i0 < int(D/sizeof(int)); i0 += WARP_SIZE) { 167 const int i = i0 + threadIdx.x; 168 169 if (i0 + WARP_SIZE <= int(D/sizeof(int)) || i < int(D/sizeof(int))) { 170 tmp_q_i32[i] = 0; 171 } 172 } 173 if (threadIdx.x < D/QK8_1) { 174 tmp_q_ds[threadIdx.x] = make_float2(0.0f, 0.0f); 175 } 176 } else { 177 const float * Q_f = (const float *) (Q + j*nb01); 178 constexpr int nthreads_quantize = D/sizeof(int) < WARP_SIZE ? D/sizeof(int) : WARP_SIZE; 179#pragma unroll 180 for (int i0 = 0; i0 < int(D/sizeof(int)); i0 += nthreads_quantize) { 181 quantize_q8_1_to_shared<float2, nthreads_quantize> 182 (Q_f + i0*sizeof(int), scale, tmp_q_i32 + i0, tmp_q_ds + i0/QI8_1); 183 } 184 } 185 } 186 187 __syncthreads(); 188 189#pragma unroll 190 for (int j = 0; j < ncols; ++j) { 191 int * tmp_q_i32 = (int *) &KQ[j*D]; 192 float2 * tmp_q_ds = (float2 *) (tmp_q_i32 + D/sizeof(int)); 193 194#pragma unroll 195 for (int i0 = 0; i0 < int(D/sizeof(int)); i0 += nthreads_KQ) { 196 const int i = i0 + (nthreads_KQ == WARP_SIZE ? threadIdx.x : threadIdx.x % nthreads_KQ); 197 198 Q_i32[j][i0/nthreads_KQ] = tmp_q_i32[i]; 199 Q_ds[j][i0/nthreads_KQ] = tmp_q_ds[i/QI8_1]; 200 } 201 } 202 203 __syncthreads(); 204 } else { 205#ifdef V_DOT2_F32_F16_AVAILABLE 206 const half2 scale_h2 = make_half2(scale, scale); 207#pragma unroll 208 for (int j = 0; j < ncols; ++j) { 209 const float2 * Q_j = (const float2 *) (Q + j*nb01); 210#pragma unroll 211 for (int i0 = 0; i0 < D/2; i0 += nthreads_KQ*cpy_ne) { 212 const int i = i0 + (nthreads_KQ == WARP_SIZE ? threadIdx.x : threadIdx.x % nthreads_KQ)*cpy_ne; 213 214 __align__(16) float2 tmp[cpy_ne] = {{0.0f, 0.0f}}; 215 if (ncols == 1 || ic0 + j < int(ne01.z)) { 216 ggml_cuda_memcpy_1<cpy_nb>(tmp, &Q_j[i]); 217 ggml_cuda_memcpy_1<cpy_nb>(tmp + cpy_ne/2, &Q_j[i + cpy_ne/2]); 218 } 219#pragma unroll 220 for (int i1 = 0; i1 < cpy_ne; ++i1) { 221 Q_reg[j][i0/nthreads_KQ + i1] = make_half2(tmp[i1].x, tmp[i1].y); 222 } 223 } 224#pragma unroll 225 for (int k = 0; k < (D/2)/nthreads_KQ; ++k) { 226 Q_reg[j][k] *= scale_h2; 227 } 228 } 229#else 230#pragma unroll 231 for (int j = 0; j < ncols; ++j) { 232 const float2 * Q_j = (const float2 *) (Q + j*nb01); 233#pragma unroll 234 for (int i0 = 0; i0 < D/2; i0 += nthreads_KQ*cpy_ne) { 235 const int i = i0 + (nthreads_KQ == WARP_SIZE ? threadIdx.x : threadIdx.x % nthreads_KQ)*cpy_ne; 236 if (ncols == 1 || ic0 + j < int(ne01.z)) { 237 ggml_cuda_memcpy_1<cpy_nb>(&Q_reg[j][i0/nthreads_KQ], &Q_j[i]); 238 ggml_cuda_memcpy_1<cpy_nb>(&Q_reg[j][i0/nthreads_KQ + cpy_ne/2], &Q_j[i + cpy_ne/2]); 239 } 240 } 241#pragma unroll 242 for (int k = 0; k < (D/2)/nthreads_KQ; ++k) { 243 Q_reg[j][k].x *= scale; 244 Q_reg[j][k].y *= scale; 245 } 246 } 247#endif // V_DOT2_F32_F16_AVAILABLE 248 } 249 250 const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11; 251 K += blockIdx.y*nthreads * nb11; 252 V += blockIdx.y*nthreads * nb21; 253 maskh += blockIdx.y*nthreads; 254 for (int k_VKQ_0 = blockIdx.y*nthreads; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*nthreads, 255 // Increment pointers after each loop: 256 K += gridDim.y*nthreads*nb11, V += gridDim.y*nthreads*nb21, maskh += gridDim.y*nthreads) { 257 258 // Calculate KQ tile and keep track of new maximum KQ values: 259 float KQ_reg[ncols]; // KQ in registers. 260 261 float KQ_max_new[ncols]; 262#pragma unroll 263 for (int j = 0; j < ncols; ++j) { 264 KQ_max_new[j] = KQ_max[j]; 265 } 266 267#pragma unroll 268 for (int i_KQ_0 = 0; i_KQ_0 < nthreads_KQ; ++i_KQ_0) { 269 const int i_KQ = threadIdx.y*WARP_SIZE + (nthreads_KQ == WARP_SIZE ? 0 : (threadIdx.x & ~(nthreads_KQ-1))) + i_KQ_0; 270 271#pragma unroll 272 for (int j = 0; j < ncols; ++j) { 273 float sum = vec_dot_KQ(K + i_KQ*nb11, Q_reg[j], Q_i32[j], Q_ds[j]); 274 sum = warp_reduce_sum<nthreads_KQ>(sum); 275 276 if (use_logit_softcap) { 277 sum = logit_softcap*tanhf(sum); 278 } 279 280 if (mask && (ncols == 1 || ic0 + j < int(ne01.z))) { 281 sum += slope*__half2float(maskh[j*ne11 + i_KQ]); 282 } 283 284 KQ_max_new[j] = fmaxf(KQ_max_new[j], sum + FATTN_KQ_MAX_OFFSET); 285 286 if ((nthreads_KQ == WARP_SIZE ? threadIdx.x : threadIdx.x % nthreads_KQ) == uint32_t(i_KQ_0)) { 287 KQ_reg[j] = sum; 288 } 289 } 290 } 291 292#pragma unroll 293 for (int j = 0; j < ncols; ++j) { 294#pragma unroll 295 for (int offset = nthreads_KQ; offset < WARP_SIZE; offset <<= 1) { 296 KQ_max_new[j] = fmaxf(KQ_max_new[j], __shfl_xor_sync(0xFFFFFFFF, KQ_max_new[j], offset, WARP_SIZE)); 297 } 298 const float KQ_max_scale = expf(KQ_max[j] - KQ_max_new[j]); 299 KQ_max[j] = KQ_max_new[j]; 300 301 KQ_reg[j] = expf(KQ_reg[j] - KQ_max[j]); 302 KQ_sum[j] = KQ_sum[j]*KQ_max_scale + KQ_reg[j]; 303 KQ[j*nthreads + tid] = KQ_reg[j]; 304 305#ifdef V_DOT2_F32_F16_AVAILABLE 306 const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale, KQ_max_scale); 307#pragma unroll 308 for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V) { 309 VKQ[j][i_VKQ_0/nthreads_V] *= KQ_max_scale_h2; 310 } 311#else 312#pragma unroll 313 for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V) { 314 VKQ[j][i_VKQ_0/nthreads_V].x *= KQ_max_scale; 315 VKQ[j][i_VKQ_0/nthreads_V].y *= KQ_max_scale; 316 } 317#endif // V_DOT2_F32_F16_AVAILABLE 318 } 319 320 ggml_cuda_syncwarp(); 321 322#pragma unroll 323 for (int k0 = 0; k0 < WARP_SIZE; k0 += V_cols_per_iter) { 324 const int k = threadIdx.y*WARP_SIZE + k0 + (nthreads_V == WARP_SIZE ? 0 : threadIdx.x / nthreads_V); 325 326#ifdef V_DOT2_F32_F16_AVAILABLE 327 half2 KQ_k[ncols]; 328#pragma unroll 329 for (int j = 0; j < ncols; ++j) { 330 KQ_k[j] = __half2half2(KQ[j*nthreads + k]); 331 } 332#pragma unroll 333 for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V*V_rows_per_thread/2) { 334 half2 tmp[V_rows_per_thread/2]; 335 if constexpr (type_V == GGML_TYPE_BF16) { 336 float2 tmp_f[V_rows_per_thread/2]; 337 dequantize_V(V + k*nb21, tmp_f, 338 2*i_VKQ_0 + (nthreads_V == WARP_SIZE ? threadIdx.x : threadIdx.x % nthreads_V)*V_rows_per_thread); 339#pragma unroll 340 for (int i_VKQ_1 = 0; i_VKQ_1 < V_rows_per_thread/2; ++i_VKQ_1) { 341 tmp[i_VKQ_1] = __float22half2_rn(tmp_f[i_VKQ_1]); 342 } 343 } else { 344 dequantize_V(V + k*nb21, tmp, 345 2*i_VKQ_0 + (nthreads_V == WARP_SIZE ? threadIdx.x : threadIdx.x % nthreads_V)*V_rows_per_thread); 346 } 347#pragma unroll 348 for (int i_VKQ_1 = 0; i_VKQ_1 < V_rows_per_thread/2; ++i_VKQ_1) { 349#pragma unroll 350 for (int j = 0; j < ncols; ++j) { 351 VKQ[j][i_VKQ_0/nthreads_V + i_VKQ_1] += tmp[i_VKQ_1]*KQ_k[j]; 352 } 353 } 354 } 355#else 356 float KQ_k[ncols]; 357#pragma unroll 358 for (int j = 0; j < ncols; ++j) { 359 KQ_k[j] = KQ[j*nthreads + k]; 360 } 361#pragma unroll 362 for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V*V_rows_per_thread/2) { 363 float2 tmp[V_rows_per_thread/2]; 364 dequantize_V(V + k*nb21, tmp, 365 2*i_VKQ_0 + (nthreads_V == WARP_SIZE ? threadIdx.x : threadIdx.x % nthreads_V)*V_rows_per_thread); 366#pragma unroll 367 for (int i_VKQ_1 = 0; i_VKQ_1 < V_rows_per_thread/2; ++i_VKQ_1) { 368#pragma unroll 369 for (int j = 0; j < ncols; ++j) { 370 VKQ[j][i_VKQ_0/nthreads_V + i_VKQ_1].x += tmp[i_VKQ_1].x*KQ_k[j]; 371 VKQ[j][i_VKQ_0/nthreads_V + i_VKQ_1].y += tmp[i_VKQ_1].y*KQ_k[j]; 372 } 373 } 374 } 375#endif // V_DOT2_F32_F16_AVAILABLE 376 } 377 } 378 379 if (sinks && blockIdx.y == 0) { 380 const float sink = ((const float *) sinks)[head]; 381 382#pragma unroll 383 for (int j0 = 0; j0 < ncols; j0 += nwarps) { 384 const int j = j0 + threadIdx.y; 385 386 if (j0 + nwarps > ncols && j >= ncols) { 387 break; 388 } 389 390 const float kqmax_new_j = fmaxf(sink, KQ_max[j]); 391 const float KQ_max_scale = expf(KQ_max[j] - kqmax_new_j); 392 KQ_max[j] = kqmax_new_j; 393 394 KQ_sum[j] = KQ_sum[j]*KQ_max_scale + (threadIdx.x == 0 ? expf(sink - KQ_max[j]) : 0.0f); 395 396#ifdef V_DOT2_F32_F16_AVAILABLE 397 const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale, KQ_max_scale); 398#pragma unroll 399 for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V) { 400 VKQ[j][i_VKQ_0/nthreads_V] *= KQ_max_scale_h2; 401 } 402#else 403#pragma unroll 404 for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V) { 405 VKQ[j][i_VKQ_0/nthreads_V].x *= KQ_max_scale; 406 VKQ[j][i_VKQ_0/nthreads_V].y *= KQ_max_scale; 407 } 408#endif // V_DOT2_F32_F16_AVAILABLE 409 } 410 } 411 412 __shared__ float KQ_max_shared[ncols][WARP_SIZE]; 413 __shared__ float KQ_sum_shared[ncols][WARP_SIZE]; 414#pragma unroll 415 for (int j = 0; j < ncols; ++j) { 416 if (threadIdx.y == 0) { 417 KQ_max_shared[j][threadIdx.x] = -FLT_MAX/2.0f; 418 KQ_sum_shared[j][threadIdx.x] = 0.0f; 419 } 420 } 421 422 __syncthreads(); 423 424#pragma unroll 425 for (int j = 0; j < ncols; ++j) { 426 if (threadIdx.x == 0) { 427 KQ_max_shared[j][threadIdx.y] = KQ_max[j]; 428 } 429 } 430 __syncthreads(); 431 432#pragma unroll 433 for (int j_VKQ = 0; j_VKQ < ncols; ++j_VKQ) { 434 if (ncols > 1 && ic0 + j_VKQ >= int(ne01.z)) { 435 break; 436 } 437 438 float kqmax_new = KQ_max_shared[j_VKQ][threadIdx.x]; 439 kqmax_new = warp_reduce_max(kqmax_new); 440 const float kqmax_scale = expf(KQ_max[j_VKQ] - kqmax_new); 441 KQ_max[j_VKQ] = kqmax_new; 442 443#ifdef V_DOT2_F32_F16_AVAILABLE 444 half2 * VKQ_tmp = (half2 *) KQ + threadIdx.y*(V_cols_per_iter*D/2) 445 + (nthreads_V == WARP_SIZE ? 0 : threadIdx.x / nthreads_V)*(D/2); 446 447 const half2 kqmax_scale_h2 = make_half2(kqmax_scale, kqmax_scale); 448#pragma unroll 449 for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V) { 450 VKQ[j_VKQ][i_VKQ_0/nthreads_V] *= kqmax_scale_h2; 451 } 452#pragma unroll 453 for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V*V_rows_per_thread/2) { 454 const int i_VKQ = i_VKQ_0 + (nthreads_V == WARP_SIZE ? threadIdx.x : threadIdx.x % nthreads_V)*(V_rows_per_thread/2); 455 456 ggml_cuda_memcpy_1<V_rows_per_thread*sizeof(half)>(VKQ_tmp + i_VKQ, &VKQ[j_VKQ][i_VKQ_0/nthreads_V]); 457 } 458#else 459 float2 * VKQ_tmp = (float2 *) KQ + threadIdx.y*(V_cols_per_iter*D/2) 460 + (nthreads_V == WARP_SIZE ? 0 : threadIdx.x / nthreads_V)*(D/2); 461 462#pragma unroll 463 for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V) { 464 VKQ[j_VKQ][i_VKQ_0/nthreads_V].x *= kqmax_scale; 465 VKQ[j_VKQ][i_VKQ_0/nthreads_V].y *= kqmax_scale; 466 } 467#pragma unroll 468 for (int i_VKQ_0 = 0; i_VKQ_0 < D/2; i_VKQ_0 += nthreads_V*V_rows_per_thread/2) { 469 const int i_VKQ = i_VKQ_0 + (nthreads_V == WARP_SIZE ? threadIdx.x : threadIdx.x % nthreads_V)*(V_rows_per_thread/2); 470 471 ggml_cuda_memcpy_1<V_rows_per_thread/2*sizeof(float)>(VKQ_tmp + i_VKQ, &VKQ[j_VKQ][i_VKQ_0/nthreads_V]); 472 ggml_cuda_memcpy_1<V_rows_per_thread/2*sizeof(float)>(VKQ_tmp + i_VKQ + V_rows_per_thread/4, &VKQ[j_VKQ][i_VKQ_0/nthreads_V + V_rows_per_thread/4]); 473 } 474#endif // V_DOT2_F32_F16_AVAILABLE 475 476 KQ_sum[j_VKQ] *= kqmax_scale; 477 KQ_sum[j_VKQ] = warp_reduce_sum(KQ_sum[j_VKQ]); 478 if (threadIdx.x == 0) { 479 KQ_sum_shared[j_VKQ][threadIdx.y] = KQ_sum[j_VKQ]; 480 } 481 482 __syncthreads(); 483 484 if (nthreads <= D || tid < D) { 485 KQ_sum[j_VKQ] = KQ_sum_shared[j_VKQ][threadIdx.x]; 486 KQ_sum[j_VKQ] = warp_reduce_sum(KQ_sum[j_VKQ]); 487 488#pragma unroll 489 for (int i0 = 0; i0 < D; i0 += nthreads) { 490 float dst_val = 0; 491#pragma unroll 492 for (int w = 0; w < nwarps; ++w) { 493#pragma unroll 494 for (int v = 0; v < V_cols_per_iter; ++v) { 495 dst_val += float(KQ[w*V_cols_per_iter*D + v*D + i0 + tid]); 496 } 497 } 498 if (gridDim.y == 1) { 499 dst_val /= KQ_sum[j_VKQ]; 500 } 501 dst[(((sequence*int(ne01.z) + ic0 + j_VKQ)*ne02 + head)*gridDim.y + blockIdx.y)*D + i0 + tid] = dst_val; 502 } 503 } 504 505 if (j_VKQ < ncols-1) { 506 __syncthreads(); 507 } 508 509 } 510 511 if (gridDim.y != 1 && tid < ncols && (ncols == 1 || ic0 + tid < int(ne01.z))) { 512 dst_meta[((sequence*int(ne01.z) + ic0 + tid)*ne02 + head)*gridDim.y + blockIdx.y] = make_float2(KQ_max[tid], KQ_sum[tid]); 513 } 514#else 515 GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale, 516 max_bias, m0, m1, n_head_log2, logit_softcap, 517 ne00, ne01, ne02, ne03, 518 nb01, nb02, nb03, 519 ne10, ne11, ne12, ne13, 520 nb11, nb12, nb13, 521 nb21, nb22, nb23, 522 ne31, ne32, ne33, 523 nb31, nb32, nb33); 524 NO_DEVICE_CODE; 525#endif // FLASH_ATTN_AVAILABLE 526} 527#ifdef __clang__ 528#pragma clang diagnostic pop 529#endif // __clang__ 530 531template <int D, int cols_per_block, ggml_type type_K, ggml_type type_V, bool use_logit_softcap> 532void ggml_cuda_flash_attn_ext_vec_case_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { 533 const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; 534 535 const int nthreads = ggml_cuda_fattn_vec_get_nthreads_host(cc); 536 const int nwarps = nthreads / WARP_SIZE; 537 fattn_kernel_t fattn_kernel = flash_attn_ext_vec<D, cols_per_block, type_K, type_V, use_logit_softcap>; 538 const bool need_f16_K = type_K == GGML_TYPE_F16; 539 const bool need_f16_V = type_V == GGML_TYPE_F16; 540 constexpr size_t nbytes_shared = 0; 541 launch_fattn<D, cols_per_block, 1>(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false, false); 542} 543 544template <int D, ggml_type type_K, ggml_type type_V> 545void ggml_cuda_flash_attn_ext_vec_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { 546 const ggml_tensor * KQV = dst; 547 const ggml_tensor * Q = dst->src[0]; 548 549 float logit_softcap; 550 memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float)); 551 552 if (Q->ne[1] == 1) { 553 constexpr int cols_per_block = 1; 554 if (logit_softcap == 0.0f) { 555 constexpr bool use_logit_softcap = false; 556 ggml_cuda_flash_attn_ext_vec_case_impl<D, cols_per_block, type_K, type_V, use_logit_softcap>(ctx, dst); 557 } else { 558 constexpr bool use_logit_softcap = true; 559 ggml_cuda_flash_attn_ext_vec_case_impl<D, cols_per_block, type_K, type_V, use_logit_softcap>(ctx, dst); 560 } 561 return; 562 } 563 564 constexpr int cols_per_block = 2; 565 if (logit_softcap == 0.0f) { 566 constexpr bool use_logit_softcap = false; 567 ggml_cuda_flash_attn_ext_vec_case_impl<D, cols_per_block, type_K, type_V, use_logit_softcap>(ctx, dst); 568 } else { 569 constexpr bool use_logit_softcap = true; 570 ggml_cuda_flash_attn_ext_vec_case_impl<D, cols_per_block, type_K, type_V, use_logit_softcap>(ctx, dst); 571 } 572} 573 574#define DECL_FATTN_VEC_CASE(D, type_K, type_V) \ 575 template void ggml_cuda_flash_attn_ext_vec_case \ 576 <D, type_K, type_V>(ggml_backend_cuda_context & ctx, ggml_tensor * dst) \ 577 578#define EXTERN_DECL_FATTN_VEC_CASES(D, type_K) \ 579 extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_F16); \ 580 extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q4_0); \ 581 extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q4_1); \ 582 extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q5_0); \ 583 extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q5_1); \ 584 extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_Q8_0); \ 585 extern DECL_FATTN_VEC_CASE(D, type_K, GGML_TYPE_BF16); \ 586 587EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_F16) 588EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q4_0) 589EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q4_1) 590EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q5_0) 591EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q5_1) 592EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_Q8_0) 593EXTERN_DECL_FATTN_VEC_CASES( 64, GGML_TYPE_BF16) 594 595EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_F16) 596EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q4_0) 597EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q4_1) 598EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q5_0) 599EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q5_1) 600EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_Q8_0) 601EXTERN_DECL_FATTN_VEC_CASES(128, GGML_TYPE_BF16) 602 603EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_F16) 604EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q4_0) 605EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q4_1) 606EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q5_0) 607EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q5_1) 608EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_Q8_0) 609EXTERN_DECL_FATTN_VEC_CASES(256, GGML_TYPE_BF16)