1#version 450
2
3#extension GL_EXT_control_flow_attributes : enable
4
5layout (push_constant) uniform parameter
6{
7 uint KX;
8 uint KY;
9 uint ne00;
10 uint ne01;
11 uint ne02;
12 uint ne12;
13 uint ne13;
14 uint nb11;
15 uint nb12;
16 uint nb13;
17 float scale;
18 float max_bias;
19 float m0;
20 float m1;
21 uint n_head_log2;
22 uint nrows_x;
23 uint has_sinks;
24} p;
25
26#include "types.glsl"
27
28layout(constant_id = 0) const uint BLOCK_SIZE = 32;
29layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
30
31layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
32layout (binding = 1) readonly buffer Y {B_TYPE data_b[];};
33layout (binding = 2) readonly buffer Z {float data_c[];};
34layout (binding = 3) buffer D {D_TYPE data_d[];};
35
36shared FLOAT_TYPE vals[BLOCK_SIZE];
37
38// num_iters is the number of BLOCK_SIZE loop iterations we need to iterate
39// over all the columns. The main function tries to pass a constant here,
40// as if it were a template function, to allow unrolling.
41void soft_max(uint num_iters) {
42 const uint tid = gl_LocalInvocationID.x;
43 const uint rowx = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
44
45 const uint32_t i03 = rowx / (p.ne01 * p.ne02);
46 const uint32_t i02 = (rowx - i03 * p.ne01 * p.ne02) / p.ne01;
47 const uint32_t i01 = rowx % p.ne01;
48
49 uint rowy_start = 0;
50 if (p.KY > 0) {
51 rowy_start = i01 * p.nb11 + (i02 % p.ne12) * p.nb12 + (i03 % p.ne13) * p.nb13;
52 }
53
54 if (rowx >= p.nrows_x) {
55 return;
56 }
57
58 float slope = 1.0f;
59
60 // ALiBi
61 if (p.max_bias > 0.0f) {
62 const uint h = (rowx / p.ne01) % p.ne02; // head index
63
64 const float base = h < p.n_head_log2 ? p.m0 : p.m1;
65 const uint exp = h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1;
66
67 slope = pow(base, exp);
68 }
69
70 // Find max
71 FLOAT_TYPE max_val = p.has_sinks == 0 ? uintBitsToFloat(0xFF800000) : data_c[i02];
72
73 // Cache values while we compute the max, so we don't need to read them
74 // again when we're ready to compute exp(x-max).
75 const uint DATA_CACHE_SIZE = 16;
76 FLOAT_TYPE data_cache[DATA_CACHE_SIZE];
77
78 [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) {
79 const uint col = col0 + tid;
80
81 FLOAT_TYPE a = FLOAT_TYPE(0);
82 if (col < p.KX) {
83 a = data_a[rowx * p.KX + col];
84 }
85
86 FLOAT_TYPE b = FLOAT_TYPE(0);
87 if (p.KY > 0 && col < p.KX) {
88 b = data_b[rowy_start + col];
89 }
90
91 FLOAT_TYPE v = a * p.scale + slope * b;
92
93 if (col < p.KX) {
94 max_val = max(max_val, v);
95 }
96
97 if (idx < DATA_CACHE_SIZE) {
98 data_cache[idx] = v;
99 }
100 }
101
102 // reduce across the workgroup
103 vals[tid] = max_val;
104 barrier();
105 [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
106 if (tid < s) {
107 vals[tid] = max(vals[tid], vals[tid + s]);
108 }
109 barrier();
110 }
111
112 max_val = vals[0];
113 barrier();
114
115 FLOAT_TYPE sum = FLOAT_TYPE(0.0f);
116
117 // Compute sum{exp(x - max)}
118 [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) {
119 const uint col = col0 + tid;
120
121 if (col >= p.KX) {
122 break;
123 }
124
125 // compute exp(a*scale+b*slope), add it to sum, and cache the new value
126 // in data_cache if possible.
127 const uint i = rowx * p.KX + col;
128 FLOAT_TYPE val;
129 if (idx < DATA_CACHE_SIZE) {
130 val = exp(data_cache[idx] - max_val);
131 } else {
132 val = exp(FLOAT_TYPE(data_a[i]) * p.scale + (p.KY > 0 ? slope * FLOAT_TYPE(data_b[rowy_start + col]) : FLOAT_TYPE(0.0f)) - max_val);
133 }
134 sum += val;
135 if (idx < DATA_CACHE_SIZE) {
136 data_cache[idx] = val;
137 } else {
138 data_d[i] = D_TYPE(val);
139 }
140 }
141
142 // reduce across the workgroup
143 vals[tid] = sum;
144 barrier();
145 [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
146 if (tid < s) {
147 vals[tid] += vals[tid + s];
148 }
149 barrier();
150 }
151 sum = vals[0];
152
153 if (p.has_sinks != 0) {
154 sum += FLOAT_TYPE(exp(FLOAT_TYPE(data_c[i02]) - max_val));
155 }
156
157 FLOAT_TYPE rcpdivisor = 1.0/sum;
158
159 [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) {
160 const uint col = col0 + tid;
161
162 if (col >= p.KX) {
163 continue;
164 }
165
166 if (idx < DATA_CACHE_SIZE) {
167 data_d[rowx*p.KX + col] = D_TYPE(data_cache[idx] * rcpdivisor);
168 } else {
169 data_d[rowx*p.KX + col] *= D_TYPE(rcpdivisor);
170 }
171 }
172}
173
174void main() {
175 // instantiate the soft_max function for several different
176 // dimensions, to allow loop unrolling
177 uint num_blocks = (p.KX + BLOCK_SIZE - 1) / BLOCK_SIZE;
178 if (num_blocks > 32) {
179 soft_max(num_blocks);
180 } else if (num_blocks > 16) {
181 soft_max(32);
182 } else if (num_blocks > 8) {
183 soft_max(16);
184 } else if (num_blocks > 4) {
185 soft_max(8);
186 } else if (num_blocks == 4) {
187 soft_max(4);
188 } else if (num_blocks == 3) {
189 soft_max(3);
190 } else if (num_blocks == 2) {
191 soft_max(2);
192 } else if (num_blocks == 1) {
193 soft_max(1);
194 }
195}