Token Flow 与 Hidden State：从 Attention 到 LM Head

一次自回归生成可以沿着一条主线理解：token id 经 Embedding 变成 hidden state，Transformer blocks 不断更新这条向量，LM Head 再把它投影成词表 logits，Sampling 选出下一个 token。

1. 一次 decode step 流过哪里

如下所示：

token id
  -> embedding lookup
  -> hidden state
  -> Attention（读历史 KV，写当前 KV）
  -> FFN or MoE
  -> final hidden state
  -> LM Head
  -> vocab logits
  -> Sampling
  -> next token id

蓝色主线始终是 hidden state；KV、router 和 sampling 是围绕它发生的读写或决策。理解这条线后，Attention、MoE 和服务指标就不会被误当成彼此独立的模块。

2. Hidden state 是什么

输入最初只是 token id：

input_ids: [B, S]

Embedding table 的形状是 [vocab_size, hidden_size]。每个 token id 查出一行后，输入变成：

hidden_states: [B, S, H]

• B：同时处理的序列数。

• S：本次 forward 的 token 数。

• H：模型的 hidden size。

Embedding 可以看作第 0 层 hidden state。此后每个 Transformer block 都接收并输出同样主维度的 tensor；“语义”不是一个单独字段，而是分布在向量各维及层间变换里。

Prefill 常处理 [B, S, H]，decode 每个序列每步通常只新增一个 token，因此逻辑形状接近 [B, 1, H]。实际 kernel 可能展平 batch/token 维或采用 packed layout，但不改变这条语义主线！

3. Attention 如何更新 hidden state

Attention 不是直接选下一个词。它从当前 hidden state 生成 Query，并让 Query 读取可见上下文的 Key/Value：

current hidden state
  -> Q/K/V projection
  -> read historical KV cache
  -> attention output
  -> residual update
  -> contextualized hidden state

在 decode 中，历史 K/V 已存入 KV Cache，当前 step 只生成并追加新 token 的 K/V；Query 仍需访问可见历史。MHA、GQA、MLA、稀疏与递推结构如何改变这一步，统一由 Attention 架构演化 解释。

4. 一个 token 在一个 MoE 层里经历什么

以常见的 pre-norm decoder block 为例，省略具体 Norm 名称和张量 layout 后，一个 token 的核心路径是：

layer input hidden state
  -> Attention：读历史 K/V，写当前 token 的 K/V
  -> residual update
  -> Router：用当前 token 的 hidden state 计算 expert scores
  -> top-k expert ids + gate weights
  -> Dispatch #1：按 expert owner 发送 hidden state（跨 rank EP 时）
  -> Expert FFN：逐 token MLP，不读取 KV Cache
  -> Combine #2：expert 输出返回原 token/rank（跨 rank EP 时）
  -> 按 gate weights 加权求和
  -> residual update
  -> layer output hidden state

这里有四个容易混淆的边界：

1. KV Cache 只属于 Attention 路径：标准 Transformer 的 Expert FFN 不读取历史 K/V，也不产生自己的 KV。
2. Router 的直接输入是当前 token 的 hidden state：它不单独读取 KV Cache 或其他 token；但该 hidden state 已经过 Attention，因此已经携带上下文信息。
3. Expert FFN 在模型语义上是逐 token 计算：一个 expert 只对收到的向量做 MLP，不直接读取其他 token。运行时会把路由到同一 expert 的 token 打包成 Grouped GEMM，以提高 GPU 利用率；这是执行批处理，不是 token 间的信息交互。
4. 两次 all-to-all 是 EP 条件路径：Expert 跨 rank 放置时，第一次 dispatch 把 hidden state 发到 expert owner，第二次 combine 把 top-k 份结果送回原 token/rank。单卡、本地 Expert 或其他通信实现不应机械地描述成两次跨卡 all-to-all。

不同模型还可能包含 shared expert、不同的 Norm/Residual 顺序或融合实现；这些不会改变“Attention 负责上下文，Expert FFN 负责逐 token 变换”的职责划分。路由、EP/TP 和通信边界见 MoE 推理。

5. LM Head 与 Sampling 做什么

最后一层 hidden state 经 LM Head 投影到词表：

final_hidden_state: [H]
LM Head weight:     [V, H]
logits:             [V]

V 个 logits 是 V 个候选 token 的分数，不是一次生成 V 个 token。Sampling 再执行 temperature、top-k、top-p 或 greedy 等策略，选出 next token：

hidden state -> vocab logits -> sampling policy -> next token id

有些模型会让 LM Head 与输入 embedding 共享权重；这改变参数存储合同，不改变“hidden state 投影为 logits”的职责。

6. Prefill 与 Decode 如何复用这条线

两阶段运行的是同一套模型块，但 token 形状和 KV 行为不同：

阶段	本次处理	KV 行为	主要输出
Prefill	prompt 的多个 token	批量建立 KV Cache	首个生成位置的状态
Decode	每条活跃序列的一个新 token	读取历史并追加当前 KV	一个 next token

Serving pipeline 还包含 queue、scheduler、batch、cache allocation 和 response streaming；它是这条模型数据流的外层。

TIP: 不要把一次 forward 图直接当成完整请求时序。

7. 看模型或模拟器图时怎么对齐

遇到配置数字时先问它属于哪一维：

• vocab_size，例如 129,280：LM Head 输出维度。
• top_k / num_experts，例如 6 / 384：MoE 路由维度。
• hidden_size：层间主数据流的向量宽度。
• num_kv_heads 或 latent 维度：KV Cache 的每 token payload。

这些数字可能同时出现在一张图里，但不能互相换算。模拟器还必须区分模型块视角与 serving 视角，详见 模拟器建模指南。

相关页面

• Attention 架构演化 — MHA、KV 共享、表示压缩、稀疏访问与递推状态。
• KV Cache — Attention 历史状态的容量与生命周期。
• MoE 推理 — Router、dispatch、expert compute 与 combine。
• 模拟器建模指南 — 把模型数据流映射到容量和吞吐合同。