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 →
21.8 KiB489 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#if defined(BFLOAT16)
12#extension GL_EXT_bfloat16 : enable
13#endif
14
15#extension GL_KHR_memory_scope_semantics : enable
16#extension GL_KHR_cooperative_matrix : enable
17#extension GL_NV_cooperative_matrix2 : enable
18#ifdef GL_NV_cooperative_matrix_decode_vector
19#extension GL_NV_cooperative_matrix_decode_vector : enable
20#endif
21#extension GL_EXT_buffer_reference : enable
22#extension GL_KHR_shader_subgroup_ballot : enable
23#extension GL_KHR_shader_subgroup_vote : enable
24#extension GL_EXT_null_initializer : enable
25
26#include "types.glsl"
27#include "flash_attn_base.glsl"
28#if !defined(BFLOAT16)
29#include "dequant_funcs_cm2.glsl"
30#endif
31
32// buffer_reference stride = sizeof(struct) = FaBlockBytesK/V.
33layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_K {
34    uint8_t raw[FaBlockBytesK];
35};
36layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_V {
37    uint8_t raw[FaBlockBytesV];
38};
39
40#if !defined(BFLOAT16)
41float16_t faDecodeK(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
42    switch (FaTypeK) {
43        case GGML_TYPE_F32:  return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock);
44        case GGML_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
45        case GGML_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
46        case GGML_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
47        case GGML_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
48        case GGML_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
49        case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
50        case GGML_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
51        default: return float16_t(0);
52    }
53}
54
55float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
56    switch (FaTypeV) {
57        case GGML_TYPE_F32:  return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock);
58        case GGML_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
59        case GGML_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
60        case GGML_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
61        case GGML_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
62        case GGML_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
63        case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
64        case GGML_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
65        default: return float16_t(0);
66    }
67}
68
69// V=4 vector decode for K/V; dispatches to per-format _v decoders.
70f16vec4 faDecodeKVector(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
71    switch (FaTypeK) {
72        case GGML_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
73        case GGML_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
74        case GGML_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
75        case GGML_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
76        case GGML_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
77        case GGML_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
78        case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
79        default: return f16vec4(0);
80    }
81}
82
83f16vec4 faDecodeVVector(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
84    switch (FaTypeV) {
85        case GGML_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
86        case GGML_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
87        case GGML_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
88        case GGML_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
89        case GGML_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
90        case GGML_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
91        case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
92        default: return f16vec4(0);
93    }
94}
95
96#ifdef GL_NV_cooperative_matrix_decode_vector
97#define FADECODEK , faDecodeK, faDecodeKVector
98#define FADECODEV , faDecodeV, faDecodeVVector
99#else
100#define FADECODEK , faDecodeK
101#define FADECODEV , faDecodeV
102#endif
103#endif
104
105layout (binding = 0) readonly buffer Q {uint8_t data_q[];};
106layout (binding = 1) readonly buffer K {uint8_t data_k[];};
107layout (binding = 2) readonly buffer V {uint8_t data_v[];};
108layout (binding = 3) readonly buffer M {uint8_t data_m[];};
109
110ACC_TYPE maxReduce(const in ACC_TYPE x, const in ACC_TYPE y) {
111    return max(x, y);
112}
113
114float16_t maxReduceFp16(const in float16_t x, const in float16_t y) {
115    return max(x, y);
116}
117
118ACC_TYPE smearReduce(const in ACC_TYPE x, const in ACC_TYPE y) {
119    return x;
120}
121
122// Replace matrix elements >= numRows or numCols with 'replace'
123ACC_TYPE replacePadding(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem, const in ACC_TYPE replace, const in uint32_t numRows, const in uint32_t numCols) {
124    if (row >= numRows || col >= numCols) {
125        return replace;
126    }
127    return elem;
128}
129
130ACC_TYPE Exp(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem)
131{
132    return exp(elem);
133}
134
135ACC_TYPE Max(const in uint32_t row, const in uint32_t col, const in ACC_TYPE elem0, const in ACC_TYPE elem1)
136{
137    return max(elem0, elem1);
138}
139
140// Store the output when doing grouped query attention.
141// Rows index by Q's dimension 2, and the first N rows are valid.
142D_TYPE perElemOpGqaStore(const in uint32_t r, const in uint32_t c, const in D_TYPE elem, const in uint32_t o_offset, const in uint32_t iq2, const in uint32_t N)
143{
144    if (r < N && c < HSV) {
145        uint32_t offset = (iq2 + r) * HSV + c;
146        data_o[o_offset + offset] = D_TYPE(elem);
147    }
148    return elem;
149}
150
151// Store O values for non-GQA split_k. Rows are tokens, not heads.
152D_TYPE perElemOpNonGqaSplitKStore(const in uint32_t r, const in uint32_t c, const in D_TYPE elem, const in uint32_t unused, const in uint32_t iq2, const in uint32_t N) {
153    uint32_t global_row = i * Br + r;
154    if (global_row < N && c < HSV) {
155        uint32_t o_off = HSV * p.ne1
156            * (split_k_index + p.k_num * (global_row + p.ne2 * iq3));
157        data_o[o_off + iq2 * HSV + c] = D_TYPE(elem);
158    }
159    return elem;
160}
161
162// Store L/M values for non-GQA split_k.
163ACC_TYPE perElemOpNonGqaSplitKStoreCol0(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t lm_base, const in uint32_t iq2, const in uint32_t N) {
164    uint32_t global_row = i * Br + r;
165    if (global_row < N && c == 0) {
166        uint32_t lm_off = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3
167            + p.ne1 * 2 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3));
168        data_o[lm_off + lm_base + iq2] = D_TYPE(elem);
169    }
170    return elem;
171}
172
173void main() {
174#ifdef NEEDS_INIT_IQ_SHMEM
175    if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) {
176        init_iq_shmem(gl_WorkGroupSize);
177    }
178#endif
179
180    init_indices();
181
182    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutQ = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
183    tensorLayoutNV<2, Clamp> tensorLayoutK = createTensorLayoutNV(2, Clamp);
184    tensorLayoutNV<2, Clamp> tensorLayoutV = createTensorLayoutNV(2, Clamp);
185
186    tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0);
187
188    const uint bs_k = fa_block_elems(FaTypeK);
189    const uint bs_v = fa_block_elems(FaTypeV);
190    tensorLayoutK = setTensorLayoutBlockSizeNV(tensorLayoutK, 1, bs_k);
191    tensorLayoutV = setTensorLayoutBlockSizeNV(tensorLayoutV, 1, bs_v);
192
193    tensorLayoutQ = setTensorLayoutDimensionNV(tensorLayoutQ, N, HSK);
194    tensorLayoutK = setTensorLayoutDimensionNV(tensorLayoutK, KV, HSK);
195    tensorLayoutV = setTensorLayoutDimensionNV(tensorLayoutV, KV, HSV);
196
197    // hint to the compiler that strides are aligned for the aligned variant of the shader
198    if (Clamp != gl_CooperativeMatrixClampModeConstantNV)
199    {
200        q_stride &= ~7;
201        if (bs_k == 1u) {
202            k_stride &= ~7;
203        }
204        if (bs_v == 1u) {
205            v_stride &= ~7;
206        }
207        m_stride &= ~7;
208    }
209    tensorLayoutQ = setTensorLayoutStrideNV(tensorLayoutQ, q_stride, 1);
210    tensorLayoutK = setTensorLayoutStrideNV(tensorLayoutK, k_stride, 1);
211    tensorLayoutV = setTensorLayoutStrideNV(tensorLayoutV, v_stride, 1);
212
213    coopmat<Q_TYPE, gl_ScopeWorkgroup, Br, HSK_pad, gl_MatrixUseAccumulator> Q;
214    coopmat<FLOAT_TYPE, gl_ScopeWorkgroup, Br, HSK_pad, gl_MatrixUseA> Qf16;
215
216    uint32_t q_offset = gqa_iq1*p.nb01*4/*sizeof(float)*/ + iq2*p.nb02+iq3*p.nb03;
217    coopMatLoadTensorNV(Q, data_q, q_offset, sliceTensorLayoutNV(tensorLayoutQ, i * Br, Br, 0, HSK_pad));
218
219    Q *= Q_TYPE(p.scale);
220    Qf16 = coopmat<FLOAT_TYPE, gl_ScopeWorkgroup, Br, HSK_pad, gl_MatrixUseA>(Q);
221
222    coopmat<O_TYPE, gl_ScopeWorkgroup, Br, HSV_pad, gl_MatrixUseAccumulator> O = coopmat<O_TYPE, gl_ScopeWorkgroup, Br, HSV_pad, gl_MatrixUseAccumulator>(0);
223
224    coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator> L, M;
225
226    // Use -FLT_MAX/2 rather than -inf to reduce the possibility of NaNs, e.g. when computing Mold-M.
227    const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF);
228
229    L = coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator>(0);
230#if defined(ACC_TYPE_MAX)
231    M = coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator>(-ACC_TYPE_MAX / ACC_TYPE(2));
232#else
233    M = coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator>(NEG_FLT_MAX_OVER_2);
234#endif
235
236    coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator> slopeMat = coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator>(1.0);
237
238    // ALiBi
239    if (p.max_bias > 0.0f) {
240        coopMatPerElementNV(slopeMat, slopeMat, perElemOpComputeSlope, iq2);
241    }
242
243    const uint32_t mo_stride = CEIL_DIV(KV, 16 * Bc);
244    // mo_offset will point to the tile starting at row i*Br and col 0
245    uint32_t mo_offset = mo_stride * i;
246
247    uint32_t m_offset = gqa_iq1*KV * 2 /*sizeof(float16_t)*/;
248    if (p.nem2 != 1 || p.nem3 != 1) {
249        m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * KV * 2 /*sizeof(float16_t)*/;
250        mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride;
251    }
252
253    uint32_t mask_opt = 0;
254    uint32_t mask_opt_idx = ~0;
255
256    [[dont_unroll]]
257    for (uint32_t j = start_j; j < end_j; ++j) {
258
259        coopmat<float16_t, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator> mv = coopmat<float16_t, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator>(0);
260        if (MASK_ENABLE) {
261
262            if (USE_MASK_OPT && mask_opt_idx != j / 16) {
263                mask_opt_idx = j / 16;
264                mask_opt = data_mask_opt[mo_offset + mask_opt_idx];
265            }
266            uint32_t mask_opt_bits = (mask_opt >> ((j % 16) * 2)) & 0x3;
267            if (mask_opt_bits == MASK_OPT_ALL_NEG_INF) {
268                // skip this block
269                continue;
270            }
271            // Only load if the block is not all zeros
272            if (mask_opt_bits != MASK_OPT_ALL_ZERO) {
273                bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0;
274
275                if (nem1_bounds_check) {
276                    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutM = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
277                    tensorLayoutM = setTensorLayoutDimensionNV(tensorLayoutM, p.nem1, KV);
278                    tensorLayoutM = setTensorLayoutStrideNV(tensorLayoutM, m_stride, 1);
279                    tensorLayoutM = setTensorLayoutClampValueNV(tensorLayoutM, 0xfc00); // -inf in float16_t
280
281                    coopmat<float16_t, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator> mvmax;
282
283                    coopMatLoadTensorNV(mv, data_m, m_offset, sliceTensorLayoutNV(tensorLayoutM, i * Br, Br, j * Bc, Bc));
284                    // skip the block if the mask is entirely -inf
285                    coopMatReduceNV(mvmax, mv, gl_CooperativeMatrixReduceRowAndColumnNV, maxReduceFp16);
286                    if (mvmax[0] <= NEG_FLT_MAX_OVER_2) {
287                        continue;
288                    }
289                } else {
290                    tensorLayoutNV<2, Clamp> tensorLayoutM = createTensorLayoutNV(2, Clamp);
291                    // Don't clamp against nem1 when GQA is enabled
292                    uint32_t m_height = p.gqa_ratio > 1 ? ~0 : p.nem1;
293                    tensorLayoutM = setTensorLayoutDimensionNV(tensorLayoutM, m_height, KV);
294                    tensorLayoutM = setTensorLayoutStrideNV(tensorLayoutM, m_stride, 1);
295
296                    coopmat<float16_t, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator> mvmax;
297
298                    coopMatLoadTensorNV(mv, data_m, m_offset, sliceTensorLayoutNV(tensorLayoutM, i * Br, Br, j * Bc, Bc));
299                    // skip the block if the mask is entirely -inf
300                    coopMatReduceNV(mvmax, mv, gl_CooperativeMatrixReduceRowAndColumnNV, maxReduceFp16);
301                    if (mvmax[0] <= NEG_FLT_MAX_OVER_2) {
302                        continue;
303                    }
304                }
305            }
306        }
307
308        coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator> S = coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator>(0);
309
310        coopmat<FLOAT_TYPE, gl_ScopeWorkgroup, HSK_pad, Bc, gl_MatrixUseB> K_T;
311
312        uint32_t k_offset = ik2*p.nb12 + ik3*p.nb13;
313        // F16: bs_k==1 (direct load). F32: bs_k==4 (vec4 / dequantFuncF32). Quantized types: bs_k==32.
314#if defined(BFLOAT16)
315        coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose);
316#else
317        const bool k_use_decode = (bs_k > 1u);
318        if (k_use_decode) {
319            coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FADECODEK);
320        } else {
321            coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose);
322        }
323#endif
324        S = coopMatMulAdd(Qf16, K_T, S);
325
326        if (LOGIT_SOFTCAP) {
327            [[unroll]]
328            for (int k = 0; k < S.length(); ++k) {
329                S[k] = ACC_TYPE(p.logit_softcap)*tanh(S[k]);
330            }
331        }
332
333        if (MASK_ENABLE) {
334            S += slopeMat*coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator>(mv);
335        }
336
337        // Clear padding elements to -inf, so they don't contribute to rowmax
338        if (Clamp != 0 &&
339            ((j + 1) * Bc > KV ||
340             (i + 1) * Br > N)) {
341
342            uint R = ((i + 1) * Br >  N) ?  (N % Br) : Br;
343            uint C = ((j + 1) * Bc > KV) ? (KV % Bc) : Bc;
344
345            coopMatPerElementNV(S, S, replacePadding, ACC_TYPE(NEG_FLT_MAX_OVER_2), R, C);
346        }
347
348        coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator> rowmax, P, rowsum, eM;
349
350        coopMatReduceNV(rowmax, S, gl_CooperativeMatrixReduceRowNV, maxReduce);
351
352        rowmax += coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator>(FATTN_KQ_MAX_OFFSET);
353
354        coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator> Mold = M;
355
356        // M = max(rowmax, Mold)
357        // P = e^(S - M)
358        // eM = e^(Mold - M)
359        coopMatPerElementNV(M, rowmax, Max, Mold);
360        coopMatPerElementNV(P, S - M, Exp);
361        coopMatPerElementNV(eM, Mold - M, Exp);
362
363        // Clear padding elements to 0, so they don't contribute to rowsum
364        if (Clamp != 0 &&
365            ((j + 1) * Bc > KV ||
366             (i + 1) * Br > N)) {
367
368            uint R = ((i + 1) * Br >  N) ?  (N % Br) : Br;
369            uint C = ((j + 1) * Bc > KV) ? (KV % Bc) : Bc;
370
371            coopMatPerElementNV(P, P, replacePadding, ACC_TYPE(0.0), R, C);
372        }
373
374        coopmat<FLOAT_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseA> P_A = coopmat<FLOAT_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseA>(P);
375
376        // compute rowsum by multiplying by matrix of all ones.
377        coopmat<FLOAT_TYPE, gl_ScopeWorkgroup, Bc, Bc, gl_MatrixUseB> One = coopmat<FLOAT_TYPE, gl_ScopeWorkgroup, Bc, Bc, gl_MatrixUseB>(1.0);
378
379        rowsum = coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator>(0.0);
380        rowsum = coopMatMulAdd(P_A, One, rowsum);
381
382        coopmat<FLOAT_TYPE, gl_ScopeWorkgroup, Bc, HSV_pad, gl_MatrixUseB> V;
383        uint32_t v_offset = iv2*p.nb22 + iv3*p.nb23;
384#if defined(BFLOAT16)
385        coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad));
386#else
387        const bool v_use_decode = (bs_v > 1u);
388        if (v_use_decode) {
389            coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad) FADECODEV);
390        } else {
391            coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad));
392        }
393#endif
394
395        L = eM*L + rowsum;
396
397        // This is the "diagonal" matrix in the paper, but since we do componentwise
398        // multiply rather than matrix multiply it has the diagonal element smeared
399        // across the row
400        coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, HSV_pad, gl_MatrixUseAccumulator> eMdiag;
401
402        // resize eM by using smear/reduce
403        coopMatReduceNV(eMdiag, eM, gl_CooperativeMatrixReduceRowNV, smearReduce);
404
405        O *= coopmat<O_TYPE, gl_ScopeWorkgroup, Br, HSV_pad, gl_MatrixUseAccumulator>(eMdiag);
406        O = coopMatMulAdd(P_A, V, O);
407    }
408
409    // If there is split_k, then the split_k resolve shader does the final
410    // division by L. Store the intermediate O value and per-row m and L values.
411    if (p.k_num > 1) {
412        coopmat<D_TYPE, gl_ScopeWorkgroup, Br, HSV_pad, gl_MatrixUseAccumulator> O_D = coopmat<D_TYPE, gl_ScopeWorkgroup, Br, HSV_pad, gl_MatrixUseAccumulator>(O);
413
414        if (p.gqa_ratio > 1) {
415            // note: O and Q have swapped coord 1,2.
416            uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3));
417            coopMatPerElementNV(O_D, O_D, perElemOpGqaStore, o_offset, iq2, N);
418
419            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));
420            coopMatPerElementNV(L, L, perElemOpStoreCol0, o_offset, iq2, N);
421            coopMatPerElementNV(M, M, perElemOpStoreCol0, o_offset + p.ne1, iq2, N);
422        } else {
423            coopMatPerElementNV(O_D, O_D, perElemOpNonGqaSplitKStore, 0u, iq2, N);
424            coopMatPerElementNV(L, L, perElemOpNonGqaSplitKStoreCol0, 0u, iq2, N);
425            coopMatPerElementNV(M, M, perElemOpNonGqaSplitKStoreCol0, p.ne1, iq2, N);
426        }
427        return;
428    }
429
430    coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, HSV_pad, gl_MatrixUseAccumulator> Ldiag;
431
432    // resize L by using smear/reduce
433    coopMatReduceNV(Ldiag, L, gl_CooperativeMatrixReduceRowNV, smearReduce);
434
435    if ((p.mask_n_head_log2 & SINK_ENABLE_BIT) != 0) {
436        coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, HSV_pad, gl_MatrixUseAccumulator> S;
437        coopMatPerElementNV(S, S, perElemOpGetSink, iq2);
438
439        coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, HSV_pad, gl_MatrixUseAccumulator> Mr;
440
441        // resize M by using smear/reduce
442        coopMatReduceNV(Mr, M, gl_CooperativeMatrixReduceRowNV, smearReduce);
443
444        // O, Ldiag, Mr all have the same type so all element locations match
445        [[unroll]] for (uint32_t i = 0; i < Ldiag.length(); ++i) {
446            ACC_TYPE sink = S[i];
447
448            ACC_TYPE ms = ACC_TYPE(1.0f);
449            ACC_TYPE vs = ACC_TYPE(1.0f);
450
451            if (sink > Mr[i]) {
452                ms = exp(Mr[i] - sink);
453
454                O[i] *= O_TYPE(ms);
455            } else {
456                vs = exp(sink - Mr[i]);
457            }
458
459            Ldiag[i] = Ldiag[i]*ms + vs;
460        }
461    }
462
463    [[unroll]]
464    for (int k = 0; k < Ldiag.length(); ++k) {
465        Ldiag[k] = (Ldiag[k] == 0.0) ? ACC_TYPE(0.0) : (ACC_TYPE(1.0) / Ldiag[k]);
466    }
467
468    coopmat<D_TYPE, gl_ScopeWorkgroup, Br, HSV_pad, gl_MatrixUseAccumulator> O_D = coopmat<D_TYPE, gl_ScopeWorkgroup, Br, HSV_pad, gl_MatrixUseAccumulator>(O);
469
470    O_D = coopmat<D_TYPE, gl_ScopeWorkgroup, Br, HSV_pad, gl_MatrixUseAccumulator>(Ldiag)*O_D;
471
472#if defined(ACC_TYPE_MAX)
473    [[unroll]] for (uint i = 0; i < O_D.length(); ++i) { O_D[i] = clamp(O_D[i], D_TYPE(-ACC_TYPE_MAX), D_TYPE(ACC_TYPE_MAX)); }
474#endif
475
476    uint32_t o_offset = gqa_iq1*p.ne1*HSV + iq3*p.ne2*p.ne1*HSV;
477
478    if (p.gqa_ratio > 1) {
479        coopMatPerElementNV(O_D, O_D, perElemOpGqaStore, o_offset, iq2, N);
480    } else {
481        tensorLayoutNV<3, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(3, gl_CooperativeMatrixClampModeConstantNV);
482        tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.ne2, p.ne1, HSV);
483
484        // permute dimensions
485        tensorViewNV<3, false, 1, 0, 2> tensorViewPermute = createTensorViewNV(3, false, 1, 0, 2);
486
487        coopMatStoreTensorNV(O_D, data_o, o_offset, sliceTensorLayoutNV(tensorLayoutD, i * Br, Br, iq2, N, 0, HSV_pad), tensorViewPermute);
488    }
489}