char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
7bb0fc18f
char/aliceai-t5-moe
1#include "common.h" 2#include "dequantize.h" 3 4// dequantize a quantized KV cache tensor to contiguous F16 before running the F16 flash attention kernels 5// - one thread per block; dispatched separately for K and V 6// - ref: https://github.com/ggml-org/llama.cpp/pull/27390 7template < 8 typename block_t, 9 short QK, 10 void (*deq_t4x4)(device const block_t *, short, thread float4x4 &)> 11kernel void kernel_flash_attn_ext_kv_f16( 12 constant ggml_metal_kargs_flash_attn_ext_kv_f16 & args, 13 device const char * x, 14 device half * x_dst, 15 uint gid [[thread_position_in_grid]]) { 16 if (gid >= (uint) args.nblocks) { 17 return; 18 } 19 20 const uint nb = args.ne0/QK; 21 const uint i0 = gid%nb; 22 uint ib = gid/nb; 23 const uint i1 = ib%args.ne1; 24 ib /= args.ne1; 25 const uint i2 = ib%args.ne2; 26 const uint i3 = ib/args.ne2; 27 28 const uint64_t offs = i0*args.nb0 + i1*args.nb1 + i2*args.nb2 + i3*args.nb3; 29 30 device const block_t * src = (device const block_t *) (x + offs); 31 device half4 * dst = (device half4 *) x_dst + (QK/4)*gid; 32 33 for (short i = 0; i < QK/16; ++i) { 34 float4x4 reg; 35 deq_t4x4(src, i, reg); 36 dst[4*i + 0] = (half4) reg[0]; 37 dst[4*i + 1] = (half4) reg[1]; 38 dst[4*i + 2] = (half4) reg[2]; 39 dst[4*i + 3] = (half4) reg[3]; 40 } 41} 42 43typedef decltype(kernel_flash_attn_ext_kv_f16<block_q8_0, 32, dequantize_q8_0>) kernel_flash_attn_ext_kv_f16_t; 44 45template [[host_name("kernel_flash_attn_ext_kv_q4_0_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16<block_q4_0, 32, dequantize_q4_0>; 46template [[host_name("kernel_flash_attn_ext_kv_q4_1_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16<block_q4_1, 32, dequantize_q4_1>; 47template [[host_name("kernel_flash_attn_ext_kv_q5_0_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16<block_q5_0, 32, dequantize_q5_0>; 48template [[host_name("kernel_flash_attn_ext_kv_q5_1_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16<block_q5_1, 32, dequantize_q5_1>; 49template [[host_name("kernel_flash_attn_ext_kv_q8_0_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16<block_q8_0, 32, dequantize_q8_0>; 50 51constant bool FC_flash_attn_ext_pad_has_mask [[function_constant(FC_FLASH_ATTN_EXT_PAD + 0)]]; 52 53constant int32_t FC_flash_attn_ext_pad_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_PAD + 25)]]; 54 55// pad the last chunk of C elements of k and v into a an extra pad buffer 56kernel void kernel_flash_attn_ext_pad( 57 constant ggml_metal_kargs_flash_attn_ext_pad & args, 58 device const char * k, 59 device const char * v, 60 device const char * mask, 61 device char * dst, 62 uint3 tgpig[[threadgroup_position_in_grid]], 63 ushort tiitg[[thread_index_in_threadgroup]], 64 ushort3 ntg[[threads_per_threadgroup]]) { 65 const int32_t C = FC_flash_attn_ext_pad_ncpsg; 66 67 device char * k_pad = dst; 68 device char * v_pad = k_pad + args.nb11*C*args.ne_12_2*args.ne_12_3; 69 device char * mask_pad = v_pad + args.nb21*C*args.ne_12_2*args.ne_12_3; 70 71 const int32_t icp = args.ne11 % C; 72 const int32_t ic0 = args.ne11 - icp; 73 74 const int32_t i1 = tgpig[0]; 75 const int32_t i2 = tgpig[1]; 76 const int32_t i3 = tgpig[2]; 77 78 if (i2 < args.ne_12_2 && i3 < args.ne_12_3) { 79 device const char * k_src = k + args.nb11*(ic0 + i1) + args.nb12*i2 + args.nb13*i3; 80 device const char * v_src = v + args.nb21*(ic0 + i1) + args.nb22*i2 + args.nb23*i3; 81 82 device char * k_dst = k_pad + args.nb11*i1 + args.nb11*C*i2 + args.nb11*C*args.ne_12_2*i3; 83 device char * v_dst = v_pad + args.nb21*i1 + args.nb21*C*i2 + args.nb21*C*args.ne_12_2*i3; 84 85 if (i1 >= icp) { 86 // here it is not important the exact value that will be used as we rely on masking out the scores in the attention 87 for (uint64_t i = tiitg; i < args.nb11; i += ntg.x) { 88 k_dst[i] = 0; 89 } 90 for (uint64_t i = tiitg; i < args.nb21; i += ntg.x) { 91 v_dst[i] = 0; 92 } 93 } else { 94 for (uint64_t i = tiitg; i < args.nb11; i += ntg.x) { 95 k_dst[i] = k_src[i]; 96 } 97 for (uint64_t i = tiitg; i < args.nb21; i += ntg.x) { 98 v_dst[i] = v_src[i]; 99 } 100 } 101 } 102 103 if (FC_flash_attn_ext_pad_has_mask) { 104 if (i2 < args.ne32 && i3 < args.ne33) { 105 for (int ib = i1; ib < args.ne31; ib += C) { 106 device const half * mask_src = (device const half *)(mask + args.nb31*ib + args.nb32*i2 + args.nb33*i3) + ic0; 107 device half * mask_dst = (device half *)(mask_pad) + C*ib + C*args.ne31*i2 + C*args.ne31*args.ne32*i3; 108 109 for (int i = tiitg; i < C; i += ntg.x) { 110 if (i >= icp) { 111 mask_dst[i] = -MAXHALF; 112 } else { 113 mask_dst[i] = mask_src[i]; 114 } 115 } 116 } 117 } 118 } 119} 120 121constant int32_t FC_flash_attn_ext_blk_nqptg [[function_constant(FC_FLASH_ATTN_EXT_BLK + 24)]]; 122constant int32_t FC_flash_attn_ext_blk_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_BLK + 25)]]; 123 124// scan the blocks of the mask that are not masked 125// 0 - masked (i.e. full of -INF, skip) 126// 1 - not masked (i.e. at least one element of the mask is not -INF) 127// 2 - all zero 128kernel void kernel_flash_attn_ext_blk( 129 constant ggml_metal_kargs_flash_attn_ext_blk & args, 130 device const char * mask, 131 device char * dst, 132 uint3 tgpig[[threadgroup_position_in_grid]], 133 ushort tiisg[[thread_index_in_simdgroup]]) { 134 // block size C x Q 135 const int32_t Q = FC_flash_attn_ext_blk_nqptg; 136 const int32_t C = FC_flash_attn_ext_blk_ncpsg; 137 138 constexpr short NW = N_SIMDWIDTH; 139 140 const int32_t i3 = tgpig[2]/args.ne32; 141 const int32_t i2 = tgpig[2]%args.ne32; 142 const int32_t i1 = tgpig[1]; 143 const int32_t i0 = tgpig[0]; 144 145 char res = i0*C + C > args.ne30 ? 1 : 0; 146 147 device const half * mask_src = (device const half *) (mask + (i1*Q)*args.nb31 + i2*args.nb32 + i3*args.nb33) + i0*C + tiisg; 148 149 // detailed check of the elements of the block 150 if ((C > NW || Q > 1) && res == 0) { 151 half mmin = MAXHALF; 152 half mmax = -MAXHALF; 153 154 FOR_UNROLL (short j = 0; j < Q; ++j) { 155 FOR_UNROLL (short ii = 0; ii < C/NW; ++ii) { 156 mmin = min(mmin, mask_src[ii*NW]); 157 mmax = max(mmax, mask_src[ii*NW]); 158 } 159 160 mask_src += args.nb31/2; 161 } 162 163 mmin = simd_min(mmin); 164 mmax = simd_max(mmax); 165 166 if (mmax > -MAXHALF) { 167 if (mmin == 0.0 && mmax == 0.0) { 168 res = 2; 169 } else { 170 res = 1; 171 } 172 } 173 } 174 175 const int32_t nblk1 = ((args.ne01 + Q - 1)/Q); 176 const int32_t nblk0 = ((args.ne30 + C - 1)/C); 177 178 if (tiisg == 0) { 179 dst[((i3*args.ne32 + i2)*nblk1 + i1)*nblk0 + i0] = res; 180 } 181} 182 183constant bool FC_flash_attn_ext_has_mask [[function_constant(FC_FLASH_ATTN_EXT + 0)]]; 184constant bool FC_flash_attn_ext_has_sinks [[function_constant(FC_FLASH_ATTN_EXT + 1)]]; 185constant bool FC_flash_attn_ext_has_bias [[function_constant(FC_FLASH_ATTN_EXT + 2)]]; 186constant bool FC_flash_attn_ext_has_scap [[function_constant(FC_FLASH_ATTN_EXT + 3)]]; 187constant bool FC_flash_attn_ext_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT + 4)]]; 188 189constant bool FC_flash_attn_ext_bc_mask [[function_constant(FC_FLASH_ATTN_EXT + 10)]]; 190 191//constant float FC_flash_attn_ext_scale [[function_constant(FC_FLASH_ATTN_EXT + 10)]]; 192//constant float FC_flash_attn_ext_max_bias [[function_constant(FC_FLASH_ATTN_EXT + 11)]]; 193//constant float FC_flash_attn_ext_logit_softcap [[function_constant(FC_FLASH_ATTN_EXT + 12)]]; 194 195constant int32_t FC_flash_attn_ext_ns10 [[function_constant(FC_FLASH_ATTN_EXT + 20)]]; 196constant int32_t FC_flash_attn_ext_ns20 [[function_constant(FC_FLASH_ATTN_EXT + 21)]]; 197constant int32_t FC_flash_attn_ext_nsg [[function_constant(FC_FLASH_ATTN_EXT + 22)]]; 198 199// ref: https://arxiv.org/pdf/2307.08691.pdf 200template< 201 typename q_t, // query types in shared memory 202 typename q4_t, 203 typename q8x8_t, 204 typename k_t, // key types in shared memory 205 typename k4x4_t, 206 typename k8x8_t, 207 typename v_t, // value types in shared memory 208 typename v4x4_t, 209 typename v8x8_t, 210 typename qk_t, // Q*K types 211 typename qk8x8_t, 212 typename s_t, // soft-max types 213 typename s2_t, 214 typename s8x8_t, 215 typename o_t, // attention accumulation types 216 typename o4_t, 217 typename o8x8_t, 218 typename kd4x4_t, // key type in device memory 219 short nl_k, 220 void (*deq_k)(device const kd4x4_t *, short, thread k4x4_t &), 221 typename vd4x4_t, // value type in device memory 222 short nl_v, 223 void (*deq_v)(device const vd4x4_t *, short, thread v4x4_t &), 224 short DK, // K head size 225 short DV, // V head size 226 short Q, // queries per threadgroup 227 short C, // cache items per threadgroup 228 short NSG> // number of simd groups 229void kernel_flash_attn_ext_impl( 230 constant ggml_metal_kargs_flash_attn_ext & args, 231 device const char * q, 232 device const char * k, 233 device const char * v, 234 device const char * mask, 235 device const char * sinks, 236 device const char * pad, 237 device const char * blk, 238 device char * dst, 239 threadgroup half * shmem_f16, 240 uint3 tgpig, 241 ushort tiisg, 242 ushort sgitg) { 243 const ushort iq3 = tgpig[2]; 244 const ushort iq2 = tgpig[1]; 245 const ushort iq1 = tgpig[0]*Q; 246 247#define NS10 (FC_flash_attn_ext_ns10) 248#define NS20 (FC_flash_attn_ext_ns20) 249 250 // note: I had some concerns that using this instead of the ugly macros above was affecting performance 251 // need to re-check carefully and if no regressions are observerd - remove the macros 252 // the concerns is that maybe using const variables requires extra registers? but not sure if the compiler 253 // is clever enough to avoid this. unfortunately, using constexpr is not possible with FC 254 //const short NS10 = FC_flash_attn_ext_ns10; 255 //const short NS20 = FC_flash_attn_ext_ns20; 256 257 constexpr short KV = 8; 258 259 constexpr short DK4 = DK/4; 260 constexpr short DK8 = DK/8; 261 constexpr short DK16 = DK/16; 262 constexpr short DV4 = DV/4; 263 //constexpr short DV8 = DV/8; 264 constexpr short DV16 = DV/16; 265 266 constexpr short PV = PAD2(DV, 64); 267 constexpr short PV4 = PV/4; 268 constexpr short PV8 = PV/8; 269 //constexpr short PV16 = PV/16; 270 271 constexpr short NW = N_SIMDWIDTH; 272 constexpr short NQ = Q/NSG; 273 constexpr short SH = 2*C; // shared memory per simdgroup (s_t == float) 274 275 constexpr short TS = 2*SH; 276 constexpr short T = DK + 2*PV; // shared memory size per query in (half) 277 278 threadgroup q_t * sq = (threadgroup q_t *) (shmem_f16 + 0*T); // holds the query data 279 threadgroup q4_t * sq4 = (threadgroup q4_t *) (shmem_f16 + 0*T); // same as above but in q4_t 280 threadgroup o_t * so = (threadgroup o_t *) (shmem_f16 + 0*T + Q*DK); // the result for all queries in 8x8 matrices (the O matrix from the paper) 281 threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + 0*T + Q*DK); 282 threadgroup s_t * ss = (threadgroup s_t *) (shmem_f16 + Q*T); // scratch buffer for attention, mask and diagonal matrix 283 threadgroup s2_t * ss2 = (threadgroup s2_t *) (shmem_f16 + Q*T); // same as above but in s2_t 284 285 threadgroup k_t * sk = (threadgroup k_t *) (shmem_f16 + sgitg*(4*16*KV) + Q*T + Q*TS); // scratch buffer to load K in shared memory 286 threadgroup k4x4_t * sk4x4 = (threadgroup k4x4_t *) (shmem_f16 + sgitg*(4*16*KV) + Q*T + Q*TS); // same as above but in k4x4_t 287 288 threadgroup v_t * sv = (threadgroup v_t *) (shmem_f16 + sgitg*(4*16*KV) + Q*T + Q*TS); // scratch buffer to load V in shared memory 289 threadgroup v4x4_t * sv4x4 = (threadgroup v4x4_t *) (shmem_f16 + sgitg*(4*16*KV) + Q*T + Q*TS); // same as above but in v4x4_t 290 291 // mask storage in shared mem 292 threadgroup half2 * sm2 = (threadgroup half2 *) (shmem_f16 + Q*T + 2*C); 293 294 // per-query mask pointers 295 device const half2 * pm2[NQ]; 296 297 FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { 298 const short j = jj*NSG + sgitg; 299 300 pm2[jj] = (device const half2 *) ((device const char *) mask + (iq1 + j)*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); 301 } 302 303 { 304 const int32_t nblk1 = ((args.ne01 + Q - 1)/Q); 305 const int32_t nblk0 = ((args.ne11 + C - 1)/C); 306 307 blk += (((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*nblk1 + iq1/Q)*nblk0; 308 } 309 310 { 311 q += iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03; 312 313 const short ikv2 = iq2/(args.ne02/args.ne_12_2); 314 const short ikv3 = iq3/(args.ne03/args.ne_12_3); 315 316 k += ikv2*args.nb12 + ikv3*args.nb13; 317 v += ikv2*args.nb22 + ikv3*args.nb23; 318 } 319 320 // load heads from Q to shared memory 321 FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { 322 const short j = jj*NSG + sgitg; 323 324 device const float4 * q4 = (device const float4 *) ((device const char *) q + j*args.nb01); 325 326 for (short i = tiisg; i < DK4; i += NW) { 327 if (iq1 + j < args.ne01) { 328 sq4[j*DK4 + i] = (q4_t) q4[i]; 329 } else { 330 sq4[j*DK4 + i] = 0; 331 } 332 } 333 } 334 335 // zero out 336 FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { 337 const short j = jj*NSG + sgitg; 338 339 for (short i = tiisg; i < DV4; i += NW) { 340 so4[j*PV4 + i] = 0; 341 } 342 343 for (short i = tiisg; i < SH; i += NW) { 344 ss[j*SH + i] = 0.0f; 345 } 346 } 347 348 threadgroup_barrier(mem_flags::mem_threadgroup); 349 350 float S[NQ] = { [0 ... NQ-1] = 0.0f }; 351 352 { 353 float M[NQ] = { [0 ... NQ-1] = -FLT_MAX/2 }; 354 355 float slope = 1.0f; 356 357 // ALiBi 358 if (FC_flash_attn_ext_has_bias) { 359 const short h = iq2; 360 361 const float base = h < args.n_head_log2 ? args.m0 : args.m1; 362 const short exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1; 363 364 slope = pow(base, exph); 365 } 366 367 // loop over the KV cache 368 // each simdgroup handles blocks of Q rows and C columns 369 for (int ic0 = 0; ; ++ic0) { 370 int ic = ic0*C; 371 if (ic >= args.ne11) { 372 break; 373 } 374 375 // the last partial chunk uses the pad buffer as source 376 if (FC_flash_attn_ext_has_kvpad && ic + C > args.ne11) { 377 k = pad; 378 v = k + args.nb11*C*args.ne_12_2*args.ne_12_3; 379 mask = v + args.nb21*C*args.ne_12_2*args.ne_12_3; 380 381 const short ikv2 = iq2/(args.ne02/args.ne_12_2); 382 const short ikv3 = iq3/(args.ne03/args.ne_12_3); 383 384 k += (ikv2 + ikv3*args.ne_12_2)*args.nb11*C; 385 v += (ikv2 + ikv3*args.ne_12_2)*args.nb21*C; 386 387 if (!FC_flash_attn_ext_has_mask) { 388 threadgroup half * sm = (threadgroup half *) (sm2); 389 390 FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { 391 const short j = jj*NSG + sgitg; 392 393 for (short i = tiisg; i < C; i += NW) { 394 if (ic + i >= args.ne11) { 395 sm[2*j*SH + i] = -MAXHALF; 396 } 397 } 398 } 399 } else { 400 FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { 401 const short j = jj*NSG + sgitg; 402 403 pm2[jj] = (device const half2 *) ((device const half *) mask + 404 (iq1 + j)*C + 405 (iq2%args.ne32)*(C*args.ne31) + 406 (iq3%args.ne33)*(C*args.ne31*args.ne32)); 407 } 408 } 409 410 ic = 0; 411 } 412 413 char blk_cur = 1; 414 415 // read the mask into shared mem 416 if (FC_flash_attn_ext_has_mask) { 417 blk_cur = blk[ic0]; 418 419 if (blk_cur == 0) { 420 FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { 421 pm2[jj] += NW; 422 } 423 424 continue; 425 } 426 427 if (blk_cur == 1) { 428 FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { 429 const short j = jj*NSG + sgitg; 430 431 if (FC_flash_attn_ext_bc_mask) { 432 sm2[j*SH + tiisg] = (iq1 + j) < args.ne31 ? pm2[jj][tiisg] : half2(-MAXHALF, -MAXHALF); 433 } else { 434 sm2[j*SH + tiisg] = pm2[jj][tiisg]; 435 } 436 437 pm2[jj] += NW; 438 } 439 } else if (blk_cur == 2) { 440 FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { 441 pm2[jj] += NW; 442 } 443 } 444 445#if 0 446 // note: old -INF block optimization - obsoleted by pre-computing non-masked blocks 447 448 threadgroup_barrier(mem_flags::mem_threadgroup); 449 450 // used to detect blocks full of -INF 451 // skip only when the entire threadgroup is masked 452 half2 smax2(-MAXHALF/2, -MAXHALF/2); 453 454 FOR_UNROLL (short j = 0; j < Q; ++j) { 455 smax2 = max(smax2, sm2[j*SH + tiisg]); 456 } 457 458 smax2 = simd_max(smax2); 459 460 if (max(smax2[0], smax2[1]) <= -MAXHALF/2) { 461 // this barrier is important 462 threadgroup_barrier(mem_flags::mem_threadgroup); 463 464 continue; 465 } 466#endif 467 } 468 469 // Q*K^T 470 // this is compile-time check, so it does not have runtime overhead 471 if (is_same<kd4x4_t, k4x4_t>::value) { 472 // we can read directly from global memory 473 device const k_t * pk = (device const k_t *) (k + ic*args.nb11); 474 threadgroup const q_t * pq = sq; 475 threadgroup s_t * ps = ss; 476 477 pk += sgitg*(8*NS10); 478 ps += sgitg*(8*1); 479 480 static_assert((C/8) % NSG == 0, ""); 481 482 constexpr short NC = (C/8)/NSG; 483 484 FOR_UNROLL (short cc = 0; cc < NC; ++cc) { 485 qk8x8_t mqk = make_filled_simdgroup_matrix<qk_t, 8>((qk_t) 0.0f); 486 487 if (DK % 16 != 0) { 488 k8x8_t mk; 489 q8x8_t mq; 490 491 FOR_UNROLL (short i = 0; i < DK8; ++i) { 492 simdgroup_barrier(mem_flags::mem_none); 493 494 simdgroup_load(mk, pk + 8*i, NS10, 0, true); 495 simdgroup_load(mq, pq + 8*i, DK); 496 497 simdgroup_barrier(mem_flags::mem_none); 498 499 simdgroup_multiply_accumulate(mqk, mq, mk, mqk); 500 } 501 } else { 502 k8x8_t mk[2]; 503 q8x8_t mq[2]; 504 505 // note: too much unroll can tank the performance for large heads 506 #pragma unroll (MIN(DK8/2, 4*NSG)) 507 for (short i = 0; i < DK8/2; ++i) { 508 simdgroup_barrier(mem_flags::mem_none); 509 510 simdgroup_load(mq[0], pq + 0*8 + 16*i, DK); 511 simdgroup_load(mq[1], pq + 1*8 + 16*i, DK); 512 513 simdgroup_load(mk[0], pk + 0*8 + 16*i, NS10, 0, true); 514 simdgroup_load(mk[1], pk + 1*8 + 16*i, NS10, 0, true); 515 516 simdgroup_barrier(mem_flags::mem_none); 517 518 simdgroup_multiply_accumulate(mqk, mq[0], mk[0], mqk); 519 simdgroup_multiply_accumulate(mqk, mq[1], mk[1], mqk); 520 } 521 } 522 523 simdgroup_store(mqk, ps, SH, 0, false); 524 525 pk += 8*(NSG*NS10); 526 ps += 8*(NSG); 527 } 528 } else { 529 // TODO: this is the quantized K cache branch - not optimized yet 530 for (short ccc = 0; ccc < (C/8)/NSG; ++ccc) { 531 const short cc = ccc*NSG + sgitg; 532 533 const short tx = tiisg%4; 534 const short ty = tiisg/4; 535 536 qk8x8_t mqk = make_filled_simdgroup_matrix<qk_t, 8>((qk_t) 0.0f); 537 538 for (short ii = 0; ii < DK16; ii += 4) { 539 device const kd4x4_t * pk4x4 = (device const kd4x4_t *) (k + ((ic + 8*cc + ty)*args.nb11)); 540 541 if (DK16%4 == 0) { 542 // the head is evenly divisible by 4*16 = 64, so no need for bound checks 543 { 544 k4x4_t tmp; 545 deq_k(pk4x4 + (ii + tx)/nl_k, (ii + tx)%nl_k, tmp); 546 sk4x4[4*ty + tx] = tmp; 547 } 548 549 simdgroup_barrier(mem_flags::mem_threadgroup); 550 551 FOR_UNROLL (short k = 0; k < 4; ++k) { 552 k8x8_t mk; 553 q8x8_t mq; 554 555 simdgroup_load(mk, sk + 16*k + 0*8, 4*16, 0, true); // transpose 556 simdgroup_load(mq, sq + (2*(ii + k) + 0)*8, DK); 557 simdgroup_multiply_accumulate(mqk, mq, mk, mqk); 558 559 simdgroup_load(mk, sk + 16*k + 1*8, 4*16, 0, true); // transpose 560 simdgroup_load(mq, sq + (2*(ii + k) + 1)*8, DK); 561 simdgroup_multiply_accumulate(mqk, mq, mk, mqk); 562 } 563 } else { 564 if (ii + tx < DK16) { 565 k4x4_t tmp; 566 deq_k(pk4x4 + (ii + tx)/nl_k, (ii + tx)%nl_k, tmp); 567 sk4x4[4*ty + tx] = tmp; 568 } 569 570 simdgroup_barrier(mem_flags::mem_threadgroup); 571 572 for (short k = 0; k < 4 && ii + k < DK16; ++k) { 573 k8x8_t mk; 574 q8x8_t mq; 575 576 simdgroup_load(mk, sk + 16*k + 0*8, 4*16, 0, true); // transpose 577 simdgroup_load(mq, sq + (2*(ii + k) + 0)*8, DK); 578 simdgroup_multiply_accumulate(mqk, mq, mk, mqk); 579 580 simdgroup_load(mk, sk + 16*k + 1*8, 4*16, 0, true); // transpose 581 simdgroup_load(mq, sq + (2*(ii + k) + 1)*8, DK); 582 simdgroup_multiply_accumulate(mqk, mq, mk, mqk); 583 } 584 } 585 } 586 587 simdgroup_store(mqk, ss + 8*cc, SH, 0, false); 588 } 589 } 590 591 threadgroup_barrier(mem_flags::mem_threadgroup); 592 593 // online softmax 594 FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { 595 const short j = jj*NSG + sgitg; 596 597 const float m = M[jj]; 598 599 // scale and apply the logitcap / mask 600 float2 s2 = ss2[j*SH/2 + tiisg]*args.scale; 601 602 if (FC_flash_attn_ext_has_scap) { 603 s2 = args.logit_softcap*precise::tanh(s2); 604 } 605 606 // mqk = mqk + slope*mask 607 if (blk_cur != 2) { 608 if (FC_flash_attn_ext_has_bias) { 609 s2 += s2_t(sm2[j*SH + tiisg])*slope; 610 } else { 611 s2 += s2_t(sm2[j*SH + tiisg]); 612 } 613 } 614 615 M[jj] = simd_max(max(M[jj], max(s2[0], s2[1]))); 616 617 const float ms = exp(m - M[jj]); 618 const float2 vs2 = exp(s2 - M[jj]); 619 620 S[jj] = S[jj]*ms + simd_sum(vs2[0] + vs2[1]); 621 622 // the P matrix from the paper (Q rows, C columns) 623 ss2[j*SH/2 + tiisg] = vs2; 624 625 if (DV4 % NW == 0) { 626 FOR_UNROLL (short ii = 0; ii < DV4/NW; ++ii) { 627 const short i = ii*NW + tiisg; 628 629 so4[j*PV4 + i] *= ms; 630 } 631 } else { 632 for (short i = tiisg; i < DV4; i += NW) { 633 so4[j*PV4 + i] *= ms; 634 } 635 } 636 } 637 638 threadgroup_barrier(mem_flags::mem_threadgroup); 639 640 // O = O + (Q*K^T)*V 641 { 642 // we can read directly from global memory 643 if (is_same<vd4x4_t, v4x4_t>::value) { 644 static_assert(PV8 % NSG == 0, ""); 645 646 constexpr short NO = PV8/NSG; 647 648 o8x8_t lo[NO]; 649 650 { 651 auto sot = so + 8*sgitg; 652 653 FOR_UNROLL (short ii = 0; ii < NO; ++ii) { 654 simdgroup_load(lo[ii], sot, PV, 0, false); 655 656 sot += 8*NSG; 657 } 658 } 659 660 { 661 device const v_t * pv = (device const v_t *) (v + ic*args.nb21); 662 663 pv += 8*sgitg; 664 665 if (DV <= 64) { 666 FOR_UNROLL (short cc = 0; cc < C/8; ++cc) { 667 s8x8_t vs; 668 simdgroup_load(vs, ss + 8*cc, SH, 0, false); 669 670 FOR_UNROLL (short ii = 0; ii < NO/2; ++ii) { 671 v8x8_t mv[2]; 672 673 simdgroup_load(mv[0], pv + 0*NSG + 16*ii*NSG, NS20, 0, false); 674 simdgroup_load(mv[1], pv + 8*NSG + 16*ii*NSG, NS20, 0, false); 675 676 simdgroup_multiply_accumulate(lo[2*ii + 0], vs, mv[0], lo[2*ii + 0]); 677 simdgroup_multiply_accumulate(lo[2*ii + 1], vs, mv[1], lo[2*ii + 1]); 678 } 679 680 pv += 8*NS20; 681 } 682 } else { 683 constexpr short NC = (C/8)/2; 684 685 FOR_UNROLL (short cc = 0; cc < NC; ++cc) { 686 s8x8_t vs[2]; 687 688 simdgroup_load(vs[0], ss + 16*cc + 0, SH, 0, false); 689 simdgroup_load(vs[1], ss + 16*cc + 8, SH, 0, false); 690 691 FOR_UNROLL (short ii = 0; ii < NO/2; ++ii) { 692 v8x8_t mv[4]; 693 694 simdgroup_load(mv[0], pv + 0*NSG + 16*ii*NSG + 0*8*NS20, NS20, 0, false); 695 simdgroup_load(mv[1], pv + 8*NSG + 16*ii*NSG + 0*8*NS20, NS20, 0, false); 696 simdgroup_load(mv[2], pv + 0*NSG + 16*ii*NSG + 1*8*NS20, NS20, 0, false); 697 simdgroup_load(mv[3], pv + 8*NSG + 16*ii*NSG + 1*8*NS20, NS20, 0, false); 698 699 simdgroup_multiply_accumulate(lo[2*ii + 0], vs[0], mv[0], lo[2*ii + 0]); 700 simdgroup_multiply_accumulate(lo[2*ii + 1], vs[0], mv[1], lo[2*ii + 1]); 701 simdgroup_multiply_accumulate(lo[2*ii + 0], vs[1], mv[2], lo[2*ii + 0]); 702 simdgroup_multiply_accumulate(lo[2*ii + 1], vs[1], mv[3], lo[2*ii + 1]); 703 } 704 705 pv += 2*8*NS20; 706 } 707 } 708 } 709 710 { 711 auto sot = so + 8*sgitg; 712 713 FOR_UNROLL (short ii = 0; ii < NO; ++ii) { 714 simdgroup_store(lo[ii], sot, PV, 0, false); 715 716 sot += 8*NSG; 717 } 718 } 719 } else { 720 // TODO: this is the quantized V cache branch - not optimized yet 721 722 const short tx = tiisg%4; 723 const short ty = tiisg/4; 724 725 for (short cc = 0; cc < C/8; ++cc) { 726 s8x8_t vs; 727 simdgroup_load(vs, ss + 8*cc, SH, 0, false); 728 729 for (short ii = 4*sgitg; ii < DV16; ii += 4*NSG) { 730 device const vd4x4_t * pv4x4 = (device const vd4x4_t *) (v + ((ic + 8*cc + ty)*args.nb21)); 731 732 if (DV16%4 == 0) { 733 // no need for bound checks 734 { 735 v4x4_t tmp; 736 deq_v(pv4x4 + (ii + tx)/nl_v, (ii + tx)%nl_v, tmp); 737 sv4x4[4*ty + tx] = tmp; 738 } 739 740 simdgroup_barrier(mem_flags::mem_threadgroup); 741 742 FOR_UNROLL (short k = 0; k < 4; ++k) { 743 v8x8_t mv[2]; 744 o8x8_t lo[2]; 745 746 simdgroup_load(mv[0], sv + 16*k + 0*8, 4*16, 0, false); 747 simdgroup_load(mv[1], sv + 16*k + 1*8, 4*16, 0, false); 748 simdgroup_load(lo[0], so + 8*(2*(ii + k) + 0), PV, 0, false); 749 simdgroup_load(lo[1], so + 8*(2*(ii + k) + 1), PV, 0, false); 750 751 simdgroup_multiply_accumulate(lo[0], vs, mv[0], lo[0]); 752 simdgroup_multiply_accumulate(lo[1], vs, mv[1], lo[1]); 753 754 simdgroup_store(lo[0], so + 8*(2*(ii + k) + 0), PV, 0, false); 755 simdgroup_store(lo[1], so + 8*(2*(ii + k) + 1), PV, 0, false); 756 } 757 } else { 758 if (ii + tx < DV16) { 759 v4x4_t tmp; 760 deq_v(pv4x4 + (ii + tx)/nl_v, (ii + tx)%nl_v, tmp); 761 sv4x4[4*ty + tx] = tmp; 762 } 763 764 simdgroup_barrier(mem_flags::mem_threadgroup); 765 766 for (short k = 0; k < 4 && ii + k < DV16; ++k) { 767 v8x8_t mv[2]; 768 o8x8_t lo[2]; 769 770 simdgroup_load(mv[0], sv + 16*k + 0*8, 4*16, 0, false); 771 simdgroup_load(mv[1], sv + 16*k + 1*8, 4*16, 0, false); 772 simdgroup_load(lo[0], so + 8*(2*(ii + k) + 0), PV, 0, false); 773 simdgroup_load(lo[1], so + 8*(2*(ii + k) + 1), PV, 0, false); 774 775 simdgroup_multiply_accumulate(lo[0], vs, mv[0], lo[0]); 776 simdgroup_multiply_accumulate(lo[1], vs, mv[1], lo[1]); 777 778 simdgroup_store(lo[0], so + 8*(2*(ii + k) + 0), PV, 0, false); 779 simdgroup_store(lo[1], so + 8*(2*(ii + k) + 1), PV, 0, false); 780 } 781 } 782 } 783 } 784 } 785 } 786 787 threadgroup_barrier(mem_flags::mem_threadgroup); 788 } 789 790 if (FC_flash_attn_ext_has_sinks) { 791 FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { 792 const short j = jj*NSG + sgitg; 793 794 const float m = M[jj]; 795 const float s = tiisg == 0 ? ((device const float *) sinks)[iq2] : -FLT_MAX/2; 796 797 M[jj] = simd_max(max(M[jj], s)); 798 799 const float ms = exp(m - M[jj]); 800 const float vs = exp(s - M[jj]); 801 802 S[jj] = S[jj]*ms + simd_sum(vs); 803 804 for (short i = tiisg; i < DV4; i += NW) { 805 so4[j*PV4 + i] *= ms; 806 } 807 } 808 } 809 } 810 811 // store to global memory 812 for (short jj = 0; jj < NQ; ++jj) { 813 const short j = jj*NSG + sgitg; 814 if (iq1 + j >= args.ne01) { 815 break; 816 } 817 818 device float4 * dst4 = (device float4 *) dst + ((uint64_t)iq3*args.ne2*args.ne1 + iq2 + (uint64_t)(iq1 + j)*args.ne1)*DV4; 819 820 const float scale = S[jj] == 0.0 ? 0.0f : 1.0f/S[jj]; 821 822 if (DV4 % NW == 0) { 823 FOR_UNROLL (short ii = 0; ii < DV4/NW; ++ii) { 824 const short i = ii*NW + tiisg; 825 826 dst4[i] = (float4) so4[j*PV4 + i]*scale; 827 } 828 } else { 829 for (short i = tiisg; i < DV4; i += NW) { 830 dst4[i] = (float4) so4[j*PV4 + i]*scale; 831 } 832 } 833 } 834 835#undef NS10 836#undef NS20 837} 838 839template< 840 typename q_t, // query types in shared memory 841 typename q4_t, 842 typename q8x8_t, 843 typename k_t, // key types in shared memory 844 typename k4x4_t, 845 typename k8x8_t, 846 typename v_t, // value types in shared memory 847 typename v4x4_t, 848 typename v8x8_t, 849 typename qk_t, // Q*K types 850 typename qk8x8_t, 851 typename s_t, // soft-max types 852 typename s2_t, 853 typename s8x8_t, 854 typename o_t, // attention accumulation types 855 typename o4_t, 856 typename o8x8_t, 857 typename kd4x4_t, // key type in device memory 858 short nl_k, 859 void (*deq_k)(device const kd4x4_t *, short, thread k4x4_t &), 860 typename vd4x4_t, // value type in device memory 861 short nl_v, 862 void (*deq_v)(device const vd4x4_t *, short, thread v4x4_t &), 863 short DK, // K head size 864 short DV, // V head size 865 short Q = OP_FLASH_ATTN_EXT_NQPSG, // queries per threadgroup 866 short C = OP_FLASH_ATTN_EXT_NCPSG> // cache items per threadgroup 867kernel void kernel_flash_attn_ext( 868 constant ggml_metal_kargs_flash_attn_ext & args, 869 device const char * q, 870 device const char * k, 871 device const char * v, 872 device const char * mask, 873 device const char * sinks, 874 device const char * pad, 875 device const char * blk, 876 device char * dst, 877 threadgroup half * shmem_f16 [[threadgroup(0)]], 878 uint3 tgpig[[threadgroup_position_in_grid]], 879 ushort tiisg[[thread_index_in_simdgroup]], 880 ushort sgitg[[simdgroup_index_in_threadgroup]]) { 881#define FWD_TMPL q_t, q4_t, q8x8_t, k_t, k4x4_t, k8x8_t, v_t, v4x4_t, v8x8_t, qk_t, qk8x8_t, s_t, s2_t, s8x8_t, o_t, o4_t, o8x8_t, kd4x4_t, nl_k, deq_k, vd4x4_t, nl_v, deq_v, DK, DV, Q, C 882#define FWD_ARGS args, q, k, v, mask, sinks, pad, blk, dst, shmem_f16, tgpig, tiisg, sgitg 883 switch (FC_flash_attn_ext_nsg) { 884 // note: disabled cases to reduce library load time 885 //case 1: kernel_flash_attn_ext_impl<FWD_TMPL, 1>(FWD_ARGS); break; 886 //case 2: kernel_flash_attn_ext_impl<FWD_TMPL, 2>(FWD_ARGS); break; 887 case 4: kernel_flash_attn_ext_impl<FWD_TMPL, 4>(FWD_ARGS); break; 888 case 8: kernel_flash_attn_ext_impl<FWD_TMPL, 8>(FWD_ARGS); break; 889 } 890#undef FWD_TMPL 891#undef FWD_ARGS 892} 893 894// TODO: this is quite ugly. in the future these types will be hardcoded in the kernel, but for now keep them as 895// template to be able to explore different combinations 896// 897#define FA_TYPES \ 898 half, half4, simdgroup_half8x8, \ 899 half, half4x4, simdgroup_half8x8, \ 900 half, half4x4, simdgroup_half8x8, \ 901 float, simdgroup_float8x8, \ 902 float, float2, simdgroup_float8x8, \ 903 float, float4, simdgroup_float8x8 904 //half, half4, simdgroup_half8x8 905 906#define FA_TYPES_BF \ 907 bfloat, bfloat4, simdgroup_bfloat8x8, \ 908 bfloat, bfloat4x4, simdgroup_bfloat8x8, \ 909 bfloat, bfloat4x4, simdgroup_bfloat8x8, \ 910 float, simdgroup_float8x8, \ 911 float, float2, simdgroup_float8x8, \ 912 half, half4, simdgroup_half8x8 913 //float, float4, simdgroup_float8x8 914 915#define FA_TYPES_F32 \ 916 half, half4, simdgroup_half8x8, \ 917 float, float4x4, simdgroup_float8x8, \ 918 float, float4x4, simdgroup_float8x8, \ 919 float, simdgroup_float8x8, \ 920 float, float2, simdgroup_float8x8, \ 921 float, float4, simdgroup_float8x8 922 //half, half4, simdgroup_half8x8 923 924typedef decltype(kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 64, 64>) flash_attn_ext_t; 925 926template [[host_name("kernel_flash_attn_ext_f32_dk32_dv32" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 32, 32>; 927template [[host_name("kernel_flash_attn_ext_f32_dk40_dv40" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 40, 40>; 928template [[host_name("kernel_flash_attn_ext_f32_dk48_dv48" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 48, 48>; 929template [[host_name("kernel_flash_attn_ext_f32_dk64_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 64, 64>; 930template [[host_name("kernel_flash_attn_ext_f32_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 72, 72>; 931template [[host_name("kernel_flash_attn_ext_f32_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 80, 80>; 932template [[host_name("kernel_flash_attn_ext_f32_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 96, 96>; 933template [[host_name("kernel_flash_attn_ext_f32_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 112, 112>; 934template [[host_name("kernel_flash_attn_ext_f32_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 128, 128>; 935template [[host_name("kernel_flash_attn_ext_f32_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 192, 192>; 936template [[host_name("kernel_flash_attn_ext_f32_dk192_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 192, 128>; 937template [[host_name("kernel_flash_attn_ext_f32_dk256_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 256, 256>; 938template [[host_name("kernel_flash_attn_ext_f32_dk320_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 320, 256>; 939template [[host_name("kernel_flash_attn_ext_f32_dk512_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 512, 512>; 940template [[host_name("kernel_flash_attn_ext_f32_dk576_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 576, 512>; 941 942template [[host_name("kernel_flash_attn_ext_f16_dk32_dv32" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 32, 32>; 943template [[host_name("kernel_flash_attn_ext_f16_dk40_dv40" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 40, 40>; 944template [[host_name("kernel_flash_attn_ext_f16_dk48_dv48" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 48, 48>; 945template [[host_name("kernel_flash_attn_ext_f16_dk64_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 64, 64>; 946template [[host_name("kernel_flash_attn_ext_f16_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 72, 72>; 947template [[host_name("kernel_flash_attn_ext_f16_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 80, 80>; 948template [[host_name("kernel_flash_attn_ext_f16_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 96, 96>; 949template [[host_name("kernel_flash_attn_ext_f16_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 112, 112>; 950template [[host_name("kernel_flash_attn_ext_f16_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 128, 128>; 951template [[host_name("kernel_flash_attn_ext_f16_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 192, 192>; 952template [[host_name("kernel_flash_attn_ext_f16_dk192_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 192, 128>; 953template [[host_name("kernel_flash_attn_ext_f16_dk256_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 256, 256>; 954template [[host_name("kernel_flash_attn_ext_f16_dk320_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 320, 256>; 955template [[host_name("kernel_flash_attn_ext_f16_dk512_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 512, 512>; 956template [[host_name("kernel_flash_attn_ext_f16_dk576_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 576, 512>; 957 958#if defined(GGML_METAL_HAS_BF16) 959template [[host_name("kernel_flash_attn_ext_bf16_dk32_dv32" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 32, 32>; 960template [[host_name("kernel_flash_attn_ext_bf16_dk40_dv40" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 40, 40>; 961template [[host_name("kernel_flash_attn_ext_bf16_dk48_dv48" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 48, 48>; 962template [[host_name("kernel_flash_attn_ext_bf16_dk64_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 64, 64>; 963template [[host_name("kernel_flash_attn_ext_bf16_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 72, 72>; 964template [[host_name("kernel_flash_attn_ext_bf16_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 80, 80>; 965template [[host_name("kernel_flash_attn_ext_bf16_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 96, 96>; 966template [[host_name("kernel_flash_attn_ext_bf16_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 112, 112>; 967template [[host_name("kernel_flash_attn_ext_bf16_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 128, 128>; 968template [[host_name("kernel_flash_attn_ext_bf16_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 192, 192>; 969template [[host_name("kernel_flash_attn_ext_bf16_dk192_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 192, 128>; 970template [[host_name("kernel_flash_attn_ext_bf16_dk256_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 256, 256>; 971template [[host_name("kernel_flash_attn_ext_bf16_dk320_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 320, 256>; 972template [[host_name("kernel_flash_attn_ext_bf16_dk512_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 512, 512>; 973template [[host_name("kernel_flash_attn_ext_bf16_dk576_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 576, 512>; 974#endif 975 976template [[host_name("kernel_flash_attn_ext_q4_0_dk32_dv32" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 32, 32>; 977template [[host_name("kernel_flash_attn_ext_q4_0_dk40_dv40" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 40, 40>; 978template [[host_name("kernel_flash_attn_ext_q4_0_dk48_dv48" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 48, 48>; 979template [[host_name("kernel_flash_attn_ext_q4_0_dk64_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 64, 64>; 980template [[host_name("kernel_flash_attn_ext_q4_0_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 72, 72>; 981template [[host_name("kernel_flash_attn_ext_q4_0_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 80, 80>; 982template [[host_name("kernel_flash_attn_ext_q4_0_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 96, 96>; 983template [[host_name("kernel_flash_attn_ext_q4_0_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 112, 112>; 984template [[host_name("kernel_flash_attn_ext_q4_0_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 128, 128>; 985template [[host_name("kernel_flash_attn_ext_q4_0_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 192, 192>; 986template [[host_name("kernel_flash_attn_ext_q4_0_dk192_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 192, 128>; 987template [[host_name("kernel_flash_attn_ext_q4_0_dk256_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 256, 256>; 988template [[host_name("kernel_flash_attn_ext_q4_0_dk320_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 320, 256>; 989template [[host_name("kernel_flash_attn_ext_q4_0_dk512_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 512, 512>; 990template [[host_name("kernel_flash_attn_ext_q4_0_dk576_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 576, 512>; 991 992template [[host_name("kernel_flash_attn_ext_q4_1_dk32_dv32" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 32, 32>; 993template [[host_name("kernel_flash_attn_ext_q4_1_dk40_dv40" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 40, 40>; 994template [[host_name("kernel_flash_attn_ext_q4_1_dk48_dv48" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 48, 48>; 995template [[host_name("kernel_flash_attn_ext_q4_1_dk64_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 64, 64>; 996template [[host_name("kernel_flash_attn_ext_q4_1_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 72, 72>; 997template [[host_name("kernel_flash_attn_ext_q4_1_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 80, 80>; 998template [[host_name("kernel_flash_attn_ext_q4_1_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 96, 96>; 999template [[host_name("kernel_flash_attn_ext_q4_1_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 112, 112>; 1000template [[host_name("kernel_flash_attn_ext_q4_1_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 128, 128>; 1001template [[host_name("kernel_flash_attn_ext_q4_1_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 192, 192>; 1002template [[host_name("kernel_flash_attn_ext_q4_1_dk192_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 192, 128>; 1003template [[host_name("kernel_flash_attn_ext_q4_1_dk256_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 256, 256>; 1004template [[host_name("kernel_flash_attn_ext_q4_1_dk320_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 320, 256>; 1005template [[host_name("kernel_flash_attn_ext_q4_1_dk512_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 512, 512>; 1006template [[host_name("kernel_flash_attn_ext_q4_1_dk576_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 576, 512>; 1007 1008template [[host_name("kernel_flash_attn_ext_q5_0_dk32_dv32" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 32, 32>; 1009template [[host_name("kernel_flash_attn_ext_q5_0_dk40_dv40" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 40, 40>; 1010template [[host_name("kernel_flash_attn_ext_q5_0_dk48_dv48" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 48, 48>; 1011template [[host_name("kernel_flash_attn_ext_q5_0_dk64_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 64, 64>; 1012template [[host_name("kernel_flash_attn_ext_q5_0_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 72, 72>; 1013template [[host_name("kernel_flash_attn_ext_q5_0_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 80, 80>; 1014template [[host_name("kernel_flash_attn_ext_q5_0_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 96, 96>; 1015template [[host_name("kernel_flash_attn_ext_q5_0_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 112, 112>; 1016template [[host_name("kernel_flash_attn_ext_q5_0_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 128, 128>; 1017template [[host_name("kernel_flash_attn_ext_q5_0_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 192, 192>; 1018template [[host_name("kernel_flash_attn_ext_q5_0_dk192_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 192, 128>; 1019template [[host_name("kernel_flash_attn_ext_q5_0_dk256_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 256, 256>; 1020template [[host_name("kernel_flash_attn_ext_q5_0_dk320_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 320, 256>; 1021template [[host_name("kernel_flash_attn_ext_q5_0_dk512_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 512, 512>; 1022template [[host_name("kernel_flash_attn_ext_q5_0_dk576_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 576, 512>; 1023 1024template [[host_name("kernel_flash_attn_ext_q5_1_dk32_dv32" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 32, 32>; 1025template [[host_name("kernel_flash_attn_ext_q5_1_dk40_dv40" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 40, 40>; 1026template [[host_name("kernel_flash_attn_ext_q5_1_dk48_dv48" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 48, 48>; 1027template [[host_name("kernel_flash_attn_ext_q5_1_dk64_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 64, 64>; 1028template [[host_name("kernel_flash_attn_ext_q5_1_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 72, 72>; 1029template [[host_name("kernel_flash_attn_ext_q5_1_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 80, 80>; 1030template [[host_name("kernel_flash_attn_ext_q5_1_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 96, 96>; 1031template [[host_name("kernel_flash_attn_ext_q5_1_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 112, 112>; 1032template [[host_name("kernel_flash_attn_ext_q5_1_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 128, 128>; 1033template [[host_name("kernel_flash_attn_ext_q5_1_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 192, 192>; 1034template [[host_name("kernel_flash_attn_ext_q5_1_dk192_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 192, 128>; 1035template [[host_name("kernel_flash_attn_ext_q5_1_dk256_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 256, 256>; 1036template [[host_name("kernel_flash_attn_ext_q5_1_dk320_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 320, 256>; 1037template [[host_name("kernel_flash_attn_ext_q5_1_dk512_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 512, 512>; 1038template [[host_name("kernel_flash_attn_ext_q5_1_dk576_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 576, 512>; 1039 1040template [[host_name("kernel_flash_attn_ext_q8_0_dk32_dv32" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 32, 32>; 1041template [[host_name("kernel_flash_attn_ext_q8_0_dk40_dv40" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 40, 40>; 1042template [[host_name("kernel_flash_attn_ext_q8_0_dk48_dv48" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 48, 48>; 1043template [[host_name("kernel_flash_attn_ext_q8_0_dk64_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 64, 64>; 1044template [[host_name("kernel_flash_attn_ext_q8_0_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 72, 72>; 1045template [[host_name("kernel_flash_attn_ext_q8_0_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 80, 80>; 1046template [[host_name("kernel_flash_attn_ext_q8_0_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 96, 96>; 1047template [[host_name("kernel_flash_attn_ext_q8_0_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 112, 112>; 1048template [[host_name("kernel_flash_attn_ext_q8_0_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 128, 128>; 1049template [[host_name("kernel_flash_attn_ext_q8_0_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 192, 192>; 1050template [[host_name("kernel_flash_attn_ext_q8_0_dk192_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 192, 128>; 1051template [[host_name("kernel_flash_attn_ext_q8_0_dk256_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 256, 256>; 1052template [[host_name("kernel_flash_attn_ext_q8_0_dk320_dv256")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 320, 256>; 1053template [[host_name("kernel_flash_attn_ext_q8_0_dk512_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 512, 512>; 1054template [[host_name("kernel_flash_attn_ext_q8_0_dk576_dv512")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 576, 512>; 1055 1056#undef FA_TYPES 1057#undef FA_TYPES_BF 1058#undef FA_TYPES_F32 1059 1060constant bool FC_flash_attn_ext_vec_has_mask [[function_constant(FC_FLASH_ATTN_EXT_VEC + 0)]]; 1061constant bool FC_flash_attn_ext_vec_has_sinks [[function_constant(FC_FLASH_ATTN_EXT_VEC + 1)]]; 1062constant bool FC_flash_attn_ext_vec_has_bias [[function_constant(FC_FLASH_ATTN_EXT_VEC + 2)]]; 1063constant bool FC_flash_attn_ext_vec_has_scap [[function_constant(FC_FLASH_ATTN_EXT_VEC + 3)]]; 1064constant bool FC_flash_attn_ext_vec_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT_VEC + 4)]]; 1065 1066//constant float FC_flash_attn_ext_vec_scale [[function_constant(FC_FLASH_ATTN_EXT_VEC + 10)]]; 1067//constant float FC_flash_attn_ext_vec_max_bias [[function_constant(FC_FLASH_ATTN_EXT_VEC + 11)]]; 1068//constant float FC_flash_attn_ext_vec_logit_softcap [[function_constant(FC_FLASH_ATTN_EXT_VEC + 12)]]; 1069 1070constant int32_t FC_flash_attn_ext_vec_ns10 [[function_constant(FC_FLASH_ATTN_EXT_VEC + 20)]]; 1071constant int32_t FC_flash_attn_ext_vec_ns20 [[function_constant(FC_FLASH_ATTN_EXT_VEC + 21)]]; 1072constant int32_t FC_flash_attn_ext_vec_nsg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 22)]]; 1073constant int32_t FC_flash_attn_ext_vec_nwg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 23)]]; 1074constant bool FC_flash_attn_ext_vec_has_sparse [[function_constant(FC_FLASH_ATTN_EXT_VEC + 5)]]; 1075 1076// compress the finite entries of each KQ mask row into a list of KV indices (ascending order), 1077// padded with -1 up to n_kv_max_padded (a multiple of OP_FLASH_ATTN_EXT_VEC_NCPSG) 1078// one threadgroup per mask row; the mask remains the single source of truth for the values 1079kernel void kernel_flash_attn_ext_vec_idx( 1080 constant ggml_metal_kargs_flash_attn_ext_vec_idx & args, 1081 device const half * mask, 1082 device int * idx, 1083 uint3 tgpig[[threadgroup_position_in_grid]], 1084 ushort tiitg[[thread_index_in_threadgroup]], 1085 ushort3 ntg[[threads_per_threadgroup]]) { 1086 constexpr short NW = N_SIMDWIDTH; 1087 constexpr short NLOCAL = 32; // max finite positions kept in registers per thread 1088 1089 const int i1 = tgpig[0]; 1090 const int i2 = tgpig[1]; 1091 const int i3 = tgpig[2]; 1092 1093 device const half * pm = (device const half *) ((device const char *) mask + i1*args.nb31 + i2*args.nb32 + i3*args.nb33); 1094 device int * pidx = idx + (((int64_t)i3*args.ne32 + i2)*args.ne31 + i1)*args.n_kv_max_padded; 1095 1096 const int n = args.ne30; 1097 const int q = n/ntg.x; 1098 const int r = n%ntg.x; 1099 1100 // each thread handles a contiguous slice of the mask row 1101 const int r0 = q*tiitg + min((int) tiitg, r); 1102 const int r1 = r0 + q + (tiitg < r ? 1 : 0); 1103 1104 // count the finite entries in the slice and keep their positions in registers (single mask read) 1105 int cnt = 0; // total finite entries in the slice 1106 int nloc = 0; // finite entries kept in registers 1107 int local[NLOCAL]; 1108 for (int i = r0; i < r1; ++i) { 1109 if (isfinite((float) pm[i])) { 1110 if (nloc < NLOCAL) { 1111 local[nloc] = i; 1112 nloc++; 1113 } 1114 cnt++; 1115 } 1116 } 1117 1118 const short sgitg = tiitg/NW; 1119 const short tiisg = tiitg%NW; 1120 1121 threadgroup int tcount[8]; 1122 1123 // simd_sum is a collective: all lanes must evaluate it 1124 const int sg_sum = simd_sum(cnt); 1125 if (tiisg == 0) { 1126 tcount[sgitg] = sg_sum; 1127 } 1128 1129 threadgroup_barrier(mem_flags::mem_threadgroup); 1130 1131 int total = 0; 1132 for (short s = 0; s < ntg.x/NW; ++s) { 1133 total += tcount[s]; 1134 } 1135 1136 // base offset of this thread's slice in the output list (exclusive scan within the simdgroup) 1137 int sg_base = 0; 1138 for (short s = 0; s < sgitg; ++s) { 1139 sg_base += tcount[s]; 1140 } 1141 1142 // exclusive prefix scan of the per-thread counts within the simdgroup 1143 int incl = cnt; 1144 for (int d = 1; d < NW; d <<= 1) { 1145 const int v = simd_shuffle_up(incl, d); 1146 if (tiisg >= d) { 1147 incl += v; 1148 } 1149 } 1150 const int base = sg_base + (incl - cnt); 1151 1152 // write the finite positions in order; if the hint is violated, keep only the first n_kv_max entries 1153 int j = 0; 1154 for (; j < nloc && base + j < args.n_kv_max; ++j) { 1155 pidx[base + j] = local[j]; 1156 } 1157 1158 // a dense mask may have more than NLOCAL finite entries in a slice; re-read the mask to write the rest 1159 if (cnt > nloc && base + nloc < args.n_kv_max) { 1160 int j2 = 0; 1161 for (int i = r0; i < r1; ++i) { 1162 if (isfinite((float) pm[i])) { 1163 if (j2 >= nloc) { 1164 pidx[base + j2] = i; 1165 } 1166 j2++; 1167 if (base + j2 >= args.n_kv_max) { 1168 break; 1169 } 1170 } 1171 } 1172 } 1173 1174 // pad the tail of the list with -1 1175 const int count = min(total, args.n_kv_max); 1176 for (int i = count + tiitg; i < args.n_kv_max_padded; i += ntg.x) { 1177 pidx[i] = -1; 1178 } 1179} 1180 1181template< 1182 typename q4_t, // query types in shared memory 1183 typename k4_t, // key types in shared memory 1184 typename v4_t, // value types in shared memory 1185 typename qk_t, // Q*K types 1186 typename s_t, // soft-max types 1187 typename s4_t, 1188 typename o4_t, // attention accumulation types 1189 typename kd4_t, // key type in device memory 1190 short nl_k, 1191 void (*deq_k_t4)(device const kd4_t *, short, thread k4_t &), 1192 typename vd4_t, // value type in device memory 1193 short nl_v, 1194 void (*deq_v_t4)(device const vd4_t *, short, thread v4_t &), 1195 short DK, // K head size 1196 short DV, // V head size 1197 short NE = 4, // head elements per thread 1198 short Q = OP_FLASH_ATTN_EXT_VEC_NQPSG, // queries per threadgroup 1199 short C = OP_FLASH_ATTN_EXT_VEC_NCPSG> // cache items per threadgroup 1200 1201kernel void kernel_flash_attn_ext_vec( 1202 constant ggml_metal_kargs_flash_attn_ext_vec & args, 1203 device const char * q, 1204 device const char * k, 1205 device const char * v, 1206 device const char * mask, 1207 device const char * sinks, 1208 device const char * pad, 1209 device char * dst, 1210 device const char * idx, 1211 threadgroup half * shmem_f16 [[threadgroup(0)]], 1212 uint3 tgpig[[threadgroup_position_in_grid]], 1213 ushort tiisg[[thread_index_in_simdgroup]], 1214 ushort sgitg[[simdgroup_index_in_threadgroup]]) { 1215 static_assert(DK % 32 == 0, "DK must be divisible by 32"); 1216 static_assert(DV % 32 == 0, "DV must be divisible by 32"); 1217 1218#define NWG (FC_flash_attn_ext_vec_nwg) 1219#define NSG (FC_flash_attn_ext_vec_nsg) 1220 1221#define NS10 (FC_flash_attn_ext_vec_ns10) 1222#define NS20 (FC_flash_attn_ext_vec_ns20) 1223 1224 const short iwg = tgpig[2]%NWG; 1225 1226 const ushort iq3 = tgpig[2]/NWG; 1227 const ushort iq2 = tgpig[1]; 1228 const ushort iq1 = tgpig[0]; 1229 1230 constexpr short DK4 = DK/4; 1231 constexpr short DV4 = DV/4; 1232 1233 constexpr short PK = PAD2(DK, 128); 1234 constexpr short PK4 = PK/4; 1235 1236 constexpr short PV = PAD2(DV, 128); 1237 constexpr short PV4 = PV/4; 1238 1239 constexpr short NW = N_SIMDWIDTH; 1240 constexpr short NL = NW/NE; // note: this can be adjusted to support different head sizes and simdgroup work loads 1241 constexpr short SH = 4*Q*C; // shared memory per simdgroup 1242 1243 static_assert(DK4 % NL == 0, "DK4 must be divisible by NL"); 1244 static_assert(DV4 % NL == 0, "DV4 must be divisible by NL"); 1245 1246 //const short T = PK + NSG*SH; // shared memory size per query in (half) 1247 1248 //threadgroup q_t * sq = (threadgroup q_t *) (shmem_f16 + 0*PK); // holds the query data 1249 threadgroup q4_t * sq4 = (threadgroup q4_t *) (shmem_f16 + 0*PK); // same as above but in q4_t 1250 threadgroup s_t * ss = (threadgroup s_t *) (shmem_f16 + sgitg*SH + Q*NSG*PK); // scratch buffer for attention 1251 threadgroup s4_t * ss4 = (threadgroup s4_t *) (shmem_f16 + sgitg*SH + Q*NSG*PK); // same as above but in s4_t 1252 threadgroup half * sm = (threadgroup half *) (shmem_f16 + sgitg*SH + 2*Q*C + Q*NSG*PK); // scratch buffer for mask 1253 threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + 2*sgitg*Q*PV + Q*NSG*PK + NSG*SH); // scratch buffer for the results 1254 1255 // store the result for all queries in shared memory (the O matrix from the paper) 1256 so4 += tiisg; 1257 1258 { 1259 q += iq1*Q*args.nb01 + iq2*args.nb02 + iq3*args.nb03; 1260 1261 const short ikv2 = iq2/(args.ne02/args.ne_12_2); 1262 const short ikv3 = iq3/(args.ne03/args.ne_12_3); 1263 1264 k += ikv2*args.nb12 + ikv3*args.nb13; 1265 v += ikv2*args.nb22 + ikv3*args.nb23; 1266 } 1267 1268 // load Q query rows to shared memory 1269 { 1270 for (short qq = 0; qq < Q; ++qq) { 1271 const int iq1_q = iq1*Q + qq; 1272 device const float4 * q4 = (device const float4 *) ((device const char *) q + qq*args.nb01); 1273 if (iq1_q < args.ne01) { 1274 for (short i = tiisg; i < PK4; i += NW) { 1275 if (i < DK4) { 1276 sq4[qq*PK4 + i] = (q4_t) q4[i]; 1277 } else { 1278 sq4[qq*PK4 + i] = (q4_t) 0.0f; 1279 } 1280 } 1281 } else { 1282 for (short i = tiisg; i < PK4; i += NW) { 1283 sq4[qq*PK4 + i] = (q4_t) 0.0f; 1284 } 1285 } 1286 } 1287 } 1288 1289 // zero out so 1290 for (short qq = 0; qq < Q; ++qq) { 1291 for (short i = 0; i < DV4/NL; ++i) { 1292 so4[qq*DV4 + i*NL] = (o4_t) 0.0f; 1293 } 1294 } 1295 1296 // zero out shared memory SH 1297 for (short i = tiisg; i < SH/4; i += NW) { 1298 ss4[i] = (s4_t) 0.0f; 1299 } 1300 1301 threadgroup_barrier(mem_flags::mem_threadgroup); 1302 1303 { 1304 float S[Q]; 1305 float M[Q]; 1306 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1307 S[qq] = 0.0f; 1308 M[qq] = -FLT_MAX/2; 1309 } 1310 1311 // thread indices inside the simdgroup 1312 const short tx = tiisg%NL; 1313 const short ty = tiisg/NL; 1314 1315 // pointer to the mask 1316 device const half * pm_base = (device const half *) (mask + iq1*Q*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); 1317 1318 // sparse indices: the list of finite mask entries per query row 1319 // the sparse path requires Q == 1 (enforced by the host) 1320 device const int * pidx = nullptr; 1321 if (FC_flash_attn_ext_vec_has_sparse) { 1322 pidx = (device const int *) idx + 1323 ((int64_t)(iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*args.ne31*args.n_kv_max_padded + (iq1%args.ne31)*args.n_kv_max_padded; 1324 } 1325 1326 float slope = 1.0f; 1327 1328 // ALiBi 1329 if (FC_flash_attn_ext_vec_has_bias) { 1330 const short h = iq2; 1331 1332 const float base = h < args.n_head_log2 ? args.m0 : args.m1; 1333 const short exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1; 1334 1335 slope = pow(base, exph); 1336 } 1337 1338 // loop over the KV cache 1339 // each simdgroup handles blocks of Q rows and C columns 1340 for (int ic0 = iwg*NSG + sgitg; ; ic0 += NWG*NSG) { 1341 int ic = ic0*C; 1342 if (ic >= args.ne11) { 1343 break; 1344 } 1345 1346 device const half * pm[Q]; 1347 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1348 // padded query rows clamp to row 0 of the mask to avoid OOB; their scores 1349 // are forced to -inf below, so the values never affect the result. 1350 pm[qq] = pm_base + ((iq1*Q + qq) < args.ne01 ? qq*(args.nb31/sizeof(half)) : -iq1*Q*(args.nb31/sizeof(half))); 1351 } 1352 1353 // the last partial chunk uses the pad buffer as source 1354 if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) { 1355 k = pad; 1356 v = k + args.nb11*C*args.ne_12_2*args.ne_12_3; 1357 mask = v + args.nb21*C*args.ne_12_2*args.ne_12_3; 1358 1359 const short ikv2 = iq2/(args.ne02/args.ne_12_2); 1360 const short ikv3 = iq3/(args.ne03/args.ne_12_3); 1361 1362 k += (ikv2 + ikv3*args.ne_12_2)*args.nb11*C; 1363 v += (ikv2 + ikv3*args.ne_12_2)*args.nb21*C; 1364 1365 if (!FC_flash_attn_ext_vec_has_mask) { 1366 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1367 if (ic + tiisg >= args.ne11) { 1368 sm[qq*C + tiisg] = -MAXHALF; 1369 } 1370 } 1371 } else { 1372 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1373 pm[qq] = (device const half *) (mask) + 1374 (iq1*Q + qq)*C + 1375 (iq2%args.ne32)*(C*args.ne31) + 1376 (iq3%args.ne33)*(C*args.ne31*args.ne32); 1377 } 1378 } 1379 1380 ic = 0; 1381 } 1382 1383 if (FC_flash_attn_ext_vec_has_mask) { 1384 if (FC_flash_attn_ext_vec_has_sparse) { 1385 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1386 const int i11 = pidx[ic + tiisg]; 1387 if ((iq1*Q + qq) < args.ne01 && i11 >= 0) { 1388 sm[qq*C + tiisg] = pm[qq][i11]; 1389 } else { 1390 sm[qq*C + tiisg] = -MAXHALF; 1391 } 1392 } 1393 } else { 1394 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1395 if ((iq1*Q + qq) < args.ne01) { 1396 sm[qq*C + tiisg] = pm[qq][ic + tiisg]; 1397 } else { 1398 sm[qq*C + tiisg] = -MAXHALF; 1399 } 1400 } 1401 } 1402 } else { 1403 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1404 if ((iq1*Q + qq) >= args.ne01) { 1405 sm[qq*C + tiisg] = -MAXHALF; 1406 } 1407 } 1408 } 1409 1410 // skip -INF mask 1411 { 1412 bool any_finite = false; 1413 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1414 if (simd_max(sm[qq*C + tiisg]) > -MAXHALF) { 1415 any_finite = true; 1416 } 1417 } 1418 if (!any_finite) { 1419 continue; 1420 } 1421 } 1422 1423 // Q*K^T 1424 { 1425 device const k4_t * pk4 = nullptr; 1426 1427 if (!FC_flash_attn_ext_vec_has_sparse) { 1428 pk4 = (device const k4_t *) (k + ic*args.nb11); 1429 1430 pk4 += ty*NS10/4 + tx; 1431 } 1432 1433 qk_t mqk[Q][C/NE]; 1434 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1435 FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { 1436 mqk[qq][cc] = 0.0f; 1437 } 1438 } 1439 1440 // each simdgroup processes Q queries and NE (NW/NL) cache elements 1441 FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { 1442 if (FC_flash_attn_ext_vec_has_sparse) { 1443 // the KV rows are gathered from the index list; -1 entries are padding 1444 const int i11 = pidx[ic + NE*cc + ty]; 1445 if (i11 >= 0) { 1446 if (is_same<kd4_t, k4_t>::value) { 1447 device const k4_t * pk4s = (device const k4_t *) (k + i11*args.nb11) + tx; 1448 FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { 1449 const k4_t k_elem = pk4s[ii*NL]; 1450 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1451 mqk[qq][cc] += dot((float4) k_elem, (float4) sq4[qq*PK4 + ii*NL + tx]); 1452 } 1453 } 1454 } else { 1455 device const kd4_t * pk = (device const kd4_t *) (k + i11*args.nb11); 1456 1457 k4_t mk; 1458 1459 FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { 1460 const short i = ii*NL + tx; 1461 1462 deq_k_t4(pk + i/nl_k, i%nl_k, mk); 1463 1464 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1465 mqk[qq][cc] += dot((float4) mk, (float4) sq4[qq*PK4 + i]); 1466 } 1467 } 1468 } 1469 } 1470 } else if (is_same<kd4_t, k4_t>::value) { 1471 FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { 1472 const k4_t k_elem = pk4[cc*NE*NS10/4 + ii*NL]; 1473 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1474 mqk[qq][cc] += dot((float4) k_elem, (float4) sq4[qq*PK4 + ii*NL + tx]); 1475 } 1476 } 1477 } else { 1478 device const kd4_t * pk = (device const kd4_t *) (k + ((ic + NE*cc + ty)*args.nb11)); 1479 1480 k4_t mk; 1481 1482 FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { 1483 const short i = ii*NL + tx; 1484 1485 deq_k_t4(pk + i/nl_k, i%nl_k, mk); 1486 1487 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1488 mqk[qq][cc] += dot((float4) mk, (float4) sq4[qq*PK4 + i]); 1489 } 1490 } 1491 } 1492 1493 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1494 if (NE == 1) { 1495 mqk[qq][cc] = simd_sum(mqk[qq][cc]); 1496 } else { 1497 // simdgroup reduce (NE = 4) 1498 // [ 0 .. 7] -> [ 0] 1499 // [ 8 .. 15] -> [ 8] 1500 // [16 .. 23] -> [16] 1501 // [24 .. 31] -> [24] 1502 if (NE <= 1) { 1503 mqk[qq][cc] += simd_shuffle_down(mqk[qq][cc], 16); 1504 } 1505 if (NE <= 2) { 1506 mqk[qq][cc] += simd_shuffle_down(mqk[qq][cc], 8); 1507 } 1508 if (NE <= 4) { 1509 mqk[qq][cc] += simd_shuffle_down(mqk[qq][cc], 4); 1510 } 1511 if (NE <= 8) { 1512 mqk[qq][cc] += simd_shuffle_down(mqk[qq][cc], 2); 1513 } 1514 if (NE <= 16) { 1515 mqk[qq][cc] += simd_shuffle_down(mqk[qq][cc], 1); 1516 } 1517 1518 // broadcast 1519 mqk[qq][cc] = simd_shuffle(mqk[qq][cc], NL*ty); 1520 } 1521 } 1522 } 1523 1524 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1525 if (FC_flash_attn_ext_vec_has_mask && 1526 !FC_flash_attn_ext_vec_has_scap && 1527 !FC_flash_attn_ext_vec_has_bias) { 1528 ss[qq*C + NE*tx + ty] = fma(mqk[qq][tx], args.scale, (qk_t) sm[qq*C + NE*tx + ty]); 1529 } else { 1530 mqk[qq][tx] *= args.scale; 1531 1532 if (FC_flash_attn_ext_vec_has_scap) { 1533 mqk[qq][tx] = args.logit_softcap*precise::tanh(mqk[qq][tx]); 1534 } 1535 1536 if (FC_flash_attn_ext_vec_has_bias) { 1537 mqk[qq][tx] += (qk_t) sm[qq*C + NE*tx + ty]*slope; 1538 } else { 1539 mqk[qq][tx] += (qk_t) sm[qq*C + NE*tx + ty]; 1540 } 1541 1542 ss[qq*C + NE*tx + ty] = mqk[qq][tx]; 1543 } 1544 } 1545 } 1546 1547 simdgroup_barrier(mem_flags::mem_threadgroup); 1548 1549 // online softmax 1550 { 1551 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1552 const float m = M[qq]; 1553 const float s = ss[qq*C + tiisg]; 1554 1555 M[qq] = simd_max(max(M[qq], s)); 1556 1557 const float ms = exp(m - M[qq]); 1558 const float vs = exp(s - M[qq]); 1559 1560 S[qq] = S[qq]*ms + simd_sum(vs); 1561 1562 // the P matrix from the paper (Q rows, C columns) 1563 ss[qq*C + tiisg] = vs; 1564 1565 // O = diag(ms)*O 1566 if ((DV4/NL % NW == 0) || ty == 0) { 1567 FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { 1568 so4[qq*DV4 + ii*NL] *= ms; 1569 } 1570 } 1571 } 1572 } 1573 1574 simdgroup_barrier(mem_flags::mem_threadgroup); 1575 1576 // O = O + (Q*K^T)*V 1577 { 1578 o4_t lo[Q][DV4/NL]; 1579 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1580 FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { 1581 lo[qq][ii] = 0.0f; 1582 } 1583 } 1584 1585 if (FC_flash_attn_ext_vec_has_sparse) { 1586 FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { 1587 // the KV rows are gathered from the index list; -1 entries are padding 1588 const int i11 = pidx[ic + NE*cc + ty]; 1589 if (i11 >= 0) { 1590 if (is_same<vd4_t, v4_t>::value) { 1591 device const v4_t * pv4 = (device const v4_t *) (v + i11*args.nb21); 1592 1593 pv4 += tx; 1594 1595 FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { 1596 const v4_t v_elem = pv4[ii*NL]; 1597 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1598 lo[qq][ii] += o4_t(float4(v_elem)*float4(ss[qq*C + cc*NE + ty])); 1599 } 1600 } 1601 } else { 1602 device const vd4_t * pv4 = (device const vd4_t *) (v + i11*args.nb21); 1603 1604 FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { 1605 const short i = ii*NL + tx; 1606 1607 v4_t mv; 1608 1609 deq_v_t4(pv4 + i/nl_v, i%nl_v, mv); 1610 1611 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1612 lo[qq][ii] += o4_t(float4(mv)*float4(ss[qq*C + cc*NE + ty])); 1613 } 1614 } 1615 } 1616 } 1617 } 1618 } else if (is_same<vd4_t, v4_t>::value) { 1619 device const v4_t * pv4 = (device const v4_t *) (v + ic*args.nb21); 1620 1621 pv4 += ty*NS20/4 + tx; 1622 1623 FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { 1624 FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { 1625 const v4_t v_elem = pv4[cc*NE*NS20/4 + ii*NL]; 1626 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1627 lo[qq][ii] += o4_t(float4(v_elem)*float4(ss[qq*C + cc*NE + ty])); 1628 } 1629 } 1630 } 1631 } else { 1632 FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { 1633 device const vd4_t * pv4 = (device const vd4_t *) (v + ((ic + NE*cc + ty)*args.nb21)); 1634 1635 FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { 1636 const short i = ii*NL + tx; 1637 1638 v4_t mv; 1639 deq_v_t4(pv4 + i/nl_v, i%nl_v, mv); 1640 1641 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1642 lo[qq][ii] += o4_t(float4(mv)*float4(ss[qq*C + NE*cc + ty])); 1643 } 1644 } 1645 } 1646 } 1647 1648 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1649 FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { 1650 if (NE > 1) { 1651 lo[qq][ii][0] += simd_shuffle_down(lo[qq][ii][0], 16); 1652 lo[qq][ii][1] += simd_shuffle_down(lo[qq][ii][1], 16); 1653 lo[qq][ii][2] += simd_shuffle_down(lo[qq][ii][2], 16); 1654 lo[qq][ii][3] += simd_shuffle_down(lo[qq][ii][3], 16); 1655 } 1656 1657 if (NE > 2) { 1658 lo[qq][ii][0] += simd_shuffle_down(lo[qq][ii][0], 8); 1659 lo[qq][ii][1] += simd_shuffle_down(lo[qq][ii][1], 8); 1660 lo[qq][ii][2] += simd_shuffle_down(lo[qq][ii][2], 8); 1661 lo[qq][ii][3] += simd_shuffle_down(lo[qq][ii][3], 8); 1662 } 1663 1664 if (NE > 4) { 1665 lo[qq][ii][0] += simd_shuffle_down(lo[qq][ii][0], 4); 1666 lo[qq][ii][1] += simd_shuffle_down(lo[qq][ii][1], 4); 1667 lo[qq][ii][2] += simd_shuffle_down(lo[qq][ii][2], 4); 1668 lo[qq][ii][3] += simd_shuffle_down(lo[qq][ii][3], 4); 1669 } 1670 1671 if (NE > 8) { 1672 lo[qq][ii][0] += simd_shuffle_down(lo[qq][ii][0], 2); 1673 lo[qq][ii][1] += simd_shuffle_down(lo[qq][ii][1], 2); 1674 lo[qq][ii][2] += simd_shuffle_down(lo[qq][ii][2], 2); 1675 lo[qq][ii][3] += simd_shuffle_down(lo[qq][ii][3], 2); 1676 } 1677 1678 if (NE > 16) { 1679 lo[qq][ii][0] += simd_shuffle_down(lo[qq][ii][0], 1); 1680 lo[qq][ii][1] += simd_shuffle_down(lo[qq][ii][1], 1); 1681 lo[qq][ii][2] += simd_shuffle_down(lo[qq][ii][2], 1); 1682 lo[qq][ii][3] += simd_shuffle_down(lo[qq][ii][3], 1); 1683 } 1684 } 1685 } 1686 1687 if ((DV4/NL % NW == 0) || ty == 0) { 1688 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1689 FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { 1690 so4[qq*DV4 + ii*NL] += lo[qq][ii]; 1691 } 1692 } 1693 } 1694 } 1695 } 1696 1697 if (FC_flash_attn_ext_vec_has_sinks && sgitg == 0 && iwg == 0) { 1698 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1699 const float m = M[qq]; 1700 const float s = tiisg == 0 ? ((device const float *) sinks)[iq2] : -FLT_MAX/2; 1701 1702 M[qq] = simd_max(max(M[qq], s)); 1703 1704 const float ms = exp(m - M[qq]); 1705 const float vs = exp(s - M[qq]); 1706 1707 S[qq] = S[qq]*ms + simd_sum(vs); 1708 1709 if ((DV4/NL % NW == 0) || ty == 0) { 1710 FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { 1711 so4[qq*DV4 + ii*NL] *= ms; 1712 } 1713 } 1714 } 1715 } 1716 1717 // these are needed for reducing the results from the simdgroups (reuse the ss buffer) 1718 if (tiisg == 0) { 1719 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1720 ss[2*qq + 0] = (s_t) S[qq]; 1721 ss[2*qq + 1] = (s_t) M[qq]; 1722 } 1723 } 1724 } 1725 1726 so4 -= tiisg; 1727 1728 threadgroup_barrier(mem_flags::mem_threadgroup); 1729 1730 // parallel reduce 1731 for (short r = NSG/2; r > 0; r >>= 1) { 1732 if (sgitg < r) { 1733 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1734 const float S0 = ss[ 2*qq + 0]; 1735 const float S1 = ss[r*(SH/2) + 2*qq + 0]; 1736 1737 const float M0 = ss[ 2*qq + 1]; 1738 const float M1 = ss[r*(SH/2) + 2*qq + 1]; 1739 1740 const float Mx = max(M0, M1); 1741 1742 const float ms0 = exp(M0 - Mx); 1743 const float ms1 = exp(M1 - Mx); 1744 1745 const float Sx = S0*ms0 + S1*ms1; 1746 1747 if (tiisg == 0) { 1748 ss[2*qq + 0] = Sx; 1749 ss[2*qq + 1] = Mx; 1750 } 1751 1752 // O_0 = diag(ms0)*O_0 + diag(ms1)*O_1 1753 for (short i = tiisg; i < DV4; i += NW) { 1754 so4[qq*DV4 + i] = so4[qq*DV4 + i]*ms0 + so4[qq*DV4 + i + r*Q*PV4]*ms1; 1755 } 1756 } 1757 } 1758 1759 threadgroup_barrier(mem_flags::mem_threadgroup); 1760 } 1761 1762 // final rescale with 1/S and store to global memory 1763 if (sgitg == 0) { 1764 const int64_t nrows = args.ne3*args.ne2*args.ne1; 1765 1766 device float4 * dst4 = (device float4 *) dst; 1767 device float * dst1 = (device float *) dst + nrows*DV*NWG; // the S and M are stored after the results 1768 1769 FOR_UNROLL (short qq = 0; qq < Q; ++qq) { 1770 const int iq1_q = iq1*Q + qq; 1771 if (iq1_q >= args.ne01) { 1772 continue; 1773 } 1774 1775 const int64_t rid = iq3*args.ne2*args.ne1 + iq2 + iq1_q*args.ne1; 1776 1777 const float Sval = NWG == 1 ? (ss[2*qq + 0] == 0.0f ? 0.0f : 1.0f/ss[2*qq + 0]) : 1.0f; 1778 1779 // interleave the workgroup data 1780 for (short i = tiisg; i < DV4; i += NW) { 1781 dst4[rid*DV4*NWG + NWG*i + iwg] = (float4) so4[qq*DV4 + i]*Sval; 1782 } 1783 1784 // store S and M 1785 if (NWG > 1) { 1786 if (tiisg == 0) { 1787 dst1[rid*(2*NWG) + 2*iwg + 0] = ss[2*qq + 0]; 1788 dst1[rid*(2*NWG) + 2*iwg + 1] = ss[2*qq + 1]; 1789 } 1790 } 1791 } 1792 } 1793 1794#undef NWG 1795#undef NSG 1796#undef NS10 1797#undef NS20 1798} 1799 1800// note: I think the s_t can be half instead of float, because the Q*K scaling is done before storing to shared mem 1801// in the other (non-vec) kernel, we need s_t to also be float because we scale during the soft_max 1802// 1803#define FA_TYPES \ 1804 half4, \ 1805 half4, \ 1806 half4, \ 1807 float, \ 1808 float, float4, \ 1809 float4 1810 1811#define FA_TYPES_F32 \ 1812 half4, \ 1813 float4, \ 1814 float4, \ 1815 float, \ 1816 float, float4, \ 1817 float4 1818 1819typedef decltype(kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 4>) flash_attn_ext_vec_t; 1820 1821template [[host_name("kernel_flash_attn_ext_vec_f32_dk32_dv32")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES_F32, float4, 1, dequantize_f32_t4, float4, 1, dequantize_f32_t4, 32, 32, 4>; 1822template [[host_name("kernel_flash_attn_ext_vec_f16_dk32_dv32")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 32, 32, 4>; 1823template [[host_name("kernel_flash_attn_ext_vec_f16_dk32_dv32_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 32, 32, 4, 2>; 1824template [[host_name("kernel_flash_attn_ext_vec_f16_dk32_dv32_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 32, 32, 4, 4>; 1825#if defined(GGML_METAL_HAS_BF16) 1826template [[host_name("kernel_flash_attn_ext_vec_bf16_dk32_dv32")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, bfloat4, 1, dequantize_bf16_t4, bfloat4, 1, dequantize_bf16_t4, 32, 32, 4>; 1827#endif 1828template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk32_dv32")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 32, 32, 4>; 1829template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk32_dv32_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 32, 32, 4, 2>; 1830template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk32_dv32_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 32, 32, 4, 4>; 1831template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk32_dv32")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 32, 32, 4>; 1832template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk32_dv32_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 32, 32, 4, 2>; 1833template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk32_dv32_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 32, 32, 4, 4>; 1834template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk32_dv32")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 32, 32, 4>; 1835template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk32_dv32_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 32, 32, 4, 2>; 1836template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk32_dv32_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 32, 32, 4, 4>; 1837template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk32_dv32")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 32, 32, 4>; 1838template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk32_dv32_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 32, 32, 4, 2>; 1839template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk32_dv32_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 32, 32, 4, 4>; 1840template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk32_dv32")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 32, 32, 4>; 1841template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk32_dv32_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 32, 32, 4, 2>; 1842template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk32_dv32_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 32, 32, 4, 4>; 1843 1844template [[host_name("kernel_flash_attn_ext_vec_f32_dk64_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES_F32, float4, 1, dequantize_f32_t4, float4, 1, dequantize_f32_t4, 64, 64, 2>; 1845template [[host_name("kernel_flash_attn_ext_vec_f16_dk64_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 64, 64, 2>; 1846template [[host_name("kernel_flash_attn_ext_vec_f16_dk64_dv64_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 64, 64, 4, 1>; 1847template [[host_name("kernel_flash_attn_ext_vec_f16_dk64_dv64_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 64, 64, 2, 2>; 1848template [[host_name("kernel_flash_attn_ext_vec_f16_dk64_dv64_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 64, 64, 4, 2>; 1849template [[host_name("kernel_flash_attn_ext_vec_f16_dk64_dv64_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 64, 64, 2, 4>; 1850template [[host_name("kernel_flash_attn_ext_vec_f16_dk64_dv64_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 64, 64, 4, 4>; 1851#if defined(GGML_METAL_HAS_BF16) 1852template [[host_name("kernel_flash_attn_ext_vec_bf16_dk64_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, bfloat4, 1, dequantize_bf16_t4, bfloat4, 1, dequantize_bf16_t4, 64, 64, 2>; 1853#endif 1854template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk64_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 64, 64, 2>; 1855template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk64_dv64_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 64, 64, 4, 1>; 1856template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk64_dv64_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 64, 64, 2, 2>; 1857template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk64_dv64_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 64, 64, 4, 2>; 1858template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk64_dv64_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 64, 64, 2, 4>; 1859template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk64_dv64_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 64, 64, 4, 4>; 1860template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk64_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 64, 64, 2>; 1861template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk64_dv64_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 64, 64, 4, 1>; 1862template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk64_dv64_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 64, 64, 2, 2>; 1863template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk64_dv64_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 64, 64, 4, 2>; 1864template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk64_dv64_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 64, 64, 2, 4>; 1865template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk64_dv64_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 64, 64, 4, 4>; 1866template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk64_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 64, 64, 2>; 1867template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk64_dv64_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 64, 64, 4, 1>; 1868template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk64_dv64_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 64, 64, 2, 2>; 1869template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk64_dv64_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 64, 64, 4, 2>; 1870template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk64_dv64_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 64, 64, 2, 4>; 1871template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk64_dv64_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 64, 64, 4, 4>; 1872template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk64_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 64, 64, 2>; 1873template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk64_dv64_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 64, 64, 4, 1>; 1874template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk64_dv64_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 64, 64, 2, 2>; 1875template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk64_dv64_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 64, 64, 4, 2>; 1876template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk64_dv64_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 64, 64, 2, 4>; 1877template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk64_dv64_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 64, 64, 4, 4>; 1878template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk64_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 64, 64, 2>; 1879template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk64_dv64_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 64, 64, 4, 1>; 1880template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk64_dv64_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 64, 64, 2, 2>; 1881template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk64_dv64_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 64, 64, 4, 2>; 1882template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk64_dv64_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 64, 64, 2, 4>; 1883template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk64_dv64_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 64, 64, 4, 4>; 1884 1885template [[host_name("kernel_flash_attn_ext_vec_f32_dk96_dv96")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES_F32, float4, 1, dequantize_f32_t4, float4, 1, dequantize_f32_t4, 96, 96, 4>; 1886template [[host_name("kernel_flash_attn_ext_vec_f16_dk96_dv96")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 96, 96, 4>; 1887template [[host_name("kernel_flash_attn_ext_vec_f16_dk96_dv96_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 96, 96, 4, 2>; 1888template [[host_name("kernel_flash_attn_ext_vec_f16_dk96_dv96_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 96, 96, 4, 4>; 1889#if defined(GGML_METAL_HAS_BF16) 1890template [[host_name("kernel_flash_attn_ext_vec_bf16_dk96_dv96")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, bfloat4, 1, dequantize_bf16_t4, bfloat4, 1, dequantize_bf16_t4, 96, 96, 4>; 1891#endif 1892template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk96_dv96")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 96, 96, 4>; 1893template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk96_dv96_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 96, 96, 4, 2>; 1894template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk96_dv96_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 96, 96, 4, 4>; 1895template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk96_dv96")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 96, 96, 4>; 1896template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk96_dv96_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 96, 96, 4, 2>; 1897template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk96_dv96_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 96, 96, 4, 4>; 1898template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk96_dv96")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 96, 96, 4>; 1899template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk96_dv96_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 96, 96, 4, 2>; 1900template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk96_dv96_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 96, 96, 4, 4>; 1901template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk96_dv96")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 96, 96, 4>; 1902template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk96_dv96_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 96, 96, 4, 2>; 1903template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk96_dv96_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 96, 96, 4, 4>; 1904template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk96_dv96")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 96, 96, 4>; 1905template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk96_dv96_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 96, 96, 4, 2>; 1906template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk96_dv96_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 96, 96, 4, 4>; 1907 1908template [[host_name("kernel_flash_attn_ext_vec_f32_dk128_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES_F32, float4, 1, dequantize_f32_t4, float4, 1, dequantize_f32_t4, 128, 128, 1>; 1909template [[host_name("kernel_flash_attn_ext_vec_f16_dk128_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 1>; 1910template [[host_name("kernel_flash_attn_ext_vec_f16_dk128_dv128_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 2, 1>; 1911template [[host_name("kernel_flash_attn_ext_vec_f16_dk128_dv128_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 4, 1>; 1912template [[host_name("kernel_flash_attn_ext_vec_f16_dk128_dv128_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 1, 2>; 1913template [[host_name("kernel_flash_attn_ext_vec_f16_dk128_dv128_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 2, 2>; 1914template [[host_name("kernel_flash_attn_ext_vec_f16_dk128_dv128_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 4, 2>; 1915template [[host_name("kernel_flash_attn_ext_vec_f16_dk128_dv128_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 1, 4>; 1916template [[host_name("kernel_flash_attn_ext_vec_f16_dk128_dv128_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 2, 4>; 1917template [[host_name("kernel_flash_attn_ext_vec_f16_dk128_dv128_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 4, 4>; 1918#if defined(GGML_METAL_HAS_BF16) 1919template [[host_name("kernel_flash_attn_ext_vec_bf16_dk128_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, bfloat4, 1, dequantize_bf16_t4, bfloat4, 1, dequantize_bf16_t4, 128, 128, 1>; 1920#endif 1921template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk128_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 128, 128, 1>; 1922template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk128_dv128_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 128, 128, 2, 1>; 1923template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk128_dv128_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 128, 128, 4, 1>; 1924template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk128_dv128_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 128, 128, 1, 2>; 1925template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk128_dv128_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 128, 128, 2, 2>; 1926template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk128_dv128_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 128, 128, 4, 2>; 1927template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk128_dv128_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 128, 128, 1, 4>; 1928template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk128_dv128_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 128, 128, 2, 4>; 1929template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk128_dv128_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 128, 128, 4, 4>; 1930template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk128_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 128, 128, 1>; 1931template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk128_dv128_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 128, 128, 2, 1>; 1932template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk128_dv128_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 128, 128, 4, 1>; 1933template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk128_dv128_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 128, 128, 1, 2>; 1934template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk128_dv128_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 128, 128, 2, 2>; 1935template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk128_dv128_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 128, 128, 4, 2>; 1936template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk128_dv128_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 128, 128, 1, 4>; 1937template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk128_dv128_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 128, 128, 2, 4>; 1938template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk128_dv128_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 128, 128, 4, 4>; 1939template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk128_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 128, 128, 1>; 1940template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk128_dv128_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 128, 128, 2, 1>; 1941template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk128_dv128_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 128, 128, 4, 1>; 1942template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk128_dv128_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 128, 128, 1, 2>; 1943template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk128_dv128_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 128, 128, 2, 2>; 1944template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk128_dv128_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 128, 128, 4, 2>; 1945template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk128_dv128_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 128, 128, 1, 4>; 1946template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk128_dv128_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 128, 128, 2, 4>; 1947template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk128_dv128_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 128, 128, 4, 4>; 1948template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk128_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 128, 128, 1>; 1949template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk128_dv128_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 128, 128, 2, 1>; 1950template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk128_dv128_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 128, 128, 4, 1>; 1951template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk128_dv128_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 128, 128, 1, 2>; 1952template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk128_dv128_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 128, 128, 2, 2>; 1953template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk128_dv128_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 128, 128, 4, 2>; 1954template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk128_dv128_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 128, 128, 1, 4>; 1955template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk128_dv128_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 128, 128, 2, 4>; 1956template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk128_dv128_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 128, 128, 4, 4>; 1957template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk128_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 128, 128, 1>; 1958template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk128_dv128_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 128, 128, 2, 1>; 1959template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk128_dv128_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 128, 128, 4, 1>; 1960template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk128_dv128_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 128, 128, 1, 2>; 1961template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk128_dv128_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 128, 128, 2, 2>; 1962template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk128_dv128_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 128, 128, 4, 2>; 1963template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk128_dv128_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 128, 128, 1, 4>; 1964template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk128_dv128_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 128, 128, 2, 4>; 1965template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk128_dv128_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 128, 128, 4, 4>; 1966 1967template [[host_name("kernel_flash_attn_ext_vec_f32_dk192_dv192")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES_F32, float4, 1, dequantize_f32_t4, float4, 1, dequantize_f32_t4, 192, 192, 2>; 1968template [[host_name("kernel_flash_attn_ext_vec_f16_dk192_dv192")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 192, 192, 2>; 1969template [[host_name("kernel_flash_attn_ext_vec_f16_dk192_dv192_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 192, 192, 4, 1>; 1970template [[host_name("kernel_flash_attn_ext_vec_f16_dk192_dv192_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 192, 192, 2, 2>; 1971template [[host_name("kernel_flash_attn_ext_vec_f16_dk192_dv192_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 192, 192, 4, 2>; 1972template [[host_name("kernel_flash_attn_ext_vec_f16_dk192_dv192_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 192, 192, 2, 4>; 1973template [[host_name("kernel_flash_attn_ext_vec_f16_dk192_dv192_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 192, 192, 4, 4>; 1974#if defined(GGML_METAL_HAS_BF16) 1975template [[host_name("kernel_flash_attn_ext_vec_bf16_dk192_dv192")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, bfloat4, 1, dequantize_bf16_t4, bfloat4, 1, dequantize_bf16_t4, 192, 192, 2>; 1976#endif 1977template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk192_dv192")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 192, 192, 2>; 1978template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk192_dv192_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 192, 192, 4, 1>; 1979template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk192_dv192_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 192, 192, 2, 2>; 1980template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk192_dv192_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 192, 192, 4, 2>; 1981template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk192_dv192_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 192, 192, 2, 4>; 1982template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk192_dv192_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 192, 192, 4, 4>; 1983template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk192_dv192")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 192, 192, 2>; 1984template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk192_dv192_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 192, 192, 4, 1>; 1985template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk192_dv192_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 192, 192, 2, 2>; 1986template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk192_dv192_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 192, 192, 4, 2>; 1987template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk192_dv192_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 192, 192, 2, 4>; 1988template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk192_dv192_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 192, 192, 4, 4>; 1989template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk192_dv192")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 192, 192, 2>; 1990template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk192_dv192_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 192, 192, 4, 1>; 1991template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk192_dv192_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 192, 192, 2, 2>; 1992template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk192_dv192_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 192, 192, 4, 2>; 1993template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk192_dv192_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 192, 192, 2, 4>; 1994template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk192_dv192_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 192, 192, 4, 4>; 1995template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk192_dv192")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 192, 192, 2>; 1996template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk192_dv192_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 192, 192, 4, 1>; 1997template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk192_dv192_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 192, 192, 2, 2>; 1998template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk192_dv192_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 192, 192, 4, 2>; 1999template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk192_dv192_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 192, 192, 2, 4>; 2000template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk192_dv192_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 192, 192, 4, 4>; 2001template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk192_dv192")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 192, 192, 2>; 2002template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk192_dv192_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 192, 192, 4, 1>; 2003template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk192_dv192_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 192, 192, 2, 2>; 2004template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk192_dv192_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 192, 192, 4, 2>; 2005template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk192_dv192_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 192, 192, 2, 4>; 2006template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk192_dv192_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 192, 192, 4, 4>; 2007 2008template [[host_name("kernel_flash_attn_ext_vec_f32_dk192_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES_F32, float4, 1, dequantize_f32_t4, float4, 1, dequantize_f32_t4, 192, 128, 2>; 2009template [[host_name("kernel_flash_attn_ext_vec_f16_dk192_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 192, 128, 2>; 2010template [[host_name("kernel_flash_attn_ext_vec_f16_dk192_dv128_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 192, 128, 4, 1>; 2011template [[host_name("kernel_flash_attn_ext_vec_f16_dk192_dv128_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 192, 128, 2, 2>; 2012template [[host_name("kernel_flash_attn_ext_vec_f16_dk192_dv128_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 192, 128, 4, 2>; 2013template [[host_name("kernel_flash_attn_ext_vec_f16_dk192_dv128_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 192, 128, 2, 4>; 2014template [[host_name("kernel_flash_attn_ext_vec_f16_dk192_dv128_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 192, 128, 4, 4>; 2015#if defined(GGML_METAL_HAS_BF16) 2016template [[host_name("kernel_flash_attn_ext_vec_bf16_dk192_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, bfloat4, 1, dequantize_bf16_t4, bfloat4, 1, dequantize_bf16_t4, 192, 128, 2>; 2017#endif 2018template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk192_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 192, 128, 2>; 2019template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk192_dv128_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 192, 128, 4, 1>; 2020template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk192_dv128_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 192, 128, 2, 2>; 2021template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk192_dv128_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 192, 128, 4, 2>; 2022template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk192_dv128_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 192, 128, 2, 4>; 2023template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk192_dv128_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 192, 128, 4, 4>; 2024template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk192_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 192, 128, 2>; 2025template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk192_dv128_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 192, 128, 4, 1>; 2026template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk192_dv128_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 192, 128, 2, 2>; 2027template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk192_dv128_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 192, 128, 4, 2>; 2028template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk192_dv128_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 192, 128, 2, 4>; 2029template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk192_dv128_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 192, 128, 4, 4>; 2030template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk192_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 192, 128, 2>; 2031template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk192_dv128_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 192, 128, 4, 1>; 2032template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk192_dv128_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 192, 128, 2, 2>; 2033template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk192_dv128_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 192, 128, 4, 2>; 2034template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk192_dv128_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 192, 128, 2, 4>; 2035template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk192_dv128_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 192, 128, 4, 4>; 2036template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk192_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 192, 128, 2>; 2037template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk192_dv128_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 192, 128, 4, 1>; 2038template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk192_dv128_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 192, 128, 2, 2>; 2039template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk192_dv128_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 192, 128, 4, 2>; 2040template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk192_dv128_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 192, 128, 2, 4>; 2041template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk192_dv128_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 192, 128, 4, 4>; 2042template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk192_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 192, 128, 2>; 2043template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk192_dv128_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 192, 128, 4, 1>; 2044template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk192_dv128_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 192, 128, 2, 2>; 2045template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk192_dv128_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 192, 128, 4, 2>; 2046template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk192_dv128_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 192, 128, 2, 4>; 2047template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk192_dv128_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 192, 128, 4, 4>; 2048 2049template [[host_name("kernel_flash_attn_ext_vec_f32_dk256_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES_F32, float4, 1, dequantize_f32_t4, float4, 1, dequantize_f32_t4, 256, 256, 1>; 2050template [[host_name("kernel_flash_attn_ext_vec_f16_dk256_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 256, 256, 1>; 2051template [[host_name("kernel_flash_attn_ext_vec_f16_dk256_dv256_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 256, 256, 2, 1>; 2052template [[host_name("kernel_flash_attn_ext_vec_f16_dk256_dv256_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 256, 256, 4, 1>; 2053template [[host_name("kernel_flash_attn_ext_vec_f16_dk256_dv256_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 256, 256, 1, 2>; 2054template [[host_name("kernel_flash_attn_ext_vec_f16_dk256_dv256_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 256, 256, 2, 2>; 2055template [[host_name("kernel_flash_attn_ext_vec_f16_dk256_dv256_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 256, 256, 4, 2>; 2056template [[host_name("kernel_flash_attn_ext_vec_f16_dk256_dv256_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 256, 256, 1, 4>; 2057template [[host_name("kernel_flash_attn_ext_vec_f16_dk256_dv256_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 256, 256, 2, 4>; 2058template [[host_name("kernel_flash_attn_ext_vec_f16_dk256_dv256_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 256, 256, 4, 4>; 2059#if defined(GGML_METAL_HAS_BF16) 2060template [[host_name("kernel_flash_attn_ext_vec_bf16_dk256_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, bfloat4, 1, dequantize_bf16_t4, bfloat4, 1, dequantize_bf16_t4, 256, 256, 1>; 2061#endif 2062template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk256_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 256, 256, 1>; 2063template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk256_dv256_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 256, 256, 2, 1>; 2064template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk256_dv256_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 256, 256, 4, 1>; 2065template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk256_dv256_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 256, 256, 1, 2>; 2066template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk256_dv256_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 256, 256, 2, 2>; 2067template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk256_dv256_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 256, 256, 4, 2>; 2068template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk256_dv256_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 256, 256, 1, 4>; 2069template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk256_dv256_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 256, 256, 2, 4>; 2070template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk256_dv256_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 256, 256, 4, 4>; 2071template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk256_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 256, 256, 1>; 2072template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk256_dv256_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 256, 256, 2, 1>; 2073template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk256_dv256_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 256, 256, 4, 1>; 2074template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk256_dv256_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 256, 256, 1, 2>; 2075template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk256_dv256_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 256, 256, 2, 2>; 2076template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk256_dv256_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 256, 256, 4, 2>; 2077template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk256_dv256_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 256, 256, 1, 4>; 2078template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk256_dv256_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 256, 256, 2, 4>; 2079template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk256_dv256_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 256, 256, 4, 4>; 2080template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk256_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 256, 256, 1>; 2081template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk256_dv256_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 256, 256, 2, 1>; 2082template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk256_dv256_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 256, 256, 4, 1>; 2083template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk256_dv256_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 256, 256, 1, 2>; 2084template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk256_dv256_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 256, 256, 2, 2>; 2085template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk256_dv256_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 256, 256, 4, 2>; 2086template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk256_dv256_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 256, 256, 1, 4>; 2087template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk256_dv256_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 256, 256, 2, 4>; 2088template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk256_dv256_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 256, 256, 4, 4>; 2089template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk256_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 256, 256, 1>; 2090template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk256_dv256_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 256, 256, 2, 1>; 2091template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk256_dv256_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 256, 256, 4, 1>; 2092template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk256_dv256_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 256, 256, 1, 2>; 2093template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk256_dv256_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 256, 256, 2, 2>; 2094template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk256_dv256_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 256, 256, 4, 2>; 2095template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk256_dv256_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 256, 256, 1, 4>; 2096template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk256_dv256_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 256, 256, 2, 4>; 2097template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk256_dv256_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 256, 256, 4, 4>; 2098template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk256_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 256, 256, 1>; 2099template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk256_dv256_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 256, 256, 2, 1>; 2100template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk256_dv256_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 256, 256, 4, 1>; 2101template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk256_dv256_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 256, 256, 1, 2>; 2102template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk256_dv256_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 256, 256, 2, 2>; 2103template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk256_dv256_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 256, 256, 4, 2>; 2104template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk256_dv256_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 256, 256, 1, 4>; 2105template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk256_dv256_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 256, 256, 2, 4>; 2106template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk256_dv256_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 256, 256, 4, 4>; 2107 2108template [[host_name("kernel_flash_attn_ext_vec_f32_dk320_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES_F32, float4, 1, dequantize_f32_t4, float4, 1, dequantize_f32_t4, 320, 256, 2>; 2109template [[host_name("kernel_flash_attn_ext_vec_f16_dk320_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 320, 256, 2>; 2110template [[host_name("kernel_flash_attn_ext_vec_f16_dk320_dv256_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 320, 256, 4, 1>; 2111template [[host_name("kernel_flash_attn_ext_vec_f16_dk320_dv256_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 320, 256, 2, 2>; 2112template [[host_name("kernel_flash_attn_ext_vec_f16_dk320_dv256_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 320, 256, 4, 2>; 2113template [[host_name("kernel_flash_attn_ext_vec_f16_dk320_dv256_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 320, 256, 2, 4>; 2114template [[host_name("kernel_flash_attn_ext_vec_f16_dk320_dv256_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 320, 256, 4, 4>; 2115#if defined(GGML_METAL_HAS_BF16) 2116template [[host_name("kernel_flash_attn_ext_vec_bf16_dk320_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, bfloat4, 1, dequantize_bf16_t4, bfloat4, 1, dequantize_bf16_t4, 320, 256, 2>; 2117#endif 2118template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk320_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 320, 256, 2>; 2119template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk320_dv256_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 320, 256, 4, 1>; 2120template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk320_dv256_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 320, 256, 2, 2>; 2121template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk320_dv256_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 320, 256, 4, 2>; 2122template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk320_dv256_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 320, 256, 2, 4>; 2123template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk320_dv256_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 320, 256, 4, 4>; 2124template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk320_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 320, 256, 2>; 2125template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk320_dv256_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 320, 256, 4, 1>; 2126template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk320_dv256_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 320, 256, 2, 2>; 2127template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk320_dv256_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 320, 256, 4, 2>; 2128template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk320_dv256_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 320, 256, 2, 4>; 2129template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk320_dv256_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 320, 256, 4, 4>; 2130template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk320_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 320, 256, 2>; 2131template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk320_dv256_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 320, 256, 4, 1>; 2132template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk320_dv256_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 320, 256, 2, 2>; 2133template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk320_dv256_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 320, 256, 4, 2>; 2134template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk320_dv256_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 320, 256, 2, 4>; 2135template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk320_dv256_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 320, 256, 4, 4>; 2136template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk320_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 320, 256, 2>; 2137template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk320_dv256_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 320, 256, 4, 1>; 2138template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk320_dv256_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 320, 256, 2, 2>; 2139template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk320_dv256_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 320, 256, 4, 2>; 2140template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk320_dv256_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 320, 256, 2, 4>; 2141template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk320_dv256_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 320, 256, 4, 4>; 2142template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk320_dv256")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 320, 256, 2>; 2143template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk320_dv256_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 320, 256, 4, 1>; 2144template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk320_dv256_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 320, 256, 2, 2>; 2145template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk320_dv256_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 320, 256, 4, 2>; 2146template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk320_dv256_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 320, 256, 2, 4>; 2147template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk320_dv256_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 320, 256, 4, 4>; 2148 2149template [[host_name("kernel_flash_attn_ext_vec_f32_dk512_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES_F32, float4, 1, dequantize_f32_t4, float4, 1, dequantize_f32_t4, 512, 512, 1>; 2150template [[host_name("kernel_flash_attn_ext_vec_f16_dk512_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 512, 512, 1>; 2151template [[host_name("kernel_flash_attn_ext_vec_f16_dk512_dv512_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 512, 512, 2, 1>; 2152template [[host_name("kernel_flash_attn_ext_vec_f16_dk512_dv512_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 512, 512, 4, 1>; 2153template [[host_name("kernel_flash_attn_ext_vec_f16_dk512_dv512_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 512, 512, 1, 2>; 2154template [[host_name("kernel_flash_attn_ext_vec_f16_dk512_dv512_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 512, 512, 2, 2>; 2155template [[host_name("kernel_flash_attn_ext_vec_f16_dk512_dv512_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 512, 512, 4, 2>; 2156template [[host_name("kernel_flash_attn_ext_vec_f16_dk512_dv512_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 512, 512, 1, 4>; 2157template [[host_name("kernel_flash_attn_ext_vec_f16_dk512_dv512_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 512, 512, 2, 4>; 2158template [[host_name("kernel_flash_attn_ext_vec_f16_dk512_dv512_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 512, 512, 4, 4>; 2159#if defined(GGML_METAL_HAS_BF16) 2160template [[host_name("kernel_flash_attn_ext_vec_bf16_dk512_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, bfloat4, 1, dequantize_bf16_t4, bfloat4, 1, dequantize_bf16_t4, 512, 512, 1>; 2161#endif 2162template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk512_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 512, 512, 1>; 2163template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk512_dv512_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 512, 512, 2, 1>; 2164template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk512_dv512_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 512, 512, 4, 1>; 2165template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk512_dv512_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 512, 512, 1, 2>; 2166template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk512_dv512_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 512, 512, 2, 2>; 2167template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk512_dv512_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 512, 512, 4, 2>; 2168template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk512_dv512_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 512, 512, 1, 4>; 2169template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk512_dv512_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 512, 512, 2, 4>; 2170template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk512_dv512_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 512, 512, 4, 4>; 2171template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk512_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 512, 512, 1>; 2172template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk512_dv512_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 512, 512, 2, 1>; 2173template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk512_dv512_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 512, 512, 4, 1>; 2174template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk512_dv512_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 512, 512, 1, 2>; 2175template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk512_dv512_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 512, 512, 2, 2>; 2176template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk512_dv512_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 512, 512, 4, 2>; 2177template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk512_dv512_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 512, 512, 1, 4>; 2178template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk512_dv512_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 512, 512, 2, 4>; 2179template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk512_dv512_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 512, 512, 4, 4>; 2180template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk512_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 512, 512, 1>; 2181template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk512_dv512_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 512, 512, 2, 1>; 2182template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk512_dv512_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 512, 512, 4, 1>; 2183template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk512_dv512_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 512, 512, 1, 2>; 2184template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk512_dv512_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 512, 512, 2, 2>; 2185template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk512_dv512_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 512, 512, 4, 2>; 2186template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk512_dv512_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 512, 512, 1, 4>; 2187template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk512_dv512_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 512, 512, 2, 4>; 2188template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk512_dv512_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 512, 512, 4, 4>; 2189template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk512_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 512, 512, 1>; 2190template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk512_dv512_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 512, 512, 2, 1>; 2191template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk512_dv512_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 512, 512, 4, 1>; 2192template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk512_dv512_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 512, 512, 1, 2>; 2193template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk512_dv512_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 512, 512, 2, 2>; 2194template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk512_dv512_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 512, 512, 4, 2>; 2195template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk512_dv512_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 512, 512, 1, 4>; 2196template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk512_dv512_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 512, 512, 2, 4>; 2197template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk512_dv512_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 512, 512, 4, 4>; 2198template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk512_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 512, 512, 1>; 2199template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk512_dv512_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 512, 512, 2, 1>; 2200template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk512_dv512_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 512, 512, 4, 1>; 2201template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk512_dv512_q2_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 512, 512, 1, 2>; 2202template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk512_dv512_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 512, 512, 2, 2>; 2203template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk512_dv512_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 512, 512, 4, 2>; 2204template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk512_dv512_q4_ne1")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 512, 512, 1, 4>; 2205template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk512_dv512_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 512, 512, 2, 4>; 2206template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk512_dv512_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 512, 512, 4, 4>; 2207 2208template [[host_name("kernel_flash_attn_ext_vec_f32_dk576_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES_F32, float4, 1, dequantize_f32_t4, float4, 1, dequantize_f32_t4, 576, 512, 2>; 2209template [[host_name("kernel_flash_attn_ext_vec_f16_dk576_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 576, 512, 2>; 2210template [[host_name("kernel_flash_attn_ext_vec_f16_dk576_dv512_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 576, 512, 4, 1>; 2211template [[host_name("kernel_flash_attn_ext_vec_f16_dk576_dv512_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 576, 512, 2, 2>; 2212template [[host_name("kernel_flash_attn_ext_vec_f16_dk576_dv512_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 576, 512, 4, 2>; 2213template [[host_name("kernel_flash_attn_ext_vec_f16_dk576_dv512_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 576, 512, 2, 4>; 2214template [[host_name("kernel_flash_attn_ext_vec_f16_dk576_dv512_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 576, 512, 4, 4>; 2215#if defined(GGML_METAL_HAS_BF16) 2216template [[host_name("kernel_flash_attn_ext_vec_bf16_dk576_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, bfloat4, 1, dequantize_bf16_t4, bfloat4, 1, dequantize_bf16_t4, 576, 512, 2>; 2217#endif 2218template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk576_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 576, 512, 2>; 2219template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk576_dv512_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 576, 512, 4, 1>; 2220template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk576_dv512_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 576, 512, 2, 2>; 2221template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk576_dv512_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 576, 512, 4, 2>; 2222template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk576_dv512_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 576, 512, 2, 4>; 2223template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk576_dv512_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 576, 512, 4, 4>; 2224template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk576_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 576, 512, 2>; 2225template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk576_dv512_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 576, 512, 4, 1>; 2226template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk576_dv512_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 576, 512, 2, 2>; 2227template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk576_dv512_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 576, 512, 4, 2>; 2228template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk576_dv512_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 576, 512, 2, 4>; 2229template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk576_dv512_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 576, 512, 4, 4>; 2230template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk576_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 576, 512, 2>; 2231template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk576_dv512_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 576, 512, 4, 1>; 2232template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk576_dv512_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 576, 512, 2, 2>; 2233template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk576_dv512_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 576, 512, 4, 2>; 2234template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk576_dv512_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 576, 512, 2, 4>; 2235template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk576_dv512_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 576, 512, 4, 4>; 2236template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk576_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 576, 512, 2>; 2237template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk576_dv512_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 576, 512, 4, 1>; 2238template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk576_dv512_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 576, 512, 2, 2>; 2239template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk576_dv512_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 576, 512, 4, 2>; 2240template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk576_dv512_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 576, 512, 2, 4>; 2241template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk576_dv512_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 576, 512, 4, 4>; 2242template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk576_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 576, 512, 2>; 2243template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk576_dv512_q1_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 576, 512, 4, 1>; 2244template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk576_dv512_q2_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 576, 512, 2, 2>; 2245template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk576_dv512_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 576, 512, 4, 2>; 2246template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk576_dv512_q4_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 576, 512, 2, 4>; 2247template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk576_dv512_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 576, 512, 4, 4>; 2248 2249 2250#undef FA_TYPES 2251#undef FA_TYPES_F32 2252 2253constant int32_t FC_flash_attn_ext_vec_reduce_DV [[function_constant(FC_FLASH_ATTN_EXT_VEC_REDUCE + 0)]]; 2254constant int32_t FC_flash_attn_ext_vec_reduce_NWG [[function_constant(FC_FLASH_ATTN_EXT_VEC_REDUCE + 1)]]; 2255 2256kernel void kernel_flash_attn_ext_vec_reduce( 2257 constant ggml_metal_kargs_flash_attn_ext_vec_reduce & args, 2258 device const char * htmp, 2259 device char * dst, 2260 uint tgpig[[threadgroup_position_in_grid]], 2261 ushort tiisg[[thread_index_in_simdgroup]], 2262 ushort sgitg[[simdgroup_index_in_threadgroup]]) { 2263#define NWG (FC_flash_attn_ext_vec_reduce_NWG) 2264#define DV (FC_flash_attn_ext_vec_reduce_DV) 2265 2266 const uint64_t rid = tgpig; 2267 2268 const short iwg = tiisg; 2269 2270 device const float * ss = (device const float *) htmp + (uint64_t)args.nrows*DV*NWG; 2271 2272 float S = ss[rid*(2*NWG) + 2*iwg + 0]; 2273 float M = ss[rid*(2*NWG) + 2*iwg + 1]; 2274 2275 const float m = simd_max(M); 2276 const float ms = exp(M - m); 2277 2278 S = simd_sum(S*ms); 2279 S = S == 0.0f ? 0.0f : 1.0f/S; 2280 2281 const short DV4 = DV/4; 2282 2283 device const float4 * htmp4 = (device const float4 *) htmp + rid*DV4*NWG; 2284 device float4 * dst4 = (device float4 *) dst + rid*DV4; 2285 2286 for (short i = sgitg; i < DV4; i += NWG) { 2287 const float4 v = simd_sum(htmp4[i*NWG + iwg]*ms); 2288 2289 if (iwg == 0) { 2290 dst4[i] = v*S; 2291 } 2292 } 2293 2294#undef NWG 2295#undef DV 2296} 2297 2298template< 2299 typename kd4x4_t, 2300 short nl_k, 2301 void (*deq_k)(device const kd4x4_t *, short, thread half4x4 &)> 2302kernel void kernel_lightning_indexer( 2303 constant ggml_metal_kargs_lightning_indexer & args, 2304 device const char * q, 2305 device const char * k, 2306 device const char * w, 2307 device const char * m, 2308 device char * dst, 2309 uint3 tgpig[[threadgroup_position_in_grid]], 2310 ushort tiitg[[thread_index_in_threadgroup]], 2311 ushort tiisg[[thread_index_in_simdgroup]], 2312 ushort sgitg[[simdgroup_index_in_threadgroup]]) { 2313 constexpr short DK = OP_LIGHTNING_INDEXER_DK; 2314 constexpr short NH = OP_LIGHTNING_INDEXER_NH; 2315 constexpr short NHPTG = OP_LIGHTNING_INDEXER_NHPTG; 2316 constexpr short NKPSG = OP_LIGHTNING_INDEXER_NKPSG; 2317 constexpr short NSG = OP_LIGHTNING_INDEXER_NSG; 2318 constexpr short NBPTG = OP_LIGHTNING_INDEXER_NBPTG; 2319 2320 constexpr short DK4 = DK/4; 2321 constexpr short DK8 = DK/8; 2322 constexpr short DK16 = DK/16; 2323 2324 constexpr short NK = NKPSG*NSG; // keys per threadgroup 2325 constexpr short NTG = 32*NSG; // threads per threadgroup 2326 2327 const int i_stream = tgpig.z; 2328 const int i_kv_0 = tgpig.x*NK; // first key of this threadgroup 2329 const int i_kv = i_kv_0 + sgitg*NKPSG; // first key of this simdgroup 2330 2331 threadgroup half sk[NK * DK16 * 16]; 2332 threadgroup half4x4 * sk4x4 = (threadgroup half4x4 *) sk; 2333 2334 for (short i = tiitg; i < NK*DK16; i += NTG) { 2335 const short ik = i/DK16; 2336 const short i16 = i%DK16; 2337 2338 half4x4 tmp; 2339 2340 if (i_kv_0 + ik < args.n_kv) { 2341 device const kd4x4_t * kr = (device const kd4x4_t *) (k + (i_kv_0 + ik)*args.nbk2 + i_stream*args.nbk3); 2342 2343 deq_k(kr + i16/nl_k, i16%nl_k, tmp); 2344 } else { 2345 FOR_UNROLL (short j = 0; j < 4; ++j) { 2346 tmp[j] = half4(0.0h); 2347 } 2348 } 2349 2350 sk4x4[i] = tmp; 2351 } 2352 2353 threadgroup_barrier(mem_flags::mem_threadgroup); 2354 2355 // K tile of this simdgroup, transposed to [DK, NKPSG] 2356 simdgroup_half8x8 mk[DK8]; 2357 2358 FOR_UNROLL (short i = 0; i < DK8; ++i) { 2359 simdgroup_load(mk[i], sk + sgitg*NKPSG*DK + 8*i, DK, 0, true); 2360 } 2361 2362 threadgroup half4 sq4[NHPTG*DK4]; 2363 threadgroup half * sq = (threadgroup half *) sq4; 2364 2365 threadgroup float sw [NHPTG]; 2366 threadgroup float sqk[NSG*NHPTG*NKPSG]; 2367 2368 const int i_batch_0 = tgpig.y*NBPTG; 2369 const int n_batch = min((int) NBPTG, args.n_batch - i_batch_0); 2370 2371 for (short ib = 0; ib < n_batch; ++ib) { 2372 const int i_batch = i_batch_0 + ib; 2373 2374 device const char * pq = q + i_batch*args.nbq2 + i_stream*args.nbq3; 2375 device const char * pw = w + i_batch*args.nbw1 + i_stream*args.nbw3; 2376 2377 float score = 0.0f; 2378 2379 FOR_UNROLL (short i_head = 0; i_head < NH; i_head += NHPTG) { 2380 // stage the Q tile [DK, NHPTG] and the (prescaled) head weights 2381 for (short i = tiitg; i < NHPTG*DK4; i += NTG) { 2382 const short ih = i/DK4; 2383 const short i4 = i%DK4; 2384 2385 device const float4 * q4 = (device const float4 *) (pq + (i_head + ih)*args.nbq1); 2386 2387 sq4[ih*DK4 + i4] = half4(q4[i4]); 2388 } 2389 2390 if (tiitg < NHPTG) { 2391 sw[tiitg] = ((device const float *) pw)[i_head + tiitg]; 2392 } 2393 2394 threadgroup_barrier(mem_flags::mem_threadgroup); 2395 2396 simdgroup_float8x8 mqk = make_filled_simdgroup_matrix<float, 8>(0.0f); 2397 2398 FOR_UNROLL (short i = 0; i < DK8; ++i) { 2399 simdgroup_half8x8 mq; 2400 2401 simdgroup_load(mq, sq + 8*i, DK, 0, false); 2402 simdgroup_multiply_accumulate(mqk, mq, mk[i], mqk); 2403 } 2404 2405 threadgroup float * pqk = sqk + sgitg*NHPTG*NKPSG; 2406 2407 simdgroup_store(mqk, pqk, NKPSG, 0, false); 2408 simdgroup_barrier(mem_flags::mem_threadgroup); 2409 2410 // one lane per key: ReLU, apply the head weight and accumulate over the head tile 2411 if (tiisg < NKPSG) { 2412 FOR_UNROLL (short ih = 0; ih < NHPTG; ++ih) { 2413 score += max(pqk[ih*NKPSG + tiisg], 0.0f)*sw[ih]; 2414 } 2415 } 2416 2417 threadgroup_barrier(mem_flags::mem_threadgroup); 2418 } 2419 2420 if (tiisg < NKPSG) { 2421 const int ik = i_kv + tiisg; 2422 if (ik < args.n_kv) { 2423 device const half * pm = (device const half *) (m + i_batch*args.nbm1 + (i_stream % args.mask_ne3)*args.nbm3); 2424 device float * pd = (device float *) (dst + i_batch*args.nb1 + i_stream*args.nb3); 2425 2426 pd[ik] = score + (float) pm[ik]; 2427 } 2428 } 2429 } 2430} 2431 2432typedef decltype(kernel_lightning_indexer<half4x4, 1, dequantize_f16>) kernel_lightning_indexer_t; 2433 2434template [[host_name("kernel_lightning_indexer_f32")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<float4x4, 1, dequantize_f32>; 2435template [[host_name("kernel_lightning_indexer_f16")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<half4x4, 1, dequantize_f16>; 2436 2437#if defined(GGML_METAL_HAS_BF16) 2438template [[host_name("kernel_lightning_indexer_bf16")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<bfloat4x4, 1, dequantize_bf16>; 2439#endif 2440 2441template [[host_name("kernel_lightning_indexer_q4_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q4_0, 2, dequantize_q4_0>; 2442template [[host_name("kernel_lightning_indexer_q4_1")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q4_1, 2, dequantize_q4_1>; 2443template [[host_name("kernel_lightning_indexer_q5_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q5_0, 2, dequantize_q5_0>; 2444template [[host_name("kernel_lightning_indexer_q5_1")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q5_1, 2, dequantize_q5_1>; 2445template [[host_name("kernel_lightning_indexer_q8_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q8_0, 2, dequantize_q8_0>;