vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

vulkan: Optimize soft_max (#10301)

Browse source 
* vulkan: Optimize soft_max

Large soft_max could already saturate memory, but small/medium sizes were
pretty slow. The bulk of the gains for them comes from using a smaller
workgroup size, and making the workgroup size match the subgroup size also
makes the barriers much cheaper.

Cache some values in locals to avoid refetching/recomputing. And stamp
out a few "template instantiations" so smaller cases will fully unroll.

Add a missing early return for OOB rows. This happens when there are more
than 512 rows and the dispatch is 512 x H.

* vulkan: Further soft_max optimizations

Restore the workgroup size of 512 case, use it for >1024.

Use unrollable loops for more iteration counts.
This commit is contained in:
Jeff Bolz 2024-11-19 01:25:17 -06:00 • committed by GitHub
parent 557924f222
commit b3e585988f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 107 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

114
ggml/src/ggml-vulkan/vulkan-shaders/soft_max.comp
View file

@ -1,6 +1,7 @@
#version 450
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_control_flow_attributes : enable
layout (push_constant) uniform parameter
{
@ -11,14 +12,13 @@ layout (push_constant) uniform parameter
float m0;
float m1;
uint n_head_log2;
uint nrows_x;
} p;
#include "types.comp"
#extension GL_EXT_control_flow_attributes : enable
#define BLOCK_SIZE 512
layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
layout(constant_id = 0) const uint BLOCK_SIZE = 32;
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
layout (binding = 1) readonly buffer Y {B_TYPE data_b[];};
@ -26,11 +26,18 @@ layout (binding = 2) buffer D {D_TYPE data_d[];};
shared FLOAT_TYPE vals[BLOCK_SIZE];
void main() {
// num_iters is the number of BLOCK_SIZE loop iterations we need to iterate
// over all the columns. The main function tries to pass a constant here,
// as if it were a template function, to allow unrolling.
void soft_max(uint num_iters) {
const uint tid = gl_LocalInvocationID.x;
const uint rowx = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
const uint rowy = rowx % p.KY;
if (rowx >= p.nrows_x) {
return;
}
float slope = 1.0f;
// ALiBi
@ -46,19 +53,37 @@ void main() {
// Find max
FLOAT_TYPE max_val = uintBitsToFloat(0xFF800000);
[[unroll]] for (uint col0 = 0; col0 < p.KX; col0 += BLOCK_SIZE) {
// Cache values while we compute the max, so we don't need to read them
// again when we're ready to compute exp(x-max).
const uint DATA_CACHE_SIZE = 16;
FLOAT_TYPE data_cache[DATA_CACHE_SIZE];
[[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) {
const uint col = col0 + tid;
if (col >= p.KX) {
break;
FLOAT_TYPE a = FLOAT_TYPE(0);
if (col < p.KX) {
a = data_a[rowx * p.KX + col];
}
max_val = max(max_val, FLOAT_TYPE(data_a[rowx * p.KX + col]) * p.scale + (p.KY > 0 ? slope * FLOAT_TYPE(data_b[rowy * p.KX + col]) : FLOAT_TYPE(0.0f)));
}
vals[tid] = max_val;
FLOAT_TYPE b = FLOAT_TYPE(0);
if (p.KY > 0 && col < p.KX) {
b = data_b[rowy * p.KX + col];
}
FLOAT_TYPE v = a * p.scale + slope * b;
max_val = max(max_val, v);
if (idx < DATA_CACHE_SIZE) {
data_cache[idx] = v;
}
}
// reduce across the workgroup
vals[tid] = max_val;
barrier();
[[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
[[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
if (tid < s) {
vals[tid] = max(vals[tid], vals[tid + s]);
}
@ -68,39 +93,80 @@ void main() {
max_val = vals[0];
barrier();
// Sum up values
vals[tid] = FLOAT_TYPE(0.0f);
FLOAT_TYPE sum = FLOAT_TYPE(0.0f);
[[unroll]] for (uint col0 = 0; col0 < p.KX; col0 += BLOCK_SIZE) {
// Compute sum{exp(x - max)}
[[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) {
const uint col = col0 + tid;
if (col >= p.KX) {
break;
}
// compute exp(a*scale+b*slope), add it to sum, and cache the new value
// in data_cache if possible.
const uint i = rowx * p.KX + col;
const FLOAT_TYPE val = exp(FLOAT_TYPE(data_a[i]) * p.scale + (p.KY > 0 ? slope * FLOAT_TYPE(data_b[rowy * p.KX + col]) : FLOAT_TYPE(0.0f)) - max_val);
vals[tid] += val;
data_d[i] = D_TYPE(val);
FLOAT_TYPE val;
if (idx < DATA_CACHE_SIZE) {
val = exp(data_cache[idx] - max_val);
} else {
val = exp(FLOAT_TYPE(data_a[i]) * p.scale + (p.KY > 0 ? slope * FLOAT_TYPE(data_b[rowy * p.KX + col]) : FLOAT_TYPE(0.0f)) - max_val);
}
sum += val;
if (idx < DATA_CACHE_SIZE) {
data_cache[idx] = val;
} else {
data_d[i] = D_TYPE(val);
}
}
// reduce across the workgroup
vals[tid] = sum;
barrier();
[[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
[[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
if (tid < s) {
vals[tid] += vals[tid + s];
}
barrier();
}
sum = vals[0];
const D_TYPE divisor = D_TYPE(vals[0]);
FLOAT_TYPE rcpdivisor = 1.0/sum;
[[unroll]] for (uint col0 = 0; col0 < p.KX; col0 += BLOCK_SIZE) {
[[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) {
const uint col = col0 + tid;
if (col >= p.KX) {
break;
continue;
}
data_d[rowx*p.KX + col] /= divisor;
if (idx < DATA_CACHE_SIZE) {
data_d[rowx*p.KX + col] = D_TYPE(data_cache[idx] * rcpdivisor);
} else {
data_d[rowx*p.KX + col] *= D_TYPE(rcpdivisor);
}
}
}
void main() {
// instantiate the soft_max function for several different
// dimensions, to allow loop unrolling
uint num_blocks = (p.KX + BLOCK_SIZE - 1) / BLOCK_SIZE;
if (num_blocks > 32) {
soft_max(num_blocks);
} else if (num_blocks > 16) {
soft_max(32);
} else if (num_blocks > 8) {
soft_max(16);
} else if (num_blocks > 4) {
soft_max(8);
} else if (num_blocks == 4) {
soft_max(4);
} else if (num_blocks == 3) {
soft_max(3);
} else if (num_blocks == 2) {
soft_max(2);
} else if (num_blocks == 1) {
soft_max(1);
}
}