Gahow Wang
9783fcf410
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>
25 lines
815 B
Rust
25 lines
815 B
Rust
pub mod activation;
|
|
pub mod attention;
|
|
pub mod embedding;
|
|
pub mod gemm;
|
|
pub mod layernorm;
|
|
pub mod rmsnorm;
|
|
pub mod rope;
|
|
pub mod softmax;
|
|
pub mod transpose;
|
|
|
|
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, decode_attention, flash_attention};
|
|
pub use embedding::embedding;
|
|
pub use gemm::{batched_matmul, matmul, GemmBackend};
|
|
pub use layernorm::layernorm;
|
|
pub use rmsnorm::{add_rmsnorm, rmsnorm};
|
|
pub use rope::{rope_inplace, RopeCache};
|
|
pub use softmax::softmax;
|
|
|
|
/// Register GPU kernels with the tensor crate. Call once at startup.
|
|
pub fn init() {
|
|
xserv_tensor::register_gpu_contiguous(strided_to_contiguous_gpu);
|
|
}
|