Skip to main content

ferrox_core/
lib.rs

1//! ferrox-core: tensor primitives, quantized matmul, RMSNorm, RoPE, and
2//! grouped-query causal attention with a simple KV cache.
3//!
4//! CPU reference implementation. The op set and naming (RMSNorm, RoPE,
5//! GQA, KV cache) follow the now-standard vocabulary popularized by
6//! llama.cpp / vLLM / candle-transformers; the actual Rust code below is
7//! written independently. See docs/THIRD_PARTY_NOTICES.md for design credit.
8//!
9//! The MoE expert-residency stack -- [`expert_store`] (the byte budget
10//! and the SSD tier), [`expert_cache`] (which experts stay resident and
11//! the copy plans that make them so), [`expert_slots`] (the bounded
12//! slot pool behind the [`expert_slots::SlotDevice`] seam),
13//! [`expert_pool`] (the CUDA side of that seam), [`expert_budget`]
14//! (bytes in, expert slot count out), [`residency`],
15//! [`placement`] and [`qstar`] -- lives together in one crate on
16//! purpose: on unified memory two independent expert budgets are the
17//! same physical RAM counted twice. [`expert_store`] is the budget
18//! holder. The policy half is ported from FreeToken (Apache-2.0); see
19//! docs/THIRD_PARTY_NOTICES.md.
20
21pub mod activation_tap;
22pub mod alibi;
23pub mod attention;
24pub mod bench_profile;
25pub mod block_sparse;
26pub mod cache;
27// The two halves of issue #27's CPU scheduling change: `cpu_pool` is
28// the persistent worker pool, `par` is the one seam every CPU parallel
29// region in this crate goes through and the switch between them.
30pub mod cpu_pool;
31pub mod csa_hca_compress;
32pub mod deepseek_v4_attention;
33pub mod expert_budget;
34pub mod expert_cache;
35// Not feature-gated: `expert_pool::split_pair` is the one part of the
36// slot-copy path a compiler cannot check and a GPU-less host can, so it
37// is compiled and tested by the ordinary `cargo test` run. Everything
38// touching cudarc inside it carries its own `cuda` gate.
39pub mod expert_pool;
40pub mod expert_slots;
41pub mod expert_store;
42pub mod gdn;
43
44/// The delta rule over a chunk of tokens: one state pass per chunk.
45pub mod gdn_chunk;
46pub mod host_memory;
47pub mod instance;
48pub mod kernel_registry;
49pub mod kv_block;
50pub mod kv_disk;
51pub mod kv_signature;
52pub mod kv_swa;
53pub mod mamba2;
54pub mod matmul;
55pub mod mla_absorbed;
56pub mod par;
57pub mod placement;
58pub mod qstar;
59pub mod recurrent_state;
60pub mod residency;
61pub mod summary_stats;
62pub mod tensor;
63pub mod threads;
64pub mod turboquant;
65pub mod vexp;
66pub mod weight_matrix;
67
68pub use attention::{
69    apply_rope_back, apply_rope_interleaved, apply_rope_interleaved_back,
70    apply_rope_interleaved_with_freq_factors, apply_rope_with_freq_factors, causal_gqa_attention,
71    causal_gqa_attention_paged, causal_gqa_attention_paged_sinks, causal_gqa_attention_prefill,
72    causal_gqa_attention_prefill_shared_kv, causal_gqa_attention_prefill_shared_kv_windowed,
73    causal_gqa_attention_sinks, causal_gqa_attention_softcap, causal_gqa_attention_windowed,
74    causal_gqa_attention_windowed_softcap, lightning_indexer_topk,
75};
76pub use cache::{
77    KvBlockPool, KvCache, KvPoolExhausted, PagedKvCache, PagedKvStore, PagedStoreExhausted,
78    SharedPagedKv,
79};
80pub use csa_hca_compress::{channel_gated_pool, compress_block};
81pub use deepseek_v4_attention::{csa_attention, hca_attention};
82pub use kernel_registry::Registry as KernelRegistry;
83pub use kv_block::{full_blocks, BlockHash, BlockHasher};
84pub use kv_disk::{
85    decode_block, encode_block, encoded_len, BlockFormatError, DiskConfig, DiskKvStore, DiskStats,
86    ReadHandle, ReadOutcome, StoreError,
87};
88pub use kv_signature::{
89    CacheSignature, KvBlock, KvDtype, SignatureError, UnverifiedBlock, BLOCK_FORMAT_VERSION,
90    READABLE_FORMAT_VERSIONS,
91};
92pub use kv_swa::{aligned_block_size, BlockLayout, BlockLayoutError};
93pub use matmul::{
94    geglu, gelu, matmul_f32, rms_norm, rms_norm_per_head, silu, situ_and_mul, softcap_inplace,
95    swiglu,
96};
97pub use tensor::Tensor;
98#[cfg(feature = "cuda")]
99pub use weight_matrix::cuda_dense_enabled;
100#[cfg(feature = "metal")]
101pub use weight_matrix::metal_dense_enabled;
102pub use weight_matrix::{
103    active_backend, cpu_int_dot_kind_supported, cuda_matvec_kind_supported, metal_matvec_kind_name,
104    metal_mul_mm_kind_supported, BatchActs, QuantKind, WeightMatrix,
105};