char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
16378d93f
char/aliceai-t5-moe
1#include "common.cuh" 2#include "fattn-common.cuh" 3#include "fattn-mma-f16.cuh" 4#include "fattn-tile.cuh" 5#include "fattn-vec.cuh" 6#include "fattn.cuh" 7 8#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) 9__launch_bounds__(256, 1) 10static __global__ void flash_attn_mask_to_sparse_indices( 11 const half * mask_ptr, int32_t * indices_ptr, const int ne30, const int n_kv_max, 12 const int64_t s31, const int64_t s33) { 13 ggml_cuda_pdl_sync(); 14 15 constexpr int values_per_lane = 8; 16 const int tid = threadIdx.x; 17 const int warp = tid / WARP_SIZE; 18 const int lane = tid % WARP_SIZE; 19 const int sequence = blockIdx.y; 20 const int query = blockIdx.x; 21 22 const half * mask = mask_ptr + sequence*s33 + query*s31; 23 int32_t * indices = indices_ptr + (int64_t(sequence)*gridDim.x + query)*n_kv_max; 24 25 __shared__ int warp_offsets[256/WARP_SIZE]; 26 __shared__ int row_count; 27 __shared__ int chunk_count; 28 29 if (tid == 0) { 30 row_count = 0; 31 } 32 __syncthreads(); 33 34 for (int i0 = 0; i0 < ne30; i0 += blockDim.x*values_per_lane) { 35 uint32_t selected_warp[values_per_lane]; 36 int warp_count = 0; 37#pragma unroll 38 for (int item = 0; item < values_per_lane; ++item) { 39 const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane; 40 const bool selected = i < ne30 && isfinite(__half2float(mask[i])); 41 selected_warp[item] = __ballot_sync(0xFFFFFFFF, selected); 42 warp_count += __popc(selected_warp[item]); 43 } 44 45 if (lane == 0) { 46 warp_offsets[warp] = warp_count; 47 } 48 __syncthreads(); 49 50 if (tid == 0) { 51 int offset = 0; 52#pragma unroll 53 for (int iw = 0; iw < 256/WARP_SIZE; ++iw) { 54 const int count = warp_offsets[iw]; 55 warp_offsets[iw] = offset; 56 offset += count; 57 } 58 chunk_count = offset; 59 } 60 __syncthreads(); 61 62 const uint32_t lane_mask = lane == 0 ? 0 : (1u << lane) - 1; 63 int warp_item_offset = 0; 64#pragma unroll 65 for (int item = 0; item < values_per_lane; ++item) { 66 const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane; 67 const int dst = row_count + warp_offsets[warp] + warp_item_offset + __popc(selected_warp[item] & lane_mask); 68 if ((selected_warp[item] & (uint32_t(1) << lane)) && dst < n_kv_max) { 69 indices[dst] = i; 70 } 71 warp_item_offset += __popc(selected_warp[item]); 72 } 73 __syncthreads(); 74 75 if (tid == 0) { 76 row_count += chunk_count; 77 } 78 __syncthreads(); 79 } 80 81 const int count = row_count; 82 for (int i = count + tid; i < n_kv_max; i += blockDim.x) { 83 indices[i] = -1; 84 } 85 __syncthreads(); 86 87 // the dependent grid reads indices, signal once the row is complete 88 ggml_cuda_pdl_lc(); 89} 90#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) 91 92void ggml_cuda_flash_attn_ext_compact_mask( 93 const ggml_tensor * mask, int32_t * indices, int32_t n_kv_max, cudaStream_t stream) { 94#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) 95 GGML_UNUSED_VARS(mask, indices, n_kv_max, stream); 96 GGML_ABORT("sparse flash attention is only supported on NVIDIA CUDA"); 97#else 98 const int64_t s31 = mask->nb[1] / sizeof(half); 99 const int64_t s33 = mask->nb[3] / sizeof(half); 100 const dim3 blocks_num(mask->ne[1], mask->ne[3], 1); 101 const dim3 block_dim(256, 1, 1); 102 const ggml_cuda_kernel_launch_params launch_params(blocks_num, block_dim, 0, stream); 103 ggml_cuda_kernel_launch(flash_attn_mask_to_sparse_indices, launch_params, 104 (const half *) mask->data, indices, int(mask->ne[0]), n_kv_max, s31, s33); 105 CUDA_CHECK(cudaGetLastError()); 106#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) 107} 108 109bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { 110#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) 111 GGML_UNUSED_VARS(ctx, dst); 112 return false; 113#else 114 const ggml_tensor * Q = dst->src[0]; 115 const ggml_tensor * K = dst->src[1]; 116 const ggml_tensor * mask = dst->src[3]; 117 const int cc = ggml_cuda_info().devices[ctx.device].cc; 118 119 float max_bias = 0.0f; 120 float logit_softcap = 0.0f; 121 memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); 122 memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); 123 124 const int32_t n_kv_max = ggml_get_op_params_i32(dst, 4); 125 return GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && 126 mask != nullptr && n_kv_max > 0 && max_bias == 0.0f && logit_softcap == 0.0f && 127 mask->ne[0] == K->ne[1] && mask->ne[1] >= Q->ne[1] && mask->ne[2] == 1 && 128 K->ne[1] >= std::max<int64_t>(4096, 2LL*n_kv_max); 129#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) 130} 131 132template <int DKQ, int DV, int ncols2> 133static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { 134 const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; 135 const ggml_tensor * Q = dst->src[0]; 136 137#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) 138 if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, 1, ncols2)) { 139 if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst)) { 140 ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 1, ncols2>(ctx, dst); 141 return; 142 } 143 } 144#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) 145 146 if constexpr (ncols2 <= 8) { 147 if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) { 148 ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 8/ncols2, ncols2>(ctx, dst); 149 return; 150 } 151 } 152 153 if constexpr (ncols2 <= 16) { 154 if (Q->ne[1] <= 16/ncols2) { 155 ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 16/ncols2, ncols2>(ctx, dst); 156 return; 157 } 158 } 159 160 if (Q->ne[1] <= 32/ncols2 || (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_TURING) || 161 (GGML_CUDA_CC_IS_AMD(cc) && DKQ > 256)) { 162 ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 32/ncols2, ncols2>(ctx, dst); 163 return; 164 } 165 166 ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 64/ncols2, ncols2>(ctx, dst); 167} 168 169template <int DKQ, int DV> 170static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { 171 const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; 172 const ggml_tensor * KQV = dst; 173 const ggml_tensor * Q = dst->src[0]; 174 const ggml_tensor * K = dst->src[1]; 175 const ggml_tensor * V = dst->src[2]; 176 const ggml_tensor * mask = dst->src[3]; 177 178 float max_bias = 0.0f; 179 memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float)); 180 181 // Edge cases like no mask, ALiBi, unpadded K/V, or misaligned addresses for large data transfers 182 // are put into the template specialization without GQA optimizations. 183 bool use_gqa_opt = mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0; 184 for (const ggml_tensor * t : {Q, K, V, mask}) { 185 if (t == nullptr || ggml_is_quantized(t->type)) { 186 continue; 187 } 188 for (size_t i = 1; i < GGML_MAX_DIMS; ++i) { 189 if (t->nb[i] % 16 != 0) { 190 use_gqa_opt = false; 191 break; 192 } 193 } 194 } 195 196 GGML_ASSERT(Q->ne[2] % K->ne[2] == 0); 197 const int gqa_ratio = Q->ne[2] / K->ne[2]; 198 199 // On Volta the GQA optimizations aren't as impactful vs. minimizing wasted compute: 200 if (cc == GGML_CUDA_CC_VOLTA) { 201 if (use_gqa_opt && gqa_ratio % 8 == 0) { 202 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 8>(ctx, dst); 203 return; 204 } 205 206 if (use_gqa_opt && gqa_ratio % 4 == 0) { 207 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 4>(ctx, dst); 208 return; 209 } 210 211 if constexpr (DKQ <= 256) { 212 if (use_gqa_opt && gqa_ratio % 2 == 0) { 213 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 2>(ctx, dst); 214 return; 215 } 216 217 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 1>(ctx, dst); 218 return; 219 } else { 220 GGML_ABORT("fatal error"); 221 } 222 } 223 224 // On RDNA it is preferable to minimize wasted compute vs. duplicate I/O for the mask. 225 if (amd_wmma_available(cc)) { 226 if (use_gqa_opt && gqa_ratio % 8 == 0) { 227 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 8>(ctx, dst); 228 return; 229 } 230 231 if (use_gqa_opt && gqa_ratio % 4 == 0) { 232 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 4>(ctx, dst); 233 return; 234 } 235 236 if (use_gqa_opt && gqa_ratio % 2 == 0) { 237 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 2>(ctx, dst); 238 return; 239 } 240 } 241 242 if (use_gqa_opt && gqa_ratio > 4) { 243 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 8>(ctx, dst); 244 return; 245 } 246 247 if (use_gqa_opt && gqa_ratio > 2) { 248 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 4>(ctx, dst); 249 return; 250 } 251 252 if (use_gqa_opt && gqa_ratio > 1) { 253 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 2>(ctx, dst); 254 return; 255 } 256 257 if constexpr (DKQ <= 256) { 258 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 1>(ctx, dst); 259 } else { 260 GGML_ABORT("fatal error"); 261 } 262} 263 264static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { 265 const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; 266 const ggml_tensor * KQV = dst; 267 const ggml_tensor * Q = dst->src[0]; 268 const ggml_tensor * K = dst->src[1]; 269 const ggml_tensor * V = dst->src[2]; 270 const ggml_tensor * mask = dst->src[3]; 271 272 switch (Q->ne[0]) { 273 case 64: 274 GGML_ASSERT(V->ne[0] == 64); 275 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2< 64, 64>(ctx, dst); 276 break; 277 case 80: 278 GGML_ASSERT(V->ne[0] == 80); 279 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2< 80, 80>(ctx, dst); 280 break; 281 case 96: 282 GGML_ASSERT(V->ne[0] == 96); 283 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2< 96, 96>(ctx, dst); 284 break; 285 case 112: 286 GGML_ASSERT(V->ne[0] == 112); 287 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2<112, 112>(ctx, dst); 288 break; 289 case 128: 290 GGML_ASSERT(V->ne[0] == 128); 291 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2<128, 128>(ctx, dst); 292 break; 293 case 192: { 294 // MiMo-V2.5 / V2.5-Pro / V2-Flash: gqa_ratio is 8 (SWA) or 16 (full attn) 295 GGML_ASSERT(V->ne[0] == 128); 296 float max_bias = 0.0f; 297 memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float)); 298 const bool use_gqa_opt = mask && max_bias == 0.0f; 299 GGML_ASSERT(use_gqa_opt); 300 GGML_ASSERT(Q->ne[2] % K->ne[2] == 0); 301 const int gqa_ratio = Q->ne[2] / K->ne[2]; 302 if (gqa_ratio % 16 == 0) { 303 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<192, 128, 16>(ctx, dst); 304 } else { 305 GGML_ASSERT(gqa_ratio % 8 == 0); 306 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<192, 128, 8>(ctx, dst); 307 } 308 } break; 309 case 256: 310 GGML_ASSERT(V->ne[0] == 256); 311 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2<256, 256>(ctx, dst); 312 break; 313 case 320: 314 // For Mistral Small 4, go straight to the ncols1 switch (ncols2=32-only build). 315 GGML_ASSERT(V->ne[0] == 256); 316 { 317 float max_bias = 0.0f; 318 memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float)); 319 320 const bool use_gqa_opt = mask && max_bias == 0.0f; 321 GGML_ASSERT(use_gqa_opt); 322 GGML_ASSERT(Q->ne[2] % K->ne[2] == 0); 323 const int gqa_ratio = Q->ne[2] / K->ne[2]; 324 GGML_ASSERT(gqa_ratio % 32 == 0); 325 326 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<320, 256, 32>(ctx, dst); 327 } 328 break; 329 case 512: 330 GGML_ASSERT(V->ne[0] == 512); 331 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2<512, 512>(ctx, dst); 332 break; 333 case 576: { 334 // For Deepseek, go straight to the ncols1 switch to avoid compiling unnecessary kernels. 335 GGML_ASSERT(V->ne[0] == 512); 336 float max_bias = 0.0f; 337 memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float)); 338 339 const bool use_gqa_opt = mask && max_bias == 0.0f; 340 GGML_ASSERT(use_gqa_opt); 341 342 GGML_ASSERT(Q->ne[2] % K->ne[2] == 0); 343 const int gqa_ratio = Q->ne[2] / K->ne[2]; 344 if (gqa_ratio == 20) { // GLM 4.7 Flash 345 if (cc >= GGML_CUDA_CC_DGX_SPARK) { 346 if (Q->ne[1] <= 8) { 347 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 16>(ctx, dst); 348 break; 349 } 350 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 4>(ctx, dst); 351 break; 352 } 353 if (cc >= GGML_CUDA_CC_BLACKWELL) { 354 if (Q->ne[1] <= 4 && K->ne[1] >= 65536) { 355 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 16>(ctx, dst); 356 break; 357 } 358 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 4>(ctx, dst); 359 break; 360 } 361 if (cc >= GGML_CUDA_CC_ADA_LOVELACE) { 362 if (Q->ne[1] <= 4) { 363 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 16>(ctx, dst); 364 break; 365 } 366 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 4>(ctx, dst); 367 break; 368 } 369 if (cc >= GGML_CUDA_CC_TURING) { 370 if (Q->ne[1] <= 4) { 371 if (K->ne[1] <= 16384) { 372 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 16>(ctx, dst); 373 break; 374 } 375 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 32>(ctx, dst); 376 break; 377 } 378 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 4>(ctx, dst); 379 break; 380 } 381 // Volta: 382 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 4>(ctx, dst); 383 } else if (gqa_ratio % 16 == 0) { 384 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 16>(ctx, dst); 385 } else { 386 ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 4>(ctx, dst); 387 } 388 } break; 389 default: 390 GGML_ABORT("fatal error"); 391 break; 392 } 393} 394 395#define FATTN_VEC_CASE(D, type_K_case, type_V_case) \ 396 if constexpr (GGML_CUDA_FA_##type_K_case##_##type_V_case) { \ 397 const bool type_K_okay = type_K == GGML_TYPE_##type_K_case || (type_K == GGML_TYPE_F32 && GGML_TYPE_##type_K_case == GGML_TYPE_F16); \ 398 const bool type_V_okay = type_V == GGML_TYPE_##type_V_case || (type_V == GGML_TYPE_F32 && GGML_TYPE_##type_V_case == GGML_TYPE_F16); \ 399 if (head_size == (D) && type_K_okay && type_V_okay) { \ 400 return ggml_cuda_flash_attn_ext_vec_case<D, GGML_TYPE_##type_K_case, GGML_TYPE_##type_V_case>; \ 401 } \ 402 } \ 403 404#define FATTN_VEC_CASES_ALL_D(type_K_case, type_V_case) \ 405 FATTN_VEC_CASE( 64, type_K_case, type_V_case) \ 406 FATTN_VEC_CASE(128, type_K_case, type_V_case) \ 407 FATTN_VEC_CASE(256, type_K_case, type_V_case) \ 408 409typedef void (* fattn_vec_case_t)(ggml_backend_cuda_context & ctx, ggml_tensor * dst); 410 411// Vector kernel for the given head size and K/V types, nullptr if its template instance was not compiled: 412static fattn_vec_case_t ggml_cuda_get_fattn_vec_case(const int64_t head_size, const ggml_type type_K, const ggml_type type_V) { 413 FATTN_VEC_CASES_ALL_D(F16, F16) 414 FATTN_VEC_CASES_ALL_D(Q4_0, F16) 415 FATTN_VEC_CASES_ALL_D(Q4_1, F16) 416 FATTN_VEC_CASES_ALL_D(Q5_0, F16) 417 FATTN_VEC_CASES_ALL_D(Q5_1, F16) 418 FATTN_VEC_CASES_ALL_D(Q8_0, F16) 419 FATTN_VEC_CASES_ALL_D(BF16, F16) 420 421 FATTN_VEC_CASES_ALL_D(F16, Q4_0) 422 FATTN_VEC_CASES_ALL_D(Q4_0, Q4_0) 423 FATTN_VEC_CASES_ALL_D(Q4_1, Q4_0) 424 FATTN_VEC_CASES_ALL_D(Q5_0, Q4_0) 425 FATTN_VEC_CASES_ALL_D(Q5_1, Q4_0) 426 FATTN_VEC_CASES_ALL_D(Q8_0, Q4_0) 427 FATTN_VEC_CASES_ALL_D(BF16, Q4_0) 428 429 FATTN_VEC_CASES_ALL_D(F16, Q4_1) 430 FATTN_VEC_CASES_ALL_D(Q4_0, Q4_1) 431 FATTN_VEC_CASES_ALL_D(Q4_1, Q4_1) 432 FATTN_VEC_CASES_ALL_D(Q5_0, Q4_1) 433 FATTN_VEC_CASES_ALL_D(Q5_1, Q4_1) 434 FATTN_VEC_CASES_ALL_D(Q8_0, Q4_1) 435 FATTN_VEC_CASES_ALL_D(BF16, Q4_1) 436 437 FATTN_VEC_CASES_ALL_D(F16, Q5_0) 438 FATTN_VEC_CASES_ALL_D(Q4_0, Q5_0) 439 FATTN_VEC_CASES_ALL_D(Q4_1, Q5_0) 440 FATTN_VEC_CASES_ALL_D(Q5_0, Q5_0) 441 FATTN_VEC_CASES_ALL_D(Q5_1, Q5_0) 442 FATTN_VEC_CASES_ALL_D(Q8_0, Q5_0) 443 FATTN_VEC_CASES_ALL_D(BF16, Q5_0) 444 445 FATTN_VEC_CASES_ALL_D(F16, Q5_1) 446 FATTN_VEC_CASES_ALL_D(Q4_0, Q5_1) 447 FATTN_VEC_CASES_ALL_D(Q4_1, Q5_1) 448 FATTN_VEC_CASES_ALL_D(Q5_0, Q5_1) 449 FATTN_VEC_CASES_ALL_D(Q5_1, Q5_1) 450 FATTN_VEC_CASES_ALL_D(Q8_0, Q5_1) 451 FATTN_VEC_CASES_ALL_D(BF16, Q5_1) 452 453 FATTN_VEC_CASES_ALL_D(F16, Q8_0) 454 FATTN_VEC_CASES_ALL_D(Q4_0, Q8_0) 455 FATTN_VEC_CASES_ALL_D(Q4_1, Q8_0) 456 FATTN_VEC_CASES_ALL_D(Q5_0, Q8_0) 457 FATTN_VEC_CASES_ALL_D(Q5_1, Q8_0) 458 FATTN_VEC_CASES_ALL_D(Q8_0, Q8_0) 459 FATTN_VEC_CASES_ALL_D(BF16, Q8_0) 460 461 FATTN_VEC_CASES_ALL_D(F16, BF16) 462 FATTN_VEC_CASES_ALL_D(Q4_0, BF16) 463 FATTN_VEC_CASES_ALL_D(Q4_1, BF16) 464 FATTN_VEC_CASES_ALL_D(Q5_0, BF16) 465 FATTN_VEC_CASES_ALL_D(Q5_1, BF16) 466 FATTN_VEC_CASES_ALL_D(Q8_0, BF16) 467 FATTN_VEC_CASES_ALL_D(BF16, BF16) 468 469 return nullptr; 470} 471 472static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { 473 const ggml_tensor * Q = dst->src[0]; 474 const ggml_tensor * K = dst->src[1]; 475 const ggml_tensor * V = dst->src[2]; 476 477 fattn_vec_case_t vec_case = ggml_cuda_get_fattn_vec_case(Q->ne[0], K->type, V->type); 478 if (vec_case == nullptr) { 479 static bool warned = false; 480 if (!warned) { 481 GGML_LOG_WARN("%s: no FlashAttention vector kernel compiled for K/V types %s-%s, converting K and V to f16 instead (slow). " 482 "Add \"%s-%s\" to GGML_CUDA_FA_QUANTS to compile it.\n", 483 __func__, ggml_type_name(K->type), ggml_type_name(V->type), ggml_type_name(K->type), ggml_type_name(V->type)); 484 warned = true; 485 } 486 vec_case = ggml_cuda_get_fattn_vec_case(Q->ne[0], GGML_TYPE_F16, GGML_TYPE_F16); 487 } 488 GGML_ASSERT(vec_case != nullptr); 489 vec_case(ctx, dst); 490} 491 492// Best FlashAttention kernel for a specific GPU: 493enum best_fattn_kernel { 494 BEST_FATTN_KERNEL_NONE = 0, 495 BEST_FATTN_KERNEL_TILE = 200, 496 BEST_FATTN_KERNEL_VEC = 100, 497 BEST_FATTN_KERNEL_MMA_F16 = 400, 498}; 499 500// K/V types for which there is a vector kernel template instance, other kernels convert these to f16: 501static bool ggml_cuda_fattn_kv_type_supported(const ggml_type type) { 502 switch (type) { 503 case GGML_TYPE_F32: 504 case GGML_TYPE_F16: 505 case GGML_TYPE_BF16: 506 case GGML_TYPE_Q4_0: 507 case GGML_TYPE_Q4_1: 508 case GGML_TYPE_Q5_0: 509 case GGML_TYPE_Q5_1: 510 case GGML_TYPE_Q8_0: 511 return true; 512 default: 513 return false; 514 } 515} 516 517static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const ggml_tensor * dst) { 518#ifndef FLASH_ATTN_AVAILABLE 519 GGML_UNUSED(device); GGML_UNUSED(dst); 520 return BEST_FATTN_KERNEL_NONE; 521#endif// FLASH_ATTN_AVAILABLE 522 523 const ggml_tensor * KQV = dst; 524 const ggml_tensor * Q = dst->src[0]; 525 const ggml_tensor * K = dst->src[1]; 526 const ggml_tensor * V = dst->src[2]; 527 const ggml_tensor * mask = dst->src[3]; 528 529 const int gqa_ratio = Q->ne[2] / K->ne[2]; 530 GGML_ASSERT(Q->ne[2] % K->ne[2] == 0); 531 532 float max_bias = 0.0f; 533 memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float)); 534 535 // The effective batch size for the kernel can be increased by gqa_ratio. 536 // The kernel versions without this optimization are also used for ALiBi, if there is no mask, or if the KV cache is not padded, 537 bool gqa_opt_applies = gqa_ratio >= 2 && mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0; 538 for (const ggml_tensor * t : {Q, K, V, mask}) { 539 if (t == nullptr || ggml_is_quantized(t->type)) { 540 continue; 541 } 542 for (size_t i = 1; i < GGML_MAX_DIMS; ++i) { 543 if (t->nb[i] % 16 != 0) { 544 gqa_opt_applies = false; 545 break; 546 } 547 } 548 } 549 550 const int cc = ggml_cuda_info().devices[device].cc; 551 552 switch (K->ne[0]) { 553 case 40: 554 case 64: 555 case 72: 556 case 80: 557 case 96: 558 case 128: 559 case 112: 560 case 256: 561 if (V->ne[0] != K->ne[0]) { 562 return BEST_FATTN_KERNEL_NONE; 563 } 564 break; 565 case 192: 566 if (V->ne[0] != 128 || !gqa_opt_applies) { 567 return BEST_FATTN_KERNEL_NONE; 568 } 569 if (gqa_ratio % 8 != 0) { 570 return BEST_FATTN_KERNEL_NONE; 571 } 572 break; 573 case 320: 574 if (V->ne[0] != 256 || !gqa_opt_applies) { 575 return BEST_FATTN_KERNEL_NONE; 576 } 577 if (gqa_ratio % 32 != 0) { 578 return BEST_FATTN_KERNEL_NONE; 579 } 580 break; 581 case 512: 582 if (V->ne[0] != K->ne[0]) { 583 return BEST_FATTN_KERNEL_NONE; 584 } 585 if (!gqa_opt_applies) { 586 return BEST_FATTN_KERNEL_NONE; 587 } 588 break; 589 case 576: 590 if (V->ne[0] != 512) { 591 return BEST_FATTN_KERNEL_NONE; 592 } 593 if (!gqa_opt_applies) { 594 return BEST_FATTN_KERNEL_NONE; 595 } 596 break; 597 default: 598 return BEST_FATTN_KERNEL_NONE; 599 } 600 601 if (!ggml_cuda_fattn_kv_type_supported(K->type) || !ggml_cuda_fattn_kv_type_supported(V->type)) { 602 return BEST_FATTN_KERNEL_NONE; 603 } 604 605 if (mask && mask->ne[2] != 1) { 606 return BEST_FATTN_KERNEL_NONE; 607 } 608 609 // For small batch sizes the vector kernel may be preferable over the kernels optimized for large batch sizes: 610 // 192 satisfies % 64 == 0 but has no vec instance (DKQ != DV); force it onto the MMA path. 611 const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && Q->ne[0] != 192 && K->ne[1] % FATTN_KQ_STRIDE == 0; 612 613 // If Turing tensor cores are available, use them: 614 if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { 615 if (can_use_vector_kernel) { 616 if (!ggml_is_quantized(K->type) && !ggml_is_quantized(V->type)) { 617 if (cc >= GGML_CUDA_CC_ADA_LOVELACE && Q->ne[1] == 1 && Q->ne[3] == 1 && !(gqa_ratio > 4 && K->ne[1] >= 8192)) { 618 return BEST_FATTN_KERNEL_VEC; 619 } 620 } else { 621 if (cc >= GGML_CUDA_CC_ADA_LOVELACE) { 622 if (Q->ne[1] <= 2) { 623 return BEST_FATTN_KERNEL_VEC; 624 } 625 } else { 626 if (Q->ne[1] == 1) { 627 return BEST_FATTN_KERNEL_VEC; 628 } 629 } 630 } 631 if (!gqa_opt_applies && Q->ne[1] == 1) { 632 return BEST_FATTN_KERNEL_VEC; 633 } 634 } 635 return BEST_FATTN_KERNEL_MMA_F16; 636 } 637 638 const int ncols2_max = Q->ne[0] == 320 ? 32 : ((Q->ne[0] == 576 || Q->ne[0] == 192) ? 16 : 8); 639 int gqa_ratio_eff = 1; 640 while (gqa_ratio % (2*gqa_ratio_eff) == 0 && gqa_ratio_eff < ncols2_max) { 641 gqa_ratio_eff *= 2; 642 } 643 644 if (volta_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { 645 if (can_use_vector_kernel && Q->ne[1] * gqa_ratio_eff <= 2) { 646 return BEST_FATTN_KERNEL_VEC; 647 } 648 if (Q->ne[1] * gqa_ratio_eff <= 16) { 649 return BEST_FATTN_KERNEL_TILE; // On Volta tensor cores are only faster for sufficiently large matrices. 650 } 651 return BEST_FATTN_KERNEL_MMA_F16; 652 } 653 654 // AMD MFMA needs a certain minimum batch size to outscale the tile kernel for large head sizes. 655 if ((amd_mfma_available(cc) && Q->ne[0] <= 256) && Q->ne[0] != 40 && Q->ne[0] != 72) { 656 if ((Q->ne[0] <= 64 && Q->ne[1] * gqa_ratio_eff > 8)) { 657 return BEST_FATTN_KERNEL_MMA_F16; 658 } 659 if ((Q->ne[0] <= 128 && Q->ne[1] * gqa_ratio_eff > 16)) { 660 return BEST_FATTN_KERNEL_MMA_F16; 661 } 662 if ((Q->ne[0] <= 256 && Q->ne[1] * gqa_ratio_eff > 64)) { 663 return BEST_FATTN_KERNEL_MMA_F16; 664 } 665 } 666 667 // AMD WMMA is faster than the tile kernel if the wide tiles with high arithmetic intensity can be utilized. 668 if ((amd_wmma_available(cc) && gqa_opt_applies && Q->ne[0] <= 256) && Q->ne[0] != 40 && Q->ne[0] != 72 && 669 Q->ne[1] * gqa_ratio_eff > (Q->ne[0] <= 128 ? 8 : 16)) { 670 return BEST_FATTN_KERNEL_MMA_F16; 671 } 672 673 // If there are no tensor cores available, use the generic tile kernel: 674 if (can_use_vector_kernel) { 675 if (!ggml_is_quantized(K->type) && !ggml_is_quantized(V->type)) { 676 if (Q->ne[1] == 1) { 677 if (!gqa_opt_applies) { 678 return BEST_FATTN_KERNEL_VEC; 679 } 680 } 681 } else { 682 if (Q->ne[1] <= 2) { 683 return BEST_FATTN_KERNEL_VEC; 684 } 685 } 686 } 687 return BEST_FATTN_KERNEL_TILE; 688} 689 690size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * dst) { 691 GGML_ASSERT(dst->op == GGML_OP_FLASH_ATTN_EXT); 692 693 const ggml_tensor * Q = dst->src[0]; 694 const ggml_tensor * K = dst->src[1]; 695 const ggml_tensor * V = dst->src[2]; 696 697 GGML_ASSERT(K != nullptr); 698 GGML_ASSERT(V != nullptr); 699 700 const best_fattn_kernel kernel = ggml_cuda_get_best_fattn_kernel(device, dst); 701 702 bool need_f16_K = false; 703 bool need_f16_V = false; 704 705 switch (kernel) { 706 case BEST_FATTN_KERNEL_TILE: 707 case BEST_FATTN_KERNEL_MMA_F16: 708 need_f16_K = true; 709 need_f16_V = true; 710 break; 711 case BEST_FATTN_KERNEL_VEC: { 712 const bool f16_fallback = ggml_cuda_get_fattn_vec_case(Q->ne[0], K->type, V->type) == nullptr; 713 need_f16_K = K->type == GGML_TYPE_F32 || f16_fallback; 714 need_f16_V = V->type == GGML_TYPE_F32 || f16_fallback; 715 } break; 716 case BEST_FATTN_KERNEL_NONE: 717 break; 718 } 719 720 const ggml_cuda_flash_attn_ext_f16_extra_data f16_extra = 721 ggml_cuda_flash_attn_ext_get_f16_extra_data(dst, need_f16_K, need_f16_V); 722 723 return f16_extra.end - (uintptr_t) dst->data; 724} 725 726void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { 727 ggml_cuda_set_device(ctx.device); 728 switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { 729 case BEST_FATTN_KERNEL_NONE: 730 GGML_ABORT("fatal error"); 731 case BEST_FATTN_KERNEL_TILE: 732 ggml_cuda_flash_attn_ext_tile(ctx, dst); 733 break; 734 case BEST_FATTN_KERNEL_VEC: 735 ggml_cuda_flash_attn_ext_vec(ctx, dst); 736 break; 737 case BEST_FATTN_KERNEL_MMA_F16: 738 ggml_cuda_flash_attn_ext_mma_f16(ctx, dst); 739 break; 740 } 741} 742 743bool ggml_cuda_flash_attn_ext_supported(int device, const ggml_tensor * dst) { 744 return ggml_cuda_get_best_fattn_kernel(device, dst) != BEST_FATTN_KERNEL_NONE; 745}