lattice-inference 0.4.2

Pure Rust transformer inference engine — safetensors loading, SIMD matmul, BGE/Qwen3 embeddings
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

//! **Stability tier**: Experimental
//!
//! This is a pure ML inference kernel with high churn, 153 `unsafe` blocks, and 22
//! `dead_code_allows`. It is NOT intended for direct use by platform or feature crates.
//! Consumers should go through `lattice-embed`. The unsafe blocks are documented in
//! `foundation/STABILITY.md §Tech Debt`. Tracking issue: #1306.
//! See `foundation/STABILITY.md` for the full policy.
//!
// ML inference kernels: many functions have >7 args by necessity (BLAS-style APIs
// where grouping into structs would require heap allocation in hot paths), and many
// loops use the index to access multiple arrays simultaneously so the
// needless_range_loop suggestion does not apply.
#![allow(clippy::too_many_arguments)]
#![allow(clippy::needless_range_loop)]
//! lattice-inference: pure Rust transformer inference for embedding models.
//!
//! Supports two architectures:
//! - **BERT/BGE** (encoder-only): bidirectional attention, mean pooling
//! - **Qwen3** (decoder-only): causal GQA with RoPE, SwiGLU, last-token pooling
//!
//! ## Module Organization
//!
//! - [`model`] — Model configs and loaders (BERT, Qwen, Qwen3.5, BitNet)
//! - [`tokenizer`] — Tokenizers (WordPiece, SentencePiece, BPE)
//! - [`weights`] — Weight storage formats (f32, f16, Q8)
//! - [`attention`] — Attention mechanisms (standard, GQA, flash, GDN)
//! - [`forward`] — Compute backends (CPU, NEON, Metal GPU, batched prefill)

// Grouped modules
pub mod attention;
pub mod forward;
pub mod model;
pub mod tokenizer;
pub mod vision;
pub mod weights;

// Standalone modules
pub mod batch;
pub mod download;
pub mod error;
pub mod generate;
pub mod grammar;
pub mod kv_cache;
pub mod lora_hook;
pub mod metrics;
#[cfg(feature = "mixture")]
pub mod mixture;
pub mod pool;
pub mod pruning;
pub mod quant;
pub mod rope;
pub mod sampling;
pub mod speculative;

#[cfg(feature = "train-backward")]
pub mod backward;

use std::path::PathBuf;

/// Default model cache directory.
pub(crate) fn default_cache_dir() -> Result<PathBuf, error::InferenceError> {
    if let Ok(path) = std::env::var("LATTICE_MODEL_CACHE") {
        return Ok(PathBuf::from(path));
    }
    let home = std::env::var("HOME").map_err(|_| {
        error::InferenceError::ModelNotFound(
            "unable to determine cache directory; HOME and LATTICE_MODEL_CACHE are unset".into(),
        )
    })?;
    Ok(PathBuf::from(home).join(".lattice").join("models"))
}

// Re-exports for public API backward compatibility
pub use crate::error::InferenceError;
pub use crate::model::{
    BertConfig, BertModel, CrossEncoderModel, LayerTimings, ProfileTimings, QwenConfig, QwenModel,
};
pub use crate::pool::BertPooling;
pub use crate::tokenizer::{
    BpeTokenizer, SentencePieceTokenizer, TokenizedInput, Tokenizer, WordPieceTokenizer,
    load_tokenizer,
};