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