gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

phase 15: decode attention kernel + fused silu_mul + fused add_rmsnorm

Browse Source 
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:
Gahow Wang
2026-05-22 19:40:56 +08:00
parent 6cc1c9332d
commit 9783fcf410
8 changed files with 387 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

191
csrc/attention/flash_attention.cu
View File

@@ -196,6 +196,177 @@ __global__ void flash_attention_bf16_kernel(
}
}
// ============================================================
// Decode Attention kernel: optimized for Q_len=1 (single-token decode).
// Parallelizes across KV sequence dimension instead of Q rows.
//
// Grid: (batch * num_q_heads, 1) — one block per Q head
// Block: 256 threads — each thread handles ceil(kv_len / 256) KV positions
// Uses online softmax reduction across threads.
// ============================================================
#define DECODE_THREADS 256
#define HEAD_DIM_MAX 128
__global__ void decode_attention_bf16_kernel(
const __nv_bfloat16* __restrict__ Q,
const __nv_bfloat16* __restrict__ K,
const __nv_bfloat16* __restrict__ V,
__nv_bfloat16* __restrict__ O,
int num_q_heads, int num_kv_heads,
int kv_len, int head_dim,
float scale
) {
int bh = blockIdx.x;
int batch_idx = bh / num_q_heads;
int q_head = bh % num_q_heads;
// GQA mapping
int heads_per_group = num_q_heads / num_kv_heads;
int kv_head = q_head / heads_per_group;
int tid = threadIdx.x;
// Pointers to this batch/head's data
// Q: [batch, num_q_heads, 1, head_dim]
const __nv_bfloat16* Q_ptr = Q + ((long long)batch_idx * num_q_heads + q_head) * head_dim;
// K/V: [batch, num_kv_heads, kv_len, head_dim]
const __nv_bfloat16* K_base = K + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim;
const __nv_bfloat16* V_base = V + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim;
__nv_bfloat16* O_ptr = O + ((long long)batch_idx * num_q_heads + q_head) * head_dim;
// Load Q vector into registers (head_dim <= 128)
float q_reg[HEAD_DIM_MAX];
for (int d = 0; d < head_dim; d++) {
q_reg[d] = __bfloat162float(Q_ptr[d]);
}
// Each thread processes a chunk of KV positions
// Thread tid handles positions: tid, tid+DECODE_THREADS, tid+2*DECODE_THREADS, ...
float local_max = -INFINITY;
float local_sum = 0.0f;
float local_O[HEAD_DIM_MAX];
for (int d = 0; d < head_dim; d++) {
local_O[d] = 0.0f;
}
for (int pos = tid; pos < kv_len; pos += DECODE_THREADS) {
// Compute dot(Q, K[pos]) * scale
const __nv_bfloat16* K_pos = K_base + pos * head_dim;
float dot = 0.0f;
for (int d = 0; d < head_dim; d++) {
dot += q_reg[d] * __bfloat162float(K_pos[d]);
}
float s = dot * scale;
// Online softmax update
float new_max = fmaxf(local_max, s);
float correction = expf(local_max - new_max);
float p = expf(s - new_max);
// Rescale running sum and O
local_sum = local_sum * correction + p;
for (int d = 0; d < head_dim; d++) {
local_O[d] = local_O[d] * correction;
}
// Accumulate V[pos] weighted by p
const __nv_bfloat16* V_pos = V_base + pos * head_dim;
for (int d = 0; d < head_dim; d++) {
local_O[d] += p * __bfloat162float(V_pos[d]);
}
local_max = new_max;
}
// --- Block-level online softmax reduction ---
// We need to combine (local_max, local_sum, local_O) across all threads.
// Strategy: reduce max, then each thread rescales, then reduce sum and O.
// Shared memory for reduction
__shared__ float smem_max[32]; // one per warp
__shared__ float smem_sum[32];
__shared__ float smem_O[HEAD_DIM_MAX]; // final output accumulator
// Step 1: Block-wide max reduction
int lane = tid & 31;
int warp_id = tid >> 5;
int num_warps = DECODE_THREADS >> 5; // 8 warps
float warp_max = local_max;
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
warp_max = fmaxf(warp_max, __shfl_down_sync(0xffffffff, warp_max, offset));
if (lane == 0) smem_max[warp_id] = warp_max;
__syncthreads();
float global_max;
if (tid == 0) {
global_max = smem_max[0];
for (int i = 1; i < num_warps; i++)
global_max = fmaxf(global_max, smem_max[i]);
smem_max[0] = global_max;
}
__syncthreads();
global_max = smem_max[0];
// Step 2: Each thread rescales its local_sum and local_O with global_max
float rescale = (local_max == -INFINITY) ? 0.0f : expf(local_max - global_max);
local_sum *= rescale;
for (int d = 0; d < head_dim; d++) {
local_O[d] *= rescale;
}
// Step 3: Reduce sum across block
float warp_sum = local_sum;
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset);
if (lane == 0) smem_sum[warp_id] = warp_sum;
__syncthreads();
float global_sum;
if (tid == 0) {
global_sum = 0.0f;
for (int i = 0; i < num_warps; i++)
global_sum += smem_sum[i];
smem_sum[0] = global_sum;
}
__syncthreads();
global_sum = smem_sum[0];
// Step 4: Reduce O across block (dimension by dimension using shared mem)
float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
// Process head_dim in chunks: each iteration reduces one dimension
// Use shared memory accumulator: each warp contributes via warp reduction + atomic
// Actually simpler: iterate over dimensions, warp reduce each, then lane0 atomicAdd to smem_O
// Initialize smem_O
for (int d = tid; d < head_dim; d += DECODE_THREADS) {
smem_O[d] = 0.0f;
}
__syncthreads();
// Each thread adds its local_O contributions via warp reduction + atomicAdd
for (int d = 0; d < head_dim; d++) {
float val = local_O[d];
// Warp-level reduction
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_down_sync(0xffffffff, val, offset);
if (lane == 0) {
atomicAdd(&smem_O[d], val);
}
}
__syncthreads();
// Thread 0..head_dim-1 write final output
for (int d = tid; d < head_dim; d += DECODE_THREADS) {
O_ptr[d] = __float2bfloat16(smem_O[d] * inv_sum);
}
}
extern "C" {
void launch_flash_attention_bf16(
@@ -222,4 +393,24 @@ void launch_flash_attention_bf16(
);
}
void launch_decode_attention_bf16(
const void* Q, const void* K, const void* V, void* O,
int batch, int num_q_heads, int num_kv_heads,
int kv_len, int head_dim,
float scale, int causal, void* stream
) {
int grid = batch * num_q_heads;
int block = DECODE_THREADS;
decode_attention_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
(const __nv_bfloat16*)Q,
(const __nv_bfloat16*)K,
(const __nv_bfloat16*)V,
(__nv_bfloat16*)O,
num_q_heads, num_kv_heads,
kv_len, head_dim,
scale
);
}
}