Token Flow 与 Hidden State：从 Attention 到 LM Head

1. 一句话

LLM 推理里，token 进入模型后会变成 hidden state。Attention、FFN/MoE 都在不断更新 hidden state。最后 LM Head 把 hidden state 映射成词表大小的 logits，Sampling 再从这些 logits 里选下一个 token。

图里蓝色是 hidden state 主数据流，黄色是 MoE routing / expert parallel，绿色表示生成 token 和 decode 继续。底部 Arrow actions 把关键箭头动作单独列出来：lookup、[B,S] -> [B,S,H]、read KV、route top-k、project vocab 和 sample。

token id
  -> embedding
  -> hidden state
  -> Attention / KV cache
  -> FFN or MoE
  -> final hidden state
  -> LM Head
  -> vocab logits
  -> Sampling
  -> next token id

这条线比单独背 Attention、MoE、LM Head 更有用，因为它把“模型结构”和“推理服务的 token flow”接起来了。

2. Hidden state 是什么

Hidden state 本质是 tensor。

更直观地说：每个 token 在模型内部不是一个词，而是一条很长的数字向量。

假设一句话被 tokenizer 切成 4 个 token：

北京 / 是 / 中国 / 的

进入模型后，可以把它想成：

北京: [0.12, -0.03, 0.77, ...]
是:   [0.44,  0.18, -0.21, ...]
中国: [0.91, -0.52, 0.08, ...]
的:   [0.05,  0.33, 0.62, ...]

如果模型的 hidden_size = 12,288，每个 token 就是一条长度为 12,288 的向量。

Embedding 可以看成第 0 层 hidden state。

模型输入一开始不是 [batch_size, seq_len, hidden_size]，而是 token id：

input_ids: [batch_size, seq_len]

例如：

[8, 4096]

Embedding table 是一张查表矩阵：

embedding_table: [vocab_size, hidden_size]

如果 vocab_size = 129,280、hidden_size = 12,288，那这张表的形状就是：

[129280, 12288]

每个 token id 去表里取一行：

13 -> embedding_table[13] -> [12288]
98 -> embedding_table[98] -> [12288]
42 -> embedding_table[42] -> [12288]

所以 embedding lookup 会把：

input_ids: [B, S]

变成：

hidden_states_l0: [B, S, H]

这就是最早的 hidden state，也就是 Transformer block 的输入。

一次推理里的 hidden state 常见形状是：

[batch_size, seq_len, hidden_size]

例如：

[8, 4096, 12288]

可以把它想成一张表：

• 每一行是一个 token。
• 每一列是一个隐藏特征。
• batch 维度表示同时处理多个请求。

Decode 阶段每步只处理当前新 token，因此常见形状可以简化成：

[batch_size, 1, hidden_size]

例如：

[8, 1, 12288]

MoE router 分发出去的，通常就是这些 token 的 hidden state 向量。

3. Attention 和 hidden state 的关系

Attention 在每层 Transformer block 里工作。它的作用不是直接选词，而是把上下文信息写入当前 token 的 hidden state。

Decode 时，Attention 会拿当前 token 的 Query 去读历史 KV cache：

current hidden state
  -> Q projection
  -> read K/V cache
  -> attention output
  -> updated hidden state

所以 Attention 的结果仍然是 hidden state，只是这条向量已经混入了上下文信息。

例如上下文是：

北京是中国的首

Attention 会让当前 token 的 hidden state 读到“北京 / 中国 / 首”这些上下文。经过多层 Attention 和 FFN/MoE 后，最终 hidden state 里会更强地表示“下一个 token 可能是 都”。

4. FFN 和 MoE 在这里做什么

Transformer block 里的 Attention 之后，通常还有 FFN。MoE 可以理解为稀疏版 FFN：不是所有 token 都走同一个 FFN，而是每个 token 根据 router 结果走 top-k 个 expert。

算法视角：

hidden state
  -> router / gate
  -> top-k experts
  -> expert FFN
  -> weighted combine
  -> updated hidden state

系统视角，如果 experts 分布在不同 GPU 上：

local hidden states
  -> All-to-All dispatch
  -> remote expert compute
  -> All-to-All combine
  -> hidden states return to original token positions

因此，“MoE 路由出去再传回来”这个理解基本正确。更精确的说法是：

• 路由出去的是 token 的 hidden state。
• expert 对 hidden state 做 FFN 计算。
• combine 按 router 权重把 expert 输出合回原 token 位置。
• 跨 GPU 部署时，dispatch/combine 往往通过 all-to-all 完成。

不是所有 MoE 都必然跨卡 all-to-all。EP=1 或 expert 刚好在本卡时，通信成本会小很多。大 MoE 生产推理通常会用 expert parallel，所以 all-to-all 是主要成本之一。

5. LM Head 和 Sampling 的关系

LM Head 在 Transformer blocks 之后。它不做上下文 attention，而是把最终 hidden state 投影到整个词表。

形状可以这样看：

final_hidden_state: [hidden_size]
LM head weight:     [vocab_size, hidden_size]
logits:             [vocab_size]

如果某个模型的 vocab_size = 129,280，那每次 decode 生成 1 个 token 时，LM Head 都会输出 129,280 个 logits。

这不是一次生成 129,280 个 token，而是给词表里的 129,280 个候选 token 各打一个分。

例如：

"都"   12.7
"府"    8.1
"市"    6.4
"."    -2.0
...

Sampling 接在 LM Head 后面：

logits
  -> temperature / top-k / top-p
  -> softmax or equivalent filtering
  -> pick next token

所以：

• LM Head 负责算所有候选 token 的分数。
• Sampling 负责按策略选出 1 个 next token。
• vocab logits 的数量等于 vocab_size。

6. Prefill 和 Decode 怎么放进这条线

Serving 视角下，同一个模型结构会跑在两个阶段：

Prefill

Prefill 把整段 prompt 一次性送进模型：

[batch_size, seq_len, hidden_size]
  -> Transformer blocks
  -> write KV cache

它并行处理 prompt token，主要目标是建立 KV cache，并得到首个可用于生成的状态。

Decode

Decode 每步生成一个新 token：

[batch_size, 1, hidden_size]
  -> read KV cache
  -> Attention / FFN or MoE
  -> LM Head
  -> Sampling
  -> next token
  -> append KV cache

Decode 会重复执行，直到生成完成。

7. 放到一张图里

上面的 SVG 是正式阅读图，下面的 Mermaid 保留为可编辑的结构草图。

flowchart LR
  A["Token IDs"] --> B["Embedding"]
  B --> C["Hidden state<br/>[batch, seq, hidden]"]
  C --> D["Attention<br/>read/write KV cache"]
  D --> E{"FFN type"}
  E --> F["Dense FFN"]
  E --> G["MoE router<br/>top-k experts"]
  G --> H["Dispatch / Expert compute / Combine"]
  F --> I["Updated hidden state"]
  H --> I
  I --> J["LM Head<br/>[hidden] -> [vocab]"]
  J --> K["Sampling"]
  K --> L["Next token"]

8. 对模拟器页面的解释口径

模型页里的 token flow 可以按两层理解：

1. Serving pipeline： Request queue -> Prefill -> KV cache -> Decode loop -> Sampling -> Response
2. Model block lens： Input tokens -> Embedding -> Attention + KV -> FFN/MoE -> Logits -> Output token

两层图讲的是同一个过程，只是观察角度不同：

• Serving pipeline 解释请求在推理服务里怎么跑。
• Model block lens 解释一次模型 forward 内部怎么更新 hidden state。

截图里的 129,280 vocab logits 属于 LM Head / Sampling 这一步。它来自模型配置里的 vocab_size，表示每次生成下一个 token 时，LM Head 会给 129,280 个词表 token 打分。

截图里的 6 / 384 experts 属于 MoE 这一步。它表示每个 token 的 hidden state 被 router 送到 384 个 routed experts 中的 6 个。

这两个数字不在同一层：

• 129,280 是输出词表维度。
• 6 / 384 是 MoE expert 路由维度。
• hidden state 的维度由 hidden_size 决定。

9. 相关页面

• MoE 推理：Expert 并行与调度机制
• CSA/HCA 注意力：DeepSeek-V4 的混合压缩稀疏机制
• GDN 与 Chunked Prefill — 混合 Attention/GDN 模型中 token state 与 chunk 边界的执行方式
• 模拟器建模指南：显存与吞吐公式
• DeepSeek-V3 Technical Report：中英对照解读