vLLM-Style Llama Inference Engine with Tensor Parallelism

A simple, offline Llama inference engine built to show how the model and the runtime fit together. It implements Llama 3.2, continuous batching, paged KV caching, custom attention kernels, and tensor parallelism without using the vLLM runtime.

Inference engine design (cropped):

Llama 3.2 with tensor parallelism details (cropped):

Full excalidraws at assets/design.excalidraw and assets/tensor-parallel-details.excalidraw.

Features and scope

This is an educational engine, not a replacement for vLLM.

Implemented

• Llama 3.2 1B and 3B with grouped-query attention, RoPE, RMSNorm, and a gated MLP.
• Hugging Face configuration, tokenization, and local safetensors loading.
• A continuous-batching scheduler with paged KV caching, block tables, slot mappings, and recompute preemption.
• PyTorch SDPA for prefill and recomputation.
• PyTorch reference and Triton kernels for paged decode attention.
• Single-node tensor parallelism with one NCCL worker per GPU.
• A synchronous LLM.generate() API with greedy decoding.

Out of scope

• An asyncio serving API for submitting requests while the engine is running, streaming results, and cancellation. The engine and scheduler support later request arrivals, but the public API only exposes synchronous offline batches.
• Models other than bias-free Llama 3.2 1B/3B checkpoints with tied embeddings, and hardware other than CUDA GPUs.
• Multi-node execution and pipeline parallelism.
• Sampling strategies other than greedy decoding.
• Chunked prefill and prefix caching.
• Quantization, speculative decoding, and CUDA graphs.
• A split-K or otherwise production-optimized Triton decode kernel.

Quick start

Set up the environment and download a Llama 3.2 checkpoint by following operations.md, then run:

from llm_inference import LLM

with LLM(
    "models/Llama-3.2-3B-Instruct",
    kv_cache_memory_bytes=64 * 1024**2,
) as llm:
    outputs = llm.generate(
        ["Paged attention is", "Continuous batching means"],
        max_new_tokens=32,
    )

for output in outputs:
    print(output.text)

Set tensor_parallel_size=2 or higher to shard the model across contiguous GPUs on one host.

How it works

1. The engine tokenizes and submits the complete offline batch.
2. The scheduler selects requests and allocates physical KV-cache blocks.
3. Each execution rank runs its local attention and MLP shards and writes to its rank-local paged KV cache.
4. Prefill and recomputation use causal SDPA and one-token decode uses a custom Triton paged attention.
5. Rank zero selects the next tokens, finished requests release their blocks, and the scheduler builds the next batch.

The scheduler can admit waiting requests while others decode, so batch membership may change on every step. Note: The public API is still offline: all requests are submitted together, with no later arrivals, streaming, or cancellation.

Tensor parallel execution

With tensor_parallel_size=1, the model runs directly in the calling process. For larger groups, the parent process keeps ownership of scheduling and spawns one worker per local GPU. It sends the same scheduled batch to every worker over multiprocessing pipes. Each worker holds its model shard and a rank-local paged KV cache. Column-parallel layers produce local output shards, while row-parallel layers use NCCL all-reductions to combine partial results. The LM head is partitioned by vocabulary: ranks send their local maximum logits and indices to rank zero, which selects each request's global maximum and returns the corresponding token IDs to the parent.

Benchmarks

On an RTX 4080, the Triton paged-attention kernel was 12.3–30.2x faster than the PyTorch reference across the tested decode workloads:

Batch	Sequence	PyTorch	Triton	Speedup
1	128	0.274 ms	0.021 ms	13.26x
1	512	0.261 ms	0.021 ms	12.60x
4	128	0.270 ms	0.021 ms	12.97x
4	512	0.340 ms	0.028 ms	12.33x
16	2,048	6.559 ms	0.217 ms	30.20x

Offline end-to-end tests used Llama 3.2 3B in BF16 on the same RTX 4080. Each cell is median whole-batch latency and includes tokenization and detokenization; model loading is excluded. vLLM 0.27.1 ran in eager mode without CUDA graphs or prefix caching.

Runtime	Batch of 1	Batch of 4	Batch of 4, longer
This engine	96.63 ms	109.63 ms	437.56 ms
vLLM 0.27.1	85.81 ms	94.32 ms	392.44 ms

Single-node tensor-parallel tests on Modal used two L4 GPUs and produced identical greedy outputs for TP=1 and TP=2. The current engine uses a vocabulary-partitioned LM head. Bold values mark the lower latency for each engine pair at the same TP size:

Runtime	Batch of 1	Batch of 4	Batch of 4, longer
This engine, TP=1	265.39 ms	265.35 ms	1,032.66 ms
vLLM 0.27.1, TP=1	219.12 ms	256.17 ms	959.61 ms
This engine, TP=2	300.95 ms	318.63 ms	1,277.39 ms
vLLM 0.27.1, TP=2	267.17 ms	339.22 ms	1,169.51 ms

At TP=2, this engine was 12.64% slower than vLLM for a single request, 6.07% faster for a batch of four, and 9.22% slower for the longer batch. Compared with its own TP=1 results, TP=2 was 13.40%, 20.08%, and 23.70% slower across the three workloads.

Local results used three warmups and the median of ten runs. The matched Modal run used one warmup and six measured samples per configuration across two balanced forward/reverse-order passes. Treat these as reproducible educational comparisons, not production claims.

Reproduce the local and Modal benchmarks with the commands in operations.md. Benchmark implementations live in benchmarks/.

Development

uv run ruff check .
uv run ruff format --check src tests benchmarks
uv run pytest

See operations.md for installation, model download, inference, testing, and benchmark instructions.

References

• Efficient Memory Management for Large Language Model Serving with PagedAttention
• FlashAttention-2
• vLLM architecture overview
• Llama 3.2 model card