357 lines
12 KiB
Plaintext
357 lines
12 KiB
Plaintext
/*
|
|
* Online Softmax CUDA Kernel
|
|
*
|
|
* Implements numerically stable online softmax computation for Flash Attention.
|
|
* This is used as a utility kernel for incremental softmax updates.
|
|
*/
|
|
|
|
#include <cuda_runtime.h>
|
|
#include <cuda_fp16.h>
|
|
#include <cooperative_groups.h>
|
|
|
|
// Constants
|
|
#define WARP_SIZE 32
|
|
#define MAX_THREADS_PER_BLOCK 1024
|
|
|
|
namespace cg = cooperative_groups;
|
|
|
|
// Online softmax state structure
|
|
struct OnlineSoftmaxState {
|
|
float m; // running max
|
|
float l; // running sum
|
|
|
|
__device__ OnlineSoftmaxState() : m(-INFINITY), l(0.0f) {}
|
|
|
|
__device__ OnlineSoftmaxState(float max_val, float sum_val) : m(max_val), l(sum_val) {}
|
|
|
|
__device__ void update(float x) {
|
|
float m_new = fmaxf(m, x);
|
|
float l_new = l * expf(m - m_new) + expf(x - m_new);
|
|
m = m_new;
|
|
l = l_new;
|
|
}
|
|
|
|
__device__ void merge(const OnlineSoftmaxState& other) {
|
|
float m_new = fmaxf(m, other.m);
|
|
float l_new = l * expf(m - m_new) + other.l * expf(other.m - m_new);
|
|
m = m_new;
|
|
l = l_new;
|
|
}
|
|
|
|
__device__ float get_log_sum_exp() const {
|
|
return m + logf(l);
|
|
}
|
|
|
|
__device__ float get_normalization() const {
|
|
return l;
|
|
}
|
|
};
|
|
|
|
// Warp-level reduction for online softmax
|
|
__device__ OnlineSoftmaxState warp_reduce_softmax(OnlineSoftmaxState state) {
|
|
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
|
|
float other_m = __shfl_down_sync(0xffffffff, state.m, offset);
|
|
float other_l = __shfl_down_sync(0xffffffff, state.l, offset);
|
|
OnlineSoftmaxState other(other_m, other_l);
|
|
state.merge(other);
|
|
}
|
|
return state;
|
|
}
|
|
|
|
// Block-level reduction for online softmax
|
|
__device__ OnlineSoftmaxState block_reduce_softmax(OnlineSoftmaxState state) {
|
|
__shared__ OnlineSoftmaxState shared_states[32]; // Max warps per block
|
|
|
|
int warp_id = threadIdx.x / WARP_SIZE;
|
|
int lane_id = threadIdx.x % WARP_SIZE;
|
|
|
|
// Warp-level reduce
|
|
state = warp_reduce_softmax(state);
|
|
|
|
// Store warp result
|
|
if (lane_id == 0) {
|
|
shared_states[warp_id] = state;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
// Block-level reduce
|
|
if (warp_id == 0) {
|
|
OnlineSoftmaxState final_state;
|
|
if (lane_id < (blockDim.x + WARP_SIZE - 1) / WARP_SIZE) {
|
|
final_state = shared_states[lane_id];
|
|
}
|
|
final_state = warp_reduce_softmax(final_state);
|
|
|
|
// Broadcast result
|
|
shared_states[0] = final_state;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
return shared_states[0];
|
|
}
|
|
|
|
// Online softmax kernel for a batch of sequences
|
|
extern "C" __global__ void online_softmax_kernel(
|
|
const float* __restrict__ scores, // Input scores [batch, seq_len, seq_len]
|
|
float* __restrict__ probs, // Output probabilities [batch, seq_len, seq_len]
|
|
float* __restrict__ lse, // Log-sum-exp [batch, seq_len]
|
|
int batch_size,
|
|
int seq_len,
|
|
unsigned int causal
|
|
) {
|
|
int batch_idx = blockIdx.x;
|
|
int seq_idx = blockIdx.y;
|
|
int tid = threadIdx.x;
|
|
|
|
if (batch_idx >= batch_size || seq_idx >= seq_len) return;
|
|
|
|
// Calculate offsets
|
|
int batch_offset = batch_idx * seq_len * seq_len;
|
|
int row_offset = batch_offset + seq_idx * seq_len;
|
|
int lse_offset = batch_idx * seq_len + seq_idx;
|
|
|
|
// Compute softmax for this sequence position
|
|
OnlineSoftmaxState state;
|
|
|
|
// First pass: compute max and sum
|
|
for (int i = tid; i < seq_len; i += blockDim.x) {
|
|
if (!causal || i <= seq_idx) {
|
|
float score = scores[row_offset + i];
|
|
state.update(score);
|
|
}
|
|
}
|
|
|
|
// Reduce across block
|
|
state = block_reduce_softmax(state);
|
|
|
|
// Second pass: compute probabilities
|
|
for (int i = tid; i < seq_len; i += blockDim.x) {
|
|
float prob = 0.0f;
|
|
if (!causal || i <= seq_idx) {
|
|
float score = scores[row_offset + i];
|
|
prob = expf(score - state.m) / state.l;
|
|
}
|
|
probs[row_offset + i] = prob;
|
|
}
|
|
|
|
// Store log-sum-exp
|
|
if (tid == 0) {
|
|
lse[lse_offset] = state.get_log_sum_exp();
|
|
}
|
|
}
|
|
|
|
// Incremental online softmax update kernel
|
|
extern "C" __global__ void online_softmax_update_kernel(
|
|
const float* __restrict__ new_scores, // New scores to add [batch, seq_len, new_len]
|
|
const float* __restrict__ old_lse, // Previous log-sum-exp [batch, seq_len]
|
|
float* __restrict__ updated_probs, // Updated probabilities [batch, seq_len, total_len]
|
|
float* __restrict__ new_lse, // Updated log-sum-exp [batch, seq_len]
|
|
int batch_size,
|
|
int seq_len,
|
|
int old_len,
|
|
int new_len,
|
|
unsigned int causal
|
|
) {
|
|
int batch_idx = blockIdx.x;
|
|
int seq_idx = blockIdx.y;
|
|
int tid = threadIdx.x;
|
|
|
|
if (batch_idx >= batch_size || seq_idx >= seq_len) return;
|
|
|
|
// Calculate offsets
|
|
int old_offset = batch_idx * seq_len * old_len + seq_idx * old_len;
|
|
int new_offset = batch_idx * seq_len * new_len + seq_idx * new_len;
|
|
int updated_offset = batch_idx * seq_len * (old_len + new_len) + seq_idx * (old_len + new_len);
|
|
int lse_offset = batch_idx * seq_len + seq_idx;
|
|
|
|
// Get old state
|
|
float old_lse_val = old_lse[lse_offset];
|
|
OnlineSoftmaxState old_state(old_lse_val - logf(expf(old_lse_val)), expf(old_lse_val));
|
|
|
|
// Compute state for new scores
|
|
OnlineSoftmaxState new_state;
|
|
for (int i = tid; i < new_len; i += blockDim.x) {
|
|
int global_pos = old_len + i;
|
|
if (!causal || global_pos <= seq_idx) {
|
|
float score = new_scores[new_offset + i];
|
|
new_state.update(score);
|
|
}
|
|
}
|
|
|
|
// Reduce new state across block
|
|
new_state = block_reduce_softmax(new_state);
|
|
|
|
// Merge old and new states
|
|
OnlineSoftmaxState merged_state = old_state;
|
|
merged_state.merge(new_state);
|
|
|
|
// Update old probabilities
|
|
float scale_factor = old_state.l * expf(old_state.m - merged_state.m) / merged_state.l;
|
|
for (int i = tid; i < old_len; i += blockDim.x) {
|
|
if (!causal || i <= seq_idx) {
|
|
updated_probs[updated_offset + i] = updated_probs[old_offset + i] * scale_factor;
|
|
} else {
|
|
updated_probs[updated_offset + i] = 0.0f;
|
|
}
|
|
}
|
|
|
|
// Compute new probabilities
|
|
for (int i = tid; i < new_len; i += blockDim.x) {
|
|
int global_pos = old_len + i;
|
|
float prob = 0.0f;
|
|
if (!causal || global_pos <= seq_idx) {
|
|
float score = new_scores[new_offset + i];
|
|
prob = expf(score - merged_state.m) / merged_state.l;
|
|
}
|
|
updated_probs[updated_offset + old_len + i] = prob;
|
|
}
|
|
|
|
// Store updated log-sum-exp
|
|
if (tid == 0) {
|
|
new_lse[lse_offset] = merged_state.get_log_sum_exp();
|
|
}
|
|
}
|
|
|
|
// Fused online softmax and value computation kernel
|
|
extern "C" __global__ void online_softmax_value_kernel(
|
|
const float* __restrict__ scores, // Input scores [batch, heads, seq_len, seq_len]
|
|
const half* __restrict__ values, // Value tensor [batch, heads, seq_len, head_dim]
|
|
half* __restrict__ output, // Output tensor [batch, heads, seq_len, head_dim]
|
|
float* __restrict__ lse, // Log-sum-exp [batch, heads, seq_len]
|
|
int batch_size,
|
|
int num_heads,
|
|
int seq_len,
|
|
int head_dim,
|
|
unsigned int causal
|
|
) {
|
|
int batch_head_idx = blockIdx.x;
|
|
int seq_idx = blockIdx.y;
|
|
int tid = threadIdx.x;
|
|
|
|
int batch_idx = batch_head_idx / num_heads;
|
|
int head_idx = batch_head_idx % num_heads;
|
|
|
|
if (batch_idx >= batch_size || head_idx >= num_heads || seq_idx >= seq_len) return;
|
|
|
|
// Calculate offsets
|
|
int scores_offset = batch_head_idx * seq_len * seq_len + seq_idx * seq_len;
|
|
int values_offset = batch_head_idx * seq_len * head_dim;
|
|
int output_offset = batch_head_idx * seq_len * head_dim + seq_idx * head_dim;
|
|
int lse_offset = batch_head_idx * seq_len + seq_idx;
|
|
|
|
// Shared memory for values
|
|
extern __shared__ half shared_values[];
|
|
|
|
// Load values into shared memory
|
|
for (int i = tid; i < seq_len * head_dim; i += blockDim.x) {
|
|
shared_values[i] = values[values_offset + i];
|
|
}
|
|
__syncthreads();
|
|
|
|
// Compute online softmax
|
|
OnlineSoftmaxState state;
|
|
for (int i = 0; i < seq_len; i++) {
|
|
if (!causal || i <= seq_idx) {
|
|
float score = scores[scores_offset + i];
|
|
state.update(score);
|
|
}
|
|
}
|
|
|
|
// Compute output values
|
|
for (int d = tid; d < head_dim; d += blockDim.x) {
|
|
float output_val = 0.0f;
|
|
|
|
for (int i = 0; i < seq_len; i++) {
|
|
if (!causal || i <= seq_idx) {
|
|
float score = scores[scores_offset + i];
|
|
float prob = expf(score - state.m) / state.l;
|
|
float value = __half2float(shared_values[i * head_dim + d]);
|
|
output_val += prob * value;
|
|
}
|
|
}
|
|
|
|
output[output_offset + d] = __float2half(output_val);
|
|
}
|
|
|
|
// Store log-sum-exp
|
|
if (tid == 0) {
|
|
lse[lse_offset] = state.get_log_sum_exp();
|
|
}
|
|
}
|
|
|
|
// Vectorized online softmax kernel using half2
|
|
extern "C" __global__ void online_softmax_half2_kernel(
|
|
const half* __restrict__ scores, // Input scores [batch, seq_len, seq_len]
|
|
half* __restrict__ probs, // Output probabilities [batch, seq_len, seq_len]
|
|
float* __restrict__ lse, // Log-sum-exp [batch, seq_len]
|
|
int batch_size,
|
|
int seq_len,
|
|
unsigned int causal
|
|
) {
|
|
int batch_idx = blockIdx.x;
|
|
int seq_idx = blockIdx.y;
|
|
int tid = threadIdx.x;
|
|
|
|
if (batch_idx >= batch_size || seq_idx >= seq_len) return;
|
|
|
|
// Calculate offsets
|
|
int batch_offset = batch_idx * seq_len * seq_len;
|
|
int row_offset = batch_offset + seq_idx * seq_len;
|
|
int lse_offset = batch_idx * seq_len + seq_idx;
|
|
|
|
// Compute softmax using vectorized operations
|
|
OnlineSoftmaxState state;
|
|
|
|
// First pass: compute max and sum using half2
|
|
for (int i = tid * 2; i < seq_len; i += blockDim.x * 2) {
|
|
half2 scores_vec = *reinterpret_cast<const half2*>(&scores[row_offset + i]);
|
|
|
|
if (!causal || i <= seq_idx) {
|
|
float score1 = __half2float(scores_vec.x);
|
|
state.update(score1);
|
|
}
|
|
|
|
if (i + 1 < seq_len && (!causal || i + 1 <= seq_idx)) {
|
|
float score2 = __half2float(scores_vec.y);
|
|
state.update(score2);
|
|
}
|
|
}
|
|
|
|
// Reduce across block
|
|
state = block_reduce_softmax(state);
|
|
|
|
// Second pass: compute probabilities using half2
|
|
for (int i = tid * 2; i < seq_len; i += blockDim.x * 2) {
|
|
half2 scores_vec = *reinterpret_cast<const half2*>(&scores[row_offset + i]);
|
|
half2 probs_vec;
|
|
|
|
if (!causal || i <= seq_idx) {
|
|
float score1 = __half2float(scores_vec.x);
|
|
float prob1 = expf(score1 - state.m) / state.l;
|
|
probs_vec.x = __float2half(prob1);
|
|
} else {
|
|
probs_vec.x = __float2half(0.0f);
|
|
}
|
|
|
|
if (i + 1 < seq_len) {
|
|
if (!causal || i + 1 <= seq_idx) {
|
|
float score2 = __half2float(scores_vec.y);
|
|
float prob2 = expf(score2 - state.m) / state.l;
|
|
probs_vec.y = __float2half(prob2);
|
|
} else {
|
|
probs_vec.y = __float2half(0.0f);
|
|
}
|
|
} else {
|
|
probs_vec.y = __float2half(0.0f);
|
|
}
|
|
|
|
*reinterpret_cast<half2*>(&probs[row_offset + i]) = probs_vec;
|
|
}
|
|
|
|
// Store log-sum-exp
|
|
if (tid == 0) {
|
|
lse[lse_offset] = state.get_log_sum_exp();
|
|
}
|
|
} |