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;
43pub mod host_memory;
44pub mod instance;
45pub mod kernel_registry;
46pub mod kv_block;
47pub mod kv_disk;
48pub mod kv_signature;
49pub mod kv_swa;
50pub mod mamba2;
51pub mod matmul;
52pub mod mla_absorbed;
53pub mod par;
54pub mod placement;
55pub mod qstar;
56pub mod recurrent_state;
57pub mod residency;
58pub mod summary_stats;
59pub mod tensor;
60pub mod threads;
61pub mod turboquant;
62pub mod vexp;
63pub mod weight_matrix;
64
65pub use attention::{
66    apply_rope_back, apply_rope_interleaved, apply_rope_interleaved_back,
67    apply_rope_interleaved_with_freq_factors, apply_rope_with_freq_factors, causal_gqa_attention,
68    causal_gqa_attention_paged, causal_gqa_attention_paged_sinks, causal_gqa_attention_prefill,
69    causal_gqa_attention_prefill_shared_kv, causal_gqa_attention_prefill_shared_kv_windowed,
70    causal_gqa_attention_sinks, causal_gqa_attention_softcap, causal_gqa_attention_windowed,
71    causal_gqa_attention_windowed_softcap, lightning_indexer_topk,
72};
73pub use cache::{
74    KvBlockPool, KvCache, KvPoolExhausted, PagedKvCache, PagedKvStore, PagedStoreExhausted,
75    SharedPagedKv,
76};
77pub use csa_hca_compress::{channel_gated_pool, compress_block};
78pub use deepseek_v4_attention::{csa_attention, hca_attention};
79pub use kernel_registry::Registry as KernelRegistry;
80pub use kv_block::{full_blocks, BlockHash, BlockHasher};
81pub use kv_disk::{
82    decode_block, encode_block, encoded_len, BlockFormatError, DiskConfig, DiskKvStore, DiskStats,
83    ReadHandle, ReadOutcome, StoreError,
84};
85pub use kv_signature::{
86    CacheSignature, KvBlock, KvDtype, SignatureError, UnverifiedBlock, BLOCK_FORMAT_VERSION,
87    READABLE_FORMAT_VERSIONS,
88};
89pub use kv_swa::{aligned_block_size, BlockLayout, BlockLayoutError};
90pub use matmul::{
91    geglu, gelu, matmul_f32, rms_norm, rms_norm_per_head, silu, situ_and_mul, softcap_inplace,
92    swiglu,
93};
94pub use tensor::Tensor;
95#[cfg(feature = "cuda")]
96pub use weight_matrix::cuda_dense_enabled;
97#[cfg(feature = "metal")]
98pub use weight_matrix::metal_dense_enabled;
99pub use weight_matrix::{
100    active_backend, cpu_int_dot_kind_supported, cuda_matvec_kind_supported, metal_matvec_kind_name,
101    metal_mul_mm_kind_supported, BatchActs, QuantKind, WeightMatrix,
102};