phase 15: decode attention kernel + fused silu_mul + fused add_rmsnorm
Three performance optimizations targeting decode throughput: 1. Decode Attention Kernel (csrc/attention/flash_attention.cu): - Specialized kernel for Q_len=1 (decode step) - 256 threads parallelize across KV sequence dimension - Online softmax with block-level warp-shuffle reduction - Replaces FA2 kernel which wasted 63/64 threads for decode - flash_attention() auto-dispatches when q_len==1 2. Fused SiLU×Mul (csrc/activation/activations.cu): - Single kernel: out = silu(gate) * up - Saves 1 HBM read + 1 HBM write per FFN layer (N elements) - Eliminates intermediate tensor allocation 3. Fused Add+RMSNorm (csrc/normalization/rmsnorm.cu): - Single kernel: (normed, sum) = (rmsnorm(x+residual), x+residual) - Saves 1 full HBM round-trip per attention block - Eliminates separate add + rmsnorm kernel pair Performance analysis: - At current short sequences (max 79 tokens), these optimizations provide marginal benefit because the bottleneck is cuBLAS GEMV overhead: 252 weight matrix reads × ~32MB each = 15.5 GB per decode step. Theoretical minimum at 1.79 TB/s = 8.7ms, actual ~78ms (9x gap). - The fused kernels and decode attention will show larger gains at longer sequences where attention and element-wise ops dominate. - Next optimization target: CUDA Graphs to eliminate kernel launch overhead, or custom GEMV kernels to replace cuBLAS for M=1. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -63,6 +63,46 @@ __global__ void rmsnorm_bf16(
|
||||
}
|
||||
}
|
||||
|
||||
// Fused Add + RMSNorm: sum_out = x + residual, normed_out = rmsnorm(sum_out, gamma, eps)
|
||||
// Each block handles one row of [hidden_size].
|
||||
__global__ void add_rmsnorm_bf16(
|
||||
const __nv_bfloat16* __restrict__ x,
|
||||
const __nv_bfloat16* __restrict__ residual,
|
||||
const __nv_bfloat16* __restrict__ gamma,
|
||||
__nv_bfloat16* __restrict__ normed_out,
|
||||
__nv_bfloat16* __restrict__ sum_out,
|
||||
int hidden_size, float eps
|
||||
) {
|
||||
int row = blockIdx.x;
|
||||
const __nv_bfloat16* x_row = x + row * hidden_size;
|
||||
const __nv_bfloat16* res_row = residual + row * hidden_size;
|
||||
__nv_bfloat16* sum_row = sum_out + row * hidden_size;
|
||||
__nv_bfloat16* norm_row = normed_out + row * hidden_size;
|
||||
|
||||
// Pass 1: compute sum = x + residual, and accumulate sum_sq
|
||||
float sum_sq = 0.0f;
|
||||
for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
|
||||
float s = __bfloat162float(x_row[i]) + __bfloat162float(res_row[i]);
|
||||
sum_row[i] = __float2bfloat16(s);
|
||||
sum_sq += s * s;
|
||||
}
|
||||
sum_sq = block_reduce_sum(sum_sq);
|
||||
|
||||
__shared__ float s_rms_inv;
|
||||
if (threadIdx.x == 0) {
|
||||
s_rms_inv = rsqrtf(sum_sq / hidden_size + eps);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Pass 2: normed_out = sum * rms_inv * gamma
|
||||
float rms_inv = s_rms_inv;
|
||||
for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
|
||||
float s = __bfloat162float(sum_row[i]);
|
||||
float g = __bfloat162float(gamma[i]);
|
||||
norm_row[i] = __float2bfloat16(s * rms_inv * g);
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" {
|
||||
|
||||
void launch_rmsnorm_f32(const void* x, const void* gamma, void* out,
|
||||
@@ -80,4 +120,15 @@ void launch_rmsnorm_bf16(const void* x, const void* gamma, void* out,
|
||||
(__nv_bfloat16*)out, hidden_size, eps);
|
||||
}
|
||||
|
||||
void launch_add_rmsnorm_bf16(const void* x, const void* residual, const void* gamma,
|
||||
void* normed_out, void* sum_out,
|
||||
int rows, int hidden_size, float eps, void* stream) {
|
||||
int block = (hidden_size < 1024) ? hidden_size : 1024;
|
||||
add_rmsnorm_bf16<<<rows, block, 0, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)x, (const __nv_bfloat16*)residual,
|
||||
(const __nv_bfloat16*)gamma,
|
||||
(__nv_bfloat16*)normed_out, (__nv_bfloat16*)sum_out,
|
||||
hidden_size, eps);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user