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 →
7.0 KiB223 linesraw
1
2layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
3
4layout (constant_id =  0) const uint32_t WorkGroupSize = 128;
5layout (constant_id =  1) const uint32_t Br = 1;
6layout (constant_id =  2) const uint32_t Bc = 32;
7layout (constant_id =  3) const uint32_t HSK = 32;
8layout (constant_id =  4) const uint32_t HSV = 32;
9layout (constant_id =  5) const uint32_t Clamp = 0;
10layout (constant_id =  6) const uint32_t D_split = 16;
11layout (constant_id =  7) const uint32_t row_split = 1;
12layout (constant_id =  8) const uint32_t SubGroupSize = 32;
13layout (constant_id =  9) const uint32_t SHMEM_STAGING = 0;
14layout (constant_id = 10) const uint32_t Flags = 0;
15layout (constant_id = 11) const uint32_t LIMIT_OCCUPANCY_SHMEM = 0;
16// ggml_type enumerant for K/V
17layout (constant_id = 12) const uint32_t FaTypeK = 0;
18layout (constant_id = 13) const uint32_t FaTypeV = 0;
19// sizeof(decode buffer): quants -> ggml block size; F32 -> 16 (decodeBufF32 vec4).
20layout (constant_id = 14) const uint32_t FaBlockBytesK = 2;
21layout (constant_id = 15) const uint32_t FaBlockBytesV = 2;
22
23const bool USE_MASK_OPT    = (Flags & 1) != 0;
24const bool MASK_ENABLE     = (Flags & 2) != 0;
25const bool LOGIT_SOFTCAP   = (Flags & 4) != 0;
26const bool OLD_AMD_WINDOWS = (Flags & 8) != 0;
27
28// Round up head sizes to a multiple of 16, for coopmat1/coopmat2 paths
29const uint32_t HSK_pad = (HSK + 15) & ~15;
30const uint32_t HSV_pad = (HSV + 15) & ~15;
31
32const bool KV_bounds_check = Clamp != 0;
33
34layout (push_constant) uniform parameter {
35    uint32_t N;
36    uint32_t KV;
37
38    uint32_t ne1;
39    uint32_t ne2;
40    uint32_t ne3;
41
42    uint32_t neq2;
43    uint32_t neq3;
44    uint32_t nek2;
45    uint32_t nek3;
46    uint32_t nev2;
47    uint32_t nev3;
48    uint32_t nem1;
49    uint32_t nem2;
50    uint32_t nem3;
51
52    uint32_t nb01;
53    uint32_t nb02;
54    uint32_t nb03;
55    uint32_t nb11;
56    uint32_t nb12;
57    uint32_t nb13;
58    uint32_t nb21;
59    uint32_t nb22;
60    uint32_t nb23;
61
62    float scale;
63    float max_bias;
64    float logit_softcap;
65
66    uint32_t mask_n_head_log2;
67    float m0;
68    float m1;
69
70    uint32_t gqa_ratio;
71    uint32_t split_kv;
72    uint32_t k_num;
73} p;
74
75#define SINK_ENABLE_BIT (1<<24)
76#define N_LOG2_MASK 0xFFFF
77
78layout (binding = 4) readonly buffer S {float data_s[];};
79
80layout (binding = 5) writeonly buffer O {D_TYPE data_o[];};
81layout (binding = 5) writeonly buffer OV4 {D_TYPEV4 data_ov4[];};
82
83layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
84
85#define MASK_OPT_ALL_NEG_INF 1
86#define MASK_OPT_ALL_ZERO 2
87
88#define BINDING_IDX_K 0
89#define BINDING_IDX_V 1
90
91#include "fa_types.glsl"
92
93#if defined(BFLOAT16)
94#define O_TYPE float
95#define O_TYPEV4 vec4
96#else
97#define O_TYPE FLOAT_TYPE
98#define O_TYPEV4 FLOAT_TYPEV4
99#endif
100
101// These can't be `const` globals because GLSL forbids function calls in global
102// const initializers, even when the spec constants would let the driver fold
103// them. Macros expand at the use site and fold after specialization.
104#define BLOCK_SIZE_K fa_block_elems(FaTypeK)
105#define BLOCK_SIZE_V fa_block_elems(FaTypeV)
106// F16 reads f16 elements directly from the binding; everything else routes
107// through dequantize4 / the MMQ helpers to unpack from the packed block layout.
108#define USE_DECODE_K (FaTypeK != GGML_TYPE_F16)
109#define USE_DECODE_V (FaTypeV != GGML_TYPE_F16)
110
111#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b))
112
113
114// Store column zero. This is used to save per-row m and L values for split_k.
115ACC_TYPE perElemOpStoreCol0(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t o_offset, const in uint32_t iq2, const in uint32_t N)
116{
117    if (r < N && c == 0) {
118        uint32_t offset = iq2 + r;
119        data_o[o_offset + offset] = D_TYPE(elem);
120    }
121    return elem;
122}
123
124// Load the slope matrix, indexed by Q's dimension 2.
125ACC_TYPE perElemOpComputeSlope(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t iq2)
126{
127    const uint32_t h = iq2 + (r % p.gqa_ratio);
128
129    uint32_t n_head_log2 = p.mask_n_head_log2 & N_LOG2_MASK;
130
131    const ACC_TYPE base = ACC_TYPE(h < n_head_log2 ? p.m0 : p.m1);
132    const int      exph = int(h < n_head_log2 ? h + 1 : 2*(h - n_head_log2) + 1);
133
134    return ACC_TYPE(pow(base, ACC_TYPE(exph)));
135}
136
137// Load the sink value, indexed by Q's dimension 2.
138ACC_TYPE perElemOpGetSink(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t iq2)
139{
140    const uint32_t h = iq2 + (r % p.gqa_ratio);
141
142    return ACC_TYPE(data_s[h]);
143}
144
145uint32_t i, N, KV, split_k_index, Tr, start_j, end_j,
146         gqa_iq1, iq2, iq3, rk2, rk3, rv2, rv3, ik2, ik3, iv2, iv3,
147         q_stride, k_stride, v_stride, m_stride;
148
149void init_indices()
150{
151    N = p.N;
152    KV = p.KV;
153
154    if (p.k_num > 1) {
155        if (p.gqa_ratio > 1) {
156            i = 0;
157            // batch and split_k share gl_WorkGroupID.x
158            gqa_iq1 = gl_WorkGroupID.x / p.k_num;
159            split_k_index = gl_WorkGroupID.x % p.k_num;
160        } else {
161            gqa_iq1 = 0;
162            split_k_index = gl_WorkGroupID.x % p.k_num;
163            i = gl_WorkGroupID.x / p.k_num;
164        }
165    } else if (p.gqa_ratio > 1) {
166        i = 0;
167        gqa_iq1 = gl_WorkGroupID.x;
168        split_k_index = 0;
169    } else {
170        i = gl_WorkGroupID.x;
171        gqa_iq1 = 0;
172        split_k_index = 0;
173    }
174
175    Tr = CEIL_DIV(N, Br);
176
177    start_j = split_k_index * p.split_kv / Bc;
178    end_j = CEIL_DIV(min(KV, (split_k_index + 1) * p.split_kv), Bc);
179
180    // When not using grouped query attention, all rows share the same iq2, equal to gl_WorkGroupID.y.
181    // When using grouped query attention, each workgroup does gqa_ratio consecutive values of iq2.
182    iq2 = gl_WorkGroupID.y * p.gqa_ratio;
183    iq3 = gl_WorkGroupID.z;
184
185    // broadcast factors
186    rk2 = p.neq2/p.nek2;
187    rk3 = p.neq3/p.nek3;
188
189    rv2 = p.neq2/p.nev2;
190    rv3 = p.neq3/p.nev3;
191
192    // k indices
193    ik3 = iq3 / rk3;
194    ik2 = iq2 / rk2;
195
196    // v indices
197    iv3 = iq3 / rv3;
198    iv2 = iq2 / rv2;
199
200    // nb?1 are already divided by the type size and are in units of elements.
201    // When using grouped query attention, Q is indexed by iq2, so the stride
202    // should be nb02 (which is in bytes).
203    q_stride = p.gqa_ratio > 1 ? (p.nb02 / 4) : p.nb01;
204    k_stride = p.nb11;
205    v_stride = p.nb21;
206    // When using grouped query attention, all rows use the same mask (stride 0).
207    // "p.gqa_ratio >> 16" is just a roundabout way of writing zero
208    // that prevents the compiler from folding the "&" through the select
209    // and breaking the alignment detection.
210    m_stride = (p.gqa_ratio > 1) ? (p.gqa_ratio >> 16) : KV;
211}
212
213// Bias applied to softmax to stay in fp16 range.
214// Based on ggml-cuda issue https://github.com/ggml-org/llama.cpp/issues/18606
215const float FATTN_KQ_MAX_OFFSET = 3.0f*0.6931f;
216
217// Store the output when doing grouped query attention.
218// Rows index by Q's dimension 2, and the first N rows are valid.
219void gqaStore(const in uint32_t r, const in uint32_t c, const in O_TYPEV4 elems, const in uint32_t o_offset, const in uint32_t iq2, const in uint32_t N)
220{
221    uint32_t offset = (iq2 + r) * HSV / 4 + c;
222    data_ov4[o_offset + offset] = D_TYPEV4(elems);
223}