High-performance layer-wise LLM inference in Rust. Memory-efficient inference for models that don't fit in GPU memory.

Rust 96.4%
Makefile 2.6%
Python 1%

Find a file

root 51af72b49a Optimize matmul and attention for better TPS Performance optimizations based on computer engineering principles: 1. Cache-blocked matmul: - 64x64 blocks tuned for L1/L2 cache (32KB/256KB) - Parallel over row blocks with rayon - 4x loop unrolling for ILP - Fall back to simple impl for small matrices 2. Parallel attention: - Heads computed in parallel with rayon - Unrolled dot products (4x) - Causal softmax only over valid positions 3. Parallel element-wise ops: - RMS norm parallel across batch - SiLU, mul, add use par_iter Results on 48-core Threadripper: - Before: >120s for 5 tokens (timeout) - After: ~16s for 5 tokens (~0.3 TPS) - user time > real time confirms parallelization works Still TODO for better TPS: - KV cache to avoid O(n²) recomputation - In-place operations to reduce allocations - SIMD intrinsics for dot products		2026-01-24 15:35:36 +00:00
python/airllm_rs	Add Python bindings with PyO3 and README	2026-01-22 09:02:48 +00:00
src	Optimize matmul and attention for better TPS	2026-01-24 15:35:36 +00:00
.gitignore	Initial commit: AirLLM-RS - Layer-wise LLM inference in Rust	2026-01-22 08:47:46 +00:00
Cargo.lock	Add HuggingFace tokenizers integration	2026-01-22 09:13:00 +00:00
Cargo.toml	Add HuggingFace tokenizers integration	2026-01-22 09:13:00 +00:00
Makefile	Add Makefile for easy commands	2026-01-22 09:20:10 +00:00
pyproject.toml	Add Python bindings with PyO3 and README	2026-01-22 09:02:48 +00:00
README.md	Add Python bindings with PyO3 and README	2026-01-22 09:02:48 +00:00

README.md

AirLLM-RS

High-performance layer-wise LLM inference in Rust with Python bindings.

Features

Layer-wise Loading: Load transformer layers one at a time, enabling inference on models that don't fit in GPU memory
Memory-mapped I/O: Zero-copy loading from safetensors files using mmap
Async Prefetching: Prefetch next layer while computing current layer
Multi-architecture: Support for Llama, Qwen, GLM-4/GLM-4.7, Mixtral
BitNet Support: Ternary weight unpacking for 1.58-bit models
Python Bindings: Easy-to-use Python API via PyO3

Installation

From PyPI (coming soon)

pip install airllm-rs

From Source

# Install maturin
pip install maturin

# Clone and build
git clone https://github.com/rileyseaburg/airllm-rs
cd airllm-rs
maturin develop --features python

Rust Library

cargo add airllm-rs

Usage

Python

from airllm_rs import AirLLM, GenerationConfig

# Load model (layer-by-layer, memory efficient)
model = AirLLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# Check model info
print(model.config())
# ModelConfig(arch=Llama, layers=32, hidden=4096, vocab=32000)

# Generate with custom config
config = GenerationConfig(
    max_new_tokens=100,
    temperature=0.7,
    top_p=0.9
)

# Input token IDs (use your tokenizer)
input_ids = [1, 15043, 29892, 920, 526, 366]  # "Hello, how are you"
output_ids = model.generate(input_ids, config)
print(output_ids)

Rust

use airllm_rs::{InferenceEngine, GenerationConfig};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Load model
    let engine = InferenceEngine::from_pretrained("./model")?;
    
    println!("Loaded: {:?}", engine.config().architecture);
    
    // Generate
    let config = GenerationConfig::default();
    let input_ids = vec![1, 15043, 29892];
    let output = engine.generate(&input_ids, &config)?;
    
    println!("Generated: {:?}", output);
    Ok(())
}

CLI

# Show model info
airllm-cli info --model ./model

# List tensors
airllm-cli tensors --model ./model --filter mlp

# Run inference (WIP: needs tokenizer)
airllm-cli run --model ./model --prompt "def fibonacci(n):"

Architecture

┌─────────────────────────────────────────────────────────────┐
│                      AirLLM-RS                              │
├─────────────────────────────────────────────────────────────┤
│  ┌─────────────┐  ┌─────────────┐  ┌─────────────────────┐  │
│  │ LayerCache  │  │ TensorStore │  │ InferenceEngine     │  │
│  │ (mmap pool) │  │ (safetensors│  │ (layer-by-layer)    │  │
│  └─────────────┘  └─────────────┘  └─────────────────────┘  │
├─────────────────────────────────────────────────────────────┤
│  ┌─────────────┐  ┌─────────────┐  ┌─────────────────────┐  │
│  │ Attention   │  │ MLP/FFN     │  │ Embeddings          │  │
│  │ (GQA/MHA)   │  │ (SwiGLU)    │  │ (+ RoPE)            │  │
│  └─────────────┘  └─────────────┘  └─────────────────────┘  │
└─────────────────────────────────────────────────────────────┘

Supported Models

Architecture	Status	Notes
Llama/Llama-2/Llama-3	✅	Full support
Mistral	✅	Uses Llama naming
Qwen/Qwen2	✅	Full support
GLM-4	✅	Including GLM-4.7 MoE
Mixtral	⚠️	MoE routing WIP
BitNet	✅	Ternary weight support

Performance

Memory usage comparison for Llama-2-7B:

Method	VRAM Required
Standard PyTorch	14GB
8-bit quantization	7GB
AirLLM-RS (layer-wise)	< 1GB

Note: Layer-wise loading trades memory for speed. Best for memory-constrained environments.

License

MIT License - see LICENSE

AirLLM - Original Python implementation
llama.cpp - C++ inference
candle - Rust ML framework