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 attention;
23pub mod bench_profile;
24pub mod block_sparse;
25pub mod cache;
26// The two halves of issue #27's CPU scheduling change: `cpu_pool` is
27// the persistent worker pool, `par` is the one seam every CPU parallel
28// region in this crate goes through and the switch between them.
29pub mod cpu_pool;
30pub mod csa_hca_compress;
31pub mod deepseek_v4_attention;
32pub mod expert_budget;
33pub mod expert_cache;
34// Not feature-gated: `expert_pool::split_pair` is the one part of the
35// slot-copy path a compiler cannot check and a GPU-less host can, so it
36// is compiled and tested by the ordinary `cargo test` run. Everything
37// touching cudarc inside it carries its own `cuda` gate.
38pub mod expert_pool;
39pub mod expert_slots;
40pub mod expert_store;
41pub mod host_memory;
42pub mod instance;
43pub mod kernel_registry;
44pub mod kv_block;
45pub mod kv_disk;
46pub mod kv_signature;
47pub mod kv_swa;
48pub mod matmul;
49pub mod par;
50pub mod placement;
51pub mod qstar;
52pub mod residency;
53pub mod summary_stats;
54pub mod tensor;
55pub mod threads;
56pub mod turboquant;
57pub mod vexp;
58pub mod weight_matrix;
59
60pub use attention::{
61    apply_rope_back, apply_rope_interleaved, apply_rope_interleaved_back,
62    apply_rope_interleaved_with_freq_factors, apply_rope_with_freq_factors, causal_gqa_attention,
63    causal_gqa_attention_paged, causal_gqa_attention_paged_sinks, causal_gqa_attention_prefill,
64    causal_gqa_attention_prefill_shared_kv, causal_gqa_attention_prefill_shared_kv_windowed,
65    causal_gqa_attention_sinks, causal_gqa_attention_softcap, causal_gqa_attention_windowed,
66    causal_gqa_attention_windowed_softcap, lightning_indexer_topk,
67};
68pub use cache::{
69    KvBlockPool, KvCache, KvPoolExhausted, PagedKvCache, PagedKvStore, PagedStoreExhausted,
70    SharedPagedKv,
71};
72pub use csa_hca_compress::{channel_gated_pool, compress_block};
73pub use deepseek_v4_attention::{csa_attention, hca_attention};
74pub use kernel_registry::Registry as KernelRegistry;
75pub use kv_block::{full_blocks, BlockHash, BlockHasher};
76pub use kv_disk::{
77    decode_block, encode_block, encoded_len, BlockFormatError, DiskConfig, DiskKvStore, DiskStats,
78    ReadHandle, ReadOutcome, StoreError,
79};
80pub use kv_signature::{
81    CacheSignature, KvBlock, KvDtype, SignatureError, UnverifiedBlock, BLOCK_FORMAT_VERSION,
82    READABLE_FORMAT_VERSIONS,
83};
84pub use kv_swa::{aligned_block_size, BlockLayout, BlockLayoutError};
85pub use matmul::{
86    geglu, gelu, matmul_f32, rms_norm, rms_norm_per_head, silu, situ_and_mul, softcap_inplace,
87    swiglu,
88};
89pub use tensor::Tensor;
90#[cfg(feature = "cuda")]
91pub use weight_matrix::cuda_dense_enabled;
92#[cfg(feature = "metal")]
93pub use weight_matrix::metal_dense_enabled;
94pub use weight_matrix::{
95    active_backend, cpu_int_dot_kind_supported, cuda_matvec_kind_supported, metal_matvec_kind_name,
96    metal_mul_mm_kind_supported, BatchActs, QuantKind, WeightMatrix,
97};