LLM内部：在MLX上使用纯Python实现Qwen3-0.6B的推理
使用MLX在约200行Python代码中逐步实现LLM推理功能：包括分词器、RMSNorm、RoPE、分组查询注意力、SwiGLU、KV缓存以及采样机制。
来源: How Modern LLMs Actually Generate Text: A Pure Python Walkthrough
输入“法国的首都是”，模型大约在900毫秒内给出“巴黎”作为答案。在这么短的时间内究竟发生了什么？是哪一系列转换将五个单词变成了预测结果？
要弄清楚这一点，最好的办法是自己构建整个推理流程，而不是使用那些隐藏了细节的库。下面的演示将使用真实的Qwen3-0.6B模型权重，通过苹果的数组框架MLX在普通Python环境中进行操作。整个脚本约有200行，可在任何苹果芯片Mac上运行。设置说明位于文末。
流程概览
Qwen3、GPT-4、Llama及其衍生模型都遵循相同的处理流程：
Text → Tokenize → Embed → [Transformer Block × 28] → LM Head → Logits → Sample → Next Token
↑ |
└-------------------------- KV Cache --------------------------┘
文本被输入后转化为数字，经过28个变换器模块处理，最终形成151,936个潜在下一个标记的概率分布。系统会从中选择一个标记并添加进去，然后循环重新开始。
该循环包含两个阶段：
- 预填充阶段。所有输入标记会一次性被处理。模型读取提示语并构建其内部状态。
- 解码阶段。标记会逐个生成。新生成的标记只需关注之前已存在的标记，这些标记的键值存储在KV缓存中，因此无需重新处理。
预填充阶段是并行运行的，而解码阶段则是顺序运行的。KV缓存正是让解码速度变快的关键。如果没有它，生成第N个标记就意味着必须重新计算之前所有N-1个标记的注意力权重。
现在逐一介绍每个阶段。
第一步：分词将文本转换为整数
模型从未见过文字，只有整数。
from tokenizers import Tokenizer
tok = Tokenizer.from_file("Qwen3-0.6B/tokenizer.json")
ids = tok.encode("The capital of France is").ids
# [785, 6722, 315, 9625, 374]
每个整数都是指向包含151,936个词元的词汇表的索引。该分词器使用字节对编码，因此词元边界并不总是与单词边界对齐。“capital”是一个词元，“France”是另一个词元，开头的空格也属于该词元。从头实现BPE并非此处重点，因为该分词器实际上是一个查找表，脚本直接使用已有的实现即可。
步骤2：嵌入将整数转换为向量
第一个实际计算步骤是将每个词元ID映射为一个1,024维的向量。
x = W["model.embed_tokens.weight"][mx.array(token_ids)] # (seq, 1024)
这只不过是表查找操作：嵌入矩阵的形状为 (151936, 1024)，每个标记对应一行，每个隐藏维度对应一列，代码会通过索引访问该矩阵。可以将这个矩阵视为模型的词典：每个标记都有一个经过学习得到的1,024位数字指纹，用于编码其含义。“巴黎”和“法国”的指纹指向方向相似，而“巴黎”和“计算机”的指纹指向方向则截然不同。
第3步：RMSNorm
在应用注意力机制之前，需要先对隐藏状态进行归一化处理。现代模型使用RMSNorm来替代旧版的LayerNorm。
简单来说，该公式会将每个向量除以其均方根值——即各元素平方值的平均值加上一个极小的常数后的平方根——然后再乘以每个特征对应的经过学习得到的权重。
def rms_norm(x, w):
return x * mx.rsqrt(mx.mean(x * x, axis=-1, keepdims=True) + EPS) * w
因此，每个向量会先根据其自身的幅度进行缩放，然后再通过为每个特征学习得到的权重进行调整。该方法无需进行均值减法，也没有偏置项，这使得它在不牺牲质量的前提下比LayerNorm更为简单。EPS（1e-6）用于防止除以零的情况，而在训练过程中学习得到的权重w则让模型能够独立地对1,024个维度中的每一个进行重新缩放。
第4步：RoPE将位置信息融入向量中
很多“从零开始构建LLM”的文章都在这一环节让读者感到困惑，但实际上它比看上去要简单。
问题所在。注意力机制是对排列不变的。如果打乱输入标记的顺序，输出结果也会是同样的打乱顺序，因此模型无法判断哪个标记是原本在前的。
旧有的解决方法。在每个标记的嵌入向量中添加位置编号。这种方法可行，但缺乏通用性，且模型在处理相对距离时存在困难。
现代的解决方法。让查询向量和键向量根据标记所在位置的角度进行旋转。这就是旋转位置嵌入，即RoPE。每一组维度都被视为平面上的一个点，并根据标记的位置来设定旋转角度。
def apply_rope(q, k, positions):
half = D_HEAD // 2 # 64
freqs = 1.0 / (THETA ** (mx.arange(half, dtype=mx.float32) / half))
angles = mx.array(positions, dtype=mx.float32)[:, None] * freqs[None, :]
cos = mx.cos(angles)[None, :, :]
sin = mx.sin(angles)[None, :, :]
def rotate(x):
x1 = x[..., :half]
x2 = x[..., half:]
return mx.concatenate([x1 * cos - x2 * sin, x1 * sin + x2 * cos], axis=-1)
return rotate(q), rotate(k)
将每对维度视为一个二维平面。位于不同位置的标记会旋转不同的角度，因此当注意力机制计算查询项与键项的点积时，结果取决于它们之间的距离，而不仅仅取决于它们所包含的内容。以下是位于位置0、1和2的三个标记，每个标记都表现为一个二维箭头，且旋转角度比前一个更大一些：
Position 0 (angle = 0°) Position 1 (angle = 24°) Position 2 (angle = 48°)
q q q
| / /
| / /
| / /
| / /
+------ x +------ x +------ x
Same content, 0° rotation Rotated 24° counterclockwise Rotated 48° counterclockwise
"I'm at the start" "I'm one step in" "I'm two steps in"
在原始图中，第三支箭的旋转角度并不比第二支更大，因此应将此图视为大致示意图而非精确图形。
Attention通过点积来比较查询箭头与键箭头。角度相同的箭头会产生较大的点积值，而角度差异极大的箭头则会产生较小的点积值。因此该点积值反映了两个标记之间的距离，无需存储任何显式的位置编号。RoPE不会将位置作为单独的信号加入。旋转操作已将位置信息融入每个向量之中。
Qwen3设置了THETA = 1,000,000，这一较高的基础频率使得模型能够处理非常长的上下文，即便不使用扩展功能也能处理多达262K个标记的上下文。
第5步：带分组查询头的Attention机制
这是整个机制的核心。理解了Attention，就等于理解了Transformer。
从MHA到MQA再到GQA
Multi-Head Attention (MHA): 16 Q heads, 16 K heads, 16 V heads — expensive
Multi-Query Attention (MQA): 16 Q heads, 1 K head, 1 V head — too aggressive
Grouped Query Attention: 16 Q heads, 8 K heads, 8 V heads — sweet spot
Qwen3-0.6B采用GQA架构，拥有16个查询头，仅有8个键值头，因此每个键值头需要为两个查询头服务。这样一来，键值缓存的大小减少了一半，而质量损失却很小。
完整的注意力机制
def attention(x, layer_idx, cache, positions):
p = f"model.layers.{layer_idx}.self_attn"
# Project to Q, K, V
q = x @ W[f"{p}.q_proj.weight"].T # (seq, 2048) = 16 heads × 128
k = x @ W[f"{p}.k_proj.weight"].T # (seq, 1024) = 8 heads × 128
v = x @ W[f"{p}.v_proj.weight"].T # (seq, 1024) = 8 heads × 128
# Split into heads
q = q.reshape(-1, N_HEADS, D_HEAD).transpose(1, 0, 2) # (16, seq, 128)
k = k.reshape(-1, N_KV, D_HEAD).transpose(1, 0, 2) # (8, seq, 128)
v = v.reshape(-1, N_KV, D_HEAD).transpose(1, 0, 2) # (8, seq, 128)
# QK-Norm (Qwen3-specific: normalize each head before RoPE)
q = rms_norm(q, W[f"{p}.q_norm.weight"])
k = rms_norm(k, W[f"{p}.k_norm.weight"])
# Apply RoPE
q, k = apply_rope(q, k, positions)
# KV Cache: append new keys/values
if cache[layer_idx] is not None:
ck, cv = cache[layer_idx]
k = mx.concatenate([ck, k], axis=1)
v = mx.concatenate([cv, v], axis=1)
cache[layer_idx] = (k, v)
# GQA: expand 8 KV heads → 16 (each KV head shared by 2 Q heads)
k = mx.repeat(k, N_HEADS // N_KV, axis=0) # (16, total, 128)
v = mx.repeat(v, N_HEADS // N_KV, axis=0)
# Scaled dot-product attention
scale = D_HEAD ** -0.5
scores = (q @ k.transpose(0, 2, 1)) * scale # (16, seq, total)
# Causal mask (only during prefill)
if seq_len > 1:
mask = mx.triu(mx.full((seq_len, total), -1e9, dtype=scores.dtype), k=1)
scores = scores + mask
# Softmax + weighted sum of V
attn = mx.softmax(scores, axis=-1)
out = attn @ v
# Merge heads + output projection
out = out.transpose(1, 0, 2).reshape(-1, N_HEADS * D_HEAD)
return out @ W[f"{p}.o_proj.weight"].T
其中有一些部分需要解释。
QK-Norm。 Qwen3在RoPE之前会对每个头的Q向量和K向量应用RMSNorm，这是一种有助于稳定训练的较新技巧。其权重q_norm和k_norm由128个数值组成，代表每个维度上的缩放系数。
KV缓存。这是推理过程中最重要的优化措施。在预填充阶段，会为每个输入标记计算K和V并存储起来。在解码阶段，只需为新的标记计算K和V即可，它们会被追加到缓存中，随后注意力机制会遍历整个历史记录。如果跳过缓存，就不得不为之前的每个标记重新计算K和V，这会使原本线性复杂度的工作变为二次复杂度。
因果掩码。在预填充阶段，位置3的标记不能看到位于其未来位置的、即位置4的标记。代码通过在得分矩阵的上三角区域加上-1e9来确保这一点。在解码阶段，由于序列长度为1，无需使用掩码，因为单个新生成的标记可以查看其之前的所有内容。
公式。注意力机制首先计算(Q乘以K的转置再除以头维度平方根)的softmax值，然后将结果与V相乘。点积用于衡量每个查询与每个键的匹配程度。除以头维度平方根可以防止得分过大导致softmax函数饱和。softmax将得分转换为概率，再与V相乘即可得到值向量的加权平均值。
第6步：SwiGLU前馈层
在经过注意力机制处理后，每个标记都会通过一个前馈网络。现代模型使用SwiGLU而非旧版的基于ReLU的版本。具体而言：首先从输入中计算出两个投影值，其中一个通过SiLU函数处理后用作门控信号，然后将这两个投影值逐元素相乘，最后通过另一个投影操作将结果还原。
def ffn(x, layer_idx):
p = f"model.layers.{layer_idx}.mlp"
gate = x @ W[f"{p}.gate_proj.weight"].T # (seq, 3072)
up = x @ W[f"{p}.up_proj.weight"].T # (seq, 3072)
silu = gate * mx.sigmoid(gate) # SiLU(x) = x · σ(x)
return (silu * up) @ W[f"{p}.down_proj.weight"].T # (seq, 1024)
前馈网络是模型存储大部分知识的地方，这也是大型密集型模型具备强大性能的原因之一。该网络会将1,024维向量扩展到3,072维，应用SiLU门函数，然后再将其投影回1,024维。这个门函数负责决定哪些信息可以通过，哪些需要阻挡，有点像神经元的守门人。从1024维扩展到3072维再缩回的过程为模型提供了存储和处理信息的空间，而正是这种门控机制使得现代前馈网络优于ReLU网络。
第7步：变换器块
一个块会整合以上所有组件：
def transformer_block(x, layer_idx, cache, positions):
p = f"model.layers.{layer_idx}"
# Pre-norm attention + residual
h = rms_norm(x, W[f"{p}.input_layernorm.weight"])
x = x + attention(h, layer_idx, cache, positions)
# Pre-norm FFN + residual
h = rms_norm(x, W[f"{p}.post_attention_layernorm.weight"])
x = x + ffn(h, layer_idx)
return x
该流程为：标准化、注意力计算、添加残差、再次标准化、前向传播，最后再次添加残差。
这些残差连接（x = x + ...）至关重要。在训练过程中，它们能让梯度直接在网络中传递；在推理时，则能让信息跳过各个模块。没有这些连接的28层模型几乎无法训练。
Qwen3-0.6B由28个这样的模块堆叠而成，每个模块都会进一步优化特征表示。到第28个模块时，隐藏状态已具备足够的输入理解能力，从而能够预测下一个标记。
第8步：语言模型头部与采样
在最后一个模块之后，还需进行一次标准化处理，随后将特征映射到词汇表大小：
x = rms_norm(x, W["model.norm.weight"])
logits = x @ W["lm_head.weight"].T # (seq, 151936)
输出结果为151,936个分数，称为logits，每个可能的标记对应一个分数。分数越高，表示该标记出现的可能性越大。对于“法国的首都是”这一短语，经过softmax处理后的前五大概率值如下：
| Token | Probability |
|-----------|-------------|
| Paris | 0.6582 |
| located | 0.0274 |
| the | 0.0233 |
| ... | 0.0157 |
| ____ | 0.0136 |
“巴黎”以65.8%的得票率胜出，这一预测相当有把握。随后，采样机制会将这些概率值转换为一个唯一的标记ID：
def sample(logits, temperature=0.0, top_p=0.95):
if temperature == 0:
return int(mx.argmax(logits)) # Greedy: always pick the best
logits = logits / temperature
probs = mx.softmax(logits)
# Top-p: keep smallest set of tokens whose cumulative prob ≥ top_p
order = mx.argsort(-probs)
sorted_p = probs[order]
cumsum = mx.cumsum(sorted_p)
cutoff = int(mx.sum(cumsum < top_p)) + 1
sorted_p = sorted_p[:cutoff]
sorted_p = sorted_p / mx.sum(sorted_p)
order = order[:cutoff]
idx = mx.random.categorical(mx.log(sorted_p[None, :]))
return int(order[idx])
温度参数控制随机性。当值为0时，概率最高的标记始终会被选中（贪婪策略）。当值为1.0时，标记会按照其概率比例被选取；而当值为0.7时，概率分布会略微向概率较高的标记集中。
Top-p，即核采样，会保留累积概率达到指定值p（例如0.95）的最小标记组。该方法可以削减概率极低的标记，同时保留一定的多样性。
第9步：生成循环
完整的循环流程包括一次预填充操作，随后进行解码：
def generate(prompt, max_tokens=100, temperature=0.0, top_p=0.95):
ids = tok.encode(prompt).ids
# Prefill: process all input tokens at once
cache = [None] * N_LAYERS
logits = forward(ids, cache, 0)
next_id = sample(logits[-1], temperature, top_p)
ids.append(next_id)
# Decode: generate one token at a time
for _ in range(max_tokens - 1):
if next_id == EOS_ID:
break
pos = len(ids) - 1
logits = forward([next_id], cache, pos)
next_id = sample(logits[-1], temperature, top_p)
ids.append(next_id)
return tok.decode(ids)
这就是整个推理循环。像 vLLM、SGLang、TensorRT-LLM 和 mlx-lm 这样的服务引擎，本质上都是对这一核心部分的优化，而该核心部分仅约 198 行代码。
脚本未涵盖的内容
该脚本涵盖了标准 Transformer 的完整推理流程。较新的模型还增加了更多内容：
- 混合注意力机制。据称 Qwen3.5 将其 75% 的注意力层替换为 Gated DeltaNet，这是一种线性注意力变体，用固定大小的循环状态替代了完整的注意力矩阵。完整注意力机制如同照相记忆，完美但成本高昂；DeltaNet 则类似压缩后的笔记，高效但存在一定损失。混合机制则能在高效处理长上下文的同时保持近乎完美的召回率。
- 多标记预测。某些模型在每次前向传播中可预测多个标记，以略微牺牲质量为代价换取大幅的速度提升。
- 量化。推理阶段会使用4位或8位权重以节省内存和计算时间。该脚本则保留原始的bfloat16权重。
- 批量处理。服务器可同时处理多个请求并共享计算资源。而该脚本是一次仅处理一个请求。
这些措施都不会改变核心算法，它们都是建立在同一基础之上的。
为何要这么做
了解推理原理能让你成为更专业的实践者。服务器发出的“KV缓存已满”警告不再令人费解，GQA所需内存的依据也容易解释。而模型出现冗长重复的内容时，你就能通过温度参数和top-p值来解决问题。
现代大型语言模型并非魔法。它其实是由28层标准化处理、注意力机制和前馈网络反复叠加而成，通过旋转方式融入位置信息，而知识则存储在门控线性层中。相关数学知识仅为本科水平。真正的巧妙之处在于工程实现，但数学则是基础。只需约200行Python代码，加上真实的权重参数和输出结果即可实现。
亲自运行
上述所有内容都可以在任何搭载Apple Silicon芯片的Mac电脑（M1或更高版本）上运行。
1. 下载模型
从Hugging Face获取Qwen3-0.6B模型（位于Qwen/Qwen3-0.6B仓库中）。只需四个文件，其余的可以跳过：
Qwen3-0.6B/
├── config.json
├── model.safetensors (1.5 GB — the actual weights)
├── tokenizer.json
└── tokenizer_config.json
将模型文件夹放在 llm_inference.py 旁边，因为该脚本会从其启动的目录来解析 ./Qwen3-0.6B/ 路径。
2. 创建虚拟环境
# Create and activate a fresh venv
python3 -m venv llm_venv
source llm_venv/bin/activate
# Install the minimum packages - just two
pip install mlx tokenizers
无需 PyTorch、transformers、accelerate 或 safetensors 包。MLX 负责数组运算和权重加载，而分词器库则负责 BPE 编码与解码。
3. 运行它
python llm_inference.py
预期输出如下：
Model: Qwen3-0.6B (28 layers, 1024 dim, 16 Q heads, 8 KV heads)
Prompt: The capital of France is
------------------------------------------------------------
Output: The capital of France is Paris. The capital of Italy is
Rome. The capital of Spain is Madrid...
Time: 0.89s
如果出现 ModuleNotFoundError: No module named 'mlx' 的错误，请确认虚拟环境已激活（执行 source llm_venv/bin/activate），并且机器为 Apple Silicon 平台。MLX 无法在 Intel Mac、Linux 或 Windows 系统上运行。
所有操作都在笔记本电脑上完成，结果会在一秒内显示出来。
相关阅读
- 现代人工智能的运作原理：LLM、提示词技术、RAG、智能体与 MCP详解 —— 用通俗语言讲解人工智能技术栈：LLM 如何预测字符、四种提示词技巧、结合嵌入技术的 RAG、ReAct 智能体循环以及 MCP 的功能。