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