Inside nanochat Part 4: Understanding Inference

Introduction

This is the fourth article in the Inside nanochat series. In the previous parts, we explored tokenization, architecture, and optimiziers. This part focuses on inference, the stage where trained language models generate text one token at a time.

Training enables the model to predict the next token, but inference is where this ability is applied to turn prompts into coherent responses. Efficient inference introduces major computational challenges. Each generated token requires a complete forward pass through the transformer, and straightforward implementations can be computationally expensive.

This article explains how nanochat achieves efficient inference through three core techniques:

• KV Cache: Removing redundant computation during autoregressive generation
• Tool Integration: Allowing models to use external functions such as calculator
• Sampling Strategies: Managing randomness and maintaining the quality of generated text

The Inference Challenge

Autoregressive language models generate text sequentially, with each token dependent on all previously generated tokens.

LLM inference example: Of course, outputs aren’t full words, they’re generally tokens. This is just a simple visualization.

A straightforward implementation performs a full forward pass through the transformer at every generation step, recalculating attention scores for all previous tokens. For a sequence of length T, this leads to O(T²) total computation, since the same attention keys and values are repeatedly recomputed.

The key to improving efficiency is KV caching, which stores these previously computed keys and values and reuses them in later steps.

KV Cache: Optimizing Inference Efficiency

A simple representation of what we need only to calcualte for the new token attention

Understanding KV Cache

In transformer attention mechanisms, each token produces three vectors: queries (Q), keys (K), and values (V). During generation:

• Queries originate from newly generated tokens
• Keys and Values are required for all tokens in the sequence (both prompt and generated)

The critical insight: once K and V are computed for a token, they remain constant. These values can be cached and reused for all subsequent generation steps.

Cache Architecture

In nanochat, the KV cache maintains a 6-dimensional tensor:

cache_shape = (L, 2, B, H, T, D)

Where:

• L = number of transformer layers
• 2 = separate storage for keys and values
• B = batch size (number of parallel sequences)
• H = number of attention heads
• T = maximum sequence length
• D = head dimension (embedding size / number of heads)

Wait… the first attention layer seems to stay mostly the same except for the new key and value pairs of the latest token. But how can we be sure that all attention layers keep the same values? After all, the outputs pass through MLP layers too???.

At first, this can be confusing. It feels like each layer should change its outputs as new tokens are added. But when we look closely at how the MLP layer works, we see that it is token independent. It applies the same operations to every token in the batch and does not mix information between tokens. It simply processes each token embedding based on its dimensions.

Therefore, if the attention outputs for previous tokens remain the same, their MLP outputs also remain the same. This means that, except for the new token, the values from earlier tokens stay identical across all layers during inference.

The complete KVCache implementation from nanochat/engine.py:

class KVCache:
    """
    Works hand-in-hand with the GPT model to maintain the KV cache.
    Note that the .pos advances automatically after the last layer of the Transformer inserts.
    """
    def __init__(self, batch_size, num_heads, seq_len, head_dim, num_layers):
        # Each of K/V is of shape (B, H, T, D) and we have one per layer of the Transformer.
        self.kv_shape = (num_layers, 2, batch_size, num_heads, seq_len, head_dim)
        self.kv_cache = None
        self.pos = 0 # current position in time in the cache
    def reset(self):
        self.pos = 0
    def get_pos(self):
        return self.pos

Lazy Initialization Strategy

An optimization approach: the cache is not allocated until the first forward pass. This design allows the cache to automatically inherit the appropriate dtype (bfloat16, float32, etc.) and device (CPU/GPU) from the model:

def insert_kv(self, layer_idx, k, v):
    # Lazy initialize the cache here because we need to know the dtype/device
    if self.kv_cache is None:
        self.kv_cache = torch.empty(self.kv_shape, dtype=k.dtype, device=k.device)

    # Insert new keys/values to the cache and return the full cache so far
    B, H, T_add, D = k.size()
    t0, t1 = self.pos, self.pos + T_add

Dynamic Memory Allocation

When generation exceeds the initial sequence length estimate, the cache dynamically expands in 1024-token increments:

# Dynamically grow the cache if needed
    if t1 > self.kv_cache.size(4):
        t_needed = t1 + 1024 # as much as we need plus buffer of 1024
        t_needed = (t_needed + 1023) & ~1023 # round up to nearest multiple of 1024
        current_shape = list(self.kv_cache.shape)
        current_shape[4] = t_needed
        self.kv_cache.resize_(current_shape)

    # Insert k, v into the cache
    self.kv_cache[layer_idx, 0, :, :, t0:t1] = k
    self.kv_cache[layer_idx, 1, :, :, t0:t1] = v

    # Return the full cached keys/values up to current position (as a view)
    key_view = self.kv_cache[layer_idx, 0, :, :, :t1]
    value_view = self.kv_cache[layer_idx, 1, :, :, :t1]

    # Increment pos after the last layer of the Transformer processes
    if layer_idx == self.kv_cache.size(0) - 1:
        self.pos = t1

    return key_view, value_view

The bitwise operation (t_needed + 1023) & ~1023 efficiently rounds up to the nearest multiple of 1024, ensuring optimal memory allocation.

Attention Mechanism Integration

The cache integrates directly into the attention mechanism in nanochat/gpt.py:

def forward(self, x, cos_sin, kv_cache):
    B, T, C = x.size()
    # Project the input to get queries, keys, and values
    q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
    k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
    v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
    # Apply Rotary Embeddings to queries and keys
    cos, sin = cos_sin
    q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
    q, k = norm(q), norm(k) # QK norm
    q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
    # Apply KV cache: insert current k,v into cache, get the full view so far
    if kv_cache is not None:
        k, v = kv_cache.insert_kv(self.layer_idx, k, v)

    Tq = q.size(2) # number of queries in this forward pass
    Tk = k.size(2) # number of keys/values in total (cached + current)

Batch Prefilling via Cache Replication

nanochat implements efficient multi-sample generation from a single prompt:

def prefill(self, other):
    """
    Prefill given another KV cache. Optionally expand along batch dim.
    This is used when we do batch 1 prefill and then want to generate
    multiple samples in parallel from there.
    """
    # 1) validate the shapes
    assert self.kv_cache is None, "Cannot prefill a non-empty KV cache"
    assert other.kv_cache is not None, "Cannot prefill with a None KV cache"

    # 2) initialize the cache
    dtype, device = other.kv_cache.dtype, other.kv_cache.device
    self.kv_cache = torch.empty(self.kv_shape, dtype=dtype, device=device)

    # 3) copy the data over
    self.kv_cache[:, :, :, :, :other.pos, :] = other.kv_cache

    # 4) update the pos
    self.pos = other.pos

This capability enables:

1. Single prompt processing with batch size 1
2. Cache cloning N times
3. Parallel generation of N different completions

This approach is essential for techniques such as best-of-N sampling and reinforcement learning from human feedback (RLHF).

Tool Integration: Extending Model Capabilities

Language models demonstrate well-documented limitations in arithmetic operations. A model may struggle with “What is 847 × 923?” despite the deterministic nature of the answer. The solution: enable tool usage.

The Calculator Implementation

nanochat implements a secure calculator for evaluating arithmetic expressions:

def use_calculator(expr):
    """Evaluate a math expression safely."""
    expr = expr.replace(",", "")
    if any([x not in "0123456789*+-/.() " for x in expr]):
        return None  # disallow non-numeric chars
    if "**" in expr:
        return None  # disallow power operator (could be expensive)
    return eval_with_timeout(expr)

def eval_with_timeout(formula, max_time=3):
    try:
        with timeout(max_time, formula):
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", SyntaxWarning)
                return eval(formula)
    except Exception as e:
        signal.alarm(0)
        return None

The 3-second timeout protection prevents malicious or computationally expensive operations from blocking the system.

Tool Use Protocol

The model learns calculator usage through designated special tokens:

Model generates: Let's calculate: <|python_start|>847*923<|python_end|>
Engine detects python block completion
Engine evaluates: 847*923 = 781681
Engine injects: <|output_start|>781681<|output_end|>
Model continues: So the answer is 781681.

This protocol establishes:

1. When to invoke tools (expressions delimited by <|python_start|> and <|python_end|>)
2. How to interpret tool outputs (delimited by <|output_start|> and <|output_end|>)

State Machine Architecture

Tool integration employs a per-sample state machine to track tool usage:

class RowState:
    # Per-row state tracking during generation
    def __init__(self, current_tokens=None):
        self.current_tokens = current_tokens or [] # Current token sequence
        self.forced_tokens = deque() # Queue of tokens to force inject
        self.in_python_block = False # Whether we are inside a python block
        self.python_expr_tokens = [] # Tokens of the current python expression
        self.completed = False # Whether this row has completed generation

State transitions occur during token generation:

# Handle tool logic
if next_token == python_start:
    state.in_python_block = True
    state.python_expr_tokens = []
elif next_token == python_end and state.in_python_block:
    state.in_python_block = False
    if state.python_expr_tokens:
        # Decode the accumulated tokens to get the expression
        expr = self.tokenizer.decode(state.python_expr_tokens)
        result = use_calculator(expr)

        if result is not None:
            # Tokenize the result and prepare to inject it
            result_tokens = self.tokenizer.encode(str(result))
            state.forced_tokens.append(output_start)
            state.forced_tokens.extend(result_tokens)
            state.forced_tokens.append(output_end)
    state.python_expr_tokens = []
elif state.in_python_block:
    # Accumulate tokens inside the python block
    state.python_expr_tokens.append(next_token)

Token Forcing Mechanism

When tool results become available, they are injected into the generation stream, overriding the model’s sampling:

# Select the next token in this row
is_forced = len(state.forced_tokens) > 0
token_masks.append(0 if is_forced else 1) # 0 = forced, 1 = sampled
next_token = state.forced_tokens.popleft() if is_forced else sampled_tokens[i]
token_column.append(next_token)

The token_masks array tracks whether tokens were sampled or forced, which is critical for training (forced tokens should not contribute to the loss function).

Multi-Tool Support

The state machine architecture naturally accommodates multiple tool invocations within a single response:

Calculate the area of a rectangle: <|python_start|>5*3<|python_end|>
<|output_start|>15<|output_end|>
Now the perimeter: <|python_start|>2*(5+3)<|python_end|>
<|output_start|>16<|output_end|>
So area is 15 and perimeter is 16.

Each <|python_end|> token triggers tool evaluation, with results queued in forced_tokens for subsequent injection.

Sampling Strategy: Controlling Generation Quality

After obtaining logits from the model, selecting the next token requires careful consideration. The sampling strategy significantly impacts generation quality, creativity, and determinism.

The Sampling Function

nanochat implements a classical sampling function with temperature and top-k filtering:

@torch.inference_mode()
def sample_next_token(logits, rng, temperature=1.0, top_k=None):
    """Sample a single next token from given logits of shape (B, vocab_size). 
    Returns (B, 1)."""
    assert temperature >= 0.0, "temperature must be non-negative"

    if temperature == 0.0:
        return torch.argmax(logits, dim=-1, keepdim=True)

    if top_k is not None:
        k = min(top_k, logits.size(-1))
        vals, idx = torch.topk(logits, k, dim=-1)
        vals = vals / temperature
        probs = F.softmax(vals, dim=-1)
        choice = torch.multinomial(probs, num_samples=1, generator=rng)
        return idx.gather(1, choice)
    else:
        logits = logits / temperature
        probs = F.softmax(logits, dim=-1)
        return torch.multinomial(probs, num_samples=1, generator=rng)

Greedy Decoding (temperature = 0.0)

if temperature == 0.0:
    return torch.argmax(logits, dim=-1, keepdim=True)

Greedy decoding consistently selects the highest-probability token. Characteristics include:

• Deterministic behavior (identical input yields identical output)
• Computational efficiency (no probability computation required)
• Potential for repetitive outputs
• Suitable for precision-critical tasks (code generation, mathematics)

Example: Given logits for “The cat sat on the ___”:

[("mat", 0.45), ("floor", 0.30), ("table", 0.15), ("car", 0.10)]
→ Consistently selects "mat"

Temperature Sampling

Temperature (τ) controls the level of randomness in token selection by scaling the logits before applying the softmax function:

logits = logits / temperature
probs = F.softmax(logits, dim=-1)
return torch.multinomial(probs, num_samples=1, generator=rng)

Temperature (τ) directly influences how confidently the model selects the next token.

• τ < 1 (e.g., 0.7): Sharpens the probability distribution and increases determinism. Before: [(“mat”, 0.45), (“floor”, 0.30), (“table”, 0.15), (“car”, 0.10)] After: [(“mat”, 0.60), (“floor”, 0.25), (“table”, 0.10), (“car”, 0.05)]
• τ = 1: Leaves the distribution unchanged. The model samples directly from its original probabilities.
• τ > 1 (e.g., 1.5): Flattens the probability distribution and increases randomness. Before: [(“mat”, 0.45), (“floor”, 0.30), (“table”, 0.15), (“car”, 0.10)] After: [(“mat”, 0.35), (“floor”, 0.30), (“table”, 0.20), (“car”, 0.15)]

Application contexts:

• Low temperature (0.5–0.7): Factual writing, code generation, translation
• Medium temperature (0.8–1.0): Conversational agents, question answering, general text
• High temperature (1.2–1.5): Creative writing, brainstorming

Top-K Filtering

Top-k filtering restricts sampling to the k most probable tokens:

if top_k is not None:
    k = min(top_k, logits.size(-1))
    vals, idx = torch.topk(logits, k, dim=-1)
    vals = vals / temperature
    probs = F.softmax(vals, dim=-1)
    choice = torch.multinomial(probs, num_samples=1, generator=rng)
    return idx.gather(1, choice)

Process:

1. Sort tokens by logit value, retain top k
2. Apply temperature scaling to these k tokens exclusively
3. Renormalize probabilities (sum to 1.0)
4. Sample from this restricted distribution

Example with k=2:

Original: [("mat", 0.45), ("floor", 0.30), ("table", 0.15), ("car", 0.10)]
Top-2:    [("mat", 0.60), ("floor", 0.40)]  # renormalized
→ "table" and "car" have zero probability of selection

Advantages:

• Prevents sampling of highly improbable tokens (reduces incoherent outputs)
• Maintains diversity among probable tokens
• Typical values: k=20 for focused generation, k=50 for creative generation

Combined Temperature and Top-K

In practice, both techniques are often combined:

# Example parameters
temperature = 0.8  # Moderately focused
top_k = 40        # Consider 40 most likely tokens
# Process:
# 1. Filter to top 40 tokens
# 2. Apply temperature scaling (slightly sharpen distribution)
# 3. Sample from resulting distribution

This approach provides:

• Safety: from top-k (eliminates improbable tokens)
• Control: from temperature (adjusts randomness level)
• Quality. from both (balances coherence and creativity)

Multi-Sample Generation: Parallel Completions

nanochat implements efficient parallel generation of multiple completions from a single prompt.

Two-Phase Generation Approach

@torch.inference_mode()
def generate(self, tokens, num_samples=1, max_tokens=None, 
             temperature=1.0, top_k=None, seed=42):
    """Generate num_samples completions with single prefill and parallel decode."""

    # 1) Run a batch 1 prefill of the prompt tokens
    m = self.model.config
    kv_model_kwargs = {
        "num_heads": m.n_kv_head, 
        "head_dim": m.n_embd // m.n_head, 
        "num_layers": m.n_layer
    }
    kv_cache_prefill = KVCache(
        batch_size=1,
        seq_len=len(tokens),
        **kv_model_kwargs,
    )
    ids = torch.tensor([tokens], dtype=torch.long, device=device)
    logits = self.model.forward(ids, kv_cache=kv_cache_prefill)
    logits = logits[:, -1, :]
    next_ids = sample_next_token(logits, rng, temperature, top_k)
    sampled_tokens = next_ids[:, 0].tolist()

Phase 1: Single Prefill

• Process entire prompt once (batch size = 1)
• Computationally expensive (O(T²) for sequence length T)
• Executed once regardless of num_samples

# 2) Replicate the KV cache for each sample/row
    kv_length_hint = (len(tokens) + max_tokens) if max_tokens \
                     else self.model.config.sequence_len
    kv_cache_decode = KVCache(
        batch_size=num_samples,
        seq_len=kv_length_hint,
        **kv_model_kwargs,
    )
    kv_cache_decode.prefill(kv_cache_prefill)
    del kv_cache_prefill # free memory

    # 3) Initialize states for each sample
    row_states = [RowState(tokens.copy()) for _ in range(num_samples)]

Phase 2: Parallel Decode

• Clone prefilled cache N times (efficient memory copy operation)
• Batch size now equals num_samples
• Each decoding step processes all samples simultaneously

Per-Sample State Management

Each sample maintains independent state:

# 4) Main generation loop
num_generated = 0
first_iteration = True
while True:
    # Stop condition: we've reached max tokens
    if max_tokens is not None and num_generated >= max_tokens:
        break
    # Stop condition: all rows are completed
    if all(state.completed for state in row_states):
        break

    # Get sampled tokens - either from prefill or from forward pass
    if first_iteration:
        sampled_tokens = [sampled_tokens[0]] * num_samples
        first_iteration = False
    else:
        logits = self.model.forward(ids, kv_cache=kv_cache_decode)
        logits = logits[:, -1, :]
        next_ids = sample_next_token(logits, rng, temperature, top_k)
        sampled_tokens = next_ids[:, 0].tolist()

Each forward pass returns a “token column” containing one token per sample:

Sample 0: [prompt...] → "Hello"
Sample 1: [prompt...] → "Hi"
Sample 2: [prompt...] → "Hey"
Sample 3: [prompt...] → "Greetings"

Early Stopping Implementation

Individual samples can complete at different times:

# Process each row: choose the next token, update state
    token_column = []
    token_masks = []
    for i, state in enumerate(row_states):
        # ... token selection and tool logic ...

        # On <|assistant_end|> or <|bos|>, mark the row as completed
        if next_token == assistant_end or next_token == bos:
            state.completed = True

When a sample generates an end token:

1. Mark as completed
2. Continue generating for remaining samples
3. Terminate loop when all samples complete

This approach is more efficient than waiting for the slowest sample to reach max_tokens.

Efficiency Analysis

Consider generating 8 completions for a 100-token prompt:

Naive approach (separate generation):

• 8 × (prefill + decode) operations
• 8 × O(T²) prefill cost = O(8T²)

Batched approach (nanochat):

• 1 × prefill operation = O(T²)
• Parallel decode = O(T) per step, with 8× memory bandwidth

For typical workloads, this achieves approximately 6–7× speedup compared to naive separate generation.

Complete Generation Flow

A comprehensive generation example:

from nanochat.engine import Engine
from nanochat.checkpoint_manager import load_model

# Load model and tokenizer
model, tokenizer, meta = load_model("base", device="cuda")
# Create engine
engine = Engine(model, tokenizer)
# Prepare prompt
prompt = "What is 15 * 23?"
bos = tokenizer.get_bos_token_id()
tokens = tokenizer.encode(prompt, prepend=bos)
# Generate 4 completions
for token_column, token_masks in engine.generate(
    tokens, 
    num_samples=4,
    max_tokens=100,
    temperature=0.7,
    top_k=50,
    seed=42
):
    # token_column contains 4 tokens (one per sample)
    for i, (token, mask) in enumerate(zip(token_column, token_masks)):
        text = tokenizer.decode([token])
        forced = "[FORCED]" if mask == 0 else ""
        print(f"Sample {i}: {text}{forced}", end="")
    print()

Internal process:

1 . Initialization:

• Create RNG with seed for reproducibility
• Retrieve special token IDs for tool use

2 . Single Prefill:

• Process prompt through model with batch size 1
• Construct initial KV cache with shape (L, 2, 1, H, T, D)
• Sample first token

3 . Cache Replication:

• Create larger cache with batch size 4
• Clone prefilled cache 4 times
• Initialize 4 independent RowState objects

4 . Decoding Loop:

• Forward pass: Generate logits for 4 samples simultaneously
• Sampling: Sample 4 tokens (one per sample)
• State machine: For each sample:
• Check for forced tokens (tool results)
• Update tool use state
• Check for completion
• Tool execution: If <|python_end|> detected:
• Decode accumulated tokens
• Evaluate expression
• Queue result tokens for forcing
• Yield: Return token column and masks

5 . Early Stopping:

• Continue until all samples reach end token or max_tokens

Example output:

Sample 0: Let's calculate: <|python_start|>15*23<|python_end|>
Sample 1: To find 15 times 23, <|python_start|>15*23<|python_end|>
Sample 2: 15 * 23 = <|python_start|>15*23<|python_end|>
Sample 3: The answer is <|python_start|>15*23<|python_end|>
[Tool evaluation occurs]
Sample 0: <|output_start|>[FORCED]345[FORCED]<|output_end|>[FORCED], so 345.
Sample 1: <|output_start|>[FORCED]345[FORCED]<|output_end|>[FORCED] which equals 345.
Sample 2: <|output_start|>[FORCED]345[FORCED]<|output_end|>[FORCED]
Sample 3: <|output_start|>[FORCED]345[FORCED]<|output_end|>[FORCED].

The [FORCED] indicators denote tokens originating from the calculator rather than model sampling.

Performance Considerations

Memory Usage

The KV cache represents the primary memory consumer during inference:

For nanochat’s base configuration:

• L = 6 layers
• B = 4 samples
• H = 6 heads
• T = 2048 max sequence
• D = 64 head dimension
• dtype = bfloat16 (2 bytes)

Memory usage scales linearly with num_samples, establishing practical limits for batched generation.

Computational Efficiency

With KV cache, each decoding step requires:

• Without cache: O(T² · d) attention over all token pairs
• With cache: O(T · d) only new queries attend to cached keys

For a 1000-token sequence, this represents a 1000× speedup in attention computation.

Batching Benefits

Processing N samples in parallel:

• Memory: N× increase (linear scaling)
• Compute: approximately 1.2N× increase (some memory bandwidth overhead)
• Wall time: approximately 1.5N× compared to serial processing (significant improvement)

The optimal configuration typically ranges from num_samples=4-8

Advanced Considerations

Speculative Decoding

An extension compatible with this architecture: employ a smaller “draft” model to generate multiple candidate tokens, then verify them with the full model in parallel. This approach can further reduce latency.

Continuous Batching

In production systems, samples can be dynamically added to or removed from the batch as they complete, maximizing GPU utilization. nanochat’s per-sample state tracking facilitates this implementation.

Multi-Step Tool Use

The tool integration extends naturally to more complex workflows:

User: Plot y = x^2 from -5 to 5
Model: <|python_start|>import matplotlib...<|python_end|>
        <|output_start|>[image data]<|output_end|>

By incorporating additional tools and expanding the state machine, support for code execution, web search, and other capabilities becomes feasible.

Conclusion

Efficient inference results from careful engineering across multiple layers:

1. KV Cache eliminates redundant computation through intelligent caching and dynamic memory management
2. Tool Integration extends model capabilities through a state machine and token forcing mechanism
3. Sampling Strategies balance creativity and coherence through temperature and top-k filtering
4. Multi-Sample Generation maximizes hardware utilization through prefill-then-parallelize architecture

These techniques transform a theoretically correct but computationally expensive generation process into a practical, real-time system capable of powering interactive applications.

References / Sources

[embed]GitHub - karpathy/nanochat: The best ChatGPT that $100 can buy. The best ChatGPT that $100 can buy. Contribute to karpathy/nanochat development by creating an account on GitHub.github.com

[embed]Speculative Decoding: How to Make Large Language Models Think Faster Without Sacrificing Quality Introduction: The Speed-Quality Dilemmamedium.com

[embed]Continuous Batching in LLM Inference Introductionmedium.com