手写系列：从零实现一个极简大模型推理引擎

前言当你调用 OpenAI API 或者本地跑通 DeepSeek 时有没有好奇过——那些动辄百亿参数的大模型到底是怎么在 GPU 上跑起来的答案藏在推理引擎里。vLLM、TensorRT-LLM、llama.cpp 这些框架动辄几十万行代码把 PagedAttention、连续批处理、量化内核层层封装但剥开外壳核心流程其实并不神秘输入文本 → Tokenize → Embed → Transformer 逐层计算 → LM Head → 采样 → 输出文本今天我们就从零实现一个极简推理引擎。不在乎性能、不追求优化只求把推理链路讲透。代码用 Python NumPy能跑通一个 2 层 Attention 的迷你模型就行。当你理解了这个最小闭环再看 vLLM 的源码、TensorRT-LLM 的算子优化就不再是黑盒了。废话不多说直接开干。一、推理引擎的核心架构拆解一个完整的大模型推理引擎从启动到返回结果依次经过以下模块┌─────────────┐ │ Tokenizer │ 文本 ↔ Token ID ├─────────────┤ │ Embedding │ Token ID → 向量 ├─────────────┤ │ Transformer │ 核心计算N 层 Decoder Layer │ ┌─────────┐│ │ │Attention││ Self-Attention KV Cache │ │ FFN ││ SwiGLU / ReLU / GELU │ └─────────┘│ ├─────────────┤ │ LM Head │ 隐藏状态 → 词表概率 ├─────────────┤ │ Sampler │ 概率 → 下一个 Token └─────────────┘这里你可能会问为什么不用 Encoder-Decoder因为当前主流的大模型GPT、LLaMA、Qwen、DeepSeek都是纯 Decoder 架构。只有 BERT 这类模型才需要 Encoder。两者的核心区别很简单-Encoder能看完整句子双向注意力适合理解任务-Decoder只能看左边因果注意力适合生成任务大模型要做生成所以只用 Decoder。接下来我们逐块实现。先定义全局配置import numpy as np import math import re from typing import List, Optional, Tuple # ─── 模型配置 ─── class ModelConfig: def __init__(self): self.vocab_size 1024 # 词表大小 self.hidden_dim 64 # 隐藏层维度 self.num_layers 2 # Decoder 层数 self.num_heads 4 # 注意力头数 self.head_dim 16 # 每头维度 self.ffn_hidden_dim 128 # FFN 中间维度 self.max_seq_len 128 # 最大序列长度 self.bos_token_id 0 # 起始符 self.eos_token_id 1 # 结束符 self.pad_token_id 2 # 填充符为了让模型可跑可验证我们把配置压得很小。实际模型如 DeepSeek-V3 的 hidden_dim7168num_layers60但架构一模一样。二、Tokenizer文本与 Token 之间的桥梁Tokenizer 是模型的第一道关卡。它的任务说起来简单把文本切成模型能理解的 Token ID 序列推理完后再拼回来。目前主流的分词器有三种类型代表模型特点BPEGPT、LLaMA、Qwen子词合并处理未登录词强UnigramSentencePiece概率模型适合多语言WordPieceBERT单字词级别我们今天实现一个简化版 BPE Tokenizer。它不依赖 SentencePiece 库纯 Python 就能跑。# ─── 简化 BPE Tokenizer ─── class SimpleTokenizer: 一个极简 BPE Tokenizer 实现。 真实 BPE 需要预训练 merges这里用模拟数据演示流程。 def __init__(self, config: ModelConfig): self.vocab_size config.vocab_size # 构建模拟词表前 100 个是单字后面是常用词 self.vocab {} self.id_to_token {} # 单字 (char-level) for i in range(100): char chr(0x4E00 i) # 部分中文字符 self.vocab[char] i self.id_to_token[i] char # 常用词填充 common_tokens [你好, 模型, 推理, 引擎, 注意, 力机, transformer, attention, token, embedding, query, key, value, softmax, ffn, layer] for i, t in enumerate(common_tokens, start100): if i self.vocab_size: self.vocab[t] i self.id_to_token[i] t # 特殊 Token self.vocab[bos] config.bos_token_id self.vocab[eos] config.eos_token_id self.vocab[pad] config.pad_token_id self.id_to_token[config.bos_token_id] bos self.id_to_token[config.eos_token_id] eos self.id_to_token[config.pad_token_id] pad def encode(self, text: str) - List[int]: 文本 → Token IDs tokens [] tokens.append(self.vocab.get(bos, 0)) # 按字切分 尝试合并为词 i 0 while i len(text): matched False # 尝试匹配最长词 (从长到短) for j in range(min(6, len(text) - i), 0, -1): word text[i:ij] if word in self.vocab: tokens.append(self.vocab[word]) i j matched True break if not matched: # 按单字处理 char text[i] tokens.append(self.vocab.get(char, self.vocab.get(pad))) i 1 tokens.append(self.vocab.get(eos, 0)) return tokens def decode(self, ids: List[int]) - str: Token IDs → 文本 text for tid in ids: if tid in self.id_to_token: t self.id_to_token[tid] if t in (bos, eos, pad): continue text t return text这里的关键设计思路是分词要保持无损可逆性。也就是说你 encode 再 decode 回去内容不应该丢失。当然真实 BPE 比这复杂得多——它需要预训练一个 merge 规则表然后用贪心匹配来切分文本。以 DeepSeek-V3 为例它使用的是 Byte-level BPE词表大小 128K支持几乎所有语言的 tokenization。训练时统计了海量文本的字节对频率逐步合并形成子词单元。三、Embedding 层把 Token ID 变成向量模型不认识数字 ID它只认识向量。Embedding 层就是一张查找表Lookup Table# ─── Embedding 层 ─── class Embedding: def __init__(self, vocab_size: int, hidden_dim: int): # 初始化词嵌入矩阵 [vocab_size, hidden_dim] # 使用 Xavier 初始化 self.weight np.random.randn(vocab_size, hidden_dim) * 0.02 def forward(self, input_ids: np.ndarray) - np.ndarray: 输入: [batch_size, seq_len] 的 Token ID 输出: [batch_size, seq_len, hidden_dim] 的向量序列 return self.weight[input_ids]看到self.weight[input_ids]这行了吗这就是 Embedding 的全部秘密。每一行是一个 Token 的向量表示。模型训练过程就是不断调整这矩阵里的每个数字让语义相近的词向量互相靠近。比如猫和狗的向量距离应该比猫和飞机更近。在实际推理引擎中Embedding 层通常是推理的第一个瓶颈——当 hidden_dim7168、batch_size 很大时这张表的内存访问量可达数百 MB。四、核心Transformer Decoder Layer这是整个推理引擎的心脏。每个 Decoder Layer 包含两个子模块输入向量 │ ▼ ┌─────────────────────────┐ │ Layer Normalization │ │ ↓ │ │ Self-Attention (带 KV) │ ← 核心计算 │ ↓ │ │ Residual │ │ Skip Connection │ └─────────────────────────┘ │ ▼ ┌─────────────────────────┐ │ Layer Normalization │ │ ↓ │ │ Feed-Forward (FFN) │ ← 知识存储 │ ↓ │ │ Residual │ │ Skip Connection │ └─────────────────────────┘ │ ▼到下一层或 LM Head我们按顺序实现每个子模块。4.1 Layer Normalization在 Transformer 出现之前主流是 Batch Normalization。但 Transformer 发现 BN 在变长序列上效果不佳于是 Layer Norm 成了标配。核心思路对每个样本的每个位置独立做归一化。# ─── Layer Normalization ─── class LayerNorm: def __init__(self, hidden_dim: int, eps: float 1e-6): self.gamma np.ones(hidden_dim) # 可学习缩放 self.beta np.zeros(hidden_dim) # 可学习偏移 self.eps eps def forward(self, x: np.ndarray) - np.ndarray: 输入: [batch_size, seq_len, hidden_dim] 输出: 相同形状 # 计算均值和方差沿最后一维 mean np.mean(x, axis-1, keepdimsTrue) var np.var(x, axis-1, keepdimsTrue) # 归一化 x_norm (x - mean) / np.sqrt(var self.eps) # 缩放和平移 return self.gamma * x_norm self.betaLLaMA 系列用的是 RMS NormRoot Mean Square Normalization它去掉了均值归一化只除 RMSclass RMSNorm: LLaMA 使用的 RMS Norm def __init__(self, hidden_dim: int, eps: float 1e-6): self.weight np.ones(hidden_dim) self.eps eps def forward(self, x: np.ndarray) - np.ndarray: rms np.sqrt(np.mean(x ** 2, axis-1, keepdimsTrue) self.eps) return x / rms * self.weightRMS Norm 比 Layer Norm 少算一次均值推理时节省约 15% 的归一化开销。别小看这 15%乘上 60 层 Decoder省出来的算力相当可观。4.2 Self-Attention因果注意力Attention 是 Transformer 的核心创新也是推理时最耗算力的瓶颈——没有之一。大模型推理使用因果注意力Causal Attention每个位置只能看到自己及之前的位置不能偷看未来。这就是为什么你在 ChatGPT 里打字时它是一字一字生成的——每生成一个 Token只能看到已输出的内容。# ─── Self-Attention ─── class SelfAttention: def __init__(self, config: ModelConfig): self.num_heads config.num_heads self.head_dim config.head_dim self.hidden_dim config.hidden_dim # QKV 投影矩阵 self.wq np.random.randn(config.hidden_dim, config.hidden_dim) * 0.02 self.wk np.random.randn(config.hidden_dim, config.hidden_dim) * 0.02 self.wv np.random.randn(config.hidden_dim, config.hidden_dim) * 0.02 # 输出投影 self.wo np.random.randn(config.hidden_dim, config.hidden_dim) * 0.02 def forward(self, x: np.ndarray, start_pos: int 0, kv_cache: Optional[Tuple] None ) - Tuple[np.ndarray, Tuple]: x: [batch_size, seq_len, hidden_dim] start_pos: 当前生成的位置偏移用于 KV Cache kv_cache: (k_cache, v_cache) 元组 返回: (输出, 更新后的 KV Cache) batch_size, seq_len, _ x.shape # 1. 计算 QKV q x self.wq # [B, S, D] k x self.wk v x self.wv # 2. 分头 q q.reshape(batch_size, seq_len, self.num_heads, self.head_dim) k k.reshape(batch_size, seq_len, self.num_heads, self.head_dim) v v.reshape(batch_size, seq_len, self.num_heads, self.head_dim) # 转置为 [B, H, S, D_head] q q.transpose(0, 2, 1, 3) k k.transpose(0, 2, 1, 3) v v.transpose(0, 2, 1, 3) # 3. KV Cache 处理 if kv_cache is not None: k_cache, v_cache kv_cache # 拼接历史 KV k np.concatenate([k_cache, k], axis2) v np.concatenate([v_cache, v], axis2) updated_kv_cache (k, v) # 4. 计算 Attention Score # S Q K^T / sqrt(d_k) scale 1.0 / math.sqrt(self.head_dim) attn (q k.transpose(0, 1, 3, 2)) * scale # 5. 因果掩码 current_seq_len k.shape[2] mask np.triu(np.ones((seq_len, current_seq_len), dtypebool), kcurrent_seq_len - seq_len 1) attn np.where(mask, float(-inf), attn) # 6. Softmax attn self._softmax(attn, dim-1) # 7. 加权求和 out attn v # [B, H, S, D_head] # 8. 合并头 out out.transpose(0, 2, 1, 3) out out.reshape(batch_size, seq_len, self.hidden_dim) # 9. 输出投影 out out self.wo return out, updated_kv_cache def _softmax(self, x: np.ndarray, dim: int -1) - np.ndarray: x_max np.max(x, axisdim, keepdimsTrue) x_exp np.exp(x - x_max) return x_exp / np.sum(x_exp, axisdim, keepdimsTrue)重点看第 3 步的 KV Cache。在生成模式下每次只新增一个 Token但 Attention 需要看到所有历史 Token。如果没有 Cache每次都要重新计算之前所有 Token 的 K 和 V复杂度从 O(L) 变为 O(L²)。KV Cache 的核心思想把每个 Token 的 KKey和 VValue缓存起来新 Token 只需计算自己的 K、V然后拼接到历史缓存之后。业界已经在用更激进的手法优化 KV Cache-MQAMulti-Query Attention多个 Query 头共享一组 Key-Value 头如 Falcon 模型-GQAGrouped-Query Attention分组的共享如 LLaMA 2/3、DeepSeek-V2/V3-MLAMulti-head Latent AttentionDeepSeek-V2 首创的低秩 KV 压缩KV Cache 仅需存储压缩后的隐向量以 DeepSeek-V3 为例MLA 将 KV Cache 大小压缩到传统 MHA 的1/8这是它能在消费级 GPU 上跑满 671B 参数量激活 37B的关键技术之一。4.3 Feed-Forward Network如果说 Attention 是信息的路由器那 FFN 就是知识的存储器。研究表明FFN 层实际上存储了模型学到的事实性知识。现代大模型普遍使用SwiGLU激活函数Shazeer, 2020它比 ReLU、GELU 效果更好# ─── SwiGLU Feed-Forward Network ─── class SwiGLU_FFN: SwiGLU: FFN(x) (xW_gate ⊙ SiLU(xW_up)) W_down 其中 ⊙ 是逐元素相乘SiLU(x) x * sigmoid(x) def __init__(self, config: ModelConfig): # 三个权重矩阵 self.w_gate np.random.randn(config.hidden_dim, config.ffn_hidden_dim) * 0.02 self.w_up np.random.randn(config.hidden_dim, config.ffn_hidden_dim) * 0.02 self.w_down np.random.randn(config.ffn_hidden_dim, config.hidden_dim) * 0.02 def silu(self, x: np.ndarray) - np.ndarray: SiLU x * sigmoid(x) return x / (1 np.exp(-x)) def forward(self, x: np.ndarray) - np.ndarray: 输入: [batch_size, seq_len, hidden_dim] 输出: [batch_size, seq_len, hidden_dim] gate x self.w_gate # 门控信号 up x self.w_up # 上投影 hidden self.silu(gate) * up # SwiGLU 激活 out hidden self.w_down # 下投影回原始维度 return outSwiGLU 比 ReLU 好在哪直观来说门控信号gate让模型能够动态选择激活哪些神经元。ReLU 是硬阈值x ≤ 0 → 0SwiGLU 是软门控sigmoid 控制在 0~1 之间信息流动更平滑。在 DeepSeek-V3 中FFN 的中间维度达 18432当 hidden_dim7168 时这意味着 FFN 占了模型总参数量的大头。DeepSeek-V3 还使用了MoEMixture of Experts架构把一个大 FFN 拆成 256 个小专家每次只激活 8 个用更少的计算量获得了更大的模型容量。4.4 完整的 Decoder Layer现在把所有子模块组合成完整的 Decoder Layer# ─── Transformer Decoder Layer ─── class DecoderLayer: def __init__(self, config: ModelConfig, layer_id: int): self.layer_id layer_id self.attention_norm RMSNorm(config.hidden_dim) self.attention SelfAttention(config) self.ffn_norm RMSNorm(config.hidden_dim) self.ffn SwiGLU_FFN(config) def forward(self, x: np.ndarray, start_pos: int 0, kv_cache: Optional[Tuple] None ) - Tuple[np.ndarray, Tuple]: # Pre-Norm Self-Attention Residual residual x x self.attention_norm.forward(x) attn_out, kv_cache self.attention.forward(x, start_pos, kv_cache) x residual attn_out # Pre-Norm FFN Residual residual x x self.ffn_norm.forward(x) ffn_out self.ffn.forward(x) x residual ffn_out return x, kv_cache注意这里的Pre-Norm结构。早期 Transformer如原始论文使用 Post-Norm先 Attention/FFN 再 NormLLaMA 之后普遍改用 Pre-Norm先 Norm 再 Attention/FFN。Pre-Norm 训练更稳定无需 warmup 也能收敛。Residual Connection残差连接是让 Transformer 能堆到 60 层甚至更深的关键——没有它梯度会随层数指数级消失。五、LM Head Sampler从隐藏状态到文本Transformer 最后一层的输出是语义丰富的隐藏状态向量但它还不是文本。我们需要两步转换LM Head投影到词表大小得到每个 Token 的概率Sampler从概率分布中采样出下一个 Token5.1 LM Head# ─── LM Head ─── class LMHead: def __init__(self, config: ModelConfig): # 线性投影: hidden_dim → vocab_size # 实践中往往与 Embedding 层共享权重Weight Tying self.weight np.random.randn(config.hidden_dim, config.vocab_size) * 0.02 def forward(self, x: np.ndarray) - np.ndarray: 输入: [batch_size, seq_len, hidden_dim] 输出: [batch_size, seq_len, vocab_size] 的 logits return x self.weight注意 self.weight和 Embedding 层self.weight[input_ids]的对应关系——很多模型如 Transformer 原始论文会让两者共享权重矩阵这叫Weight Tying。训练时 Embedding 和 LM Head 使用同一个矩阵可以减少参数量并提升效果。5.2 采样策略有了每个 Token 的 logits未归一化的分数我们需要把它变成具体的 Token ID。不同的采样策略直接影响生成质量# ─── 采样器 ─── class Sampler: def __init__(self, config: ModelConfig): self.vocab_size config.vocab_size def sample(self, logits: np.ndarray, temperature: float 1.0, top_k: int 40, top_p: float 0.9) - int: 从 logits 中采样下一个 Token logits: [vocab_size] # 1. Temperature 缩放 if temperature 1e-6: # 贪心解码 return int(np.argmax(logits)) logits logits / temperature # 2. Softmax 转概率 probs self._softmax(logits) # 3. Top-K 过滤 if top_k 0 and top_k self.vocab_size: indices np.argpartition(logits, -top_k)[-top_k:] mask np.ones(self.vocab_size, dtypebool) mask[indices] False probs[mask] 0.0 probs probs / probs.sum() # 4. Top-P (Nucleus) 过滤 if top_p 1.0: sorted_idx np.argsort(probs)[::-1] cumsum np.cumsum(probs[sorted_idx]) cutoff_idx np.searchsorted(cumsum, top_p) mask np.ones(self.vocab_size, dtypebool) mask[sorted_idx[:cutoff_idx 1]] False probs[mask] 0.0 probs probs / probs.sum() # 5. 从过滤后的分布采样 return int(np.random.choice(self.vocab_size, pprobs)) def _softmax(self, x: np.ndarray) - np.ndarray: x x - np.max(x) exp_x np.exp(x) return exp_x / np.sum(exp_x)三个关键参数的作用参数作用越低→越高→Temperature控制随机性确定性强保守随机性强有创意Top-K只保留概率最高的 K 个 Token输出稳定多样性增加Top-P保留累计概率达 P 的最小 Token 集同 Top-K 但动态更多选择空间实际使用中temperature0.7, top_k40, top_p0.9是一个不错的起点。如果你想要最确定的输出比如代码生成可以设temperature0启用贪心解码。六、组装完整的推理引擎所有模块就位我们把它们串联起来# ─── 极简推理引擎 ─── class MiniInferenceEngine: 一个完整的极简大模型推理引擎 支持 Prefill Decode 两阶段推理 def __init__(self, config: ModelConfig): self.config config self.tokenizer SimpleTokenizer(config) self.embedding Embedding(config.vocab_size, config.hidden_dim) self.layers [DecoderLayer(config, i) for i in range(config.num_layers)] self.norm RMSNorm(config.hidden_dim) self.lm_head LMHead(config) self.sampler Sampler(config) # KV Cache: 每层一对 (k, v) self.kv_caches [None] * config.num_layers def reset_kv_cache(self): 重置 KV Cache新序列时调用 self.kv_caches [None] * config.num_layers def forward(self, tokens: np.ndarray, start_pos: int 0) - np.ndarray: 单次前向传播 tokens: [batch_size, seq_len] start_pos: 当前序列的起始位置 返回: [batch_size, seq_len, vocab_size] logits # Embedding h self.embedding.forward(tokens) # [B, S, D] # 逐层 Transformer for i, layer in enumerate(self.layers): h, self.kv_caches[i] layer.forward(h, start_pos, self.kv_caches[i]) # 最终 Norm h self.norm.forward(h) # LM Head logits self.lm_head.forward(h) # [B, S, V] return logits def generate(self, prompt: str, max_new_tokens: int 32, temperature: float 0.7, top_k: int 40, top_p: float 0.9) - str: 完整的文本生成流程 分为两个阶段: Phase 1 ─ Prefill: 处理整个 Prompt生成第一个新 Token Phase 2 ─ Decode: 逐 Token 生成利用 KV Cache # 1. Tokenize input_ids self.tokenizer.encode(prompt) tokens np.array([input_ids], dtypenp.int32) self.reset_kv_cache() generated input_ids.copy() # 2. Phase 1: Prefill一次处理整个 Prompt logits self.forward(tokens, start_pos0) next_logit logits[0, -1, :] # 取最后一个位置的 logits # 采样第一个新 Token next_token self.sampler.sample(next_logit, temperature, top_k, top_p) generated.append(next_token) # 3. Phase 2: Decode逐 Token 生成 for _ in range(max_new_tokens - 1): cur_tokens np.array([[next_token]], dtypenp.int32) logits self.forward(cur_tokens, start_poslen(generated) - 1) next_logit logits[0, -1, :] next_token self.sampler.sample(next_logit, temperature, top_k, top_p) generated.append(next_token) # 遇到 EOS 提前终止 if next_token self.config.eos_token_id: break # 4. Decode 回文本 return self.tokenizer.decode(generated)注意Prefill 和 Decode 的分离-Prefill 阶段一次性计算 Prompt 中所有 Token 的隐藏状态填充 KV Cache。计算量大但并行度高。-Decode 阶段每步只算一个 Token利用 KV Cache 避免重复计算。计算量小但有串行依赖。在实际框架中Prefill 和 Decode 通常会使用不同的 CUDA Kernel 来实现最优性能vLLM 的 PagedAttention、TensorRT-LLM 的 inflight batching 都是围绕这个差异做优化。七、跑起来让我们写个 main 函数验证一切是否正常def main(): config ModelConfig() engine MiniInferenceEngine(config) prompt 你好模型 output engine.generate(prompt, max_new_tokens16) print(fPrompt: {prompt}) print(f输出: {output}) # 看看 KV Cache 大小 total_kv_elements 0 for i, cache in enumerate(engine.kv_caches): if cache is not None: k, v cache total_kv_elements k.size v.size print(f\nKV Cache 总元素数: {total_kv_elements}) print(fKV Cache 形状: k{engine.kv_caches[0][0].shape}, fv{engine.kv_caches[0][1].shape}) if __name__ __main__: main()由于权重是随机初始化的输出不会特别有意义就像没有训练过的婴儿在胡说八道。但这个流程——输入文本 → Tokenize → Embed → Transformer × 2 → LM Head → 采样 → 输出文本和 LLaMA 3 70B 跑在 A100 上的流程一模一样。区别只是参数规模和优化手段。八、从极简到工业级推理引擎的进阶之路我们的 MiniInferenceEngine 展示了推理引擎的核心骨架。但真正的生产级引擎还需要解决以下问题8.1 连续批处理Continuous Batching我们的引擎每次只能处理一个请求。vLLM 首创的Continuous Batching允许多个请求在同一批次中不同步地生成——有的刚 Prefill、有的在 Decode交错执行以提高 GPU 利用率。8.2 PagedAttention 与内存管理KV Cache 是内存大户。一个 70B 模型在 2048 序列长度下单请求的 KV Cache 需要约 2.3GB 显存。50 个并发请求就是 115GB。vLLM 的PagedAttention借鉴操作系统的虚拟内存机制把 KV Cache 切分成固定大小的块Page按需分配消除了碎片。这能将显存利用率从 30% 提升到 95% 以上。8.3 量化推理FP16 的权重重 2 字节/参数70B 模型就是 140GB远超单卡容量。量化将权重压缩到 INT40.5 字节/参数70B 模型只需 35GB一张 A100 就能放下。主流量化方法包括-GPTQ基于 Hessian 矩阵的后训练量化-AWQ基于激活值的感知量化-GGUFllama.cpp 使用的格式支持多级量化-FP8NVIDIA 新硬件原生支持的 8 位浮点8.4 Speculative Decoding投机解码大模型生成时每步只能出一个 TokenGPU 利用率仅 10-30%。投机解码的思路是用一个更快的草稿模型生成多个候选 Token然后让大模型批量验证。如果草稿模型准确率高一次 Forward 就能验证多个 Token速度提升 2-3 倍。Google Medusa、DeepSeek 的 Self-Speculative Decoding 都是这个思路。8.5 前缀缓存Prefix Caching系统提示词System Prompt通常是共享的。如果每次请求都重新计算那前缀部分的 KV Cache 显然是被浪费了。前缀缓存将相同前缀的 KV Cache 缓存下来新请求命中时直接复用。在 RAG 场景中这能节省 60-80% 的计算量。总结我们从零实现了一个完整的极简推理引擎核心链路只有 6 步Tokenize → Embed → Transformer × N → LM Head → Sample每步背后都是明确的数学和工程原理模块核心原理工业级方案TokenizerBPE/Unigram 分词SentencePiece、tiktokenEmbedding查找表 位置编码RoPE、ALiBiSelf-AttentionScaled Dot-Product AttentionFlashAttention、PagedAttentionFFN门控线性单元SwiGLU、MoEKV Cache空间换时间MLA、GQASamplerTop-K Top-P对比解码、Mirostat看完这篇文章你应该能理解- 为什么大模型一字一字地生成因果注意力- 为什么推理上下文越长越慢KV Cache 读取瓶颈- 为什么量化能省这么多显存参数位宽压缩- 为什么 vLLM 比原生推理框架快那么多PagedAttention 连续批处理推理引擎的优化理念我们无法在一篇文章中穷尽但核心架构就这 6 步。理解了骨架再去看 vLLM、TensorRT-LLM 的源码你会发现——它们做的事我们的引擎都做了只是做得更极致。动手实践建议如果想进一步深入推荐以下实操路径先跑通我们的 MiniInferenceEngine把本文的完整代码粘贴到一个 Python 文件中确保能执行。虽然输出无意义但能帮你理解完整的推理流程。替换真实权重可以从 HuggingFace 下载 LLaMA 3.2-1B 的权重写一个转换脚本把 PyTorch 权重加载到我们引擎的参数中。这时就能看到真正的文本输出了。学习 FlashAttention在 Self-Attention 中替换为 FlashAttention 的 Tiling 算法观察内存占用和速度的变化。阅读生产级框架源码推荐按以下顺序阅读llama.cpp纯 C 实现代码量最小且逻辑清晰vLLMPython CUDA重点关注 PagedAttention 和 SchedulerTensorRT-LLM学习 Inflight Batching 和插件化算子设计动手加入量化支持在引擎中加入 INT8 或 INT4 量化推理体验参数精度和生成质量之间的权衡。相关阅读推荐如果你对系列其他文章感兴趣建议配合阅读- 手写系列从零实现 FlashAttention — 深入理解 Attention 的 IO 优化- 手写系列从零实现 KV Cache 量化推理引擎 — 本文 KV Cache 部分的实战延伸- 手写系列从零实现 MoE 架构 — 深入 DeepSeek-V3 的核心技术对于 DeepSeek 系列模型的推理部署实战推荐深度阅读DeepSeek 模型高性能推理与部署纯技术指南这篇文章深入分析了 DeepSeek-V2 的 MLA 注意力机制、MoE 的路由策略以及在实际部署中遇到的内存管理和推理延迟优化技巧是本文理论知识的实战延伸。本文为手写系列的第 14 篇。系列已涵盖 FlashAttention、MoE、RLHF、低代码引擎、SQL 引擎、前缀缓存、量化优化等主题。如果你想从零搭建一个完整的深度学习推理体系建议从第一篇开始顺序阅读。用技术深度见证成长。✍️