Building an LLM Inference Engine from Scratch: Deep Dive into Transformer Inference and System Optimization

Project Overview

This open-source project is developed by ashwinvijayakumar24 (GitHub repo: llm_inference_engine, Release date: June 5, 2026). It aims to build an LLM inference engine from scratch and deeply解析 the core principles of production-grade inference systems. The project covers key technologies including Transformer forward propagation implementation, KV cache mechanism, continuous batching, PagedAttention, and CUDA kernel optimization. Its goal is to achieve efficient inference for the Llama3.2 1B model on NVIDIA A100/H100/H200 GPUs, and conduct benchmark comparisons with HuggingFace Transformers and llama.cpp, providing a practical guide for developers.

Background and Motivation

In current LLM development, most developers rely on model.generate() or mature frameworks (such as vLLM, llama.cpp) for inference, but lack in-depth understanding of underlying system details (Transformer forward propagation, KV cache management, batch scheduling, CUDA optimization). This project helps developers master the core principles of production-grade inference by building a complete inference engine from scratch. The project uses Python for high-level orchestration and CUDA C++ for writing high-performance kernels, focusing on efficient inference of the Llama3.2 1B model on NVIDIA GPUs.

Tech Stack and Architecture

• Allowed underlying libraries: NumPy/PyTorch tensor storage, cuBLAS (GEMM operations), HuggingFace Tokenizer (Rust backend), safetensors weight parsing.
• Self-implemented components: Transformer layer logic (RoPE positional encoding, GQA grouped query attention, RMSNorm normalization, SwiGLU feedforward network), KV cache management (naive continuous buffer → paged block management), continuous batching scheduler, quantization paths (int8/int4), custom CUDA attention kernel.
• Development environment: Apple M4 MacBook Air (correctness verification, Python engine development); NVIDIA PACE cluster (CUDA kernel development, performance benchmarking using A100/H100/H200 GPUs and Slurm scheduling).

Core Implementation Details

1. Transformer forward propagation: Manually implement the complete process (word embedding → multi-layer Transformer → LM head → sampling). Key details include RoPE positional encoding, GQA grouped query attention, RMSNorm normalization, and SwiGLU feedforward network.
2. KV cache management: Evolve from naive continuous buffer (pre-allocated maximum length, memory waste) to paged block management (PagedAttention style, eliminate fragmentation, support variable-length sequences).
3. Continuous batching: Allow new requests to join and completed requests to leave the current batch, improving GPU utilization; it is recommended to first implement a single-thread scheduler to verify correctness, then perform parallel optimization.
4. Quantization optimization: Plan to implement int8 (halve memory) and int4 (reduce memory to 1/4) weight quantization. Evaluation metrics include memory usage, perplexity, and throughput.
5. CUDA attention kernel: Development process: NumPy reference → CUDA code → Python bindings → performance profiling → compare with PyTorch baseline; focus on optimizing memory access patterns for the decoding phase (sequence length 1).

Development Plan and Benchmarking

• Phased development:
  1. Environment setup (CMake+nanobind, weight parsing, Tokenizer verification);
  2. Correctness verification (NumPy implementation of forward propagation, compare hidden states with HF);
  3. Usability improvement (naive KV cache, sampling strategy, CLI/HTTP endpoints);
  4. Baseline benchmark (performance test framework, compare with HF/llama.cpp);
  5. Differentiated optimization (quantization → continuous batching → paged KV → CUDA kernel → speculative decoding);
  6. Project polishing (documentation improvement, performance charts).
• Benchmark metrics: Throughput (prefill/decode), TTFT (time to first token), ITL (inter-token latency), peak memory usage, throughput vs batch size, memory vs context length; comparison targets are HuggingFace Transformers and llama.cpp.

Potential Risks and Mitigation

1. Logits mismatch with HF: Compare hidden states layer by layer,排查 RoPE parameters, RMSNorm epsilon position, attention scaling factors, etc.
2. PACE Slurm queue waiting: Complete Python/correctness work on Mac, only use PACE for CUDA development.
3. CUDA kernel correctness: First use the NumPy reference version to verify output, then perform performance tests.
4. Scope creep: Focus on 2-3 high-quality features, avoid half-finished products.

Project Value and Future Directions

• Value: Educational significance (understand the inner mechanism of inference), engineering practice (modular system design), interview preparation (resume highlight), full-stack technical depth (from algorithm to hardware).
• Future directions: Multi-GPU/tensor parallel inference, Flash Attention fused kernel, LoRA adapter hot-swapping, complete speculative decoding implementation.