2 Generating text with a pre-trained LLM

This chapter establishes the practical foundation for building a reasoning model by starting from a pre-trained base LLM rather than an instruction-tuned assistant. It explains that using a base model makes it easier to isolate the effects of later reasoning methods, and it walks through the setup needed to follow along: installing the required Python packages, choosing between a pip-based workflow and the faster, more convenient uv-based workflow, and selecting suitable hardware. The chapter emphasizes that although the model used is relatively small, it is still fully capable of generating coherent text, and the examples are designed to run reasonably well even on a CPU.

A major focus is preparing text for the model and understanding how generation works at the token level. The reader loads a tokenizer, performs an encode-decode round trip to see how input text becomes token IDs and then back into readable text, and observes how words may be split into subword pieces. After that, the chapter loads the Qwen3 0.6B base model, describes its main architectural components at a high level, and shows that the model can already produce sensible next-token predictions. The core idea of autoregressive generation is introduced: the model processes the whole prompt, produces scores for the next token, and the final prediction is selected and appended one token at a time.

The chapter then turns that idea into working code by implementing a minimal text-generation loop that repeatedly feeds the growing sequence back into the model, decodes the chosen token, and stops on an end-of-sequence token when appropriate. To make the process measurable, it adds simple runtime and memory reporting, showing that naive generation is slow but functional. Finally, it introduces two important speedups: KV caching, which reuses attention keys and values instead of recomputing them each step, and PyTorch model compilation, which reduces overhead by optimizing the computation graph. Together, these techniques deliver large performance gains while preserving the same output, and they prepare the reader for more efficient evaluation and reasoning workflows in later chapters.

A mental model depicting the four main stages of developing a reasoning model. This chapter focuses on stage 1, loading a conventional LLM and implementing the text generation functionality.

An overview depicting an LLM generating a response (output text) given a user query (input text)

Installing and using the uv Python package and project manager via the macOS terminal

An overview of the Lightning AI GPU cloud platform in a web browser. The interface supports Python scripts, Jupyter notebooks, terminal access, and lets users switch between CPU and various GPU types based on their compute needs.

A simplified illustration of how an LLM receives input data and generates output. The user-provided text is tokenized into IDs using the tokenizer's encode method, which are then processed by the LLM to generate output token IDs. These are decoded back into human-readable text using the tokenizer's decode method.

A demonstration of the round-trip tokenization process using a tokenizer. The user-provided input text is first converted into token IDs using the encode method, and then accurately reconstructed back into the original text using the decode method.

An overview of the four key stages in developing a reasoning model in this book. This section focuses on loading pre-trained LLM in Stage 1.

Overview of the Qwen3 0.6B model architecture, as instantiated earlier. Input text is tokenized and passed through an embedding layer, followed by 28 repeated transformer blocks. Each block contains grouped-query attention, feedforward layers, and RMS normalization. The model ends with a final normalization and linear output layer. Arrows show the data flow through the model.

An overview of the four key stages in developing a reasoning model in this book. This section explains the main concept behind text generation in LLMs, which allows us to implement a text generation function for using the pre-trained LLM in the remainder of this chapter.

An illustration of the sequential (autoregressive) text generation in LLMs. At each iteration, the model generates the next token based on the input and previously generated tokens, which are cumulatively fed back into the model to produce coherent output.

The LLM processes the full input sequence and produces one next-token prediction per input position. These predictions are conceptually shifted one step ahead, because each position predicts the following token. The illustration is position-aligned for intuition, but the outputs are not literally the input prompt shifted to the right. During generation, we keep only the final prediction and use it to continue the input prompt one token at a time.

A closer look at how the raw scores output by an LLM, in a single text generation iteration, are converted into a token ID and its corresponding text representation.

An overview of the four key stages in developing a reasoning model in this book. In this section we implement a text generation function for the pre-trained LLM.

An illustration of sequential (autoregressive) text generation in large language models (LLMs), with token IDs shown explicitly. At each iteration, the model generates the next token based on the original input and all previously generated tokens. The predicted token is added to the sequence in both its textual and token ID form.

An overview of the four key stages in developing a reasoning model in this book. This section builds on pre-trained LLM and the basic text generation function we coded earlier and applies KV caching to speed up execution.

Illustration of how a KV cache improves efficiency during autoregressive text generation. Instead of reprocessing the entire input sequence at each step, the KV cache stores intermediate representations so that the LLM can reuse them to generate the next token. This eliminates the need to concatenate the generated token with prior inputs in each subsequent iteration.

An overview of the four key stages in developing a reasoning model in this book. This section builds on pre-trained LLM and the basic text generation function we coded earlier, including KV caching, and adds model compilation to speed up the execution speed even further.

• Using LLMs to generate text involves multiple key steps:
• Setting up the coding environment to run LLM code and install necessary dependencies.
• Loading a pre-trained base LLM (such as Qwen3 0.6B), which will be extended with reasoning capabilities in later chapters.
• Initializing and using a tokenizer, which converts text input into token IDs and decodes output back to human-readable form.
• Text generation in LLMs follows a sequential (autoregressive) process, where the model generates one token at a time by predicting the next most likely token.
• The speed and efficiency of text generation can be improved through:
• KV caching, which stores intermediate states to avoid recomputing previously encountered input tokens at each step.
• Model compilation using torch.compile, which optimizes runtime performance.
• This chapter lays the technical foundation for reasoning capabilities in upcoming chapters by implementing a functional, efficient text generation pipeline using a pre-trained base LLM.

FAQ