---
date: '2025-10-02'
description: prefill and decode constraints with KV caching and sampling
id: '7'
modified: 2026-06-07 01:18:28 GMT-04:00
seealso:
  - '[[thoughts/Transformers|Transformers]]'
  - '[[thoughts/LLMs|LLMs]]'
  - '[[thoughts/vllm|vLLM]]'
  - '[[thoughts/Continuous batching|continuous batching]]'
  - '[[thoughts/paged attention|paged attention]]'
  - '[[thoughts/prefix caching|prefix caching]]'
  - '[[thoughts/KV compression|KV compression]]'
  - '[[thoughts/Speculative decoding|speculative decoding]]'
  - '[[thoughts/quantization|quantization]]'
socials:
  link: https://tsfm.ca/lecture-seven
tags:
  - ml
  - tsfm
title: 'lecture seven: inference and sampling'
created: '2025-10-02'
published: '2025-10-02'
pageLayout: default
slug: thoughts/tsfm/7
permalink: https://aarnphm.xyz/thoughts/tsfm/7.md
generator:
  quartz: v4.6.0
  hostedProvider: Cloudflare
  baseUrl: aarnphm.xyz
full: https://aarnphm.xyz/llms-full.txt
---
## inference

autoregressive inference has two phases. during prefill, the model processes the prompt tokens in parallel and writes their keys and values to the KV cache. during decode, the model generates one token at a time and reads the cache for the earlier context \[@scalingbook; @280922\].

let $s$ be the current context length, $L$ the number of layers, $H_{\mathrm{KV}}$ the number of key-value heads, $d_h$ the head dimension, $B$ the batch size, and $b$ the bytes per cache element. the KV cache occupies

$$
M_{\mathrm{KV}}=2bBLH_{\mathrm{KV}}sd_h.
$$

the factor of $2$ accounts for keys and values. multi-query and grouped-query attention reduce $H_{\mathrm{KV}}$, which shrinks the cache before paging or compression.

## arithmetic intensity

arithmetic intensity is the number of floating-point operations performed per byte moved to or from memory.

$$
\mathrm{AI}=\frac{\mathrm{FLOPs}}{\mathrm{bytes}}.
$$

for hardware with peak compute $C$ FLOP/s and memory bandwidth $W$ bytes/s, the roofline threshold is

$$
\mathrm{AI}_{\mathrm{threshold}}=\frac{C}{W}.
$$

a kernel is compute-bound above this threshold and memory-bound below it. the lecture panel uses an H100 value of $989.5$ TFLOP/s and $3.35$ TB/s, which gives

$$
\frac{989.5}{3.35}\approx 295\ \text{FLOPs/byte}.
$$

for attention with $q$ query tokens and $s$ cached tokens, a compact estimate is

$$
\mathrm{AI}_{\mathrm{attn}}\approx\frac{sq}{s+q}.
$$

during prefill, $q\approx s$, so one block of prompt tokens reuses the same weights and attention data. during decode, $q=1$, so attention reads a long cache to produce one token. decode therefore tends to be limited by memory bandwidth at small batches. larger batches reuse the model weights across more requests, while each request still brings its own KV cache.

## sampling

given logits $z_i$ and temperature $\tau$, the next-token distribution is

$$
p_i(\tau)=\frac{\exp(z_i/\tau)}{\sum_j\exp(z_j/\tau)}.
$$

top-$k$ sampling keeps the $k$ tokens with the largest probabilities. top-$p$ sampling keeps the smallest probability-sorted set whose cumulative probability is at least $p$ \[@holtzman2020curiouscaseneuraltext\]. the sampler renormalizes the remaining probabilities before drawing one token.

![[thoughts/quantization#floating point]]