oxicuda-lm 0.2.0

Large language model inference primitives for OxiCUDA: BPE tokenizer, transformer layers with KV cache, GPT-2 and LLaMA architectures — pure Rust, zero CUDA SDK dependency.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

//! Transformer layer building blocks.
//!
//! # Modules
//!
//! | Module | Contents |
//! |--------|----------|
//! | [`attention`] | Multi-head attention (MHA/GQA) with KV-cache |
//! | [`embedding`] | Token embedding, learned positional embedding, RoPE |
//! | [`ffn`]       | MLP (GELU) and SwiGLU feed-forward networks |
//! | [`mod@flash_attention`] | Tiled online-softmax FlashAttention CPU reference |
//! | [`norm`]      | RMSNorm and LayerNorm |
//! | [`rope_scaling`] | Long-context RoPE scaling (Linear / NTK-aware / YaRN) |
//! | [`transformer`] | GPT-2 and LLaMA transformer blocks; `PastKvCache` |

pub mod attention;
pub mod embedding;
pub mod ffn;
pub mod flash_attention;
pub mod norm;
pub mod rope_scaling;
pub mod transformer;

pub use attention::{LayerKvCache, MultiHeadAttention};
pub use embedding::{LearnedPositionalEmbedding, RotaryEmbedding, TokenEmbedding};
pub use ffn::{MlpFfn, SwiGluFfn, gelu, silu};
pub use flash_attention::{FlashAttentionConfig, flash_attention};
pub use norm::{LayerNorm, RmsNorm};
pub use rope_scaling::{RopeScaling, RopeScalingKind};
pub use transformer::{GptBlock, LlamaBlock, PastKvCache};