TwinEdgeAI

↑ Course home

← Previous: Week 1, What Are Agentic Systems?  ·  Next: Week 2b, Retrieval-Augmented Generation →

Engineering AI Agents · Week 2

LLMs for Builders

What you need to know about language models to build agents on top of them: decoding, attention and context length, training, inference, structured output, and context engineering.

Author: Grok Bot Version: 1.0, October 28, 2026 Study time: about 9 hours

CreditThis coursework is adapted from, and inspired by, Stanford CS329Z: Engineering AI Agents (Diyi Yang, Michael Ryan, John Yang), Week 2, "LLMs for Builders" (Monday, September 28). The topic list and the three required readings come from the course's public schedule [25]. I checked the topic order against the public lecture slides [26]. The notes, lab, quiz and assignment are original writing, and this pack does not reproduce the slides. It is not an official Stanford document. I added extra sources of my own. Each is marked "added" in the references.

1. Week overview and learning objectives

Week 1 gave you the map of agentic systems. This week opens the box in the middle: the language model.

You do not need to train a model to build an agent. But you do need a working picture of how one behaves. Why does the same prompt give different answers? Why does a long chat get worse? Why is JSON sometimes broken? Why does one model cost ten times more than another? The answers come from how models are built and how they run.

The course puts it this way: agent builders work at the top of the "LLM stack", but they still need to know how the layers below behave, because those layers set limits on what the top can do [26]. This week is about those limits and how to design around them.

Learning objectives

By the end of this week you should be able to:

1. Explain how an autoregressive language model turns a prompt into text, one token at a time, and why sums of log probabilities are used in practice.
2. Choose a decoding setting (greedy, temperature, top-k, top-p) for a given agent task, and say why [5].
3. Describe scaled dot-product attention, the causal mask and the KV cache, and say why attention cost grows with the square of the context length [1].
4. Compare full softmax attention, linear attention and hybrid designs, and say what each means for long contexts [2] [18] [19].
5. Trace the path from pretraining to a deployed assistant (SFT, RLHF, DPO, verifiable rewards), and name what each stage changes [8] [9].
6. Explain test-time scaling: chain of thought, repeated sampling, voting and speculative decoding, and what each costs [11] [13] [14] [15].
7. Pick a structured-output method (prompt and retry, JSON mode, schema-constrained decoding) and handle its failure cases [22] [24].
8. Apply context engineering: budget the window, clear old tool output, compact, and place key facts well [3] [21].
9. Build and test small versions of all of the above in Python, without an API key.

Plan for the week

Part	What you do	Time
Reading	Three required readings, plus optional ones (reading guide)	3.5 h
Lecture notes	Work through section 2. Redraw the diagrams from memory.	2 h
Lab	Run, read and extend the toolkit	1.5 h
Discussion and quiz	Talk it through with a partner, then take the quiz	1 h
Assignment	Small graded project (section 7)	1 h to 2 h

2. Lecture notes

These notes follow the topic order of the lecture [25] [26]. Citations in square brackets link to the references.

2.1 What a language model is

A language model gives a probability to the next token. A token is a small piece of text: a word, part of a word, or a symbol. Given the tokens so far, the model returns one probability for every token in its vocabulary. These sum to 1.

To write text, the model picks a token, adds it to the input, and asks again. This is autoregressive generation. The probability of a whole sequence is the product of each step's probability, given everything before it.

Training uses the same numbers. The loss for one position is minus the log of the probability the model gave to the token that was really there. My own example: if the model gave the right token 0.25, the loss is −ln(0.25) ≈ 1.39. If it gave 0.9, the loss is about 0.11. Training pushes the right token's probability up, across trillions of positions.

Why this matters for builders. A language model is a generative model of text. It can sample. It does not "look up" an answer. Everything else in this week follows from that: the answer is a draw from a distribution, you can change how you draw, and every token costs time and memory.

2.2 Decoding: greedy, temperature, top-k, top-p

A decoding strategy is the rule for picking the next token from the distribution.

• Greedy: always take the most likely token. It is fast and repeatable on the same inputs. It can get stuck in loops. Holtzman et al. found that pushing for high likelihood gives text that is "bland and strangely repetitive" [5].
• Temperature (T): divide the scores by T before the softmax. Low T makes the top token even more likely. High T flattens the distribution. As T approaches 0 you get greedy.
• Top-k: keep only the k most likely tokens, then renormalize and sample.
• Top-p (nucleus): keep the smallest set of top tokens whose total probability reaches p, then renormalize and sample. The set size changes with the step: small when the model is sure, large when it is not [5].
• Beam search: keep several candidate sequences and extend the best ones. It costs more time and memory [26].

Setting	Effect	Typical use (my advice, not a rule)
Temperature near 0, or greedy	Most repeatable. Least varied.	Tool calls, JSON, code, math, classification
Temperature about 0.7 to 1	More variety. More chance of odd picks.	Drafting, brainstorming, many diverse samples
Top-p 0.9 to 0.95	Cuts the unreliable tail while keeping variety	Open-ended text [5]
Top-k small	Hard limit on how many choices exist	A simple guard when the tail is noisy

For agents. A tool call that is slightly wrong breaks the loop. Use low temperature for the steps that must be exact. Use higher temperature when you want diverse candidates to vote on or rank (see 2.6).

2.3 Transformer attention

Nearly every modern LLM is built on the Transformer [1]. The key part is attention. It lets each token read information from other tokens in the context. For a gentle visual tour, read Alammar's illustrated guide [4].

Scaled dot-product attention

Each token's vector is turned into three vectors by learned matrices: a query (what I am looking for), a key (what I offer) and a value (what I pass on). For one token:

1. Compare its query with every key by dot product. This gives one score per token.
2. Divide the scores by √d_k, where d_k is the key size. Vaswani et al. say that for large d_k the dot products grow large and push the softmax into regions with extremely small gradients. The scaling counters that [1].
3. Apply softmax. Now the scores are weights that sum to 1.
4. Take the weighted sum of the value vectors. That is the output for this token.

In matrix form this is softmax(QK^T/√d_k)V [1].

The pieces around attention

• Causal mask. In a text generator, a token may only attend to itself and earlier tokens. This makes training match generation, because at generation time the future does not exist yet. The lab test test_changing_a_future_token_does_not_change_the_past checks it.
• Multi-head attention. Run several attention heads in parallel, each with its own matrices, then join the results. The base model in the original paper used 8 heads with d_k = 64 [1].
• The block. Each layer wraps attention with a feed-forward network, residual connections and layer norm. Stack many layers, then use a final linear layer and softmax to get next-token probabilities [1] [4].
• Positions. Attention alone has no sense of order, so the model adds position information to its inputs [1]. How this is done matters for context length (2.4).

The original paper used an encoder and a decoder. Most chat and agent models today use only the decoder side. They are called decoder-only models [26].

Why Transformers won, and what it costs

Table 1 of the paper compares layer types. Self-attention needs O(1) sequential steps per layer, so GPUs can process all positions at once. A recurrent layer needs O(n) sequential steps. But self-attention costs O(n²·d) per layer, where n is the sequence length [1]. Double the context and the attention work goes up about four times. This one fact drives much of the rest of this week.

2.4 Linear and hybrid attention, and context length

Linear attention: the Transformer as an RNN

Katharopoulos et al. rewrite attention as a dot product of kernel feature maps and use the associativity of matrix products. This drops the cost from O(N²) to O(N) [2]. Here is the idea in plain words:

• Softmax attention builds an n-by-n score table before it touches the values.
• If you drop the softmax and apply a positive feature map φ to queries and keys instead, you can group the math as φ(Q) × (φ(K)^T V). Compute the second part first. It is a small fixed-size matrix.
• For text generation with a causal mask, that small matrix is a running sum. Each new token adds its own term and then reads the sum. The paper keeps two such sums: an attention memory and a normalizer memory. That is an RNN with a fixed-size state [2].

The paper reports linear transformers with similar quality to vanilla ones on its tests, and up to 4000× faster on autoregressive prediction of very long sequences [2]. The lab builds both forms. A test shows they give the same numbers, using the paper's feature map elu(x)+1.

Hybrid models

Squeezing the past into a fixed state costs recall. The Mamba authors write that earlier subquadratic designs, including linear attention, did not match attention on language, and identify a key weakness: they could not do content-based reasoning. Mamba makes the state-space parameters depend on the input to fix this [18]. Jamba goes a step further and interleaves Transformer layers with Mamba layers, and adds mixture-of-experts. The authors report high throughput, a small memory footprint, and strong results up to 256K tokens of context [19]. The lecture's summary is a good rule: in practice, people mix global and linear attention [26].

Context length is not one number

• What the model supports is limited by how it was trained, especially its position handling [26]. Anthropic notes that models see fewer long sequences than short ones in training, and that techniques like position-encoding interpolation stretch a model to longer contexts with some loss in position understanding [3].
• What the model uses well is often less. RULER tested 17 long-context models. Almost all lost accuracy as length grew. Only about half kept satisfactory performance at 32K tokens, even though all claimed 32K or more [20].
• Where facts sit matters. In "Lost in the Middle", performance was often highest when the relevant fact was at the start or end of the input, and fell when it was in the middle [21].

Making attention run faster and fit in memory

• KV cache. When generating, the keys and values of earlier tokens do not change. The model stores them and only computes the new token's. This saves compute but costs memory that grows with context [26].
• FlashAttention computes exact attention but is "IO-aware": it uses tiling to cut reads and writes between GPU high-bandwidth memory and on-chip memory. It is a faster way to get the same answer, not an approximation [17].
• PagedAttention (in vLLM) manages the KV cache like operating-system virtual memory, cutting waste from fragmentation. The authors report 2 to 4 times higher throughput than earlier serving systems at the same latency [16].

KV cache size. My own arithmetic. The model below is made up, not a real product.

Setting	Value
Layers, KV heads, head size, precision	32, 8, 128, 2 bytes (fp16)
Bytes per token = 2 (K and V) × 32 × 8 × 128 × 2	131,072 (128 KiB)
At 32,768 tokens (one request)	4 GiB
At 128 requests at once	512 GiB

The lab function kv_cache_bytes does this sum. Real models differ in layout, so use your model's own numbers. The point is the shape: memory grows with context length and with the number of requests. That is why serving long contexts is expensive, and why inference cost shows up in your agent's bill.

Why this matters for agents. Agents run long. Their contexts fill with tool output. Knowing that attention is costly and that recall fades helps you pick a model, set limits, and keep the context small (2.8) [26].

2.5 From pretraining to post-training

You will rarely train a model from scratch. But the steps that made it decide how it behaves, so you should know them.

Pretraining

The model reads a very large amount of text and learns to predict the next token. The loss is the cross-entropy from 2.1. Two things matter most: how much compute you spend, and what data you feed it.

• Scaling laws. Kaplan et al. found that loss falls as a power law with model size, data size and compute, over many orders of magnitude [6].
• Compute-optimal training. Hoffmann et al. (the "Chinchilla" paper) trained over 400 models and concluded that models of the time were undertrained. For a fixed budget, model size and training tokens should grow together: double the model, double the tokens [7].
• Data. The lecture shows that web text is filtered and de-duplicated, and that mixing in more of your target domain (math, code) helps [26].

Post-training

• Instruction tuning (SFT). Train on prompt and good-answer pairs. The model learns to follow instructions rather than only to continue text. In the InstructGPT recipe, labelers wrote demonstrations, and GPT-3 was fine-tuned on them [8].
• RLHF. Humans rank model outputs. A reward model learns those rankings. Reinforcement learning then tunes the model toward higher reward. The InstructGPT paper reports that outputs from the 1.3B InstructGPT model were preferred to those of the 175B GPT-3 in human tests, despite having 100 times fewer parameters [8].
• DPO. Direct Preference Optimization reaches the same goal with a simple classification-style loss. It does not need a separate reward model or sampling from the model during fine-tuning. The authors say it is stable and light, and matches or beats PPO-based RLHF on their tasks [9].
• Verifiable rewards (RLVR). When the answer can be checked by a program (unit tests, exact match, a proof checker), you can use that check as the reward. This works well for math and code. You must make sure the checker cannot be fooled [26]. DeepSeekMath's GRPO is one RL method used for math reasoning [10].

For agents. The lecture also covers training aimed at agent tasks themselves, such as training on coding-agent trajectories [26]. When you pick a model, ask what it was post-trained for: chat, tool use, long tasks, or reasoning. A model's post-training often matters more to your agent than its raw size.

2.6 Inference and test-time scaling

How inference runs

Generation has two phases [26]:

1. Prefill. The model reads the whole prompt at once and builds the KV cache. All prompt tokens can be processed in parallel.
2. Decode. The model produces one token at a time. Each new token must read the cache. This phase is usually limited by how fast memory can be read, not by raw arithmetic.

Two practical results: a long prompt makes the first token slower (prefill), and a long answer makes the whole reply slower (decode). Price lists often charge for input and output tokens at different rates. Check your provider's page.

Speculative decoding

A small, fast draft model guesses several tokens ahead. The big model then checks all of them in one parallel pass and keeps the ones it agrees with. Leviathan et al. show this gives the same output distribution as the big model alone, with a 2 to 3 times speedup on T5-XXL [15]. Gains depend on how predictable the text is: boilerplate and code are easier to guess than fresh prose [26].

Test-time scaling: spend more compute at answer time

Pretraining spends compute once. Test-time scaling spends more compute on each question. There are a few ways.

Method	What it does	What to remember
Chain of thought	The model writes intermediate steps before the answer [11]	Costs more output tokens. Helps on hard multi-step tasks.
Reasoning models	Trained to produce long reasoning on their own, often with an "effort" setting [26]	More effort means more latency and cost. Match effort to the task.
Repeated sampling	Draw many answers, keep one that passes a check [13]	Needs a good verifier (tests, a checker).
Self-consistency	Sample many reasoning paths, take the most common answer [12]	Works when there is no verifier, but gains can plateau [13].
Search against a verifier	Use a reward model to guide or pick among candidates [14]	Best method depends on how hard the prompt is.

Evidence worth knowing.

• Brown et al. found that "coverage", the share of problems solved by any sample, rises with the number of samples over four orders of magnitude. On SWE-bench Lite, DeepSeek-Coder-V2-Instruct went from 15.9% with one sample to 56% with 250 samples. Without an automatic verifier, majority voting and reward models plateaued after several hundred samples [13].
• Snell et al. found that the best way to use test-time compute depends on prompt difficulty. A compute-optimal, per-prompt strategy improved efficiency by over 4 times compared with plain best-of-N. In a FLOPs-matched test, on problems where a smaller model had non-trivial success, extra test-time compute could beat a model 14 times larger [14].

For agents. Agents are loops, so cost multiplies: steps × samples × tokens. Use cheap settings for easy steps. Spend on the hard steps, and only where a check can tell good from bad. This links straight to Week 1's point about cost-aware evaluation.

2.7 Structured I/O and constrained decoding

Software parts need a contract. A planner hands work to a tool. A model hands data to a database. If the model's output is "almost JSON", the next step fails. So agent builders spend a lot of effort on making outputs machine-readable. The course frames this as the shared contract between parts of a compound system [26].

A ladder of methods, from weakest to strongest guarantee

1. Ask nicely. Put the format in the prompt and show an example. No guarantee.
2. Validate and retry. Parse the reply. Check it against a schema. If it fails, send the errors back and ask again. This is the loop in the lab. It always works eventually if the model can do the task, but costs extra calls.
3. JSON mode. The API makes the model emit valid JSON. OpenAI's guide says JSON mode ensures valid JSON but does not ensure that it matches your schema [24].
4. Schema-constrained output. OpenAI's Structured Outputs feature makes replies adhere to a supplied JSON Schema, and the guide recommends it over JSON mode when possible [24].
5. Constrained decoding in your own stack. You control the decoder. Before each sampling step, you block every token that would break the format [22].

Other ways to constrain. Instead of constraining tokens, you can constrain the choice: a small classifier head on top of embeddings that picks one of N labels, or a tool-calling interface with a fixed list of tools. The lecture covers both as ways to constrain the model [26]. In my words: the more you constrain a step, the more predictable it is, and the less flexible it is.

2.8 Context engineering

Anthropic's article defines context as the set of tokens included when sampling from an LLM. Context engineering is the work of curating and maintaining the best set of tokens during inference, including everything that lands there beyond the prompt [3]. Prompt engineering is about writing instructions. Context engineering is about the whole state: system prompt, tools, examples, message history, retrieved data, and tool results.

Why not put everything in?

Three reasons, all from earlier sections:

• Context rot. As the number of tokens grows, the model's ability to recall information from the context drops. Anthropic describes this as a gradient, not a cliff [3]. RULER measured similar drops [20].
• Attention is a budget. Every token competes for attention, and n tokens make about n² pairs [3] (2.3).
• Position matters. Facts in the middle of a long context are used less well than facts at the ends [21].

The guiding principle from the article: find the smallest possible set of high-signal tokens that makes the desired outcome likely [3].

Practical rules

1. System prompt: the right altitude. Not brittle if-else logic. Not vague hints. Specific enough to guide, flexible enough to give strong heuristics. Start minimal, test, then add based on failures [3].
2. Few tools, clear tools. A common failure is a bloated tool set with overlapping jobs. If a human cannot say which tool fits, the model will not do better [3]. The lecture notes that a few tools can be enough: one example agent uses only read, write, edit and bash [26].
3. Canonical examples, not a pile of edge cases. Pick a few diverse, clear examples. "For an LLM, examples are the pictures worth a thousand words" [3].
4. Retrieve just in time. Keep light references (file paths, queries, links) and load content with tools when needed, instead of loading everything up front. A hybrid is common: load some data first for speed, explore the rest at runtime [3]. Next week covers retrieval in depth [25].
5. Compact. When near the limit, summarize and restart with the summary. Anthropic calls compaction the first lever for long-term coherence. The hard part is deciding what to keep: too aggressive and you lose something you need later. Tune on real traces, first for recall, then for precision [3].
6. Clear old tool results. Anthropic calls this "one of the safest lightest touch forms of compaction". Once a tool was called deep in the history, the agent rarely needs the raw output again [3].
7. Take notes outside the window. Let the agent write a NOTES file or a to-do list and read it back later [3].
8. Use sub-agents. A sub-agent can use tens of thousands of tokens exploring, then return a short summary (often 1,000 to 2,000 tokens, per the article) [3].
9. Put key facts at the edges. Because of the "lost in the middle" effect, place the instruction and the most relevant evidence near the start or end [21]. The lab function edges_first does this.
10. Protect the cache. Append new content at the end. Editing earlier text changes everything after it, and cached keys and values no longer apply [26]. Many providers offer prompt caching. Check your provider's rules.
11. Measure. There is no single right way to build context. Make a small eval, change one thing, and compare [26] [3].

2.9 Week 2 in one page

Topic	Core idea	What to do as a builder
Language model	Predicts next-token probabilities. Text is a sample.	Expect variation. Check outputs with code where it matters.
Decoding	Temperature, top-k, top-p reshape the draw [5]	Low for exact steps. Higher for diverse candidates.
Attention	Each token reads earlier tokens. Cost is quadratic in length [1]	Keep contexts short. Expect long calls to cost more.
Linear and hybrid	Fixed-size state trades recall for speed [2] [19]	Test long-context recall on your task [20].
Training	Pretrain, then SFT, preference tuning, verifiable RL [8] [9]	Choose models by post-training fit. Watch for sycophancy.
Inference	Prefill, decode, KV cache. Test-time compute buys accuracy [13] [14]	Spend compute where a check can tell good from bad.
Structured I/O	Constraints guarantee shape, not truth [22] [24]	Constrain, validate, handle refusal and truncation.
Context engineering	Smallest set of high-signal tokens [3]	Budget, clear, compact, take notes, split into sub-agents.

3. Reading guide

Three readings are required. Ten more are optional. For each one you get: why it matters, the key ideas, and questions to think about as you read. Read the questions first. Take short notes as you go. Pick optional readings by your interest: the first three optional ones are good for most people.

Also cited in the notes, not assigned: Kaplan et al. [6], Rafailov et al. [9], Shao et al. [10], Wei et al. [11], Wang et al. [12], Brown et al. [13], Gu and Dao [18], Hsieh et al. [20], and the Outlines code [23].

4. Hands-on lab: an LLM toolkit with a mock model

Time: about 90 minutes. Goal: build small, testable versions of four ideas from this week: decoding, attention, constrained output, and a context budget. You do not need an API key. Mock models stand in for real ones.

What you will learn

• How temperature, top-k and top-p change a distribution, with numbers you can print.
• How attention weights are built, why the causal mask matters, and why linear attention can be run as an RNN.
• The difference between retrying until the output is valid and masking so it cannot be invalid.
• How a context budget works: count, clear, summarize, keep what is pinned.

Setup

1. You need Python 3.9 or newer. There is nothing to install. The code uses only the standard library.
2. Save the code below as llm_toolkit.py. Or download it: llm_toolkit.py.
3. Save the tests as test_llm_toolkit.py in the same folder. Or download: test_llm_toolkit.py. The test code is listed after the exercises.
4. Run python llm_toolkit.py for the demo and python test_llm_toolkit.py for the tests.

The code: llm_toolkit.py

"""llm_toolkit.py - Week 2 lab: LLMs for builders. Standard library only.

Four small tools, each with a mock model so no API key is needed:
  1. Decoding: greedy, temperature, top-k, top-p over a toy distribution
  2. Attention: scaled dot-product, causal, and linear attention (as an RNN)
  3. Constrained decoding: token masking, plus a validate-and-retry loop
  4. Context budget: count, clear old tool output, compact, reorder

Run `python llm_toolkit.py` for a demo. Run `python test_llm_toolkit.py` for tests.
"""
import json
import math
import random

# ---------------------------------------------------------------- 1. decoding

def softmax(logits, temperature=1.0):
    """Turn scores into probabilities. Lower temperature = sharper."""
    if temperature <= 0:
        raise ValueError("temperature must be > 0 (use greedy for 0)")
    scaled = [x / temperature for x in logits]
    m = max(scaled)
    exps = [math.exp(x - m) for x in scaled]
    total = sum(exps)
    return [e / total for e in exps]


def entropy(probs):
    """Average surprise in bits. 0 means the model is certain."""
    return -sum(p * math.log2(p) for p in probs if p > 0)


def greedy(probs):
    """Index of the most likely token. Ties go to the lowest index."""
    return max(range(len(probs)), key=lambda i: (probs[i], -i))


def top_k_filter(probs, k):
    """Keep the k most likely tokens, zero the rest, renormalize."""
    if k < 1:
        raise ValueError("k must be >= 1")
    keep = set(sorted(range(len(probs)), key=lambda i: (-probs[i], i))[:k])
    return _renorm([p if i in keep else 0.0 for i, p in enumerate(probs)])


def top_p_filter(probs, p):
    """Nucleus: keep the smallest set of top tokens whose mass reaches p."""
    if not 0 < p <= 1:
        raise ValueError("p must be in (0, 1]")
    order = sorted(range(len(probs)), key=lambda i: (-probs[i], i))
    keep, mass = set(), 0.0
    for i in order:
        keep.add(i)
        mass += probs[i]
        if mass >= p - 1e-12:
            break
    return _renorm([q if i in keep else 0.0 for i, q in enumerate(probs)])


def _renorm(probs):
    total = sum(probs)
    return [q / total for q in probs]


def sample_index(probs, rng):
    """Draw one index from a probability list."""
    r, acc = rng.random(), 0.0
    for i, q in enumerate(probs):
        acc += q
        if r < acc:
            return i
    return max(i for i, q in enumerate(probs) if q > 0)


def decode_step(logits, rng, temperature=1.0, top_k=None, top_p=None):
    """One decoding step. temperature=0 means greedy."""
    if temperature == 0:
        return greedy(logits)
    probs = softmax(logits, temperature)
    if top_k is not None:
        probs = top_k_filter(probs, top_k)
    if top_p is not None:
        probs = top_p_filter(probs, top_p)
    return sample_index(probs, rng)

# ---------------------------------------------------------------- 2. attention

def matmul(a, b):
    bt = list(zip(*b))
    return [[sum(x * y for x, y in zip(row, col)) for col in bt] for row in a]


def transpose(a):
    return [list(r) for r in zip(*a)]


def attention(Q, K, V, causal=True):
    """Scaled dot-product attention. Returns (output, weights).
    Q, K: n x d lists. V: n x dv. Row i of weights sums to 1."""
    d = len(Q[0])
    scores = matmul(Q, transpose(K))
    n = len(scores)
    weights = []
    for i in range(n):
        row = [s / math.sqrt(d) for s in scores[i]]
        if causal:
            row = [x if j <= i else float("-inf") for j, x in enumerate(row)]
        m = max(row)
        e = [0.0 if x == float("-inf") else math.exp(x - m) for x in row]
        t = sum(e)
        weights.append([x / t for x in e])
    return matmul(weights, V), weights


def feature_map(x):
    """elu(x) + 1, the positive feature map used for linear attention."""
    return [v + 1.0 if v > 0 else math.exp(v) for v in x]


def linear_attention_parallel(Q, K, V):
    """Causal linear attention written like softmax attention: O(n^2)."""
    n = len(Q)
    out = []
    for i in range(n):
        qf = feature_map(Q[i])
        ws = [sum(a * b for a, b in zip(qf, feature_map(K[j]))) for j in range(i + 1)]
        z = sum(ws)
        out.append([sum(ws[j] * V[j][c] for j in range(i + 1)) / z for c in range(len(V[0]))])
    return out


def linear_attention_recurrent(Q, K, V):
    """The same result as an RNN. Keeps a fixed-size state: O(n) time, O(1) memory."""
    d, dv = len(Q[0]), len(V[0])
    S = [[0.0] * dv for _ in range(d)]   # sum of phi(k) v^T
    z = [0.0] * d                        # sum of phi(k)
    out = []
    for q, k, v in zip(Q, K, V):
        kf, qf = feature_map(k), feature_map(q)
        for a in range(d):
            z[a] += kf[a]
            for c in range(dv):
                S[a][c] += kf[a] * v[c]
        den = sum(qf[a] * z[a] for a in range(d))
        out.append([sum(qf[a] * S[a][c] for a in range(d)) / den for c in range(dv)])
    return out


def kv_cache_bytes(layers, kv_heads, head_dim, seq_len, bytes_per_value=2, batch=1):
    """Size of the key-value cache: 2 (K and V) x layers x heads x dim x tokens."""
    return 2 * layers * kv_heads * head_dim * seq_len * bytes_per_value * batch

# ------------------------------------------------- 3. constrained decoding

class MockJSONModel:
    """Pretends to be an LLM that is bad at JSON at first.
    It replays `replies` in order. Once it runs out, it repeats the last one."""
    def __init__(self, replies):
        self.replies, self.calls, self.prompts = list(replies), 0, []

    def complete(self, prompt):
        self.prompts.append(prompt)
        r = self.replies[min(self.calls, len(self.replies) - 1)]
        self.calls += 1
        return r


_TYPES = {"string": str, "integer": int, "number": (int, float), "boolean": bool,
          "array": list, "object": dict, "null": type(None)}


def validate(value, schema, path="$"):
    """Tiny JSON-Schema subset: type, enum, required, properties, items.
    Returns a list of error strings. Empty list means valid."""
    errs = []
    t = schema.get("type")
    if t:
        ok = isinstance(value, _TYPES[t]) and not (t in ("integer", "number") and isinstance(value, bool))
        if t == "integer" and isinstance(value, float):
            ok = False
        if not ok:
            return [f"{path}: expected {t}, got {type(value).__name__}"]
    if "enum" in schema and value not in schema["enum"]:
        errs.append(f"{path}: {value!r} is not one of {schema['enum']}")
    if isinstance(value, dict):
        for key in schema.get("required", []):
            if key not in value:
                errs.append(f"{path}: missing required key '{key}'")
        for key, sub in schema.get("properties", {}).items():
            if key in value:
                errs += validate(value[key], sub, f"{path}.{key}")
    if isinstance(value, list) and "items" in schema:
        for i, item in enumerate(value):
            errs += validate(item, schema["items"], f"{path}[{i}]")
    return errs


def extract_json(text):
    """Pull the first {...} block out of a reply. Models often add chatter or code fences."""
    start = text.find("{")
    if start < 0:
        raise ValueError("no JSON object found")
    depth, in_str, esc = 0, False, False
    for i in range(start, len(text)):
        c = text[i]
        if in_str:
            esc = (c == "\\") and not esc
            if c == '"' and not esc:
                in_str = False
        elif c == '"':
            in_str = True
        elif c == "{":
            depth += 1
        elif c == "}":
            depth -= 1
            if depth == 0:
                return json.loads(text[start:i + 1])
    raise ValueError("unterminated JSON object")


def generate_json(model, task, schema, max_retries=3):
    """Ask, parse, validate. On failure, send the errors back and retry.
    Returns (value, attempts). Raises RuntimeError when retries run out."""
    prompt = f"{task}\nReply with JSON only. Schema: {json.dumps(schema)}"
    for attempt in range(1, max_retries + 2):
        reply = model.complete(prompt)
        try:
            value = extract_json(reply)
            errs = validate(value, schema)
        except ValueError as e:      # json.JSONDecodeError is a ValueError
            errs = [f"not valid JSON: {e}"]
        if not errs:
            return value, attempt
        prompt = (f"{task}\nReply with JSON only. Schema: {json.dumps(schema)}\n"
                  f"Your last reply was rejected:\n- " + "\n- ".join(errs) + "\nTry again.")
    raise RuntimeError(f"no valid JSON after {max_retries + 1} attempts; last errors: {errs}")


class TrieConstraint:
    """Allow only token sequences that spell one of `allowed` strings.
    This is the idea behind constrained decoding: at each step, only legal
    next tokens keep their score. Real systems build such masks from a regex
    or grammar over the model's real vocabulary [Willard and Louf 2023]."""
    def __init__(self, allowed, tokenize):
        self.paths = {tuple(tokenize(s)) for s in allowed}

    def legal_tokens(self, prefix):
        n = len(prefix)
        return {p[n] for p in self.paths if len(p) > n and p[:n] == tuple(prefix)}

    def is_complete(self, prefix):
        return tuple(prefix) in self.paths


def constrained_decode(score_fn, vocab, constraint, rng, temperature=1.0, max_len=20):
    """Sample tokens, but set every illegal token's score to -inf first.
    score_fn(prefix) -> list of logits aligned with vocab (a mock 'model')."""
    prefix = []
    while not constraint.is_complete(prefix) and len(prefix) < max_len:
        legal = constraint.legal_tokens(prefix)
        if not legal:
            raise RuntimeError("dead end: no legal token")
        logits = score_fn(prefix)
        masked = [x if vocab[i] in legal else float("-inf") for i, x in enumerate(logits)]
        probs = softmax(masked, temperature)
        prefix.append(vocab[sample_index(probs, rng)])
    return "".join(prefix)

# ------------------------------------------------------ 4. context budget

def count_tokens(text):
    """Toy token counter: 1 word = 1 token. Real tokenizers differ. Swap yours in."""
    return len(text.split())


def total_tokens(messages, counter=count_tokens):
    return sum(counter(m["text"]) for m in messages)


def clear_old_tool_results(messages, keep_last=2):
    """Lightest compaction: replace old tool outputs with a short stub."""
    tool_idx = [i for i, m in enumerate(messages) if m["role"] == "tool"]
    old = set(tool_idx[:-keep_last]) if keep_last else set(tool_idx)
    return [dict(m, text="[tool result cleared]") if i in old and not m.get("pinned") else m
            for i, m in enumerate(messages)]


def mock_summarize(messages):
    """Stand-in for an LLM summary: first six words of each message."""
    return "Summary: " + " | ".join(" ".join(m["text"].split()[:6]) for m in messages)


def compact(messages, limit, reserve_output=0, summarize=mock_summarize,
            keep_recent=3, counter=count_tokens):
    """Fit messages into limit - reserve_output tokens.
    Step 1: clear old tool results. Step 2: summarize the oldest unpinned
    messages into one. Pinned messages and the last `keep_recent` stay.
    Returns (messages, report)."""
    room = limit - reserve_output
    report = {"before": total_tokens(messages, counter), "steps": []}
    cur = messages
    if total_tokens(cur, counter) > room:
        cur = clear_old_tool_results(cur)
        report["steps"].append("cleared old tool results")
    if total_tokens(cur, counter) > room:
        tail = cur[-keep_recent:] if keep_recent else []
        head = cur[:len(cur) - len(tail)]
        pinned = [m for m in head if m.get("pinned")]
        loose = [m for m in head if not m.get("pinned")]
        if loose:
            cur = pinned + [{"role": "system", "text": summarize(loose)}] + tail
            report["steps"].append(f"summarized {len(loose)} messages")
    report["after"] = total_tokens(cur, counter)
    report["fits"] = report["after"] <= room
    return cur, report


def edges_first(items_by_relevance):
    """Reorder so the most relevant items sit at the start and the end,
    and the least relevant sit in the middle (see 'Lost in the Middle')."""
    front, back = [], []
    for i, item in enumerate(items_by_relevance):
        (front if i % 2 == 0 else back).append(item)
    return front + back[::-1]

# ---------------------------------------------------------------- 5. extras

def coverage(p_single, n):
    """Chance that at least one of n independent samples is right."""
    return 1 - (1 - p_single) ** n


def majority_vote(answers):
    """Most common answer. Ties go to the one seen first."""
    counts = {}
    for a in answers:
        counts[a] = counts.get(a, 0) + 1
    return max(counts, key=lambda a: counts[a])

# ---------------------------------------------------------------- demo

def _demo():
    rng = random.Random(7)
    words = ["to", "as", "the", "a", "two", "how"]
    logits = [2.0, 0.5, 0.3, 0.2, -0.5, -1.0]
    print("== 1. decoding (toy distribution) ==")
    for T in (0.5, 1.0, 2.0):
        p = softmax(logits, T)
        print(f"T={T}: " + " ".join(f"{w}={q:.2f}" for w, q in zip(words, p)) + f"  entropy={entropy(p):.2f} bits")
    p1 = softmax(logits)
    print("top_k=2 :", [round(x, 2) for x in top_k_filter(p1, 2)])
    print("top_p=.8:", [round(x, 2) for x in top_p_filter(p1, 0.8)])
    print("greedy  :", words[greedy(p1)])
    draws = [words[decode_step(logits, rng, 1.0, top_p=0.9)] for _ in range(8)]
    print("8 draws, T=1, top_p=0.9:", draws)

    print("\n== 2. attention ==")
    Q = [[1.0, 0.0], [0.0, 1.0], [2.0, 0.0]]
    K = [[1.0, 0.0], [0.0, 1.0], [1.0, 0.0]]
    V = [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]]
    out, w = attention(Q, K, V)
    print("causal weights row 3:", [round(x, 3) for x in w[2]])
    a = linear_attention_parallel(Q, K, V)
    b = linear_attention_recurrent(Q, K, V)
    print("linear: parallel == recurrent:", all(abs(x - y) < 1e-9 for r, s in zip(a, b) for x, y in zip(r, s)))
    kv = kv_cache_bytes(32, 8, 128, 32768)
    print(f"KV cache for a made-up 32-layer model at 32k tokens: {kv / 2**30:.1f} GiB")

    print("\n== 3. constrained decoding ==")
    schema = {"type": "object", "required": ["name", "intent"], "properties": {
        "name": {"type": "string"}, "intent": {"type": "string", "enum": ["meeting", "intro", "follow-up"]}}}
    model = MockJSONModel(['Sure! {"name": "Sarah", "intent": "chat"}',
                           '{"name": "Sarah"}',
                           '```json\n{"name": "Sarah", "intent": "meeting"}\n```'])
    value, attempts = generate_json(model, "Extract contact info: I'm Sarah. Meet Thursday?", schema)
    print(f"valid after {attempts} attempts:", value)
    vocab = ["meet", "ing", "in", "tro", "follow", "-up", "chat", "xyz"]
    tok = lambda s: {"meeting": ["meet", "ing"], "intro": ["in", "tro"], "follow-up": ["follow", "-up"]}[s]
    cons = TrieConstraint(["meeting", "intro", "follow-up"], tok)
    bad_model = lambda prefix: [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 5.0, 5.0]  # loves illegal tokens
    outs = {constrained_decode(bad_model, vocab, cons, rng) for _ in range(20)}
    print("20 constrained samples, all legal:", sorted(outs))

    print("\n== 4. context budget ==")
    msgs = [{"role": "system", "text": "You are a careful coding agent.", "pinned": True},
            {"role": "user", "text": "Fix the failing test in the billing module please"},
            {"role": "tool", "text": "log line " * 40},
            {"role": "assistant", "text": "I will open the test file and read it now"},
            {"role": "tool", "text": "file contents " * 30},
            {"role": "assistant", "text": "The bug is a wrong rounding mode"},
            {"role": "tool", "text": "diff output " * 10},
            {"role": "user", "text": "Great. Now add a regression test."}]
    new, rep = compact(msgs, limit=120, reserve_output=20)
    print(rep)
    print("order for 5 docs ranked 1..5:", edges_first([1, 2, 3, 4, 5]))


if __name__ == "__main__":
    _demo()

Run the demo

$ python llm_toolkit.py
== 1. decoding (toy distribution) ==
T=0.5: to=0.89 as=0.04 the=0.03 a=0.02 two=0.01 how=0.00  entropy=0.69 bits
T=1.0: to=0.59 as=0.13 the=0.11 a=0.10 two=0.05 how=0.03  entropy=1.87 bits
T=2.0: to=0.36 as=0.17 the=0.15 a=0.14 two=0.10 how=0.08  entropy=2.40 bits
top_k=2 : [0.82, 0.18, 0.0, 0.0, 0.0, 0.0]
top_p=.8: [0.71, 0.16, 0.13, 0.0, 0.0, 0.0]
greedy  : to
8 draws, T=1, top_p=0.9: ['to', 'to', 'as', 'to', 'to', 'to', 'to', 'to']

== 2. attention ==
causal weights row 3: [0.446, 0.108, 0.446]
linear: parallel == recurrent: True
KV cache for a made-up 32-layer model at 32k tokens: 4.0 GiB

== 3. constrained decoding ==
valid after 3 attempts: {'name': 'Sarah', 'intent': 'meeting'}
20 constrained samples, all legal: ['follow-up', 'intro', 'meeting']

== 4. context budget ==
{'before': 198, 'steps': ['cleared old tool results', 'summarized 4 messages'], 'after': 64, 'fits': True}
order for 5 docs ranked 1..5: [1, 3, 5, 4, 2]

Read the code in four parts

1. Decoding. softmax turns scores into probabilities and takes a temperature. The two filters zero out tokens and renormalize. decode_step chains them. At temperature 0 it returns the greedy pick. Run the demo and watch entropy rise as temperature rises: a flatter distribution has more uncertainty [5].
2. Attention. attention is softmax(QK^T/√d)V with an optional causal mask [1]. The two linear-attention functions use the feature map elu(x)+1 from [2]. The parallel one builds weights like softmax attention. The recurrent one keeps only two running sums, a matrix S and a vector z. They match, and that is the paper's main point.
3. Constrained output. generate_json is the validate-and-retry loop. It sends the validator's errors back in the next prompt. constrained_decode is different: it sets illegal tokens to −∞ before sampling, using a trie of allowed token paths. This is a toy version of the idea in [22]. The demo uses a model that loves illegal tokens, and it still only ever produces legal output.
4. Context budget. compact first clears old tool results, then summarizes the oldest unpinned messages. Pinned messages and the last few turns are kept [3]. edges_first orders items so the best ones sit at the start and end [21].

Use a real model (optional)

Anything with a complete(prompt) method returning text can replace MockJSONModel in generate_json. Wrap your provider's client in a small class. I tested only the mock path. If your provider offers schema-constrained output [24], try it, and compare how often it still needs a retry.

Exercises

1. See the sharpening. Print softmax(logits, T) for T = 0.2, 0.5, 1, 2, 5. At what T does the top token's share fall below 0.4 for the demo logits? What does entropy do?
2. Top-p versus top-k. Make a very sharp logits list (one big score) and a very flat one. How many tokens does top_p_filter(p=0.9) keep in each? How many does top_k_filter(k=3) keep? Explain the difference in your own words [5].
3. Sample 1,000 times. Count how often each token appears with T = 1 and with T = 0.5. Compare with the printed probabilities.
4. Break causality. Call attention(..., causal=False) and change a later token. Which earlier outputs change? Then try the same with causal=True.
5. Check the RNN claim. Add a counter to linear_attention_recurrent to show the state size stays fixed as the sequence grows. How does that compare with kv_cache_bytes for the same length [2]?
6. Retry budget. Make a MockJSONModel that is right 50% of the time, in a random order. Over 1,000 runs, what is the average number of attempts? (My own arithmetic: if each try succeeds with probability 0.5 and tries are independent, the average is 2.) Does your measured number match?
7. Constrain a number. Extend TrieConstraint to allow only the strings "0" to "99". What breaks if tokens are multi-digit? What does this tell you about why real grammars are built over the model's token vocabulary [22]?
8. Better compaction. Replace mock_summarize with a rule that keeps any message containing "bug" or "decision" in full. Add a test that fails if a "decision" message is lost.
9. Position test. Write 20 short facts and one question. Build prompts with the key fact at position 1, 10 and 20. Use a real model if you have one. Does accuracy change [21]? If you have no model, write down your prediction first.

The tests: test_llm_toolkit.py

import math
import random
import unittest

from llm_toolkit import *

LOGITS = [2.0, 0.5, 0.3, 0.2, -0.5, -1.0]
SCHEMA = {"type": "object", "required": ["name", "intent"], "properties": {
    "name": {"type": "string"},
    "intent": {"type": "string", "enum": ["meeting", "intro", "follow-up"]}}}


class TestDecoding(unittest.TestCase):
    def test_softmax_sums_to_one_and_temperature_sharpens(self):
        lo, hi = softmax(LOGITS, 0.5), softmax(LOGITS, 2.0)
        self.assertAlmostEqual(sum(lo), 1.0)
        self.assertGreater(lo[0], hi[0])
        self.assertLess(entropy(lo), entropy(hi))

    def test_greedy_and_temperature_zero(self):
        self.assertEqual(greedy(softmax(LOGITS)), 0)
        rng = random.Random(1)
        self.assertEqual({decode_step(LOGITS, rng, temperature=0) for _ in range(20)}, {0})

    def test_top_k_keeps_exactly_k(self):
        p = top_k_filter(softmax(LOGITS), 2)
        self.assertEqual(sum(1 for x in p if x > 0), 2)
        self.assertAlmostEqual(sum(p), 1.0)

    def test_top_p_keeps_smallest_nucleus(self):
        p = top_p_filter(softmax(LOGITS), 0.8)
        self.assertEqual([i for i, x in enumerate(p) if x > 0], [0, 1, 2])
        self.assertEqual(sum(1 for x in top_p_filter(softmax(LOGITS), 1.0) if x > 0), len(LOGITS))
        self.assertEqual(sum(1 for x in top_p_filter(softmax(LOGITS), 0.01) if x > 0), 1)

    def test_filtered_sampling_never_picks_removed_tokens(self):
        rng = random.Random(3)
        seen = {decode_step(LOGITS, rng, 1.0, top_k=2) for _ in range(300)}
        self.assertTrue(seen <= {0, 1})

    def test_bad_arguments(self):
        with self.assertRaises(ValueError):
            softmax(LOGITS, 0)
        with self.assertRaises(ValueError):
            top_p_filter([0.5, 0.5], 0)


class TestAttention(unittest.TestCase):
    Q = [[1.0, 0.0], [0.0, 1.0], [2.0, 0.0]]
    K = [[1.0, 0.0], [0.0, 1.0], [1.0, 0.0]]
    V = [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]]

    def test_rows_sum_to_one(self):
        _, w = attention(self.Q, self.K, self.V)
        for row in w:
            self.assertAlmostEqual(sum(row), 1.0)

    def test_causal_mask_blocks_the_future(self):
        _, w = attention(self.Q, self.K, self.V, causal=True)
        self.assertEqual(w[0][1:], [0.0, 0.0])
        self.assertEqual(w[1][2], 0.0)
        _, w2 = attention(self.Q, self.K, self.V, causal=False)
        self.assertGreater(w2[0][2], 0.0)

    def test_changing_a_future_token_does_not_change_the_past(self):
        base, _ = attention(self.Q, self.K, self.V)
        K2 = [r[:] for r in self.K]; K2[2] = [9.0, 9.0]
        V2 = [r[:] for r in self.V]; V2[2] = [9.0, 9.0]
        new, _ = attention(self.Q, K2, V2)
        self.assertEqual(base[:2], new[:2])

    def test_linear_attention_parallel_equals_recurrent(self):
        a = linear_attention_parallel(self.Q, self.K, self.V)
        b = linear_attention_recurrent(self.Q, self.K, self.V)
        for r, s in zip(a, b):
            for x, y in zip(r, s):
                self.assertAlmostEqual(x, y)

    def test_kv_cache_arithmetic(self):
        self.assertEqual(kv_cache_bytes(32, 8, 128, 32768), 4 * 2**30)
        self.assertEqual(kv_cache_bytes(1, 1, 1, 10, bytes_per_value=2, batch=3), 2 * 10 * 2 * 3)


class TestConstrained(unittest.TestCase):
    def test_validate_catches_each_kind_of_error(self):
        self.assertEqual(validate({"name": "A", "intent": "intro"}, SCHEMA), [])
        self.assertIn("missing required key 'intent'", validate({"name": "A"}, SCHEMA)[0])
        self.assertIn("not one of", validate({"name": "A", "intent": "chat"}, SCHEMA)[0])
        self.assertIn("expected string", validate({"name": 5, "intent": "intro"}, SCHEMA)[0])
        self.assertTrue(validate(True, {"type": "integer"}))

    def test_extract_json_handles_chatter_and_fences(self):
        self.assertEqual(extract_json('Sure! {"a": {"b": "}"}} thanks'), {"a": {"b": "}"}})
        self.assertEqual(extract_json('```json\n{"a": 1}\n```'), {"a": 1})
        with self.assertRaises(ValueError):
            extract_json("no json here")

    def test_retry_loop_recovers_and_feeds_errors_back(self):
        m = MockJSONModel(['{"name": "S", "intent": "chat"}', 'oops', '{"name": "S", "intent": "intro"}'])
        value, attempts = generate_json(m, "task", SCHEMA)
        self.assertEqual((value["intent"], attempts), ("intro", 3))
        self.assertIn("rejected", m.prompts[1])
        self.assertIn("not one of", m.prompts[1])

    def test_retry_loop_gives_up(self):
        m = MockJSONModel(["never json"])
        with self.assertRaises(RuntimeError):
            generate_json(m, "task", SCHEMA, max_retries=2)
        self.assertEqual(m.calls, 3)

    def test_constrained_decoding_is_always_legal(self):
        vocab = ["meet", "ing", "in", "tro", "follow", "-up", "chat", "xyz"]
        table = {"meeting": ["meet", "ing"], "intro": ["in", "tro"], "follow-up": ["follow", "-up"]}
        cons = TrieConstraint(list(table), lambda s: table[s])
        loves_illegal = lambda prefix: [0.0] * 6 + [9.0, 9.0]
        rng = random.Random(5)
        outs = {constrained_decode(loves_illegal, vocab, cons, rng) for _ in range(100)}
        self.assertTrue(outs <= set(table))
        self.assertGreater(len(outs), 1)


class TestContext(unittest.TestCase):
    def msgs(self):
        return [{"role": "system", "text": "be careful always", "pinned": True},
                {"role": "user", "text": "fix the test"},
                {"role": "tool", "text": "x " * 50},
                {"role": "assistant", "text": "reading the file now"},
                {"role": "tool", "text": "y " * 40},
                {"role": "assistant", "text": "found the bug"},
                {"role": "tool", "text": "z " * 5},
                {"role": "user", "text": "now add a test"}]

    def test_no_change_when_it_fits(self):
        m = self.msgs()
        new, rep = compact(m, limit=1000)
        self.assertEqual(new, m)
        self.assertEqual(rep["steps"], [])

    def test_clearing_tool_results_keeps_the_last_two(self):
        out = clear_old_tool_results(self.msgs())
        tools = [m["text"] for m in out if m["role"] == "tool"]
        self.assertEqual(tools[0], "[tool result cleared]")
        self.assertNotEqual(tools[1], "[tool result cleared]")

    def test_compaction_fits_and_keeps_pinned_and_recent(self):
        m = self.msgs()
        new, rep = compact(m, limit=40, reserve_output=5)
        self.assertTrue(rep["fits"])
        self.assertLess(rep["after"], rep["before"])
        self.assertEqual(new[0]["text"], "be careful always")
        self.assertEqual(new[-1]["text"], "now add a test")
        self.assertTrue(any(x["text"].startswith("Summary:") for x in new))

    def test_edges_first_puts_best_at_the_ends(self):
        out = edges_first([1, 2, 3, 4, 5])
        self.assertEqual((out[0], out[-1]), (1, 2))
        self.assertEqual(out[len(out) // 2], 5)
        self.assertEqual(sorted(out), [1, 2, 3, 4, 5])


class TestSampling(unittest.TestCase):
    def test_coverage_and_vote(self):
        self.assertAlmostEqual(coverage(0.2, 1), 0.2)
        self.assertAlmostEqual(coverage(0.2, 10), 1 - 0.8 ** 10)
        self.assertEqual(majority_vote(["a", "b", "b", "c"]), "b")
        self.assertEqual(majority_vote(["x", "y"]), "x")


if __name__ == "__main__":
    unittest.main(verbosity=1)

$ python test_llm_toolkit.py
.....................
----------------------------------------------------------------------
Ran 21 tests in 0.002s

OK

5. Discussion questions

Use these in a study group or a class. Try to answer with an example and a source. Give each about 5 to 10 minutes.

1. Holtzman et al. found that the most likely text is not the best text [5]. Why might "most likely" be a poor goal for stories but a good goal for a tool call? Where is the line?
2. Attention cost is quadratic in context length [1]. Linear attention fixes the cost but weakens recall [2] [20]. For which agent jobs would you accept the weaker recall to get cheaper long contexts? For which not?
3. InstructGPT showed a 1.3B model preferred over a 175B base model [8]. What does this say about choosing a model for your agent by parameter count alone?
4. DPO removes the separate reward model and the RL loop [9]. What do you gain and lose? Why might a team still use a verifiable-reward setup for code [26]?
5. Brown et al. saw large gains from repeated sampling when a verifier exists, and plateaus when it does not [13]. Name an agent task where you can build a cheap verifier, and one where you cannot. What would you do in the second case?
6. Snell et al. say test-time compute should be spent per prompt, based on difficulty [14]. How could an agent estimate difficulty at run time? What if it guesses wrong?
7. OpenAI says Structured Outputs guarantees the schema but can still contain mistakes [24]. Write down three errors a perfectly valid JSON reply can still contain. Which of them would your code catch?
8. Retry loops and constrained decoding both give valid output. Compare cost, latency, and what each does to the model's answer quality. When do you want both [22]?
9. Anthropic says to find the smallest set of high-signal tokens [3]. Take a real agent prompt from your work. Cut it in half. What did you remove, and how would you test that nothing important was lost?
10. "Lost in the middle" says position matters [21]. If you must pass 30 retrieved documents, how would you order them? Would you pass 30 at all?
11. You can compact a context, write notes to a file, or split work across sub-agents [3]. Pick one long task and choose between them. State what you would risk losing in each.

6. Quiz (10 questions)

Pick one answer for each question, then press Check answers. Check answers 0 of 10 answered.

1. A model gives each of 200 tokens a probability of 0.01. Why do real systems add log probabilities instead of multiplying the probabilities [26]?
   • Logs make each probability larger, so the model is more confident
   • Probabilities cannot legally be multiplied in a neural network
   • The product is so small (about 10⁻⁴⁰⁰) that a 64-bit float would underflow to 0, but the sum of logs stays a usable number
   • Logs make the model generate more accurate text
2. What happens to the next-token distribution when you lower the temperature?
   • It gets sharper, so the top token gets more of the probability, and in the limit it becomes greedy decoding
   • It gets flatter, so rare tokens are chosen more often
   • The model's weights change
   • The context window shrinks
3. What is the main difference between top-k and top-p (nucleus) sampling [5]?
   • Top-k uses randomness. Top-p does not.
   • They are the same method with different names
   • Top-p keeps the least likely tokens
   • Top-k always keeps k tokens. Top-p keeps the smallest set whose probabilities add up to p, so the set size changes from step to step.
4. Why does a text-generating Transformer use a causal mask in self-attention [1]?
   • To make the attention scores larger
   • To stop a token from using information from later tokens, which will not exist at generation time
   • To remove the need for positional information
   • To reduce the vocabulary size
5. Which statement best describes linear attention as shown in "Transformers are RNNs" [2]?
   • It keeps the full n-by-n softmax table but uses less memory for it
   • It stores a key and a value for every past token, forever
   • It replaces softmax with a feature map so keys and values can be summed into a fixed-size running state, giving O(N) cost and an RNN-like form
   • It needs no queries
6. My own arithmetic setup: a made-up model has 24 layers, 8 KV heads, head size 64, stores values in 2 bytes, and holds 1,000 tokens. The KV cache is 2 (K and V) × layers × heads × head size × bytes × tokens. About how big is it?
   • About 49 KB
   • About 24.6 MB
   • About 49 GB
   • About 49 MB
7. What did the InstructGPT paper report about model size and human preference [8]?
   • Outputs from the 1.3B InstructGPT model were preferred to those from the 175B GPT-3, which has 100 times more parameters
   • Larger models were always preferred, with no exceptions
   • Human feedback made models less truthful
   • Human feedback was only useful for translation
8. Which is the best summary of the finding in Snell et al. on test-time compute [14]?
   • Always sample as many answers as you can afford
   • Test-time compute never helps
   • The best way to spend test-time compute depends on how hard the prompt is, and an adaptive per-prompt plan beat plain best-of-N by over 4 times in efficiency
   • Test-time compute only works for translation
9. An API offers both "JSON mode" and "Structured Outputs". What is the key difference, and what must you still handle in code [24]?
   • There is no difference. Both always return the right answer.
   • JSON mode ensures valid JSON but not your schema. Structured Outputs follows the schema. You still must handle refusals, truncated replies and wrong values.
   • JSON mode follows the schema. Structured Outputs only returns text.
   • Structured Outputs removes the need to test your agent
10. What is the guiding principle of context engineering in Anthropic's article [3]?
    • Fill the whole context window so the model has every fact
    • Always use the longest system prompt you can write
    • Never use tools, because tool results use tokens
    • Find the smallest possible set of high-signal tokens that makes the desired outcome most likely

Answer key

Scoring: 9 to 10 means you are ready for Week 3. 7 to 8: re-read the sections linked in the answers. Below 7: redo the reading guide questions first. When you print this page, check the quiz first if you want your marks on the paper copy. The answer key is included when you print.

7. Graded assignment and rubric

Title: Decode, constrain, budget: three small experiments.
Time: 1 to 2 hours. Total: 100 points. Work: solo.

The task

Extend the lab code and write a short memo (one to two pages). Mock models are fine. You may use a real model if you have a key, but say clearly which parts used it.

Part A. Beam search versus greedy (30 points)

• Add a beam_search function. Give it a small made-up table of next-token probabilities (a dictionary that maps the last token to a distribution), a beam width, and a length. It should return the best sequence by total log probability [26].
• Build an example table where greedy gives a lower total probability than beam search. Show both sequences and both scores (use sums of logs).
• Add at least three unit tests: width 1 equals greedy, a wider beam is never worse in total log probability on your table, and an edge case (for example length 0 or a one-token vocabulary).
• In two sentences, say why beam search costs more time and memory, and why Holtzman et al. would warn against maximizing likelihood for open-ended text [5].

Part B. A stronger JSON loop (30 points)

• Write eight mock replies that fail in different ways: chatter around the JSON, a code fence, a missing key, a wrong enum value, a wrong type, cut-off JSON, a refusal sentence, and one correct reply.
• Run each one through generate_json. Report in a table: attempts used and the final result.
• Extend the loop so it detects a refusal sentence (such as "I can't help with that") and stops with a clear error. It should not waste retries on it [24]. Add tests.
• Report the expected number of attempts if each try is valid with probability s, and tries are independent. State it as my own arithmetic and show your working. Compare it with a measured average over 1,000 mock runs. Say why the two may differ.

Part C. A context budget memo (25 points)

• Write a 30-message trace for an imagined coding agent. Mark at least five messages as important (the task, a decision, a bug found, and so on). Mark some as pinned.
• Run compact with a limit that forces both steps (clearing, then summarizing). Show the report.
• Write a test that checks every important fact is still present after compaction. If the toy summarizer loses one, improve the summarizer or the policy until the test passes. Describe what you changed and why [3].
• Choose between compaction, notes and sub-agents for a four-hour migration task. Give three sentences of reasons and one risk [3].

Part D. Settings and honesty (15 points)

• For each of three tasks (a JSON tool call, a bug fix with a test suite, a brainstorm of ten product names), choose temperature, top-p, and whether to use repeated sampling. Give one reason each, with a source from this week [5] [13] [14].
• State plainly what your experiments do not show. For example: the mock distributions are made up, and word count is not a real token count.

What to hand in

• Your llm_toolkit.py (or your own file names) and tests. All tests must pass.
• A memo in PDF or Markdown with the tables from Parts A to C and your answers to Part D.
• A sources list for any reading you cite.

Rubric

Criterion	Full marks	Partial	Points
A. Beam search	Correct function. An example where beam beats greedy, with scores shown. Three or more passing tests. Clear two-sentence explanation.	Function works but no clear example, or tests are thin.	30
B. JSON loop	Eight varied failure cases. Refusal handling with tests. Table of attempts. Expected-attempts working shown and compared with measurement.	Cases are not varied, or the arithmetic is not compared with data.	30
C. Context budget	Trace has marked important facts. Compaction runs both steps. Fact-survival test passes. Reasoned choice among compaction, notes and sub-agents.	No survival test, or the choice has no reasons or risk.	25
D. Settings and honesty	Three sensible choices with sources. Limits of the experiments are stated plainly.	Choices are not justified, or limits are not stated.	15
Total			100

Grade bands: 90 to 100 excellent. 75 to 89 good. 60 to 74 needs work. Under 60: redo and resubmit. Plain writing is rewarded: short sentences, no filler. Label your own arithmetic as your own.

8. Glossary

Term	Meaning
Attention	A step where each token builds its output as a weighted mix of other tokens' values. The weights come from comparing queries with keys [1].
Autoregressive	Generating one token at a time, each conditioned on the tokens before it.
Beam search	A decoding method that keeps several candidate sequences and extends the best ones. It costs more time and memory than greedy decoding [26].
Causal mask	A mask that stops each position from attending to later positions [1].
Chain of thought	Having the model write intermediate reasoning steps before the final answer [11].
Compaction	Summarizing a context that is near its limit and restarting with the summary [3].
Constrained decoding	Blocking tokens that would break a required format before each sampling step [22].
Context engineering	Curating and maintaining the best set of tokens in the context during inference [3].
Context rot	The drop in recall as the number of tokens in the context grows [3].
Context window	The set of tokens the model can read in one call, including the prompt, history, tool results and the reply.
Coverage	The share of problems solved by at least one of many samples [13].
Decoder-only	A Transformer that uses only the decoder side of the original design. Most chat and agent models are of this kind [26].
Decoding strategy	The rule for picking the next token from the model's distribution (greedy, temperature, top-k, top-p and others).
DPO	Direct Preference Optimization. A way to tune a model on preference data with a simple classification loss, without a separate reward model or RL loop [9].
Feature map (φ)	A function applied to queries and keys in linear attention. The paper uses elu(x) + 1 [2].
FlashAttention	An exact attention algorithm that is fast because it reduces memory reads and writes on the GPU [17].
Greedy decoding	Always picking the single most likely token.
Hybrid model	A model that mixes layer types, such as full attention and state-space layers [19].
Instruction tuning (SFT)	Supervised training on prompts paired with good answers [8].
JSON mode	An API setting that makes the model emit valid JSON, without guaranteeing a particular schema [24].
KV cache	Stored keys and values of earlier tokens, reused so each new token needs only its own new ones. It grows with context length [26].
Linear attention	Attention rewritten with a feature map so its cost is linear in sequence length. Its causal form is an RNN with a fixed-size state [2].
Logit	A raw score for a token, before softmax turns the scores into probabilities.
Nucleus (top-p) sampling	Sampling from the smallest set of top tokens whose total probability reaches p [5].
PagedAttention	A way to manage the KV cache in pages, like virtual memory, to cut wasted memory [16].
Prefill and decode	The two phases of inference. Prefill reads the whole prompt at once. Decode produces tokens one by one [26].
Query, key, value	The three vectors made from each token for attention: what it looks for, what it offers, and what it passes on [1].
RLHF	Reinforcement learning from human feedback. Humans rank outputs, a reward model learns the ranking, and RL tunes the model toward it [8].
RLVR	Reinforcement learning with verifiable rewards. The reward is a program check, such as unit tests or exact match [26].
Scaling law	A smooth rule linking loss to model size, data and compute [6] [7].
Self-consistency	Sampling many reasoning paths and keeping the most common answer [12].
Speculative decoding	A small draft model guesses tokens and the large model checks them in parallel, giving the same output distribution faster [15].
State-space model (SSM)	A sequence layer that carries a fixed-size state forward. Mamba is one example [18].
Structured Outputs	An API feature that makes replies follow a supplied JSON Schema [24].
Temperature	A number that divides the logits before softmax. Low is sharper. High is flatter.
Test-time scaling	Spending more compute per question at answer time, for example with longer reasoning or more samples [14].
Token	A piece of text the model reads and writes: a word, a part of a word, or a symbol.
Top-k	Keeping only the k most likely tokens before sampling.

9. References

Every link below was opened and checked on October 28, 2026. Items 1 to 3 are the required readings for this week. Items 4 to 24 are sources I added from my own research. Items 25 and 26 are the course pages. For papers on arXiv I give the arXiv identifier and the year it was first posted. I give a venue only where the course page names one.

• [1] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. "Attention Is All You Need." NeurIPS, 2017. arXiv:1706.03762. https://arxiv.org/abs/1706.03762 (required)
• [2] Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention." ICML, 2020. arXiv:2006.16236. https://arxiv.org/abs/2006.16236 (required)
• [3] Rajasekaran, P., Dixon, E., Ryan, C., and Hadfield, J. (Anthropic Applied AI team). "Effective Context Engineering for AI Agents." Anthropic, 2025. https://www.anthropic.com/engineering/effective-context-engineering-for-ai-agents (the address redirects to the same article under anthropic.com/research) (required)
• [4] Alammar, J. "The Illustrated Transformer." Blog post. https://jalammar.github.io/illustrated-transformer/ (added)
• [5] Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y. "The Curious Case of Neural Text Degeneration." arXiv:1904.09751, 2019. https://arxiv.org/abs/1904.09751 (added)
• [6] Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. "Scaling Laws for Neural Language Models." arXiv:2001.08361, 2020. https://arxiv.org/abs/2001.08361 (added)
• [7] Hoffmann, J., Borgeaud, S., Mensch, A., and 19 others. "Training Compute-Optimal Large Language Models." arXiv:2203.15556, 2022. https://arxiv.org/abs/2203.15556 (added)
• [8] Ouyang, L., Wu, J., Jiang, X., and 17 others. "Training Language Models to Follow Instructions with Human Feedback." arXiv:2203.02155, 2022. https://arxiv.org/abs/2203.02155 (added)
• [9] Rafailov, R., Sharma, A., Mitchell, E., Ermon, S., Manning, C. D., and Finn, C. "Direct Preference Optimization: Your Language Model is Secretly a Reward Model." arXiv:2305.18290, 2023. https://arxiv.org/abs/2305.18290 (added)
• [10] Shao, Z., Wang, P., Zhu, Q., and 8 others. "DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models." arXiv:2402.03300, 2024. https://arxiv.org/abs/2402.03300 (added)
• [11] Wei, J., Wang, X., Schuurmans, D., Bosma, M., Ichter, B., Xia, F., Chi, E., Le, Q., and Zhou, D. "Chain-of-Thought Prompting Elicits Reasoning in Large Language Models." arXiv:2201.11903, 2022. https://arxiv.org/abs/2201.11903 (added)
• [12] Wang, X., Wei, J., Schuurmans, D., Le, Q., Chi, E., Narang, S., Chowdhery, A., and Zhou, D. "Self-Consistency Improves Chain of Thought Reasoning in Language Models." arXiv:2203.11171, 2022. https://arxiv.org/abs/2203.11171 (added)
• [13] Brown, B., Juravsky, J., Ehrlich, R., Clark, R., Le, Q. V., Ré, C., and Mirhoseini, A. "Large Language Monkeys: Scaling Inference Compute with Repeated Sampling." arXiv:2407.21787, 2024. https://arxiv.org/abs/2407.21787 (added)
• [14] Snell, C., Lee, J., Xu, K., and Kumar, A. "Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters." arXiv:2408.03314, 2024. https://arxiv.org/abs/2408.03314 (added)
• [15] Leviathan, Y., Kalman, M., and Matias, Y. "Fast Inference from Transformers via Speculative Decoding." arXiv:2211.17192, 2022. https://arxiv.org/abs/2211.17192 (added)
• [16] Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. "Efficient Memory Management for Large Language Model Serving with PagedAttention." arXiv:2309.06180, 2023. https://arxiv.org/abs/2309.06180 (added)
• [17] Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." arXiv:2205.14135, 2022. https://arxiv.org/abs/2205.14135 (added)
• [18] Gu, A., and Dao, T. "Mamba: Linear-Time Sequence Modeling with Selective State Spaces." arXiv:2312.00752, 2023. https://arxiv.org/abs/2312.00752 (added)
• [19] Lieber, O., Lenz, B., Bata, H., and 19 others. "Jamba: A Hybrid Transformer-Mamba Language Model." arXiv:2403.19887, 2024. https://arxiv.org/abs/2403.19887 (added)
• [20] Hsieh, C.-P., Sun, S., Kriman, S., Acharya, S., Rekesh, D., Jia, F., Zhang, Y., and Ginsburg, B. "RULER: What's the Real Context Size of Your Long-Context Language Models?" arXiv:2404.06654, 2024. https://arxiv.org/abs/2404.06654 (added)
• [21] Liu, N. F., Lin, K., Hewitt, J., Paranjape, A., Bevilacqua, M., Petroni, F., and Liang, P. "Lost in the Middle: How Language Models Use Long Contexts." arXiv:2307.03172, 2023. https://arxiv.org/abs/2307.03172 (added)
• [22] Willard, B. T., and Louf, R. "Efficient Guided Generation for Large Language Models." arXiv:2307.09702, 2023. https://arxiv.org/abs/2307.09702 (added)
• [23] dottxt-ai. "Outlines: Structured Outputs." Open-source code repository, GitHub. https://github.com/dottxt-ai/outlines (added)
• [24] OpenAI. "Structured model outputs." OpenAI API documentation (opened October 28, 2026). https://platform.openai.com/docs/guides/structured-outputs (redirects to developers.openai.com/api/docs/guides/structured-outputs) (added)
• [25] Yang, D., Ryan, M., and Yang, J. (instructors). "CS329Z: Engineering AI Agents." Stanford University, Fall 2026. Course site and schedule: https://cs329z.stanford.edu/
• [26] Yang, D., Ryan, M., and Yang, J. "LLMs for Builders." CS329Z lecture slides (Lecture 2, September 28), Stanford University. https://web.stanford.edu/class/cs329z/slides/lecture02.pdf. Used only to check topic coverage and to point to what the lecture covers. This pack does not reproduce the slides.