lattice-inference 0.6.0

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
83
84
85
86
87
88
89
90
91
92

//! Qwen3.5-2B text generation model.
//!
//! Hybrid architecture: 18 GatedDeltaNet (linear attention) + 6 standard GQA layers.
//! Pattern: [linear, linear, linear, full] x 6 = 24 layers.
//!
//! Key features:
//! - Autoregressive text generation with temperature/top-k/top-p sampling
//! - Tied embeddings (lm_head = embed_tokens^T)
//! - Partial RoPE (25% of head_dim) on full-attention layers
//! - GatedDeltaNet recurrent state for linear layers, KV cache for full layers

mod cache;
mod debug;
pub(crate) mod detokenize;
mod eval;
mod forward;
mod generation;
mod loading;
mod model;
mod moe;
mod norm;
mod rerank;
mod sampling;
pub(crate) mod stop_strings;
mod weights;

pub use eval::{PerplexityConfig, PerplexityReport};
/// Re-exported for the Metal Q4 perplexity harness in
/// [`crate::forward::metal_qwen35`]; the CPU forward path consumes them
/// directly inside [`eval`].
#[cfg(all(target_os = "macos", feature = "metal-gpu"))]
pub(crate) use eval::{log_softmax_nll, run_strided_perplexity};

#[cfg(test)]
mod tests;

/// Test-only tiny zero-weight model construction (ADR-080 C2 round 2, codex
/// finding #2), gated behind the `test-utils` Cargo feature so it never
/// ships in a normal build -- see the module doc comment for why this can't
/// just be `#[cfg(test)]` like the library's own internal test fixtures.
#[cfg(feature = "test-utils")]
pub mod test_support;

pub use model::Qwen35Model;
pub use weights::ModelWeights;

pub(crate) use cache::{ForwardScratch, KvCache, resize};
pub(crate) use detokenize::decode_tokens;
// Re-exported so that quantized CPU generate helpers (cpu_q8, cpu_f16,
// neon_forward) can share the same typed grammar-not-set guard without
// duplicating the predicate or the error message (#397/#398).
pub(crate) use generation::check_grammar_not_set;
// Sibling guard for `logprobs` on the same unwired paths (#585).
pub(crate) use generation::check_logprobs_not_set;
// Sibling guards for `stop_strings` / `reasoning_budget` on the same unwired
// paths (ADR-080 C3, #783).
pub(crate) use generation::{check_reasoning_budget_not_set, check_stop_strings_not_set};
// Shared backend-neutral decode-policy struct (reasoning-budget accounting +
// logprobs formatting), consumed by the Metal streaming loops in
// `crate::forward::metal_qwen35` so the same bookkeeping isn't re-duplicated
// across the CPU/Metal boundary (ADR-080 C3). Only the Metal (`mod inner`,
// gated identically) consumer needs the re-export; `generation.rs` itself
// uses `DecodePolicy` directly within its own module.
#[cfg(all(target_os = "macos", feature = "metal-gpu"))]
pub(crate) use generation::{DecodePolicy, StepOutcome, StopCheckOutcome};
// Sibling guard for `enable_mtp` on the cross-turn prefix-cache path, which
// has no MTP draft/verify wiring (codex round-2 medium #4, PR #787). Only
// that Metal-only path needs it, same gate as `DecodePolicy`/`StepOutcome`.
#[cfg(all(target_os = "macos", feature = "metal-gpu"))]
pub(crate) use generation::check_mtp_not_requested;
pub(crate) use norm::qwen35_rms_norm;
pub(crate) use sampling::sample_token;
pub(crate) use weights::{
    AttentionWeights, CommonLayerWeights, DenseFfnWeights, FeedForwardWeights,
    FullAttentionLayerWeights,
};

#[cfg(test)]
pub(crate) use weights::{MoeLayerWeights, MoeRouter, RoutedExperts, SharedExpert};

#[cfg(test)]
pub use detokenize::bytes_to_unicode;
// Needed by all generate paths (cpu_q8, cpu_f16, neon_forward, batch_prefill)
// and by tests. `pub(crate)` keeps it out of the public API surface.
pub(crate) use generation::should_stop_token;

/// Exposed for consumers that need to drive a per-layer coverage check
/// over a Qwen3.5 checkpoint without going through the model loader —
/// e.g., the QuaRot offline converter (ADR-044 step 3c) iterating
/// rotation rules against an actual safetensors file. Originally
/// `#[cfg(test)]`-only; promoted in step 3b.
pub use loading::qwen_required_tensor_names;