AI/ML 详解：推理优化

KV cache 就是自回归生成的 memoization——和动态规划里"算过的子问题别再算一遍"完全同构。后端世界对应：流式聚合的中间状态（Flink 的 keyed state）、数据库的物化视图、HTTP 长连接的会话上下文。没有它，每生成一个 token 都要把之前所有 token 的 attention 重新跑一遍——O(n²) 变 O(n³)。

LLM 推理是自回归的：生成第 100 个 token 时，attention 要算它对前 99 个 token 的相关性。Attention 的公式是 softmax(Q·KT/√d)·V——其中 Q(Query) 是当前 token 的"提问向量"，K(Key) 是历史 token 的"标签向量"，V(Value) 是历史 token 的"内容向量"。直觉：Q 像查询，K 像索引，V 像被查的数据。

关键观察：第 1-99 个 token 的 K 和 V 在生成第 100 个 token 时完全不变——它们只依赖各自位置之前的内容。所以 K/V 算一次就能永久复用，只有新 token 的 Q/K/V 需要新算。把所有历史 K/V 存起来，就是 KV cache。

这就是为什么 LLM 推理的瓶颈不在算力，而在显存带宽——一个反直觉但极其重要的事实。后续三个优化（speculative decoding、continuous batching、quantization）本质都在"想办法多榨点带宽出来"。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tok = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
m = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", torch_dtype=torch.float16).cuda()

input_ids = tok("The capital of France is", return_tensors="pt").input_ids.cuda()
past_key_values = None  # ← KV cache，初始为空

for _ in range(20):
    with torch.no_grad():
        out = m(input_ids=input_ids, past_key_values=past_key_values, use_cache=True)
    next_id = out.logits[:, -1, :].argmax(-1, keepdim=True)
    past_key_values = out.past_key_values  # ← 累积 cache
    input_ids = next_id  # ← 下一轮只喂"新 token"，前面的从 cache 取
    if next_id.item() == tok.eos_token_id: break
# 关掉 use_cache 让你直观感受差距——同样输出 20 token，慢 10-50 倍

和 CPU 的分支预测 + 推测执行（speculative execution）是同一个套路：用一个便宜的"小模型"先抢跑几步，再让"大模型"一次性并行验证。验证通过就赚到，验证失败就回滚。后端类比：乐观锁（先写后查，冲突再重试）、predicate pushdown（先用便宜的过滤器砍掉大部分数据）。

Decode 阶段每次只算 1 个新 token，但每次都得把整个 70B 模型的参数从显存搬到 GPU 算一遍——参数搬运 = 140 GB（FP16），即使 H100 也要花约 50ms。瓶颈是带宽不是算力，意味着 GPU 算力其实大量闲置。Speculative decoding 的洞察：既然 forward pass 一次的带宽成本是固定的，那一次验证 5 个候选 token几乎不增加成本——免费的并行。

机制三步走：

• ① Draft（草稿）：用一个 1B 量级的小模型连续生成 K 个 token（K 通常 4-8）。小模型快 10-20 倍，串行跑 K 步成本很低；
• ② Verify（验证）：把这 K 个 token 一次性喂给大模型，一次 forward pass 同时算出每个位置大模型自己会预测什么；
• ③ Accept / Reject：从头比对，连续相同的部分接受；第一个不同的位置用大模型的预测替换、丢弃后面的草稿。

性能数学：设接受率 α，每轮 K 个草稿，平均产出 (1-αK+1)/(1-α) 个 token。α=0.7、K=5 时约 2.9 token / 大模型调用——理论 2.9x 加速。实测 1.5-3x，剩余开销在小模型本身。Medusa / EAGLE 等变体把"小模型"换成在大模型上加几个轻量 head，省掉独立小模型的开销，加速比可达 3-4x。

# HuggingFace transformers 内置 assisted_generation 即 speculative decoding
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

big   = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-70B", torch_dtype=torch.float16, device_map="auto")
draft = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B",  torch_dtype=torch.float16, device_map="auto")
tok   = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-70B")

prompt = tok("Write a haiku about caching:", return_tensors="pt").to("cuda")

# assistant_model 就是"草稿模型"——大小模型必须用同一家族（共享 tokenizer + 词表）
out = big.generate(
    **prompt,
    assistant_model=draft,       # ← 关键参数
    max_new_tokens=100,
    do_sample=False,
)
print(tok.decode(out[0]))
# 输出和不开 assistant_model 时完全一致，速度 1.5-3x。
# vLLM / TensorRT-LLM / SGLang 在生产端有更优实现（PagedAttn + speculative 融合）

从阻塞式线程池升级到异步事件循环的范式跃迁——和 Node.js / Netty / asyncio 是同一个内核思想。传统 static batching 像"木桶效应"：整批等最慢那个请求生成完才能换下一批；continuous batching 像 OS 调度器，谁先完成就立刻让出位置，新请求填进来。vLLM 2023 论文把这件事推向主流，让吞吐量直接涨 5-23 倍。

LLM 请求的输出长度方差极大：有人要 10 token，有人要 2000 token。传统做法是"request-level batching"——攒 8 个请求一起跑，但因为 batch 同步前进，所有人都得等最慢那个 2000 token 跑完。中间 GPU 大量空转，吞吐拉胯。

Continuous batching（也叫 iteration-level scheduling）把调度粒度从"整个请求"降到"每一步 decode"：

• 每生成一个 token 后调度器都重新审视一次 batch；
• 完成的请求立刻退出，释放它在 KV cache 里的内存；
• 等待队列里的新请求立刻填进空位，开始 prefill；
• 同一个 batch 里 prefill 和 decode 阶段可以混跑（chunked prefill 进一步把长 prefill 拆成小块插到 decode 间隙）。

实现挑战在 KV cache 内存管理：传统做法给每个请求预分配最大长度的 KV 显存，浪费 60-80%。vLLM 的 PagedAttention 把 KV 切成固定 4KB 的"页"按需分配（直接借用 OS 虚拟内存的分页思想），让 continuous batching 真正可用。这是 2023 年系统侧最重要的 LLM 服务突破，今天 vLLM / SGLang / TensorRT-LLM / TGI 全都用类似机制。

# vLLM 把 continuous batching 默认开启——你只管批量 submit 请求即可
from vllm import LLM, SamplingParams

llm = LLM(model="meta-llama/Llama-3.1-8B-Instruct",
          gpu_memory_utilization=0.9,
          max_num_seqs=256)  # ← 最大并发请求数，调度器自动 in-flight 调度

prompts = [
    "Explain caching in one paragraph.",   # 输出短
    "Write a 1000-word essay on consciousness.",  # 输出长
    "What is 2+2?",                          # 极短
    # ... 几十几百个混合长度请求一起喂
]
params = SamplingParams(max_tokens=1024, temperature=0.7)

outputs = llm.generate(prompts, params)
# 短请求 1-2 秒就返回不必等长的，长请求继续跑——同一 batch 持续被新请求填满
# 同样硬件下，吞吐对比 HuggingFace transformers 通常 10x+，
# 这是为什么所有生产 LLM serving 都用 vLLM/SGLang/TensorRT-LLM 而不是裸 transformers

量化 = 神经网络的 JPEG 压缩——主动接受可控的有损精度损失换取显存和速度的几何级数收益。后端类比：定点数代替浮点数（嵌入式系统传统手法）、列存压缩（Parquet 用 dictionary encoding 把 string 压成 int）、protobuf 的 varint（小数字用更少字节）。所有这些都在押注："数据的分布特性允许我们用更少 bit 还原原貌。"

Llama-3 70B 原版 FP16 占 140 GB 显存——单张 H100（80 GB）装不下。但 H100 卖 ¥30 万。量化到 INT4 后只占 35 GB，一张消费级 RTX 4090（24 GB）配上 CPU offload 就能跑，成本差 100 倍。这就是量化为什么是本地 / 边缘部署 LLM 的入场券。

核心公式：x_int = round((x_float - zero) / scale)，反量化 x_float ≈ x_int * scale + zero。直觉：把浮点数线性映射到一个小整数区间——scale 是"每步多大"，zero 是"哪里是零点"。FP16 的 [-65504, 65504] 压到 INT8 的 [-128, 127] 就丢精度，但模型权重的分布是钟形曲线，大部分集中在零附近，损失远比想象中小。

权重量化（weight-only） vs 激活也量化（W·A）是两条不同路线：仅权重量化省显存 + 加快内存搬运（因为 LLM 推理瓶颈正是带宽），实现简单；激活也量化能用上 GPU 的 INT8 tensor core 真正提速但精度更难控。社区共识：本地推理用 weight-only INT4（AWQ/GPTQ），云端高吞吐用 W8A8 / FP8。

# 用 bitsandbytes 直接加载 4-bit 量化模型——一行配置搞定
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

qcfg = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",        # normal-float 4，QLoRA 论文提出，对正态分布最优
    bnb_4bit_use_double_quant=True,    # 量化"量化常数"，再省一点
)
m = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    quantization_config=qcfg,
    device_map="auto",
)
tok = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")

ids = tok("Define entropy in one sentence.", return_tensors="pt").to("cuda")
print(tok.decode(m.generate(**ids, max_new_tokens=100)[0]))
# 8B 模型从 16GB 降到 5GB——M2 Mac/RTX 3060 都能跑，质量损失通常<2%

• vLLM / PagedAttention 论文 (Kwon et al. 2023) — continuous batching + KV 分页的奠基性工作，必读
• Speculative Decoding 原论文 (Leviathan et al. 2023, Google) — 数学证明 + 算法细节
• Medusa (2024) — 不用独立小模型的 speculative decoding 变体，加速比 2-3x
• LLM.int8() (Dettmers 2022) — 大模型量化的开创性工作，QLoRA 作者前传
• QLoRA / NF4 (Dettmers et al. 2023) — 4-bit 量化 + LoRA 微调的标杆方案
• AWQ (Lin et al. 2023, MIT) — Activation-aware Weight Quantization，本地部署主流
• llama.cpp + GGUF — CPU/Mac 本地推理事实标准
• Lilian Weng — Large Transformer Model Inference Optimization — 系统性综述，工程视角最佳入门
• Kipply — Transformer Inference Arithmetic — 第一性原理推算延迟/显存/带宽，理解 LLM 推理必看
• Anyscale — How continuous batching enables 23x throughput — 工程视角最清晰的图文解释