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

22
crates/xserv-kernels/src/activation.rs
View File

@@ -12,6 +12,7 @@ unsafe extern "C" {
fn launch_add_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_mul_f32(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_mul_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_silu_mul_bf16(gate: *const c_void, up: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
}
fn dispatch_unary(x: &Tensor, f32_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void),
@@ -67,3 +68,24 @@ pub fn scale(x: &Tensor, scale_val: f32) -> Tensor {
pub fn add(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_add_f32, launch_add_bf16) }
pub fn mul(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_mul_f32, launch_mul_bf16) }
/// Fused SiLU×Mul: out = silu(gate) * up (BF16 only)
/// Saves one HBM read + one HBM write compared to separate silu + mul.
pub fn silu_mul(gate: &Tensor, up: &Tensor) -> Tensor {
assert_eq!(gate.shape(), up.shape());
assert!(gate.is_contiguous() && up.is_contiguous());
assert!(matches!(gate.device(), Device::Cuda(_)));
assert_eq!(gate.dtype(), DType::BF16, "silu_mul requires BF16");
let out = Tensor::zeros(gate.shape(), gate.dtype(), gate.device());
let n = gate.numel() as i32;
unsafe {
launch_silu_mul_bf16(
gate.data_ptr() as *const c_void,
up.data_ptr() as *const c_void,
out.data_ptr() as *mut c_void,
n,
std::ptr::null_mut(),
);
}
out
}

56
crates/xserv-kernels/src/attention.rs
View File

@@ -16,6 +16,12 @@ unsafe extern "C" {
q_len: i32, kv_len: i32, head_dim: i32,
scale: f32, causal: i32, stream: *mut c_void,
);
fn launch_decode_attention_bf16(
q: *const c_void, k: *const c_void, v: *const c_void, o: *mut c_void,
batch: i32, num_q_heads: i32, num_kv_heads: i32,
kv_len: i32, head_dim: i32,
scale: f32, causal: i32, stream: *mut c_void,
);
}
fn apply_causal_mask(scores: &Tensor, offset: usize) {
@@ -81,7 +87,52 @@ pub fn attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tensor {
batched_matmul(&weights, v)
}
/// Decode Attention — optimized for single-token decode (q_len=1).
///
/// q: [batch, num_q_heads, 1, head_dim] BF16, contiguous, GPU
/// k: [batch, num_kv_heads, kv_len, head_dim] BF16, contiguous, GPU
/// v: [batch, num_kv_heads, kv_len, head_dim] BF16, contiguous, GPU
///
/// Returns: [batch, num_q_heads, 1, head_dim] BF16
pub fn decode_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Tensor {
assert_eq!(q.ndim(), 4);
assert_eq!(q.shape()[2], 1, "decode_attention requires q_len == 1");
let batch = q.shape()[0];
let num_q_heads = q.shape()[1];
let head_dim = q.shape()[3];
let num_kv_heads = k.shape()[1];
let kv_len = k.shape()[2];
let scale = 1.0 / (head_dim as f32).sqrt();
let output = Tensor::zeros(
&[batch, num_q_heads, 1, head_dim],
DType::BF16,
q.device(),
);
unsafe {
launch_decode_attention_bf16(
q.data_ptr() as *const c_void,
k.data_ptr() as *const c_void,
v.data_ptr() as *const c_void,
output.data_ptr() as *mut c_void,
batch as i32,
num_q_heads as i32,
num_kv_heads as i32,
kv_len as i32,
head_dim as i32,
scale,
1, // causal (always 1 for decode)
std::ptr::null_mut(),
);
}
output
}
/// Flash Attention 2 — O(1) extra memory, supports GQA natively.
/// Auto-dispatches to decode_attention when q_len == 1.
///
/// q: [batch, num_q_heads, q_len, head_dim] BF16, contiguous, GPU
/// k: [batch, num_kv_heads, kv_len, head_dim] BF16, contiguous, GPU
@@ -109,6 +160,11 @@ pub fn flash_attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tens
assert!(num_q_heads % num_kv_heads == 0, "num_q_heads must be divisible by num_kv_heads");
assert!(head_dim <= 128, "flash_attention supports head_dim up to 128");
// Dispatch to specialized decode kernel for single-token generation
if q_len == 1 {
return decode_attention(q, k, v);
}
let scale = 1.0 / (head_dim as f32).sqrt();
let output = Tensor::zeros(
&[batch, num_q_heads, q_len, head_dim],

6
crates/xserv-kernels/src/lib.rs
View File

@@ -8,13 +8,13 @@ pub mod rope;
pub mod softmax;
pub mod transpose;
pub use activation::{add, gelu, mul, scale, silu};
pub use activation::{add, gelu, mul, scale, silu, silu_mul};
pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
pub use attention::{attention, flash_attention};
pub use attention::{attention, decode_attention, flash_attention};
pub use embedding::embedding;
pub use gemm::{batched_matmul, matmul, GemmBackend};
pub use layernorm::layernorm;
pub use rmsnorm::rmsnorm;
pub use rmsnorm::{add_rmsnorm, rmsnorm};
pub use rope::{rope_inplace, RopeCache};
pub use softmax::softmax;

39
crates/xserv-kernels/src/rmsnorm.rs
View File

@@ -6,6 +6,9 @@ unsafe extern "C" {
rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
fn launch_rmsnorm_bf16(x: *const c_void, gamma: *const c_void, out: *mut c_void,
rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
fn launch_add_rmsnorm_bf16(x: *const c_void, residual: *const c_void, gamma: *const c_void,
normed_out: *mut c_void, sum_out: *mut c_void,
rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
}
pub fn rmsnorm(x: &Tensor, gamma: &Tensor, eps: f32) -> Tensor {
@@ -34,3 +37,39 @@ pub fn rmsnorm(x: &Tensor, gamma: &Tensor, eps: f32) -> Tensor {
}
out
}
/// Fused Add + RMSNorm: computes sum = x + residual, then normed = rmsnorm(sum, gamma, eps).
/// Returns (normed, sum). BF16 only.
/// Saves one kernel launch and one full HBM round-trip per layer.
pub fn add_rmsnorm(x: &Tensor, residual: &Tensor, gamma: &Tensor, eps: f32) -> (Tensor, Tensor) {
assert!(x.ndim() >= 1);
assert_eq!(x.shape(), residual.shape());
assert!(x.is_contiguous() && residual.is_contiguous() && gamma.is_contiguous());
assert!(matches!(x.device(), Device::Cuda(_)));
assert_eq!(x.dtype(), DType::BF16, "add_rmsnorm requires BF16");
assert_eq!(residual.dtype(), DType::BF16);
assert_eq!(gamma.dtype(), DType::BF16);
let hidden_size = *x.shape().last().unwrap();
assert_eq!(gamma.shape(), &[hidden_size]);
let rows = x.numel() / hidden_size;
let normed_out = Tensor::zeros(x.shape(), DType::BF16, x.device());
let sum_out = Tensor::zeros(x.shape(), DType::BF16, x.device());
unsafe {
launch_add_rmsnorm_bf16(
x.data_ptr() as *const c_void,
residual.data_ptr() as *const c_void,
gamma.data_ptr() as *const c_void,
normed_out.data_ptr() as *mut c_void,
sum_out.data_ptr() as *mut c_void,
rows as i32,
hidden_size as i32,
eps,
std::ptr::null_mut(),
);
}
(normed_out, sum_out)
}