nano-vllm 用千行代码拆解 vLLM 核心,是读懂大模型推理最快的捷径。

1. 介绍

前文介绍的 ModelRunner 把模型当黑盒——run_model 直接 model(input_ids, positions) 拿到 hidden,再过 lm_head 得 logits。

本篇打开这个黑盒,看 Qwen3 的骨架:一串 token id 进去,怎么经词嵌入、28 层 decoder、最后一道归一化,输出能算 logits 的 hidden。

2. 总览

Qwen3 的骨架只有四段:词嵌入 embed_tokens、28 层 Qwen3DecoderLayer 堆叠、最后一道 RMSNorm、输出头 lm_head。token id 自上向下计算，最终变成 logits。

Qwen3-0.6B 的关键配置:

配置项	值	含义
hidden_size	1024	隐藏状态向量的维度
num_hidden_layers	28	堆叠的 decoder 层数
num_attention_heads	16	query 头数
num_key_value_heads	8	KV 头数（少于 q 头，GQA）
head_dim	128	每个头的维度
intermediate_size	3072	MLP 中间层宽度
vocab_size	151936	词表大小
tie_word_embeddings	True	lm_head 与 embed_tokens 共享权重

3. Qwen3Model 骨架

Qwen3Model 把骨架装在一起,forward 把四段串起来跑一遍。

import torch
from torch import nn
from transformers import Qwen3Config

from nanovllm.layers.layernorm import RMSNorm                       
from nanovllm.layers.embed_head import VocabParallelEmbedding, ParallelLMHead
from nanovllm.models.qwen3 import Qwen3Attention, Qwen3MLP         


class Qwen3Model(nn.Module):
    def __init__(self, config: Qwen3Config):
        super().__init__()
        # 词嵌入：token id → [hidden] 向量
        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size, config.hidden_size)
        # 28 层 decoder 堆叠，结构完全一样
        self.layers = nn.ModuleList(
            [Qwen3DecoderLayer(config)
             for _ in range(config.num_hidden_layers)])
        # 最后一道归一化
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, input_ids, positions):
        # 查表得初始 hidden [N, hidden]
        hidden_states = self.embed_tokens(input_ids)
        # 残差连接
        residual = None
        # 顺序执行 28 个 DecoderLayer，上一层 layer 输出是下一次 layer 输入
        for layer in self.layers:
            hidden_states, residual = layer(
                positions, hidden_states, residual)
        # final norm 把最后一层的残差并入并归一化
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states

3.1 embed_tokens

模型的入口,本质是一张 [vocab_size, hidden_size] 的查找表:输入一串 token id,逐个查出对应的 hidden_size 维向量,拼成初始 hidden:

hidden_states = self.embed_tokens(input_ids)

为什么先转成向量:token id 只是个编号,像身份证号——能唯一指认一个词,号码本身却不含信息,5 号和 6 号不会因此更相近。而模型只对连续向量做计算(相加、投影、注意力),没法拿编号直接算。embedding 给每个词发一张「多维画像」:每一维刻画一点属性(训练时模型自己学的),语义相近的词画像也相近——词与词的远近、关系,这才能用数学表达。

3.2 Qwen3DecoderLayer

骨架里 self.layers 是 28 个结构完全相同的 Qwen3DecoderLayer,模型绝大部分参数与计算都在这里。一层按 input_layernorm → self_attn → post_attention_layernorm → mlp 的顺序走一遍,每个子层的输出都经残差加回主干。三种部件先看作用,内部留到后续篇幅:

• self_attn(自注意力):让每个 token 环视它之前出现的所有 token,谁相关就多看谁,按相关度把信息汇总过来——token 之间的沟通、对上下文的理解,靠的就是它。打个比方:像开会时环顾全场,谁说的跟我有关,就多记谁的话。
• mlp(前馈):注意力管"交流",mlp 管"消化"——每个 token 各自把刚收到的信息深加工一遍(先摊到更宽的 3072 维算,再收回 1024)。打个比方:把头脑风暴的记录整理成自己的结论。
• layernorm(RMSNorm):子层开算前的"校准",把输入的数值规模拉回统一范围,让 attn / mlp 在稳定的尺度上计算，防止多层计算导致数值过大或过小。

class Qwen3DecoderLayer(nn.Module):
    def __init__(self, config: Qwen3Config):
        super().__init__()
        # 注意力层
        self.self_attn = Qwen3Attention(
            hidden_size=config.hidden_size,
            num_heads=config.num_attention_heads,
            num_kv_heads=config.num_key_value_heads,
            max_position=config.max_position_embeddings,
            rms_norm_eps=config.rms_norm_eps,
            qkv_bias=getattr(config, 'attention_bias', True),
            head_dim=getattr(config, 'head_dim', None),
            rope_theta=getattr(config, "rope_theta", 1000000),
            rope_scaling=getattr(config, "rope_scaling", None),
        )
        # 前馈层
        self.mlp = Qwen3MLP(
            hidden_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
        )
        # 两道 RMSNorm：attn 前、mlp 前
        self.input_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, positions, hidden_states, residual):
        if residual is None:
            # 归一化（第 0 层走这个if分支）
            hidden_states, residual = \
                self.input_layernorm(hidden_states), hidden_states
        else:
            # 相加 + 归一化融成一步
            hidden_states, residual = \
                self.input_layernorm(hidden_states, residual)
        # 注意力在归一化副本上算，输出待下一道 norm 时并入主干
        hidden_states = self.self_attn(positions, hidden_states)
        # post_norm：把 attn 输出并入主干 + 归一化
        hidden_states, residual = \
            self.post_attention_layernorm(hidden_states, residual)
        # MLP 在归一化副本上算
        hidden_states = self.mlp(hidden_states)
        # 把 (mlp 输出, 主干) 交给下一层（或 final norm）合并
        return hidden_states, residual

3.3 norm

28 层跑完,主干最后一层的输出还没归一化过。self.norm 这道 RMSNorm 给它收尾,把数值调回稳定范围,再交给 lm_head 算 logits:

hidden_states, _ = self.norm(hidden_states, residual)

它和 decoder layer 层内那两道是同一种 RMSNorm,区别只在位置——在所有层之后、输出之前。

3.4 残差连接

残差连接把每个子层的输出加回一条贯穿全程的主干,这条主干就是残差流。

是什么:一条从 embed_tokens 输出起、贯穿 28 层、直到 final norm 的主干流,宽度恒为 hidden_size(=1024)。每一层不直接覆写它,而是在它的归一化副本上算出一份「增量」,再加回主干。

打个比方:像一条传送带。每个工位(子层)从带上取一份副本加工,把成品放回带子;带子本身一路不断、越积越厚,后面的工位始终看得到最初的货。

为什么需要:28 层堆叠,若每层直接覆写输出,底层信息在深处会衰减。残差连接(Transformer 标配)让每层只学「对主干的修正」,信息直通到顶。

解决了什么:让深层堆叠稳定——主干恒在,子层只做加法。

怎么解决:每个子层做 主干 += 子层(norm(主干))。

回看 3.2 的 forward:在 nano-vllm 里,这个「加」被融进下一个 norm——两参版 RMSNorm(hidden, residual) 一次做完「相加 + 归一化」,返回 (归一化结果, 新残差)。于是 residual 必须作为单独张量在层间显式传递,而不是在层内就地相加,省一遍对整条主干的读写。

主干宽度为何恒为 1024:注意力内部会投到 2048、MLP 内部投到 3072,但 o_proj / down_proj 都把它收回 1024 才加回主干,所以主干一路不变。

4. Qwen3ForCausalLM：输出 logits

Qwen3Model 算出 hidden 后还差一步:把每个位置的 [1024] hidden 投到 [151936] 词表分数(logits)。这一步是 lm_head,封装在最外层的 Qwen3ForCausalLM 里。

class Qwen3ForCausalLM(nn.Module):
    # HF 的分散权重名 → 合并参数名（权重加载用，后续章节介绍）
    packed_modules_mapping = {
        "q_proj": ("qkv_proj", "q"),
        "k_proj": ("qkv_proj", "k"),
        "v_proj": ("qkv_proj", "v"),
        "gate_proj": ("gate_up_proj", 0),
        "up_proj": ("gate_up_proj", 1),
    }

    def __init__(self, config: Qwen3Config):
        super().__init__()
        self.model = Qwen3Model(config)                    # 骨架主体
        self.lm_head = ParallelLMHead(                     # 输出头
            config.vocab_size, config.hidden_size)
        if config.tie_word_embeddings:
            # 共享权重：lm_head 直接指向 embed 的同一块存储
            self.lm_head.weight.data = \
                self.model.embed_tokens.weight.data

    def forward(self, input_ids, positions):
        return self.model(input_ids, positions)            # 得 hidden

    def compute_logits(self, hidden_states):
        return self.lm_head(hidden_states)                 # hidden → logits

lm_head：从 hidden 到 logits

主干过完 final norm 得到 hidden(每个位置一个 [1024] 向量),lm_head 是一个无 bias 的线性层,把 [1024] 投成整张词表的分数 [151936]:

logits = hidden @ lm_head.weight.T　# 每个位置 [1024] → [151936]，权重 [151936, 1024]

即给词表里每个词打一个分(logit),分最高的(或按 sampler 策略采样的)就是预测的下一个词。(自回归只用序列最后一个位置来预测下一个词,实际只对那一位算 logits。)

tie_word_embeddings(权重共享)

lm_head 与 embed_tokens 共用同一份 [vocab, hidden] 权重矩阵 E。E 的第 i 行 E[i] 是词 i 的"释义向量"。

打个比方(查字典):E 是一本字典,每个词对应一条释义向量。embed_tokens 是正向查:给一个词(token id i),翻到那一行,取出释义向量 E[i];lm_head 是反向查:拿模型输出的向量 h,跟字典里每个词的释义逐一比对(点积 h·E[j]),谁最像,谁的分(logit)就最高。

为什么能共享:两步查的是同一本字典——「词 → 向量」与「向量 → 词」互为正反。一个词在入口的释义向量,本就该是出口拿来比对的那一行;共享 E 不只省参数,语义上也自洽。数学上 lm_head 算的 logits = h @ Eᵀ,正是「h 跟 E 的每一行做点积」,与 embed 的「取 E 的一行」严丝合缝。

省了多少:[151936, 1024] ≈ 1.56 亿参数,占 0.6B 模型相当比例,共享省下一整本字典的显存。

怎么实现:self.lm_head.weight.data = self.model.embed_tokens.weight.data,让两者指向同一块存储(集成验证用 data_ptr() 相等佐证)。

compute_logits 正是 L11 run_model 的第二步:先 model(...) 拿 hidden,再 model.compute_logits(hidden) 得 logits。

5. 集成验证

加载真实 Qwen3-0.6B,打印模块树看清骨架,跑一次 forward 确认残差流宽度,再验证 lm_head 与 embed 确实共享权重。

import torch.distributed as dist
from modelscope import snapshot_download
from nanovllm.config import Config
from nanovllm.engine.sequence import Sequence
from nanovllm.sampling_params import SamplingParams
from nanovllm.utils.context import reset_context
from transformers import AutoTokenizer

# 复用 L11 ModelRunner
from topic11_model_runner import ModelRunner

# 单卡环境（同 L11）
torch.cuda.set_device(0)
if not dist.is_initialized():
    dist.init_process_group(
        "nccl", "tcp://localhost:2335", world_size=1, rank=0)

model_path = snapshot_download("Qwen/Qwen3-0.6B")
config = Config(model_path, enforce_eager=True, max_model_len=4096)
runner = ModelRunner(config)
model = runner.model           # 正是上文读的 Qwen3ForCausalLM
tokenizer = AutoTokenizer.from_pretrained(model_path)

Downloading Model from https://www.modelscope.cn to directory: /DATA/disk5/cache/modelscope/models/Qwen/Qwen3-0.6B


2026-06-01 16:58:27,884 - modelscope - INFO - Target directory already exists, skipping creation.

# 打印模块树
print(model)

Qwen3ForCausalLM(
  (model): Qwen3Model(
    (embed_tokens): VocabParallelEmbedding()
    (layers): ModuleList(
      (0-27): 28 x Qwen3DecoderLayer(
        (self_attn): Qwen3Attention(
          (qkv_proj): QKVParallelLinear()
          (o_proj): RowParallelLinear()
          (rotary_emb): RotaryEmbedding()
          (attn): Attention()
          (q_norm): RMSNorm()
          (k_norm): RMSNorm()
        )
        (mlp): Qwen3MLP(
          (gate_up_proj): MergedColumnParallelLinear()
          (down_proj): RowParallelLinear()
          (act_fn): SiluAndMul()
        )
        (input_layernorm): RMSNorm()
        (post_attention_layernorm): RMSNorm()
      )
    )
    (norm): RMSNorm()
  )
  (lm_head): ParallelLMHead()
)

# prompt → token ids（chat template，同 L11/L12）
msgs = [{"role": "user", "content": "你是谁"}]
text = tokenizer.apply_chat_template(
    msgs, tokenize=False,
    add_generation_prompt=True, enable_thinking=False,
)
prompt_ids = tokenizer(text).input_ids

# 造一条 Sequence，铺成 prefill 张量（set_context 由 prepare_prefill 完成）
seq = Sequence(prompt_ids, SamplingParams(temperature=0.6))
seq.num_scheduled_tokens = len(seq)
seq.block_table = list(range(seq.num_blocks))
input_ids, positions = runner.prepare_prefill([seq])

# 跑骨架主干：embed → 28 层 → final norm
# 推理不需梯度
with torch.inference_mode():
    hidden = model.model(input_ids, positions)
print("hidden shape :", tuple(hidden.shape))   # [N, 1024] 残差流宽度恒 1024
reset_context()

# 验证 tie weights：lm_head 与 embed 指向同一块存储
shared = (model.lm_head.weight.data_ptr()
          == model.model.embed_tokens.weight.data_ptr())
print("tie weights  :", shared)                # True

hidden shape : (14, 1024)
tie weights  : True

6. 小结

Qwen3 的骨架是四段:embed_tokens → 28×Qwen3DecoderLayer → final RMSNorm → lm_head。

贯穿其中的是残差流:一条宽度恒为 1024 的主干,每个子层在它的归一化副本上算增量、加回主干。nano-vllm 把这个「加」融进下一个 norm,所以 residual 作为单独张量在层间显式传递。

tie_word_embeddings 让 lm_head 与 embed_tokens 共享一份词表权重,省下约 1.56 亿参数的显存。

下一篇详细介绍 RMSNorm 与残差相加怎么融成一个 add_rms。