Tokamak: A Minimalist LLM Inference Engine Built from Scratch

Original Author/Maintainer: WeichunK Source Platform: GitHub Original Link: https://github.com/WeichunK/tokamak Publication Time: 2026-07-12T18:21:40Z

Tokamak is a lightweight LLM inference engine implemented from scratch using PyTorch. It fully covers core optimization techniques including KV caching, paged attention, continuous batching, and speculative decoding, and has been benchmarked against vLLM for performance. The project uses minimal code to demonstrate the underlying principles of the inference mechanism, making it suitable for scenarios like learning and prototype validation.

LLM inference optimization is an active area in the AI infrastructure field. The expansion of model scales brings challenges to the efficient utilization of hardware resources. Existing frameworks (such as vLLM and TensorRT-LLM) have complex code and steep learning curves. Tokamak was created to showcase core technologies through a minimalist implementation. Its name is derived from the nuclear fusion device 'Tokamak', symbolizing the controlled release of the power of large models.

1. KV Caching: Caches the Key/Value vectors of previous tokens to avoid redundant computations, reducing the time complexity of generating new tokens from O(n²) to O(n).
2. Paged Attention: Inspired by vLLM's technology, it divides KV caches into fixed-size pages, similar to virtual memory management, improving GPU memory utilization and supporting more concurrent requests.
3. Continuous Batching: Dynamically adds/removes requests to maximize GPU utilization without waiting for the entire batch to complete.
4. Speculative Decoding: Uses a lightweight draft model to generate candidate tokens, which are then verified in parallel by the large model. Multiple tokens are accepted at once, increasing throughput.

Tokamak provides performance benchmarking comparisons with vLLM, evaluating metrics such as throughput, latency, and memory usage under the same hardware and model parameters. This open and transparent comparison helps users understand actual performance, provides data support for framework optimization, and facilitates researchers to reproduce industrial-grade inference functions.

• Uses the PyTorch framework, which offers high code readability and ease of modification/debugging, making it suitable for algorithm experiments and teaching.
• Has a clear code structure, with core components (KV cache management, attention calculation, etc.) encapsulated in modules that can be studied individually or used in combination.
• Minimalist design with a codebase much smaller than production-grade frameworks, making it an ideal starting point for learning LLM inference optimization.

• Teaching and Learning: Provides a concise and complete reference implementation for students and researchers.
• Algorithm Prototype Validation: Enables rapid experimentation with new inference optimization algorithms, which can then be migrated to production frameworks after validation.
• Edge Deployment: Its lightweight nature makes it suitable for deploying small models on resource-constrained edge devices.
• Customization Needs: Serves as a foundation for secondary development to meet specific needs not covered by existing frameworks.

Tokamak implements core modern LLM inference technologies with minimal code, providing learning resources and an experimental platform for the community. It demonstrates that lightweight implementations can cover key optimization techniques, lowering the entry barrier to the field of inference optimization. As large model technology evolves, there is still room for innovation in inference optimization, and Tokamak helps more people participate in the research and improvement of underlying technologies.