char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
816164100
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_float16 : require 7#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require 8 9#if defined(BFLOAT16) 10#extension GL_EXT_bfloat16 : enable 11#endif 12 13#extension GL_KHR_shader_subgroup_basic : enable 14#extension GL_KHR_shader_subgroup_arithmetic : enable 15#extension GL_KHR_shader_subgroup_vote : enable 16#extension GL_KHR_memory_scope_semantics : enable 17#extension GL_KHR_cooperative_matrix : enable 18 19#include "types.glsl" 20#include "flash_attn_base.glsl" 21#if !defined(BFLOAT16) 22#include "flash_attn_dequant.glsl" 23#endif 24 25// These need to be supported N,M values for a MatBc x MatBr x 16 coopmatmuladd 26const uint32_t MatBr = 16; 27const uint32_t MatBc = 16; 28 29const uint32_t rows_per_thread = Br / row_split; 30const uint32_t cols_per_iter = gl_WorkGroupSize.x / row_split; 31const uint32_t cols_per_thread = Bc / cols_per_iter; 32 33 34layout (binding = 0) readonly buffer Q {float data_q[];}; 35layout (binding = 0) readonly buffer QV4 {vec4 data_qv4[];}; 36layout (binding = 1) readonly buffer K {FLOAT_TYPE data_k[];}; 37layout (binding = 1) readonly buffer KV4 {FLOAT_TYPEV4 data_kv4[];}; 38layout (binding = 2) readonly buffer V {FLOAT_TYPE data_v[];}; 39layout (binding = 2) readonly buffer VV4 {FLOAT_TYPEV4 data_vv4[];}; 40layout (binding = 3) readonly buffer M {float16_t data_m[];}; 41 42shared float tmpsh[row_split]; 43 44const uint32_t qstride = HSK_pad / 4 + 2; 45shared FLOAT_TYPEV4 Qf[Br * qstride]; 46 47const uint psh_stride = Br / 4 + 2; 48shared FLOAT_TYPEV4 Psh[Bc * psh_stride]; 49 50// Avoid padding for hsk==256 to make it fit in 48KB shmem. 51const uint32_t sfshstride = (HSK <= 128) ? (Br / 4 + 2) : Br / 4; 52shared ACC_TYPEV4 sfsh[Bc * sfshstride]; 53 54const uint32_t D_pad = HSK_pad > HSV_pad ? HSK_pad : HSV_pad; 55const uint32_t kvsh_stride = (SHMEM_STAGING != 0 ? D_pad : MatBr) / 4 + 2; 56const uint v_cols = MatBc / 4 * row_split; // total cols, 4 vec4s per MatBc * number of subgroups 57const uint vsh_stride = v_cols; 58shared FLOAT_TYPEV4 kvsh[(kvsh_stride >= vsh_stride) ? (Bc * kvsh_stride) : (Bc * vsh_stride)]; 59 60const uint32_t osh_stride = row_split * MatBr / 4; 61shared O_TYPEV4 pvsh[MatBc * osh_stride]; 62 63shared ACC_TYPE slope[Br]; 64 65void main() { 66#ifdef NEEDS_INIT_IQ_SHMEM 67 if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) { 68 init_iq_shmem(gl_WorkGroupSize); 69 } 70#endif 71 72 init_indices(); 73 74 const uint32_t tid = gl_LocalInvocationIndex; 75 76 const uint32_t threads_per_rowgroup = gl_WorkGroupSize.x / row_split; 77 const uint32_t d_per_thread = (HSV/4 + threads_per_rowgroup - 1) / threads_per_rowgroup; 78 const uint32_t row_tid = gl_LocalInvocationIndex / threads_per_rowgroup; 79 const uint32_t col_tid = gl_LocalInvocationIndex % threads_per_rowgroup; 80 81#define tile_row(r) (row_tid * rows_per_thread + (r)) 82 83 // Zero-initialize shared memory for Q/K when HSK is not a multiple of 16 (HSK_pad > HSK). 84 if ((HSK % 16) != 0) { 85 [[unroll]] for (uint i = 0; i < Br * qstride; i += gl_WorkGroupSize.x) { 86 if (i + tid < Br * qstride) { 87 Qf[i + tid] = FLOAT_TYPEV4(0); 88 } 89 } 90 barrier(); 91 } 92 93 uint32_t q_offset = gqa_iq1*p.nb01 + (iq2*p.nb02+iq3*p.nb03) / 4; 94 95 [[unroll]] for (uint32_t idx = 0; idx < Br * HSK / 4; idx += gl_WorkGroupSize.x) { 96 uint32_t d = (idx + tid) % (HSK / 4); 97 uint32_t r = (idx + tid) / (HSK / 4); 98 if (r < Br && d < HSK / 4 && 99 i * Br + r < N) { 100 Qf[r * qstride + d] = FLOAT_TYPEV4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale); 101 } 102 } 103 barrier(); 104 105 O_TYPEV4 Of[rows_per_thread][d_per_thread]; 106 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 107 [[unroll]] for (uint32_t d = 0; d < d_per_thread; ++d) { 108 Of[r][d] = O_TYPEV4(0.0); 109 } 110 } 111 112 float Lf[rows_per_thread], Mf[rows_per_thread]; 113 114 // Use -FLT_MAX/2 rather than -inf to reduce the possibility of NaNs, e.g. when computing Mold-M. 115 const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF); 116 117 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 118 Lf[r] = 0; 119 Mf[r] = NEG_FLT_MAX_OVER_2; 120 } 121 122 // ALiBi 123 if (p.max_bias > 0.0f) { 124 if (tid < Br) { 125 uint r = tid; 126 slope[r] = perElemOpComputeSlope(r, col_tid, ACC_TYPE(0), iq2); 127 } 128 } else { 129 if (tid < Br) { 130 uint r = tid; 131 slope[r] = ACC_TYPE(1.0); 132 } 133 } 134 135 const uint32_t mo_stride = CEIL_DIV(KV, 16 * Bc); 136 // mo_offset will point to the tile starting at row i*Br and col 0 137 uint32_t mo_offset = mo_stride * i; 138 139 // FaBlockBytesK/V == 2 for f16 (sizeof f16) and == 16 for f32 (vec4) and == ggml block size for quants. 140 uint32_t k_offset = (ik2*p.nb12 + ik3*p.nb13) / FaBlockBytesK; 141 uint32_t v_offset = (iv2*p.nb22 + iv3*p.nb23) / FaBlockBytesV; 142 uint32_t m_offset = gqa_iq1*KV; 143 if (p.nem2 != 1 || p.nem3 != 1) { 144 m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * KV; 145 mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride; 146 } 147 148 uint32_t mask_opt = 0; 149 uint32_t mask_opt_idx = ~0; 150 uint32_t mask_opt_bits = 0; 151 f16vec4 mask_cache[Bc * Br / 4 / WorkGroupSize]; 152 153 [[dont_unroll]] 154 for (uint32_t j = start_j; j < end_j; ++j) { 155 156 [[unroll]] for (uint32_t idx = 0; idx < mask_cache.length(); ++idx) { 157 mask_cache[idx] = f16vec4(0); 158 } 159 160 if (MASK_ENABLE) { 161 if (USE_MASK_OPT && mask_opt_idx != j / 16) { 162 mask_opt_idx = j / 16; 163 mask_opt = data_mask_opt[mo_offset + mask_opt_idx]; 164 } 165 mask_opt_bits = (mask_opt >> ((j % 16) * 2)) & 0x3; 166 if (mask_opt_bits == MASK_OPT_ALL_NEG_INF) { 167 // skip this block 168 continue; 169 } 170 // Only load if the block is not all zeros 171 if (mask_opt_bits != MASK_OPT_ALL_ZERO) { 172 bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0; 173 174 float max_mask = NEG_FLT_MAX_OVER_2; 175 [[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) { 176 uint32_t c = (idx + tid) / (Br / 4); 177 uint32_t r = (idx + tid) % (Br / 4); 178 if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { 179 if ((!KV_bounds_check || j * Bc + c < KV)) { 180 f16vec4 m; 181 if (!nem1_bounds_check || i * Br + r * 4 + 3 < p.nem1) { 182 m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], 183 data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], 184 data_m[m_offset + (i * Br + r * 4 + 2) * m_stride + (j * Bc + c)], 185 data_m[m_offset + (i * Br + r * 4 + 3) * m_stride + (j * Bc + c)]); 186 max_mask = max(max(max(max(max_mask, float(m[0])), float(m[1])), float(m[2])), float(m[3])); 187 } else if (i * Br + r * 4 + 2 < p.nem1) { 188 m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], 189 data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], 190 data_m[m_offset + (i * Br + r * 4 + 2) * m_stride + (j * Bc + c)], 191 0.0); 192 max_mask = max(max(max(max_mask, float(m[0])), float(m[1])), float(m[2])); 193 } else if (i * Br + r * 4 + 1 < p.nem1) { 194 m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], 195 data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], 196 0.0, 197 0.0); 198 max_mask = max(max(max_mask, float(m[0])), float(m[1])); 199 } else if (i * Br + r * 4 < p.nem1) { 200 m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], 201 0.0, 202 0.0, 203 0.0); 204 max_mask = max(max_mask, float(m[0])); 205 } else { 206 m = f16vec4(0.0); 207 } 208 mask_cache[idx / WorkGroupSize] = m; 209 } 210 } 211 } 212 // skip the block if the mask is entirely -inf 213 bool all_less = subgroupAll(max_mask <= NEG_FLT_MAX_OVER_2); 214 barrier(); 215 if (gl_SubgroupInvocationID == 0) { 216 tmpsh[gl_SubgroupID] = all_less ? NEG_FLT_MAX_OVER_2 : 0.0f; 217 } 218 barrier(); 219 [[unroll]] for (uint s = 0; s < gl_NumSubgroups; ++s) { 220 max_mask = max(max_mask, tmpsh[s]); 221 } 222 if (max_mask <= NEG_FLT_MAX_OVER_2) { 223 continue; 224 } 225 } 226 } 227 228 if (SHMEM_STAGING != 0) { 229 [[unroll]] for (uint32_t idx = 0; idx < Bc * HSK_pad / 4; idx += gl_WorkGroupSize.x) { 230 uint32_t d = (idx + tid) % (HSK_pad / 4); 231 uint32_t c = (idx + tid) / (HSK_pad / 4); 232 if (idx + gl_WorkGroupSize.x <= Bc * HSK_pad / 4 || c < Bc) { 233 FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0); 234 if ((!KV_bounds_check || j * Bc + c < KV) && (HSK == HSK_pad || d < HSK / 4)) { 235#if !defined(BFLOAT16) 236 if (USE_DECODE_K) { 237 uint coord = (j * Bc + c) * k_stride * BLOCK_SIZE_K + 4 * d; 238 uint ib = coord / BLOCK_SIZE_K; 239 uint iqs = (coord % BLOCK_SIZE_K); 240 K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); 241 } else 242#endif 243 { 244 K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c) * k_stride / 4 + d]); 245 } 246 } 247 248 kvsh[c * kvsh_stride + d] = K_Tf; 249 } 250 } 251 barrier(); 252 } 253 254 // K * Q^T -> S^T: Bc x HSK_pad * HSK_pad x Br -> Bc x Br 255 // Bc split across workgroup (four subgroups), loop over HSK in chunks of 16: 16 x 16 * 16 x 16 -> 16 x 16 256 // This is written transposed in order to allow for N being 8 if implementations need it 257 coopmat<ACC_TYPE, gl_ScopeSubgroup, MatBc, MatBr, gl_MatrixUseAccumulator> SfMat = coopmat<ACC_TYPE, gl_ScopeSubgroup, MatBc, MatBr, gl_MatrixUseAccumulator>(0); 258 coopmat<FLOAT_TYPE, gl_ScopeSubgroup, MatBc, 16, gl_MatrixUseA> KMat; 259 coopmat<FLOAT_TYPE, gl_ScopeSubgroup, 16, MatBr, gl_MatrixUseB> QMat; 260 261 [[unroll]] for (uint32_t d = 0; d < HSK_pad / 16; ++d) { 262 // If SHMEM_STAGING is set, a Bc * HSK_pad size tile of K is loaded to shmem 263 // If not, K is loaded directly from global memory if aligned, otherwise 264 // staged through a Bc * MatBr size staging buffer. 265 // If K is a quant type, then it is always staged for dequantization. 266 if (SHMEM_STAGING == 0) { 267 // For quants we always need to dequant into kvsh; for f16/bf16 we can load 268 // directly from global memory when alignment / bounds allow it. 269 const bool stage_k = USE_DECODE_K || KV_bounds_check || d * 16 + 16 > HSK; 270 if (stage_k) { 271 barrier(); 272 [[unroll]] for (uint32_t idx = 0; idx < Bc * MatBr / 4; idx += gl_WorkGroupSize.x) { 273 uint32_t col_vec = (idx + tid) % (MatBr / 4); 274 uint32_t row = (idx + tid) / (MatBr / 4); 275 if (idx + tid < Bc * MatBr / 4) { 276 FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0); 277 if ((!KV_bounds_check || j * Bc + row < KV) && (HSK == HSK_pad || d * 16 + col_vec * 4 < HSK)) { 278#if !defined(BFLOAT16) 279 if (USE_DECODE_K) { 280 uint coord = (j * Bc + row) * k_stride * BLOCK_SIZE_K + d * 16 + col_vec * 4; 281 uint ib = coord / BLOCK_SIZE_K; 282 uint iqs = (coord % BLOCK_SIZE_K); 283 K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); 284 } else 285#endif 286 { 287 K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + row) * k_stride / 4 + d * 16 / 4 + col_vec]); 288 } 289 } 290 291 kvsh[row * kvsh_stride + col_vec] = K_Tf; 292 } 293 } 294 barrier(); 295 } 296 297 if (stage_k) { 298 uint coord = (gl_SubgroupID * MatBc) * kvsh_stride; 299 coopMatLoad(KMat, kvsh, coord, kvsh_stride, gl_CooperativeMatrixLayoutRowMajor); 300 } else { 301 const uint coord = k_offset / 4 + (j * Bc + gl_SubgroupID * MatBc) * k_stride / 4 + d * 16 / 4; 302 coopMatLoad(KMat, data_kv4, coord, k_stride / 4, gl_CooperativeMatrixLayoutRowMajor); 303 } 304 } else { 305 uint coord = (gl_SubgroupID * MatBc) * kvsh_stride + d * 16 / 4; 306 coopMatLoad(KMat, kvsh, coord, kvsh_stride, gl_CooperativeMatrixLayoutRowMajor); 307 } 308 309 coopMatLoad(QMat, Qf, d * 16 / 4, qstride, gl_CooperativeMatrixLayoutColumnMajor); 310 311 SfMat = coopMatMulAdd(KMat, QMat, SfMat); 312 } 313 314 uint coord = gl_SubgroupID * MatBc * sfshstride; 315 coopMatStore(SfMat, sfsh, coord, sfshstride, gl_CooperativeMatrixLayoutRowMajor); 316 barrier(); 317 318 if (LOGIT_SOFTCAP) { 319 [[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) { 320 uint32_t c = (idx + tid) / (Br / 4); 321 uint32_t r = (idx + tid) % (Br / 4); 322 if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { 323 sfsh[c * sfshstride + r] = ACC_TYPEV4(p.logit_softcap * tanh(sfsh[c * sfshstride + r])); 324 } 325 } 326 barrier(); 327 } 328 329 if (MASK_ENABLE && mask_opt_bits != MASK_OPT_ALL_ZERO) { 330 [[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) { 331 uint32_t c = (idx + tid) / (Br / 4); 332 uint32_t r = (idx + tid) % (Br / 4); 333 if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { 334 if (!KV_bounds_check || j * Bc + c < KV) { 335 // Mask nem1 bounds check is handled when loading masks 336 ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); 337 ACC_TYPEV4 slopes = ACC_TYPEV4(slope[r * 4], slope[r * 4 + 1], slope[r * 4 + 2], slope[r * 4 + 3]); 338 sfsh[c * sfshstride + r] += slopes * masks; 339 } 340 } 341 } 342 barrier(); 343 } 344 345 float eMf[rows_per_thread]; 346 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 347 const uint r_vec = tile_row(r) / 4; 348 const uint r_comp = tile_row(r) % 4; 349 350 float rowmaxf = NEG_FLT_MAX_OVER_2; 351 [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { 352 if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { 353 continue; 354 } 355 rowmaxf = max(rowmaxf, float(sfsh[r_vec + (c * cols_per_iter + col_tid) * sfshstride][r_comp])); 356 } 357 rowmaxf += FATTN_KQ_MAX_OFFSET; 358 float Moldf = Mf[r]; 359 360 // Compute max across the row 361 rowmaxf = subgroupMax(rowmaxf); 362 363 // M = max(rowmax, Mold) 364 // P = e^(S - M) 365 // eM = e^(Mold - M) 366 Mf[r] = max(rowmaxf, Moldf); 367 eMf[r] = exp(Moldf - Mf[r]); 368 369 Lf[r] = eMf[r]*Lf[r]; 370 } 371 372 [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { 373 const uint d_local = d0 / threads_per_rowgroup; 374 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 375 Of[r][d_local] = O_TYPE(eMf[r]) * Of[r][d_local]; 376 } 377 } 378 379 // Calculate and store Pf in Psh 380 [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { 381 const uint col = c * cols_per_iter + col_tid; 382 383 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; r += 4) { 384 const uint row = tile_row(r); 385 if (KV_bounds_check && j * Bc + col >= KV) { 386 Psh[col * psh_stride + row / 4] = FLOAT_TYPEV4(0.0f); 387 } else { 388 const vec4 mfvec = vec4(Mf[r], Mf[r + 1], Mf[r + 2], Mf[r + 3]); 389 const FLOAT_TYPEV4 Pf = FLOAT_TYPEV4(exp(vec4(sfsh[row / 4 + col * sfshstride]) - mfvec)); 390 [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { 391 Lf[r + vec_idx] += Pf[vec_idx]; 392 } 393 Psh[col * psh_stride + row / 4] = Pf; 394 } 395 } 396 } 397 398 if (SHMEM_STAGING != 0) { 399 [[unroll]] for (uint32_t idx = 0; idx < Bc * HSV_pad / 4; idx += gl_WorkGroupSize.x) { 400 uint32_t d = (idx + tid) % (HSV_pad / 4); 401 uint32_t c = (idx + tid) / (HSV_pad / 4); 402 if (idx + gl_WorkGroupSize.x <= Bc * HSV_pad / 4 || c < Bc) { 403 FLOAT_TYPEV4 V_Tf = FLOAT_TYPEV4(0); 404 if ((!KV_bounds_check || j * Bc + c < KV) && (HSV == HSV_pad || d < HSV / 4)) { 405#if !defined(BFLOAT16) 406 if (USE_DECODE_V) { 407 uint coord = (j * Bc + c) * v_stride * BLOCK_SIZE_V + 4 * d; 408 uint ib = coord / BLOCK_SIZE_V; 409 uint iqs = (coord % BLOCK_SIZE_V); 410 V_Tf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); 411 } else 412#endif 413 { 414 V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c) * v_stride / 4 + d]); 415 } 416 } 417 418 kvsh[c * kvsh_stride + d] = V_Tf; 419 } 420 } 421 } 422 barrier(); 423 424 const uint num_hsv_tiles = (HSV + MatBc * row_split - 1) / (MatBc * row_split); // round up 425 426 // Each subgroup handles HSV/4 columns 427 [[unroll]] for (uint32_t hsv_tile = 0; hsv_tile < num_hsv_tiles; ++hsv_tile) { 428 const uint hsv_offset = (hsv_tile * row_split + gl_SubgroupID) * 16; 429 430 coopmat<O_TYPE, gl_ScopeSubgroup, MatBc, MatBr, gl_MatrixUseAccumulator> PVMat = coopmat<O_TYPE, gl_ScopeSubgroup, MatBc, MatBr, gl_MatrixUseAccumulator>(0); 431 432 // Preload V tiles for [Bc, 16 * num subgroups] 433 const uint v_rows = Bc; 434 const uint v_total = v_rows * v_cols; 435 const uint v_loads_per_thread = v_total / gl_WorkGroupSize.x; 436 437 // If SHMEM_STAGING is set, a Bc * HSV_pad size tile of V is loaded to shmem. 438 // If not, V is loaded directly from global memory if aligned, otherwise 439 // staged through a Bc * MatBr size staging buffer. 440 // If V is a quant type, then it is always staged for dequantization. 441 if (SHMEM_STAGING == 0) { 442 // For quants we always preload via kvsh. For f16/bf16 we only preload when 443 // alignment / bounds force it (otherwise we coopMatLoad direct from data_vv4). 444 const bool stage_v = USE_DECODE_V || KV_bounds_check; 445 if (stage_v) { 446 [[unroll]] for (uint32_t i = 0; i < v_loads_per_thread; ++i) { 447 const uint idx = i * gl_WorkGroupSize.x + tid; 448 const uint row = idx / v_cols; 449 const uint col = idx % v_cols; 450 451 const uint v_row = j * Bc + row; 452 const uint v_col = hsv_tile * MatBc * row_split + col * 4; 453 454 const uint coord = v_row * v_stride * BLOCK_SIZE_V + v_col; 455 const uint ib = coord / BLOCK_SIZE_V; 456 const uint iqs = coord % BLOCK_SIZE_V; 457 458 if (!KV_bounds_check || (v_row < KV && v_col < HSV)) { 459#if !defined(BFLOAT16) 460 if (USE_DECODE_V) { 461 kvsh[row * vsh_stride + col] = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); 462 } else 463#endif 464 { 465 kvsh[row * vsh_stride + col] = data_vv4[(v_offset + v_row * v_stride + v_col) / 4]; 466 } 467 } else { 468 kvsh[row * vsh_stride + col] = FLOAT_TYPEV4(0.0f); 469 } 470 } 471 } 472 } 473 barrier(); 474 475 const uint o_offset = gl_SubgroupID * MatBr / 4; 476 477 if (hsv_offset < HSV_pad) { 478 [[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) { 479 coopMatLoad(KMat, Psh, bc_chunk * MatBc * psh_stride, psh_stride, gl_CooperativeMatrixLayoutColumnMajor); 480 481 if (SHMEM_STAGING == 0) { 482 if (!USE_DECODE_V && !KV_bounds_check) { 483 // F16/BF16 values can be loaded directly from global memory 484 const uint v_tile_row = j * Bc + bc_chunk * MatBc; 485 const uint v_tile_offset = v_offset / 4 + v_tile_row * v_stride / 4 + hsv_offset / 4; 486 coopMatLoad(QMat, data_vv4, v_tile_offset, v_stride / 4, gl_CooperativeMatrixLayoutRowMajor); 487 } else { 488 const uint v_tile_offset = bc_chunk * MatBr * v_cols + gl_SubgroupID * (MatBc / 4); 489 coopMatLoad(QMat, kvsh, v_tile_offset, vsh_stride, gl_CooperativeMatrixLayoutRowMajor); 490 } 491 } else { 492 const uint v_tile_offset = bc_chunk * MatBc * kvsh_stride + (hsv_tile * row_split + gl_SubgroupID) * (MatBc / 4); 493 coopMatLoad(QMat, kvsh, v_tile_offset, kvsh_stride, gl_CooperativeMatrixLayoutRowMajor); 494 } 495 496 PVMat = coopMatMulAdd(KMat, QMat, PVMat); 497 } 498 499 // Store PVMat to pvsh and load into Of 500 coopMatStore(PVMat, pvsh, o_offset, osh_stride, gl_CooperativeMatrixLayoutRowMajor); 501 } 502 503 barrier(); 504 505 const uint hsv_per_tile = row_split * MatBc; 506 const uint hsv_base = hsv_tile * hsv_per_tile; 507 const uint d_values_per_tile = hsv_per_tile / 4; 508 509 const uint d_start = hsv_tile * d_values_per_tile; 510 const uint d_end = min(d_start + d_values_per_tile, HSV / 4); 511 512 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 513 const uint row = tile_row(r); 514 515 [[unroll]] for (uint32_t d_local = 0; d_local < d_per_thread; ++d_local) { 516 const uint d = d_local * threads_per_rowgroup + col_tid; 517 const uint hsv_col = 4 * d; 518 519 if (hsv_col >= hsv_base && hsv_col < hsv_base + hsv_per_tile && hsv_col < HSV) { 520 const uint local_hsv = (hsv_col - hsv_base) / 4; 521 Of[r][d_local] += pvsh[row * osh_stride + local_hsv]; 522 } 523 } 524 } 525 } 526 527 barrier(); 528 } 529 530 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 531 Lf[r] = subgroupAdd(Lf[r]); 532 } 533 534 // If there is split_k, then the split_k resolve shader does the final 535 // division by L. Store the intermediate O value and per-row m and L values. 536 if (p.k_num > 1) { 537 if (p.gqa_ratio > 1) { 538 // note: O and Q have swapped coord 1,2. 539 uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)) / 4; 540 541 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 542 if (tile_row(r) < N) { 543 [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { 544 const uint d = d0 + col_tid; 545 if (d >= HSV/4) break; 546 const uint d_local = d0 / threads_per_rowgroup; 547 gqaStore(tile_row(r), d, Of[r][d_local], o_offset, iq2, N); 548 } 549 } 550 } 551 552 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)); 553 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 554 if (tile_row(r) < N) { 555 perElemOpStoreCol0(tile_row(r), 0u, ACC_TYPE(Lf[r]), o_offset, iq2, N); 556 perElemOpStoreCol0(tile_row(r), 0u, ACC_TYPE(Mf[r]), o_offset + p.ne1, iq2, N); 557 } 558 } 559 } else { 560 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 561 const uint row = tile_row(r); 562 const uint global_row = i * Br + row; 563 564 if (global_row < N) { 565 uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)) / 4; 566 567 [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { 568 const uint d = d0 + col_tid; 569 if (d >= HSV/4) break; 570 data_ov4[o_offset + iq2 * HSV/4 + d] = D_TYPEV4(Of[r][d/threads_per_rowgroup]); 571 } 572 } 573 574 if (global_row < N && col_tid == 0) { 575 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)); 576 data_o[lm_offset + iq2] = D_TYPE(Lf[r]); 577 data_o[lm_offset + p.ne1 + iq2] = D_TYPE(Mf[r]); 578 } 579 } 580 } 581 582 return; 583 } 584 585 if ((p.mask_n_head_log2 & SINK_ENABLE_BIT) != 0) { 586 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 587 float sink = perElemOpGetSink(tile_row(r), 0u, ACC_TYPE(0), iq2); 588 589 float ms = 1.0f; 590 float vs = 1.0f; 591 592 if (sink > Mf[r]) { 593 ms = exp(Mf[r] - sink); 594 595 [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { 596 const uint d_local = d0 / threads_per_rowgroup; 597 Of[r][d_local] *= O_TYPE(ms); 598 } 599 } else { 600 vs = exp(sink - Mf[r]); 601 } 602 603 Lf[r] = Lf[r]*ms + vs; 604 } 605 } 606 607 float Lfrcp[rows_per_thread]; 608 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 609 Lfrcp[r] = (Lf[r] == 0.0) ? 0.0 : (1.0 / Lf[r]); 610 } 611 612 [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { 613 const uint d_local = d0 / threads_per_rowgroup; 614 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 615 Of[r][d_local] *= O_TYPE(Lfrcp[r]); 616#if defined(FLOAT_TYPE_MAX) 617 Of[r][d_local] = clamp(Of[r][d_local], -FLOAT_TYPE_MAX, FLOAT_TYPE_MAX); 618#endif 619 } 620 } 621 622 uint32_t o_offset = (gqa_iq1*p.ne1*HSV + iq3*p.ne2*p.ne1*HSV) / 4; 623 624 if (p.gqa_ratio > 1) { 625 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 626 if (tile_row(r) < N) { 627 [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { 628 const uint d = d0 + col_tid; 629 if (d >= HSV / 4) break; 630 const uint d_local = d0 / threads_per_rowgroup; 631 gqaStore(tile_row(r), d, Of[r][d_local], o_offset, iq2, N); 632 } 633 } 634 } 635 } else { 636 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { 637 if (i * Br + tile_row(r) < N) { 638 [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { 639 const uint d = d0 + col_tid; 640 if (d >= HSV / 4) break; 641 const uint d_local = d0 / threads_per_rowgroup; 642 data_ov4[o_offset + (iq2 * HSV + (i * Br + tile_row(r)) * p.ne1 * HSV) / 4 + d] = D_TYPEV4(Of[r][d_local]); 643 } 644 } 645 } 646 } 647}