tiny-llm: A Lightweight Transformer Inference Engine Implemented in Pure CUDA C++

The open-source project tiny-llm provides a high-performance Transformer inference engine built from scratch, implemented in pure CUDA C++. It supports W8A16 quantization, KV cache management, and optimized kernels, aiming to address the needs for lightweight and controllable inference deployment on edge devices.

Large language model (LLM) inference deployment faces a dilemma: cloud APIs have high latency, significant privacy risks, and uncontrollable costs; while mainstream local solutions like llama.cpp and vLLM are too heavy for resource-constrained edge devices. Scenarios such as embedded systems, mobile devices, and IoT gateways require more lightweight and controllable inference solutions that balance performance and resource overhead.

Design Philosophy: Implemented in pure CUDA C++ with no framework dependencies, achieving extreme lightweight (size at MB level), full controllability (fine-grained optimization), and transparent performance characteristics (easy to tune). Core Technologies:

• W8A16 quantization: 8-bit weights and 16-bit activations, using symmetric quantization + per-channel scaling, improving inference speed by 40-60% with perplexity loss ≤2%;
• KV cache management: Pre-allocated continuous memory, tensor reuse, dynamic expansion, reducing generation overhead;
• Optimized CUDA kernels: FlashAttention-style attention computation, fused operators, vectorized access, Warp-level optimization.

Modular Layer Design: Bottom-layer CUDA kernel library, middle-layer computation graph engine, upper-layer model definition layer, facilitating the extension of new models; Memory Pool Management: Pre-allocate memory blocks before inference, reuse via offsets, reducing cudaMalloc/cudaFree overhead; Asynchronous Execution Pipeline: Parallel CPU preprocessing and GPU computation, using CUDA stream/event collaboration to maximize hardware utilization.

On NVIDIA Jetson AGX Orin, compared to an equivalent PyTorch implementation, tiny-llm reduces the inference latency of 7B models by 3-5x and memory usage by over 60%; the advantage is more significant in batch scenarios. Its performance is similar to the llama.cpp CUDA backend, but it offers more flexibility in custom operators and extensibility.

Suitable scenarios:

• Embedded AI devices (drones, robots, smart cameras): Offline local LLM;
• Edge gateways: Low-latency AI services, protecting data privacy;
• Mobile applications: Local AI capabilities, avoiding network dependency;
• Research and teaching: Clear code structure, facilitating understanding of the underlying mechanisms of Transformer inference.

Limitations: Currently only supports Decoder-only Transformers; complex structures like MoE are to be developed; only supports NVIDIA GPUs. Future Plans: Introduce 4-bit quantization, support multi-GPU parallelism, add AMD ROCm support, and develop model conversion tools (import from HuggingFace).

By returning to the basics and simplifying dependencies, tiny-llm achieves excellent performance in edge deployment scenarios, providing a lightweight and controllable technical solution for running LLMs in resource-constrained environments, which is worth developers' attention.