char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
5112b9738
char/aliceai-t5-moe
1#ifndef SRC0 2#define SRC0 src0 3#endif 4#ifndef SRC1 5#define SRC1 src1 6#endif 7 8#ifdef VEC 9#define VEC_SIZE 4 10#define SHMEM_TYPE vec4<f16> 11#define DST_TYPE vec4<f32> 12#define SRC0_TYPE vec4<SRC0_INNER_TYPE> 13#define SRC1_TYPE vec4<SRC1_INNER_TYPE> 14 15fn store_shmem(val: vec4<f16>, idx: u32) { 16 shmem[idx] = val.x; 17 shmem[idx + 1] = val.y; 18 shmem[idx + 2] = val.z; 19 shmem[idx + 3] = val.w; 20} 21#endif // VEC 22 23#ifdef SCALAR 24#define VEC_SIZE 1 25#define SHMEM_TYPE f16 26#define DST_TYPE f32 27#define SRC0_TYPE SRC0_INNER_TYPE 28#define SRC1_TYPE SRC1_INNER_TYPE 29 30fn store_shmem(val: f16, idx: u32) { 31 shmem[idx] = val; 32} 33#endif // SCALAR 34 35#define QUANT_SHMEM shmem 36#define QUANT_OUT_TYPE f16 37#include "quant_inner_loops.tmpl" 38 39#ifdef INIT_SRC0_SHMEM_FLOAT 40fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { 41 for (var elem_idx = thread_id * VEC_SIZE; elem_idx < TILE_SRC0_SHMEM; elem_idx += TOTAL_WORKGROUP_SIZE * VEC_SIZE) { 42 let tile_m = elem_idx / TILE_K; 43 let tile_k = elem_idx % TILE_K; 44 let global_m = offset_m + tile_m; 45 let global_k = k_outer + tile_k; 46 let src0_idx = batch_offset + global_m * params.stride_01 + global_k; 47 let src0_val = select( // taking a slight performance hit to avoid oob 48 SRC0_TYPE(0.0), 49 SRC0[src0_idx/VEC_SIZE], 50 global_m < params.m && global_k < params.k); 51 store_shmem(SHMEM_TYPE(src0_val), elem_idx); 52 } 53} 54#endif // INIT_SRC0_SHMEM_FLOAT 55 56#ifndef MUL_MAT_ID 57#ifdef INIT_SRC1_SHMEM_FLOAT 58fn init_shmem_src1(thread_id: u32, batch_offset: u32, offset_n: u32, k_outer: u32) { 59 for (var elem_idx = thread_id * VEC_SIZE; elem_idx < TILE_SRC1_SHMEM; elem_idx += TOTAL_WORKGROUP_SIZE * VEC_SIZE) { 60 let tile_n = elem_idx / TILE_K; 61 let tile_k = elem_idx % TILE_K; 62 let global_n = offset_n + tile_n; 63 let global_k = k_outer + tile_k; 64 let src1_idx = batch_offset + global_n * params.stride_11 + global_k; 65 let src1_val = select( 66 SRC1_TYPE(0.0), 67 SRC1[src1_idx/VEC_SIZE], 68 global_n < params.n && global_k < params.k); 69 store_shmem(SHMEM_TYPE(src1_val), TILE_SRC0_SHMEM + elem_idx); 70 } 71} 72#endif // INIT_SRC1_SHMEM_FLOAT 73#endif 74 75#ifdef INIT_SRC0_SHMEM_Q1_0 76const BLOCK_SIZE = 128u; 77const BLOCK_SIZE_BYTES = 18u; 78const NQ = 8u; // 8 weights (1 byte of qs) per thread per iteration 79 80fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { 81 for (var i = thread_id * NQ; i < TILE_SRC0_SHMEM; i += TOTAL_WORKGROUP_SIZE * NQ) { 82 let tile_m = i / TILE_K; 83 let tile_k_start = i % TILE_K; 84 let global_m = offset_m + tile_m; 85 let global_k_start = k_outer + tile_k_start; 86 87 if (global_m >= params.m) { 88 break; 89 } 90 91 let block_k = global_k_start / BLOCK_SIZE; 92 let byte_in_block = (global_k_start % BLOCK_SIZE) / 8u; 93 let src0_idx = batch_offset + global_m * params.stride_01 + block_k; 94 let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; 95 let d = load_f16_at_src0(block_byte_base); 96 let q_byte = load_u32_at_src0(block_byte_base + 2u + byte_in_block) & 0xFFu; 97 98 for (var bit = 0u; bit < NQ; bit++) { 99 let global_k = global_k_start + bit; 100 if (global_k < params.k) { 101 shmem[i + bit] = select(-d, d, ((q_byte >> bit) & 1u) != 0u); 102 } 103 } 104 } 105} 106#endif // INIT_SRC0_SHMEM_Q1_0 107 108// legacy-quants 109#if defined(INIT_SRC0_SHMEM_Q4_0) || defined(INIT_SRC0_SHMEM_Q4_1) || defined(INIT_SRC0_SHMEM_Q5_0) || defined(INIT_SRC0_SHMEM_Q5_1) || defined(INIT_SRC0_SHMEM_Q8_0) || defined(INIT_SRC0_SHMEM_Q8_1) || defined(INIT_SRC0_SHMEM_MXFP4) 110const BLOCK_SIZE = 32u; 111// the number of blocks per k-tile. Note that this currently only works if TILE_K is a multiple of BLOCK_SIZE, which may need to be rethought for larger quantized types. 112override BLOCKS_K = TILE_K/BLOCK_SIZE; 113const NQ = 16u; 114#if defined(INIT_SRC0_SHMEM_Q8_0) || defined(INIT_SRC0_SHMEM_Q8_1) 115const BYTES_PER_THREAD = 16u; // NQ(16) weights use 16 bytes of q 116#else 117const BYTES_PER_THREAD = 8u; // NQ(16) weights use 8 bytes of q 118#endif 119const BYTES_PER_INNER_LOOP = 4u; // == sizeof(q_packed) 120 121fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { 122 for (var i = thread_id * NQ; i < TILE_SRC0_SHMEM; i += TOTAL_WORKGROUP_SIZE * NQ) { 123 let block_idx = i / BLOCK_SIZE; 124 let block_offset = (i % BLOCK_SIZE) / NQ; 125 let shmem_idx = block_idx * BLOCK_SIZE + block_offset * BYTES_PER_THREAD; 126 127 let tile_m = block_idx / BLOCKS_K; 128 let global_m = offset_m + tile_m; 129 let block_k = block_idx % BLOCKS_K; 130 let global_block_k = k_outer / BLOCK_SIZE + block_k; 131 132 if (global_m < params.m && global_block_k < params.k / BLOCK_SIZE) { 133 let src0_idx = batch_offset + global_m * params.stride_01 + global_block_k; 134 135#if defined(INIT_SRC0_SHMEM_Q4_0) 136 let block_byte_base = src0_idx * 18u; // BLOCK_SIZE_BYTES = 18u; 137 let d = load_f16_at_src0(block_byte_base); 138 139 // load NQ(16) weights 140 for (var j = 0u; j < BYTES_PER_THREAD / BYTES_PER_INNER_LOOP; j += 1) { 141 let q_byte_offset = block_byte_base + 2u + block_offset * BYTES_PER_THREAD + j * BYTES_PER_INNER_LOOP; 142 let q_packed = load_u32_at_src0(q_byte_offset); 143 dequant_q4_0_packed_to_shmem(q_packed, d, shmem_idx + j * BYTES_PER_INNER_LOOP); 144 } 145#endif // INIT_SRC0_SHMEM_Q4_0 146 147#if defined(INIT_SRC0_SHMEM_Q4_1) 148 let block_byte_base = src0_idx * 20u; // BLOCK_SIZE_BYTES = 20u; 149 let dm = unpack2x16float(load_u32_at_src0_aligned(block_byte_base)); 150 let d = f16(dm[0]); 151 let m = f16(dm[1]); 152 153 // load NQ(16) weights 154 for (var j = 0u; j < BYTES_PER_THREAD / BYTES_PER_INNER_LOOP; j += 1) { 155 let q_byte_offset = block_byte_base + 4u + block_offset * BYTES_PER_THREAD + j * BYTES_PER_INNER_LOOP; 156 let q_packed = load_u32_at_src0(q_byte_offset); 157 158 for (var k = 0u; k < BYTES_PER_INNER_LOOP; k++) { 159 let q_byte = get_byte(q_packed, k); 160 let q_lo = f16(q_byte & 0xF) * d + m; 161 let q_hi = f16((q_byte >> 4) & 0xF) * d + m; 162 shmem[shmem_idx + j * BYTES_PER_INNER_LOOP + k] = q_lo; 163 shmem[shmem_idx + j * BYTES_PER_INNER_LOOP + k + 16u] = q_hi; 164 } 165 } 166#endif // INIT_SRC0_SHMEM_Q4_1 167 168#if defined(INIT_SRC0_SHMEM_Q5_0) 169 let block_byte_base = src0_idx * 22u; // BLOCK_SIZE_BYTES = 22u; 170 171 let d = load_f16_at_src0(block_byte_base); 172 let qh_packed = load_u32_at_src0(block_byte_base + 2u); 173 174 // load NQ(16) weights 175 for (var j = 0u; j < BYTES_PER_THREAD / BYTES_PER_INNER_LOOP; j += 1) { 176 let q_byte_offset = block_byte_base + 6u + block_offset * BYTES_PER_THREAD + j * BYTES_PER_INNER_LOOP; 177 let q_packed = load_u32_at_src0(q_byte_offset); 178 179 for (var k = 0u; k < BYTES_PER_INNER_LOOP; k++) { 180 let q_byte = get_byte(q_packed, k); 181 182 let byte_idx = block_offset * BYTES_PER_THREAD + j * BYTES_PER_INNER_LOOP + k; 183 let qh_hi = (qh_packed >> (byte_idx + 12u)) & 0x10; 184 let q_hi = (f16(((q_byte >> 4) & 0xF) | qh_hi) - 16.0) * d; 185 let qh_lo = ((qh_packed >> byte_idx) << 4) & 0x10; 186 let q_lo = (f16((q_byte & 0xF) | qh_lo) - 16.0) * d; 187 shmem[shmem_idx + j * BYTES_PER_INNER_LOOP + k] = q_lo; 188 shmem[shmem_idx + j * BYTES_PER_INNER_LOOP + k + 16u] = q_hi; 189 } 190 } 191#endif // INIT_SRC0_SHMEM_Q5_0 192 193#if defined(INIT_SRC0_SHMEM_Q5_1) 194 let block_byte_base = src0_idx * 24u; // BLOCK_SIZE_BYTES = 24u; 195 196 let dm = unpack2x16float(load_u32_at_src0_aligned(block_byte_base)); 197 let d = f16(dm[0]); 198 let m = f16(dm[1]); 199 let qh_packed = load_u32_at_src0_aligned(block_byte_base + 4u); 200 201 // load NQ(16) weights 202 for (var j = 0u; j < BYTES_PER_THREAD / BYTES_PER_INNER_LOOP; j += 1) { 203 let q_byte_offset = block_byte_base + 8u + block_offset * BYTES_PER_THREAD + j * BYTES_PER_INNER_LOOP; 204 let q_packed = load_u32_at_src0_aligned(q_byte_offset); 205 206 for (var k = 0u; k < BYTES_PER_INNER_LOOP; k++) { 207 let q_byte = get_byte(q_packed, k); 208 209 let byte_idx = block_offset * BYTES_PER_THREAD + j * BYTES_PER_INNER_LOOP + k; 210 let qh_hi = (qh_packed >> (byte_idx + 12u)) & 0x10; 211 let q_hi = f16(((q_byte >> 4) & 0xF) | qh_hi) * d + m; 212 let qh_lo = ((qh_packed >> byte_idx) << 4) & 0x10; 213 let q_lo = f16((q_byte & 0xF) | qh_lo) * d + m; 214 shmem[shmem_idx + j * BYTES_PER_INNER_LOOP + k] = q_lo; 215 shmem[shmem_idx + j * BYTES_PER_INNER_LOOP + k + 16u] = q_hi; 216 } 217 } 218#endif // INIT_SRC0_SHMEM_Q5_1 219 220#if defined(INIT_SRC0_SHMEM_Q8_0) 221 let block_byte_base = src0_idx * 34u; // BLOCK_SIZE_BYTES = 34u; 222 let d = load_f16_at_src0(block_byte_base); 223 224 // load NQ(16) weights 225 for (var j = 0u; j < BYTES_PER_THREAD / BYTES_PER_INNER_LOOP; j += 1) { 226 let q_byte_offset = block_byte_base + 2u + block_offset * BYTES_PER_THREAD + j * BYTES_PER_INNER_LOOP; 227 let q_packed = load_u32_at_src0(q_byte_offset); 228 dequant_q8_0_packed_to_shmem(q_packed, d, shmem_idx + j * BYTES_PER_INNER_LOOP); 229 } 230#endif // INIT_SRC0_SHMEM_Q8_0 231 232#if defined(INIT_SRC0_SHMEM_Q8_1) 233 let block_byte_base = src0_idx * 36u; // BLOCK_SIZE_BYTES = 36u; 234 let dm = unpack2x16float(load_u32_at_src0_aligned(block_byte_base)); 235 let d = f16(dm[0]); 236 let m = f16(dm[1]); 237 238 // load NQ(16) weights 239 for (var j = 0u; j < BYTES_PER_THREAD / BYTES_PER_INNER_LOOP; j += 1) { 240 let q_byte_offset = block_byte_base + 4u + block_offset * BYTES_PER_THREAD + j * BYTES_PER_INNER_LOOP; 241 let q_packed = load_u32_at_src0(q_byte_offset); 242 for (var k = 0u; k < BYTES_PER_INNER_LOOP; k++) { 243 let q_byte = get_byte_i32(q_packed, k); 244 let q_val = f16(q_byte) * d + m; 245 shmem[shmem_idx + j * BYTES_PER_INNER_LOOP + k] = q_val; 246 } 247 } 248#endif // INIT_SRC0_SHMEM_Q8_1 249 250#if defined(INIT_SRC0_SHMEM_MXFP4) 251 let block_byte_base = src0_idx * 17u; // BLOCK_SIZE_BYTES = 17u; 252 let eu8 = get_byte(load_u32_at_src0_aligned(block_byte_base), block_byte_base & 3u); 253 let e = ldexp(1.0, i32(eu8) - 128); 254 255 // load NQ(16) weights 256 for (var j = 0u; j < BYTES_PER_THREAD / BYTES_PER_INNER_LOOP; j += 1) { 257 let q_byte_offset = block_byte_base + 1u + block_offset * BYTES_PER_THREAD + j * BYTES_PER_INNER_LOOP; 258 let q_packed = load_u32_at_src0(q_byte_offset); 259 for (var k = 0u; k < BYTES_PER_INNER_LOOP; k++) { 260 let q_byte = get_byte(q_packed, k); 261 let q_hi = f32(kvalues_mxfp4[(q_byte >> 4) & 0xF]) * e; 262 let q_lo = f32(kvalues_mxfp4[q_byte & 0xF]) * e; 263 shmem[shmem_idx + j * BYTES_PER_INNER_LOOP + k] = f16(q_lo); 264 shmem[shmem_idx + j * BYTES_PER_INNER_LOOP + k + 16u] = f16(q_hi); 265 } 266 } 267#endif // INIT_SRC0_SHMEM_MXFP4 268 } 269 } 270} 271#endif // legacy-quants 272 273#if defined(INIT_SRC0_SHMEM_NVFP4) 274const BLOCK_SIZE = 64u; 275const BLOCK_SIZE_BYTES = 36u; 276const SUB_BLOCK_SIZE = 16u; // elements sharing one UE4M3 scale 277const NQ = 16u; 278const BYTES_PER_THREAD = 8u; 279const BYTES_PER_INNER_LOOP = 4u; 280 281fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { 282 for (var i = thread_id * NQ; i < TILE_SRC0_SHMEM; i += TOTAL_WORKGROUP_SIZE * NQ) { 283 let tile_m = i / TILE_K; 284 let tile_k_start = i % TILE_K; 285 let global_m = offset_m + tile_m; 286 let global_k_start = k_outer + tile_k_start; 287 288 if (global_m >= params.m) { 289 break; 290 } 291 292 let block_k = global_k_start / BLOCK_SIZE; 293 let sub_block = (global_k_start % BLOCK_SIZE) / SUB_BLOCK_SIZE; 294 let src0_idx = batch_offset + global_m * params.stride_01 + block_k; 295 296 let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; 297 let d_byte_base = block_byte_base; 298 let qs_byte_base = block_byte_base + 4u; 299 300 let d = ue4m3_to_fp32(get_byte(load_u32_at_src0_aligned(d_byte_base), sub_block)) * 0.5; 301 302 for (var j = 0u; j < BYTES_PER_THREAD / BYTES_PER_INNER_LOOP; j++) { 303 let q_packed = load_u32_at_src0_aligned(qs_byte_base + sub_block * 8u + j * 4u); 304 for (var k = 0u; k < BYTES_PER_INNER_LOOP; k++) { 305 let q_byte = get_byte(q_packed, k); 306 shmem[i + j * BYTES_PER_INNER_LOOP + k] = f16(f32(kvalues_mxfp4[q_byte & 0xF]) * d); 307 shmem[i + j * BYTES_PER_INNER_LOOP + k + 8u] = f16(f32(kvalues_mxfp4[(q_byte >> 4) & 0xF]) * d); 308 } 309 } 310 } 311} 312#endif // INIT_SRC0_SHMEM_NVFP4 313 314// k-quants 315#if defined(INIT_SRC0_SHMEM_Q2_K) || defined(INIT_SRC0_SHMEM_Q3_K) || defined(INIT_SRC0_SHMEM_Q4_K) || defined(INIT_SRC0_SHMEM_Q5_K) || defined(INIT_SRC0_SHMEM_Q6_K) 316const BLOCK_SIZE = 256u; 317const NQ = 4u; 318 319fn store_shmem_kquants(val: vec4<f16>, idx: u32) { 320 shmem[idx] = val.x; 321 shmem[idx + 1] = val.y; 322 shmem[idx + 2] = val.z; 323 shmem[idx + 3] = val.w; 324} 325 326fn load_byte_at_src0_aligned(byte_offset: u32) -> u32 { 327 return get_byte(load_u32_at_src0_aligned(byte_offset), byte_offset % 4u); 328} 329 330fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { 331 for (var elem_idx = thread_id * NQ; elem_idx < TILE_SRC0_SHMEM; elem_idx += TOTAL_WORKGROUP_SIZE * NQ) { 332 let tile_m = elem_idx / TILE_K; 333 let tile_k = elem_idx % TILE_K; 334 335 let global_m = offset_m + tile_m; 336 let global_k = k_outer + tile_k; 337 338 if (global_m >= params.m || global_k >= params.k) { 339 store_shmem_kquants(vec4<f16>(f16(0.0), f16(0.0), f16(0.0), f16(0.0)), elem_idx); 340 continue; 341 } 342 343 let block_k = global_k / BLOCK_SIZE; 344 let k_in_block = global_k % BLOCK_SIZE; // k_in_block % 4 == 0; 345 346 let src0_idx = batch_offset + global_m * params.stride_01 + block_k; 347 348#if defined(INIT_SRC0_SHMEM_Q2_K) 349 let block_byte_base = src0_idx * 84u; // BLOCK_SIZE_BYTES = 84u; 350 let scales_byte_base = block_byte_base; 351 let qs_byte_base = block_byte_base + 16u; 352 let dm_byte_base = block_byte_base + 80u; 353 354 let d_packed = unpack2x16float(load_u32_at_src0_aligned(dm_byte_base)); 355 let d = f16(d_packed[0]); 356 let dmin = f16(d_packed[1]); 357 358 let chunk = k_in_block / 128u; 359 let pos_in_chunk = k_in_block % 32u; 360 let sub_block = k_in_block / 16u; 361 let shift_phase = (k_in_block % 128u) / 32u; 362 363 // whole 2 bits (4 elems) 364 let qs_word = load_u32_at_src0_aligned(qs_byte_base + 32u * chunk + 1u * pos_in_chunk); 365 let qs_vec4 = vec4<f16>( 366 f16((qs_word >> (2u * shift_phase + 0u)) & 0x3u), 367 f16((qs_word >> (2u * shift_phase + 8u)) & 0x3u), 368 f16((qs_word >> (2u * shift_phase + 16u)) & 0x3u), 369 f16((qs_word >> (2u * shift_phase + 24u)) & 0x3u), 370 ); 371 372 let scale = load_byte_at_src0_aligned(scales_byte_base + sub_block); 373 374 let dl = d * f16(scale & 0xFu); 375 let ml = dmin * f16(scale >> 4u); 376 377 store_shmem_kquants(qs_vec4 * dl - ml, elem_idx); 378#endif // INIT_SRC0_SHMEM_Q2_K 379 380#if defined(INIT_SRC0_SHMEM_Q3_K) 381 let block_byte_base = src0_idx * 110u; // BLOCK_SIZE_BYTES = 110u; 382 let hmask_byte_base = block_byte_base + 0u; 383 let qs_byte_base = block_byte_base + 32u; 384 let scales_byte_base = block_byte_base + 96u; 385 386 let d_all = load_f16_at_src0(block_byte_base + 108u); 387 388 let chunk = k_in_block / 128u; 389 let pos_in_chunk = k_in_block % 32u; 390 let sub_block = k_in_block / 16u; 391 let shift_phase = (k_in_block % 128u) / 32u; 392 393 let hmask_block = pos_in_chunk; 394 let hmask_shift_phase = k_in_block / 32u; 395 396 // low 2 bits (4 elems) 397 let q_lo2_word = load_u32_at_src0(qs_byte_base + 32u * chunk + 1u * hmask_block); 398 let q_lo2_vec4 = vec4<f16>( 399 f16((q_lo2_word >> (2u * shift_phase + 0u)) & 3u), 400 f16((q_lo2_word >> (2u * shift_phase + 8u)) & 3u), 401 f16((q_lo2_word >> (2u * shift_phase + 16u)) & 3u), 402 f16((q_lo2_word >> (2u * shift_phase + 24u)) & 3u) 403 ); 404 405 // high 1 bit (4 elems) 406 let q_hi1_word = load_u32_at_src0(hmask_byte_base + pos_in_chunk); 407 let q_hi1_vec4 = vec4<f16>( 408 f16(select(4.0, 0.0, ((q_hi1_word >> (1u * hmask_shift_phase + 0u)) & 1u) == 1u)), 409 f16(select(4.0, 0.0, ((q_hi1_word >> (1u * hmask_shift_phase + 8u)) & 1u) == 1u)), 410 f16(select(4.0, 0.0, ((q_hi1_word >> (1u * hmask_shift_phase + 16u)) & 1u) == 1u)), 411 f16(select(4.0, 0.0, ((q_hi1_word >> (1u * hmask_shift_phase + 24u)) & 1u) == 1u)) 412 ); 413 414 let q_vec4 = q_lo2_vec4 - q_hi1_vec4; 415 416 let scale_low4 = (load_byte_at_src0_aligned(scales_byte_base + (sub_block % 8u)) >> (4u * (sub_block / 8u))) & 0xFu; 417 let scale_hi2 = (load_byte_at_src0_aligned(scales_byte_base + 8u + (sub_block % 4u)) >> (2u * (sub_block / 4u))) & 3u; 418 let dl = d_all * (f16((scale_hi2 << 4u) | scale_low4) - 32.0); 419 420 store_shmem_kquants(dl * q_vec4, elem_idx); 421#endif // INIT_SRC0_SHMEM_Q3_K 422 423#if defined(INIT_SRC0_SHMEM_Q4_K) 424 let block_byte_base = src0_idx * 144u; // BLOCK_SIZE_BYTES = 144u; 425 let dm_byte_base = block_byte_base + 0u; 426 let scale_byte_base = block_byte_base + 4u; 427 let qs_byte_base = block_byte_base + 16u; 428 429 let dm = unpack2x16float(load_u32_at_src0_aligned(dm_byte_base)); 430 let d = f16(dm[0]); 431 let dmin = f16(dm[1]); 432 433 let chunk = k_in_block / 64u; 434 let pos_in_chunk = (k_in_block % 64u) % 32u; 435 let sub_block = k_in_block / 32u; 436 let shift_phase = sub_block & 1u; 437 438 // whole 4 bits (4 elems) 439 let qs_word = load_u32_at_src0_aligned(qs_byte_base + 32u * chunk + 1u * pos_in_chunk); 440 let qs_vec4 = vec4<f16>( 441 f16((qs_word >> (4u * shift_phase + 0u)) & 0xFu), 442 f16((qs_word >> (4u * shift_phase + 8u)) & 0xFu), 443 f16((qs_word >> (4u * shift_phase + 16u)) & 0xFu), 444 f16((qs_word >> (4u * shift_phase + 24u)) & 0xFu) 445 ); 446 447 var sc: u32; 448 var mn: u32; 449 450 if (sub_block < 4u) { 451 let sc_byte = get_byte(load_u32_at_src0_aligned(scale_byte_base), sub_block % 4u); 452 let min_byte = get_byte(load_u32_at_src0_aligned(scale_byte_base + 4), sub_block % 4u); 453 sc = sc_byte & 63u; 454 mn = min_byte & 63u; 455 } else { 456 let sc_min_lo = get_byte(load_u32_at_src0_aligned(scale_byte_base + 8), (sub_block + 4u) % 4u); 457 let sc_hi = get_byte(load_u32_at_src0_aligned(scale_byte_base), (sub_block - 4u) % 4u); 458 let min_hi = get_byte(load_u32_at_src0_aligned(scale_byte_base + 4), sub_block % 4u); 459 sc = (sc_min_lo & 0xFu) | ((sc_hi >> 6u) << 4u); 460 mn = (sc_min_lo >> 4u) | ((min_hi >> 6u) << 4u); 461 } 462 463 let dl = d * f16(sc); 464 let ml = dmin * f16(mn); 465 466 store_shmem_kquants(dl * qs_vec4 - vec4(ml, ml, ml, ml), elem_idx); 467#endif // INIT_SRC0_SHMEM_Q4_K 468 469#if defined(INIT_SRC0_SHMEM_Q5_K) 470 let block_byte_base = src0_idx * 176u; // BLOCK_SIZE_BYTES = 176u; 471 let dm_byte_base = block_byte_base + 0u; 472 let scale_byte_base = block_byte_base + 4u; 473 let qh_byte_base = block_byte_base + 16u; 474 let qs_byte_base = block_byte_base + 48u; 475 476 let dm = unpack2x16float(load_u32_at_src0_aligned(dm_byte_base)); 477 let d = f16(dm[0]); 478 let dmin = f16(dm[1]); 479 480 let chunk = k_in_block / 64u; 481 let pos_in_chunk = (k_in_block % 64u) % 32u; 482 let sub_block = k_in_block / 32u; 483 let shift_phase = sub_block & 1u; 484 485 let qh_block = k_in_block % 32u; 486 let qh_shift_phase = sub_block; 487 488 // low 4 bits (4 elems) 489 let qs_word = load_u32_at_src0_aligned(qs_byte_base + 32u * chunk + 1u * pos_in_chunk); 490 let qs_lo4_vec4 = vec4<f16>( 491 f16((qs_word >> (4u * shift_phase + 0u)) & 0xFu), 492 f16((qs_word >> (4u * shift_phase + 8u)) & 0xFu), 493 f16((qs_word >> (4u * shift_phase + 16u)) & 0xFu), 494 f16((qs_word >> (4u * shift_phase + 24u)) & 0xFu) 495 ); 496 497 // high 1 bit (4 elems) 498 let qh_word = load_u32_at_src0_aligned(qh_byte_base + qh_block); 499 let qh_vec4 = vec4<f16>( 500 f16(select(0.0, 16.0, ((qh_word >> (1u * qh_shift_phase + 0u)) & 1u) == 1u)), 501 f16(select(0.0, 16.0, ((qh_word >> (1u * qh_shift_phase + 8u)) & 1u) == 1u)), 502 f16(select(0.0, 16.0, ((qh_word >> (1u * qh_shift_phase + 16u)) & 1u) == 1u)), 503 f16(select(0.0, 16.0, ((qh_word >> (1u * qh_shift_phase + 24u)) & 1u) == 1u)) 504 ); 505 506 var sc: u32; 507 var mn: u32; 508 509 if (sub_block < 4u) { 510 let sc_byte = get_byte(load_u32_at_src0_aligned(scale_byte_base), sub_block % 4u); 511 let min_byte = get_byte(load_u32_at_src0_aligned(scale_byte_base + 4), sub_block % 4u); 512 sc = sc_byte & 63u; 513 mn = min_byte & 63u; 514 } else { 515 let sc_min_lo = get_byte(load_u32_at_src0_aligned(scale_byte_base + 8), (sub_block + 4u) % 4u); 516 let sc_hi = get_byte(load_u32_at_src0_aligned(scale_byte_base), (sub_block - 4u) % 4u); 517 let min_hi = get_byte(load_u32_at_src0_aligned(scale_byte_base + 4), sub_block % 4u); 518 sc = (sc_min_lo & 0xFu) | ((sc_hi >> 6u) << 4u); 519 mn = (sc_min_lo >> 4u) | ((min_hi >> 6u) << 4u); 520 } 521 522 let dl = d * f16(sc); 523 let ml = dmin * f16(mn); 524 525 store_shmem_kquants((qh_vec4 + qs_lo4_vec4) * dl - vec4<f16>(ml, ml, ml, ml), elem_idx); 526#endif // INIT_SRC0_SHMEM_Q5_K 527 528#if defined(INIT_SRC0_SHMEM_Q6_K) 529 let block_byte_base = src0_idx * 210u; // BLOCK_SIZE_BYTES = 210u; 530 let ql_byte_base = block_byte_base; 531 let qh_byte_base = block_byte_base + 128u; 532 let scales_byte_base = block_byte_base + 192u; 533 let d_byte_base = block_byte_base + 208u; 534 535 let d = load_f16_at_src0(d_byte_base); 536 537 let chunk = k_in_block / 128u; 538 let ql_pos_in_chunk = (k_in_block % 128u) % 64u; 539 let qh_pos_in_chunk = (k_in_block % 128u) % 32u; 540 let sub_block = k_in_block / 16u; 541 let ql_shift_phase = (k_in_block % 128u) / 64u; 542 let qh_shift_phase = (k_in_block % 128u) / 32u; 543 544 // low 4 bits (4 elems) 545 let ql_word = load_u32_at_src0(ql_byte_base + 64u * chunk + 1u * ql_pos_in_chunk); 546 let ql_lo4_vec4 = vec4<u32>( 547 (ql_word >> (4u * ql_shift_phase + 0u)) & 0xFu, 548 (ql_word >> (4u * ql_shift_phase + 8u)) & 0xFu, 549 (ql_word >> (4u * ql_shift_phase + 16u)) & 0xFu, 550 (ql_word >> (4u * ql_shift_phase + 24u)) & 0xFu 551 ); 552 553 // hi 2 bits (4 elems) 554 let qh_word = load_u32_at_src0(qh_byte_base + 32u * chunk + 1u * qh_pos_in_chunk); 555 let qh_hi2_vec4 = vec4<u32>( 556 ((qh_word >> (2u * qh_shift_phase + 0u)) & 0x3u) << 4u, 557 ((qh_word >> (2u * qh_shift_phase + 8u)) & 0x3u) << 4u, 558 ((qh_word >> (2u * qh_shift_phase + 16u)) & 0x3u) << 4u, 559 ((qh_word >> (2u * qh_shift_phase + 24u)) & 0x3u) << 4u, 560 ); 561 562 let q_vec4 = vec4<f16>(qh_hi2_vec4 | ql_lo4_vec4) - vec4<f16>(32.0, 32.0, 32.0, 32.0); 563 564 let scale_byte = scales_byte_base + 1u * sub_block; 565 let scale_word = load_u32_at_src0_aligned(scale_byte); 566 let scale = get_byte_i32(scale_word, scale_byte & 3u); 567 568 store_shmem_kquants(d * q_vec4 * f16(scale), elem_idx); 569#endif // INIT_SRC0_SHMEM_Q6_K 570 } 571} 572#endif // k-quants 573 574#if defined(INIT_SRC0_SHMEM_IQ4_NL) 575const BLOCK_SIZE = 32u; 576const BLOCK_SIZE_BYTES = 18u; 577const NQ = 4u; 578 579fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { 580 for (var elem_idx = thread_id * NQ; elem_idx < TILE_SRC0_SHMEM; elem_idx += NQ * TOTAL_WORKGROUP_SIZE) { 581 let tile_m = elem_idx / TILE_K; 582 let tile_k = elem_idx % TILE_K; 583 let global_m = offset_m + tile_m; 584 let global_k = k_outer + tile_k; 585 586 if (global_m >= params.m || global_k >= params.k) { 587 shmem[elem_idx] = f16(0.0); 588 continue; 589 } 590 591 let block_k = global_k / BLOCK_SIZE; 592 let k_in_block = global_k % BLOCK_SIZE; // k_in_block % 4 == 0; 593 594 let src0_idx = batch_offset + global_m * params.stride_01 + block_k; 595 596 let block_byte_base = src0_idx * BLOCK_SIZE_BYTES; 597 let d_byte_base = block_byte_base + 0u; 598 let qs_byte_base = block_byte_base + 2u; 599 600 let d = load_f16_at_src0(d_byte_base); 601 602 let id_qtr = (k_in_block % 16u) / 4u; 603 let shift_phase = k_in_block / 16u; 604 605 let qs_u32 = load_u32_at_src0(qs_byte_base + 4u * id_qtr); 606 607 shmem[elem_idx + 0u] = d * f16(kvalues_iq4nl[(qs_u32 >> ( 0u + 4u * shift_phase)) & 0xFu]); 608 shmem[elem_idx + 1u] = d * f16(kvalues_iq4nl[(qs_u32 >> ( 8u + 4u * shift_phase)) & 0xFu]); 609 shmem[elem_idx + 2u] = d * f16(kvalues_iq4nl[(qs_u32 >> (16u + 4u * shift_phase)) & 0xFu]); 610 shmem[elem_idx + 3u] = d * f16(kvalues_iq4nl[(qs_u32 >> (24u + 4u * shift_phase)) & 0xFu]); 611 } 612} 613#endif // INIT_SRC0_SHMEM_IQ4_NL 614 615// i-quants (super block size: 256) 616#if defined(INIT_SRC0_SHMEM_IQ4_XS) || defined(INIT_SRC0_SHMEM_IQ1_S) || defined(INIT_SRC0_SHMEM_IQ1_M) || defined(INIT_SRC0_SHMEM_IQ2_XXS) \ 617|| defined(INIT_SRC0_SHMEM_IQ2_XS) || defined(INIT_SRC0_SHMEM_IQ2_S) || defined(INIT_SRC0_SHMEM_IQ3_XXS) || defined(INIT_SRC0_SHMEM_IQ3_S) 618const BLOCK_SIZE = 256u; 619const NQ = 16u; 620 621fn store_shmem_iquants(val: vec4<f16>, idx: u32) { 622 shmem[idx] = val.x; 623 shmem[idx + 1] = val.y; 624 shmem[idx + 2] = val.z; 625 shmem[idx + 3] = val.w; 626} 627 628fn load_byte_at_src0_aligned(byte_offset: u32) -> u32 { 629 return get_byte(load_u32_at_src0_aligned(byte_offset), byte_offset % 4u); 630} 631 632#if defined(INIT_SRC0_SHMEM_IQ1_M) || defined(INIT_SRC0_SHMEM_IQ1_S) 633fn create_iq_gw4(dl: f32, gw: u32, shift_base: u32, delta: f32) -> vec4<f16> { 634 return vec4<f16>( 635 f16(dl * (f32((bitcast<i32>(((gw >> (shift_base + 0u)) & 3u) << 30u) >> 30u)) + delta)), 636 f16(dl * (f32((bitcast<i32>(((gw >> (shift_base + 2u)) & 3u) << 30u) >> 30u)) + delta)), 637 f16(dl * (f32((bitcast<i32>(((gw >> (shift_base + 4u)) & 3u) << 30u) >> 30u)) + delta)), 638 f16(dl * (f32((bitcast<i32>(((gw >> (shift_base + 6u)) & 3u) << 30u) >> 30u)) + delta)), 639 ); 640} 641#endif 642 643#if defined(INIT_SRC0_SHMEM_IQ4_XS) 644fn create_iq_gw4(dl: f16, qs_u32: u32, shift_phase: u32) -> vec4<f16> { 645 return vec4<f16>( 646 dl * f16(kvalues_iq4nl[(qs_u32 >> (4 * shift_phase + 0u)) & 0xFu]), 647 dl * f16(kvalues_iq4nl[(qs_u32 >> (4 * shift_phase + 8u)) & 0xFu]), 648 dl * f16(kvalues_iq4nl[(qs_u32 >> (4 * shift_phase + 16u)) & 0xFu]), 649 dl * f16(kvalues_iq4nl[(qs_u32 >> (4 * shift_phase + 24u)) & 0xFu]), 650 ); 651} 652#endif 653 654#if defined(INIT_SRC0_SHMEM_IQ2_XXS) 655fn create_iq_gw4(ig: u32, grid_phase: u32) -> vec4<f32> { 656 return vec4<f32>( 657 f32(get_byte(iq2xxs_grid[(ig + grid_phase + 0u) / 4u], (ig + grid_phase + 0u) % 4u)), 658 f32(get_byte(iq2xxs_grid[(ig + grid_phase + 1u) / 4u], (ig + grid_phase + 1u) % 4u)), 659 f32(get_byte(iq2xxs_grid[(ig + grid_phase + 2u) / 4u], (ig + grid_phase + 2u) % 4u)), 660 f32(get_byte(iq2xxs_grid[(ig + grid_phase + 3u) / 4u], (ig + grid_phase + 3u) % 4u)), 661 ); 662} 663#endif 664 665#if defined(INIT_SRC0_SHMEM_IQ2_XS) 666fn create_iq_gw4(ig: u32, grid_phase: u32) -> vec4<f32> { 667 return vec4<f32>( 668 f32(get_byte(iq2xs_grid[(ig + grid_phase + 0u) / 4u], (ig + grid_phase + 0u) % 4u)), 669 f32(get_byte(iq2xs_grid[(ig + grid_phase + 1u) / 4u], (ig + grid_phase + 1u) % 4u)), 670 f32(get_byte(iq2xs_grid[(ig + grid_phase + 2u) / 4u], (ig + grid_phase + 2u) % 4u)), 671 f32(get_byte(iq2xs_grid[(ig + grid_phase + 3u) / 4u], (ig + grid_phase + 3u) % 4u)), 672 ); 673} 674#endif 675 676#if defined(INIT_SRC0_SHMEM_IQ2_S) 677fn create_iq_gw4(ig: u32, grid_phase: u32) -> vec4<f32> { 678 return vec4<f32>( 679 f32(get_byte(iq2s_grid[(ig + grid_phase + 0u) / 4u], (ig + grid_phase + 0u) % 4u)), 680 f32(get_byte(iq2s_grid[(ig + grid_phase + 1u) / 4u], (ig + grid_phase + 1u) % 4u)), 681 f32(get_byte(iq2s_grid[(ig + grid_phase + 2u) / 4u], (ig + grid_phase + 2u) % 4u)), 682 f32(get_byte(iq2s_grid[(ig + grid_phase + 3u) / 4u], (ig + grid_phase + 3u) % 4u)), 683 ); 684} 685#endif 686 687#if defined(INIT_SRC0_SHMEM_IQ3_XXS) 688fn create_iq_gw4(ig: u32) -> vec4<f32> { 689 return vec4<f32>( 690 f32(get_byte(iq3xxs_grid[ig], 0)), 691 f32(get_byte(iq3xxs_grid[ig], 1)), 692 f32(get_byte(iq3xxs_grid[ig], 2)), 693 f32(get_byte(iq3xxs_grid[ig], 3)), 694 ); 695} 696#endif 697 698#if defined(INIT_SRC0_SHMEM_IQ3_S) 699fn create_iq_gw4(ig: u32) -> vec4<f32> { 700 return vec4<f32>( 701 f32(get_byte(iq3s_grid[ig], 0)), 702 f32(get_byte(iq3s_grid[ig], 1)), 703 f32(get_byte(iq3s_grid[ig], 2)), 704 f32(get_byte(iq3s_grid[ig], 3)), 705 ); 706} 707#endif 708 709#if defined(INIT_SRC0_SHMEM_IQ2_XXS) || defined(INIT_SRC0_SHMEM_IQ2_XS) || defined(INIT_SRC0_SHMEM_IQ2_S) \ 710|| defined(INIT_SRC0_SHMEM_IQ3_XXS) || defined(INIT_SRC0_SHMEM_IQ3_S) 711fn create_iq2_m4(signs: u32, mask_phase: u32) -> vec4<f32> { 712 return vec4<f32>( 713 select(1.0, -1.0, (get_byte(kmask_iq2xs[mask_phase], 0) & signs) != 0u), 714 select(1.0, -1.0, (get_byte(kmask_iq2xs[mask_phase], 1) & signs) != 0u), 715 select(1.0, -1.0, (get_byte(kmask_iq2xs[mask_phase], 2) & signs) != 0u), 716 select(1.0, -1.0, (get_byte(kmask_iq2xs[mask_phase], 3) & signs) != 0u), 717 ); 718} 719#endif 720 721fn init_shmem_src0(thread_id: u32, batch_offset: u32, offset_m: u32, k_outer: u32) { 722 for (var elem_idx = thread_id * NQ; elem_idx < TILE_SRC0_SHMEM; elem_idx += NQ * TOTAL_WORKGROUP_SIZE) { 723 let tile_m = elem_idx / TILE_K; 724 let tile_k = elem_idx % TILE_K; 725 let global_m = offset_m + tile_m; 726 let global_k = k_outer + tile_k; 727 728 if (global_m >= params.m || global_k >= params.k) { 729 let zero_vec4 = vec4<f16>(f16(0.0), f16(0.0), f16(0.0), f16(0.0)); 730 store_shmem_iquants(zero_vec4, elem_idx + 0u); 731 store_shmem_iquants(zero_vec4, elem_idx + 4u); 732 store_shmem_iquants(zero_vec4, elem_idx + 8u); 733 store_shmem_iquants(zero_vec4, elem_idx + 12u); 734 continue; 735 } 736 737 let block_k = global_k / BLOCK_SIZE; 738 let k_in_block = global_k % BLOCK_SIZE; // k_in_block % 16 == 0; 739 740 let src0_idx = batch_offset + global_m * params.stride_01 + block_k; 741 742#if defined(INIT_SRC0_SHMEM_IQ4_XS) 743 let block_byte_base = src0_idx * 136u; // BLOCK_SIZE_BYTES = 136u; 744 let d_byte_base = block_byte_base + 0u; 745 let scales_l_byte_base = block_byte_base + 4u; 746 let qs_byte_base = block_byte_base + 8u; 747 748 let d_scales_h = load_u32_at_src0_aligned(d_byte_base); 749 let d = bitcast<vec2<f16>>(d_scales_h).x; 750 let scales_h = d_scales_h >> 16u; 751 752 let sub_block = k_in_block / 32u; 753 let phase = (k_in_block / NQ) % 2u; 754 755 let scales_l_u32 = load_u32_at_src0_aligned(scales_l_byte_base); 756 let ls_lo = (get_byte(scales_l_u32, sub_block / 2u) >> (4u * (sub_block % 2u))) & 0xFu; 757 let ls_hi = ((scales_h >> (2u * sub_block)) & 3u) << 4u; 758 let dl = d * f16(i32(ls_lo | ls_hi) - 32); 759 760 let qs_0_3_u32 = load_u32_at_src0_aligned(qs_byte_base + 16u * sub_block + 0u); 761 let qs_4_7_u32 = load_u32_at_src0_aligned(qs_byte_base + 16u * sub_block + 4u); 762 let qs_8_11_u32 = load_u32_at_src0_aligned(qs_byte_base + 16u * sub_block + 8u); 763 let qs_12_15_u32 = load_u32_at_src0_aligned(qs_byte_base + 16u * sub_block + 12u); 764 765 store_shmem_iquants(create_iq_gw4(dl, qs_0_3_u32, phase), elem_idx + 0u); 766 store_shmem_iquants(create_iq_gw4(dl, qs_4_7_u32, phase), elem_idx + 4u); 767 store_shmem_iquants(create_iq_gw4(dl, qs_8_11_u32, phase), elem_idx + 8u); 768 store_shmem_iquants(create_iq_gw4(dl, qs_12_15_u32, phase), elem_idx + 12u); 769#endif // INIT_SRC0_SHMEM_IQ4_XS 770 771#if defined(INIT_SRC0_SHMEM_IQ1_S) 772 let block_byte_base = src0_idx * 50u; // BLOCK_SIZE_BYTES = 50u; 773 let d_byte_base = block_byte_base + 0u; 774 let qs_byte_base = block_byte_base + 2u; 775 let qh_byte_base = block_byte_base + 34u; 776 777 let d = load_f16_as_f32_at_src0(d_byte_base); 778 779 let sub_block = k_in_block / 32u; 780 let phase = (k_in_block / NQ) % 2u; 781 782 let qh_u16 = load_u32_at_src0(qh_byte_base + sub_block * 2u) & 0xFFFFu; 783 let qs_u16 = load_u32_at_src0(qs_byte_base + sub_block * 4u + phase * 2u) & 0xFFFFu; 784 785 let dl = d * (2.0 * f32((qh_u16 >> 12u) & 7u) + 1.0); 786 let delta = select(IQ1_DELTA, -IQ1_DELTA, (qh_u16 & 0x8000u) != 0u); 787 788 let gp0_grid_id = ((qs_u16 & 0xFFu) | (((qh_u16 >> (phase * 6u)) & 7u) << 8u)) * 8u; 789 let gp1_grid_id = (((qs_u16 >> 8) & 0xFFu) | (((qh_u16 >> (phase * 6u + 3u)) & 7u) << 8u)) * 8u; 790 791 let gp0_gw = iq1_grid[(gp0_grid_id) / 16u]; 792 let gp1_gw = iq1_grid[(gp1_grid_id) / 16u]; 793 794 let gp0_shift_base = (gp0_grid_id % 16u) * 2u; 795 let gp1_shift_base = (gp1_grid_id % 16u) * 2u; 796 797 store_shmem_iquants(create_iq_gw4(dl, gp0_gw, gp0_shift_base + 0u, delta), elem_idx + 0u); 798 store_shmem_iquants(create_iq_gw4(dl, gp0_gw, gp0_shift_base + 8u, delta), elem_idx + 4u); 799 store_shmem_iquants(create_iq_gw4(dl, gp1_gw, gp1_shift_base + 0u, delta), elem_idx + 8u); 800 store_shmem_iquants(create_iq_gw4(dl, gp1_gw, gp1_shift_base + 8u, delta), elem_idx + 12u); 801#endif // INIT_SRC0_SHMEM_IQ1_S 802 803#if defined(INIT_SRC0_SHMEM_IQ1_M) 804 let block_byte_base = src0_idx * 56u; // BLOCK_SIZE_BYTES = 56u; 805 let qs_byte_base = block_byte_base + 0u; 806 let qh_byte_base = block_byte_base + 32u; 807 let scales_byte_base = block_byte_base + 48u; 808 809 let scales0 = load_u32_at_src0_aligned(scales_byte_base); 810 let scales1 = load_u32_at_src0_aligned(scales_byte_base + 4u); 811 let scale_packed = ((scales0 >> 12u) & 0xFu) | 812 ((scales0 >> 24u) & 0x00F0u) | 813 ((scales1 >> 4u) & 0x0F00u) | 814 ((scales1 >> 16u) & 0xF000u); 815 let d = f32(bitcast<vec2<f16>>(scale_packed).x); 816 817 let sub_block = k_in_block / 32u; 818 let phase = (k_in_block / NQ) % 2u; 819 820 let scale_u32 = select(scales0, scales1, sub_block >= 4u); 821 let scale_u3 = (scale_u32 >> (16u * ((sub_block / 2u) % 2u) + 6u * (sub_block % 2u) + 3u * phase)) & 0x7u; 822 let dl = d * f32(2u * scale_u3 + 1u); 823 824 let qh_u8 = (load_u32_at_src0_aligned(qh_byte_base + 4u * (sub_block / 2u)) >> (16u * (sub_block % 2u) + 8u * phase)) & 0xFFu; 825 let qs_u16 = (load_u32_at_src0_aligned(qs_byte_base + 4u * sub_block) >> (16u * phase)) & 0xFFFFu; 826 827 let gp0_grid_id = ((qs_u16 & 0xFFu) | ((qh_u8 & 7u) << 8u)) * 8u; 828 let gp0_delta = select(IQ1_DELTA, -IQ1_DELTA, (qh_u8 & 0x8u) != 0u); 829 830 let gp1_grid_id = (((qs_u16 >> 8u) & 0xFFu) | (((qh_u8 >> 4u) & 7u) << 8u)) * 8u; 831 let gp1_delta = select(IQ1_DELTA, -IQ1_DELTA, (qh_u8 & 0x80u) != 0u); 832 833 let gp0_gw = iq1_grid[(gp0_grid_id) / 16u]; 834 let gp1_gw = iq1_grid[(gp1_grid_id) / 16u]; 835 836 let gp0_shift_base = (gp0_grid_id % 16u) * 2u; 837 let gp1_shift_base = (gp1_grid_id % 16u) * 2u; 838 839 store_shmem_iquants(create_iq_gw4(dl, gp0_gw, gp0_shift_base + 0u, gp0_delta), elem_idx + 0u); 840 store_shmem_iquants(create_iq_gw4(dl, gp0_gw, gp0_shift_base + 8u, gp0_delta), elem_idx + 4u); 841 store_shmem_iquants(create_iq_gw4(dl, gp1_gw, gp1_shift_base + 0u, gp1_delta), elem_idx + 8u); 842 store_shmem_iquants(create_iq_gw4(dl, gp1_gw, gp1_shift_base + 8u, gp1_delta), elem_idx + 12u); 843#endif // INIT_SRC0_SHMEM_IQ1_M 844 845#if defined(INIT_SRC0_SHMEM_IQ2_XXS) 846 let block_byte_base = src0_idx * 66u; // BLOCK_SIZE_BYTES = 66u; 847 let d_byte_base = block_byte_base + 0u; 848 let qs_byte_base = block_byte_base + 2u; 849 850 let d = load_f16_as_f32_at_src0(d_byte_base); 851 852 let sub_block = k_in_block / 32u; 853 let phase = (k_in_block / NQ) % 2u; 854 855 let aux0 = load_u32_at_src0(qs_byte_base + 8u * sub_block + 0u); 856 let aux1 = load_u32_at_src0(qs_byte_base + 8u * sub_block + 4u); 857 let db = d * (0.5 + f32(aux1 >> 28u)) * 0.25; 858 859 let gp0_ig = get_byte(aux0, 2u * phase + 0u) * 8u; 860 let gp1_ig = get_byte(aux0, 2u * phase + 1u) * 8u; 861 862 let gp0_is = (aux1 >> (14u * phase + 0u)) & 127u; 863 let gp1_is = (aux1 >> (14u * phase + 7u)) & 127u; 864 865 let gp0_signs = get_byte(ksigns_iq2xs[gp0_is / 4u], gp0_is % 4u); 866 let gp1_signs = get_byte(ksigns_iq2xs[gp1_is / 4u], gp1_is % 4u); 867 868 let m_0_3_val4 = create_iq2_m4(gp0_signs, 0); 869 let m_4_7_val4 = create_iq2_m4(gp0_signs, 1); 870 let m_8_11_val4 = create_iq2_m4(gp1_signs, 0); 871 let m_12_15_val4 = create_iq2_m4(gp1_signs, 1); 872 873 let gw_0_3_val4 = create_iq_gw4(gp0_ig, 0); 874 let gw_4_7_val4 = create_iq_gw4(gp0_ig, 4); 875 let gw_8_11_val4 = create_iq_gw4(gp1_ig, 0); 876 let gw_12_15_val4 = create_iq_gw4(gp1_ig, 4); 877 878 store_shmem_iquants(vec4<f16>(db * m_0_3_val4 * gw_0_3_val4), elem_idx + 0u); 879 store_shmem_iquants(vec4<f16>(db * m_4_7_val4 * gw_4_7_val4), elem_idx + 4u); 880 store_shmem_iquants(vec4<f16>(db * m_8_11_val4 * gw_8_11_val4), elem_idx + 8u); 881 store_shmem_iquants(vec4<f16>(db * m_12_15_val4 * gw_12_15_val4), elem_idx + 12u); 882#endif // INIT_SRC0_SHMEM_IQ2_XXS 883 884#if defined(INIT_SRC0_SHMEM_IQ2_XS) 885 let block_byte_base = src0_idx * 74u; // BLOCK_SIZE_BYTES = 74u; 886 let d_byte_base = block_byte_base + 0u; 887 let qs_byte_base = block_byte_base + 2u; 888 let scales_byte_base = block_byte_base + 66u; 889 890 let d = load_f16_as_f32_at_src0(d_byte_base); 891 892 let sub_block = k_in_block / 32u; 893 let phase = (k_in_block / NQ) % 2u; 894 895 let scale = (load_byte_at_src0_aligned(scales_byte_base + 1u * sub_block) >> (4u * phase)) & 0xFu; 896 let db = d * (0.5 + f32(scale)) * 0.25; 897 898 let qs_u32 = load_u32_at_src0(qs_byte_base + 8u * sub_block + 4u * phase); 899 900 let gp0_ig = (qs_u32 & 0x1FFu) * 8u; 901 let gp1_ig = ((qs_u32 >> 16u) & 0x1FFu) * 8u; 902 903 let gp0_is = (qs_u32 >> 9u) & 0x7Fu; 904 let gp1_is = (qs_u32 >> 25u) & 0x7Fu; 905 906 let gp0_signs = get_byte(ksigns_iq2xs[gp0_is / 4u], gp0_is % 4u); 907 let gp1_signs = get_byte(ksigns_iq2xs[gp1_is / 4u], gp1_is % 4u); 908 909 let m_0_3_val4 = create_iq2_m4(gp0_signs, 0); 910 let m_4_7_val4 = create_iq2_m4(gp0_signs, 1); 911 let m_8_11_val4 = create_iq2_m4(gp1_signs, 0); 912 let m_12_15_val4 = create_iq2_m4(gp1_signs, 1); 913 914 let gw_0_3_val4 = create_iq_gw4(gp0_ig, 0); 915 let gw_4_7_val4 = create_iq_gw4(gp0_ig, 4); 916 let gw_8_11_val4 = create_iq_gw4(gp1_ig, 0); 917 let gw_12_15_val4 = create_iq_gw4(gp1_ig, 4); 918 919 store_shmem_iquants(vec4<f16>(db * m_0_3_val4 * gw_0_3_val4), elem_idx + 0u); 920 store_shmem_iquants(vec4<f16>(db * m_4_7_val4 * gw_4_7_val4), elem_idx + 4u); 921 store_shmem_iquants(vec4<f16>(db * m_8_11_val4 * gw_8_11_val4), elem_idx + 8u); 922 store_shmem_iquants(vec4<f16>(db * m_12_15_val4 * gw_12_15_val4), elem_idx + 12u); 923#endif // INIT_SRC0_SHMEM_IQ2_XS 924 925#if defined(INIT_SRC0_SHMEM_IQ2_S) 926 let block_byte_base = src0_idx * 82u; // BLOCK_SIZE_BYTES = 82u; 927 let d_byte_base = block_byte_base + 0u; 928 let qs_byte_base = block_byte_base + 2u; 929 let qh_byte_base = block_byte_base + 66u; 930 let scales_byte_base = block_byte_base + 74u; 931 932 let d = load_f16_as_f32_at_src0(d_byte_base); 933 934 let sub_block = k_in_block / 32u; 935 let phase = (k_in_block / NQ) % 2u; 936 937 let scale = (load_byte_at_src0_aligned(scales_byte_base + 1u * sub_block) >> (4u * phase)) & 0xFu; 938 let db = d * (0.5 + f32(scale)) * 0.25; 939 940 let qs_u16 = load_u32_at_src0(qs_byte_base + 4u * sub_block + 2u * phase) & 0xFFFFu; 941 let signs_u16 = load_u32_at_src0(qs_byte_base + 32u + 4u * sub_block + 2u * phase) & 0xFFFFu; 942 let qh_u4 = (load_byte_at_src0_aligned(qh_byte_base + 1u * sub_block) >> (4u * phase)) & 0xFu; 943 944 let gp0_ig = ((qs_u16 & 0xFFu) | ((qh_u4 & 0x3u) << 8u)) * 8u; 945 let gp1_ig = (((qs_u16 >> 8u) & 0xFFu) | ((qh_u4 & 0xCu) << 6u)) * 8u; 946 947 let gp0_signs = get_byte(signs_u16, 0); 948 let gp1_signs = get_byte(signs_u16, 1); 949 950 let m_0_3_val4 = create_iq2_m4(gp0_signs, 0); 951 let m_4_7_val4 = create_iq2_m4(gp0_signs, 1); 952 let m_8_11_val4 = create_iq2_m4(gp1_signs, 0); 953 let m_12_15_val4 = create_iq2_m4(gp1_signs, 1); 954 955 let gw_0_3_val4 = create_iq_gw4(gp0_ig, 0); 956 let gw_4_7_val4 = create_iq_gw4(gp0_ig, 4); 957 let gw_8_11_val4 = create_iq_gw4(gp1_ig, 0); 958 let gw_12_15_val4 = create_iq_gw4(gp1_ig, 4); 959 960 store_shmem_iquants(vec4<f16>(db * m_0_3_val4 * gw_0_3_val4), elem_idx + 0u); 961 store_shmem_iquants(vec4<f16>(db * m_4_7_val4 * gw_4_7_val4), elem_idx + 4u); 962 store_shmem_iquants(vec4<f16>(db * m_8_11_val4 * gw_8_11_val4), elem_idx + 8u); 963 store_shmem_iquants(vec4<f16>(db * m_12_15_val4 * gw_12_15_val4), elem_idx + 12u); 964#endif // INIT_SRC0_SHMEM_IQ2_S 965 966#if defined(INIT_SRC0_SHMEM_IQ3_XXS) 967 let block_byte_base = src0_idx * 98u; // BLOCK_SIZE_BYTES = 98u; 968 let d_byte_base = block_byte_base + 0u; 969 let qs_byte_base = block_byte_base + 2u; 970 971 let d = load_f16_as_f32_at_src0(d_byte_base); 972 973 let sub_block = k_in_block / 32u; 974 let phase = (k_in_block / NQ) % 2u; 975 976 let qs_u32 = load_u32_at_src0(qs_byte_base + 8u * sub_block + 4u * phase); 977 let sign_u32 = load_u32_at_src0(qs_byte_base + 64u + 4u * sub_block); 978 let db = d * (0.5 + f32(sign_u32 >> 28u)) * 0.5; 979 980 let ig_0_3 = get_byte(qs_u32, 0); 981 let ig_4_7 = get_byte(qs_u32, 1); 982 let ig_8_11 = get_byte(qs_u32, 2); 983 let ig_12_15 = get_byte(qs_u32, 3); 984 985 let gp0_is = (sign_u32 >> (14u * phase + 0u)) & 0x7Fu; 986 let gp1_is = (sign_u32 >> (14u * phase + 7u)) & 0x7Fu; 987 988 let gp0_signs = get_byte(ksigns_iq2xs[gp0_is / 4u], gp0_is % 4u); 989 let gp1_signs = get_byte(ksigns_iq2xs[gp1_is / 4u], gp1_is % 4u); 990 991 let m_0_3_val4 = create_iq2_m4(gp0_signs, 0); 992 let m_4_7_val4 = create_iq2_m4(gp0_signs, 1); 993 let m_8_11_val4 = create_iq2_m4(gp1_signs, 0); 994 let m_12_15_val4 = create_iq2_m4(gp1_signs, 1); 995 996 let gw_0_3_val4 = create_iq_gw4(ig_0_3); 997 let gw_4_7_val4 = create_iq_gw4(ig_4_7); 998 let gw_8_11_val4 = create_iq_gw4(ig_8_11); 999 let gw_12_15_val4 = create_iq_gw4(ig_12_15); 1000 1001 store_shmem_iquants(vec4<f16>(db * m_0_3_val4 * gw_0_3_val4), elem_idx + 0u); 1002 store_shmem_iquants(vec4<f16>(db * m_4_7_val4 * gw_4_7_val4), elem_idx + 4u); 1003 store_shmem_iquants(vec4<f16>(db * m_8_11_val4 * gw_8_11_val4), elem_idx + 8u); 1004 store_shmem_iquants(vec4<f16>(db * m_12_15_val4 * gw_12_15_val4), elem_idx + 12u); 1005#endif // INIT_SRC0_SHMEM_IQ3_XXS 1006 1007#if defined(INIT_SRC0_SHMEM_IQ3_S) 1008 let block_byte_base = src0_idx * 110u; // BLOCK_SIZE_BYTES = 110u; 1009 let d_byte_base = block_byte_base + 0u; 1010 let qs_byte_base = block_byte_base + 2u; 1011 let qh_byte_base = block_byte_base + 66u; 1012 let signs_byte_base = block_byte_base + 74u; 1013 let scales_byte_base = block_byte_base + 106u; 1014 1015 let d = load_f16_as_f32_at_src0(d_byte_base); 1016 1017 let sub_block = k_in_block / 32u; 1018 let phase = (k_in_block / NQ) % 2u; 1019 1020 let scale = (load_byte_at_src0_aligned(scales_byte_base + 1u * (sub_block / 2u)) >> (4u * (sub_block % 2u))) & 0xFu; 1021 let db = d * (1.0 + 2.0 * f32(scale)); 1022 1023 let qs_u32 = load_u32_at_src0(qs_byte_base + 8u * sub_block + 4u * phase); 1024 let qh_u4 = (load_byte_at_src0_aligned(qh_byte_base + 1u * sub_block) >> (4u * phase)) & 0xFu; 1025 let signs_u16 = (load_u32_at_src0(signs_byte_base + 4u * sub_block + 2u * phase)) & 0xFFFFu; 1026 1027 let ig_0_3 = ((qs_u32 >> 0u) & 0xFFu) | ((qh_u4 & 0x1u) << 8u); 1028 let ig_4_7 = ((qs_u32 >> 8u) & 0xFFu) | ((qh_u4 & 0x2u) << 7u); 1029 let ig_8_11 = ((qs_u32 >> 16u) & 0xFFu) | ((qh_u4 & 0x4u) << 6u); 1030 let ig_12_15 = ((qs_u32 >> 24u) & 0xFFu) | ((qh_u4 & 0x8u) << 5u); 1031 1032 let gp0_signs = get_byte(signs_u16, 0); 1033 let gp1_signs = get_byte(signs_u16, 1); 1034 1035 let m_0_3_val4 = create_iq2_m4(gp0_signs, 0); 1036 let m_4_7_val4 = create_iq2_m4(gp0_signs, 1); 1037 let m_8_11_val4 = create_iq2_m4(gp1_signs, 0); 1038 let m_12_15_val4 = create_iq2_m4(gp1_signs, 1); 1039 1040 let gw_0_3_val4 = create_iq_gw4(ig_0_3); 1041 let gw_4_7_val4 = create_iq_gw4(ig_4_7); 1042 let gw_8_11_val4 = create_iq_gw4(ig_8_11); 1043 let gw_12_15_val4 = create_iq_gw4(ig_12_15); 1044 1045 store_shmem_iquants(vec4<f16>(db * m_0_3_val4 * gw_0_3_val4), elem_idx + 0u); 1046 store_shmem_iquants(vec4<f16>(db * m_4_7_val4 * gw_4_7_val4), elem_idx + 4u); 1047 store_shmem_iquants(vec4<f16>(db * m_8_11_val4 * gw_8_11_val4), elem_idx + 8u); 1048 store_shmem_iquants(vec4<f16>(db * m_12_15_val4 * gw_12_15_val4), elem_idx + 12u); 1049#endif // INIT_SRC0_SHMEM_IQ3_S 1050 } 1051} 1052#endif // i-quants (super block size: 256)