Autoregressive Streaming Inference and KV-Cache Memory Optimization

5.2 The Fatal Trap of Naive Inference: The $O(N^2)$ Goldfish Memory Disaster

How do beginners typically implement this loop? They pass the entire growing sentence back into the model on every single token step:

5.2.1 The Typist with Compulsive Amnesia

Picture a typist with severe compulsive amnesia:

• To type the 1st word, he reads the prompt and types 1 word;

===================================================================================
                       朴素自回归推理的重复计算噩梦
===================================================================================
生成第 1 个字 ("公"):
  打字员朗读上文: "今天天气真好，我们一起去"                  ──> 算了 11 个字的特征

生成第 2 个字 ("园"):
  打字员从头重读: "今天天气真好，我们一起去公"                ──> 算了 12 个字的特征

生成第 3 个字 ("散"):
  打字员再次从头重读: "今天天气真好，我们一起去公园"          ──> 算了 13 个字的特征
  ...
生成第 100 个字:
  打字员又从头重读: [ 前面全部 110 个字全部重读一遍！ ]       ──> 算了 111 个字的特征
===================================================================================

5.2.2 The Horrific Computational Bill

When generating a 500-token poem or essay, naive re-reading requires computing token representations:

A typist with standard memory only processes the single newly arrived token at each step, requiring merely $500$ evaluations! The naive approach burns a staggering 260 times more computation for zero gain!

| Inference Architecture | Input Provided per Iteration | Time Complexity | Compute for 500 Tokens | User Experience | | Naive Uncached Inference | Full historical context (entire sequence of length $T$) | $O(N^2)$ (Quadratic Explosion) | $130,500$ forward token passes | Severe lag as output grows; pauses between tokens stretch to seconds. | | KV-Cache Accelerated (Production) | Only the single newest generated token (length strictly 1) | $O(N)$ (Linear Constant) | Strictly $500$ forward token passes (260x speedup!) | Silky smooth, instantaneous streaming regardless of sequence length. |

This explains why amateur inference scripts run fast for the first 5 tokens but grind GPUs to a halt as sequences lengthen.

5.3 The Breakthrough: Intuitive KV-Cache Sticky Note Analogy

How do we eliminate this redundant recomputation? The golden industry solution is KV-Cache (Key-Value Cache).

5.3.1 Why Do Historical Tokens Never Need Recomputation?

Recall the self-attention formula from Chapter 02:

Reflect on what these three matrices represent:

• $Q$ (Query): What the current token wants to ask;
• $K$ (Key): What semantic tags each historical token advertises;
• $V$ (Value): What substantive information each historical token carries.

When generating the 100th token:

• What is new? Only the 100th token was just generated, so we solely need to compute Query $q_{100}$!
• What is old? The preceding 99 tokens are already etched in stone! The keys and values ($k_1, v_1, \dots, k_{99}, v_{99}$) were calculated seconds ago!
• The Iron Law of Causality: Because the model is causal autoregressive, future tokens cannot alter past representations. Therefore, $k_1$ and $v_1$ for token 1 will never change across the entire timeline!

5.3.2 The Sticky-Note Ledger Metaphor

Since historical $K$ and $V$ vectors are immutable, why recompute them? We simply store them in a persistent buffer (the Cache), appending newly computed $K$ and $V$ vectors like sticky notes into a ledger!

The Miracle Occurs: Whether generating token 10 or token 1,000, every forward step evaluates exactly 1 token! Typewriter latency remains perfectly constant throughout!

5.3.3 The VRAM Ledger: How Much Memory Does 0.04B KV-Cache Actually Consume?

Engineers often hear that KV-Cache inflates VRAM. For our 0.04B golden architecture, how much memory is actually required for a full 1,024-token context? Let us calculate the exact ledger for Batch Size = 1:

• Total network layers: $n_{\text{layers}} = 12$;
• Using Grouped-Query Attention (GQA), each layer has only $n_{kv} = 2$ KV heads (slashing KV memory by 75% compared to 8-head MHA!);
• Head feature dimension: $d_{\text{head}} = 64$;
• Half-precision floating point (fp16, 2 bytes per float);
• Both $K$ and $V$ tensors must be cached.

KV Cache Size Generated per Token:

Total VRAM Footprint for Full 1,024-Token Context:

[!TIP]

A Truly Stunning Result: Even when generating a complete 1,024-token essay, the entire KV-Cache for our 0.04B model requires merely 6.29 MB of VRAM! Slashed by 75% relative to standard Multi-Head Attention (25.2 MB), it runs at blazing speeds on mobile phones, Raspberry Pis, and embedded edge devices with zero pressure!

5.4 Decoding Sampling Strategies: Guiding the Model to Speak Coherently

When the neural network finishes computing, it outputs not characters, but an array of 4,096 unnormalized numbers called Logits. Converting these 4,096 raw scores into the chosen token is determined entirely by the Sampling Strategy.

Before inspecting line-by-line derivations, let us review the four golden hyperparameters:

| Sampling Hyperparameter | Recommended Value | Core Governing Physics | High vs. Low Setting Comparison | | Repetition Penalty (repetition_penalty) | 1.1 ~ 1.2 | Discounting logits of recently generated tokens | 1.0: No penalty, easily trapped in infinite repetition loops;<br>1.2: Effectively breaks repetitive cycles, boosting vocabulary diversity. | | Sampling Temperature (temperature) | 0.7 ~ 0.8 | Scaling logit margins to control probability flatness | Low (0.1~0.5): Deterministic, conservative, strict rhyme and meter;<br>High (1.0~1.5): Imaginative, creative, but overly high values produce incoherent gibberish. | | Top-K Truncation (top_k) | 40 ~ 50 | Hard cutoff: retaining only the top K highest-scoring tokens | 0: No cutoff;<br>50: Safely amputates the bottom 4,000+ irrelevant low-probability tail tokens. | | Nucleus Sampling (top_p) | 0.8 ~ 0.9 | Dynamic probability mass: keeping candidates up to cumulative probability P | Dynamically adaptive: Expands candidate pool for open prompts; contracts pool for rigid idioms. Superior to static Top-K. |

Now let us walk through a concrete numerical example to dissect every step of the sampling function!

Hypothetical Scenario

Suppose the model's vocabulary contains just 4 candidate words with raw logits:

• "Beijing": 4.0 points
• "Shanghai": 3.0 points
• "Guangzhou": 2.0 points
• "Mars": 0.1 points (unrelated low-score token)

5.4.1 Step 1: Repetition Penalty — Forbidding Endless Echoes

If the model recently generated "Beijing", we must penalize its score:

核心逻辑代码（在 src/generate.py 中，我们还特意对标点符号与特殊符进行豁免，防止标点因惩罚产生交替震荡）：

if repetition_penalty != 1.0 and generated_ids:
    exempt_ids = {0, 1, 2, 3, 14, 262, 263, 273, 274, 275, 534, 1227, 1916}
    for tid in set(generated_ids):
        if tid in exempt_ids:
            continue
        if logits[tid] < 0:
            logits[tid] *= repetition_penalty
        else:
            logits[tid] /= repetition_penalty

If score > 0: logit / 1.2 $\to$ reduces score ($4.0 / 1.2 \approx 3.33$);

• If score < 0: logit * 1.2 $\to$ makes the negative score even more negative ($-2.0 \times 1.2 = -2.4$);
• "北京" 的得分是正数 4.0，执行 4.0 / 1.2 = 3.33！它的得分被生生削弱了；
• Notice: Why not use simple subtraction logit - penalty? Because division scales proportionally with magnitude without destroying relative rankings among other unpenalized tokens!
• This sharply suppresses the probability of repeating recent words, curing repetitive loops.

5.4.2 Step 2: Temperature Scaling — Injecting Creativity vs. Preserving Rigor

The Core Scaling Operation:

logits = logits / temperature

Why does dividing by a single scalar alter model behavior so profoundly? Let us run numbers:

Case A: High Temperature ($T = 2.0$, Unleashing Imagination)

• "Beijing": $4.0 / 2.0 = \mathbf{2.0}$
• "Shanghai": $3.0 / 2.0 = \mathbf{1.5}$
• "Guangzhou": $2.0 / 2.0 = \mathbf{1.0}$
• "Mars": $0.1 / 2.0 = \mathbf{0.05}$
• Impact: The original wide gap between 4 and 2 compresses down to 2 and 1! Probabilities flatten across candidates, encouraging unexpected, highly imaginative combinations (though excessive values risk gibberish).

Case B: Low Temperature ($T = 0.5$, Calm and Deterministic)

• "Beijing": $4.0 / 0.5 = \mathbf{8.0}$
• "Shanghai": $3.0 / 0.5 = \mathbf{6.0}$
• "Guangzhou": $2.0 / 0.5 = \mathbf{4.0}$
• Impact: Gaps are exponentially magnified! Top candidates dominate completely while low-scoring candidates vanish, yielding confident, strict, and highly deterministic outputs.

5.4.3 Step 3: Top-P (Nucleus) Sampling — Why the Right-Shift Is Vital

This is a piece of logic that puzzles many beginners. Examining standard nucleus implementations:

sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

# 下面这两行到底在干嘛？！
sorted_indices_to_remove = cumulative_probs > top_p
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
sorted_indices_to_remove[..., 0] = 0

Let us dissect why this right-shift trick is indispensable:

Concrete Numerical Simulation

Suppose Softmax yields these sorted probabilities:

• "Beijing": Probability 60%
• "Shanghai": Probability 25%
• "Guangzhou": Probability 10%
• "Mars": Probability 5%

We set top_p = 0.8 (80%), intending to sample exclusively within the top 80% credible mass while trimming the long tail ("Mars").

1. Computing Cumulative Probabilities (cumulative_probs):

• "Beijing": $60\%$
• "Shanghai": $60\% + 25\% = \mathbf{85\%}$ (Crossing the 80% threshold!)
• "Guangzhou": $85\% + 10\% = 95\%$
• "Mars": $95\% + 5\% = 100\%$

2. What Happens If We Naively Filter with cumulative_probs > 0.8?

• "Beijing" (60%): $\le 0.8 \to$ Retained
• "Shanghai" (85%): $> 0.8 \to$ Erroneously Flagged for Removal!
• "Guangzhou" (95%): Removed
• "Mars" (100%): Removed
• The Fatal Flaw: "Shanghai" was the exact candidate needed to satisfy our 80% threshold! Removing it leaves only "Beijing" (60%), failing to cover the requested 80% probability mass!

3. The Elegant Right-Shift Solution

To safely preserve the threshold-crossing boundary token ("Shanghai"), engineers shift the removal mask right by one position:

原始累加是否超过 0.8:  [ False,   True,   True,   True ]
                                   ↘       ↘       ↘
向右平移一位后结果:    [   ?,   False,   True,   True ]
强制将第一名设为 False: [ False,  False,   True,   True ]
                         (北京)   (上海)   (广州)  (火星)
                          保留     保留     剔除    剔除！

Behold: "Shanghai" is preserved! The candidate set contains "Beijing" (60%) + "Shanghai" (25%) = 85%, comfortably covering the 80% target while cleanly eliminating tail tokens "Guangzhou" and "Mars"!

sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() # 向右平移
sorted_indices_to_remove[..., 0] = 0 # 无论如何，绝对不能把第一名剔除掉

This demonstrates the mathematical elegance embedded in production AI systems.

5.5 Assembling the Complete Autoregressive Sampling Operator

Combining all components into a modular sampling function:

import torch
import torch.nn.functional as F
from typing import List, Optional

def sample_next_token(
    logits: torch.Tensor,
    temperature: float = 0.8,
    top_k: int = 50,
    top_p: float = 0.9,
    repetition_penalty: float = 1.1,
    generated_ids: Optional[List[int]] = None
) -> int:
    """
    自回归核心采样函数：将模型打分转化为最终选出的单字 ID
    """
    # 保证操作不影响原张量，且展平为一维 [Vocab_Size]
    logits = logits.clone().squeeze()

    # 1. 重复惩罚：削减历史出现过文字的分数，拒绝复读机
    if repetition_penalty != 1.0 and generated_ids:
        for tid in set(generated_ids):
            if logits[tid] < 0:
                logits[tid] *= repetition_penalty
            else:
                logits[tid] /= repetition_penalty

    # 2. 极端贪心搜索保护：如果温度无限接近 0，直接取第一名
    if temperature <= 1e-4:
        return torch.argmax(logits, dim=-1).item()

    # 3. 温度缩放：平滑或陡峭化概率曲线
    logits = logits / temperature

    # 4. Top-K 截断：只保留前 K 个最高分的强力候选
    if top_k > 0:
        top_k = min(top_k, logits.size(-1))
        # 找出第 K 名的分数，凡是比它小的全部赋为负无穷大 (-inf)
        k_th_val = torch.topk(logits, top_k)[0][..., -1, None]
        indices_to_remove = logits < k_th_val
        logits[indices_to_remove] = -float("Inf")

    # 5. Top-P 核采样：动态累加概率，切除长尾荒谬候选
    if top_p < 1.0:
        # 从大到小排序
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        # 计算累加概率和
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

        # 找到超出阈值的截断点
        sorted_indices_to_remove = cumulative_probs > top_p
        # 巧妙右移，保留临界词
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        # 将被剔除的索引还原映射回原 logits，并置为负无穷大
        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        logits[indices_to_remove] = -float("Inf")

    # 6. 转为概率分布，并进行轮盘赌抽奖 (Multinomial 依概率采样)
    probs = F.softmax(logits, dim=-1)
    next_token_id = torch.multinomial(probs, num_samples=1).item()
    return next_token_id

5.6 Implementing the Streaming Generator: TextGenerator

We encapsulate tokenizer, model, KV-Cache, and sampling into a cohesive class: TextGenerator, enabling streaming output in terminal. Source located at src/generate.py.

Implementation Walkthrough:

class TextGenerator:
    def __init__(self, model, tokenizer, device: str = "cpu"):
        self.model = model.to(device).eval() # 切换到评估模式，关闭 Dropout
        self.tokenizer = tokenizer
        self.device = device

    @torch.no_grad() # 推理阶段绝不计算梯度，节约海量显存与计算时间
    def generate(
        self,
        prompt: str,
        max_new_tokens: int = 128,
        temperature: float = 0.8,
        top_k: int = 50,
        top_p: float = 0.9,
        repetition_penalty: float = 1.1,
        stream: bool = True
    ) -> str:
        # 第一步：把人类输入的提示词编码成数字序列
        prompt_ids = self.tokenizer.encode(prompt, add_bos=True)
        generated = list(prompt_ids)

        if stream:
            print(prompt, end="", flush=True)

        past_key_values = None
        input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=self.device)

        # 第二步【首字填充 Prefill 阶段】：
        # 一次性将完整的 Prompt 送入模型，产生初始的 KV-Cache 便签条，并得到第 1 个新字
        logits, _, past_key_values = self.model(input_ids, past_key_values=None, start_pos=0)
        next_token_logits = logits[0, -1, :] # 提取 Prompt 最后一个位置的预测得分

        # 采样出第 1 个生成的词
        next_id = sample_next_token(
            next_token_logits,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            generated_ids=generated
        )

        generated.append(next_id)
        if stream:
            # 实时将这个词翻译回字符，并打在终端屏幕上！
            print(self.tokenizer.decode([next_id], skip_special_tokens=False), end="", flush=True)

        start_pos = len(prompt_ids)

        # 第三步【增量自回归 Decode 阶段】：
        # 依靠 KV-Cache，每次只喂入最新生成的 1 个字！
        for _ in range(max_new_tokens - 1):
            # 如果生成了句子结束符 </s>，说明文章写完了，优雅退出
            if next_id == self.tokenizer.eos_token_id:
                break
            # 如果长度超出了模型最大上限，停止生成
            if start_pos >= self.model.config.max_seq_len:
                break

            # 核心提速秘诀：单字输入！形状仅仅为 [1, 1]
            step_input = torch.tensor([[next_id]], dtype=torch.long, device=self.device)
            logits, _, past_key_values = self.model(
                step_input,
                past_key_values=past_key_values, # 传入并持续追加便利贴
                start_pos=start_pos
            )
            next_token_logits = logits[0, -1, :]

            # 采样下一个字
            next_id = sample_next_token(
                next_token_logits,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                generated_ids=generated
            )

            generated.append(next_id)
            start_pos += 1

            if stream:
                print(self.tokenizer.decode([next_id], skip_special_tokens=False), end="", flush=True)

        if stream:
            print() # 换行收尾

        # 返回全部新生成的文字内容
        return self.tokenizer.decode(generated[len(prompt_ids):])

5.7 Core Takeaways & Hands-on Verification

Through this guide, you have mastered the foundational mechanics powering ChatGPT:

• Essence of Autoregression: A recursive solitaire loop feeding newly generated tokens back as inputs;
• Magic of KV-Cache: Caching immutable past $K, V$ states to conquer the $O(N^2)$ computational explosion;
• The Art of Sampling:
• For rigorous tasks, lower temperature toward 0.1;
• For poetry and creative writing, set temperature to 0.8 with Top-P 0.9 to balance flair and coherence;
• To break repetitive phrasing, enable repetition penalty 1.1.

💡 Hands-on Verification Experiments

Experiment with hyperparameters in scripts/run_generate.py:

• Set temperature = 2.0 and observe how generation becomes wildly divergent;
• Set repetition_penalty = 1.0 (zero penalty) and watch how readily the model falls into repeating lines;

In the next chapter, we integrate the entire system in an end-to-end rehearsal—Chapter 06 | 0.04B Model End-to-End Practice & Telemetry Diagnostics!