char-fork/llama.cpp

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

Jeff Bolzvulkan: combine duplicated fastdiv functions, rename the one optimizing small divs (#27526)c9ca51c1f

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
4.4 KiB135 linesraw
1#version 450
2
3#extension GL_EXT_control_flow_attributes : enable
4
5#ifdef USE_SUBGROUPS
6#extension GL_KHR_shader_subgroup_basic : enable
7#extension GL_KHR_shader_subgroup_arithmetic : enable
8#endif
9
10#include "types.glsl"
11#include "utils.glsl"
12
13layout (push_constant) uniform parameter
14{
15    uint32_t ne00;
16    uint32_t ne01;
17    uint32_t nb00;
18    uint32_t nb01;
19    uint32_t a_offset;
20    uint32_t n_experts;
21    uint32_t hoist_row_ids;
22    uint32_t ne00mp;
23    uint32_t ne00L;
24} p;
25
26#define BLOCK_SIZE 256
27
28layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
29
30layout (binding = 0) readonly buffer A {uint data_a[];};
31layout (binding = 1) writeonly buffer D {uint data_d[];};
32
33shared uint vals[BLOCK_SIZE];
34shared uint offsets[BLOCK_SIZE];
35shared uint cursors[BLOCK_SIZE];
36
37// data_d layout when p.hoist_row_ids is set:
38//   [0,              n_experts)   per-expert row count
39//   [n_experts,    2*n_experts)   per-expert start offset into the row id region
40//   [2*n_experts]                 total row count
41//   [2*n_experts + 1,         )   row ids grouped by expert, packed as (i01 << 16) | (i00 & 0xffff)
42// Otherwise only data_d[expert_id] is written, holding that expert's row count.
43void main() {
44    const uint expert_id = gl_WorkGroupID.x;
45    const uint num_elements = p.ne00 * p.ne01;
46    const uint tid = gl_LocalInvocationID.x;
47
48    if (p.hoist_row_ids != 0) {
49        if (tid < p.n_experts) {
50            vals[tid] = 0;
51        }
52        barrier();
53
54        for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
55            const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
56            const uint i00 = idx - i01 * p.ne00;
57            const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
58            if (expert < p.n_experts) {
59                atomicAdd(vals[expert], 1);
60            }
61        }
62        barrier();
63
64#ifdef USE_SUBGROUPS
65        if (gl_SubgroupID == 0) {
66            // pad the trip count so the subgroup ops stay in uniform control flow
67            const uint n_experts_padded = (p.n_experts + gl_SubgroupSize - 1) & ~(gl_SubgroupSize - 1);
68            uint base = 0;
69            for (uint expert = gl_SubgroupInvocationID; expert < n_experts_padded; expert += gl_SubgroupSize) {
70                const bool in_range = expert < p.n_experts;
71                const uint count = in_range ? vals[expert] : 0;
72                const uint offset = base + subgroupExclusiveAdd(count);
73                if (in_range) {
74                    data_d[expert] = count;
75                    data_d[p.n_experts + expert] = offset;
76                    offsets[expert] = offset;
77                    cursors[expert] = 0;
78                }
79                base += subgroupAdd(count);
80            }
81            if (subgroupElect()) {
82                data_d[2 * p.n_experts] = base;
83            }
84        }
85#else
86        if (tid == 0) {
87            uint offset = 0;
88            for (uint expert = 0; expert < p.n_experts; ++expert) {
89                const uint count = vals[expert];
90                data_d[expert] = count;
91                data_d[p.n_experts + expert] = offset;
92                offsets[expert] = offset;
93                cursors[expert] = 0;
94                offset += count;
95            }
96            data_d[2 * p.n_experts] = offset;
97        }
98#endif
99        barrier();
100
101        for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
102            const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
103            const uint i00 = idx - i01 * p.ne00;
104            const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
105            if (expert < p.n_experts) {
106                const uint row = atomicAdd(cursors[expert], 1);
107                const uint packed_row_id = (i01 << 16) | (i00 & 0xffffu);
108                data_d[2 * p.n_experts + 1 + offsets[expert] + row] = packed_row_id;
109            }
110        }
111        return;
112    }
113
114    uint count = 0;
115    for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
116        const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
117        const uint i00 = idx - i01 * p.ne00;
118        const uint a = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
119
120        count += uint(a == expert_id);
121    }
122
123    vals[tid] = count;
124    barrier();
125    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
126        if (tid < s) {
127            vals[tid] += vals[tid + s];
128        }
129        barrier();
130    }
131
132    if (tid == 0) {
133        data_d[expert_id] = vals[0];
134    }
135}