lattice-inference 0.7.2

Pure Rust transformer inference engine — safetensors loading, SIMD matmul, BGE/Qwen3 embeddings
Documentation

lattice-inference — Contributor Map

Stability tier: Experimental. This crate has 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 stability policy.

If you are trying to generate embeddings in an application, start with lattice-embed, not this crate. The README you are reading is a contributor's map for people opening PRs on the inference engine itself.


Module Map

Grouped modules

These modules contain submodules (directories with their own mod.rs).

Module Responsibility
attention Attention kernel variants: standard MHA, GQA, flash (CPU tiled), flash-causal (decoder prefill), GDN (gated decay network), GDN-fused, native sparse, and differential. Exposes AttentionTag for typed dispatch.
forward Compute backends: scalar CPU, NEON (AArch64), Metal (macOS MSL), WGPU, Q8/f16 CPU kernels, tiled GEMM (metal_gemm, gpu_gemm), batched prefill, and BitNet kernel. Exports matmul_bt, rms_norm, silu_inplace, elementwise_mul from the CPU submodule.
model Model configs and loaders: BertModel/BertConfig, QwenModel/QwenConfig, Qwen35Model, CrossEncoderModel, BitNet config. Each submodule owns its safetensors load path and forward-pass dispatch.
tokenizer Tokenizer implementations: WordPieceTokenizer (BERT), BpeTokenizer (Qwen3/GPT-2 byte-level), SentencePieceTokenizer. Shared Tokenizer trait and load_tokenizer auto-detect helper live in tokenizer/common.rs.
vision Qwen3-VL vision encoder (ADR-049): ViT forward pass, patch preprocessing, MLP merger, MultimodalInput type. CPU-only; Metal GPU path is deferred.
weights Weight storage formats (Tensor2D), safetensors memory-mapped loading, f32/f16/Q8/Q4 weight structs, sharded-checkpoint support.

Standalone modules

Module Responsibility
batch Continuous batching engine (ADR-048): BatchWorker, FifoScheduler, Sequence, SequenceManager, chunked prefill interleaved with decode steps.
download Model file cache and conditional download (ensure_model_files). Disabled unless the download feature is enabled.
error InferenceError — the crate's single error type, re-exported at the crate root. Stable; adding variants is backward-compatible.
grammar Grammar-constrained decoding (ADR-046): GrammarEngine compiles a JSON Schema or GBNF spec to a byte-level PDA, then masks logits via mask_logits and advances state via advance each decode step.
kv_cache Two KV cache implementations: FlatKVCache (contiguous pre-allocated, single request) and PagedKVCache (256-token pages, LRU eviction, multi-model serving). Includes PrefixPageCache for prefix reuse.
lora_hook LoraHook trait: defines the apply(layer_idx, module, x, output) contract that the forward pass calls after each linear projection. NoopLoraHook is the zero-overhead default.
metrics Inference metrics (ADR-061): MetricsMode enum and LayerMetrics schema. Zero overhead when MetricsMode::Off.
pool Pooling strategies for producing a single embedding vector: BertPooling (Mean or CLS), mean_pool, cls_pool, last_token_pool (Qwen3), l2_normalize.
pruning ShortGPT block-influence scorer (ADR-060 D2 seed): BlockInfluenceAccumulator and BlockInfluence. Standalone math; the calibration pipeline hooks are future work.
quant Quantization pre-transforms: quarot (Walsh-Hadamard rotation primitives, ADR-044). Not yet wired into the Q4 weight path.
rope Rotary position embeddings: RopeTable precomputes sin/cos tables and applies them in-place to Q and K slices before attention.
sampling Token sampling for generation: Sampler, SamplingConfig (temperature, top-k, top-p, repetition penalty).
speculative N-gram prompt-lookup speculative decoding: NgramSpeculator::speculate + verify_draft (low-level) and generate_with_speculation (high-level wrapper).

Feature-gated modules

Module Feature flag Responsibility
mixture mixture Adapter routing and mixture-of-LoRA support layered on top of lora_hook.
backward train-backward Reverse-mode autodiff for LoRA training: tape, ops, attention_gqa, gradcheck. Submodule of lattice-tune's training loop.

Run an Example

Qwen3.5 text generation demo (src/bin/qwen35_generate.rs):

cargo run --release -p lattice-inference --bin qwen35_generate -- \
  --model-dir PATH --prompt "Hello" --max-tokens 64 --seed 42

Qwen3 embedding benchmark (examples/bench_embedding.rs, registered in Cargo.toml as an [[example]] target — despite an outdated --bin reference in the file's own header comment, run it with --example):

cargo run --release -p lattice-inference --example bench_embedding --features f16

For the concise forward-pass trace (safetensors load → tokenize → prefill → decode-loop sampling), see docs/architecture.md § Forward Pass Pipeline. A deeper Qwen3.5-specific walkthrough lives in docs/forward-pass.md.


Supported Architectures

Architecture Submodule Notes
BERT / BGE (encoder-only) model/bert.rs Bidirectional attention, WordPiece tokenizer, mean or CLS pooling. Covers BGE-small/base/large, E5-small/base, all-MiniLM, paraphrase-multilingual-MiniLM.
Qwen3 / Qwen3.5 (decoder-only) model/qwen.rs, model/qwen35/ Causal GQA, RoPE, RMSNorm, SwiGLU FFN, BPE tokenizer, last-token pooling. Covers Qwen3-Embedding-0.6B and 4B.
BitNet model/bitnet_config.rs 1.58-bit quantized weight architecture.
Qwen3-VL (vision) vision/ ViT patch encoder + MLP merger for multimodal input; CPU-only (ADR-049 v0 scope).
Cross-encoder model/cross_encoder.rs BERT-style reranker with a classification head on top of the CLS token.

Where to Start for Common PR Types

New model architecture. Add a submodule under model/ following the pattern in model/qwen.rs (config struct, weights struct, load function, encode/forward method). Wire the new architecture into lattice-embed's EmbeddingModel enum when ready.

Attention kernel change. Work in attention/ and forward/. A new kernel goes in its own file under attention/ (e.g. attention/flash_causal.rs), registers a variant on AttentionTag, and dispatches from the relevant forward/ backend. See docs/architecture.md for the full forward-pass walkthrough before touching this path.

Sampling or generation. sampling.rs owns the shared Sampler and SamplingConfig primitives. Canonical Qwen3.5 generation lives in model/qwen35/generation.rs, with backend-specific entry points under forward/; grammar-constrained decoding uses the shared grammar module.

LoRA adapter injection. The LoraHook trait in lora_hook.rs is the injection point. The training-side implementation lives in crates/tune. Backward-pass gradients are behind the train-backward feature in backward/.

Quantization. Q4 and Q8 weight structs are in weights/q4_weights.rs and weights/q8_weights.rs. Pre-transform math (QuaRot) is in the quant/quarot/ module (hadamard.rs, convert.rs, pipeline.rs). Wiring QuaRot into the Q4 load path is tracked in ADR-044.


Cross-References

  • docs/architecture.md — crate dependency graph, design philosophy, and forward-pass pipeline.
  • AGENTS.md — contribution policy, SIMD verification rules, agent attribution requirements.
  • docs/adr/INDEX.md — full list of architecture decision records by crate.