Metadata-Version: 2.5
Name: liteinfer
Version: 0.6.0
Summary: A lightweight, hackable LLM inference engine built from scratch.
Project-URL: Repository, https://github.com/ValeGian/liteinfer
Author: Valerio
License: Apache-2.0
Keywords: cuda,inference,llm,pytorch,transformers,vllm
Classifier: Development Status :: 2 - Pre-Alpha
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: accelerate
Requires-Dist: huggingface-hub>=0.24
Requires-Dist: numpy>=1.26
Requires-Dist: safetensors>=0.4
Requires-Dist: torch>=2.4
Requires-Dist: transformers>=4.45
Provides-Extra: bench
Requires-Dist: vllm==0.28.0; extra == 'bench'
Provides-Extra: dev
Requires-Dist: pyright>=1.1.380; extra == 'dev'
Requires-Dist: pytest-cov>=5.0; extra == 'dev'
Requires-Dist: pytest-xdist>=3.5; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Requires-Dist: ruff>=0.6; extra == 'dev'
Description-Content-Type: text/markdown

# liteinfer

[![PyPI](https://img.shields.io/pypi/v/liteinfer)](https://pypi.org/project/liteinfer/)

A lightweight, hackable LLM inference engine built from scratch — designed
to make state-of-the-art inference techniques (paged KV cache, prefix
caching, tensor parallelism, `torch.compile`, CUDA graphs, …) easy to read,
test, and benchmark.

**New to inference serving?** [**Inside liteinfer**](https://claude.ai/code/artifact/52a9e43f-c529-4a70-af07-d55dbffb1cbf)
is an illustrated walkthrough of how the engine works and why each design was
chosen — prefill and decode, KV caching, paged memory, and continuous batching —
following the measured progression that produced this codebase, including the
steps that lost. It assumes PyTorch, not serving experience.

## Goals

1. **Fast offline inference** — throughput in the same league as vLLM on a single node.
2. **Readable codebase** — clean, minimal, well-structured. The core engine should fit in your head.
3. **Optimization suite** — a clear place for each technique (prefix caching, TP, `torch.compile`, CUDA graphs, …) with isolated, testable implementations.
4. **HuggingFace compatibility** — load any compatible HF model from the Hub or a local safetensors directory.

## Status

v0 — greedy/sampled inference on local safetensors via continuous batching over
a paged KV cache. `AsyncLLM` is the native interface (async context manager plus
streaming); `LLM` is a synchronous facade for offline batch use. See
[`docs/milestones.md`](docs/milestones.md) for what is in, and
[`docs/roadmap.md`](docs/roadmap.md) for what is queued.

## Installation

```bash
pip install liteinfer
```

For development or benchmark comparisons:

```bash
git clone https://github.com/ValeGian/liteinfer.git
cd liteinfer
python -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"
```

## Quick start

```python
from liteinfer import LLM, SamplingParams

llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct")
params = SamplingParams(temperature=0.8, max_tokens=128)

outputs = llm.generate(["Explain paged attention in one paragraph."], params)
print(outputs[0].text)
```

```python
# Async continuous batching with streaming
import asyncio
from liteinfer import AsyncLLM, SamplingParams

async def main():
    async with AsyncLLM("meta-llama/Llama-3.2-1B-Instruct") as llm:
        # Batch — all requests processed concurrently
        results = await llm.generate(prompts, SamplingParams(max_tokens=128))

        # Streaming — per-token events
        async for event in llm.stream("Explain paged attention.", SamplingParams(max_tokens=128)):
            print(event.text, end="\r")

asyncio.run(main())
```

The attention kernel is chosen for the device unless you name
one — on CUDA that is `paged`, a Triton kernel that reads the KV pool in place
instead of gathering it, so the decode step stops growing with context. Pass
`attn_implementation="sdpa"` or `"eager"` to override; naming a kernel that
cannot run here is an error rather than a silent downgrade.

## Repository layout

```
liteinfer/
├── liteinfer/             # Library source
│   ├── llm.py             # Synchronous LLM facade
│   ├── async_llm.py       # AsyncLLM: the engine's native interface
│   ├── config.py          # EngineConfig
│   ├── engine/            # Orchestration: scheduler, sequence, model runner, metrics
│   ├── models/            # Loader + per-architecture implementations (layers included)
│   ├── cache/             # Paged KV cache: block pool + slot addressing
│   └── sampling/          # SamplingParams + Sampler
├── tests/                 # Unit / integration / e2e tests
├── benchmarks/            # Benchmark matrix, harness, and report
└── pyproject.toml
```

Each module's `__init__.py` documents the contract it owns.

## Architecture (brief)

One engine: continuous batching over a paged KV cache.

- **`AsyncLLM`** — asyncio API: `await llm.generate(...)`, or `async for` to
  stream tokens. **`LLM`** is a synchronous facade over it for offline batch use.
- **`ContinuousScheduler`** — fills empty batch slots from the waiting queue on
  every step and evicts finished sequences individually.
- **`ContinuousModelRunner`** — runs one prefill or decode forward pass.
  `torch.compile`, CUDA graph capture and tensor parallelism plug in here.
- **`ContinuousKVCache`** — per-sequence blocks drawn from a shared `BlockPool`;
  `slot_table` maps logical token positions to physical slots, so a whole batch
  is read or written with a single indexing op.
- **`models/attention.py`** — the attention kernel, one per
  `attn_implementation`. `sdpa` (default) never materialises the score matrix;
  `eager` writes it out in plain matmuls, which reads better and is the parity
  reference, but caps the prompt length that fits in memory. `paged` is the fast
  decode path: a Triton kernel that reads the KV pool through the slot table
  instead of gathering it, so the decode step stops growing with context. The
  engine picks between them from the device.

Sampling is a separate stage so strategies (greedy, top-p, …) can be swapped
without touching the engine. `stats` records a `StepMetrics` per forward pass.

Each of these components is explained, with diagrams and interactive figures, in
[Inside liteinfer](https://claude.ai/code/artifact/52a9e43f-c529-4a70-af07-d55dbffb1cbf).

Earlier designs — no cache, `DynamicCache`, plain-tensor cache, static
batching — were measured against each other and then removed; the numbers live
in [`docs/benchmarks.md`](docs/benchmarks.md) and the reasoning in
[`docs/milestones.md`](docs/milestones.md).

## Testing

`liteinfer` is **test-first**: every feature ships with the tests that
pin its contract.

```bash
pytest                              # full suite — runs sequentially, GPU-safe
pytest -m "not gpu and not slow"    # fast suite (no model downloads, no GPU)
pytest -m gpu                       # GPU tests only — sequential, never use -n auto
pytest -n auto                      # parallel mode — CPU-only tests only
pytest tests/unit/                  # one directory
```

> **GPU tests**: always run sequentially. The e2e tests (`tests/e2e/`) load
> real models onto GPU. Running them in parallel (`-n auto`) will cause OOM or
> cross-process interference.

Test layout:

- `tests/unit/` — single-component tests. CPU-only, fast, no model loading.
- `tests/integration/` — multiple components wired together (still no HF download).
- `tests/e2e/` — load a small real model and verify generation against `transformers`.

See `tests/README.md` for conventions.

## Performance

A40 · Llama-3.2-1B-Instruct · ISL 128 · OSL 256 · vs vLLM 0.28.0 at matched batch width.

| | liteinfer | vLLM | |
|---|---:|---:|---|
| Throughput, B=32 | **3,374 tok/s** | 4,397 tok/s | 1.3× behind |
| Throughput, ISL 1024 / OSL 1024 | **2,535 tok/s** | 3,243 tok/s | 1.3× behind |
| Decode step (ITL p50) | **6.5 ms** | 5.2 ms | 1.3× behind |
| Time to first token (p50) | 14.7 ms | 22.4 ms | 1.5× ahead\* |
| Long prompts (ISL 2048) | runs | runs | eager attention cannot |

\* At this prompt length TTFT is mostly fixed per-call API overhead rather than
prefill compute, so read it as offline round-trip latency, not prefill speed.

Rows 1-4 are ISL 128 / OSL 256 unless stated, on the kernel the engine picks
for an A40 — `paged`, the Triton decode kernel — with the decode forward
replayed from a CUDA graph. Where either precondition does not hold the engine
falls back rather than failing.

**This cycle was the launch overhead, and it was most of the gap.** A decode
step issued ~700 kernels to do ~7 ms of GPU work, at a flat 18.7 µs of wall per
launch whatever the batch width — so the GPU sat idle 41-52% of every step and
32× the tokens bought 1.29× the GPU time. Replaying the forward from a captured
graph took the decode step **13.3 → 6.5 ms (2.0×)** and throughput **1,997 →
3,374 tok/s (1.7×)**, closing the gap to vLLM from 2.6× to 1.3× on the step. What
remains is arithmetic rather than overhead: liteinfer's kernels are 1.9× off the
memory roofline where vLLM's whole step is 1.5× off, and that is fusion —
[`docs/roadmap.md`](docs/roadmap.md) carries the numbers, including three
optimisations that were built, measured and reverted.

**Prompts of different lengths no longer pay for each other.** A prefill batch
was padded to its longest prompt — on real prompt lengths, a batch of 32
computed **13.4× the positions it kept**, and the mask that hid them also ruled
FlashAttention out. Packed end to end with `cu_seqlens`, the same mixed-length
workload runs **2,009 → 3,337 tok/s (1.7×)**, and **4.8×** when outputs are short
enough that prefill is most of the work. On prompts that are all the same length
it is 1.00×, which is why no number in this file moved until the benchmark
grew a mixed-length shape to measure it on.

**One of those reverts came back.** Split-K decode lost by 6% in #34 because the
extra kernel launch it adds per layer cost more than the kernel saved, on a step
that was then 47% host time. A captured step is 97.5% GPU, so the same kernel now
pays: at one request and a 15,360-token context the decode step is **12.0 → 7.5 ms
(1.6×)**, and at ISL 3584 it is 7.8 → 6.7 ms. The win grows with context because a
single request's attention runs on 8 programs of an 84-SM device and that is the
part of the step context lengthens — 2% of a short step, 65% of a 32k one. At
ISL 128 it is 1.00×, and from 12 concurrent sequences up the batch already fills
the device and the kernel is the unsplit one.
Throughput figures are certified against `vllm bench throughput` to within 1%. Full tables, methodology, and per-milestone deltas:
[`docs/benchmarks.md`](docs/benchmarks.md) · [live dashboard](https://valegian.github.io/liteinfer/).

## Benchmarking

Every liteinfer configuration is measured against the one it improves on, and
against vLLM, on byte-identical prompts.

```bash
pip install -e ".[dev,bench]"

# 1. Build a dataset (once per model + ISL/OSL)
bench dataset --model meta-llama/Llama-3.2-1B-Instruct --isl 128 --osl 256 -n 200

DS=benchmarks/datasets/isl128_osl256_n200_meta_llama_llama_3_2_1b_instruct.json

# 2. Run every config, in both modes
bench run --all --dataset "$DS" --mode throughput -n 200
bench run --all --dataset "$DS" --mode latency    -n 50

# 3. Report
bench report --out docs/index.html
```

Engines: `liteinfer`, `vllm`. See [`benchmarks/README.md`](benchmarks/README.md)
for the config matrix and controls, and [`docs/benchmarks.md`](docs/benchmarks.md)
for methodology and results.

Every performance change is measured against the config it improves on. A
general win replaces that config; a win scoped to a precondition (MoE,
quantization, long context) joins it instead — see "Shipping an improvement" in
[`docs/roadmap.md`](docs/roadmap.md).

## Roadmap

High-level direction:

- [x] Greedy/sampled generation from local safetensors
- [x] Paged KV cache
- [x] Continuous batching (async, streaming) — superseded static batching, now removed
- [ ] prefix caching
- [ ] `torch.compile` and CUDA graphs for decode
- [ ] Tensor parallelism (single node)
- [ ] Speculative decoding

Detailed, fine-grained backlog with scope and parity-test notes lives
in [`docs/roadmap.md`](docs/roadmap.md). Achieved milestones are
tracked separately in [`docs/milestones.md`](docs/milestones.md).

## License

Apache-2.0.
