char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Ruben Ortlamvulkan: use spec constant for matrix matrix multiplication A-type (#25773)91f6a6cf3

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
31.5 KiB779 linesraw
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_int8 : require
8#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
9#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
10
11#extension GL_KHR_memory_scope_semantics : enable
12#extension GL_KHR_cooperative_matrix : enable
13#extension GL_NV_cooperative_matrix2 : enable
14#ifdef GGML_VULKAN_COOPMAT2_DECODE_VECTOR
15#extension GL_NV_cooperative_matrix_decode_vector : enable
16#endif
17#extension GL_EXT_buffer_reference : enable
18#extension GL_KHR_shader_subgroup_ballot : enable
19#extension GL_KHR_shader_subgroup_vote : enable
20#ifdef DATA_A_BF16
21#extension GL_EXT_bfloat16 : enable
22#endif
23
24#include "ggml_type_ids.glsl"
25
26#ifdef MULMAT_QUANT
27layout (constant_id = 7) const uint MmTypeA = 0;
28layout (constant_id = 8) const uint MmABlockBytes = 2;
29#endif
30
31#include "types.glsl"
32#include "utils.glsl"
33
34layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
35
36#define IS_MUL_MM2 1
37
38layout (constant_id = 0) const uint BLOCK_SIZE = 256;
39layout (constant_id = 1) const uint BM = 64;
40layout (constant_id = 2) const uint BN = 64;
41layout (constant_id = 3) const uint BK = 16;  // Assumed to be 32 if working with a quant
42
43layout (constant_id = 4) const bool enable_smaller_matrices = false;
44const uint BNover2 = enable_smaller_matrices ? (BN / 2) : BN;
45const uint BNover4 = enable_smaller_matrices ? (BN / 4) : BN;
46layout (constant_id = 5) const uint ALIGNED = 0;
47layout (constant_id = 6) const uint subgroup_size = 32;
48
49#ifdef MULMAT_QUANT
50
51uint mm_quant_k() {
52    switch (MmTypeA) {
53    case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: case GGML_TYPE_Q5_1:
54    case GGML_TYPE_Q8_0:
55        return 32u;
56    case GGML_TYPE_Q1_0:
57        return 128u;
58    case GGML_TYPE_Q2_0:
59        return 64u;
60    default:
61        return 256u;
62    }
63}
64#endif
65
66layout (push_constant) uniform parameter
67{
68    uint M;
69    uint N;
70    uint K;
71    uint stride_a;
72    uint stride_b;
73    uint stride_d;
74
75    uint batch_stride_a;
76    uint batch_stride_b;
77    uint batch_stride_d;
78
79#ifdef MUL_MAT_ID
80    uint nei0;
81    uint nei1;
82    uint nbi1;
83    uint ne11;
84    uint n_experts;
85    uint hoist_row_ids;
86#else
87    uint base_work_group_z;
88    uint num_batches;
89    uint k_split;
90    uint ne02;
91    uint ne12;
92    uint broadcast2;
93    uint broadcast3;
94    // N dimension for the B matrix can be >= p.N
95    uint padded_N;
96#endif
97} p;
98
99
100#ifndef MULMAT_QUANT
101layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
102#else
103layout (binding = 0) readonly buffer A {uint8_t data_a[];};
104#endif
105layout (binding = 1) readonly buffer B {B_TYPE data_b[];};
106layout (binding = 2) writeonly buffer D {D_TYPE data_d[];};
107#if defined(MUL_MAT_ID) && defined(GGML_VULKAN_COOPMAT2_DECODE_VECTOR)
108layout (binding = 1) readonly buffer B4 {B_TYPEV4 data_b_v4[];};
109#endif
110
111#if defined(MULMAT_QUANT) || QUANT_K > 1
112#include "dequant_funcs_cm2.glsl"
113#ifndef MULMAT_QUANT
114// Per-type path: use the alias set by dequant_funcs_cm2.glsl
115#if defined(dequantFuncA_v) && defined(GGML_VULKAN_COOPMAT2_DECODE_VECTOR)
116#define DECODEFUNCA , dequantFuncA, dequantFuncA_v
117#else
118#define DECODEFUNCA , dequantFuncA
119#endif
120#else
121layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufA {
122    uint8_t raw[MmABlockBytes];
123};
124
125float16_t mmDecodeA(const in decodeBufA bl_in, const in uint blockCoords[2], const in uint coordInBlock[2]) {
126    switch (MmTypeA) {
127    case GGML_TYPE_Q1_0:  return dequantFuncQ1_0 (decodeBufQ1_0 (bl_in), blockCoords, coordInBlock);
128    case GGML_TYPE_Q2_0:  return dequantFuncQ2_0 (decodeBufQ2_0 (bl_in), blockCoords, coordInBlock);
129    case GGML_TYPE_Q4_0:  return dequantFuncQ4_0 (decodeBufQ4_0 (bl_in), blockCoords, coordInBlock);
130    case GGML_TYPE_Q4_1:  return dequantFuncQ4_1 (decodeBufQ4_1 (bl_in), blockCoords, coordInBlock);
131    case GGML_TYPE_Q5_0:  return dequantFuncQ5_0 (decodeBufQ5_0 (bl_in), blockCoords, coordInBlock);
132    case GGML_TYPE_Q5_1:  return dequantFuncQ5_1 (decodeBufQ5_1 (bl_in), blockCoords, coordInBlock);
133    case GGML_TYPE_Q8_0:  return dequantFuncQ8_0 (decodeBufQ8_0 (bl_in), blockCoords, coordInBlock);
134    case GGML_TYPE_Q2_K:  return dequantFuncQ2_K (decodeBufQ2_K (bl_in), blockCoords, coordInBlock);
135    case GGML_TYPE_Q3_K:  return dequantFuncQ3_K (decodeBufQ3_K (bl_in), blockCoords, coordInBlock);
136    case GGML_TYPE_Q6_K:  return dequantFuncQ6_K (decodeBufQ6_K (bl_in), blockCoords, coordInBlock);
137    case GGML_TYPE_TQ1_0: return dequantFuncTQ1_0(decodeBufTQ1_0(bl_in), blockCoords, coordInBlock);
138    case GGML_TYPE_TQ2_0: return dequantFuncTQ2_0(decodeBufTQ2_0(bl_in), blockCoords, coordInBlock);
139    default:              return float16_t(0);
140    }
141}
142
143#ifdef GGML_VULKAN_COOPMAT2_DECODE_VECTOR
144f16vec4 mmDecodeA_v(const in decodeBufA bl_in, const in uint blockCoords[2], const in uint coordInBlock[2]) {
145    switch (MmTypeA) {
146    case GGML_TYPE_Q1_0:  return dequantFuncQ1_0_v (decodeBufQ1_0 (bl_in), blockCoords, coordInBlock);
147    case GGML_TYPE_Q2_0:  return dequantFuncQ2_0_v (decodeBufQ2_0 (bl_in), blockCoords, coordInBlock);
148    case GGML_TYPE_Q4_0:  return dequantFuncQ4_0_v (decodeBufQ4_0 (bl_in), blockCoords, coordInBlock);
149    case GGML_TYPE_Q4_1:  return dequantFuncQ4_1_v (decodeBufQ4_1 (bl_in), blockCoords, coordInBlock);
150    case GGML_TYPE_Q5_0:  return dequantFuncQ5_0_v (decodeBufQ5_0 (bl_in), blockCoords, coordInBlock);
151    case GGML_TYPE_Q5_1:  return dequantFuncQ5_1_v (decodeBufQ5_1 (bl_in), blockCoords, coordInBlock);
152    case GGML_TYPE_Q8_0:  return dequantFuncQ8_0_v (decodeBufQ8_0 (bl_in), blockCoords, coordInBlock);
153    case GGML_TYPE_Q2_K:  return dequantFuncQ2_K_v (decodeBufQ2_K (bl_in), blockCoords, coordInBlock);
154    case GGML_TYPE_Q3_K:  return dequantFuncQ3_K_v (decodeBufQ3_K (bl_in), blockCoords, coordInBlock);
155    case GGML_TYPE_Q6_K:  return dequantFuncQ6_K_v (decodeBufQ6_K (bl_in), blockCoords, coordInBlock);
156    case GGML_TYPE_TQ1_0: return dequantFuncTQ1_0_v(decodeBufTQ1_0(bl_in), blockCoords, coordInBlock);
157    case GGML_TYPE_TQ2_0: return dequantFuncTQ2_0_v(decodeBufTQ2_0(bl_in), blockCoords, coordInBlock);
158    default:              return f16vec4(0);
159    }
160}
161#define DECODEFUNCA , mmDecodeA, mmDecodeA_v
162#else
163#define DECODEFUNCA , mmDecodeA
164#endif
165#endif
166#else
167#define DECODEFUNCA
168#endif
169
170#if !defined(fetch_scales)
171#define fetch_scales(a, b, c, d, e, f)
172#endif
173#if !defined(store_scales)
174#define store_scales(a)
175#endif
176
177#if defined(DATA_A_BF16)
178#define MAT_TYPE bfloat16_t
179#else
180#define MAT_TYPE FLOAT_TYPE
181#endif
182
183#ifdef MUL_MAT_ID
184layout (binding = 3) readonly buffer IDS {int data_ids[];};
185layout (binding = 4) readonly buffer Counts {int data_expert_count[];};
186
187shared u16vec4 row_ids[BN];
188
189layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufB {
190   B_TYPE b[];
191};
192
193uint _ne1;
194shared uvec4 ballots_sh[BLOCK_SIZE / subgroup_size];
195
196B_TYPE decodeFuncB(const in decodeBufB bl, const in uint blockCoords[2], const in uint coordInBlock[2])
197{
198    const uint row_i = blockCoords[0];
199
200    const u16vec4 row_idx = row_ids[row_i];
201#if defined(GGML_VULKAN_COOPMAT2_DECODE_VECTOR)
202    // The decode-vector path gives B a K-dimension tensor-layout block size of BK.
203    const uint k = blockCoords[1] * BK + coordInBlock[1];
204#else
205    const uint k = blockCoords[1];
206#endif
207    B_TYPE ret = data_b[row_idx.y * p.batch_stride_b + row_idx.x * p.stride_b + k];
208
209    return ret;
210}
211
212#if defined(GGML_VULKAN_COOPMAT2_DECODE_VECTOR)
213B_TYPEV4 decodeFuncB_v(const in decodeBufB bl, const in uint blockCoords[2], const in uint coordInBlock[2])
214{
215    const uint row_i = blockCoords[0];
216
217    const u16vec4 row_idx = row_ids[row_i];
218    const uint k = blockCoords[1] * BK + coordInBlock[1];
219    const uint base = row_idx.y * p.batch_stride_b + row_idx.x * p.stride_b + k;
220
221    return data_b_v4[base >> 2];
222}
223#define DECODEFUNCB , decodeFuncB, decodeFuncB_v
224#else
225#define DECODEFUNCB , decodeFuncB
226#endif
227
228D_TYPE perElemOpD(const in uint32_t r, const in uint32_t c, const in D_TYPE elem, const in uint32_t ir, const in uint32_t ic)
229{
230    uint dr = ir * BM + r;
231    uint dc = ic * BN + c;
232
233    if (dr < p.M && dc < _ne1) {
234        uint row_i = c;
235        const u16vec4 row_idx = row_ids[row_i];
236        data_d[row_idx.y * p.batch_stride_d + row_idx.z * p.stride_d + dr] = elem;
237    }
238    return elem;
239}
240
241void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) {
242    _ne1 = 0;
243    uint num_elements = p.nei1 * p.nei0;
244    uint nei0shift = findLSB(p.nei0);
245
246    uint ids[16];
247    uint iter = 0;
248
249    uint expert_count = data_expert_count[expert_idx];
250
251    for (uint j = 0; j < num_elements; j += BLOCK_SIZE) {
252        // prefetch up to 16 elements
253        if (iter == 0) {
254            [[unroll]] for (uint k = 0; k < 16; ++k) {
255                uint i = j + gl_LocalInvocationIndex + k*BLOCK_SIZE;
256                bool in_range = i < num_elements;
257                uint ii1;
258                if (nei0_is_pow2) {
259                    ii1 = i >> nei0shift;
260                } else {
261                    ii1 = i / p.nei0;
262                }
263                uint ii0 = i - ii1 * p.nei0;
264                ids[k] = in_range ? data_ids[ii1*p.nbi1 + ii0] : 0;
265            }
266        }
267        uint i = j + gl_LocalInvocationIndex;
268        bool in_range = i < num_elements;
269        uint ii1;
270        if (nei0_is_pow2) {
271            ii1 = i >> nei0shift;
272        } else {
273            ii1 = i / p.nei0;
274        }
275        uint ii0 = i - ii1 * p.nei0;
276        uint id = ids[iter++];
277        uvec4 ballot = subgroupBallot(in_range && id == expert_idx);
278
279        if (gl_SubgroupInvocationID == 0) {
280            ballots_sh[gl_SubgroupID] = ballot;
281        }
282        barrier();
283
284        uint subgroup_base = 0;
285        uint total = 0;
286        for (uint k = 0; k < gl_NumSubgroups; ++k) {
287            if (k == gl_SubgroupID) {
288                subgroup_base = total;
289            }
290            total += subgroupBallotBitCount(ballots_sh[k]);
291        }
292        barrier();
293
294        uint idx = subgroup_base + subgroupBallotExclusiveBitCount(ballot);
295        if (in_range && id == expert_idx && _ne1 + idx >= ic * BN && _ne1 + idx < (ic + 1) * BN) {
296            row_ids[_ne1 + idx - ic * BN] = u16vec4(fastmod(ii0, p.ne11), ii1, ii0, 0);
297        }
298        _ne1 += total;
299        iter &= 15;
300        if (_ne1 >= (ic + 1) * BN || _ne1 == expert_count) {
301            break;
302        }
303    }
304    barrier();
305}
306
307void load_row_ids_hoisted(uint expert_idx, uint ic) {
308    _ne1 = uint(data_expert_count[expert_idx]);
309
310    const uint tile_begin = ic * BN;
311    const uint tile_count = tile_begin < _ne1 ? min(BN, _ne1 - tile_begin) : 0;
312    const uint expert_offset = uint(data_expert_count[p.n_experts + expert_idx]);
313    const uint row_ids_offset = 2 * p.n_experts + 1 + expert_offset + tile_begin;
314
315    for (uint i = gl_LocalInvocationIndex; i < tile_count; i += BLOCK_SIZE) {
316        const uint packed_row_id = uint(data_expert_count[row_ids_offset + i]);
317        const uint ii0 = packed_row_id & 0xffffu;
318        const uint ii1 = packed_row_id >> 16;
319        row_ids[i] = u16vec4(fastmod(ii0, p.ne11), ii1, ii0, 0);
320    }
321    barrier();
322}
323#endif
324
325#ifdef MULMAT_QUANT
326#include "iq_shmem_init.glsl"
327#endif
328
329void main() {
330    const uint tid = gl_LocalInvocationIndex;
331    const uint ic = gl_WorkGroupID.y;
332
333#ifdef MUL_MAT_ID
334    const uint expert_idx = gl_WorkGroupID.z;
335    if (ic * BN >= data_expert_count[expert_idx]) {
336        return;
337    }
338    // initialize to row 0 so we don't need to bounds check
339    if (tid < BN) {
340        row_ids[tid] = u16vec4(0);
341    }
342#if !defined(NEEDS_INIT_IQ_SHMEM)
343    barrier();
344#endif
345#endif
346
347#if defined(NEEDS_INIT_IQ_SHMEM) || defined(MULMAT_QUANT)
348    init_iq_shmem(gl_WorkGroupSize);
349#endif
350
351#ifndef MUL_MAT_ID
352    const uint batch_idx = gl_WorkGroupID.z + p.base_work_group_z;
353
354    const uint i13 = batch_idx / p.ne12;
355    const uint i12 = batch_idx % p.ne12;
356
357    const uint i03 = i13 / p.broadcast3;
358    const uint i02 = i12 / p.broadcast2;
359
360    const uint batch_idx_a = i03 * p.ne02 + i02;
361#endif
362
363    const uint blocks_m = (p.M + BM - 1) / BM;
364    const uint ir = gl_WorkGroupID.x % blocks_m;
365    const uint ik = gl_WorkGroupID.x / blocks_m;
366
367#ifdef MUL_MAT_ID
368    if (p.hoist_row_ids != 0) {
369        load_row_ids_hoisted(expert_idx, ic);
370    } else if (bitCount(p.nei0) == 1) {
371        load_row_ids(expert_idx, true, ic);
372    } else {
373        load_row_ids(expert_idx, false, ic);
374    }
375
376    // Workgroup has no work
377    if (ic * BN >= _ne1) return;
378#endif
379
380#ifdef MUL_MAT_ID
381    uint start_k = 0;
382    const uint end_k = p.K;
383#else
384    uint start_k = ik * p.k_split;
385    const uint end_k = min(p.K, (ik + 1) * p.k_split);
386#endif
387
388#ifdef MULMAT_QUANT
389    const uint qk = mm_quant_k();
390#else
391    const uint qk = QUANT_K;
392#endif
393
394#ifdef MUL_MAT_ID
395    uint pos_a = expert_idx * (p.batch_stride_a / qk);
396    uint pos_b = 0;
397#else
398    uint pos_a = batch_idx_a * (p.batch_stride_a / qk);
399    uint pos_b = batch_idx * p.batch_stride_b;
400    uint pos_d = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * p.num_batches;
401#endif
402
403#ifdef MULMAT_QUANT
404    // pos_a is a byte offset into the raw buffer; strides stay in block units
405    pos_a *= MmABlockBytes;
406#endif
407
408    uint stride_a = p.stride_a / qk;
409    uint stride_b = p.stride_b;
410
411    // Hint to the compiler that values are aligned (want 16B alignment).
412    // Quants are always block-aligned, no alignment needed.
413    if (ALIGNED != 0) {
414#if !defined(MULMAT_QUANT) && QUANT_K == 1
415        stride_a &= ~7;
416#endif
417        stride_b &= ~7;
418    }
419
420    // Create layouts for both clamped and unclamped accesses
421    tensorLayoutNV<2> tensorLayoutA = createTensorLayoutNV(2);
422    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutAClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
423    tensorLayoutNV<2> tensorLayoutB = createTensorLayoutNV(2);
424#ifndef MUL_MAT_ID
425    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutBClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
426#endif
427    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
428
429    tensorLayoutA = setTensorLayoutBlockSizeNV(tensorLayoutA, 1, qk);
430    tensorLayoutAClamp = setTensorLayoutBlockSizeNV(tensorLayoutAClamp, 1, qk);
431#if defined(MUL_MAT_ID) && defined(GGML_VULKAN_COOPMAT2_DECODE_VECTOR)
432    tensorLayoutB = setTensorLayoutBlockSizeNV(tensorLayoutB, 1, BK);
433#endif
434
435    // Use end_k rather than p.K as the dimension because that's what
436    // we need to bound check against when using split_k.
437    tensorLayoutA = setTensorLayoutDimensionNV(tensorLayoutA, p.M, end_k);
438#ifdef MUL_MAT_ID
439    // MUL_MAT_ID pads each B row to stride_b so partial K tiles read zeros without clamping.
440    tensorLayoutB = setTensorLayoutDimensionNV(tensorLayoutB, BN, p.stride_b);
441#else
442    // Bounds check B against padded_N, but bounds check D against N.
443    tensorLayoutB = setTensorLayoutDimensionNV(tensorLayoutB, p.padded_N, end_k);
444#endif
445    tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.N, p.M);
446    tensorLayoutAClamp = setTensorLayoutDimensionNV(tensorLayoutAClamp, p.M, end_k);
447#ifndef MUL_MAT_ID
448    tensorLayoutBClamp = setTensorLayoutDimensionNV(tensorLayoutBClamp, p.padded_N, end_k);
449#endif
450
451    tensorLayoutD = setTensorLayoutStrideNV(tensorLayoutD, p.stride_d, 1);
452
453    tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0);
454
455#if !defined(MUL_MAT_ID)
456
457    const uint START_ALIGN_K = 256;
458    // For Qi_K (block size 256), unroll whole 256 element tiles.
459    // For legacy quants (block size 32), unroll 8x.
460    const uint UNROLL_K = (qk == 256) ? 256 : (BK * 8);
461    const uint unroll_count = UNROLL_K / BK;
462
463    // Detect a fast path where all loads are entirely in bounds and no clamping is required
464    if ((ir + 1) * BM <= p.M && (ic + 1) * BN <= p.padded_N && (start_k % START_ALIGN_K) == 0 && (end_k % BK) == 0 &&
465#if !defined(MULMAT_QUANT) && QUANT_K == 1
466        (stride_a % 8) == 0 &&
467#endif
468        (stride_b % 8) == 0) {
469        // Hint to the compiler that values are aligned (want 16B alignment)
470        start_k &= ~(START_ALIGN_K-1);
471        stride_b &= ~7;
472#if !defined(MULMAT_QUANT) && QUANT_K == 1
473        stride_a &= ~7;
474#endif
475
476        tensorLayoutA = setTensorLayoutStrideNV(tensorLayoutA, stride_a, 1);
477        tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1);
478
479        uint k_iters = (end_k - start_k) / UNROLL_K;
480        uint block_k = start_k;
481
482        // fetch scale values for a tile of quants. These will be copied into shared memory.
483        // The fetches and stores are pipelined to hide the latency.
484        fetch_scales(ir * BM, pos_a, stride_a, start_k, tid, true);
485
486        if (enable_smaller_matrices && ic * BN + BNover4 >= p.N) {
487            coopmat<ACC_TYPE, gl_ScopeWorkgroup, BM, BNover4, gl_MatrixUseAccumulator> sum = coopmat<ACC_TYPE, gl_ScopeWorkgroup, BM, BNover4, gl_MatrixUseAccumulator>(0.0);
488            for (uint i = 0; i < k_iters; ++i) {
489
490                store_scales(tid);
491                if (block_k + UNROLL_K < end_k) {
492                    fetch_scales(ir * BM, pos_a, stride_a, block_k + UNROLL_K, tid, true);
493                }
494
495                // Manually partial unroll
496                [[unroll]] for (uint j = 0; j < unroll_count; ++j) {
497                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> mat_a;
498                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BK, BNover4, gl_MatrixUseB> mat_b;
499
500                    coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA);
501                    coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BNover4, block_k, BK), tensorViewTranspose);
502
503                    sum = coopMatMulAdd(mat_a, mat_b, sum);
504                    block_k += BK;
505                }
506            }
507            // Do any remaining iterations that were not unrolled
508            if (block_k < end_k) {
509                store_scales(tid);
510            }
511            while (block_k < end_k) {
512                coopmat<MAT_TYPE, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> mat_a;
513                coopmat<MAT_TYPE, gl_ScopeWorkgroup, BK, BNover4, gl_MatrixUseB> mat_b;
514
515                coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA);
516                coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BNover4, block_k, BK), tensorViewTranspose);
517
518                sum = coopMatMulAdd(mat_a, mat_b, sum);
519                block_k += BK;
520            }
521#if defined(ACC_TYPE_MAX)
522            [[unroll]] for (uint i = 0; i < sum.length(); ++i) { sum[i] = clamp(sum[i], -ACC_TYPE_MAX, ACC_TYPE_MAX); }
523#endif
524
525            coopmat<D_TYPE, gl_ScopeWorkgroup, BM, BNover4, gl_MatrixUseAccumulator> mat_d = coopmat<D_TYPE, gl_ScopeWorkgroup, BM, BNover4, gl_MatrixUseAccumulator>(sum);
526
527            coopMatStoreTensorNV(mat_d, data_d, pos_d, sliceTensorLayoutNV(tensorLayoutD, ic * BN, BNover4, ir * BM, BM), tensorViewTranspose);
528            return;
529        } else if (enable_smaller_matrices && ic * BN + BNover2 >= p.N) {
530            coopmat<ACC_TYPE, gl_ScopeWorkgroup, BM, BNover2, gl_MatrixUseAccumulator> sum = coopmat<ACC_TYPE, gl_ScopeWorkgroup, BM, BNover2, gl_MatrixUseAccumulator>(0.0);
531            for (uint i = 0; i < k_iters; ++i) {
532
533                store_scales(tid);
534                if (block_k + UNROLL_K < end_k) {
535                    fetch_scales(ir * BM, pos_a, stride_a, block_k + UNROLL_K, tid, true);
536                }
537
538                // Manually partial unroll
539                [[unroll]] for (uint j = 0; j < unroll_count; ++j) {
540                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> mat_a;
541                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BK, BNover2, gl_MatrixUseB> mat_b;
542
543                    coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA);
544                    coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BNover2, block_k, BK), tensorViewTranspose);
545
546                    sum = coopMatMulAdd(mat_a, mat_b, sum);
547                    block_k += BK;
548                }
549            }
550            // Do any remaining iterations that were not unrolled
551            if (block_k < end_k) {
552                store_scales(tid);
553            }
554            while (block_k < end_k) {
555                coopmat<MAT_TYPE, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> mat_a;
556                coopmat<MAT_TYPE, gl_ScopeWorkgroup, BK, BNover2, gl_MatrixUseB> mat_b;
557
558                coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA);
559                coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BNover2, block_k, BK), tensorViewTranspose);
560
561                sum = coopMatMulAdd(mat_a, mat_b, sum);
562                block_k += BK;
563            }
564#if defined(ACC_TYPE_MAX)
565            [[unroll]] for (uint i = 0; i < sum.length(); ++i) { sum[i] = clamp(sum[i], -ACC_TYPE_MAX, ACC_TYPE_MAX); }
566#endif
567
568            coopmat<D_TYPE, gl_ScopeWorkgroup, BM, BNover2, gl_MatrixUseAccumulator> mat_d = coopmat<D_TYPE, gl_ScopeWorkgroup, BM, BNover2, gl_MatrixUseAccumulator>(sum);
569
570            coopMatStoreTensorNV(mat_d, data_d, pos_d, sliceTensorLayoutNV(tensorLayoutD, ic * BN, BNover2, ir * BM, BM), tensorViewTranspose);
571            return;
572        } else {
573            coopmat<ACC_TYPE, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum = coopmat<ACC_TYPE, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0.0);
574
575            for (uint i = 0; i < k_iters; ++i) {
576
577                store_scales(tid);
578                if (block_k + UNROLL_K < end_k) {
579                    fetch_scales(ir * BM, pos_a, stride_a, block_k + UNROLL_K, tid, true);
580                }
581
582                // Manually partial unroll
583                [[unroll]] for (uint j = 0; j < unroll_count; ++j) {
584                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> mat_a;
585                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> mat_b;
586
587                    coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA);
588                    coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose);
589
590                    sum = coopMatMulAdd(mat_a, mat_b, sum);
591                    block_k += BK;
592                }
593            }
594            // Do any remaining iterations that were not unrolled
595            if (block_k < end_k) {
596                store_scales(tid);
597            }
598            while (block_k < end_k) {
599                coopmat<MAT_TYPE, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> mat_a;
600                coopmat<MAT_TYPE, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> mat_b;
601
602                coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA);
603                coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, ic * BN, BN, block_k, BK), tensorViewTranspose);
604
605                sum = coopMatMulAdd(mat_a, mat_b, sum);
606                block_k += BK;
607            }
608#if defined(ACC_TYPE_MAX)
609            [[unroll]] for (uint i = 0; i < sum.length(); ++i) { sum[i] = clamp(sum[i], -ACC_TYPE_MAX, ACC_TYPE_MAX); }
610#endif
611
612            coopmat<D_TYPE, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> mat_d = coopmat<D_TYPE, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(sum);
613
614            coopMatStoreTensorNV(mat_d, data_d, pos_d, sliceTensorLayoutNV(tensorLayoutD, ic * BN, BN, ir * BM, BM), tensorViewTranspose);
615            return;
616        }
617    } else
618#endif // !defined(MUL_MAT_ID)
619    {
620        tensorLayoutA = setTensorLayoutStrideNV(tensorLayoutA, stride_a, 1);
621
622        tensorLayoutAClamp = setTensorLayoutStrideNV(tensorLayoutAClamp, stride_a, 1);
623
624        tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1);
625
626#ifndef MUL_MAT_ID
627        tensorLayoutBClamp = setTensorLayoutStrideNV(tensorLayoutBClamp, stride_b, 1);
628#endif
629
630        uint k_iters = (end_k - start_k + BK - 1) / BK;
631
632        fetch_scales(ir * BM, pos_a, stride_a, start_k, tid, false);
633        store_scales(tid);
634
635#ifdef MUL_MAT_ID
636        if (enable_smaller_matrices && ic * BN + BNover4 >= _ne1) {
637            coopmat<ACC_TYPE, gl_ScopeWorkgroup, BM, BNover4, gl_MatrixUseAccumulator> sum;
638            sum = coopmat<ACC_TYPE, gl_ScopeWorkgroup, BM, BNover4, gl_MatrixUseAccumulator>(0.0);
639
640            [[dont_unroll]]
641            for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) {
642
643                if ((block_k % qk) == 0) {
644                    store_scales(tid);
645                }
646                if (block_k + BK < end_k && ((block_k + BK) % qk) == 0) {
647                    fetch_scales(ir * BM, pos_a, stride_a, block_k + BK, tid, false);
648                }
649
650                if ((ir + 1) * BM <= p.M && block_k + BK <= end_k) {
651                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> mat_a;
652                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BK, BNover4, gl_MatrixUseB> mat_b;
653
654                    coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA);
655                    coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, 0, BNover4, block_k, BK), tensorViewTranspose DECODEFUNCB);
656
657                    sum = coopMatMulAdd(mat_a, mat_b, sum);
658                } else {
659                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> mat_a;
660                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BK, BNover4, gl_MatrixUseB> mat_b;
661
662                    coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutAClamp, ir * BM, BM, block_k, BK) DECODEFUNCA);
663                    coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, 0, BNover4, block_k, BK), tensorViewTranspose DECODEFUNCB);
664
665                    sum = coopMatMulAdd(mat_a, mat_b, sum);
666                }
667            }
668#if defined(ACC_TYPE_MAX)
669            [[unroll]] for (uint i = 0; i < sum.length(); ++i) { sum[i] = clamp(sum[i], -ACC_TYPE_MAX, ACC_TYPE_MAX); }
670#endif
671
672            // Convert from ACC_TYPE to D_TYPE
673            coopmat<D_TYPE, gl_ScopeWorkgroup, BM, BNover4, gl_MatrixUseAccumulator> mat_d;
674            mat_d = coopmat<D_TYPE, gl_ScopeWorkgroup, BM, BNover4, gl_MatrixUseAccumulator>(sum);
675
676            // Call callback to store each element, remapping row through shared memory
677            coopMatPerElementNV(mat_d, mat_d, perElemOpD, ir, ic);
678            return;
679        }
680        if (enable_smaller_matrices && ic * BN + BNover2 >= _ne1) {
681            coopmat<ACC_TYPE, gl_ScopeWorkgroup, BM, BNover2, gl_MatrixUseAccumulator> sum;
682            sum = coopmat<ACC_TYPE, gl_ScopeWorkgroup, BM, BNover2, gl_MatrixUseAccumulator>(0.0);
683
684            [[dont_unroll]]
685            for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) {
686
687                if ((block_k % qk) == 0) {
688                    store_scales(tid);
689                }
690                if (block_k + BK < end_k && ((block_k + BK) % qk) == 0) {
691                    fetch_scales(ir * BM, pos_a, stride_a, block_k + BK, tid, false);
692                }
693
694                if ((ir + 1) * BM <= p.M && block_k + BK <= end_k) {
695                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> mat_a;
696                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BK, BNover2, gl_MatrixUseB> mat_b;
697
698                    coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA);
699                    coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, 0, BNover2, block_k, BK), tensorViewTranspose DECODEFUNCB);
700
701                    sum = coopMatMulAdd(mat_a, mat_b, sum);
702                } else {
703                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> mat_a;
704                    coopmat<MAT_TYPE, gl_ScopeWorkgroup, BK, BNover2, gl_MatrixUseB> mat_b;
705
706                    coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutAClamp, ir * BM, BM, block_k, BK) DECODEFUNCA);
707                    coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, 0, BNover2, block_k, BK), tensorViewTranspose DECODEFUNCB);
708
709                    sum = coopMatMulAdd(mat_a, mat_b, sum);
710                }
711            }
712#if defined(ACC_TYPE_MAX)
713            [[unroll]] for (uint i = 0; i < sum.length(); ++i) { sum[i] = clamp(sum[i], -ACC_TYPE_MAX, ACC_TYPE_MAX); }
714#endif
715
716            // Convert from ACC_TYPE to D_TYPE
717            coopmat<D_TYPE, gl_ScopeWorkgroup, BM, BNover2, gl_MatrixUseAccumulator> mat_d;
718            mat_d = coopmat<D_TYPE, gl_ScopeWorkgroup, BM, BNover2, gl_MatrixUseAccumulator>(sum);
719
720            // Call callback to store each element, remapping row through shared memory
721            coopMatPerElementNV(mat_d, mat_d, perElemOpD, ir, ic);
722            return;
723        }
724#endif
725        coopmat<ACC_TYPE, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum;
726        sum = coopmat<ACC_TYPE, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0.0);
727
728        [[dont_unroll]]
729        for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) {
730
731            if ((block_k % qk) == 0) {
732                store_scales(tid);
733            }
734            if (block_k + BK < end_k && ((block_k + BK) % qk) == 0) {
735                fetch_scales(ir * BM, pos_a, stride_a, block_k + BK, tid, false);
736            }
737
738            if ((ir + 1) * BM <= p.M && block_k + BK <= end_k) {
739                coopmat<MAT_TYPE, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> mat_a;
740                coopmat<MAT_TYPE, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> mat_b;
741
742                coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutA, ir * BM, BM, block_k, BK) DECODEFUNCA);
743#ifdef MUL_MAT_ID
744                coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, 0, BN, block_k, BK), tensorViewTranspose DECODEFUNCB);
745#else
746                coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutBClamp, ic * BN, BN, block_k, BK), tensorViewTranspose);
747#endif
748
749                sum = coopMatMulAdd(mat_a, mat_b, sum);
750            } else {
751                coopmat<MAT_TYPE, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> mat_a;
752                coopmat<MAT_TYPE, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> mat_b;
753
754                coopMatLoadTensorNV(mat_a, data_a, pos_a, sliceTensorLayoutNV(tensorLayoutAClamp, ir * BM, BM, block_k, BK) DECODEFUNCA);
755#ifdef MUL_MAT_ID
756                coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutB, 0, BN, block_k, BK), tensorViewTranspose DECODEFUNCB);
757#else
758                coopMatLoadTensorNV(mat_b, data_b, pos_b, sliceTensorLayoutNV(tensorLayoutBClamp, ic * BN, BN, block_k, BK), tensorViewTranspose);
759#endif
760
761                sum = coopMatMulAdd(mat_a, mat_b, sum);
762            }
763        }
764#if defined(ACC_TYPE_MAX)
765        [[unroll]] for (uint i = 0; i < sum.length(); ++i) { sum[i] = clamp(sum[i], -ACC_TYPE_MAX, ACC_TYPE_MAX); }
766#endif
767
768        // Convert from ACC_TYPE to D_TYPE
769        coopmat<D_TYPE, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> mat_d;
770        mat_d = coopmat<D_TYPE, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(sum);
771
772#ifdef MUL_MAT_ID
773        // Call callback to store each element, remapping row through shared memory
774        coopMatPerElementNV(mat_d, mat_d, perElemOpD, ir, ic);
775#else
776        coopMatStoreTensorNV(mat_d, data_d, pos_d, sliceTensorLayoutNV(tensorLayoutD, ic * BN, BN, ir * BM, BM), tensorViewTranspose);
777#endif
778    }
779}