haagenti_neural/lib.rs
1//! Neural Compression
2//!
3//! 10x compression using learned codebooks for model weights.
4//!
5//! # Architecture
6//!
7//! ```text
8//! ┌─────────────────────────────────────────────────────────────────┐
9//! │ Neural Compression │
10//! ├─────────────────────────────────────────────────────────────────┤
11//! │ │
12//! │ Traditional: weights.safetensors (16GB) │
13//! │ ↓ LZ4 │
14//! │ weights.hct (8GB) — 2:1 compression │
15//! │ │
16//! │ Neural: weights.safetensors (16GB) │
17//! │ ↓ Learned Encoder │
18//! │ weights.nct (1.6GB) — 10:1 compression │
19//! │ ↓ Learned Decoder (GPU) │
20//! │ weights (reconstructed) │
21//! │ │
22//! │ ┌──────────────────────────────────────────────────────────┐ │
23//! │ │ Compression Pipeline │ │
24//! │ │ │ │
25//! │ │ Encoder (offline, CPU): │ │
26//! │ │ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │
27//! │ │ │ Weight │ -> │ Vector │ -> │ Codebook│ -> indices │ │
28//! │ │ │ Tensor │ │ Quantize│ │ Lookup │ │ │
29//! │ │ └─────────┘ └─────────┘ └─────────┘ │ │
30//! │ │ │ │
31//! │ │ Decoder (runtime, GPU): │ │
32//! │ │ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │
33//! │ │ │ Indices │ -> │ Codebook│ -> │ Residual│ -> weights │ │
34//! │ │ │ │ │ Lookup │ │ Refine │ │ │
35//! │ │ └─────────┘ └─────────┘ └─────────┘ │ │
36//! │ └──────────────────────────────────────────────────────────┘ │
37//! └─────────────────────────────────────────────────────────────────┘
38//! ```
39//!
40//! # Layer-Specific Codebooks
41//!
42//! Different layer types use different codebook sizes:
43//! - Attention Q/K: 4096 centroids, 12-bit indices
44//! - Attention V/O: 8192 centroids, 13-bit indices
45//! - FFN: 2048 centroids, 11-bit indices
46//! - Normalization: 256 centroids, 8-bit (lossless)
47
48mod codebook;
49mod decoder;
50mod encoder;
51mod error;
52mod format;
53mod residual;
54mod training;
55
56pub use codebook::{Codebook, CodebookConfig, CodebookStats, LayerCodebook, LayerType};
57pub use decoder::{DecoderConfig, GpuDecoder, NeuralDecoder};
58pub use encoder::{EncodedTensor, EncoderConfig, NeuralEncoder};
59pub use error::{NeuralError, Result};
60pub use format::{NctFile, NctHeader, NctMetadata};
61pub use residual::{RefinerConfig, ResidualRefiner};
62pub use training::{CodebookTrainer, TrainingConfig};
63
64/// File extension for neural compressed tensors
65pub const NCT_EXTENSION: &str = "nct";
66
67/// Magic bytes for NCT format
68pub const NCT_MAGIC: [u8; 4] = [0x4E, 0x43, 0x54, 0x00]; // "NCT\0"
69
70/// Default compression ratio target
71pub const DEFAULT_COMPRESSION_RATIO: f32 = 10.0;
72
73/// Prelude for common imports
74pub mod prelude {
75 pub use super::{Codebook, NctFile, NeuralDecoder, NeuralEncoder, Result};
76}