From 5aef3742d645dc187f39f22ab9f8a26bfad49324 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Mon, 15 Jun 2026 15:44:09 +0800 Subject: [PATCH] ops: transformer op fwd/bwd CUDA kernels + Tensor wrappers add/mul/add_bias(+sum_rows)/rms_norm/silu/rope/softmax/cross_entropy, each with its analytic backward, in csrc/ops/nn.cu (inlined warp/block reductions). FFI declarations + nn.cu in build.rs (no_cuda gated). Tensor gains the matching thin wrappers; DType grows I32 for cross-entropy targets. Co-Authored-By: Claude Opus 4.8 --- crates/xtrain-cuda/build.rs | 1 + crates/xtrain-cuda/src/ffi.rs | 114 +++++++++ crates/xtrain-tensor/src/dtype.rs | 14 ++ crates/xtrain-tensor/src/tensor.rs | 328 ++++++++++++++++++++++++++ csrc/ops/nn.cu | 358 +++++++++++++++++++++++++++++ 5 files changed, 815 insertions(+) create mode 100644 csrc/ops/nn.cu diff --git a/crates/xtrain-cuda/build.rs b/crates/xtrain-cuda/build.rs index 9d29fa3..aa0b99b 100644 --- a/crates/xtrain-cuda/build.rs +++ b/crates/xtrain-cuda/build.rs @@ -32,6 +32,7 @@ fn main() { .file("../../csrc/test/vecadd.cu") .file("../../csrc/ops/elementwise.cu") .file("../../csrc/ops/gemm.cu") + .file("../../csrc/ops/nn.cu") .compile("xtrain_cuda_kernels"); } diff --git a/crates/xtrain-cuda/src/ffi.rs b/crates/xtrain-cuda/src/ffi.rs index a13a3e3..179d299 100644 --- a/crates/xtrain-cuda/src/ffi.rs +++ b/crates/xtrain-cuda/src/ffi.rs @@ -63,6 +63,120 @@ unsafe extern "C" { ); } +// Transformer / autograd op kernels (csrc/ops/nn.cu). Forward + backward for the +// ops the Phase T4 tape engine needs. All F32, row-major, contiguous. +#[cfg(not(no_cuda))] +unsafe extern "C" { + // Elementwise: out = a + b ; out = a * b. + pub fn launch_add_f32(a: *const f32, b: *const f32, out: *mut f32, n: i32, s: CudaStream); + pub fn launch_mul_f32(a: *const f32, b: *const f32, out: *mut f32, n: i32, s: CudaStream); + + // Broadcast bias add: out[r,c] = x[r,c] + bias[c]. x:[rows,cols], bias:[cols]. + pub fn launch_add_bias_f32( + x: *const f32, + bias: *const f32, + out: *mut f32, + rows: i32, + cols: i32, + s: CudaStream, + ); + // Column-sum (over rows): dbias[c] = sum_r dout[r,c]. Bias backward. + pub fn launch_sum_rows_f32( + dout: *const f32, + dbias: *mut f32, + rows: i32, + cols: i32, + s: CudaStream, + ); + + // RMSNorm forward: writes y[rows,cols] and inv_rms[rows] (cached for bwd). + pub fn launch_rms_norm_f32( + x: *const f32, + gamma: *const f32, + y: *mut f32, + inv_rms: *mut f32, + rows: i32, + cols: i32, + eps: f32, + s: CudaStream, + ); + pub fn launch_rms_norm_dx_f32( + x: *const f32, + gamma: *const f32, + dy: *const f32, + inv_rms: *const f32, + dx: *mut f32, + rows: i32, + cols: i32, + s: CudaStream, + ); + pub fn launch_rms_norm_dgamma_f32( + x: *const f32, + dy: *const f32, + inv_rms: *const f32, + dgamma: *mut f32, + rows: i32, + cols: i32, + s: CudaStream, + ); + + // SiLU: y = x*sigmoid(x); backward dx. + pub fn launch_silu_f32(x: *const f32, y: *mut f32, n: i32, s: CudaStream); + pub fn launch_silu_dx_f32(x: *const f32, dy: *const f32, dx: *mut f32, n: i32, s: CudaStream); + + // RoPE (rotate_half), x:[tokens,heads,head_dim], position = token index. + pub fn launch_rope_f32( + x: *const f32, + y: *mut f32, + tokens: i32, + heads: i32, + head_dim: i32, + theta: f32, + s: CudaStream, + ); + pub fn launch_rope_dx_f32( + dy: *const f32, + dx: *mut f32, + tokens: i32, + heads: i32, + head_dim: i32, + theta: f32, + s: CudaStream, + ); + + // Row-wise softmax + Jacobian backward. + pub fn launch_softmax_f32(x: *const f32, y: *mut f32, rows: i32, cols: i32, s: CudaStream); + pub fn launch_softmax_dx_f32( + y: *const f32, + dy: *const f32, + dx: *mut f32, + rows: i32, + cols: i32, + s: CudaStream, + ); + + // Cross-entropy: fwd writes probs[rows,cols] + per-row loss[rows]; + // bwd dx = scale*(probs - onehot). + pub fn launch_cross_entropy_fwd_f32( + x: *const f32, + target: *const i32, + probs: *mut f32, + loss: *mut f32, + rows: i32, + cols: i32, + s: CudaStream, + ); + pub fn launch_cross_entropy_dx_f32( + probs: *const f32, + target: *const i32, + dx: *mut f32, + rows: i32, + cols: i32, + scale: f32, + s: CudaStream, + ); +} + // cuBLAS — used ONLY as a correctness reference for the hand-written GEMM in // tests. Declared (and linked, see build.rs) only when CUDA is compiled in. #[cfg(not(no_cuda))] diff --git a/crates/xtrain-tensor/src/dtype.rs b/crates/xtrain-tensor/src/dtype.rs index fece395..f373d43 100644 --- a/crates/xtrain-tensor/src/dtype.rs +++ b/crates/xtrain-tensor/src/dtype.rs @@ -7,18 +7,22 @@ #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum DType { F32, + /// 32-bit signed integers. Used for cross-entropy targets (token ids). + I32, } impl DType { pub fn size_bytes(self) -> usize { match self { DType::F32 => 4, + DType::I32 => 4, } } pub fn name(self) -> &'static str { match self { DType::F32 => "f32", + DType::I32 => "i32", } } } @@ -45,3 +49,13 @@ impl TensorDType for f32 { v as f32 } } + +impl TensorDType for i32 { + const DTYPE: DType = DType::I32; + fn to_f64(self) -> f64 { + self as f64 + } + fn from_f64(v: f64) -> Self { + v as i32 + } +} diff --git a/crates/xtrain-tensor/src/tensor.rs b/crates/xtrain-tensor/src/tensor.rs index 6cbb8d7..db4fcaa 100644 --- a/crates/xtrain-tensor/src/tensor.rs +++ b/crates/xtrain-tensor/src/tensor.rs @@ -247,6 +247,334 @@ impl Tensor { let db = a.transpose_2d().matmul(dc); // [K,M] @ [M,N] = [K,N] (da, db) } + + // --- Transformer / autograd op primitives (the T4 kernels) --- + // + // Each is a thin, contiguous-F32-on-GPU wrapper over a kernel in + // csrc/ops/nn.cu. The autograd `Var` layer (xtrain-autodiff) builds nodes on + // top of these; the analytic backwards are derived in docs/03-autograd-engine.md. + + /// Elementwise `out = self + other` (same shape). + #[cfg(not(no_cuda))] + pub fn add(&self, other: &Tensor) -> Self { + self.check_binary(other, "add"); + let out = Tensor::zeros(&self.shape, DType::F32, self.device()); + unsafe { + xtrain_cuda::ffi::launch_add_f32( + self.data_ptr() as *const f32, + other.data_ptr() as *const f32, + out.data_ptr() as *mut f32, + self.numel() as i32, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("add sync failed"); + out + } + + /// Elementwise `out = self * other` (same shape, Hadamard product). + #[cfg(not(no_cuda))] + pub fn mul(&self, other: &Tensor) -> Self { + self.check_binary(other, "mul"); + let out = Tensor::zeros(&self.shape, DType::F32, self.device()); + unsafe { + xtrain_cuda::ffi::launch_mul_f32( + self.data_ptr() as *const f32, + other.data_ptr() as *const f32, + out.data_ptr() as *mut f32, + self.numel() as i32, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("mul sync failed"); + out + } + + /// Broadcast bias add: `out[r,c] = self[r,c] + bias[c]`. + /// `self`:[rows,cols], `bias`:[cols]. + #[cfg(not(no_cuda))] + pub fn add_bias(&self, bias: &Tensor) -> Self { + assert_eq!(self.ndim(), 2, "add_bias requires 2D input"); + assert_eq!(bias.ndim(), 1, "bias must be 1D"); + assert_eq!(self.shape[1], bias.shape[0], "bias len != cols"); + let (rows, cols) = (self.shape[0], self.shape[1]); + let out = Tensor::zeros(&self.shape, DType::F32, self.device()); + unsafe { + xtrain_cuda::ffi::launch_add_bias_f32( + self.data_ptr() as *const f32, + bias.data_ptr() as *const f32, + out.data_ptr() as *mut f32, + rows as i32, + cols as i32, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("add_bias sync failed"); + out + } + + /// Column-sum over rows: `out[c] = sum_r self[r,c]`. This is the bias + /// backward (sum the upstream grad over the broadcast dim). `self`:[rows,cols] + /// → [cols]. + #[cfg(not(no_cuda))] + pub fn sum_rows(&self) -> Self { + assert_eq!(self.ndim(), 2, "sum_rows requires 2D input"); + let (rows, cols) = (self.shape[0], self.shape[1]); + let out = Tensor::zeros(&[cols], DType::F32, self.device()); + unsafe { + xtrain_cuda::ffi::launch_sum_rows_f32( + self.data_ptr() as *const f32, + out.data_ptr() as *mut f32, + rows as i32, + cols as i32, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("sum_rows sync failed"); + out + } + + /// RMSNorm forward: `y[r,c] = x[r,c] * inv_rms[r] * gamma[c]` with + /// `inv_rms = rsqrt(mean(x²) + eps)`. `self`:[rows,cols], `gamma`:[cols]. + /// Returns `(y, inv_rms)`; `inv_rms`:[rows] is cached for backward. + #[cfg(not(no_cuda))] + pub fn rms_norm(&self, gamma: &Tensor, eps: f32) -> (Tensor, Tensor) { + assert_eq!(self.ndim(), 2, "rms_norm requires 2D input"); + assert_eq!(gamma.ndim(), 1, "gamma must be 1D"); + assert_eq!(self.shape[1], gamma.shape[0], "gamma len != cols"); + let (rows, cols) = (self.shape[0], self.shape[1]); + let y = Tensor::zeros(&self.shape, DType::F32, self.device()); + let inv_rms = Tensor::zeros(&[rows], DType::F32, self.device()); + unsafe { + xtrain_cuda::ffi::launch_rms_norm_f32( + self.data_ptr() as *const f32, + gamma.data_ptr() as *const f32, + y.data_ptr() as *mut f32, + inv_rms.data_ptr() as *mut f32, + rows as i32, + cols as i32, + eps, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("rms_norm sync failed"); + (y, inv_rms) + } + + /// RMSNorm backward. Inputs are the forward `x`, `gamma`, upstream `dy`, and + /// the cached `inv_rms`. Returns `(dx, dgamma)`. + #[cfg(not(no_cuda))] + pub fn rms_norm_backward( + x: &Tensor, + gamma: &Tensor, + dy: &Tensor, + inv_rms: &Tensor, + ) -> (Tensor, Tensor) { + let (rows, cols) = (x.shape[0], x.shape[1]); + let dx = Tensor::zeros(&[rows, cols], DType::F32, x.device()); + let dgamma = Tensor::zeros(&[cols], DType::F32, x.device()); + unsafe { + xtrain_cuda::ffi::launch_rms_norm_dx_f32( + x.data_ptr() as *const f32, + gamma.data_ptr() as *const f32, + dy.data_ptr() as *const f32, + inv_rms.data_ptr() as *const f32, + dx.data_ptr() as *mut f32, + rows as i32, + cols as i32, + std::ptr::null_mut(), + ); + xtrain_cuda::ffi::launch_rms_norm_dgamma_f32( + x.data_ptr() as *const f32, + dy.data_ptr() as *const f32, + inv_rms.data_ptr() as *const f32, + dgamma.data_ptr() as *mut f32, + rows as i32, + cols as i32, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("rms_norm_backward sync failed"); + (dx, dgamma) + } + + /// SiLU forward: `y = x * sigmoid(x)`, elementwise. + #[cfg(not(no_cuda))] + pub fn silu(&self) -> Self { + assert_eq!(self.dtype, DType::F32, "silu only supports F32"); + let out = Tensor::zeros(&self.shape, DType::F32, self.device()); + unsafe { + xtrain_cuda::ffi::launch_silu_f32( + self.data_ptr() as *const f32, + out.data_ptr() as *mut f32, + self.numel() as i32, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("silu sync failed"); + out + } + + /// SiLU backward: `dx = dy * (sig + x*sig*(1-sig))`, `sig = sigmoid(x)`. + /// Inputs are the forward `x` and upstream `dy`. + #[cfg(not(no_cuda))] + pub fn silu_backward(x: &Tensor, dy: &Tensor) -> Self { + let dx = Tensor::zeros(&x.shape, DType::F32, x.device()); + unsafe { + xtrain_cuda::ffi::launch_silu_dx_f32( + x.data_ptr() as *const f32, + dy.data_ptr() as *const f32, + dx.data_ptr() as *mut f32, + x.numel() as i32, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("silu_backward sync failed"); + dx + } + + /// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; the position + /// of each token is its row index. Returns the rotated tensor. + #[cfg(not(no_cuda))] + pub fn rope(&self, theta: f32) -> Self { + assert_eq!(self.ndim(), 3, "rope requires [tokens,heads,head_dim]"); + let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]); + assert_eq!(head_dim % 2, 0, "head_dim must be even"); + let out = Tensor::zeros(&self.shape, DType::F32, self.device()); + unsafe { + xtrain_cuda::ffi::launch_rope_f32( + self.data_ptr() as *const f32, + out.data_ptr() as *mut f32, + tokens as i32, + heads as i32, + head_dim as i32, + theta, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("rope sync failed"); + out + } + + /// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an + /// orthogonal map, so it needs no cached forward values, only `theta`. + #[cfg(not(no_cuda))] + pub fn rope_backward(dy: &Tensor, theta: f32) -> Self { + let (tokens, heads, head_dim) = (dy.shape[0], dy.shape[1], dy.shape[2]); + let dx = Tensor::zeros(&dy.shape, DType::F32, dy.device()); + unsafe { + xtrain_cuda::ffi::launch_rope_dx_f32( + dy.data_ptr() as *const f32, + dx.data_ptr() as *mut f32, + tokens as i32, + heads as i32, + head_dim as i32, + theta, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("rope_backward sync failed"); + dx + } + + /// Row-wise safe softmax over the last dim. `self`:[rows,cols]. + #[cfg(not(no_cuda))] + pub fn softmax(&self) -> Self { + assert_eq!(self.ndim(), 2, "softmax requires 2D input"); + let (rows, cols) = (self.shape[0], self.shape[1]); + let out = Tensor::zeros(&self.shape, DType::F32, self.device()); + unsafe { + xtrain_cuda::ffi::launch_softmax_f32( + self.data_ptr() as *const f32, + out.data_ptr() as *mut f32, + rows as i32, + cols as i32, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("softmax sync failed"); + out + } + + /// Softmax backward (Jacobian): `dx[r,c] = y[r,c]*(dy[r,c] - sum_c'(dy*y))`. + /// Inputs are the forward output `y` and upstream `dy`. + #[cfg(not(no_cuda))] + pub fn softmax_backward(y: &Tensor, dy: &Tensor) -> Self { + let (rows, cols) = (y.shape[0], y.shape[1]); + let dx = Tensor::zeros(&y.shape, DType::F32, y.device()); + unsafe { + xtrain_cuda::ffi::launch_softmax_dx_f32( + y.data_ptr() as *const f32, + dy.data_ptr() as *const f32, + dx.data_ptr() as *mut f32, + rows as i32, + cols as i32, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("softmax_backward sync failed"); + dx + } + + /// Cross-entropy forward over logits `self`:[rows,cols] with one I32 target + /// per row. Returns `(probs, loss)` where `probs`:[rows,cols] is the softmax + /// (cached for backward) and `loss`:[rows] is the per-row negative log-likelihood. + #[cfg(not(no_cuda))] + pub fn cross_entropy(&self, target: &Tensor) -> (Tensor, Tensor) { + assert_eq!(self.ndim(), 2, "cross_entropy requires 2D logits"); + assert_eq!(target.dtype, DType::I32, "target must be I32"); + assert_eq!(target.numel(), self.shape[0], "one target per row"); + let (rows, cols) = (self.shape[0], self.shape[1]); + let probs = Tensor::zeros(&self.shape, DType::F32, self.device()); + let loss = Tensor::zeros(&[rows], DType::F32, self.device()); + unsafe { + xtrain_cuda::ffi::launch_cross_entropy_fwd_f32( + self.data_ptr() as *const f32, + target.data_ptr() as *const i32, + probs.data_ptr() as *mut f32, + loss.data_ptr() as *mut f32, + rows as i32, + cols as i32, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("cross_entropy sync failed"); + (probs, loss) + } + + /// Cross-entropy backward: `dx = scale * (probs - onehot(target))`. With + /// `scale = upstream / rows`, this is the gradient of the mean per-row loss. + #[cfg(not(no_cuda))] + pub fn cross_entropy_backward(probs: &Tensor, target: &Tensor, scale: f32) -> Self { + let (rows, cols) = (probs.shape[0], probs.shape[1]); + let dx = Tensor::zeros(&probs.shape, DType::F32, probs.device()); + unsafe { + xtrain_cuda::ffi::launch_cross_entropy_dx_f32( + probs.data_ptr() as *const f32, + target.data_ptr() as *const i32, + dx.data_ptr() as *mut f32, + rows as i32, + cols as i32, + scale, + std::ptr::null_mut(), + ); + } + xtrain_cuda::device::synchronize().expect("cross_entropy_backward sync failed"); + dx + } + + // Shared validation for same-shape binary elementwise ops. + #[cfg(not(no_cuda))] + fn check_binary(&self, other: &Tensor, op: &str) { + assert_eq!(self.dtype, DType::F32, "{op} only supports F32"); + assert_eq!(other.dtype, DType::F32, "{op} only supports F32"); + assert_eq!(self.shape(), other.shape(), "{op} shape mismatch"); + assert_eq!(self.device(), other.device(), "{op} device mismatch"); + assert!( + self.is_contiguous() && other.is_contiguous(), + "{op} requires contiguous tensors" + ); + } } impl std::fmt::Debug for Tensor { diff --git a/csrc/ops/nn.cu b/csrc/ops/nn.cu new file mode 100644 index 0000000..e8dd32a --- /dev/null +++ b/csrc/ops/nn.cu @@ -0,0 +1,358 @@ +// Forward + backward CUDA kernels for the transformer ops the autograd engine +// (Phase T4) needs: elementwise add/mul, broadcast bias add + its row-sum +// backward, RMSNorm, SiLU, RoPE, row-wise softmax, and cross-entropy. +// +// All F32, row-major, contiguous. Forward kernels mirror xserv +// (docs/04-transformer-kernels.md, docs/05-attention.md); the backward kernels +// are new (xserv is inference-only). Reduction helpers are inlined here so this +// file is self-contained (no shared header), matching the existing csrc/ layout. + +#include + +extern "C" { + +// --- Warp / block reductions (sum + max), block handles one row --- + +__device__ __forceinline__ float warp_reduce_sum(float v) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + v += __shfl_down_sync(0xffffffff, v, off); + return v; +} + +__device__ __forceinline__ float warp_reduce_max(float v) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + v = fmaxf(v, __shfl_down_sync(0xffffffff, v, off)); + return v; +} + +__device__ __forceinline__ float block_reduce_sum(float v) { + __shared__ float shared[32]; + int lane = threadIdx.x & 31; + int warp = threadIdx.x >> 5; + int nwarps = (blockDim.x + 31) >> 5; + v = warp_reduce_sum(v); + if (lane == 0) shared[warp] = v; + __syncthreads(); + v = (threadIdx.x < nwarps) ? shared[threadIdx.x] : 0.0f; + if (warp == 0) v = warp_reduce_sum(v); + // broadcast warp-0 lane-0 result to whole block + __shared__ float bcast; + if (threadIdx.x == 0) bcast = v; + __syncthreads(); + return bcast; +} + +__device__ __forceinline__ float block_reduce_max(float v) { + __shared__ float shared[32]; + int lane = threadIdx.x & 31; + int warp = threadIdx.x >> 5; + int nwarps = (blockDim.x + 31) >> 5; + v = warp_reduce_max(v); + if (lane == 0) shared[warp] = v; + __syncthreads(); + v = (threadIdx.x < nwarps) ? shared[threadIdx.x] : -INFINITY; + if (warp == 0) v = warp_reduce_max(v); + __shared__ float bcast; + if (threadIdx.x == 0) bcast = v; + __syncthreads(); + return bcast; +} + +// ===================================================================== +// Elementwise add / mul (same-shape) +// ===================================================================== + +__global__ void add_k(const float* a, const float* b, float* out, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) out[i] = a[i] + b[i]; +} +void launch_add_f32(const float* a, const float* b, float* out, int n, void* s) { + int blk = 256, grid = (n + blk - 1) / blk; + add_k<<>>(a, b, out, n); +} + +__global__ void mul_k(const float* a, const float* b, float* out, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) out[i] = a[i] * b[i]; +} +void launch_mul_f32(const float* a, const float* b, float* out, int n, void* s) { + int blk = 256, grid = (n + blk - 1) / blk; + mul_k<<>>(a, b, out, n); +} + +// ===================================================================== +// Broadcast bias add: out[r,c] = x[r,c] + bias[c] (x:[rows,cols]) +// Backward for bias is a column-sum (sum over rows): dbias[c] = sum_r dout[r,c]. +// ===================================================================== + +__global__ void add_bias_k(const float* x, const float* bias, float* out, + int rows, int cols) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < rows * cols) out[i] = x[i] + bias[i % cols]; +} +void launch_add_bias_f32(const float* x, const float* bias, float* out, + int rows, int cols, void* s) { + int n = rows * cols, blk = 256, grid = (n + blk - 1) / blk; + add_bias_k<<>>(x, bias, out, rows, cols); +} + +// dbias[c] = sum_r dout[r,c]. One block per column, threads stride over rows. +__global__ void sum_rows_k(const float* dout, float* dbias, int rows, int cols) { + int col = blockIdx.x; + float acc = 0.0f; + for (int r = threadIdx.x; r < rows; r += blockDim.x) + acc += dout[r * cols + col]; + acc = block_reduce_sum(acc); + if (threadIdx.x == 0) dbias[col] = acc; +} +void launch_sum_rows_f32(const float* dout, float* dbias, int rows, int cols, void* s) { + int blk = 256; + sum_rows_k<<>>(dout, dbias, rows, cols); +} + +// ===================================================================== +// RMSNorm: y[r,c] = x[r,c] * inv_rms[r] * gamma[c], inv_rms = rsqrt(mean(x²)+eps) +// x:[rows,cols], gamma:[cols]. Forward also writes inv_rms[rows] for backward. +// ===================================================================== + +__global__ void rms_norm_k(const float* x, const float* gamma, float* y, + float* inv_rms, int rows, int cols, float eps) { + int r = blockIdx.x; + const float* xr = x + r * cols; + float* yr = y + r * cols; + float ss = 0.0f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) ss += xr[c] * xr[c]; + ss = block_reduce_sum(ss); + float ir = rsqrtf(ss / cols + eps); + if (threadIdx.x == 0) inv_rms[r] = ir; + for (int c = threadIdx.x; c < cols; c += blockDim.x) + yr[c] = xr[c] * ir * gamma[c]; +} +void launch_rms_norm_f32(const float* x, const float* gamma, float* y, + float* inv_rms, int rows, int cols, float eps, void* s) { + int blk = cols < 1024 ? cols : 1024; + if (blk < 32) blk = 32; + rms_norm_k<<>>(x, gamma, y, inv_rms, rows, cols, eps); +} + +// RMSNorm backward. +// Let g[c] = dy[r,c]*gamma[c], ir = inv_rms[r], n = cols. +// dx[r,c] = ir*g[c] - x[r,c]*ir³/n * sum_c(g[c]*x[r,c]) +// dgamma[c] = sum_r dy[r,c] * x[r,c] * ir (accumulated across rows) +__global__ void rms_norm_dx_k(const float* x, const float* gamma, const float* dy, + const float* inv_rms, float* dx, int rows, int cols) { + int r = blockIdx.x; + const float* xr = x + r * cols; + const float* dyr = dy + r * cols; + float* dxr = dx + r * cols; + float ir = inv_rms[r]; + float dot = 0.0f; // sum_c g[c]*x[c] + for (int c = threadIdx.x; c < cols; c += blockDim.x) + dot += dyr[c] * gamma[c] * xr[c]; + dot = block_reduce_sum(dot); + float coeff = ir * ir * ir / (float)cols * dot; + for (int c = threadIdx.x; c < cols; c += blockDim.x) + dxr[c] = ir * dyr[c] * gamma[c] - xr[c] * coeff; +} +void launch_rms_norm_dx_f32(const float* x, const float* gamma, const float* dy, + const float* inv_rms, float* dx, int rows, int cols, void* s) { + int blk = cols < 1024 ? cols : 1024; + if (blk < 32) blk = 32; + rms_norm_dx_k<<>>(x, gamma, dy, inv_rms, dx, rows, cols); +} + +// dgamma[c] = sum_r dy[r,c] * x[r,c] * inv_rms[r]. One block per column. +__global__ void rms_norm_dgamma_k(const float* x, const float* dy, const float* inv_rms, + float* dgamma, int rows, int cols) { + int col = blockIdx.x; + float acc = 0.0f; + for (int r = threadIdx.x; r < rows; r += blockDim.x) + acc += dy[r * cols + col] * x[r * cols + col] * inv_rms[r]; + acc = block_reduce_sum(acc); + if (threadIdx.x == 0) dgamma[col] = acc; +} +void launch_rms_norm_dgamma_f32(const float* x, const float* dy, const float* inv_rms, + float* dgamma, int rows, int cols, void* s) { + int blk = 256; + rms_norm_dgamma_k<<>>(x, dy, inv_rms, dgamma, rows, cols); +} + +// ===================================================================== +// SiLU: y = x * sigmoid(x). Backward: dx = dy * (sig + x*sig*(1-sig)). +// ===================================================================== + +__global__ void silu_k(const float* x, float* y, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { float xv = x[i]; y[i] = xv / (1.0f + expf(-xv)); } +} +void launch_silu_f32(const float* x, float* y, int n, void* s) { + int blk = 256, grid = (n + blk - 1) / blk; + silu_k<<>>(x, y, n); +} + +__global__ void silu_dx_k(const float* x, const float* dy, float* dx, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + float xv = x[i]; + float sig = 1.0f / (1.0f + expf(-xv)); + dx[i] = dy[i] * (sig + xv * sig * (1.0f - sig)); + } +} +void launch_silu_dx_f32(const float* x, const float* dy, float* dx, int n, void* s) { + int blk = 256, grid = (n + blk - 1) / blk; + silu_dx_k<<>>(x, dy, dx, n); +} + +// ===================================================================== +// RoPE (rotate_half layout). x:[tokens, heads, head_dim]; position = token index. +// y[i] = x[i]*cos - x[i+h]*sin +// y[i+h] = x[i+h]*cos + x[i]*sin (i in [0,half), h=half_dim) +// freq[i] = theta^(-2i/head_dim); angle = pos*freq[i]. +// Backward is the inverse (transpose) rotation: apply +angle's transpose ≡ -angle. +// dx[i] = dy[i]*cos + dy[i+h]*sin +// dx[i+h] = dy[i+h]*cos - dy[i]*sin +// ===================================================================== + +__global__ void rope_k(const float* x, float* y, int heads, int head_dim, float theta) { + int tok = blockIdx.x; + int head = blockIdx.y; + int half = head_dim / 2; + int i = threadIdx.x; + if (i >= half) return; + float freq = powf(theta, -(float)(2 * i) / (float)head_dim); + float angle = (float)tok * freq; + float c = cosf(angle), sn = sinf(angle); + int base = (tok * heads + head) * head_dim; + float x0 = x[base + i], x1 = x[base + i + half]; + y[base + i] = x0 * c - x1 * sn; + y[base + i + half] = x1 * c + x0 * sn; +} +void launch_rope_f32(const float* x, float* y, int tokens, int heads, + int head_dim, float theta, void* s) { + dim3 grid(tokens, heads); + int blk = head_dim / 2; + rope_k<<>>(x, y, heads, head_dim, theta); +} + +__global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim, float theta) { + int tok = blockIdx.x; + int head = blockIdx.y; + int half = head_dim / 2; + int i = threadIdx.x; + if (i >= half) return; + float freq = powf(theta, -(float)(2 * i) / (float)head_dim); + float angle = (float)tok * freq; + float c = cosf(angle), sn = sinf(angle); + int base = (tok * heads + head) * head_dim; + float d0 = dy[base + i], d1 = dy[base + i + half]; + dx[base + i] = d0 * c + d1 * sn; + dx[base + i + half] = d1 * c - d0 * sn; +} +void launch_rope_dx_f32(const float* dy, float* dx, int tokens, int heads, + int head_dim, float theta, void* s) { + dim3 grid(tokens, heads); + int blk = head_dim / 2; + rope_dx_k<<>>(dy, dx, heads, head_dim, theta); +} + +// ===================================================================== +// Row-wise safe softmax. x:[rows,cols] → y. Backward (Jacobian): +// dx[r,c] = y[r,c] * (dy[r,c] - sum_c'(dy[r,c']*y[r,c'])) +// ===================================================================== + +__global__ void softmax_k(const float* x, float* y, int rows, int cols) { + int r = blockIdx.x; + const float* xr = x + r * cols; + float* yr = y + r * cols; + float m = -INFINITY; + for (int c = threadIdx.x; c < cols; c += blockDim.x) m = fmaxf(m, xr[c]); + m = block_reduce_max(m); + float sum = 0.0f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) { + float e = expf(xr[c] - m); + yr[c] = e; + sum += e; + } + sum = block_reduce_sum(sum); + float inv = 1.0f / sum; + for (int c = threadIdx.x; c < cols; c += blockDim.x) yr[c] *= inv; +} +void launch_softmax_f32(const float* x, float* y, int rows, int cols, void* s) { + int blk = cols < 1024 ? cols : 1024; + if (blk < 32) blk = 32; + softmax_k<<>>(x, y, rows, cols); +} + +__global__ void softmax_dx_k(const float* y, const float* dy, float* dx, + int rows, int cols) { + int r = blockIdx.x; + const float* yr = y + r * cols; + const float* dyr = dy + r * cols; + float* dxr = dx + r * cols; + float dot = 0.0f; // sum_c dy*y + for (int c = threadIdx.x; c < cols; c += blockDim.x) dot += dyr[c] * yr[c]; + dot = block_reduce_sum(dot); + for (int c = threadIdx.x; c < cols; c += blockDim.x) + dxr[c] = yr[c] * (dyr[c] - dot); +} +void launch_softmax_dx_f32(const float* y, const float* dy, float* dx, + int rows, int cols, void* s) { + int blk = cols < 1024 ? cols : 1024; + if (blk < 32) blk = 32; + softmax_dx_k<<>>(y, dy, dx, rows, cols); +} + +// ===================================================================== +// Cross-entropy over logits x:[rows,cols] with int target per row. +// Forward writes per-row loss[r] = -log(softmax(x)[target]) and the softmax +// probs[r,:] (cached for backward). Backward: dx[r,c] = (probs[r,c]-onehot)/rows +// (mean reduction; the *rows scale folds the 1/rows of mean loss into dx). +// ===================================================================== + +__global__ void cross_entropy_fwd_k(const float* x, const int* target, + float* probs, float* loss, int rows, int cols) { + int r = blockIdx.x; + const float* xr = x + r * cols; + float* pr = probs + r * cols; + float m = -INFINITY; + for (int c = threadIdx.x; c < cols; c += blockDim.x) m = fmaxf(m, xr[c]); + m = block_reduce_max(m); + float sum = 0.0f; + for (int c = threadIdx.x; c < cols; c += blockDim.x) { + float e = expf(xr[c] - m); + pr[c] = e; + sum += e; + } + sum = block_reduce_sum(sum); + float inv = 1.0f / sum; + for (int c = threadIdx.x; c < cols; c += blockDim.x) pr[c] *= inv; + if (threadIdx.x == 0) { + int t = target[r]; + loss[r] = -logf(pr[t]); + } +} +void launch_cross_entropy_fwd_f32(const float* x, const int* target, + float* probs, float* loss, int rows, int cols, void* s) { + int blk = cols < 1024 ? cols : 1024; + if (blk < 32) blk = 32; + cross_entropy_fwd_k<<>>(x, target, probs, loss, rows, cols); +} + +// dx[r,c] = scale * (probs[r,c] - [c==target]). scale = upstream/rows. +__global__ void cross_entropy_dx_k(const float* probs, const int* target, + float* dx, int rows, int cols, float scale) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= rows * cols) return; + int r = i / cols, c = i % cols; + float g = probs[i] - (c == target[r] ? 1.0f : 0.0f); + dx[i] = g * scale; +} +void launch_cross_entropy_dx_f32(const float* probs, const int* target, + float* dx, int rows, int cols, float scale, void* s) { + int n = rows * cols, blk = 256, grid = (n + blk - 1) / blk; + cross_entropy_dx_k<<>>(probs, target, dx, rows, cols, scale); +} + +} // extern "C"