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 →
4.6 KiB151 linesraw
1#version 450
2
3#extension GL_EXT_control_flow_attributes : require
4#extension GL_EXT_shader_16bit_storage : require
5#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
6#extension GL_KHR_shader_subgroup_basic : enable
7#if USE_SUBGROUP_ADD
8#extension GL_KHR_shader_subgroup_arithmetic : enable
9#endif
10
11#define BINDING_IDX_K 0u
12
13#include "types.glsl"
14#include "fa_types.glsl"
15#define FaTypeV GGML_TYPE_F32
16
17layout(constant_id = 0) const uint FaTypeK = GGML_TYPE_F32;
18layout(constant_id = 1) const uint FaBlockBytesK = 4;
19layout(constant_id = 2) const uint SUBGROUP_SIZE = 32;
20
21#include "flash_attn_dequant.glsl"
22
23// one workgroup computes one output element, one invocation per head element
24#define HEAD_SIZE 128
25
26layout(local_size_x = HEAD_SIZE, local_size_y = 1, local_size_z = 1) in;
27
28layout(binding = 0) readonly buffer QBuf { float q[]; };
29layout(binding = 1) readonly buffer KBufF16 { float16_t k_f16[]; };
30layout(binding = 1) readonly buffer KBufF32 { float k_f32[]; };
31layout(binding = 1) readonly buffer KBufBF16 { uint16_t k_bf16[]; };
32layout(binding = 2) readonly buffer WBuf { float weights[]; };
33layout(binding = 3) readonly buffer MBuf { float16_t mask[]; };
34layout(binding = 4) writeonly buffer DstBuf { float dst[]; };
35
36layout(push_constant) uniform PushConstants {
37    uint n_kv;
38    uint n_heads;
39    uint n_tokens;
40    uint n_streams;
41    uint n_masks;
42    uint dispatch_x;
43    uint q_nb1;
44    uint q_nb2;
45    uint q_nb3;
46    uint k_nb2;
47    uint k_nb3;
48    uint w_nb1;
49    uint w_nb3;
50    uint m_nb1;
51    uint m_nb3;
52    uint d_nb1;
53    uint d_nb3;
54};
55
56shared float k_row[HEAD_SIZE];
57
58#if USE_SUBGROUP_ADD
59shared float sg_partials[HEAD_SIZE / SUBGROUP_SIZE];
60#else
61shared float partials[HEAD_SIZE];
62#endif
63
64void main() {
65    const uint tid = gl_LocalInvocationID.x;
66    const uint output_idx = gl_WorkGroupID.y * dispatch_x + gl_WorkGroupID.x;
67    const uint n_outputs = n_kv * n_tokens * n_streams;
68
69    if (fa_type_needs_shmem(FaTypeK)) {
70        init_iq_shmem(gl_WorkGroupSize);
71    }
72
73    if (output_idx >= n_outputs) {
74        return;
75    }
76
77    const uint ik = output_idx % n_kv;
78    const uint ts = output_idx / n_kv;
79    const uint t = ts % n_tokens;
80    const uint s = ts / n_tokens;
81    const uint k_offset = ik * k_nb2 + s * k_nb3;
82
83    // k strides come in as bytes, so scale them down to the view being indexed
84    const uint k_block_elems = fa_block_elems(FaTypeK);
85    const uint k_elem_bytes = FaBlockBytesK / k_block_elems;
86
87    if (FaTypeK == GGML_TYPE_F16) {
88        k_row[tid] = float(k_f16[k_offset / k_elem_bytes + tid]);
89    } else if (FaTypeK == GGML_TYPE_F32) {
90        k_row[tid] = k_f32[k_offset / k_elem_bytes + tid];
91    } else if (FaTypeK == GGML_TYPE_BF16) {
92        k_row[tid] = bf16_to_fp32(uint(k_bf16[k_offset / k_elem_bytes + tid]));
93    } else if (4 * tid < HEAD_SIZE) {
94        const uint coord = 4 * tid;
95        const uint ib = coord / k_block_elems;
96        const uint iqs = coord % k_block_elems;
97        const vec4 values = dequantize4(ib, iqs, k_offset / FaBlockBytesK, BINDING_IDX_K);
98        k_row[coord + 0] = values.x;
99        k_row[coord + 1] = values.y;
100        k_row[coord + 2] = values.z;
101        k_row[coord + 3] = values.w;
102    }
103    barrier();
104
105    const float k_val = k_row[tid];
106
107    float score = 0.0;
108    for (uint h = 0; h < n_heads; ++h) {
109        const float prod = q[h * q_nb1 + t * q_nb2 + s * q_nb3 + tid] * k_val;
110
111#if USE_SUBGROUP_ADD
112        const float sg_sum = subgroupAdd(prod);
113        if (gl_SubgroupInvocationID == 0) {
114            sg_partials[gl_SubgroupID] = sg_sum;
115        }
116        barrier();
117
118        if (tid == 0) {
119            float sum = 0.0;
120            [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) {
121                sum += sg_partials[i];
122            }
123            score += max(sum, 0.0) * weights[h + t * w_nb1 + s * w_nb3];
124        }
125        // the reads above must complete before the next iteration overwrites sg_partials
126        barrier();
127#else
128        partials[tid] = prod;
129        barrier();
130
131        [[unroll]] for (uint stride = HEAD_SIZE / 2; stride > 0; stride >>= 1) {
132            if (tid < stride) {
133                partials[tid] += partials[tid + stride];
134            }
135            barrier();
136        }
137
138        if (tid == 0) {
139            score += max(partials[0], 0.0) * weights[h + t * w_nb1 + s * w_nb3];
140        }
141        // the read of partials[0] above must complete before the next iteration
142        // overwrites partials[tid]
143        barrier();
144#endif
145    }
146
147    if (tid == 0) {
148        const uint mask_offset = ik + t * m_nb1 + (s % n_masks) * m_nb3;
149        dst[ik + t * d_nb1 + s * d_nb3] = score + float(mask[mask_offset]);
150    }
151}