LLM Inference Optimization & GPU Kernel Engineering

A deep-dive repository into the core systems mechanics of modern LLM serving engines (such as vLLM and TensorRT-LLM). This project implements foundational inference primitives from scratch—ranging from KV-cache state management and multi-timeline speculative decoding to custom JIT-compiled Triton GPU kernels.

---

Repository Structure

├── speculative_decoding/     # Custom Speculative Decoding Engine with KV-Cache Threading
└── triton_kernels/           # Custom JIT-Compiled Triton Kernels (Tiled Online Softmax)

---

Part 1: Speculative Decoding Engine with Multi-Timeline KV-Cache

Standard autoregressive generation incurs heavy $O(N^2)$ prefill penalties if sequences are reprocessed from scratch. This engine implements a complete draft-verify speculative loop between a 0.5B draft model and a 4-bit quantized 1.5B target model.

Key Engineering Features:

• KV-Cache Threading & Delta Routing: Tracks distinct draft and target caches, feeding only the token delta (current_ids[:, kv_len:]) to eliminate redundant context processing.
• Alternate Timeline Pruning (crop_cache): Physically iterates through layer tensors to truncate memory states back to the exact rejection point.
• Rigorous Benchmarking: Employs median-of-N filtering to isolate true hardware performance from OS/GPU timing noise.

Performance Results (Qwen-2.5 0.5B -> 1.5B-Instruct, RTX 3050):

• Sweet Spot ($k=2$): Achieved a 69.1% acceptance rate resulting in a 1.36x throughput speedup over baseline.
• The Decay Curve ($k=8$): Acceptance dropped to 39.1% due to draft-target distribution drift, highlighting the trade-off of higher lookahead horizons.

---

Part 2: Fused & Tiled Triton Softmax Kernel

Standard PyTorch operations are often memory-bandwidth bound, launching multiple separate kernels that thrash GPU high-bandwidth memory (HBM). This module explores custom GPU programming via OpenAI Triton to fuse operations and handle wide tensor constraints.

Key Engineering Features:

• Naive vs. Tiled Architecture: Analyzed why naive single-block row kernels hit hardware limits and suffer from register spilling at massive widths ($w = 65,536$).
• Online Softmax Algorithm: Implemented a two-pass chunked streaming kernel featuring running max and running sum with dynamic scale corrections (the core math behind FlashAttention).
• Correctness Validation: Verified using strict floating-point comparisons (torch.allclose), achieving max differences down to $\approx 5.8e-11$ across shapes up to $(2, 65536)$.

Performance Results (Tensor Width = 65,536):

• Naive Triton: OOM / Slow due to block resource exhaustion.
• Tiled Triton vs. PyTorch: Outperformed native PyTorch execution, achieving a 1.30x speedup (28.884 ms vs 37.408 ms) by maintaining safe on-chip memory utilization.

---

Getting Started

Prerequisites & Installation

Ensure you have an NVIDIA GPU with CUDA drivers installed, then set up your environment:

python -m venv .venv
.venv\Scripts\Activate
pip install torch transformers accelerate bitsandbytes triton-windows

Running the Speculative Decoding Benchmark

cd speculative_decoding
python benchmark.py

Running the Triton Softmax Benchmark

cd triton_kernels
python benchmark_softmax.py

```"# inference-engine-from-scratch"