KV Cache Growth, PRMs vs ORMs & Exercises

The KV cache wanted to grow. It wanted to grow at every token, on every layer, in every head. I told it about quantization. It grew anyway, just more slowly.

KV, Cache-Aware AI Agent

Fun Fact: Why ORMs Lost the Vibes War

Outcome Reward Models scored only the final answer, so they often gave full credit to lucky guesses and zero credit to a near-miss with perfect intermediate work. Process Reward Models, by scoring each step, turned what looked like 'this kid is bad at math' into 'this kid forgot to carry the one'. The pedagogy was already known to every grade-school teacher in the world. The interesting thing is that it took a few billion parameters of reasoning research to rediscover it.

Having framed why test-time compute matters in Section 8.1, we now look at the two engineering realities that govern production deployment: how the KV cache explodes under long reasoning traces, and how PRMs and ORMs differ in how they score those traces.

8.1.7 KV cache Growth in Reasoning Models

One of the most consequential engineering challenges with reasoning models is memory. Every token the model generates, including every thinking token, must be stored in the KV cache during the generation pass. For standard chat models generating 200 to 500 tokens, this is manageable. For reasoning models generating 5,000 to 50,000 thinking tokens, it becomes a serious infrastructure problem. (For the technical mechanics of KV cache and its relationship to inference optimization, see Section 9.3.)

8.1.7.1 The Memory Arithmetic

The KV cache stores key and value tensors for every token in the context. For a transformer with $L$ layers, $H$ key-value heads, and head dimension $d_h$, the KV cache memory per token is:

The factor of 2 accounts for both keys and values. For a model using Grouped Query Attention (GQA), the number of KV heads is smaller than the number of query heads, which reduces this cost proportionally.

8.1.7.2 Serving Implications

These numbers have direct consequences for how reasoning model deployments are architected:

• Throughput collapse: A GPU that serves 200 concurrent standard chat requests can serve as few as 2 to 5 concurrent reasoning requests at maximum thinking depth. This means the effective cost per reasoning request is not just the token count difference but also the opportunity cost of blocked GPU capacity.
• Prefill/decode phase imbalance: In standard inference, the prefill phase (processing the input prompt) is fast, and the decode phase (generating output tokens one by one) is the bottleneck. In reasoning models, the decode phase is enormously extended by thinking tokens, making decode the dominant cost by far.
• Memory-bound vs. compute-bound: With tens of thousands of tokens in the KV cache, attention computation becomes memory-bandwidth-bound. The GPU spends most of its time reading KV cache tensors from HBM rather than performing matrix multiplications. FlashAttention and GQA (see Section 9.3) are therefore especially valuable for reasoning models.

8.1.7.3 Mitigation Strategies

Production deployments use several strategies to manage KV cache growth in reasoning models:

• KV cache eviction after reasoning completes: Once the model emits the separator token (transitioning from thinking to answer), the thinking tokens are no longer needed for future generation. Their KV cache entries can be freed, recovering GPU memory before serving the next request. This is not possible during generation, but can happen between requests.
• Thinking token compression: Research systems have explored compressing or summarizing the KV cache of thinking tokens into a smaller representation. This is an active research area and not yet widely deployed.
• Budget caps: Most reasoning model APIs expose a "reasoning budget" or "thinking effort" parameter. Setting thinking_tokens=1024 caps the maximum KV cache growth and prevents runaway costs on simple queries. OpenAI's o-series exposes reasoning_effort (low, medium, high); DeepSeek R1 supports explicit max_tokens caps on the thinking section. Always set a budget cap for production deployments.
• Routing: The most effective mitigation is simply not sending easy queries to reasoning models. A difficulty router (see Section 4.2) prevents most requests from incurring thinking token overhead at all.

import anthropic
client = anthropic.Anthropic()
# Using Claude's extended thinking with a budget cap to limit KV cache growth.
# budget_tokens limits the maximum thinking tokens, directly capping memory usage.
response = client.messages.create(
    model="claude-opus-4-5",
    max_tokens=16000,
    thinking={
    "type": "enabled",
    "budget_tokens": 2048 # cap thinking tokens: reduces memory from ~640 MB to ~26 MB
    # for a 70B-equivalent model. Default can be up to 10000.
    },
    messages=[{
    "role": "user",
    "content": "Prove that the square root of 2 is irrational."
    }]
    )
# Response includes both thinking blocks and text blocks
for block in response.content:
    if block.type == "thinking":
        print(f"Thinking ({len(block.thinking)} chars): {block.thinking[:200]}...")
    elif block.type == "text":
        print(f"Answer: {block.text}")

8.1.8 Process Reward Models vs. Outcome Reward Models

Every reasoning strategy that involves evaluation (best-of-N, tree search, and RL training of reasoning models) requires a reward model. Not all reward models are created equal. The choice between an outcome reward model (ORM) and a process reward model (PRM) fundamentally affects what can be achieved.

8.1 Outcome Reward Models

An ORM scores the final answer only. Given a problem and a complete reasoning trace plus final answer, the ORM outputs a single scalar: how good is this answer? ORMs are relatively easy to train because labels are at the answer level: for math problems, you can automatically verify whether the final numerical answer is correct. For code, you can run tests. The feedback is binary and abundant.

The limitation of ORMs is sparse signal. If a 50-step reasoning chain arrives at the wrong answer, the ORM assigns a low score to the entire chain, but provides no information about which of the 50 steps went wrong. This makes ORMs poorly suited for tree search (you cannot prune a bad branch until you reach the end of it) and for training (the gradient signal does not flow back to the specific reasoning step that caused the error).

8.1.8.2 Process Reward Models

A PRM scores each individual reasoning step. Given a problem and the reasoning trace up to step $k$, the PRM outputs a score for whether step $k$ was correct and productive. This provides dense, step-level supervision: the model receives a signal at every step, not just at the end.

The foundational work is Lightman et al. (2023), "Let's Verify Step by Step," which introduced the PRM800K dataset: 800,000 step-level human annotations on math reasoning traces. Their key finding was that PRMs substantially outperform ORMs on hard math benchmarks, particularly when used for best-of-N selection. Given N candidate solutions, a PRM can score each step of each solution and select the one with the highest minimum step score (the weakest link), which is a more reliable quality signal than the final-answer score alone.

8.1.8.3 Automated PRM Training

The cost of step-level human annotation has motivated research into automated PRM training. Two approaches have shown promise:

• Monte Carlo rollout estimation: For a given partial reasoning trace ending at step $k$, generate many complete solutions (rollouts) from step $k$ forward. The fraction of rollouts that arrive at the correct final answer estimates the probability that step $k$ is on a path to success. This probability is used as the step-level label. The DeepSeek-R1 team used this technique to generate PRM training data at scale.
• LLM-as-critic: Use a separate language model to evaluate whether each reasoning step is logically valid. This is cheaper than human annotation but introduces the risk of the critic having the same biases as the generator. Best practice is to use a different model family for the critic than for the generator.

For the RL training context (how PRMs and ORMs connect to PPO, GRPO, and DPO during reasoning model training), see Section 18.1.

Exercises