warpllama

From-scratch Llama inference engine with hand-written Triton & CUDA kernels.

warpllama runs Llama 3.2 on a single GPU. It reimplements the model from scratch in Python and replaces the memory- and compute-bound ops with hand-written Triton kernels (and CUDA for the attention hot path) to push decode throughput up and memory down: a fast, lean, readable engine you can run, benchmark, and build on.

Every optimization is measured, never assumed: profile → reason about the roofline → fuse → measure → repeat. Nothing lands without a number in bench/results/.

Target: TBD tok/s decode @ batch 1 on A100 80GB, N× faster than HF transformers and within Y% of gpt-fast. (Filled in as milestones land.)

See SPEC.md for the full scope and docs/writeup.md for the roofline analysis behind each speedup.

Results

Every benchmark is captured as JSON in bench/results/ with the GPU, library versions, and commit hash, then rendered here by make table.

Milestone	decode tok/s (bs=1)	TTFT (ms)	peak mem	vs M0	Δ perplexity
HF baseline	104.92	11.43	2.52 GiB	1.0×	ref
from-scratch forward	-	-	-	-	-
static KV cache	-	-	-	-	-
fused Triton norm/rope/mlp	-	-	-	-	-
fused attention (hero)	-	-	-	-	-
int4 weight-only quant	-	-	-	-	-
CUDA graphs	-	-	-	-	-

Current numbers are the dev baseline on an A40 (unsloth/Llama-3.2-1B, bf16) - the GPU that's reliably available on RunPod. The M0→M6 progression is measured on this same A40 so the speedups are apples-to-apples; a final headline run on A100/H100 with the official meta-llama/Llama-3.2-1B weights lands the "vs gpt-fast" number. M0 sits at ~37% of the A40 weight-bandwidth roofline (~282 tok/s) - that gap is the budget.

Quickstart

pip install -e .
# Baseline + harness (M0)
make bench MODEL=meta-llama/Llama-3.2-1B
# Correctness (logits match + perplexity + golden output)
make correctness
# Render the results table from bench/results/*.json
make table

Milestones

• HF baseline + benchmark/correctness harness
• from-scratch Llama forward (PyTorch), logits match HF
• static KV cache + single-token decode
• fused Triton kernels: RMSNorm, RoPE, SwiGLU
• fused attention (Triton → hand-written CUDA hero kernel)
• int8 → int4 weight-only quantization (W4A16 dequant-fused matmul)
• CUDA graphs (+ speculative decoding, stretch)

Dev environment

Built to run on RunPod. Use a CUDA -devel image so nvcc is available to compile the CUDA kernel, and a persistent Network Volume for the model weights. See infra/ and SPEC.md.