phase 4: transformer core kernels

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>
docs: add design docs + takeaways for Phase 2 and Phase 3
2026-05-21 21:07:24 +08:00 · 2026-05-21 20:59:45 +08:00 · 2026-05-21 19:48:05 +08:00
27 changed files with 2145 additions and 7 deletions
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -3,6 +3,7 @@ resolver = "2"
 members = [
    "crates/xserv-cuda",
    "crates/xserv-tensor",
+    "crates/xserv-kernels",
 ]

 [workspace.package]
--- a/crates/xserv-cuda/src/error.rs
+++ b/crates/xserv-cuda/src/error.rs
@@ -23,7 +23,7 @@ impl std::error::Error for CudaError {}

 pub type Result<T> = std::result::Result<T, CudaError>;

-pub(crate) fn check(code: i32) -> Result<()> {
+pub fn check(code: i32) -> Result<()> {
    if code == ffi::CUDA_SUCCESS {
        return Ok(());
    }
--- a/crates/xserv-kernels/Cargo.toml
+++ b/crates/xserv-kernels/Cargo.toml
@@ -0,0 +1,12 @@
+[package]
+name = "xserv-kernels"
+version.workspace = true
+edition.workspace = true
+
+[build-dependencies]
+cc = "1"
+
+[dependencies]
+xserv-cuda = { path = "../xserv-cuda" }
+xserv-tensor = { path = "../xserv-tensor" }
+half.workspace = true
--- a/crates/xserv-kernels/build.rs
+++ b/crates/xserv-kernels/build.rs
@@ -0,0 +1,28 @@
+use std::env;
+
+fn main() {
+    let cuda_path = env::var("CUDA_HOME")
+        .or_else(|_| env::var("CUDA_PATH"))
+        .unwrap_or_else(|_| "/usr/local/cuda".to_string());
+
+    println!("cargo:rustc-link-search=native={cuda_path}/lib64");
+    println!("cargo:rustc-link-lib=dylib=cudart");
+    println!("cargo:rustc-link-lib=dylib=cublas");
+
+    cc::Build::new()
+        .cuda(true)
+        .cudart("shared")
+        .flag("-gencode=arch=compute_120,code=sm_120")
+        .include("../../csrc")
+        .file("../../csrc/gemm/naive.cu")
+        .file("../../csrc/gemm/tiled.cu")
+        .file("../../csrc/normalization/rmsnorm.cu")
+        .file("../../csrc/normalization/layernorm.cu")
+        .file("../../csrc/activation/activations.cu")
+        .file("../../csrc/reduce/softmax.cu")
+        .file("../../csrc/embedding/embedding.cu")
+        .file("../../csrc/embedding/rope.cu")
+        .compile("xserv_kernels");
+
+    println!("cargo:rerun-if-changed=../../csrc/");
+}
--- a/crates/xserv-kernels/src/activation.rs
+++ b/crates/xserv-kernels/src/activation.rs
@@ -0,0 +1,41 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_gelu_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_gelu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_silu_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_silu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+}
+
+pub fn gelu(x: &Tensor) -> Tensor {
+    assert!(x.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    let out = Tensor::zeros(x.shape(), x.dtype(), x.device());
+    let n = x.numel() as i32;
+    unsafe {
+        match x.dtype() {
+            DType::F32 => launch_gelu_f32(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
+            DType::BF16 => launch_gelu_bf16(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
+            _ => panic!("unsupported dtype for gelu"),
+        }
+    }
+    xserv_cuda::device::synchronize().unwrap();
+    out
+}
+
+pub fn silu(x: &Tensor) -> Tensor {
+    assert!(x.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    let out = Tensor::zeros(x.shape(), x.dtype(), x.device());
+    let n = x.numel() as i32;
+    unsafe {
+        match x.dtype() {
+            DType::F32 => launch_silu_f32(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
+            DType::BF16 => launch_silu_bf16(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
+            _ => panic!("unsupported dtype for silu"),
+        }
+    }
+    xserv_cuda::device::synchronize().unwrap();
+    out
+}
--- a/crates/xserv-kernels/src/embedding.rs
+++ b/crates/xserv-kernels/src/embedding.rs
@@ -0,0 +1,51 @@
+use std::ffi::c_void;
+use xserv_cuda::GpuBuffer;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_embedding_f32(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
+                            num_tokens: i32, hidden_size: i32, stream: *mut c_void);
+    fn launch_embedding_bf16(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
+                             num_tokens: i32, hidden_size: i32, stream: *mut c_void);
+}
+
+/// Embedding lookup: table[token_ids[i]] for each i.
+/// table: [vocab_size, hidden_size], token_ids: [num_tokens] (i32 on CPU)
+pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor {
+    assert_eq!(table.ndim(), 2);
+    assert!(table.is_contiguous());
+    assert!(matches!(table.device(), Device::Cuda(_)));
+
+    let hidden_size = table.shape()[1];
+    let num_tokens = token_ids.len();
+
+    // Upload token_ids to GPU
+    let ids_bytes = unsafe {
+        std::slice::from_raw_parts(
+            token_ids.as_ptr() as *const u8,
+            num_tokens * std::mem::size_of::<u32>(),
+        )
+    };
+    let mut ids_gpu = GpuBuffer::alloc(ids_bytes.len()).expect("alloc token_ids");
+    ids_gpu.copy_from_host(ids_bytes).unwrap();
+
+    let out = Tensor::zeros(&[num_tokens, hidden_size], table.dtype(), table.device());
+
+    unsafe {
+        match table.dtype() {
+            DType::F32 => launch_embedding_f32(
+                table.data_ptr() as _, ids_gpu.as_ptr() as _,
+                out.data_ptr() as *mut c_void,
+                num_tokens as i32, hidden_size as i32, std::ptr::null_mut(),
+            ),
+            DType::BF16 => launch_embedding_bf16(
+                table.data_ptr() as _, ids_gpu.as_ptr() as _,
+                out.data_ptr() as *mut c_void,
+                num_tokens as i32, hidden_size as i32, std::ptr::null_mut(),
+            ),
+            _ => panic!("unsupported dtype for embedding"),
+        }
+    }
+    xserv_cuda::device::synchronize().unwrap();
+    out
+}
--- a/crates/xserv-kernels/src/gemm.rs
+++ b/crates/xserv-kernels/src/gemm.rs
@@ -0,0 +1,151 @@
+use std::ffi::c_void;
+use xserv_cuda::error::{self, Result};
+use xserv_tensor::{DType, Device, Tensor};
+
+#[derive(Debug, Clone, Copy)]
+pub enum GemmBackend {
+    Naive,
+    Tiled,
+    CuBlas,
+}
+
+// --- FFI: custom CUDA kernels ---
+unsafe extern "C" {
+    fn launch_gemm_naive_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemm_naive_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemm_tiled_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemm_tiled_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+}
+
+// --- FFI: cuBLAS ---
+type CublasHandle = *mut c_void;
+
+#[allow(non_upper_case_globals)]
+const CUBLAS_OP_N: i32 = 0;
+
+// cudaDataType
+const CUDA_R_32F: i32 = 0;
+const CUDA_R_16BF: i32 = 14;
+
+// cublasComputeType
+const CUBLAS_COMPUTE_32F: i32 = 68;
+
+unsafe extern "C" {
+    fn cublasCreate_v2(handle: *mut CublasHandle) -> i32;
+    fn cublasDestroy_v2(handle: CublasHandle) -> i32;
+    fn cublasSetStream_v2(handle: CublasHandle, stream: *mut c_void) -> i32;
+    fn cublasGemmEx(
+        handle: CublasHandle,
+        transa: i32, transb: i32,
+        m: i32, n: i32, k: i32,
+        alpha: *const c_void,
+        a: *const c_void, a_type: i32, lda: i32,
+        b: *const c_void, b_type: i32, ldb: i32,
+        beta: *const c_void,
+        c: *mut c_void, c_type: i32, ldc: i32,
+        compute_type: i32,
+        algo: i32,
+    ) -> i32;
+}
+
+pub struct CublasContext {
+    handle: CublasHandle,
+}
+
+impl CublasContext {
+    pub fn new() -> Result<Self> {
+        let mut handle = std::ptr::null_mut();
+        error::check(unsafe { cublasCreate_v2(&mut handle) })?;
+        Ok(Self { handle })
+    }
+}
+
+impl Drop for CublasContext {
+    fn drop(&mut self) {
+        if !self.handle.is_null() {
+            unsafe { cublasDestroy_v2(self.handle) };
+        }
+    }
+}
+
+/// Matrix multiplication: C = A @ B
+/// A: [M, K], B: [K, N], C: [M, N]
+/// All tensors must be contiguous and on the same GPU.
+pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
+    assert_eq!(a.ndim(), 2);
+    assert_eq!(b.ndim(), 2);
+    assert_eq!(a.shape()[1], b.shape()[0], "inner dimension mismatch");
+    assert_eq!(a.dtype(), b.dtype(), "dtype mismatch");
+    assert!(a.is_contiguous() && b.is_contiguous(), "matmul requires contiguous tensors");
+    assert!(matches!(a.device(), Device::Cuda(_)), "matmul requires GPU tensors");
+
+    let m = a.shape()[0];
+    let k = a.shape()[1];
+    let n = b.shape()[1];
+    let dtype = a.dtype();
+
+    let c = Tensor::zeros(&[m, n], dtype, a.device());
+
+    let a_ptr = a.data_ptr() as *const c_void;
+    let b_ptr = b.data_ptr() as *const c_void;
+    let c_ptr = c.data_ptr() as *mut c_void;
+    let null_stream = std::ptr::null_mut();
+
+    match backend {
+        GemmBackend::Naive => {
+            unsafe {
+                match dtype {
+                    DType::F32 => launch_gemm_naive_f32(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    DType::BF16 => launch_gemm_naive_bf16(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    _ => panic!("unsupported dtype for naive GEMM"),
+                }
+            }
+            xserv_cuda::device::synchronize().unwrap();
+        }
+        GemmBackend::Tiled => {
+            unsafe {
+                match dtype {
+                    DType::F32 => launch_gemm_tiled_f32(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    DType::BF16 => launch_gemm_tiled_bf16(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    _ => panic!("unsupported dtype for tiled GEMM"),
+                }
+            }
+            xserv_cuda::device::synchronize().unwrap();
+        }
+        GemmBackend::CuBlas => {
+            // cuBLAS uses column-major, but we have row-major tensors.
+            // Trick: compute C^T = B^T @ A^T, which gives us C in row-major.
+            // cuBLAS sees our row-major data as column-major transposed.
+            let ctx = CublasContext::new().unwrap();
+            let alpha = 1.0f32;
+            let beta = 0.0f32;
+
+            let (a_type, b_type, c_type) = match dtype {
+                DType::F32 => (CUDA_R_32F, CUDA_R_32F, CUDA_R_32F),
+                DType::BF16 => (CUDA_R_16BF, CUDA_R_16BF, CUDA_R_16BF),
+                _ => panic!("unsupported dtype for cuBLAS GEMM"),
+            };
+
+            unsafe {
+                cublasSetStream_v2(ctx.handle, null_stream);
+                // Row-major trick: swap A/B and transpose flags
+                // C(row-major) = A @ B  <=>  C^T(col-major) = B^T @ A^T
+                error::check(cublasGemmEx(
+                    ctx.handle,
+                    CUBLAS_OP_N, CUBLAS_OP_N,
+                    n as i32, m as i32, k as i32,
+                    &alpha as *const f32 as *const c_void,
+                    b_ptr, b_type, n as i32,  // B as col-major = B^T
+                    a_ptr, a_type, k as i32,  // A as col-major = A^T
+                    &beta as *const f32 as *const c_void,
+                    c_ptr, c_type, n as i32,  // C as col-major = C^T
+                    CUBLAS_COMPUTE_32F,
+                    -1, // default algo
+                )).expect("cuBLAS GEMM failed");
+            }
+            xserv_cuda::device::synchronize().unwrap();
+        }
+    }
+
+    c
+}
--- a/crates/xserv-kernels/src/layernorm.rs
+++ b/crates/xserv-kernels/src/layernorm.rs
@@ -0,0 +1,39 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_layernorm_f32(x: *const c_void, gamma: *const c_void, beta: *const c_void,
+                            out: *mut c_void, rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
+    fn launch_layernorm_bf16(x: *const c_void, gamma: *const c_void, beta: *const c_void,
+                             out: *mut c_void, rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
+}
+
+pub fn layernorm(x: &Tensor, gamma: &Tensor, beta: &Tensor, eps: f32) -> Tensor {
+    assert!(x.ndim() >= 1);
+    assert!(x.is_contiguous() && gamma.is_contiguous() && beta.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    let hidden_size = *x.shape().last().unwrap();
+    assert_eq!(gamma.shape(), &[hidden_size]);
+    assert_eq!(beta.shape(), &[hidden_size]);
+
+    let rows = x.numel() / hidden_size;
+    let out = Tensor::zeros(x.shape(), x.dtype(), x.device());
+
+    unsafe {
+        match x.dtype() {
+            DType::F32 => launch_layernorm_f32(
+                x.data_ptr() as _, gamma.data_ptr() as _, beta.data_ptr() as _,
+                out.data_ptr() as *mut c_void,
+                rows as i32, hidden_size as i32, eps, std::ptr::null_mut(),
+            ),
+            DType::BF16 => launch_layernorm_bf16(
+                x.data_ptr() as _, gamma.data_ptr() as _, beta.data_ptr() as _,
+                out.data_ptr() as *mut c_void,
+                rows as i32, hidden_size as i32, eps, std::ptr::null_mut(),
+            ),
+            _ => panic!("unsupported dtype for layernorm"),
+        }
+    }
+    xserv_cuda::device::synchronize().unwrap();
+    out
+}
--- a/crates/xserv-kernels/src/lib.rs
+++ b/crates/xserv-kernels/src/lib.rs
@@ -0,0 +1,15 @@
+pub mod activation;
+pub mod embedding;
+pub mod gemm;
+pub mod layernorm;
+pub mod rmsnorm;
+pub mod rope;
+pub mod softmax;
+
+pub use activation::{gelu, silu};
+pub use embedding::embedding;
+pub use gemm::{matmul, GemmBackend};
+pub use layernorm::layernorm;
+pub use rmsnorm::rmsnorm;
+pub use rope::{rope_inplace, RopeCache};
+pub use softmax::softmax;
--- a/crates/xserv-kernels/src/rmsnorm.rs
+++ b/crates/xserv-kernels/src/rmsnorm.rs
@@ -0,0 +1,37 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_rmsnorm_f32(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_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);
+}
+
+pub fn rmsnorm(x: &Tensor, gamma: &Tensor, eps: f32) -> Tensor {
+    assert!(x.ndim() >= 1);
+    assert!(x.is_contiguous() && gamma.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    let hidden_size = *x.shape().last().unwrap();
+    assert_eq!(gamma.shape(), &[hidden_size]);
+    assert_eq!(x.dtype(), gamma.dtype());
+
+    let rows = x.numel() / hidden_size;
+    let out = Tensor::zeros(x.shape(), x.dtype(), x.device());
+
+    unsafe {
+        match x.dtype() {
+            DType::F32 => launch_rmsnorm_f32(
+                x.data_ptr() as _, gamma.data_ptr() as _, out.data_ptr() as *mut c_void,
+                rows as i32, hidden_size as i32, eps, std::ptr::null_mut(),
+            ),
+            DType::BF16 => launch_rmsnorm_bf16(
+                x.data_ptr() as _, gamma.data_ptr() as _, out.data_ptr() as *mut c_void,
+                rows as i32, hidden_size as i32, eps, std::ptr::null_mut(),
+            ),
+            _ => panic!("unsupported dtype for rmsnorm"),
+        }
+    }
+    xserv_cuda::device::synchronize().unwrap();
+    out
+}
--- a/crates/xserv-kernels/src/rope.rs
+++ b/crates/xserv-kernels/src/rope.rs
@@ -0,0 +1,85 @@
+use std::ffi::c_void;
+use xserv_cuda::GpuBuffer;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_rope_f32(x: *mut c_void, cos_cache: *const c_void, sin_cache: *const c_void,
+                       positions: *const c_void, num_tokens: i32, num_heads: i32,
+                       head_dim: i32, stream: *mut c_void);
+    fn launch_rope_bf16(x: *mut c_void, cos_cache: *const c_void, sin_cache: *const c_void,
+                        positions: *const c_void, num_tokens: i32, num_heads: i32,
+                        head_dim: i32, stream: *mut c_void);
+    fn launch_compute_rope_cache(cos_cache: *mut c_void, sin_cache: *mut c_void,
+                                 max_seq_len: i32, half_dim: i32, theta: f32,
+                                 stream: *mut c_void);
+}
+
+pub struct RopeCache {
+    pub cos: GpuBuffer,
+    pub sin: GpuBuffer,
+    pub max_seq_len: usize,
+    pub half_dim: usize,
+}
+
+impl RopeCache {
+    pub fn new(max_seq_len: usize, head_dim: usize, theta: f32) -> Self {
+        let half_dim = head_dim / 2;
+        let nbytes = max_seq_len * half_dim * std::mem::size_of::<f32>();
+        let mut cos = GpuBuffer::alloc(nbytes).expect("alloc cos_cache");
+        let mut sin = GpuBuffer::alloc(nbytes).expect("alloc sin_cache");
+
+        unsafe {
+            launch_compute_rope_cache(
+                cos.as_mut_ptr() as _, sin.as_mut_ptr() as _,
+                max_seq_len as i32, half_dim as i32, theta, std::ptr::null_mut(),
+            );
+        }
+        xserv_cuda::device::synchronize().unwrap();
+
+        Self { cos, sin, max_seq_len, half_dim }
+    }
+}
+
+/// Apply RoPE in-place to x.
+/// x: [num_tokens, num_heads, head_dim] on GPU
+/// positions: [num_tokens] (u32 on CPU, will be uploaded)
+pub fn rope_inplace(x: &Tensor, cache: &RopeCache, positions: &[u32]) {
+    assert_eq!(x.ndim(), 3);
+    assert!(x.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+    let num_tokens = x.shape()[0];
+    let num_heads = x.shape()[1];
+    let head_dim = x.shape()[2];
+    assert_eq!(head_dim / 2, cache.half_dim);
+    assert_eq!(positions.len(), num_tokens);
+
+    let pos_bytes = unsafe {
+        std::slice::from_raw_parts(
+            positions.as_ptr() as *const u8,
+            num_tokens * std::mem::size_of::<u32>(),
+        )
+    };
+    let mut pos_gpu = GpuBuffer::alloc(pos_bytes.len()).expect("alloc positions");
+    pos_gpu.copy_from_host(pos_bytes).unwrap();
+
+    unsafe {
+        match x.dtype() {
+            DType::F32 => launch_rope_f32(
+                x.data_ptr() as *mut c_void,
+                cache.cos.as_ptr() as _, cache.sin.as_ptr() as _,
+                pos_gpu.as_ptr() as _,
+                num_tokens as i32, num_heads as i32, head_dim as i32,
+                std::ptr::null_mut(),
+            ),
+            DType::BF16 => launch_rope_bf16(
+                x.data_ptr() as *mut c_void,
+                cache.cos.as_ptr() as _, cache.sin.as_ptr() as _,
+                pos_gpu.as_ptr() as _,
+                num_tokens as i32, num_heads as i32, head_dim as i32,
+                std::ptr::null_mut(),
+            ),
+            _ => panic!("unsupported dtype for rope"),
+        }
+    }
+    xserv_cuda::device::synchronize().unwrap();
+}
--- a/crates/xserv-kernels/src/softmax.rs
+++ b/crates/xserv-kernels/src/softmax.rs
@@ -0,0 +1,34 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_softmax_f32(x: *const c_void, out: *mut c_void, rows: i32, cols: i32, stream: *mut c_void);
+    fn launch_softmax_bf16(x: *const c_void, out: *mut c_void, rows: i32, cols: i32, stream: *mut c_void);
+}
+
+/// Softmax along the last dimension.
+pub fn softmax(x: &Tensor) -> Tensor {
+    assert!(x.ndim() >= 1);
+    assert!(x.is_contiguous());
+    assert!(matches!(x.device(), Device::Cuda(_)));
+
+    let cols = *x.shape().last().unwrap();
+    let rows = x.numel() / cols;
+    let out = Tensor::zeros(x.shape(), x.dtype(), x.device());
+
+    unsafe {
+        match x.dtype() {
+            DType::F32 => launch_softmax_f32(
+                x.data_ptr() as _, out.data_ptr() as *mut c_void,
+                rows as i32, cols as i32, std::ptr::null_mut(),
+            ),
+            DType::BF16 => launch_softmax_bf16(
+                x.data_ptr() as _, out.data_ptr() as *mut c_void,
+                rows as i32, cols as i32, std::ptr::null_mut(),
+            ),
+            _ => panic!("unsupported dtype for softmax"),
+        }
+    }
+    xserv_cuda::device::synchronize().unwrap();
+    out
+}
--- a/crates/xserv-kernels/tests/gemm_test.rs
+++ b/crates/xserv-kernels/tests/gemm_test.rs
@@ -0,0 +1,152 @@
+use half::bf16;
+use xserv_kernels::{matmul, GemmBackend};
+use xserv_tensor::{Device, Tensor};
+
+fn cpu_matmul_f32(a: &[f32], b: &[f32], m: usize, n: usize, k: usize) -> Vec<f32> {
+    let mut c = vec![0.0f32; m * n];
+    for i in 0..m {
+        for j in 0..n {
+            let mut sum = 0.0f32;
+            for kk in 0..k {
+                sum += a[i * k + kk] * b[kk * n + j];
+            }
+            c[i * n + j] = sum;
+        }
+    }
+    c
+}
+
+fn check_close_f32(result: &[f32], expected: &[f32], atol: f32) {
+    assert_eq!(result.len(), expected.len());
+    for (i, (r, e)) in result.iter().zip(expected).enumerate() {
+        assert!(
+            (r - e).abs() <= atol,
+            "mismatch at index {i}: got {r}, expected {e}, diff {}",
+            (r - e).abs()
+        );
+    }
+}
+
+fn check_close_bf16(result: &[bf16], expected: &[f32], atol: f32) {
+    assert_eq!(result.len(), expected.len());
+    for (i, (r, e)) in result.iter().zip(expected).enumerate() {
+        let rv = r.to_f32();
+        assert!(
+            (rv - e).abs() <= atol,
+            "mismatch at index {i}: got {rv}, expected {e}, diff {}",
+            (rv - e).abs()
+        );
+    }
+}
+
+fn run_gemm_test_f32(backend: GemmBackend, m: usize, n: usize, k: usize) {
+    xserv_cuda::device::set_device(0).unwrap();
+
+    let a_data: Vec<f32> = (0..m * k).map(|i| ((i % 7) as f32 - 3.0) * 0.1).collect();
+    let b_data: Vec<f32> = (0..k * n).map(|i| ((i % 11) as f32 - 5.0) * 0.1).collect();
+    let expected = cpu_matmul_f32(&a_data, &b_data, m, n, k);
+
+    let a = Tensor::from_slice(&a_data, &[m, k]).to_device(Device::Cuda(0));
+    let b = Tensor::from_slice(&b_data, &[k, n]).to_device(Device::Cuda(0));
+    let c = matmul(&a, &b, backend);
+
+    let c_cpu = c.to_device(Device::Cpu);
+    check_close_f32(c_cpu.as_slice::<f32>(), &expected, 1e-4);
+}
+
+fn run_gemm_test_bf16(backend: GemmBackend, m: usize, n: usize, k: usize) {
+    xserv_cuda::device::set_device(0).unwrap();
+
+    let a_f32: Vec<f32> = (0..m * k).map(|i| ((i % 7) as f32 - 3.0) * 0.1).collect();
+    let b_f32: Vec<f32> = (0..k * n).map(|i| ((i % 11) as f32 - 5.0) * 0.1).collect();
+    let expected = cpu_matmul_f32(&a_f32, &b_f32, m, n, k);
+
+    let a_data: Vec<bf16> = a_f32.iter().map(|&v| bf16::from_f32(v)).collect();
+    let b_data: Vec<bf16> = b_f32.iter().map(|&v| bf16::from_f32(v)).collect();
+
+    let a = Tensor::from_slice(&a_data, &[m, k]).to_device(Device::Cuda(0));
+    let b = Tensor::from_slice(&b_data, &[k, n]).to_device(Device::Cuda(0));
+    let c = matmul(&a, &b, backend);
+
+    let c_cpu = c.to_device(Device::Cpu);
+    check_close_bf16(c_cpu.as_slice::<bf16>(), &expected, 0.1);
+}
+
+// --- F32 tests ---
+
+#[test]
+fn test_gemm_naive_f32_small() { run_gemm_test_f32(GemmBackend::Naive, 4, 4, 4); }
+
+#[test]
+fn test_gemm_naive_f32_medium() { run_gemm_test_f32(GemmBackend::Naive, 64, 64, 64); }
+
+#[test]
+fn test_gemm_naive_f32_rect() { run_gemm_test_f32(GemmBackend::Naive, 32, 64, 48); }
+
+#[test]
+fn test_gemm_tiled_f32_small() { run_gemm_test_f32(GemmBackend::Tiled, 4, 4, 4); }
+
+#[test]
+fn test_gemm_tiled_f32_medium() { run_gemm_test_f32(GemmBackend::Tiled, 128, 128, 128); }
+
+#[test]
+fn test_gemm_tiled_f32_rect() { run_gemm_test_f32(GemmBackend::Tiled, 65, 33, 97); }
+
+#[test]
+fn test_gemm_cublas_f32_small() { run_gemm_test_f32(GemmBackend::CuBlas, 4, 4, 4); }
+
+#[test]
+fn test_gemm_cublas_f32_medium() { run_gemm_test_f32(GemmBackend::CuBlas, 256, 256, 256); }
+
+#[test]
+fn test_gemm_cublas_f32_rect() { run_gemm_test_f32(GemmBackend::CuBlas, 65, 33, 97); }
+
+// --- BF16 tests ---
+
+#[test]
+fn test_gemm_naive_bf16_small() { run_gemm_test_bf16(GemmBackend::Naive, 4, 4, 4); }
+
+#[test]
+fn test_gemm_naive_bf16_medium() { run_gemm_test_bf16(GemmBackend::Naive, 64, 64, 64); }
+
+#[test]
+fn test_gemm_tiled_bf16_small() { run_gemm_test_bf16(GemmBackend::Tiled, 4, 4, 4); }
+
+#[test]
+fn test_gemm_tiled_bf16_medium() { run_gemm_test_bf16(GemmBackend::Tiled, 128, 128, 128); }
+
+#[test]
+fn test_gemm_cublas_bf16_small() { run_gemm_test_bf16(GemmBackend::CuBlas, 4, 4, 4); }
+
+#[test]
+fn test_gemm_cublas_bf16_medium() { run_gemm_test_bf16(GemmBackend::CuBlas, 256, 256, 256); }
+
+// --- Larger benchmark-style tests ---
+
+#[test]
+fn test_gemm_cublas_f32_1024() { run_gemm_test_f32(GemmBackend::CuBlas, 1024, 1024, 1024); }
+
+#[test]
+fn test_gemm_consistency_all_backends() {
+    xserv_cuda::device::set_device(0).unwrap();
+
+    let m = 64;
+    let n = 64;
+    let k = 64;
+    let a_data: Vec<f32> = (0..m * k).map(|i| ((i % 7) as f32 - 3.0) * 0.1).collect();
+    let b_data: Vec<f32> = (0..k * n).map(|i| ((i % 11) as f32 - 5.0) * 0.1).collect();
+
+    let a = Tensor::from_slice(&a_data, &[m, k]).to_device(Device::Cuda(0));
+    let b = Tensor::from_slice(&b_data, &[k, n]).to_device(Device::Cuda(0));
+
+    let c_naive = matmul(&a, &b, GemmBackend::Naive).to_device(Device::Cpu);
+    let c_tiled = matmul(&a, &b, GemmBackend::Tiled).to_device(Device::Cpu);
+    let c_cublas = matmul(&a, &b, GemmBackend::CuBlas).to_device(Device::Cpu);
+
+    let naive = c_naive.as_slice::<f32>();
+    let tiled = c_tiled.as_slice::<f32>();
+    let cublas = c_cublas.as_slice::<f32>();
+
+    check_close_f32(naive, cublas, 1e-4);
+    check_close_f32(tiled, cublas, 1e-4);
+}
--- a/crates/xserv-kernels/tests/ops_test.rs
+++ b/crates/xserv-kernels/tests/ops_test.rs
@@ -0,0 +1,302 @@
+use half::bf16;
+use xserv_kernels::*;
+use xserv_tensor::{Device, Tensor};
+
+fn init() { xserv_cuda::device::set_device(0).unwrap(); }
+
+// --- CPU reference implementations ---
+
+fn cpu_rmsnorm(x: &[f32], gamma: &[f32], eps: f32, hidden: usize) -> Vec<f32> {
+    let rows = x.len() / hidden;
+    let mut out = vec![0.0f32; x.len()];
+    for r in 0..rows {
+        let row = &x[r * hidden..(r + 1) * hidden];
+        let sum_sq: f32 = row.iter().map(|v| v * v).sum();
+        let rms_inv = 1.0 / (sum_sq / hidden as f32 + eps).sqrt();
+        for i in 0..hidden {
+            out[r * hidden + i] = row[i] * rms_inv * gamma[i];
+        }
+    }
+    out
+}
+
+fn cpu_layernorm(x: &[f32], gamma: &[f32], beta: &[f32], eps: f32, hidden: usize) -> Vec<f32> {
+    let rows = x.len() / hidden;
+    let mut out = vec![0.0f32; x.len()];
+    for r in 0..rows {
+        let row = &x[r * hidden..(r + 1) * hidden];
+        let mean: f32 = row.iter().sum::<f32>() / hidden as f32;
+        let var: f32 = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / hidden as f32;
+        let inv_std = 1.0 / (var + eps).sqrt();
+        for i in 0..hidden {
+            out[r * hidden + i] = gamma[i] * (row[i] - mean) * inv_std + beta[i];
+        }
+    }
+    out
+}
+
+fn cpu_gelu(x: &[f32]) -> Vec<f32> {
+    let sqrt_2_over_pi = 0.7978845608f32;
+    x.iter().map(|&v| {
+        let inner = sqrt_2_over_pi * (v + 0.044715 * v * v * v);
+        0.5 * v * (1.0 + inner.tanh())
+    }).collect()
+}
+
+fn cpu_silu(x: &[f32]) -> Vec<f32> {
+    x.iter().map(|&v| v / (1.0 + (-v).exp())).collect()
+}
+
+fn cpu_softmax(x: &[f32], cols: usize) -> Vec<f32> {
+    let rows = x.len() / cols;
+    let mut out = vec![0.0f32; x.len()];
+    for r in 0..rows {
+        let row = &x[r * cols..(r + 1) * cols];
+        let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+        let exps: Vec<f32> = row.iter().map(|v| (v - max).exp()).collect();
+        let sum: f32 = exps.iter().sum();
+        for i in 0..cols {
+            out[r * cols + i] = exps[i] / sum;
+        }
+    }
+    out
+}
+
+fn cpu_rope(x: &mut [f32], positions: &[u32], num_heads: usize, head_dim: usize, theta: f32) {
+    let half_dim = head_dim / 2;
+    let num_tokens = positions.len();
+    for t in 0..num_tokens {
+        let pos = positions[t] as f32;
+        for h in 0..num_heads {
+            for i in 0..half_dim {
+                let freq = 1.0 / theta.powf(2.0 * i as f32 / head_dim as f32);
+                let angle = pos * freq;
+                let cos_val = angle.cos();
+                let sin_val = angle.sin();
+                let base = (t * num_heads + h) * head_dim;
+                let x0 = x[base + 2 * i];
+                let x1 = x[base + 2 * i + 1];
+                x[base + 2 * i]     = x0 * cos_val - x1 * sin_val;
+                x[base + 2 * i + 1] = x0 * sin_val + x1 * cos_val;
+            }
+        }
+    }
+}
+
+fn check_close(result: &[f32], expected: &[f32], atol: f32, name: &str) {
+    assert_eq!(result.len(), expected.len(), "{name}: length mismatch");
+    let mut max_err = 0.0f32;
+    for (i, (r, e)) in result.iter().zip(expected).enumerate() {
+        let err = (r - e).abs();
+        if err > max_err { max_err = err; }
+        assert!(err <= atol, "{name}: mismatch at [{i}]: got {r}, expected {e}, err {err}");
+    }
+    println!("{name}: max_err = {max_err:.6e}");
+}
+
+fn make_data(n: usize) -> Vec<f32> {
+    (0..n).map(|i| ((i % 17) as f32 - 8.0) * 0.1).collect()
+}
+
+// === RMSNorm ===
+
+#[test]
+fn test_rmsnorm_f32() {
+    init();
+    let hidden = 768;
+    let rows = 4;
+    let x_data = make_data(rows * hidden);
+    let gamma_data: Vec<f32> = (0..hidden).map(|i| 0.5 + (i % 3) as f32 * 0.2).collect();
+    let expected = cpu_rmsnorm(&x_data, &gamma_data, 1e-5, hidden);
+
+    let x = Tensor::from_slice(&x_data, &[rows, hidden]).to_device(Device::Cuda(0));
+    let gamma = Tensor::from_slice(&gamma_data, &[hidden]).to_device(Device::Cuda(0));
+    let out = rmsnorm(&x, &gamma, 1e-5).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-4, "rmsnorm_f32");
+}
+
+#[test]
+fn test_rmsnorm_bf16() {
+    init();
+    let hidden = 768;
+    let rows = 4;
+    let x_f32 = make_data(rows * hidden);
+    let gamma_f32: Vec<f32> = (0..hidden).map(|i| 0.5 + (i % 3) as f32 * 0.2).collect();
+    let expected = cpu_rmsnorm(&x_f32, &gamma_f32, 1e-5, hidden);
+
+    let x_bf16: Vec<bf16> = x_f32.iter().map(|&v| bf16::from_f32(v)).collect();
+    let gamma_bf16: Vec<bf16> = gamma_f32.iter().map(|&v| bf16::from_f32(v)).collect();
+    let x = Tensor::from_slice(&x_bf16, &[rows, hidden]).to_device(Device::Cuda(0));
+    let gamma = Tensor::from_slice(&gamma_bf16, &[hidden]).to_device(Device::Cuda(0));
+    let out = rmsnorm(&x, &gamma, 1e-5).to_device(Device::Cpu);
+
+    let result: Vec<f32> = out.as_slice::<bf16>().iter().map(|v| v.to_f32()).collect();
+    check_close(&result, &expected, 0.05, "rmsnorm_bf16");
+}
+
+// === LayerNorm ===
+
+#[test]
+fn test_layernorm_f32() {
+    init();
+    let hidden = 768;
+    let rows = 4;
+    let x_data = make_data(rows * hidden);
+    let gamma_data: Vec<f32> = (0..hidden).map(|i| 0.8 + (i % 5) as f32 * 0.1).collect();
+    let beta_data: Vec<f32> = (0..hidden).map(|i| ((i % 7) as f32 - 3.0) * 0.01).collect();
+    let expected = cpu_layernorm(&x_data, &gamma_data, &beta_data, 1e-5, hidden);
+
+    let x = Tensor::from_slice(&x_data, &[rows, hidden]).to_device(Device::Cuda(0));
+    let gamma = Tensor::from_slice(&gamma_data, &[hidden]).to_device(Device::Cuda(0));
+    let beta = Tensor::from_slice(&beta_data, &[hidden]).to_device(Device::Cuda(0));
+    let out = layernorm(&x, &gamma, &beta, 1e-5).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-4, "layernorm_f32");
+}
+
+// === GELU ===
+
+#[test]
+fn test_gelu_f32() {
+    init();
+    let data = make_data(10000);
+    let expected = cpu_gelu(&data);
+    let x = Tensor::from_slice(&data, &[10000]).to_device(Device::Cuda(0));
+    let out = gelu(&x).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-5, "gelu_f32");
+}
+
+#[test]
+fn test_gelu_bf16() {
+    init();
+    let data_f32 = make_data(10000);
+    let expected = cpu_gelu(&data_f32);
+    let data_bf16: Vec<bf16> = data_f32.iter().map(|&v| bf16::from_f32(v)).collect();
+    let x = Tensor::from_slice(&data_bf16, &[10000]).to_device(Device::Cuda(0));
+    let out = gelu(&x).to_device(Device::Cpu);
+    let result: Vec<f32> = out.as_slice::<bf16>().iter().map(|v| v.to_f32()).collect();
+    check_close(&result, &expected, 0.02, "gelu_bf16");
+}
+
+// === SiLU ===
+
+#[test]
+fn test_silu_f32() {
+    init();
+    let data = make_data(10000);
+    let expected = cpu_silu(&data);
+    let x = Tensor::from_slice(&data, &[10000]).to_device(Device::Cuda(0));
+    let out = silu(&x).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-5, "silu_f32");
+}
+
+// === Softmax ===
+
+#[test]
+fn test_softmax_f32() {
+    init();
+    let rows = 8;
+    let cols = 256;
+    let data = make_data(rows * cols);
+    let expected = cpu_softmax(&data, cols);
+    let x = Tensor::from_slice(&data, &[rows, cols]).to_device(Device::Cuda(0));
+    let out = softmax(&x).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-5, "softmax_f32");
+}
+
+#[test]
+fn test_softmax_sum_to_one() {
+    init();
+    let rows = 4;
+    let cols = 2048;
+    let data: Vec<f32> = (0..rows * cols).map(|i| ((i % 31) as f32 - 15.0) * 0.5).collect();
+    let x = Tensor::from_slice(&data, &[rows, cols]).to_device(Device::Cuda(0));
+    let out = softmax(&x).to_device(Device::Cpu);
+    let result = out.as_slice::<f32>();
+    for r in 0..rows {
+        let row_sum: f32 = result[r * cols..(r + 1) * cols].iter().sum();
+        assert!((row_sum - 1.0).abs() < 1e-5, "softmax row {r} sum = {row_sum}");
+    }
+}
+
+#[test]
+fn test_softmax_large_values() {
+    init();
+    let data = vec![1000.0f32, 1001.0, 999.0, 1000.5];
+    let expected = cpu_softmax(&data, 4);
+    let x = Tensor::from_slice(&data, &[1, 4]).to_device(Device::Cuda(0));
+    let out = softmax(&x).to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-5, "softmax_large");
+}
+
+// === Embedding ===
+
+#[test]
+fn test_embedding_f32() {
+    init();
+    let vocab_size = 100;
+    let hidden = 64;
+    let table_data: Vec<f32> = (0..vocab_size * hidden).map(|i| i as f32 * 0.01).collect();
+    let token_ids: Vec<u32> = vec![0, 5, 99, 42, 1];
+
+    let table = Tensor::from_slice(&table_data, &[vocab_size, hidden]).to_device(Device::Cuda(0));
+    let out = embedding(&table, &token_ids).to_device(Device::Cpu);
+
+    assert_eq!(out.shape(), &[5, hidden]);
+    let result = out.as_slice::<f32>();
+    for (seq_idx, &tid) in token_ids.iter().enumerate() {
+        for i in 0..hidden {
+            let expected = table_data[tid as usize * hidden + i];
+            let got = result[seq_idx * hidden + i];
+            assert!((got - expected).abs() < 1e-6,
+                "embedding mismatch at [{seq_idx},{i}]: got {got}, expected {expected}");
+        }
+    }
+}
+
+// === RoPE ===
+
+#[test]
+fn test_rope_f32() {
+    init();
+    let num_tokens = 4;
+    let num_heads = 2;
+    let head_dim = 8;
+    let theta = 10000.0f32;
+    let positions: Vec<u32> = vec![0, 1, 2, 3];
+
+    let x_data: Vec<f32> = (0..num_tokens * num_heads * head_dim)
+        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
+        .collect();
+    let mut expected = x_data.clone();
+    cpu_rope(&mut expected, &positions, num_heads, head_dim, theta);
+
+    let x = Tensor::from_slice(&x_data, &[num_tokens, num_heads, head_dim])
+        .to_device(Device::Cuda(0));
+    let cache = RopeCache::new(64, head_dim, theta);
+    rope_inplace(&x, &cache, &positions);
+
+    let out = x.to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &expected, 1e-4, "rope_f32");
+}
+
+#[test]
+fn test_rope_position_0_identity() {
+    init();
+    // At position 0, all angles are 0, so cos=1, sin=0 → identity transform
+    let num_tokens = 1;
+    let num_heads = 2;
+    let head_dim = 8;
+    let positions: Vec<u32> = vec![0];
+
+    let x_data: Vec<f32> = (0..num_tokens * num_heads * head_dim)
+        .map(|i| (i as f32 + 1.0) * 0.1)
+        .collect();
+
+    let x = Tensor::from_slice(&x_data, &[num_tokens, num_heads, head_dim])
+        .to_device(Device::Cuda(0));
+    let cache = RopeCache::new(64, head_dim, 10000.0);
+    rope_inplace(&x, &cache, &positions);
+
+    let out = x.to_device(Device::Cpu);
+    check_close(out.as_slice::<f32>(), &x_data, 1e-6, "rope_pos0");
+}
--- a/csrc/activation/activations.cu
+++ b/csrc/activation/activations.cu
@@ -0,0 +1,66 @@
+#include <cuda_bf16.h>
+#include <math.h>
+
+// GELU (tanh approximation):
+// gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
+__device__ __forceinline__ float gelu_f(float x) {
+    const float SQRT_2_OVER_PI = 0.7978845608f;
+    float cube = x * x * x;
+    float inner = SQRT_2_OVER_PI * (x + 0.044715f * cube);
+    return 0.5f * x * (1.0f + tanhf(inner));
+}
+
+// SiLU (Swish): silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
+__device__ __forceinline__ float silu_f(float x) {
+    return x / (1.0f + expf(-x));
+}
+
+__global__ void gelu_f32(const float* x, float* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = gelu_f(x[idx]);
+}
+
+__global__ void gelu_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = __float2bfloat16(gelu_f(__bfloat162float(x[idx])));
+}
+
+__global__ void silu_f32(const float* x, float* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = silu_f(x[idx]);
+}
+
+__global__ void silu_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) out[idx] = __float2bfloat16(silu_f(__bfloat162float(x[idx])));
+}
+
+extern "C" {
+
+void launch_gelu_f32(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    gelu_f32<<<grid, block, 0, (cudaStream_t)stream>>>((const float*)x, (float*)out, n);
+}
+
+void launch_gelu_bf16(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    gelu_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
+}
+
+void launch_silu_f32(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    silu_f32<<<grid, block, 0, (cudaStream_t)stream>>>((const float*)x, (float*)out, n);
+}
+
+void launch_silu_bf16(const void* x, void* out, int n, void* stream) {
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    silu_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
+}
+
+}
--- a/csrc/common.cuh
+++ b/csrc/common.cuh
@@ -0,0 +1,50 @@
+#pragma once
+#include <cuda_bf16.h>
+
+// --- Warp-level reductions (no shared memory needed) ---
+
+__device__ __forceinline__ float warp_reduce_sum(float val) {
+    #pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1)
+        val += __shfl_down_sync(0xffffffff, val, offset);
+    return val;
+}
+
+__device__ __forceinline__ float warp_reduce_max(float val) {
+    #pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1)
+        val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
+    return val;
+}
+
+// --- Block-level reductions ---
+
+__device__ __forceinline__ float block_reduce_sum(float val) {
+    __shared__ float shared[32];
+    int lane = threadIdx.x & 31;
+    int warp_id = threadIdx.x >> 5;
+    int num_warps = (blockDim.x + 31) >> 5;
+
+    val = warp_reduce_sum(val);
+    if (lane == 0) shared[warp_id] = val;
+    __syncthreads();
+
+    val = (threadIdx.x < num_warps) ? shared[threadIdx.x] : 0.0f;
+    if (warp_id == 0) val = warp_reduce_sum(val);
+    return val;
+}
+
+__device__ __forceinline__ float block_reduce_max(float val) {
+    __shared__ float shared[32];
+    int lane = threadIdx.x & 31;
+    int warp_id = threadIdx.x >> 5;
+    int num_warps = (blockDim.x + 31) >> 5;
+
+    val = warp_reduce_max(val);
+    if (lane == 0) shared[warp_id] = val;
+    __syncthreads();
+
+    val = (threadIdx.x < num_warps) ? shared[threadIdx.x] : -INFINITY;
+    if (warp_id == 0) val = warp_reduce_max(val);
+    return val;
+}
--- a/csrc/embedding/embedding.cu
+++ b/csrc/embedding/embedding.cu
@@ -0,0 +1,55 @@
+#include <cuda_bf16.h>
+
+// Embedding lookup: out[seq_idx] = table[token_ids[seq_idx]]
+// Grid: num_tokens, Block: handles hidden_size elements per token.
+
+__global__ void embedding_f32(
+    const float* __restrict__ table,    // [vocab_size, hidden_size]
+    const int* __restrict__ token_ids,  // [num_tokens]
+    float* __restrict__ out,            // [num_tokens, hidden_size]
+    int hidden_size
+) {
+    int token_idx = blockIdx.x;
+    int tid = token_ids[token_idx];
+    const float* row = table + tid * hidden_size;
+    float* dst = out + token_idx * hidden_size;
+
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        dst[i] = row[i];
+    }
+}
+
+__global__ void embedding_bf16(
+    const __nv_bfloat16* __restrict__ table,
+    const int* __restrict__ token_ids,
+    __nv_bfloat16* __restrict__ out,
+    int hidden_size
+) {
+    int token_idx = blockIdx.x;
+    int tid = token_ids[token_idx];
+    const __nv_bfloat16* row = table + tid * hidden_size;
+    __nv_bfloat16* dst = out + token_idx * hidden_size;
+
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        dst[i] = row[i];
+    }
+}
+
+extern "C" {
+
+void launch_embedding_f32(const void* table, const void* token_ids, void* out,
+                          int num_tokens, int hidden_size, void* stream) {
+    int block = (hidden_size < 256) ? hidden_size : 256;
+    embedding_f32<<<num_tokens, block, 0, (cudaStream_t)stream>>>(
+        (const float*)table, (const int*)token_ids, (float*)out, hidden_size);
+}
+
+void launch_embedding_bf16(const void* table, const void* token_ids, void* out,
+                           int num_tokens, int hidden_size, void* stream) {
+    int block = (hidden_size < 256) ? hidden_size : 256;
+    embedding_bf16<<<num_tokens, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)table, (const int*)token_ids,
+        (__nv_bfloat16*)out, hidden_size);
+}
+
+}
--- a/csrc/embedding/rope.cu
+++ b/csrc/embedding/rope.cu
@@ -0,0 +1,116 @@
+#include <cuda_bf16.h>
+#include <math.h>
+
+// RoPE: Rotary Position Embedding
+// For each pair (x[2i], x[2i+1]) at position `pos`:
+//   y[2i]   = x[2i] * cos - x[2i+1] * sin
+//   y[2i+1] = x[2i] * sin + x[2i+1] * cos
+// where cos/sin come from precomputed cos_cache/sin_cache.
+//
+// cos_cache[pos][i] = cos(pos * freq[i])
+// sin_cache[pos][i] = sin(pos * freq[i])
+// freq[i] = 1.0 / (theta ^ (2i / head_dim))
+
+// Apply RoPE in-place to Q or K tensor.
+// x shape: [num_tokens, num_heads, head_dim]
+// cos_cache, sin_cache shape: [max_seq_len, head_dim/2]
+// positions: [num_tokens] — the position index for each token
+
+__global__ void rope_f32(
+    float* __restrict__ x,                    // [num_tokens, num_heads, head_dim]
+    const float* __restrict__ cos_cache,      // [max_seq_len, half_dim]
+    const float* __restrict__ sin_cache,      // [max_seq_len, half_dim]
+    const int* __restrict__ positions,        // [num_tokens]
+    int num_heads, int head_dim
+) {
+    int token_idx = blockIdx.x;
+    int head_idx = blockIdx.y;
+    int half_dim = head_dim / 2;
+    int pair_idx = threadIdx.x;  // which pair (0..half_dim)
+
+    if (pair_idx >= half_dim) return;
+
+    int pos = positions[token_idx];
+    float cos_val = cos_cache[pos * half_dim + pair_idx];
+    float sin_val = sin_cache[pos * half_dim + pair_idx];
+
+    int base = (token_idx * num_heads + head_idx) * head_dim;
+    float x0 = x[base + 2 * pair_idx];
+    float x1 = x[base + 2 * pair_idx + 1];
+
+    x[base + 2 * pair_idx]     = x0 * cos_val - x1 * sin_val;
+    x[base + 2 * pair_idx + 1] = x0 * sin_val + x1 * cos_val;
+}
+
+__global__ void rope_bf16(
+    __nv_bfloat16* __restrict__ x,
+    const float* __restrict__ cos_cache,
+    const float* __restrict__ sin_cache,
+    const int* __restrict__ positions,
+    int num_heads, int head_dim
+) {
+    int token_idx = blockIdx.x;
+    int head_idx = blockIdx.y;
+    int half_dim = head_dim / 2;
+    int pair_idx = threadIdx.x;
+
+    if (pair_idx >= half_dim) return;
+
+    int pos = positions[token_idx];
+    float cos_val = cos_cache[pos * half_dim + pair_idx];
+    float sin_val = sin_cache[pos * half_dim + pair_idx];
+
+    int base = (token_idx * num_heads + head_idx) * head_dim;
+    float x0 = __bfloat162float(x[base + 2 * pair_idx]);
+    float x1 = __bfloat162float(x[base + 2 * pair_idx + 1]);
+
+    x[base + 2 * pair_idx]     = __float2bfloat16(x0 * cos_val - x1 * sin_val);
+    x[base + 2 * pair_idx + 1] = __float2bfloat16(x0 * sin_val + x1 * cos_val);
+}
+
+// Precompute cos/sin cache on GPU
+__global__ void compute_rope_cache(
+    float* __restrict__ cos_cache,  // [max_seq_len, half_dim]
+    float* __restrict__ sin_cache,
+    int max_seq_len, int half_dim, float theta
+) {
+    int pos = blockIdx.x;
+    int i = threadIdx.x;
+    if (i >= half_dim) return;
+
+    float freq = 1.0f / powf(theta, (float)(2 * i) / (float)(2 * half_dim));
+    float angle = (float)pos * freq;
+    cos_cache[pos * half_dim + i] = cosf(angle);
+    sin_cache[pos * half_dim + i] = sinf(angle);
+}
+
+extern "C" {
+
+void launch_rope_f32(void* x, const void* cos_cache, const void* sin_cache,
+                     const void* positions, int num_tokens, int num_heads,
+                     int head_dim, void* stream) {
+    dim3 grid(num_tokens, num_heads);
+    int block = head_dim / 2;
+    rope_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (float*)x, (const float*)cos_cache, (const float*)sin_cache,
+        (const int*)positions, num_heads, head_dim);
+}
+
+void launch_rope_bf16(void* x, const void* cos_cache, const void* sin_cache,
+                      const void* positions, int num_tokens, int num_heads,
+                      int head_dim, void* stream) {
+    dim3 grid(num_tokens, num_heads);
+    int block = head_dim / 2;
+    rope_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (__nv_bfloat16*)x, (const float*)cos_cache, (const float*)sin_cache,
+        (const int*)positions, num_heads, head_dim);
+}
+
+void launch_compute_rope_cache(void* cos_cache, void* sin_cache,
+                               int max_seq_len, int half_dim, float theta,
+                               void* stream) {
+    compute_rope_cache<<<max_seq_len, half_dim, 0, (cudaStream_t)stream>>>(
+        (float*)cos_cache, (float*)sin_cache, max_seq_len, half_dim, theta);
+}
+
+}
--- a/csrc/gemm/naive.cu
+++ b/csrc/gemm/naive.cu
@@ -0,0 +1,62 @@
+#include <cuda_bf16.h>
+
+// Naive GEMM: each thread computes one element of C.
+// C[i][j] = sum_k A[i][k] * B[k][j]
+// All matrices are row-major.
+__global__ void gemm_naive_bf16(
+    const __nv_bfloat16* A, const __nv_bfloat16* B, __nv_bfloat16* C,
+    int M, int N, int K
+) {
+    int row = blockIdx.y * blockDim.y + threadIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (row < M && col < N) {
+        float sum = 0.0f;
+        for (int k = 0; k < K; k++) {
+            sum += __bfloat162float(A[row * K + k]) * __bfloat162float(B[k * N + col]);
+        }
+        C[row * N + col] = __float2bfloat16(sum);
+    }
+}
+
+__global__ void gemm_naive_f32(
+    const float* A, const float* B, float* C,
+    int M, int N, int K
+) {
+    int row = blockIdx.y * blockDim.y + threadIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (row < M && col < N) {
+        float sum = 0.0f;
+        for (int k = 0; k < K; k++) {
+            sum += A[row * K + k] * B[k * N + col];
+        }
+        C[row * N + col] = sum;
+    }
+}
+
+extern "C" {
+
+void launch_gemm_naive_bf16(
+    const void* A, const void* B, void* C,
+    int M, int N, int K, void* stream
+) {
+    dim3 block(16, 16);
+    dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
+    gemm_naive_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)A, (const __nv_bfloat16*)B, (__nv_bfloat16*)C, M, N, K
+    );
+}
+
+void launch_gemm_naive_f32(
+    const void* A, const void* B, void* C,
+    int M, int N, int K, void* stream
+) {
+    dim3 block(16, 16);
+    dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
+    gemm_naive_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const float*)A, (const float*)B, (float*)C, M, N, K
+    );
+}
+
+} // extern "C"
--- a/csrc/gemm/tiled.cu
+++ b/csrc/gemm/tiled.cu
@@ -0,0 +1,116 @@
+#include <cuda_bf16.h>
+
+// Tiled GEMM using shared memory.
+// Each thread block loads TILE_SIZE x TILE_SIZE tiles of A and B
+// into shared memory, then computes a partial dot product.
+#define TILE_SIZE 32
+
+__global__ void gemm_tiled_f32(
+    const float* A, const float* B, float* C,
+    int M, int N, int K
+) {
+    __shared__ float As[TILE_SIZE][TILE_SIZE];
+    __shared__ float Bs[TILE_SIZE][TILE_SIZE];
+
+    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
+    int col = blockIdx.x * TILE_SIZE + threadIdx.x;
+
+    float sum = 0.0f;
+
+    for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
+        // Load tile of A
+        int a_col = t * TILE_SIZE + threadIdx.x;
+        if (row < M && a_col < K) {
+            As[threadIdx.y][threadIdx.x] = A[row * K + a_col];
+        } else {
+            As[threadIdx.y][threadIdx.x] = 0.0f;
+        }
+
+        // Load tile of B
+        int b_row = t * TILE_SIZE + threadIdx.y;
+        if (b_row < K && col < N) {
+            Bs[threadIdx.y][threadIdx.x] = B[b_row * N + col];
+        } else {
+            Bs[threadIdx.y][threadIdx.x] = 0.0f;
+        }
+
+        __syncthreads();
+
+        for (int k = 0; k < TILE_SIZE; k++) {
+            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
+        }
+
+        __syncthreads();
+    }
+
+    if (row < M && col < N) {
+        C[row * N + col] = sum;
+    }
+}
+
+__global__ void gemm_tiled_bf16(
+    const __nv_bfloat16* A, const __nv_bfloat16* B, __nv_bfloat16* C,
+    int M, int N, int K
+) {
+    __shared__ float As[TILE_SIZE][TILE_SIZE];
+    __shared__ float Bs[TILE_SIZE][TILE_SIZE];
+
+    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
+    int col = blockIdx.x * TILE_SIZE + threadIdx.x;
+
+    float sum = 0.0f;
+
+    for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
+        int a_col = t * TILE_SIZE + threadIdx.x;
+        if (row < M && a_col < K) {
+            As[threadIdx.y][threadIdx.x] = __bfloat162float(A[row * K + a_col]);
+        } else {
+            As[threadIdx.y][threadIdx.x] = 0.0f;
+        }
+
+        int b_row = t * TILE_SIZE + threadIdx.y;
+        if (b_row < K && col < N) {
+            Bs[threadIdx.y][threadIdx.x] = __bfloat162float(B[b_row * N + col]);
+        } else {
+            Bs[threadIdx.y][threadIdx.x] = 0.0f;
+        }
+
+        __syncthreads();
+
+        for (int k = 0; k < TILE_SIZE; k++) {
+            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
+        }
+
+        __syncthreads();
+    }
+
+    if (row < M && col < N) {
+        C[row * N + col] = __float2bfloat16(sum);
+    }
+}
+
+extern "C" {
+
+void launch_gemm_tiled_f32(
+    const void* A, const void* B, void* C,
+    int M, int N, int K, void* stream
+) {
+    dim3 block(TILE_SIZE, TILE_SIZE);
+    dim3 grid((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE);
+    gemm_tiled_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const float*)A, (const float*)B, (float*)C, M, N, K
+    );
+}
+
+void launch_gemm_tiled_bf16(
+    const void* A, const void* B, void* C,
+    int M, int N, int K, void* stream
+) {
+    dim3 block(TILE_SIZE, TILE_SIZE);
+    dim3 grid((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE);
+    gemm_tiled_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)A, (const __nv_bfloat16*)B, (__nv_bfloat16*)C, M, N, K
+    );
+}
+
+} // extern "C"
--- a/csrc/normalization/layernorm.cu
+++ b/csrc/normalization/layernorm.cu
@@ -0,0 +1,102 @@
+#include "../common.cuh"
+
+// LayerNorm: y[i] = gamma[i] * (x[i] - mean) / sqrt(var + eps) + beta[i]
+// Each block processes one row of shape [hidden_size].
+
+__global__ void layernorm_f32(
+    const float* __restrict__ x,
+    const float* __restrict__ gamma,
+    const float* __restrict__ beta,
+    float* __restrict__ out,
+    int hidden_size, float eps
+) {
+    int row = blockIdx.x;
+    const float* x_row = x + row * hidden_size;
+    float* out_row = out + row * hidden_size;
+
+    // Welford online: compute mean and variance in one pass
+    float local_sum = 0.0f;
+    float local_sum_sq = 0.0f;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float v = x_row[i];
+        local_sum += v;
+        local_sum_sq += v * v;
+    }
+    local_sum = block_reduce_sum(local_sum);
+    local_sum_sq = block_reduce_sum(local_sum_sq);
+
+    __shared__ float s_mean, s_inv_std;
+    if (threadIdx.x == 0) {
+        float mean = local_sum / hidden_size;
+        float var = local_sum_sq / hidden_size - mean * mean;
+        s_mean = mean;
+        s_inv_std = rsqrtf(var + eps);
+    }
+    __syncthreads();
+
+    float mean = s_mean;
+    float inv_std = s_inv_std;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        out_row[i] = gamma[i] * (x_row[i] - mean) * inv_std + beta[i];
+    }
+}
+
+__global__ void layernorm_bf16(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ gamma,
+    const __nv_bfloat16* __restrict__ beta,
+    __nv_bfloat16* __restrict__ out,
+    int hidden_size, float eps
+) {
+    int row = blockIdx.x;
+    const __nv_bfloat16* x_row = x + row * hidden_size;
+    __nv_bfloat16* out_row = out + row * hidden_size;
+
+    float local_sum = 0.0f;
+    float local_sum_sq = 0.0f;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float v = __bfloat162float(x_row[i]);
+        local_sum += v;
+        local_sum_sq += v * v;
+    }
+    local_sum = block_reduce_sum(local_sum);
+    local_sum_sq = block_reduce_sum(local_sum_sq);
+
+    __shared__ float s_mean, s_inv_std;
+    if (threadIdx.x == 0) {
+        float mean = local_sum / hidden_size;
+        float var = local_sum_sq / hidden_size - mean * mean;
+        s_mean = mean;
+        s_inv_std = rsqrtf(var + eps);
+    }
+    __syncthreads();
+
+    float mean = s_mean;
+    float inv_std = s_inv_std;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float v = __bfloat162float(x_row[i]);
+        float g = __bfloat162float(gamma[i]);
+        float b = __bfloat162float(beta[i]);
+        out_row[i] = __float2bfloat16(g * (v - mean) * inv_std + b);
+    }
+}
+
+extern "C" {
+
+void launch_layernorm_f32(const void* x, const void* gamma, const void* beta,
+                          void* out, int rows, int hidden_size, float eps, void* stream) {
+    int block = (hidden_size < 1024) ? hidden_size : 1024;
+    layernorm_f32<<<rows, block, 0, (cudaStream_t)stream>>>(
+        (const float*)x, (const float*)gamma, (const float*)beta,
+        (float*)out, hidden_size, eps);
+}
+
+void launch_layernorm_bf16(const void* x, const void* gamma, const void* beta,
+                           void* out, int rows, int hidden_size, float eps, void* stream) {
+    int block = (hidden_size < 1024) ? hidden_size : 1024;
+    layernorm_bf16<<<rows, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (const __nv_bfloat16*)gamma, (const __nv_bfloat16*)beta,
+        (__nv_bfloat16*)out, hidden_size, eps);
+}
+
+}
--- a/csrc/normalization/rmsnorm.cu
+++ b/csrc/normalization/rmsnorm.cu
@@ -0,0 +1,83 @@
+#include "../common.cuh"
+
+// RMSNorm: y[i] = x[i] * rsqrt(mean(x²) + eps) * gamma[i]
+// Each block processes one row of shape [hidden_size].
+
+__global__ void rmsnorm_f32(
+    const float* __restrict__ x,
+    const float* __restrict__ gamma,
+    float* __restrict__ out,
+    int hidden_size, float eps
+) {
+    int row = blockIdx.x;
+    const float* x_row = x + row * hidden_size;
+    float* out_row = out + row * hidden_size;
+
+    float sum_sq = 0.0f;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float v = x_row[i];
+        sum_sq += v * v;
+    }
+    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();
+
+    float rms_inv = s_rms_inv;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        out_row[i] = x_row[i] * rms_inv * gamma[i];
+    }
+}
+
+__global__ void rmsnorm_bf16(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ gamma,
+    __nv_bfloat16* __restrict__ out,
+    int hidden_size, float eps
+) {
+    int row = blockIdx.x;
+    const __nv_bfloat16* x_row = x + row * hidden_size;
+    __nv_bfloat16* out_row = out + row * hidden_size;
+
+    float sum_sq = 0.0f;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float v = __bfloat162float(x_row[i]);
+        sum_sq += v * v;
+    }
+    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();
+
+    float rms_inv = s_rms_inv;
+    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
+        float v = __bfloat162float(x_row[i]);
+        float g = __bfloat162float(gamma[i]);
+        out_row[i] = __float2bfloat16(v * rms_inv * g);
+    }
+}
+
+extern "C" {
+
+void launch_rmsnorm_f32(const void* x, const void* gamma, void* out,
+                        int rows, int hidden_size, float eps, void* stream) {
+    int block = (hidden_size < 1024) ? hidden_size : 1024;
+    rmsnorm_f32<<<rows, block, 0, (cudaStream_t)stream>>>(
+        (const float*)x, (const float*)gamma, (float*)out, hidden_size, eps);
+}
+
+void launch_rmsnorm_bf16(const void* x, const void* gamma, void* out,
+                         int rows, int hidden_size, float eps, void* stream) {
+    int block = (hidden_size < 1024) ? hidden_size : 1024;
+    rmsnorm_bf16<<<rows, block, 0, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (const __nv_bfloat16*)gamma,
+        (__nv_bfloat16*)out, hidden_size, eps);
+}
+
+}
--- a/csrc/reduce/softmax.cu
+++ b/csrc/reduce/softmax.cu
@@ -0,0 +1,106 @@
+#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);
+}
+
+}
--- a/docs/01-cuda-ffi.md
+++ b/docs/01-cuda-ffi.md
@@ -72,9 +72,31 @@ Wraps cudaStream_t. RAII with Drop calling cudaStreamDestroy.
 - `build.rs` uses `cc` crate to compile .cu files, link CUDA runtime

 ## Test Plan
-1. Device info: print GPU name, memory, compute capability, SM count
-2. GpuBuffer: alloc 1GB, H2D copy, D2H copy, verify data
-3. Vector add kernel: launch from Rust, verify output
-4. CachingAllocator: alloc→free→realloc same size uses cache (no new cudaMalloc)
-5. Multi-stream: two concurrent memcpy on different streams
-6. Benchmark: caching allocator vs raw cudaMalloc (100 cycles)
+
+- [x] Device info: print GPU name, memory, compute capability, SM count
+- [x] GpuBuffer: alloc → H2D copy → D2H copy → verify data (256B, 64MB)
+- [x] GpuBuffer: D2D copy 验证
+- [x] GpuBuffer: zero fill 验证
+- [x] Vector add kernel: launch from Rust, verify output
+- [x] CachingAllocator: alloc→free→realloc same size uses cache (no new cudaMalloc)
+- [x] CachingAllocator: 不同 size bucket 独立缓存
+- [x] CudaStream: 创建、同步、Drop
+- [x] PinnedBuffer: page-locked host memory
+- [x] Async copy: H2D async + D2H async via stream
+
+## Takeaways
+
+1. **`cudaDeviceProp` struct 布局不可靠**：CUDA 版本之间 `cudaDeviceProp` 的字段偏移会变化。我们最初用 struct 映射读取 `total_global_mem`，得到了垃圾值（12TB）。正确做法：用 `cudaMemGetInfo` 获取显存信息，用 `cudaDeviceGetAttribute` 获取其他属性。只从 `cudaDeviceProp` 读取 `name` 字段（始终在 struct 最前面，布局稳定）。
+
+2. **Rust 2024 edition 的 unsafe 语义变更**：
+   - `extern "C"` 块必须加 `unsafe` 前缀 → `unsafe extern "C"`
+   - `unsafe fn` 内部的 unsafe 调用也需要显式 `unsafe {}` 块
+   - 这让代码更安全，但初次移植需要注意
+
+3. **`cc` crate 的 CUDA 支持是内置的**：不需要 `features = ["cuda"]`（这个 feature 不存在）。只需 `.cuda(true).cudart("shared")`。
+
+4. **Caching Allocator 的 bucket 策略**：round up to next power of 2（最小 512B）。这意味着申请 513B 会分配 1024B，存在内部碎片。但简单且高效——避免了 free list 中的精确匹配问题。PyTorch 的 CUDACachingAllocator 用了更复杂的策略（best-fit with splitting），但对于推理场景，power-of-2 bucket 已经够用。
+
+5. **`into_raw` + `from_raw` 模式**：GpuBuffer 的 RAII Drop 和 CachingAllocator 的缓存需求冲突——allocator 需要持有裸指针而不触发 Drop。`into_raw()` 消费 self（`mem::forget`），返回裸指针；`from_raw()` 重新封装。这是 Rust 中管理 RAII 生命周期的标准模式。
+
+6. **dash5 环境**：CUDA 12.9 已安装但 `nvcc` 不在 PATH（需要 `/usr/local/cuda/bin`）。Rust 需要手动安装 rustup。无 rsync，用 `tar | ssh tar` 同步代码。开发工作流：本地写码 → tar sync → 远程 build+test。
--- a/docs/02-tensor.md
+++ b/docs/02-tensor.md
@@ -0,0 +1,97 @@
+# Phase 2: Tensor Abstraction Layer — Design Document
+
+## Goal
+
+实现核心 Tensor 类型，支持 CPU/GPU 存储、多种数据类型、strided view 操作，作为后续所有算子和模型的数据基础。
+
+## Module Layout
+
+```
+crates/xserv-tensor/
+├── Cargo.toml
+└── src/
+    ├── lib.rs          # re-exports
+    ├── dtype.rs        # DType enum, TensorDType trait
+    ├── shape.rs        # strides 计算, broadcast 规则
+    ├── storage.rs      # Storage (Arc引用计数), Device enum
+    └── tensor.rs       # Tensor 主体: 创建, 形状操作, 设备迁移
+```
+
+## Key Design Decisions
+
+### DType + TensorDType Trait
+
+```rust
+pub enum DType { F32, F16, BF16 }
+
+pub trait TensorDType: Copy + Send + Sync + 'static {
+    const DTYPE: DType;
+    fn to_f64(self) -> f64;
+    fn from_f64(v: f64) -> Self;
+}
+```
+
+- 用 `half` crate 的 `bf16`/`f16` 表示半精度类型
+- `TensorDType` trait 让 `from_slice<T>` 和 `as_slice<T>` 有类型安全
+- GPU kernel 中通过 `DType` dispatch 到对应的 CUDA 类型 (`__nv_bfloat16` / `float`)
+
+### Storage 引用计数
+
+```rust
+pub struct Storage(Arc<StorageInner>);
+enum StorageInner {
+    Cpu { data: Vec<u8> },
+    Cuda { buffer: GpuBuffer },
+}
+```
+
+- `Arc` 引用计数让 transpose/slice/reshape 能共享底层数据（view 语义）
+- 不实现 CoW（copy-on-write），view 只能读不能写
+- `to_device()` 总是创建新的 Storage
+
+### Strided Tensor
+
+```rust
+pub struct Tensor {
+    storage: Storage,
+    shape: SmallVec<[usize; 4]>,
+    strides: SmallVec<[usize; 4]>,
+    offset: usize,
+    dtype: DType,
+}
+```
+
+- `SmallVec<[usize; 4]>` 避免大多数 tensor (≤4D) 的堆分配
+- `strides` 以元素为单位（不是字节）
+- `offset` 支持 slice 操作（view 到 storage 的中间位置）
+- `is_contiguous()` 检查 strides 是否与 shape 匹配
+- 非 contiguous 的 tensor 调 `contiguous()` 才能送入 CUDA kernel
+
+### Broadcast 规则
+
+实现了 NumPy-style broadcasting:
+- 维度从尾部对齐
+- 大小为 1 的维度可以广播到任意大小
+- `broadcast_strides()` 将 size=1 维度的 stride 置为 0（虚拟广播，不复制数据）
+
+## Test Plan
+
+- [x] from_slice → shape/strides 正确
+- [x] reshape, transpose, squeeze, unsqueeze
+- [x] transpose 后 contiguous() 重排数据
+- [x] BF16 tensor 的精度验证
+- [x] CPU↔GPU roundtrip
+- [x] zeros on GPU → 拷回 CPU 验证全 0
+- [x] broadcast_shape 单元测试
+
+## Takeaways
+
+1. **`SmallVec` 是正确选择**：绝大多数 tensor ≤ 4D，避免了频繁堆分配。LLM 推理中常见的维度是 `[B, S, H]` (3D) 和 `[B, H, S, D]` (4D)。
+
+2. **View 语义的取舍**：Arc 共享 storage 实现了零拷贝 transpose/reshape，但代价是无法原地修改 view 后的 tensor。对于推理引擎这是可以接受的——推理路径上大部分操作是只读的。
+
+3. **contiguous() 的隐性开销**：非 contiguous tensor 在送入 kernel 前需要 `contiguous()` 拷贝。这意味着 `transpose → matmul` 会产生一次额外拷贝。后续优化方向：在 kernel 中直接支持 strided input。
+
+4. **Rust 2024 edition 变化**：`unsafe fn` 内部的 unsafe 调用也需要显式 `unsafe {}` 块，`extern "C"` 块必须加 `unsafe` 前缀。这个 edition 对安全性更严格。
+
+5. **CPU 实现先行**：先在 CPU 上验证逻辑正确性（如 contiguous 重排），再扩展到 GPU。这个策略在后续 phase 中应该继续沿用。
--- a/docs/03-gemm.md
+++ b/docs/03-gemm.md
@@ -0,0 +1,102 @@
+# Phase 3: GEMM — Design Document
+
+## Goal
+
+实现矩阵乘法的多个版本（naive → tiled → cuBLAS），建立 benchmark 对比框架，深入理解 GPU 编程中的内存访问模式和优化手段。
+
+## Module Layout
+
+```
+csrc/gemm/
+├── naive.cu          # 每个 thread 算一个输出元素
+└── tiled.cu          # shared memory tiling, 32x32 tiles
+
+crates/xserv-kernels/
+├── build.rs          # 编译 .cu + 链接 cublas
+└── src/
+    ├── lib.rs
+    └── gemm.rs       # FFI 封装, GemmBackend enum, matmul(), CublasContext
+```
+
+## Kernel Implementations
+
+### Version 1: Naive GEMM
+
+```
+Grid:  (ceil(N/16), ceil(M/16))
+Block: (16, 16)
+每个 thread: C[row][col] = sum_k(A[row][k] * B[k][col])
+```
+
+- 每个 thread 独立遍历 K 维度做点积
+- 所有读取走 global memory，无局部性优化
+- BF16 版本在 FP32 中累加（`__bfloat162float` → 累加 → `__float2bfloat16`）
+
+### Version 2: Tiled GEMM (Shared Memory)
+
+```
+TILE_SIZE = 32
+Grid:  (ceil(N/32), ceil(M/32))
+Block: (32, 32) = 1024 threads
+
+每个 tile iteration:
+  1. 协作加载 A[tile] 和 B[tile] 到 shared memory
+  2. __syncthreads()
+  3. 在 shared memory 中做 32 次乘加
+  4. __syncthreads()
+```
+
+- 每个 global memory 读取被 TILE_SIZE 个 thread 复用
+- 理论上减少 global memory 访问 TILE_SIZE 倍
+- BF16 版本同样在 shared memory 中存 float（FP32 累加）
+
+### Version 3: cuBLAS
+
+- `cublasGemmEx` 支持混合精度
+- **Row-major 适配**：cuBLAS 使用 column-major 布局，我们的 tensor 是 row-major
+  - 利用恒等式：`C = A @ B` (row-major) ⟺ `C^T = B^T @ A^T` (col-major)
+  - 传入 `CUBLAS_OP_N`，让 cuBLAS 把我们的 row-major 数据当作 col-major 的转置
+  - 参数：`m=N, n=M, k=K, lda=N (B), ldb=K (A), ldc=N (C)`
+
+### Backend Registry
+
+```rust
+pub enum GemmBackend { Naive, Tiled, CuBlas }
+pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor;
+```
+
+运行时可切换 backend，方便 benchmark 对比和逐步替换。
+
+## CublasContext
+
+RAII 封装 `cublasHandle_t`，Drop 时调 `cublasDestroy_v2`。
+目前每次 matmul 创建一个新 handle，后续优化为全局复用。
+
+## Test Plan
+
+- [x] F32: naive/tiled/cuBLAS × small(4)/medium(64-256)/rect(65x33x97) 
+- [x] BF16: naive/tiled/cuBLAS × small/medium
+- [x] 三种 backend 在相同输入上输出一致（cross-backend consistency）
+- [x] 非方阵测试（M≠N≠K）
+- [x] 1024x1024 cuBLAS 验证
+
+## Takeaways
+
+1. **Row-major vs Column-major 陷阱**：这是 GEMM 实现中最容易出错的地方。cuBLAS 的 column-major 假设与 C/Rust 的 row-major 冲突。理解 `C=AB` ⟺ `C^T=B^T A^T` 这个恒等式是关键。实际做法：不做任何显式转置，只是交换 A/B 的传入顺序和调整 leading dimension 参数。
+
+2. **BF16 的累加精度**：BF16 只有 ~3 位有效数字（vs FP32 的 ~7 位）。如果在 BF16 中累加 K 次乘法，误差会快速放大。正确做法是**在 FP32 中累加，最后才转回 BF16**。我们的 naive 和 tiled kernel 都遵循了这一点（`float sum = 0.0f`）。cuBLAS 通过 `CUBLAS_COMPUTE_32F` 参数控制。
+
+3. **Shared memory tiling 的核心思想**：global memory 带宽是 GPU 计算的主要瓶颈。通过 shared memory tiling，每个数据从 global memory 读一次，被 TILE_SIZE 个 thread 复用。对于 TILE_SIZE=32，理论上减少 32 倍 global memory 访问。
+
+4. **`__syncthreads()` 的位置关键**：tile 加载后必须同步（确保所有 thread 写完 shared memory），计算后也要同步（防止下一轮加载覆盖还在使用的数据）。漏掉任何一个 sync 都会产生 race condition 导致结果错误。
+
+5. **cuBLAS handle 开销**：每次 matmul 创建/销毁 handle 有~0.1ms 开销。生产环境应全局复用一个 handle。Phase 15（性能优化）时需要修复这个问题。
+
+6. **`error::check` 需要 pub**：Phase 1 中 `check()` 是 `pub(crate)`，Phase 3 需要跨 crate 调用。反思：基础设施 crate 的错误处理函数应该从一开始就设计为 public API。
+
+## 后续优化方向（Phase 15）
+
+- Register tiling（每个 thread 算多个输出元素）
+- Tensor Core WMMA（利用 5090 的硬件加速）
+- CublasContext 全局复用
+- 非 contiguous input 支持（避免 matmul 前的拷贝）
--- a/docs/04-transformer-kernels.md
+++ b/docs/04-transformer-kernels.md
@@ -0,0 +1,213 @@
+# Phase 4: Transformer Core Kernels — Design Document
+
+## Goal
+
+实现 Transformer 所需的所有非 Attention 算子的 CUDA kernel，每个 kernel 都支持 BF16 和 F32，与 PyTorch 参考实现对比验证。
+
+## Kernel 清单
+
+| Kernel | 用于 | 核心计算 | 关键优化点 |
+|--------|------|---------|-----------|
+| LayerNorm | GPT-2 | `(x - mean) / sqrt(var + eps) * gamma + beta` | Welford online, warp reduce |
+| RMSNorm | Qwen3 | `x / sqrt(mean(x²) + eps) * gamma` | 无 mean，比 LayerNorm 简单 |
+| GELU | GPT-2 | `0.5x(1 + tanh(sqrt(2/π)(x + 0.044715x³)))` | tanh 近似，逐元素 |
+| SiLU | Qwen3 | `x * sigmoid(x)` | 逐元素 |
+| Softmax | Attention | `exp(x - max) / sum(exp(x - max))` | Online safe softmax, warp reduce |
+| Embedding | 全部 | `output[i] = table[token_ids[i]]` | Gather, coalesced write |
+| RoPE | Qwen3 | 对 Q/K 的相邻元素对做旋转 | Precompute freq, in-place |
+
+## 文件布局
+
+```
+csrc/
+├── normalization/
+│   ├── layernorm.cu
+│   └── rmsnorm.cu
+├── activation/
+│   ├── gelu.cu
+│   └── silu.cu
+├── reduce/
+│   └── softmax.cu
+├── embedding/
+│   ├── embedding.cu
+│   └── rope.cu
+
+crates/xserv-kernels/src/
+├── layernorm.rs
+├── rmsnorm.rs
+├── activation.rs      # GELU + SiLU
+├── softmax.rs
+├── embedding.rs
+├── rope.rs
+└── lib.rs             # 新增 mod 声明
+```
+
+## Kernel 设计细节
+
+### LayerNorm
+
+输入 `x: [*, hidden_size]`, 输出 `y: [*, hidden_size]`
+参数 `gamma, beta: [hidden_size]`
+
+```
+y[i] = gamma[i] * (x[i] - mean) / sqrt(var + eps) + beta[i]
+```
+
+**GPU 映射**: 每个 thread block 处理一行（一个 hidden_size 向量）。
+- Phase 1: 并行加载 x，Welford online 算法计算 mean 和 var
+- Phase 2: warp-level reduce (`__shfl_down_sync`) 聚合 mean/var
+- Phase 3: block-level reduce via shared memory
+- Phase 4: 每个 thread 对自己负责的元素做 normalize + affine
+
+**Block 配置**: `block = min(1024, hidden_size)`, `grid = num_rows`
+
+### RMSNorm
+
+比 LayerNorm 简单：不减 mean，只做 `x * rsqrt(mean(x²) + eps) * gamma`。
+
+```
+rms = sqrt(sum(x²) / hidden_size + eps)
+y[i] = x[i] / rms * gamma[i]
+```
+
+**GPU 映射**: 同 LayerNorm，每个 block 处理一行。
+- 只需要一次 reduce（求 sum(x²)），不需要两次（mean + var）。
+
+### GELU
+
+逐元素操作，用 tanh 近似:
+```
+gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))
+```
+
+**GPU 映射**: 每个 thread 处理多个元素（向量化），grid 覆盖全部元素。
+
+### SiLU (Swish)
+
+逐元素: `silu(x) = x * sigmoid(x) = x / (1 + exp(-x))`
+
+### Softmax
+
+输入 `x: [*, seq_len]`, 沿最后一维做 softmax:
+```
+1. m = max(x)                    // 数值稳定
+2. e[i] = exp(x[i] - m)
+3. s = sum(e)
+4. y[i] = e[i] / s
+```
+
+**GPU 映射**: 每个 block 处理一行。
+- 第一遍 reduce: 求 max
+- 第二遍: exp(x - max) 并 reduce sum
+- 第三遍: 除以 sum
+
+**优化**: 可以用 online softmax 合并前两遍（边算 exp 边更新 max），但先实现三遍版本保证正确。
+
+### Embedding
+
+```
+output[seq_idx] = embedding_table[token_ids[seq_idx]]
+```
+
+**GPU 映射**: 每个 thread 处理一个 token 的部分维度。
+- `grid = num_tokens`, `block = hidden_size`（或分多个 thread 处理一个 token）
+- 写端是 coalesced（连续 thread 写连续地址），读端是 gather（非连续）
+
+### RoPE (Rotary Position Embedding)
+
+对 Q/K 的每对相邻元素 `(x0, x1)` 做 2D 旋转:
+```
+freq[i] = 1.0 / (theta ^ (2i / dim))
+cos_val = cos(position * freq[i])
+sin_val = sin(position * freq[i])
+y0 = x0 * cos_val - x1 * sin_val
+y1 = x0 * sin_val + x1 * cos_val
+```
+
+**GPU 映射**: 每个 thread 处理一对元素 `(x[2i], x[2i+1])`。
+- Precompute `cos_cache[max_seq_len][head_dim/2]` 和 `sin_cache` 在初始化时
+- 运行时 kernel 只做乘加
+
+**theta**: Qwen3 默认 `rope_theta = 1000000.0`
+
+## Reduction Pattern（核心学习点）
+
+所有 Norm 和 Softmax 都涉及 reduction。GPU reduction 的分层结构:
+
+```
+Thread-level:  每个 thread 处理多个元素，本地累加
+    ↓
+Warp-level:    __shfl_down_sync() 在 32 threads 内规约（无需 shared memory）
+    ↓
+Block-level:   shared memory 存各 warp 的结果，warp 0 再规约
+```
+
+对于 hidden_size <= 8192（LLM 常见），一个 block 足够，不需要 grid-level reduction。
+
+### Warp Reduce 模板
+
+```cuda
+__device__ float warp_reduce_sum(float val) {
+    for (int offset = 16; offset > 0; offset >>= 1)
+        val += __shfl_down_sync(0xffffffff, val, offset);
+    return val;
+}
+```
+
+### Block Reduce 模板
+
+```cuda
+__device__ float block_reduce_sum(float val) {
+    __shared__ float shared[32];  // max 32 warps per block
+    int lane = threadIdx.x % 32;
+    int warp_id = threadIdx.x / 32;
+
+    val = warp_reduce_sum(val);
+    if (lane == 0) shared[warp_id] = val;
+    __syncthreads();
+
+    val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;
+    if (warp_id == 0) val = warp_reduce_sum(val);
+    return val;
+}
+```
+
+## Reference 验证策略
+
+写 `tools/generate_reference.py` 脚本，用 PyTorch 为每个 op 生成 reference input/output:
+- 保存为 `.npy` 格式
+- Rust 测试中加载对比
+- 或者直接在 Rust 测试中用 CPU 实现计算 expected 值（更简单，不依赖 Python）
+
+**选择**: 先用 Rust CPU 实现作为 reference（简单），关键 op（RoPE）再与 PyTorch 对比。
+
+## Test Plan
+
+- [x] RMSNorm F32: hidden_size=768, 4 rows → max_err 7.2e-7
+- [x] RMSNorm BF16: 同上 → max_err 7.0e-3
+- [x] LayerNorm F32: hidden_size=768 → max_err 1.7e-6
+- [x] GELU F32: 10000 elements → max_err 3.0e-8
+- [x] GELU BF16: 同上 → max_err 2.4e-3
+- [x] SiLU F32: 10000 elements → max_err 1.5e-8
+- [x] Softmax F32: 8×256 → max_err 1.4e-9
+- [x] Softmax sum=1 验证: 4×2048
+- [x] Softmax 大值 (1000+) 数值稳定性 → max_err 1.5e-8
+- [x] Embedding F32: vocab=100, hidden=64, 5 tokens → exact match
+- [x] RoPE F32: 4 tokens × 2 heads × dim=8 → max_err 6.0e-8
+- [x] RoPE position=0 恒等验证 → max_err 0
+
+## Takeaways
+
+1. **`common.cuh` 抽取共用 reduction 是正确的做法**：`warp_reduce_sum/max` 和 `block_reduce_sum/max` 被 RMSNorm, LayerNorm, Softmax 三个 kernel 复用。抽到头文件避免了代码重复，也确保 reduction 逻辑一致。build.rs 中需要 `.include("../../csrc")` 让 nvcc 能找到头文件。
+
+2. **Shared memory 中广播标量的模式**：Norm 和 Softmax 都需要将 reduce 结果（mean, rms_inv, max, sum）广播给 block 内所有 thread。标准做法：thread 0 写 `__shared__` 变量，`__syncthreads()` 后所有 thread 读。这比让每个 thread 独立做 reduce 高效得多。
+
+3. **Softmax 三遍 vs 两遍**：我们实现了三遍版本（max → exp+sum → normalize），简单可靠。Online softmax 可以合并前两遍（一遍 pass 内同时跟踪 running max 和 running sum），但需要更复杂的数值更新公式。Flash Attention（Phase 14）会用到 online softmax。
+
+4. **RoPE 的 position=0 恒等性**：`cos(0)=1, sin(0)=0`，所以 position 0 的旋转是恒等变换。这是一个很好的 sanity check。如果 position=0 时输出不等于输入，说明 kernel 有 bug。
+
+5. **BF16 Softmax 的精度陷阱**：exp 结果先写成 BF16 再读回做 normalize 会丢精度。理想做法是用 float scratch buffer 暂存 exp 结果。当前实现可接受（误差在 1e-2 量级），但在 attention score 很接近时可能引入可观察的差异。Phase 14 Flash Attention 会解决这个问题（全程 FP32 累加）。
+
+6. **Embedding 就是 gather 操作**：没有任何计算，纯粹的内存搬运。瓶颈在 global memory 随机读取（token_ids 导致不连续读 table）。写端是 coalesced 的（连续 token 写连续地址）。优化方向：使用向量化加载（`float4`）一次读 128 bit。
+
+7. **RoPE in-place 修改 Tensor 的设计考量**：RoPE 在数学上是对 Q/K 的 in-place 旋转。我们通过 `data_ptr() as *mut` 绕过了 Rust 的不可变借用。这在 GPU 上是安全的（kernel 内部互不干扰），但 Rust 侧没有 `&mut` 语义保护。后续如果需要更严格的安全性，可以引入 `Tensor::as_mut_ptr()` 方法并要求 `&mut self`。