Skip to content

Repository files navigation

RunIT — LLM Inference Engine

From-scratch LLM inference engine in Rust + Metal for Apple Silicon

macOS 14+ Rust 1.78+ Metal GPU llama.cpp match MIT License

No Python. No PyTorch. No llama.cpp dependency. Pure Rust + Metal shaders.


✨ Highlights

┌─────────────────────────────────────────────────────────────────┐
│  🎯 100% token-match with llama.cpp on Qwen2.5-0.5B Q8_0      │
│  ⚡ 233 tok/sec (Q8_0, 3.2× speedup) — matches llama.cpp       │
│  🧠 Qwen2 + Qwen3 + LLaMA + Mistral + OLMoE architectures    │
│  🦀 Pure Rust — zero Python runtime                            │
│  🔧 15 Metal shaders, 85+ GPU kernels, fused Q4/Q8 dequant    │
│  📦 GGUF native — Q4_0, Q8_0, Q4K, Q5K, Q6K, F16, BF16      │
│  🌐 OpenAI-compatible HTTP server                              │
│  🔬 f32 precision pipeline for research-grade accuracy         │
└─────────────────────────────────────────────────────────────────┘

📊 Accuracy: Token-Perfect Match with llama.cpp

15/15 position accuracy

Every test produces identical output to llama.cpp (greedy, temp=0):

Prompt llama.cpp RunIT Engine Match
"What is 2+2?" 2 + 2 equals 4. 2 + 2 equals 4. ✅
"Capital of France?" The capital of France is Paris. The capital of France is Paris. ✅
"Capital of Japan?" The capital of Japan is Tokyo. The capital of Japan is Tokyo. ✅
"What is machine learning?" Machine learning is a subset of AI... (92 tok) Machine learning is a subset of AI... (92 tok) ✅
"Translate to French" Bonjour, comment ça va ? Bonjour, comment ça va ? ✅
"Write a poem about the moon" The moon, a celestial sight... The moon, a celestial sight... ✅

📋 Full per-position logit comparison: docs/COMPARISON.md

⚡ Performance

Performance comparison — Before & After

Per-prompt decode speed

🚀 3.2× Speedup: 73 → 233 tok/sec (matches llama.cpp)

Metric v0.0.1 v0.2.0 v0.3.0 v0.7.0 (Q8_0) llama.cpp
Decode tok/sec 73 161 224 233 ✅ 230
Avg latency 13.94 ms 6.28 ms 4.73 ms 4.29 ms ~4.35 ms
Quantization Q8_0 Q8_0 Q8_0 Q8_0 Q8_0
Output quality ✅ Perfect ✅ Perfect ✅ Perfect ✅ Perfect ✅

Measured on Apple M4 Pro · Qwen2.5-0.5B Q8_0 · greedy decode · 50 tokens

The engine also supports Q4_0 fused GEMV (258 tok/sec) for larger models (3B+) where 4-bit quantization produces good output. On 0.5B, Q4_0 loses too much precision.

Optimization Journey

Version Optimization tok/sec Speedup Key Technique
v0.0.1 Kahan GEMV + f32 weights 73 1.0× Correctness-first
v0.2.0 simd_sum + half4 + f16 dequant 161 2.2× Hardware SIMD reduction
v0.3.0 Fused Q8_0 GEMV 224 3.1× On-the-fly dequant (1.06 B/elem)
v0.6.0 Kernel fusion (QKV+bias, FFN, RoPE) 228 3.1× 242 dispatches (was 387)
v0.7.0 + Fused Q4_0 GEMV 233 (Q8) 3.2× Q4_0 ready for larger models

What Makes It Fast

Optimization Impact Files
Fused Q8_0 GEMV 1.06 bytes/elem — 47% less BW than f16 (primary path) gemv_q8_0.metal
Fused Q4_0 GEMV 0.56 bytes/elem — for larger models (3B+) at Q4_0 gemv_q4_0.metal
Fused QKV+bias 6 dispatches → 1 per layer (Qwen path) fused_qkv.metal
Fused gate+up+silu 3 dispatches → 1 per layer fused_ffn.metal
simd_sum() reduction 1-cycle hardware sum replaces 124-step Kahan all GEMV shaders
Vectorized half4/float4 4× fewer loads, perfect 256-byte coalescing gemv.metal
Multi-row GEMV 4 rows/TG (128 threads) for better occupancy gemv_multirow.metal

📋 Full roadmap: docs/implementation-plan.md


🚀 Quickstart

Prerequisites

  • macOS 14+ with Apple Silicon (M1/M2/M3/M4)
  • Rust 1.78+ (curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh)
  • Xcode Command Line Tools (xcode-select --install)

Download a Model

# Qwen2.5-0.5B Q8_0 (recommended for testing — 675 MB)
huggingface-cli download Qwen/Qwen2.5-0.5B-Instruct-GGUF \
    qwen2.5-0.5b-instruct-q8_0.gguf --local-dir ~/models/

# Qwen2.5-Coder-7B Q4_K_M (for coding tasks — 4.7 GB)
huggingface-cli download Qwen/Qwen2.5-Coder-7B-Instruct-GGUF \
    qwen2.5-coder-7b-instruct-q4_k_m.gguf --local-dir ~/models/

# Download tokenizer
huggingface-cli download Qwen/Qwen2.5-0.5B-Instruct \
    tokenizer.json --local-dir ~/models/

Build & Run

# Build
cargo build --release -p bare-metal-engine

# Generate text
./target/release/generate ~/models/qwen2.5-0.5b-instruct-q8_0.gguf \
    --tokenizer ~/models/tokenizer.json \
    --prompt '<|im_start|>user
What is 2+2?<|im_end|>
<|im_start|>assistant
' \
    --tokens 100

# Start OpenAI-compatible server
./target/release/serve ~/models/qwen2.5-0.5b-instruct-q8_0.gguf \
    ~/models/tokenizer.json --port 8080

# Query the server
curl http://localhost:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"messages":[{"role":"user","content":"Hello!"}],"max_tokens":64}'

CLI Options

generate <model.gguf> [OPTIONS]

Options:
  --tokenizer <path>      Path to tokenizer.json
  --prompt <text>         Input prompt
  --tokens N              Max tokens to generate (default: 30)
  --temperature T         0.0 = greedy (default: 0.0)
  --top-p P               Nucleus sampling (default: 0.9)
  --top-k K               Top-K sampling (default: 50)
  --rep-penalty F         Repetition penalty (default: 1.1)
  --seed S                RNG seed (default: 42)
  --tq                    Enable TurboQuant KV cache compression

🏗️ Architecture

Per-prompt decode speed

Crate Structure

Crate Description
bare-metal-gguf GGUF parser — mmap, tensor metadata, zero-copy views
bare-metal-tokenizer HuggingFace tokenizer wrapper
bare-metal-kernels Metal shaders + Rust dispatch layer
bare-metal-engine Model config, forward pass, KV cache, server

Metal GPU Kernels

Kernel File Description
gemv_q8_0_f32in_f32out gemv_q8_0.metal Fused Q8_0 GEMV — 47% less bandwidth (primary)
gemv_q4_0_f32in_f32out gemv_q4_0.metal Fused Q4_0 GEMV — for larger models at 4-bit
gemv_f16w_f32in_f32out gemv.metal GEMV with simd_sum + vectorized half4 loads
gemv_q4k_f16 gemv_q4k.metal Fused Q4_K dequant + GEMV
decode_attention_f32 attention.metal Decode attention with online softmax
flash_attention_f16 attention.metal Tiled FlashAttention-2 for prefill
rms_norm_f32_f32_f32g norm.metal Full f32 RMSNorm
rope_inplace_f32 rope.metal Non-interleaved RoPE (Qwen2/LLaMA)
silu_mul_f32 activation.metal SwiGLU activation (fused)
dequant_q4k_f16 dequant.metal Q4_K_M GPU dequantization
fused_qkv_bias_q8_0_f32 fused_qkv.metal Q+K+V+bias in ONE dispatch
fused_ffn_q4_0_f32 fused_ffn.metal gate+up+silu fused
+ 60 more various See docs/KERNELS.md for full reference

🔬 Precision Pipeline

RunIT uses an f32 residual stream with fused quantized weight kernels for maximum speed without quality loss. Weights are read in their native packed format (Q4_0/Q8_0) and dequantized on-the-fly during the GEMV dot product — no intermediate buffer needed:

Token Embedding (f32 lookup)
    │
    ▼
RMSNorm (f32 in → f32 out, f32 gamma)
    │
    ▼
Q/K/V Projections (fused QKV+bias, Q8_0/Q4_0 on-the-fly dequant)
    │                    └── simd_sum + vectorized half4 loads
    ▼
QK-Norm (optional, Qwen3: per-head RMSNorm on Q and K)
    │
    ▼
RoPE (f32, non-interleaved pairing)
    │
    ▼
KV Cache (f32, flat layout)
    │
    ▼
Decode Attention (f32 Q·K, f32 softmax, f32 output)
    │
    ▼
FFN: gate/up → SiLU → down (f16 weights, f32 accumulation)
    │
    ▼
lm_head (f16 weights × f32 input → f32 logits)

Weights dequantized from Q8_0/Q5K/Q6K to f16 (2 bytes/element) instead of f32 (4 bytes). Quantized formats have ≤8 significant bits — f16's 10-bit mantissa loses nothing.


📦 Supported Models & Quantizations

Architecture Status Models Tested tok/sec
Qwen2 ✅ Full support Qwen2.5-0.5B (Q8_0) 233
Qwen3 ✅ Full support Qwen3-0.6B (Q8_0) 176
LLaMA ✅ Full support LLaMA-2, LLaMA-3 —
Mistral ✅ Full support Mistral-7B —
OLMoE ✅ MoE support OLMoE-1B-7B (Q4K_M) —
Quantization Status Notes
Q8_0 ✅ Fused GEMV On-the-fly dequant, 1.06 B/elem — primary path
Q4_0 ✅ Fused GEMV On-the-fly dequant, 0.56 B/elem — fastest (3B+ models)
Q4_K_M ✅ Fused Q4K GPU dequant for aligned K, f16 fallback
Q6_K ✅ Supported CPU dequant to f16
Q5_K ✅ Supported CPU dequant to f16
Q5_0 ✅ Supported CPU dequant to f16
Q8_K ✅ Supported CPU dequant to f16
F16 ✅ Native Zero-copy if page-aligned
BF16 ✅ Converted BF16 → f16 at load time

📈 Development Roadmap

Phase Feature Status
0 GGUF parser (mmap, tensors, metadata) ✅ Done
1 Metal GPU kernels (GEMV, RoPE, RMSNorm, Attention) ✅ Done
2 Model loading + config ✅ Done
3 Transformer forward pass + KV cache ✅ Done
4 Q4_K_M GPU dequantization ✅ Done
5 TurboQuant KV-cache compression (3-4 bit) ✅ Done
6 Text I/O — tokenizer, sampling ✅ Done
7 OpenAI-compatible HTTP server ✅ Done
8 Prefill batching (GEMM kernel) ✅ Done
9 f32 precision pipeline ✅ Done
10 RoPE fix + llama.cpp parity ✅ Done
11 GEMV optimization (simd_sum, vectorize, f16 dequant) ✅ Done — 2.2×
12 Fused Q8_0 GEMV (on-the-fly dequant, 47% less BW) ✅ Done — 3.1×
13 Kernel fusion (QKV+bias, FFN, RoPE) + multi-row ✅ Done
14 Fused Q4_0 GEMV (for larger models at 4-bit) ✅ Done
15 Profiling + speculative decoding infra ✅ Done
16 Qwen3 support (QK-norm + head_dim detection) ✅ Done
17 7B model support + benchmarks 🔜 Next
18 PagedAttention KV cache + continuous batching 📋 Planned

🐛 Major Bug Fix: RoPE Dimension Pairing

The engine had a critical bug where the Rotary Position Embedding used interleaved pairing (0,1), (2,3)... instead of the correct non-interleaved pairing (i, i+d/2) used by Qwen2/LLaMA/Mistral.

Why it was hidden: At position 0, RoPE is the identity rotation (cos(0)=1, sin(0)=0) regardless of pairing — so single-token tests always passed. The bug only manifested at position ≥ 1, scrambling the positional encoding and corrupting multi-token attention.

Before (wrong):  pairs (0,1) (2,3) (4,5) ... → scrambled positions at pos≥1
After (correct): pairs (0,32) (1,33) (2,34) ... → perfect llama.cpp match

📋 Full investigation: docs/COMPARISON.md · docs/STATUS_REPORT.md


📂 Repository Layout

RunIT/
├── crates/
│   ├── gguf/           # GGUF file parser (mmap, zero-copy)
│   ├── tokenizer/      # HuggingFace tokenizer wrapper
│   ├── kernels/        # Metal shaders + Rust dispatch
│   │   ├── shaders/    # .metal source → .metallib at build time
│   │   │   ├── gemv.metal          # 14 GEMV variants (simd_sum + half4)
│   │   │   ├── gemv_q4_0.metal     # 7 fused Q4_0 GEMV (0.56 B/elem) 🏆
│   │   │   ├── gemv_q8_0.metal     # 7 fused Q8_0 GEMV (1.06 B/elem)
│   │   │   ├── gemv_q4k.metal      # 6 fused Q4K GEMV
│   │   │   ├── fused_qkv.metal     # Q+K+V+bias in 1 dispatch
│   │   │   ├── fused_ffn.metal     # gate+up+silu in 1 dispatch
│   │   │   ├── fused_rope.metal    # QK RoPE in 1 dispatch
│   │   │   ├── gemv_multirow.metal # Multi-row GEMV (4 rows/TG)
│   │   │   ├── attention.metal     # FlashAttention-2 + decode
│   │   │   ├── gemm.metal          # simdgroup_matrix GEMM + tiled
│   │   │   ├── norm.metal          # RMSNorm (4 variants)
│   │   │   ├── rope.metal          # RoPE (f16/f32/batch)
│   │   │   ├── activation.metal    # SwiGLU, add, argmax, KV scatter
│   │   │   ├── dequant.metal       # Q4K GPU dequantization
│   │   │   └── turboquant.metal    # TurboQuant KV compression
│   │   └── src/        # MetalContext, dispatch, error types
│   └── engine/         # Model loader, forward pass, server
│       └── src/
│           ├── forward.rs       # Executor + 8 quantization decoders
│           ├── kv_cache.rs      # f32 flat KV cache
│           ├── sampler.rs       # temp/top-k/top-p/rep-penalty
│           ├── server/          # OpenAI-compatible HTTP (axum, SSE)
│           └── bin/
│               ├── generate.rs  # CLI text generation + benchmark
│               └── serve.rs     # HTTP server binary
├── docs/
│   ├── assets/          # Logo, charts, diagrams
│   ├── ARCHITECTURE.md  # System design and data-flow
│   ├── COMPARISON.md    # llama.cpp comparison + test results
│   ├── KERNELS.md       # Metal kernel reference
│   └── STATUS_REPORT.md # Bug fixes + precision improvements
└── README.md

🔧 Development

# Check all crates
cargo check --workspace

# Run tests
cargo test --workspace

# Build release
cargo build --release -p bare-metal-engine

# Debug: per-position logits during prompt processing
DEBUG_LOGITS=1 ./target/release/generate <model.gguf> \
    --tokenizer <tokenizer.json> --tokens 10 \
    --prompt '<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n'

# Debug: per-layer hidden state dump
DEBUG_LAYERS=1 ./target/release/generate <model.gguf> \
    --tokenizer <tokenizer.json> --tokens 1 --prompt 'Hello'

📚 References & Further Reading

  • Awesome LLM Inference — Curated collection of LLM inference papers, frameworks, and optimization techniques. Covers attention mechanisms, quantization, speculative decoding, KV cache management, and hardware-specific optimizations that inform RunIT’s design.
  • FlashAttention-2 — Tiled attention with online softmax (implemented in attention.metal)
  • TurboQuant — KV-cache compression to 3–4 bits (implemented in turboquant.metal)
  • GGUF Format — Model file format specification
  • llama.cpp — Reference C++ implementation used for accuracy validation
  • Metal Shading Language Spec — Apple’s GPU programming language

📄 License

MIT


Built with 🦀 Rust and ⚡ Metal
For Apple Silicon, from the ground up

About

From-scratch LLM inference engine in Rust + Metal for Apple Silicon

Resources

Stars

7 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages