Inference Engine

A lightweight, high-throughput LLM inference engine optimized for Apple Silicon (MLX). It features a custom Paged KV Cache to eliminate physical memory fragmentation and an asynchronous Continuous Batching Scheduler to maximize hardware saturation under concurrent loads.

By dynamically managing virtual memory and interleaving request phases, the engine scales throughput efficiently under client saturation, bypassing the O(N²) recomputation bottlenecks of naive sequential serving.

---

Performance Profiling

Below is the throughput scaling profile comparing this engine against a standard naive sequential baseline, using Qwen2.5-1.5B-Instruct-4bit.

Key Metrics & Analysis

• Concurrency 1: Even at single-request loads, the paged layout introduces negligible overhead while matching or exceeding baseline speeds.
• Concurrency 16 (Saturation): Achieves a ~5.5x throughput improvement over the baseline. While the sequential baseline hits a hard performance ceiling due to quadratic history recomputation, this engine successfully scales utilization across the hardware's unified memory bandwidth.

---

Core Components

Paged KV Cache (src/paged_cache.py)

Dynamically divides the Key-Value tensors into fixed-size physical blocks (16 tokens) allocated across a pre-managed tensor pool. Eliminates the need for contiguous memory allocation per request, preventing virtual memory fragmentation and ensuring O(1) page lookups during attention execution.

Continuous Batching Scheduler (src/scheduler.py)

Operates an asynchronous iteration loop that interleaves the multi-token Prefill phase of incoming requests with the single-token Decode phase of ongoing requests. Ensures the GPU/Neural Engine rarely idles by packing independent token generation sequences into dynamic execution matrices.

Asynchronous Server API (src/server.py)

Exposes a concurrent HTTP endpoint via FastAPI that handles incoming payload ingestion, drives background scheduler ticks, and recycles physical block handles immediately upon request termination.

---

Setup & Execution

1. Installation

Clone the repository and install dependencies inside a virtual environment:

python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt

2. Running the Server

Launch the asynchronous serving infrastructure on localhost:

python -m uvicorn src.server:app --host 0.0.0.0 --port 8000

Test the engine's active generation capability via curl:

curl -X POST "http://localhost:8000/generate" \
-H "Content-Type: application/json" \
-d '{"prompt": "What is the capital of France?", "max_tokens": 20}'

3. Running the Benchmark Suite

To regenerate the performance profiling line graph:

PYTHONPATH=src python benchmark/benchmark.py

---

Roadmap

• Prefix/prompt caching for shared system prompts across requests
• Speculative decoding with a lightweight draft model
• Dynamic page eviction policies under memory pressure
• Multi-model serving support

---

License

MIT