entrenar/autograd/cuda_forward/
mod.rs

1//! CUDA-accelerated forward kernels for autograd
2//!
3//! This module wraps trueno-gpu forward kernels for GPU-accelerated forward passes.
4//! Provides 10-100x speedup over CPU ndarray implementations.
5//!
6//! # Architecture (SPEC-FT-001 v3.0.0)
7//!
8//! ```text
9//! entrenar autograd
10//!     └── cuda_forward (this module)
11//!             └── trueno-gpu/kernels
12//!                     └── PTX generation + CUDA driver
13//! ```
14//!
15//! # Available Kernels
16//!
17//! - `relu_forward` - ReLU activation
18//! - `softmax_forward` - Numerically stable softmax with warp shuffle
19//! - `batched_softmax_forward` - Row-wise batched softmax for attention
20//! - `layer_norm_forward` - Fused layer normalization
21//! - `rms_norm_forward` - RMS normalization (LLaMA-style)
22//! - `gelu_forward` - GELU activation
23//! - `silu_forward` - SiLU/Swish activation
24//! - `gemm_forward` - Matrix multiplication
25//! - `batched_4d_gemm_forward` - 4D batched GEMM for multi-head attention
26//! - `residual_add_forward` - Element-wise addition (residual connections)
27//! - `elementwise_mul_forward` - Element-wise multiplication
28//! - `scale_forward` - Scalar multiplication
29
30mod activations;
31pub mod bf16_cast;
32mod cache;
33mod elementwise;
34mod matmul;
35pub mod matmul_f16;
36mod normalization;
37#[cfg(test)]
38mod tests;
39
40pub use activations::{
41    batched_softmax_forward, gelu_forward, relu_forward, silu_forward, softmax_forward,
42};
43pub use bf16_cast::{bf16_slice_to_f32, f32_slice_to_bf16};
44#[cfg(feature = "cuda")]
45pub use bf16_cast::{cast_bf16_to_f32_gpu, cast_f32_to_bf16_gpu, cast_f32_to_f16_gpu};
46#[cfg(feature = "cuda")]
47pub(crate) use cache::set_forward_cublas_stream;
48pub use cache::{
49    init_forward_kernel_cache, pre_warm_forward_kernels, pre_warm_lora_backward_kernels,
50    set_cublas_workspace,
51};
52pub use elementwise::{
53    batched_to_interleaved_forward, batched_transpose_forward, elementwise_mul_forward,
54    expand_kv_heads, inplace_add_gpu, interleaved_to_batched_forward, residual_add_forward,
55    scale_forward,
56};
57#[cfg(feature = "cuda")]
58pub use matmul::gemm_forward_bf16;
59pub use matmul::{
60    batched_4d_gemm_forward, fused_swiglu_forward, gemm_forward, gemm_forward_bt,
61    gemm_nf4_backward_a, gemm_nf4_forward, gemm_nf4_gate_up_forward, gemm_nf4_tc_backward_a,
62    gemm_nf4_tc_forward,
63};
64#[cfg(feature = "cuda")]
65pub(crate) use matmul::{
66    cublas_gemm_backward_a, cublas_gemm_backward_a_accumulate, cublas_gemm_backward_b,
67};
68#[cfg(feature = "cuda")]
69pub use matmul::{gemm_nf4_backward_a_cublas, gemm_nf4_dequant_cublas};
70#[cfg(feature = "cuda")]
71pub use matmul_f16::{gemm_f16_to_f32_backward_a, gemm_f16_to_f32_forward, gemm_forward_f16};
72pub use normalization::{
73    batched_rope_neox_backward, batched_rope_neox_forward, fused_residual_rmsnorm_forward,
74    layer_norm_forward, per_head_rmsnorm_forward, rms_norm_forward, rms_norm_forward_with_eps,
75    rope_neox_forward,
76};
entrenar/autograd/cuda_forward/mod.rs

entrenar/autograd/cuda_forward/
mod.rs
