Gahow Wang
c8e8153702
CUDA kernels (csrc/): - common.cuh: shared warp_reduce_sum/max, block_reduce_sum/max - normalization/rmsnorm.cu: RMSNorm (F32 + BF16) - normalization/layernorm.cu: LayerNorm with Welford (F32 + BF16) - activation/activations.cu: GELU tanh-approx + SiLU (F32 + BF16) - reduce/softmax.cu: safe softmax, 3-pass (F32 + BF16) - embedding/embedding.cu: gather lookup (F32 + BF16) - embedding/rope.cu: RoPE in-place + precomputed cos/sin cache (F32 + BF16) Rust wrappers (xserv-kernels/src/): - rmsnorm.rs, layernorm.rs, activation.rs, softmax.rs, embedding.rs, rope.rs - RopeCache struct with GPU-side precomputation Tests: 12 new tests (ops_test.rs), all passing with good precision: - F32: max_err 1e-6 ~ 1e-9 - BF16: max_err 2e-3 ~ 7e-3 Total: 29 kernel tests + 27 prior = 56 tests passing Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
107 lines
3.1 KiB
Plaintext
107 lines
3.1 KiB
Plaintext
#include "../common.cuh"
|
|
|
|
// Safe softmax along the last dimension.
|
|
// Each block handles one row of length `cols`.
|
|
// Three-pass: 1) find max, 2) exp + sum, 3) normalize.
|
|
|
|
__global__ void softmax_f32(
|
|
const float* __restrict__ x,
|
|
float* __restrict__ out,
|
|
int cols
|
|
) {
|
|
int row = blockIdx.x;
|
|
const float* x_row = x + row * cols;
|
|
float* out_row = out + row * cols;
|
|
|
|
// Pass 1: find max
|
|
float local_max = -INFINITY;
|
|
for (int i = threadIdx.x; i < cols; i += blockDim.x) {
|
|
local_max = fmaxf(local_max, x_row[i]);
|
|
}
|
|
float row_max = block_reduce_max(local_max);
|
|
|
|
__shared__ float s_max;
|
|
if (threadIdx.x == 0) s_max = row_max;
|
|
__syncthreads();
|
|
row_max = s_max;
|
|
|
|
// Pass 2: exp and sum
|
|
float local_sum = 0.0f;
|
|
for (int i = threadIdx.x; i < cols; i += blockDim.x) {
|
|
float e = expf(x_row[i] - row_max);
|
|
out_row[i] = e;
|
|
local_sum += e;
|
|
}
|
|
float row_sum = block_reduce_sum(local_sum);
|
|
|
|
__shared__ float s_inv_sum;
|
|
if (threadIdx.x == 0) s_inv_sum = 1.0f / row_sum;
|
|
__syncthreads();
|
|
float inv_sum = s_inv_sum;
|
|
|
|
// Pass 3: normalize
|
|
for (int i = threadIdx.x; i < cols; i += blockDim.x) {
|
|
out_row[i] *= inv_sum;
|
|
}
|
|
}
|
|
|
|
__global__ void softmax_bf16(
|
|
const __nv_bfloat16* __restrict__ x,
|
|
__nv_bfloat16* __restrict__ out,
|
|
int cols
|
|
) {
|
|
int row = blockIdx.x;
|
|
const __nv_bfloat16* x_row = x + row * cols;
|
|
__nv_bfloat16* out_row = out + row * cols;
|
|
|
|
float local_max = -INFINITY;
|
|
for (int i = threadIdx.x; i < cols; i += blockDim.x) {
|
|
local_max = fmaxf(local_max, __bfloat162float(x_row[i]));
|
|
}
|
|
float row_max = block_reduce_max(local_max);
|
|
|
|
__shared__ float s_max;
|
|
if (threadIdx.x == 0) s_max = row_max;
|
|
__syncthreads();
|
|
row_max = s_max;
|
|
|
|
// We need float scratch for exp values. Reuse out (write bf16 in pass 3).
|
|
// Use registers to hold exp values during sum pass instead.
|
|
float local_sum = 0.0f;
|
|
for (int i = threadIdx.x; i < cols; i += blockDim.x) {
|
|
float e = expf(__bfloat162float(x_row[i]) - row_max);
|
|
// Temporarily store exp in output as bf16 (slight precision loss, acceptable)
|
|
out_row[i] = __float2bfloat16(e);
|
|
local_sum += e;
|
|
}
|
|
float row_sum = block_reduce_sum(local_sum);
|
|
|
|
__shared__ float s_inv_sum;
|
|
if (threadIdx.x == 0) s_inv_sum = 1.0f / row_sum;
|
|
__syncthreads();
|
|
float inv_sum = s_inv_sum;
|
|
|
|
for (int i = threadIdx.x; i < cols; i += blockDim.x) {
|
|
float e = __bfloat162float(out_row[i]);
|
|
out_row[i] = __float2bfloat16(e * inv_sum);
|
|
}
|
|
}
|
|
|
|
extern "C" {
|
|
|
|
void launch_softmax_f32(const void* x, void* out, int rows, int cols, void* stream) {
|
|
int block = (cols < 1024) ? cols : 1024;
|
|
if (block < 32) block = 32;
|
|
softmax_f32<<<rows, block, 0, (cudaStream_t)stream>>>(
|
|
(const float*)x, (float*)out, cols);
|
|
}
|
|
|
|
void launch_softmax_bf16(const void* x, void* out, int rows, int cols, void* stream) {
|
|
int block = (cols < 1024) ? cols : 1024;
|
|
if (block < 32) block = 32;
|
|
softmax_bf16<<<rows, block, 0, (cudaStream_t)stream>>>(
|
|
(const __nv_bfloat16*)x, (__nv_bfloat16*)out, cols);
|
|
}
|
|
|
|
}
|