aparler/c-llm
1
0
Fork 0
mirror of https://github.com/karpathy/llm.c.git synced 2026-07-28 20:35:09 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

online softmax

Browse source
This commit is contained in:
Erik Schultheis 2024-04-12 12:47:48 +03:00
parent de98a9a44a
commit 8e87aac341
1 changed files with 16 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

35
dev/cuda/attention_forward.cu
View file

@ -24,6 +24,7 @@ this turns out to be ~20X faster than (1) nice
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <float.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <cooperative_groups.h>
@ -330,45 +331,41 @@ __global__ void softmax_forward_kernel5(float* out, float inv_temperature, const
// one row of inp, i.e. inp[idx, :] of shape (T,)
const float* x = inp + idx * T;
// reduce to max
float maxval = -INFINITY;
// not INF, so we don't get NaNs accidentally when subtracting two values.
float maxval = -FLT_MAX;
float sumval = 0.0f;
const float4* x_vec = reinterpret_cast<const float4*>(x);
for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) {
float4 v = x_vec[i];
float old_maxval = maxval;
for(int k = 0; k < 4; ++k) {
maxval = fmaxf(maxval, vec_at(v, k));
}
}
if(4*pos_by_4 + warp.thread_rank() <= own_pos) {
maxval = fmaxf(maxval, x[4*pos_by_4 + warp.thread_rank()]);
}
float offset = cg::reduce(warp, maxval, cg::greater<float>{});
// compute exp and sum
float sumval = 0.0f;
for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) {
// subtract max for numerical stability
float4 v = x_vec[i];
sumval *= expf(inv_temperature * (old_maxval - maxval));
for(int k = 0; k < 4; ++k) {
sumval += expf(inv_temperature * (vec_at(v, k) - offset));
sumval += expf(inv_temperature * (vec_at(v, k) - maxval));
}
}
if(4*pos_by_4 + warp.thread_rank() <= own_pos) {
sumval += expf(inv_temperature * (x[4*pos_by_4 + warp.thread_rank()] - offset));
float old_maxval = maxval;
maxval = fmaxf(maxval, x[4*pos_by_4 + warp.thread_rank()]);
sumval *= expf(inv_temperature * (old_maxval - maxval));
sumval += expf(inv_temperature * (x[4*pos_by_4 + warp.thread_rank()] - maxval));
}
float sum = cg::reduce(warp, sumval, cg::plus<float>{});
float global_maxval = cg::reduce(warp, maxval, cg::greater<float>{});
sumval *= expf(inv_temperature * (maxval - global_maxval));
float sum = cg::reduce(warp, sumval, cg::plus<float>{});
float norm = 1.f / sum;
// divide the whole row by the sum
for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) {
// recalculation is faster than doing the round-trip through memory.
float ev = expf(inv_temperature * (__ldcs(x + i) - offset));
float ev = expf(inv_temperature * (__ldcs(x + i) - global_maxval));
__stcs(out + idx * T + i, ev * norm);
}
}