# infer-inference

Inference engine for transformer models. Forward pass, BLAS-fused attention, hardware-accelerated matmul backends, and pluggable FFN routing.

## Overview

This crate runs transformer forward passes with Apple Accelerate (AMX) and optional Metal GPU acceleration. It uses `infer-vindex` for gate KNN (sparse feature selection) and `infer-models` for weight loading and architecture definitions.

```rust
use infer_inference::InferenceModel;

// Load a model
let model = InferenceModel::load("google/gemma-3-4b-it")?;

// Run inference
let result = infer_inference::predict(
    model.weights(), model.tokenizer(), &token_ids, 5,
);
println!("Top prediction: {} ({:.1}%)", result.predictions[0].0, result.predictions[0].1 * 100.0);
```

## Key Components

| Module | Purpose |
|--------|---------|
| `attention/` | BLAS-fused GQA attention: block, GQA, GPU dispatch, RoPE |
| `forward/` | Forward pass: embed, layer, predict, PLE (per-layer embeddings), trace |
| `ffn/` | FFN evaluation: dense, sparse, highway, route-guided (experimental backends deprecated) |
| `layer_graph/` | Layer graphs + prediction pipeline: `pipeline_layer` (shared FullPipelineLayer construction), `predict` (entry points), `generate` (token loop), `logits` (KNN logits), `prefill` (KV cache) |
| `residual.rs` | RMS norm, layer norm |
| `trace/` | Residual stream decomposition and tiered storage |
| `vindex/walk_ffn.rs` | WalkFfn: mmap'd FFN — faster than dense (517ms vs 535ms) |
| `capture.rs` | Residual stream vector capture for probing |
| `walker/` | Weight-level graph walkers (no forward pass) |
| `model.rs` | Model loading (re-exports from infer-models) |

## Compute Backend

All GPU pipeline operations use `infer_compute::ComputeBackend`:

```rust
use infer_compute::{default_backend, ComputeBackend};

let backend = default_backend();  // Auto-selects CPU or Metal, calibrates
println!("Using: {} ({})", backend.name(), backend.device_info());
```

The inference crate builds `FullPipelineLayer` structs (per-layer architecture params + quantized weights) and passes them to `backend.decode_token()` or `backend.prefill_q4()`. All model-specific behavior (norm type, activation, head_dim, RoPE base) is parameterized per-layer — no model-type branching in the compute path.

**CPU path**: BLAS matmul via Apple Accelerate (AMX). Used for attention in `predict_honest`.
**GPU path** (`--features metal`): Q4_K/Q8 Metal shaders with KV cache. Used for decode and prefill.

```bash
# Build with Metal GPU support
cargo build --release -p infer-inference --features metal
```

## BLAS-Fused Attention

The attention kernel uses BLAS `gemv` inside an online-softmax loop. For each query position:

1. `scores = K[0..=qi] @ Q[qi]` (BLAS gemv, AMX-accelerated)
2. Scale + optional softcap + two-pass softmax (f64 accumulation)
3. `output = V[0..=qi]^T @ softmax_scores` (BLAS gemv)

Never allocates the `[seq, seq]` attention matrix. At Gemma-3's head_dim=256, **1.6x faster** than the materialized path. Supports GQA, softcap (Gemma2), attention weight capture.

## WalkFfn

The WalkFfn replaces the dense down projection with a zero-copy mmap read from the vindex:

1. Gate + up projections from model weights (exact, same as dense)
2. GEGLU activation (exact, same as dense)
3. Down projection from mmap'd `down_features.bin` (zero-copy, feature-major)
4. Result is identical to dense FFN — **and faster** (517ms vs 535ms)

The mmap'd feature-major layout has better page cache behavior than the safetensors weight layout.

Build the required vindex files:
```bash
cargo run --release -p infer-vindex --example convert_gates_f32 -- path/to/vindex
cargo run --release -p infer-vindex --example build_down_features -- path/to/vindex
cargo run --release -p infer-vindex --example build_up_features -- path/to/vindex
```

### Walk-only mode

Drop FFN weights — 16.6GB → 5.5GB:

```rust
let model = InferenceModel::load_walk_only("google/gemma-3-4b-it")?;
// Frees 10.7 GB of FFN tensors. Requires down_features.bin + up_features.bin.
```

### Server

```bash
cargo run --release -p infer-server -- path/to/vindex --port 8080

curl -X POST http://localhost:8080/v1/infer \
  -H "Content-Type: application/json" \
  -d '{"prompt": "The capital of France is", "top": 5, "mode": "walk"}'
```

## Examples

```bash
# Walk inference benchmark (dense vs walk vs HNSW, needs model + vindex)
cargo run --release -p infer-inference --example bench_walk_inference -- \
  --model google/gemma-3-4b-it --vindex path/to/vindex

# Walk boundary sweep (correctness proof across all 34 layers)
cargo run --release -p infer-inference --example bench_walk_boundary_sweep -- \
  --model google/gemma-3-4b-it --vindex path/to/vindex

# Fused attention demo and benchmark
cargo run --release -p infer-inference --example demo_attention
cargo run --release -p infer-inference --example bench_attention

# Backend demo and benchmark (CPU vs Metal)
cargo run --release -p infer-inference --example demo_backend --features metal
cargo run --release -p infer-inference --example bench_backend --features metal

# Full inference benchmark (needs model weights)
cargo run --release -p infer-inference --example bench_inference

# End-to-end inference demo (needs model weights)
cargo run --release -p infer-inference --example demo_inference

# Clustering and pair matching demos
cargo run -p infer-inference --example demo_clustering
cargo run -p infer-inference --example demo_pair_matching
```

### Vindex tools

```bash
# Convert gate vectors from f16 to f32 (zero-copy mmap)
cargo run --release -p infer-vindex --example convert_gates_f32 -- path/to/vindex

# Build feature-major down vectors (contiguous per-feature layout)
cargo run --release -p infer-vindex --example build_down_features -- path/to/vindex
```

## Tests

```bash
# Inference tests (96 tests)
cargo test -p infer-inference

# HNSW tests
cargo test -p infer-vindex --test test_hnsw --release

# Individual test suites
cargo test -p infer-inference --test test_fused_attention   # 23 tests
cargo test -p infer-inference --test test_backend           # 13 tests
cargo test -p infer-inference --test test_modules           # 15 tests
cargo test -p infer-inference --test test_trace             # 14 tests
cargo test -p infer-inference --test test_walkers           # 12 tests
cargo test -p infer-inference --test test_walker_utils      # 10 tests
```

| Area | Tests | Coverage |
|------|-------|----------|
| Backend (ComputeBackend) | 13 | Shape, correctness, batch, Metal vs CPU |
| Fused attention | 23 | GQA, softcap, capture, reference agreement, edge cases |
| FFN + modules | 15 | SiLU, GELU, dense, highway, multi-position |
| Trace stores | 14 | Write/read, tiers, boundaries, additive property |
| Walkers | 12 | Weight/attention walkers, vector extractor |
| Utils | 10 | Top-k, rounding, entity sorting |
| Unit (lib) | 9 | Core module tests |

## Crate Dependencies

```
infer-models      ModelWeights, architecture traits, quant
infer-compute     ComputeBackend, Q4 matvec, Metal GPU (used by vindex + walk_ffn Q4 paths)
infer-vindex      VectorIndex, gate KNN, adaptive residency, Q4 gates
infer-core        Graph, Edge, algorithms (knowledge graph engine)
    ↓
infer-inference   Forward pass, attention, backends, WalkFfn
```

> **Note:** The GPU pipeline paths (`predict_honest`, `predict_pipeline`) use `infer_compute::ComputeBackend`
> with `FullPipelineLayer` structs that carry per-layer architecture params from `infer_models::ModelArchitecture`.
> The CPU attention and FFN paths use direct BLAS calls via ndarray. Both paths produce identical results.

### Vindex module structure

| Module | Responsibility |
|--------|---------------|
| `types` | FeatureMeta, GateIndex trait, WalkHit, callbacks |
| `core` | VectorIndex struct, constructors, loading, accessors |
| `gate` | Gate KNN: search, batch, scores, HNSW, warmup |
| `walk` | Walk FFN data: mmap'd down/up feature-major vectors |
| `hnsw` | HNSW graph index |
| `mutate` | INSERT/DELETE mutations |
| `router` | MoE expert routing |

## Documentation

| Doc | Content |
|-----|---------|
| [PERFORMANCE.md](PERFORMANCE.md) | Component breakdown, cross-crate comparison, Ollama reference |
| [ROADMAP.md](ROADMAP.md) | Planned optimizations, completed items |
| [docs/adr/001](docs/adr/001-fused-attention.md) | BLAS-fused online softmax attention |
| [docs/adr/002](docs/adr/002-walk-ffn.md) | WalkFfn — zero-copy mmap'd down projection |
| [docs/adr/003](docs/adr/003-cached-layer-graph.md) | Cached layer graph for template-fixed layers |
| [docs/adr/004](docs/adr/004-predict-honest.md) | predict_honest — production pipeline with per-layer params |
| [docs/adr/005](docs/adr/005-per-layer-graph.md) | PerLayerGraph — adaptive per-layer strategy |

## License

Apache-2.0