# Build A Reasoning Model Scratch 2: Loading a Base Model, Text Generation, and KV Caching

## Executive summary

This session provides an in-depth, hands-on guide to working with Large Language Models (LLMs) from scratch using PyTorch. The process covers loading pre-trained models (e.g., Qwen3), understanding the tokenization and autoregressive text generation process, and critically, implementing advanced optimization techniques. Key focus areas include utilizing KV caching to drastically reduce inference latency and applying `torch.compile` for graph optimization, which are essential skills for deploying high-performance LLM services.

## Key takeaways

- LLMs Generate Text Autoregressively: Text generation is not a single step; the model generates one token at a time (autoregressive process). The output of each step is appended to the input context for the next iteration. This iterative nature makes text generation computationally expensive.
- Greedy Decoding and Token Selection: The basic method for selecting the next token involves Greedy Decoding, which selects the token with the highest score (using `torch.argmax`) from the model's output logits. The process relies on the tokenizer to convert these IDs back into readable text.
- KV Caching for Inference Optimization: To improve performance, KV caching stores intermediate Key and Value tensors computed during attention mechanisms. Instead of recomputing these values in every step, they are retrieved from the cache, significantly reducing computational overhead and improving throughput (e.g., from 4 to 28 tokens/second).
- Model Compilation with `torch.compile`: The `torch.compile()` feature optimizes the PyTorch computation graph by fusing operations, reducing overhead and improving execution speed. This is a powerful optimization technique for deployment but requires careful handling due to potential compatibility issues.

## Technical details

- Model Loading & Architecture: The session uses the Qwen3 model (a 6 billion parameter model) and the `reasoning-from-scratch` library for loading pre-trained weights. The process involves initializing the architecture, loading weights from disk, and setting up the device (CPU/MPS/CUDA).
- Tokenization: Text is broken down into subword tokens via a tokenizer's `encode` method (text to token IDs) and then converted back using the `decode` method (token IDs to text). The process requires converting Python lists of IDs into PyTorch tensors.
- Inference Mode: When running inference, it is crucial to wrap the model call in `torch.no_grad()` (or setting the model to evaluation mode) to prevent PyTorch from building a computation graph and consuming excessive memory.
- Text Generation Loop: The core generation function iteratively feeds the extended input tensor back into the model, extracting only the last generated token ID at each step. The process must handle end-of-sequence tokens (e.g., `end_tok` ID) to gracefully stop output.
- Performance Benchmarking: Benchmarking involves measuring total elapsed time and calculating the throughput (tokens per second). Optimization is achieved by comparing naive generation vs. KV caching vs. compilation, demonstrating significant speedups.
- Device Compatibility: The code attempts to automatically select the optimal device: CUDA GPU > MPS (Apple Silicon) > Intel GPU > CPU. The use of `torch.cuda` or `mps` must be managed carefully for maximum compatibility.

## Practical implications

- Understanding the difference between training and inference modes is critical for optimizing resource usage in production LLM services.
- Implementing KV caching is a mandatory optimization step when deploying LLMs to reduce latency and increase throughput, especially for long-context generation.
- Knowing how to use `torch.compile` allows engineers to significantly boost the performance of PyTorch models by fusing operations into optimized kernels.
- The ability to benchmark model speed (tokens/second) across different hardware (CPU vs. GPU vs. MPS) is vital for cost and performance planning in cloud deployments.

## Topics

Large Language Models (LLMs), PyTorch, Tokenization, Autoregressive Generation, KV Caching, Model Compilation, Inference Optimization, Reasoning Book, LLMs from Scratch book, Reasoning GitHub repo, KV cache from scratch article

Source: https://www.youtube.com/watch?v=BJua0yjO5dk