UniLLM

A modular LLM inference runtime written in Rust.

UniLLM provides a unified, type-safe interface for running large language models across 47 architectures. It is built around three composable abstractions -- TensorCore, ModelCore, and WeightLoaderCore -- that let you load weights in any format, run inference on any device, and add new model architectures with minimal boilerplate.

Quick Start

git clone https://github.com/cognisoc/unillm.git
cd unillm
cargo check            # verify compilation
cargo test --workspace # run all tests

Run inference

# Generate text (downloads TinyLlama on first run, ~600MB)
cargo run --bin unillm -p unillm-runtime -- generate --prompt "Explain gravity"

# Use a different model
cargo run --bin unillm -p unillm-runtime -- generate --model llama2:7b --prompt "Hello"

# List cached models
cargo run --bin unillm -p unillm-runtime -- models

Supported Models

UniLLM implements 47 model architectures across 10 categories:

Category	Models
Core LLMs	LLaMA, Qwen, Gemma, Phi, DeepSeek, Mistral, Mixtral
GPT Family	GPT-2, GPT-J, GPT-NeoX, OPT, BLOOM, MPT
Code	StarCoder, CodeLlama
MoE	DeepSeek-MoE, DBRX, Grok, Arctic, Jamba
RWKV / Linear Attention	RWKV-4, RWKV-6, RecurrentGemma
Vision-Language	Qwen2-VL, Phi-3-Vision, InternVL, CogVLM, Idefics, Florence, LLaVA, CLIP
Audio / Speech	Wav2Vec2, HuBERT, MusicGen, Encodec, Whisper
Encoder	BERT, T5
Specialized	Mamba, MiniCPM, OLMo, Granite
Additional	Yi, Falcon, Baichuan, InternLM, ChatGLM

All models share the same Model trait and are configured through the model_config! macro.

Architecture

UniLLM is organized into three layers:

TensorCore -- Device-agnostic tensor operations (CPU, CUDA, Metal). All ops go through ops_fn::operation().
ModelCore -- Universal Model trait with forward() and generate(). Configuration via model_config! macro.
WeightLoaderCore -- Format-agnostic weight loading for SafeTensors, GGUF, and PyTorch files.

Adding a model

model_config!(MyModelConfig {
    vocab_size: usize = 32000,
    hidden_size: usize = 4096,
    num_hidden_layers: usize = 32,
});

impl Model for MyModel {
    type Config = MyModelConfig;

    fn forward(&self, inputs: &ModelInputs) -> Result<ModelOutputs> {
        // model-specific forward pass
    }
}

Project Structure

crates/
  runtime/       Core inference runtime (tensor ops, model trait, weight loading, 47 models)
  inference/     High-level inference engine and batching
  kv/            Hybrid KV cache (RadixAttention + PagedAttention)
  scheduler/     Request scheduling with continuous batching
docs/            Architecture docs, API reference, developer guide

Development

cargo check                       # type-check the workspace
cargo test --workspace            # run all tests
cargo test --lib -p unillm-runtime # test the runtime crate
cargo clippy --workspace          # lint
cargo fmt --all                   # format
cargo build --release             # optimized build

See docs/developer_guide.md for a full development setup guide.

Documentation

Roadmap -- Current status, what works, what's next
Architecture -- Three-layer system design
API Reference -- Detailed API docs
Developer Guide -- Getting started with development

License

Apache-2.0 -- see LICENSE for details.

unillm-runtime 0.1.0