char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
91f6a6cf3
char/aliceai-t5-moe
1#version 450 2 3#extension GL_EXT_control_flow_attributes : enable 4#extension GL_EXT_shader_16bit_storage : require 5 6#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require 7 8#ifdef FLOAT16 9#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require 10#extension GL_EXT_shader_subgroup_extended_types_float16 : require 11#endif 12 13#ifdef MMQ 14#extension GL_EXT_integer_dot_product : require 15#extension GL_KHR_shader_subgroup_clustered : require 16 17#include "mul_mmq_shmem_types.glsl" 18#endif 19 20#extension GL_KHR_shader_subgroup_shuffle : enable 21#extension GL_KHR_shader_subgroup_vote : enable 22 23#include "types.glsl" 24#include "dot_product_funcs.glsl" 25#include "flash_attn_base.glsl" 26#include "flash_attn_dequant.glsl" 27 28const uint32_t HSK_per_thread = HSK / D_split; 29const uint32_t HSV_per_thread = HSV / D_split; 30 31const uint32_t rows_per_thread = Br / row_split; 32const uint32_t cols_per_iter = WorkGroupSize / D_split / row_split; 33const uint32_t cols_per_thread = Bc / cols_per_iter; 34const uint32_t num_subgroups = SubGroupSize == 0 ? 0 : WorkGroupSize / SubGroupSize; 35 36 37layout (binding = 0) readonly buffer Q {float data_q[];}; 38layout (binding = 0) readonly buffer QV4 {vec4 data_qv4[];}; 39layout (binding = 1) readonly buffer K {float16_t data_k[];}; 40layout (binding = 1) readonly buffer KV4 {f16vec4 data_kv4[];}; 41layout (binding = 2) readonly buffer V {float16_t data_v[];}; 42layout (binding = 2) readonly buffer VV4 {f16vec4 data_vv4[];}; 43layout (binding = 3) readonly buffer M {float16_t data_m[];}; 44 45// If SubGroupSize is set to 0 then only use shmem reductions 46const uint32_t tmpsh_size = (SubGroupSize > 0) ? (row_split == 1 ? num_subgroups * D_split : num_subgroups) : WorkGroupSize; 47shared float tmpsh[tmpsh_size]; 48shared FLOAT_TYPEV4 tmpshv4[tmpsh_size]; 49 50const uint32_t masksh_stride = Br + 1; 51shared FLOAT_TYPE masksh[Bc * masksh_stride]; 52 53#ifndef MMQ 54const uint32_t qf_stride = HSK / 4 + 1; 55shared FLOAT_TYPEV4 Qf[Br * qf_stride]; 56#else 57 58const uint32_t qf_stride = HSK / 32; 59shared block_b_cache Qf[Br * qf_stride]; 60#endif 61 62#ifndef MMQ 63const uint32_t D = HSK > HSV ? HSK : HSV; 64#else 65const uint32_t D = HSV; 66#endif 67const uint32_t kvsh_stride = D / 4 + 1; 68shared FLOAT_TYPEV4 kvsh[SHMEM_STAGING != 0 ? Bc * kvsh_stride : 1]; 69 70#ifdef MMQ 71 72shared block_a_cache kblocksh[SHMEM_STAGING != 0 ? Bc * qf_stride : 1]; 73#endif 74 75shared vec4 occupancy_limiter[LIMIT_OCCUPANCY_SHMEM > 0 ? LIMIT_OCCUPANCY_SHMEM : 1]; 76 77#ifdef MMQ 78#include "flash_attn_mmq_funcs.glsl" 79#endif 80 81void main() { 82#ifdef NEEDS_INIT_IQ_SHMEM 83 if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) { 84 init_iq_shmem(gl_WorkGroupSize); 85 } 86#endif 87 88 init_indices(); 89 90 const uint32_t tid = gl_LocalInvocationIndex; 91 const uint32_t threads_per_rowgroup = gl_WorkGroupSize.x / row_split; 92 const uint32_t row_tid = gl_LocalInvocationIndex / threads_per_rowgroup; 93 const uint32_t rowgroup_tid = gl_LocalInvocationIndex % threads_per_rowgroup; 94 const uint32_t d_tid = gl_LocalInvocationIndex % D_split; 95 const uint32_t col_tid = (gl_LocalInvocationIndex % threads_per_rowgroup) / D_split; 96 97 if (LIMIT_OCCUPANCY_SHMEM > 0) { 98 // This just exists to avoid the occupancy_limiter array getting optimized out 99 occupancy_limiter[tid] = vec4(tid); 100 101 barrier(); 102 103 if (occupancy_limiter[tid] == vec4(99999.0)) { 104 data_ov4[0] = D_TYPEV4(occupancy_limiter[tid]); 105 } 106 } 107 108#define tile_row(r) (row_tid * rows_per_thread + (r)) 109 110 uint32_t q_offset = gqa_iq1*p.nb01 + (iq2*p.nb02 + iq3*p.nb03) / 4; 111 112 [[unroll]] for (uint32_t idx = 0; idx < Br * HSK / 4; idx += gl_WorkGroupSize.x) { 113 uint32_t d = (idx + tid) % (HSK / 4); 114 uint32_t r = (idx + tid) / (HSK / 4); 115 const bool is_in_bounds = r < Br && d < HSK / 4 && i * Br + r < N; 116#ifndef MMQ 117 if (is_in_bounds) { 118 Qf[r * qf_stride + d] = FLOAT_TYPEV4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale); 119 } 120#else 121 const uint buf_ib = r * qf_stride + d / 8; 122 const uint buf_iqs = d % 8; 123 124 vec4 vals = is_in_bounds ? data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale : vec4(0.0f); 125 const vec4 abs_vals = abs(vals); 126 127 const float thread_max = max(max(abs_vals.x, abs_vals.y), max(abs_vals.z, abs_vals.w)); 128 const float amax = subgroupClusteredMax(thread_max, 8); 129 const float qd = amax / 127.0f; 130 const float qd_inv = qd != 0.0f ? 1.0f / qd : 0.0f; 131 vals = round(vals * qd_inv); 132 133 Qf[buf_ib].qs[buf_iqs] = pack32(i8vec4(vals)); 134 135 // Q8_0 K only needs (qd, _); the asymmetric Q4_*/Q5_* family also stores 136 // the row-sum scaled by qd, used in k_dot_correction. 137 if (FaTypeK == GGML_TYPE_Q8_0) { 138 if (buf_iqs == 0) { 139 Qf[buf_ib].ds = FLOAT_TYPEV2(qd, 0.0f); 140 } 141 } else { 142 const float thread_sum = vals.x + vals.y + vals.z + vals.w; 143 const float sum = subgroupClusteredAdd(thread_sum, 8); 144 145 if (buf_iqs == 0) { 146 Qf[buf_ib].ds = FLOAT_TYPEV2(qd, sum * qd); 147 } 148 } 149#endif 150 } 151 barrier(); 152 153 FLOAT_TYPEV4 Of[rows_per_thread][HSV_per_thread / 4]; 154 [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { 155 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 156 Of[r][d] = FLOAT_TYPEV4(0.0); 157 } 158 } 159 160 float Lf[rows_per_thread], Mf[rows_per_thread]; 161 162 // Use -FLT_MAX/2 rather than -inf to reduce the possibility of NaNs, e.g. when computing Mold-M. 163 const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF); 164 165 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 166 Lf[r] = 0; 167 Mf[r] = NEG_FLT_MAX_OVER_2; 168 } 169 170 ACC_TYPE slope[rows_per_thread]; 171 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 172 slope[r] = ACC_TYPE(1.0); 173 } 174 175 // ALiBi 176 if (p.max_bias > 0.0f) { 177 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 178 slope[r] = perElemOpComputeSlope(tile_row(r), col_tid, ACC_TYPE(0), iq2); 179 } 180 } 181 182 const uint32_t mo_stride = CEIL_DIV(KV, 16 * Bc); 183 // mo_offset will point to the tile starting at row i*Br and col 0 184 uint32_t mo_offset = mo_stride * i; 185 186 // FaBlockBytesK/V == 2 for f16, 16 for f32, ggml block byte size for quants. 187 uint32_t k_offset = (ik2*p.nb12 + ik3*p.nb13) / FaBlockBytesK; 188 uint32_t v_offset = (iv2*p.nb22 + iv3*p.nb23) / FaBlockBytesV; 189 uint32_t m_offset = gqa_iq1*KV; 190 if (p.nem2 != 1 || p.nem3 != 1) { 191 m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * KV; 192 mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride; 193 } 194 195 uint32_t mask_opt = 0; 196 uint32_t mask_opt_idx = ~0; 197 uint32_t mask_opt_bits = 0; 198 199 [[dont_unroll]] 200 for (uint32_t j = start_j; j < end_j; ++j) { 201 if (MASK_ENABLE) { 202 if (USE_MASK_OPT && mask_opt_idx != j / 16) { 203 mask_opt_idx = j / 16; 204 mask_opt = data_mask_opt[mo_offset + mask_opt_idx]; 205 } 206 mask_opt_bits = (mask_opt >> ((j % 16) * 2)) & 0x3; 207 if (mask_opt_bits == MASK_OPT_ALL_NEG_INF) { 208 // skip this block 209 continue; 210 } 211 // Only load if the block is not all zeros 212 if (mask_opt_bits != MASK_OPT_ALL_ZERO) { 213 bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0; 214 215 float max_mask = NEG_FLT_MAX_OVER_2; 216 barrier(); 217 [[unroll]] for (uint32_t idx = 0; idx < Bc * Br; idx += gl_WorkGroupSize.x) { 218 uint32_t c = (idx + tid) % Bc; 219 uint32_t r = (idx + tid) / Bc; 220 if (idx + tid < Bc * Br) { 221 if ((!KV_bounds_check || j * Bc + c < KV) && (!nem1_bounds_check || i * Br + r < p.nem1)) { 222 FLOAT_TYPE m = FLOAT_TYPE(data_m[m_offset + (i * Br + r) * m_stride + (j * Bc + c)]); 223 masksh[c * masksh_stride + r] = m; 224 max_mask = max(max_mask, float(m)); 225 } else { 226 masksh[c * masksh_stride + r] = FLOAT_TYPE(0); 227 } 228 } 229 } 230 // skip the block if the mask is entirely -inf 231 bool all_less = subgroupAll(max_mask <= NEG_FLT_MAX_OVER_2); 232 barrier(); 233 if (gl_SubgroupInvocationID == 0) { 234 tmpsh[gl_SubgroupID] = all_less ? NEG_FLT_MAX_OVER_2 : 0.0f; 235 } 236 barrier(); 237 [[unroll]] for (uint s = 0; s < gl_NumSubgroups; ++s) { 238 max_mask = max(max_mask, tmpsh[s]); 239 } 240 if (max_mask <= NEG_FLT_MAX_OVER_2) { 241 continue; 242 } 243 } 244 } 245 246 ACC_TYPE Sf[rows_per_thread][cols_per_thread]; 247 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 248 [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { 249 Sf[r][c] = ACC_TYPE(0.0); 250 } 251 } 252 253 if (SHMEM_STAGING != 0) { 254 barrier(); 255#ifndef MMQ 256 [[unroll]] for (uint32_t idx = 0; idx < Bc * HSK / 4; idx += gl_WorkGroupSize.x) { 257 uint32_t d = (idx + tid) % (HSK / 4); 258 uint32_t c = (idx + tid) / (HSK / 4); 259 if (idx + gl_WorkGroupSize.x <= Bc * HSK / 4 || c < Bc) { 260 FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0); 261 if (!KV_bounds_check || j * Bc + c < KV) { 262 if (USE_DECODE_K) { 263 uint coord = (j * Bc + c) * k_stride * BLOCK_SIZE_K + 4 * d; 264 uint ib = coord / BLOCK_SIZE_K; 265 uint iqs = (coord % BLOCK_SIZE_K); 266 K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); 267 } else { 268 K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c) * k_stride / 4 + d]); 269 } 270 } 271 272 kvsh[c * kvsh_stride + d] = K_Tf; 273 } 274 } 275#else // MMQ 276 const uint ints_per_block = 8u / fa_quant_r_mmq(FaTypeK); 277 const uint quant_iters = Bc * HSK / 32 * ints_per_block; 278 [[unroll]] for (uint32_t idx = 0; idx < quant_iters; idx += gl_WorkGroupSize.x) { 279 const uint32_t iqs = (idx + tid) % ints_per_block; 280 const uint32_t ib = (idx + tid) / ints_per_block; 281 const uint32_t c = ib / (HSK / 32); 282 const uint32_t block = ib % (HSK / 32); 283 if (idx + gl_WorkGroupSize.x <= quant_iters || c < Bc) { 284 const uint buf_ib = c * qf_stride + block; 285 if (!KV_bounds_check || j * Bc + c < KV) { 286 const uint global_ib = (j * Bc + c) * k_stride + block; 287 k_block_to_shmem(buf_ib, global_ib, iqs, k_offset); 288 } else { 289 k_block_to_shmem_zero(buf_ib, iqs); 290 } 291 } 292 } 293#endif // MMQ 294 barrier(); 295 } 296 297#ifndef MMQ 298 // More d iterations means Q register caching becomes relevant 299 // Few iterations means the additional registers needed are worse than the speed-up from caching 300 if (HSK_per_thread / 4 > 4) { 301 [[unroll]] for (uint32_t d = 0; d < HSK_per_thread / 4; ++d) { 302 FLOAT_TYPEV4 Q_cache[rows_per_thread]; 303 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 304 Q_cache[r] = Qf[tile_row(r) * qf_stride + d * D_split + d_tid]; 305 } 306 307 [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { 308 if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { 309 continue; 310 } 311 312 FLOAT_TYPEV4 K_Tf; 313 if (SHMEM_STAGING != 0) { 314 K_Tf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)]; 315 } else if (USE_DECODE_K) { 316 uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid); 317 uint ib = coord / BLOCK_SIZE_K; 318 uint iqs = (coord % BLOCK_SIZE_K); 319 K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); 320 } else { 321 K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * k_stride / 4 + d * D_split + d_tid]); 322 } 323 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 324 Sf[r][c] = dot_product(Q_cache[r], K_Tf, Sf[r][c]); 325 } 326 } 327 } 328 } else { 329 [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { 330 if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { 331 continue; 332 } 333 334 [[unroll]] for (uint32_t d = 0; d < HSK_per_thread / 4; ++d) { 335 FLOAT_TYPEV4 K_Tf; 336 if (SHMEM_STAGING != 0) { 337 K_Tf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)]; 338 } else if (USE_DECODE_K) { 339 uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid); 340 uint ib = coord / BLOCK_SIZE_K; 341 uint iqs = (coord % BLOCK_SIZE_K); 342 K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); 343 } else { 344 K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * k_stride / 4 + d * D_split + d_tid]); 345 } 346 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 347 Sf[r][c] = dot_product(Qf[tile_row(r) * qf_stride + d * D_split + d_tid], K_Tf, Sf[r][c]); 348 } 349 } 350 } 351 } 352#else // MMQ 353 const uint hsk4 = HSK_per_thread / 4; 354 const uint d_per_step = (hsk4 % 8 == 0) ? 8 : 355 (hsk4 % 4 == 0) ? 4 : 356 (hsk4 % 2 == 0) ? 2 : 1; 357 358 [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { 359 if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { 360 continue; 361 } 362 363 [[unroll]] for (uint32_t d_block = 0; d_block < HSK_per_thread / 4; d_block += d_per_step) { 364 int32_t k_quants[d_per_step]; 365 ACC_TYPEV2 k_dm; 366 367 // Q4_*/Q5_* take the block-8 fast path when one step covers a full 368 // block; Q8_0 always goes through the per-int get_k_qs* helpers 369 // (its qs is byte-packed, not nibble-packed). 370 const bool block8_fast = (d_per_step == 8) && (FaTypeK != GGML_TYPE_Q8_0); 371 372 if (SHMEM_STAGING != 0) { 373 const uint k_block_idx = (d_tid * (HSK_per_thread / 4) + d_block) / 8; 374 const uint buf_ib = (c * cols_per_iter + col_tid) * qf_stride + k_block_idx; 375 k_dm = ACC_TYPEV2(kblocksh[buf_ib].dm); 376 377 if (block8_fast) { 378 const bool has_qh = (FaTypeK == GGML_TYPE_Q5_0) || (FaTypeK == GGML_TYPE_Q5_1); 379 [[unroll]] for (uint32_t d = 0; d < 4; d++) { 380 uint vui = kblocksh[buf_ib].qs[d]; 381 k_quants[d ] = int32_t( vui & 0x0F0F0F0F); 382 k_quants[d + 4] = int32_t((vui >> 4) & 0x0F0F0F0F); 383 if (has_qh) { 384 uint qh_lo = (kblocksh[buf_ib].qh >> (d * 4)) & 0xF; 385 uint qh_hi = (kblocksh[buf_ib].qh >> (d * 4 + 16)) & 0xF; 386 k_quants[d ] |= int32_t((qh_lo * 0x02040810u) & 0x10101010u); 387 k_quants[d + 4] |= int32_t((qh_hi * 0x02040810u) & 0x10101010u); 388 } 389 } 390 } else { 391 [[unroll]] for (uint32_t d = 0; d < d_per_step; d++) { 392 k_quants[d] = get_k_qs_shmem(buf_ib, (d_tid * (HSK_per_thread / 4) + d_block) % 8 + d); 393 } 394 } 395 } else { 396 const uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d_tid * (HSK_per_thread / 4) + d_block); 397 const uint ib = coord / BLOCK_SIZE_K; 398 const uint iqs = (coord % BLOCK_SIZE_K); 399 400 k_dm = ACC_TYPEV2(get_k_scale(ib, k_offset)); 401 402 if (block8_fast) { 403 fa_k_qs_block8 blk = get_k_qs_block8(ib, k_offset); 404 [[unroll]] for (uint32_t d = 0; d < 8; d++) { 405 k_quants[d] = blk.qs[d]; 406 } 407 } else { 408 [[unroll]] for (uint32_t d = 0; d < d_per_step; d++) { 409 k_quants[d] = get_k_qs(ib, iqs + d * 4, k_offset); 410 } 411 } 412 } 413 414 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 415 const uint qib = tile_row(r) * qf_stride + (d_tid * (HSK_per_thread / 4) + d_block) / 8; 416 const uint qiqs = (d_tid * (HSK_per_thread / 4) + d_block) % 8; 417 418 int32_t acc = 0; 419 [[unroll]] for (uint32_t d = 0; d < d_per_step; d++) { 420 acc += dotPacked4x8EXT(Qf[qib].qs[qiqs + d], k_quants[d]); 421 } 422 423 Sf[r][c] += ACC_TYPE(acc) * ACC_TYPE(Qf[qib].ds.x) * k_dm.x; 424 if ((d_tid * (HSK_per_thread / 4) + d_block) % 8 == 0) { 425 Sf[r][c] += k_dot_correction(qib, k_dm); 426 } 427 } 428 } 429 } 430#endif // MMQ 431 432 [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { 433 // Compute sum across the D_split 434 [[unroll]] for (uint s = D_split / 2; s > 0; s >>= 1) { 435 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 436 Sf[r][c] += subgroupShuffleXor(Sf[r][c], s); 437 } 438 } 439 } 440 441 if (LOGIT_SOFTCAP) { 442 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 443 [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { 444 Sf[r][c] = ACC_TYPE(p.logit_softcap * tanh(Sf[r][c])); 445 } 446 } 447 } 448 449 if (MASK_ENABLE && mask_opt_bits != MASK_OPT_ALL_ZERO) { 450 [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { 451 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 452 FLOAT_TYPE mvf = masksh[(c * cols_per_iter + col_tid) * masksh_stride + tile_row(r)]; 453 454 Sf[r][c] += slope[r]*mvf; 455 } 456 } 457 } 458 459 float eMf[rows_per_thread]; 460 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 461 float rowmaxf = NEG_FLT_MAX_OVER_2; 462 [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { 463 if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { 464 continue; 465 } 466 rowmaxf = max(rowmaxf, float(Sf[r][c])); 467 } 468 rowmaxf += FATTN_KQ_MAX_OFFSET; 469 float Moldf = Mf[r]; 470 471 // M = max(rowmax, Mold) 472 // P = e^(S - M) 473 // eM = e^(Mold - M) 474 Mf[r] = max(rowmaxf, Moldf); 475 eMf[r] = exp(Moldf - Mf[r]); 476 Lf[r] = eMf[r]*Lf[r]; 477 } 478 479 [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { 480 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 481 Of[r][d] = FLOAT_TYPE(eMf[r]) * Of[r][d]; 482 } 483 } 484 485 if (SHMEM_STAGING != 0) { 486 barrier(); 487 [[unroll]] for (uint32_t idx = 0; idx < Bc * HSV / 4; idx += gl_WorkGroupSize.x) { 488 uint32_t d = (idx + tid) % (HSV / 4); 489 uint32_t c = (idx + tid) / (HSV / 4); 490 if (idx + gl_WorkGroupSize.x <= Bc * HSV / 4 || c < Bc) { 491 FLOAT_TYPEV4 V_Tf = FLOAT_TYPEV4(0); 492 if (!KV_bounds_check || j * Bc + c < KV) { 493 if (USE_DECODE_V) { 494 uint coord = (j * Bc + c) * v_stride * BLOCK_SIZE_V + 4 * d; 495 uint ib = coord / BLOCK_SIZE_V; 496 uint iqs = (coord % BLOCK_SIZE_V); 497 V_Tf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); 498 } else { 499 V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c) * v_stride / 4 + d]); 500 } 501 } 502 503 kvsh[c * kvsh_stride + d] = V_Tf; 504 } 505 } 506 barrier(); 507 } 508 509 [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { 510 if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { 511 continue; 512 } 513 514 FLOAT_TYPE Pf[rows_per_thread]; 515 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 516 Pf[r] = FLOAT_TYPE(exp(float(Sf[r][c]) - Mf[r])); 517 Lf[r] += Pf[r]; 518 } 519 520 [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { 521 FLOAT_TYPEV4 Vf; 522 if (SHMEM_STAGING != 0) { 523 Vf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)]; 524 } else if (USE_DECODE_V) { 525 uint coord = (j * Bc + c * cols_per_iter + col_tid) * v_stride * BLOCK_SIZE_V + 4 * (d * D_split + d_tid); 526 uint ib = coord / BLOCK_SIZE_V; 527 uint iqs = (coord % BLOCK_SIZE_V); 528 Vf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); 529 } else { 530 Vf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * v_stride / 4 + d * D_split + d_tid]); 531 } 532 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 533 Of[r][d] += FLOAT_TYPEV4(Pf[r] * Vf); 534 } 535 } 536 } 537 } 538 539 // prevent race on tmpsh 540 barrier(); 541 542 // reduce across threads 543 544 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 545 float rowmaxf = Mf[r]; 546 547 // Compute max across the row 548 if (SubGroupSize > 0) { 549 [[unroll]] for (uint s = D_split; s < SubGroupSize; s *= 2) { 550 rowmaxf = max(rowmaxf, subgroupShuffleXor(rowmaxf, s)); 551 } 552 if (row_split == 1) { 553 // Reduce inside workgroup with shmem 554 barrier(); 555 if (gl_SubgroupInvocationID == d_tid) { 556 tmpsh[gl_SubgroupID * D_split + d_tid] = rowmaxf; 557 } 558 barrier(); 559 rowmaxf = tmpsh[d_tid]; 560 [[unroll]] for (uint32_t s = 1; s < num_subgroups; ++s) { 561 rowmaxf = max(rowmaxf, tmpsh[s * D_split + d_tid]); 562 } 563 } 564 } else { 565 barrier(); 566 tmpsh[tid] = rowmaxf; 567 barrier(); 568 [[unroll]] for (int s = int(threads_per_rowgroup) / 2; s >= D_split; s >>= 1) { 569 if (rowgroup_tid < s) { 570 tmpsh[tid] = max(tmpsh[tid], tmpsh[tid ^ s]); 571 } 572 barrier(); 573 } 574 rowmaxf = tmpsh[row_tid * threads_per_rowgroup + d_tid]; 575 } 576 577 float Moldf = Mf[r]; 578 579 // M = max(rowmax, Mold) 580 // eM = e^(Mold - M) 581 Mf[r] = max(rowmaxf, Moldf); 582 float eMf = exp(Moldf - Mf[r]); 583 584 Lf[r] = eMf*Lf[r]; 585 586 // Compute sum across the row 587 if (SubGroupSize > 0) { 588 [[unroll]] for (uint s = D_split; s < SubGroupSize; s *= 2) { 589 Lf[r] += subgroupShuffleXor(Lf[r], s); 590 } 591 if (row_split == 1) { 592 barrier(); 593 if (gl_SubgroupInvocationID == d_tid) { 594 tmpsh[gl_SubgroupID * D_split + d_tid] = Lf[r]; 595 } 596 barrier(); 597 Lf[r] = tmpsh[d_tid]; 598 [[unroll]] for (uint32_t s = 1; s < num_subgroups; ++s) { 599 Lf[r] += tmpsh[s * D_split + d_tid]; 600 } 601 } 602 } else { 603 barrier(); 604 tmpsh[tid] = Lf[r]; 605 barrier(); 606 [[unroll]] for (int s = int(threads_per_rowgroup) / 2; s >= D_split; s >>= 1) { 607 if (rowgroup_tid < s) { 608 tmpsh[tid] = tmpsh[tid] + tmpsh[tid ^ s]; 609 } 610 barrier(); 611 } 612 Lf[r] = tmpsh[row_tid * threads_per_rowgroup + d_tid]; 613 } 614 615 [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { 616 Of[r][d] = FLOAT_TYPE(eMf) * Of[r][d]; 617 618 if (SubGroupSize > 0) { 619 [[unroll]] for (uint s = D_split; s < SubGroupSize; s *= 2) { 620 if (!OLD_AMD_WINDOWS) { 621 Of[r][d] += subgroupShuffleXor(Of[r][d], s); 622 } else { 623 // Something about f16vec4 subgroupShuffleXor is broken on AMD Windows RDNA2 and below. 624 // Shuffle full vec4 as workaround. 625 // See https://github.com/ggml-org/llama.cpp/issues/19881#issuecomment-3958643697 626 Of[r][d] += FLOAT_TYPEV4(subgroupShuffleXor(vec4(Of[r][d]), s)); 627 } 628 } 629 if (row_split == 1) { 630 barrier(); 631 if (gl_SubgroupInvocationID == d_tid) { 632 tmpshv4[gl_SubgroupID * D_split + d_tid] = Of[r][d]; 633 } 634 barrier(); 635 Of[r][d] = tmpshv4[d_tid]; 636 [[unroll]] for (uint32_t s = 1; s < num_subgroups; ++s) { 637 Of[r][d] += tmpshv4[s * D_split + d_tid]; 638 } 639 } 640 } else { 641 barrier(); 642 tmpshv4[tid] = Of[r][d]; 643 barrier(); 644 [[unroll]] for (int s = int(threads_per_rowgroup) / 2; s >= D_split; s >>= 1) { 645 if (rowgroup_tid < s) { 646 Of[r][d] += tmpshv4[tid ^ s]; 647 tmpshv4[tid] = Of[r][d]; 648 } 649 barrier(); 650 } 651 Of[r][d] = tmpshv4[row_tid * threads_per_rowgroup + d_tid]; 652 } 653 } 654 } 655 656 657 // If there is split_k, then the split_k resolve shader does the final 658 // division by L. Store the intermediate O value and per-row m and L values. 659 if (p.k_num > 1) { 660 if (p.gqa_ratio > 1) { 661 // note: O and Q have swapped coord 1,2. 662 uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)) / 4; 663 664 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 665 const uint row = tile_row(r); 666 if (row < N) { 667 [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { 668 gqaStore(row, d * D_split + d_tid, Of[r][d], o_offset, iq2, N); 669 } 670 } 671 } 672 673 o_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)); 674 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 675 const uint row = tile_row(r); 676 if (row < N) { 677 perElemOpStoreCol0(row, 0u, ACC_TYPE(Lf[r]), o_offset, iq2, N); 678 perElemOpStoreCol0(row, 0u, ACC_TYPE(Mf[r]), o_offset + p.ne1, iq2, N); 679 } 680 } 681 } else { 682 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 683 const uint row = tile_row(r); 684 const uint global_row = i * Br + row; 685 686 if (global_row < N) { 687 uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)) / 4; 688 689 [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { 690 data_ov4[o_offset + iq2 * HSV/4 + d * D_split + d_tid] = D_TYPEV4(Of[r][d]); 691 } 692 } 693 694 if (global_row < N && d_tid == 0 && col_tid == 0) { 695 uint32_t lm_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)); 696 data_o[lm_offset + iq2] = D_TYPE(Lf[r]); 697 data_o[lm_offset + p.ne1 + iq2] = D_TYPE(Mf[r]); 698 } 699 } 700 } 701 return; 702 } 703 704 if ((p.mask_n_head_log2 & SINK_ENABLE_BIT) != 0) { 705 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 706 float sink = perElemOpGetSink(tile_row(r), 0u, ACC_TYPE(0), iq2); 707 708 float ms = 1.0f; 709 float vs = 1.0f; 710 711 if (sink > Mf[r]) { 712 ms = exp(Mf[r] - sink); 713 714 [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { 715 Of[r][d] *= FLOAT_TYPE(ms); 716 } 717 } else { 718 vs = exp(sink - Mf[r]); 719 } 720 721 Lf[r] = Lf[r]*ms + vs; 722 } 723 } 724 725 float Lfrcp[rows_per_thread]; 726 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 727 Lfrcp[r] = (Lf[r] == 0.0) ? 0.0 : (1.0 / Lf[r]); 728 } 729 730 [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { 731 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 732 Of[r][d] *= FLOAT_TYPE(Lfrcp[r]); 733#if defined(FLOAT_TYPE_MAX) 734 Of[r][d] = clamp(Of[r][d], -FLOAT_TYPE_MAX, FLOAT_TYPE_MAX); 735#endif 736 } 737 } 738 739 uint32_t o_offset = (gqa_iq1*p.ne1*HSV + iq3*p.ne2*p.ne1*HSV) / 4; 740 741 if (p.gqa_ratio > 1) { 742 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 743 const uint row = tile_row(r); 744 if (row < N) { 745 [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { 746 gqaStore(row, d * D_split + d_tid, Of[r][d], o_offset, iq2, N); 747 } 748 } 749 } 750 } else { 751 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 752 const uint row = tile_row(r); 753 if (i * Br + row < N) { 754 [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { 755 data_ov4[o_offset + (iq2 * HSV + (i * Br + row) * p.ne1 * HSV) / 4 + d * D_split + d_tid] = D_TYPEV4(Of[r][d]); 756 } 757 } 758 } 759 } 760}