LLaMa 2 7B

                    ┌──────────────────────┐
                    │  下一个 token 概率     │
                    └──────────▲───────────┘
                               │
                    ┌──────────┴───────────┐
                    │  Linear (vocab head) │     ← (4096, 32000)
                    └──────────▲───────────┘
                               │
                    ┌──────────┴───────────┐
                    │   RMSNorm (final)    │
                    └──────────▲───────────┘
                               │
              ┌────────────────┴────────────────┐
              │   LLaMA Block × 32              │
              │  ┌───────────────────────────┐  │
              │  │  Add residual             │  │
              │  └─────────▲─────────────────┘  │
              │            │                    │
              │  ┌─────────┴─────────────────┐  │
              │  │  FFN (SwiGLU)             │  │
              │  │  hidden 4096 → 11008 →    │  │
              │  │  4096                     │  │
              │  └─────────▲─────────────────┘  │
              │            │                    │
              │  ┌─────────┴─────────────────┐  │
              │  │  RMSNorm (pre-FFN)        │  │
              │  └─────────▲─────────────────┘  │
              │            │                    │
              │  ┌─────────┴─────────────────┐  │
              │  │  Add residual             │  │
              │  └─────────▲─────────────────┘  │
              │            │                    │
              │  ┌─────────┴─────────────────┐  │
              │  │  Multi-Head Attention     │  │
              │  │  (causal, with RoPE)      │  │
              │  │  32 heads × 128 dim       │  │
              │  └─────────▲─────────────────┘  │
              │            │                    │
              │  ┌─────────┴─────────────────┐  │
              │  │  RMSNorm (pre-attn)       │  │
              │  └─────────▲─────────────────┘  │
              └────────────┼────────────────────┘
                           │
                    ┌──────┴───────┐
                    │  Embedding   │     ← (32000, 4096)
                    │  (no PE)     │
                    └──────▲───────┘
                           │
                    ┌──────┴───────┐
                    │  Token IDs   │
                    └──────────────┘

词表大小 V                     = 32,000
隐藏维度 d_model               = 4,096
层数 N                         = 32
头数 h                         = 32
每头维度 d_k                   = 4096 / 32 = 128
FFN 中间维度 d_ff              = 11,008（约 2.7×，因为 SwiGLU 结构特殊）
上下文长度 max_seq_len         = 4,096
参数量                          ≈ 7B

参数量

Embedding layer: 32000 * 4096 = 131,072,000

32 Decoder blocks:      6,476,267,520
One Decoder block:        202,383,360
  RMSNorm layer:                4,096
  Attention layer:         67,108,864
    WQ: 4096 * 4096 =      16,777,216
    WK: 4096 * 4096 =      16,777,216
    WV: 4096 * 4096 =      16,777,216
    WO: 4096 * 4096 =      16,777,216
  RMSNorm layer:                4,096
  FFN layer:              135,266,304
    Wgate: 4096  * 11008 = 45,088,768
    Wup:   4096  * 11008 = 45,088,768
    Wdown: 11008 * 4096  = 45,088,768

Final RMSNorm: 4,096

LM head: 4096 * 32000 = 131,072,000

Total:
Embedding            131,072,000
32 decoder blcks   6,476,267,520
Final RMSNorm              4,096
LM head              131,072,000
                   -------------
                   6,738,415,616

三件套

RMSNorm

RMSNorm = Root Mean Square Normalization —— LayerNorm 的简化版。

LayerNorm 做的事：

$LayerNorm (x) = γ \cdot \frac{x - μ}{\sqrt{σ^{2} + ϵ}} + β$

其中 μ 是均值，σ 是标准差，γ,β 是可学习参数。

两步：先去均值（centering），再除标准差（scaling）。

RMSNorm 对每个 token 的向量 $x$ 独立计算：

$RMSNorm (x) = γ \cdot \frac{x}{\sqrt{\frac{1}{d} \sum_{i} x_{i}^{2} + ϵ}}$

其中 $γ$ 是向量，每个维度都是可学习的缩放因子。

只做一步：只 scaling，不 centering。

为什么这样改？

实验发现，LayerNorm 里"去均值"那一步对效果贡献很小，但占了大约一半的计算量。

砍掉去均值以后：
- 计算更快（约 7-64% 的加速，取决于硬件和实现）
- 效果几乎不掉（在很多任务上甚至略好）
所以 RMSNorm 是工程优化——不是数学上更"正确"，是性价比更高。

LLaMA、Qwen、DeepSeek 等都用 RMSNorm——成了现代大模型的标配。
提出者

Zhang & Sennrich, 2019 — "Root Mean Square Layer Normalization"

提出时没火，直到 LLaMA 采用后才广泛流行。和 RoPE 一样的故事：LLaMA 的标志性效应让一些"小众但有用"的技术成为标准。

RoPE

2023 年 2 月，Meta 发布 LLaMA 1，选了 RoPE 作为位置编码。

为什么需要 RoPE

正弦余弦 PE：直接把位置向量加到 embedding 上。问题：

是绝对位置编码（位置 0、1、2...）
模型要"间接"学相对位置规律
外推性差（超过训练长度就不准）

RoPE 的目标：让 attention 直接看到相对位置 $j - i$ ，而不是绝对位置。

几何直觉

想象一个 2 维平面上的向量 $q$ 。如果把它旋转角度 $θ$ ——

原始 q:           旋转 θ 后的 q:
   ↑                    ↗
   |                   /
   |__→               /__

旋转矩阵：

$R (θ) = [\begin{matrix} \cos θ & - \sin θ \sin θ & \cos θ \end{matrix}]$

关键性质：

$R (θ_{1}) q \cdot R (θ_{2}) k = q \cdot R (θ_{2} - θ_{1}) k$

也就是——两个旋转后的向量做点积，结果只取决于旋转角度的差。

RoPE 怎么用这个性质

给位置 $i$ 的 query 旋转角度 $i \cdot θ$ ：

$q_{i}^{'} = R (i θ) \cdot q_{i}$

给位置 $j$ 的 key 旋转角度 $j \cdot θ$ ：

$k_{j}^{'} = R (j θ) \cdot k_{j}$

那么——

$q_{i}^{'} \cdot k_{j}^{'} = q_{i} \cdot R ((j - i) θ) \cdot k_{j}$

点积只依赖 $j - i$ （相对位置）——绝对位置 $i, j$ 自动消掉。

高维怎么做

LLaMA 的 head 维度是 128——不是 2 维。RoPE 把 128 维两两分组——每组 2 个维度做一次旋转，每组用不同的旋转频率。

128 维向量 → 64 对 → 64 个 2D 旋转

第 0 组（维度 0,1）：旋转频率 ω_0 = 1
第 1 组（维度 2,3）：旋转频率 ω_1 = θ_base^(-2/128)
第 2 组（维度 4,5）：旋转频率 ω_2 = θ_base^(-4/128)
...
第 63 组（维度 126,127）：旋转频率 ω_63 = θ_base^(-126/128)

LLaMA 用 $θ_{base} = 10000$ （和正弦 PE 一样）。

位置 pos 在第 k 组的旋转角度为 $p o s \cdot ω_{k}, ω_{k} = θ_{b a s e}^{- 2 k / d}$ 即 $p o s \cdot 10000^{- 2 k / d}$ 。

所以"高频组"快速旋转区分相邻位置，"低频组"慢速旋转区分远距离——和正弦 PE 的多频率思想一致。

实际上，$pos \cdot 10000^{-2k/d}$ 这个角度也被用在 PE 里面。具体来说，对位置 pos 的第 k 对，我们生成 sin 和 cos 的值，然后加到这一对上。

而 RoPE 则是用这个角度，对位置 pos 的第 k 对，进行旋转。

实际例子

Prompt：[BOS, The, capital, of, France]（5 个 token）

位置：position_id = 0, 1, 2, 3, 4

为了画面清爽——把 LLaMA 实际的 head_dim = 128 简化为 head_dim = 8。其他不变：

$θ_{base} = 10000$
把 8 维分成 4 对：(0,1), (2,3), (4,5), (6,7)

Step 1：算每对维度的旋转频率

每对的频率：$\omega_k = \theta_{\text{base}}^{-2k/d}$，其中 $d$ 是 head_dim。

第 0 对（维度 0,1）：ω_0 = 10000^(-0/8) = 10000^0       = 1.0
第 1 对（维度 2,3）：ω_1 = 10000^(-2/8) = 10000^(-0.25) ≈ 0.1
第 2 对（维度 4,5）：ω_2 = 10000^(-4/8) = 10000^(-0.5)  ≈ 0.01
第 3 对（维度 6,7）：ω_3 = 10000^(-6/8) = 10000^(-0.75) ≈ 0.001

频率从高到低——前面的对快速旋转，后面的对慢速旋转。

Step 2：算每个位置在每对上的旋转角度

旋转角度 =

'_' allowed only in math mode

$\text{position_id} \times \omega_k$

。

位置 0 (BOS):     [0×1.0, 0×0.1, 0×0.01, 0×0.001] = [0, 0, 0, 0]
位置 1 (The):     [1×1.0, 1×0.1, 1×0.01, 1×0.001] = [1.0, 0.1, 0.01, 0.001]
位置 2 (capital): [2×1.0, 2×0.1, 2×0.01, 2×0.001] = [2.0, 0.2, 0.02, 0.002]
位置 3 (of):      [3×1.0, 3×0.1, 3×0.01, 3×0.001] = [3.0, 0.3, 0.03, 0.003]
位置 4 (France):  [4×1.0, 4×0.1, 4×0.01, 4×0.001] = [4.0, 0.4, 0.04, 0.004]

每个位置有 4 个旋转角度——对应 4 对维度。 Step 3：算每个位置每对的 cos/sin 为了节省空间——只展示位置 1 (The) 的：

位置 1，第 0 对：角度 1.0   → cos(1.0)   ≈ 0.540,  sin(1.0)   ≈ 0.841
位置 1，第 1 对：角度 0.1   → cos(0.1)   ≈ 0.995,  sin(0.1)   ≈ 0.100
位置 1，第 2 对：角度 0.01  → cos(0.01)  ≈ 1.000,  sin(0.01)  ≈ 0.010
位置 1，第 3 对：角度 0.001 → cos(0.001) ≈ 1.000,  sin(0.001) ≈ 0.001

注意—— 高频对（第 0 对）：sin/cos 值有较大变化（0.540/0.841） 低频对（第 3 对）：sin/cos 几乎不变（≈1.0/0.001）这就是为什么 RoPE 能区分相邻位置（高频对差异大）也能区分远距离位置（低频对累积差异）。 Step 4：假设算出来的 Q 和 K 每个 token 经过

W^{Q}

投影得到 8 维 query 向量。假设位置 1 (The) 的 query 是：

$q_{The} = [\underset{第0对}{\underset{⏟}{0.5, 0.3}}, \underset{第1对}{\underset{⏟}{- 0.2, 0.4}}, \underset{第2对}{\underset{⏟}{0.1, - 0.3}}, \underset{第3对}{\underset{⏟}{0.6, 0.2}}]$

把它按对分组：

$q_{The} \to (0.5, 0.3), (- 0.2, 0.4), (0.1, - 0.3), (0.6, 0.2)$

每对就是一个 2D 向量——对它做 2D 旋转。

Step 5：对位置 1 的 q 应用 RoPE

第 0 对 $(0.5, 0.3)$ ，旋转角度 1.0：

$[\begin{matrix} q_{0}^{'} q_{1}^{'} \end{matrix}] = [\begin{matrix} \cos (1.0) & - \sin (1.0) \sin (1.0) & \cos (1.0) \end{matrix}] [\begin{matrix} 0.5 0.3 \end{matrix}] = [\begin{matrix} 0.540 & - 0.841 0.841 & 0.540 \end{matrix}] [\begin{matrix} 0.5 0.3 \end{matrix}]$

$= [\begin{matrix} 0.540 \times 0.5 - 0.841 \times 0.3 0.841 \times 0.5 + 0.540 \times 0.3 \end{matrix}] = [\begin{matrix} 0.018 0.583 \end{matrix}]$

第 1 对 $(- 0.2, 0.4)$ ，旋转角度 0.1：

$[\begin{matrix} q_{2}^{'} q_{3}^{'} \end{matrix}] = [\begin{matrix} 0.995 & - 0.100 0.100 & 0.995 \end{matrix}] [\begin{matrix} - 0.2 0.4 \end{matrix}] = [\begin{matrix} - 0.239 0.378 \end{matrix}]$

第 2 对 $(0.1, - 0.3)$ ，旋转角度 0.01：

$[\begin{matrix} q_{4}^{'} q_{5}^{'} \end{matrix}] = [\begin{matrix} 1.000 & - 0.010 0.010 & 1.000 \end{matrix}] [\begin{matrix} 0.1 - 0.3 \end{matrix}] = [\begin{matrix} 0.103 - 0.299 \end{matrix}]$

第 3 对 $(0.6, 0.2)$ ，旋转角度 0.001：

$[\begin{matrix} q_{6}^{'} q_{7}^{'} \end{matrix}] = [\begin{matrix} 1.000 & - 0.001 0.001 & 1.000 \end{matrix}] [\begin{matrix} 0.6 0.2 \end{matrix}] = [\begin{matrix} 0.600 0.201 \end{matrix}]$

合起来——位置 1 经过 RoPE 后的 query：

$q_{The}^{'} = [0.018, 0.583, - 0.239, 0.378, 0.103, - 0.299, 0.600, 0.201]$

Step 6：观察"高频对变化大、低频对变化小"

对比原始和旋转后的 q：

维度对	旋转前	旋转后	变化
第 0 对（高频）	$(0.5, 0.3)$	$(0.018, 0.583)$	大变化
第 1 对	$(- 0.2, 0.4)$	$(- 0.239, 0.378)$	中等变化
第 2 对	$(0.1, - 0.3)$	$(0.103, - 0.299)$	小变化
第 3 对（低频）	$(0.6, 0.2)$	$(0.600, 0.201)$	几乎无变化

这就是 RoPE 的"多频率"思想——

高频对快速旋转——即使位置差 1 也能明显区分
低频对慢速旋转——只有位置差很大时才能累积出明显差异

Step 7：所有位置的 RoPE 应用情况

位置 0 (BOS)——所有旋转角度都是 0—— $q_{BOS}$ 完全不变。

位置 4 (France)——旋转角度 $[4.0, 0.4, 0.04, 0.004]$ ——

第 0 对的旋转角度 4.0 是个比较大的角（约 229°）——$\cos(4.0) \approx -0.654$, $\sin (4.0) \approx - 0.757$ ——和位置 1 的旋转完全不同。

具体不算了，原理是一样的——每对独立做 2D 旋转。

Step 8：K 的处理完全一样

K 也按位置做相同的旋转——同一个位置的 q 和 k 用相同的旋转角度。

具体说

位置 1 的 k：用旋转角度 $[1.0, 0.1, 0.01, 0.001]$ 旋转
位置 2 的 k：用旋转角度 $[2.0, 0.2, 0.02, 0.002]$ 旋转
...

Step 9：算完旋转后做 attention

                  旋转角度 pos * theta^{-2k/d}
qBOS     kBOS     0*10000^(-0/8) 0*10000^(-2/8) 0*10000^(-4/8) 0*10000^(-6/8)
qThe     kThe     1*10000^(-0/8) 1*10000^(-2/8) 1*10000^(-4/8) 1*10000^(-6/8)
qcapital kcapital 2*10000^(-0/8) 2*10000^(-2/8) 2*10000^(-4/8) 2*10000^(-6/8)
qof      kof      3*10000^(-0/8) 3*10000^(-2/8) 3*10000^(-4/8) 3*10000^(-6/8)
qFrance  kFrance  4*10000^(-0/8) 4*10000^(-2/8) 4*10000^(-4/8) 4*10000^(-6/8)

经过 RoPE 后，所有 q 和 k 都旋转完毕——然后做正常的 attention：

$scores = softmax! (\frac{Q^{'} K^{' ⊤}}{\sqrt{d_{k}}}) V$

关键：算 $q_{i}^{'} \cdot k_{j}^{'}$ 时——

$q_{i}^{'} \cdot k_{j}^{'} = R (i θ) q_{i} \cdot R (j θ) k_{j} = q_{i} \cdot R ((j - i) θ) k_{j}$

点积只依赖位置差 $j - i$ ——这就是 RoPE 让 attention 看到相对位置的全部魔法。

         kBOS   kThe   kcapital  kof    kFrance
qBOS     0 - 0  1 - 0  2 - 0     3 - 0  4 - 0
qThe     0 - 1  1 - 1  2 - 1     3 - 1  4 - 1
qcapital 0 - 2  1 - 2  2 - 2     3 - 2  4 - 2
qof      0 - 3  1 - 3  2 - 3     3 - 3  4 - 3
qFrance  0 - 4  1 - 4  2 - 4     3 - 4  4 - 4

在 attention 流程里的位置

input → embedding → 第 1 层
                       ↓
                    RMSNorm
                       ↓
                    × W_Q, W_K, W_V → Q, K, V
                       ↓
                    RoPE（旋转 Q 和 K）   ← RoPE 在这里
                       ↓
                    attention 计算
                       ↓
                    ...

注意——

只对 Q 和 K 做旋转（V 不动，因为 V 不参与位置匹配）
每一层 attention 都做一次 RoPE——不是只在 input 层加一次（这和正弦 PE 不同）

三个关键优势

直接编码相对位置

模型不需要"间接"学相对位置规律——RoPE 把它直接植入了点积里。
外推性好

旋转操作对所有位置都有定义——超过训练长度也能稳定工作（实践中需要一些技巧，比如 NTK-aware scaling）。
不增加参数

RoPE 是无参数的（只有几个固定的旋转频率）——比 learned PE 省参数。

一个具体的工程细节

LLaMA 实现里 RoPE 是这样的——

# 预先计算所有位置、所有维度对的 cos 和 sin
# shape: (max_seq_len, head_dim/2)
cos_cached, sin_cached = precompute_rope(max_seq_len, head_dim, base=10000)

def apply_rope(q, k, position_ids):
    # q, k shape: (batch, n_heads, seq_len, head_dim)
    cos = cos_cached[position_ids]  # 查表
    sin = sin_cached[position_ids]

    # 把 q 的每两个相邻维度看作一对，做旋转
    q_rotated = rotate_pairs(q, cos, sin)
    k_rotated = rotate_pairs(k, cos, sin)

    return q_rotated, k_rotated

def rotate_pairs(x, cos, sin):
    # llama 实现里不会真的做矩阵乘法，这样可以避免矩阵乘法的开销
    ...
    # 每对 (x_i, x_{i+1}) 做旋转
    x_rotated_i    = x_i * cos[k] - x_i+1 * sin[k]
    x_rotated_i+1  = x_i * cos[k] - x_i+1 * sin[k]

预计算 cos/sin 表——避免每次 forward 都重新算三角函数。这是个工程优化，不影响数学。

SwiGLU

整体结构

x → ┌─ × W_gate → SiLU  ─┐
    │                    × → × W_down → 输出
    └─ × W_up ───────────┘

公式：

$FFN (x) = W_{down} \cdot (SiLU (W_{gate} \cdot x) \otimes (W_{up} \cdot x))$

其中 $\otimes$ 是逐元素相乘（element-wise multiplication）。可以看到没有 bias。

形状变化：

x:           (batch, seq_len, 4096)

W_gate:      (4096, 11008)        ← 升维
W_up:        (4096, 11008)        ← 升维（独立的另一个矩阵）

gate:        (batch, seq_len, 11008)
up:          (batch, seq_len, 11008)

SiLU(gate):  (batch, seq_len, 11008)
activated:   SiLU(gate) ⊗ up    (batch, seq_len, 11008)

W_down:      (11008, 4096)        ← 降维
output:      (batch, seq_len, 4096)

三个矩阵——比原版多一个。

细节解释

a. SiLU 是什么？

SiLU（Sigmoid Linear Unit）= $x \cdot σ (x)$ ——其中 $σ$ 是 sigmoid 函数。

也叫 Swish——

$SiLU (x) = x \cdot \frac{1}{1 + e^{- x}}$

关键差别：

ReLU：$x < 0$ 时严格输出 0
SiLU：$x < 0$ 时输出小的负值（不会完全死掉），且处处可导

实践中——SiLU 比 ReLU 略好一点，主要因为它"平滑"——梯度信号更稳定。

b. Gate 机制是什么？

这是 SwiGLU 的核心——也是它和原版 FFN 的本质差别。

gate 这个词来自 LSTM——意思是"控制信息流通的开关"。

具体看 SwiGLU 在做什么——

gate = SiLU(W_gate · x)        ← 这是"开关信号"，每个维度一个值
up = W_up · x                   ← 这是"实际内容"

output_inner = gate ⊗ up        ← 逐元素相乘

直觉：

$W_{up}$ 把 x 投到一个高维空间——存放"备选信息"
$W_{gate}$ 也把 x 投到一个高维空间——产出"哪些维度该被激活"的信号
SiLU 把 gate 信号变成 0~1 之间的"开关程度"
逐元素相乘——每个维度的内容被它对应的 gate 信号"过滤"

类比——

原版 FFN：所有信息都过一遍 ReLU——粗暴的"要么全保留，要么全丢"
SwiGLU：每个维度有自己的"开关"——精细控制"哪些信息让过、让多少过"

这种 gating 机制让模型表达力更强——能学到更复杂的"条件激活"模式。

c. 为什么中间维度是 11008 而不是 16384

原版 FFN 的中间维度 = $4 \times d_{model}$ （这是 Transformer 论文的经验值，效果好）。

SwiGLU 比原版多一个矩阵——如果还用 $4 \times 4096 = 16384$ 作中间维度——参数量会比原版 FFN 多 50%。

为了保持参数量大致相同——LLaMA 把中间维度调整为 $\frac{2}{3} \times 4 \times d_{model} = \frac{8}{3} \times 4096 \approx 10923$ ，再向上对齐到 256 的整数倍——得到 11008。

具体计算——

FFN 类型	矩阵数	中间维度	参数量
原版（ReLU）	2	16384	$2 \times 4096 \times 16384 \approx 134 M$
SwiGLU	3	11008	$3 \times 4096 \times 11008 \approx 135 M$

几乎一样——这是为了"等参数量比较"才这样调的。

d. 为什么没有 bias?

LLaMA 去掉了所有 Linear 层的 bias 项——

原版：$y = Wx + b$
LLaMA：$y = Wx$

为什么？

bias 项参数量小（相对于 W）但占了内存
实验发现去掉 bias 几乎不影响效果
简化代码、节省参数——纯工程优化

不只是 FFN——LLaMA 的 attention 投影矩阵 $W_{Q}, W_{K}, W_{V}, W_{O}$ 也都没有 bias。

SwiGLU 的来源

Shazeer, 2020 "GLU Variants Improve Transformer"——

Noam Shazeer（Transformer 论文作者之一，后来创办 Character.AI）系统比较了各种 GLU 变体——

ReGLU（用 ReLU）
GEGLU（用 GELU）
SwiGLU（用 SiLU/Swish） ⭐
...

发现 SwiGLU 在大多数任务上效果最好。

但论文里有一句很有意思的话——

"We offer no explanation as to why these architectures seem to work; we attribute their success, as all else, to divine benevolence."

意思是——"我们没法解释为什么 SwiGLU 更好——一切归功于神的恩赐"。

这是研究界的玩笑话——意思是 SwiGLU 是纯实证发现，没有清晰的理论解释——但它就是 work。

LLaMA 选 SwiGLU 也是基于实证经验——不是因为有什么理论保证。

Prefill

Prefill 的意义是什么

prefill 的意义就是产生上下文信息，更精确地说：

Prefill 让 prompt 里的每个 token，通过 32 层 attention + FFN 的处理，得到包含完整上下文信息的"上下文化向量"——并且把每一层每个位置的 K/V 都存进 KV cache，供后续 decode 阶段使用。

它做了两件事：

让最后一个位置（"is"）产生一个"知道整个 prompt"的输出向量——用来预测第一个新 token
把 KV cache 填好——这是它的"额外副产品"，但对后续 decode 至关重要

具体例子

Prompt："The capital of France is"

Tokenize 后（假设用 LLaMA 2 的 tokenizer）：

"The"      → token id  450
"capital"  → token id  7483
"of"       → token id  310
"France"   → token id  3444
"is"       → token id  338

input_ids（HuggingFace 输入格式）：

input_ids = [[1, 450, 7483, 310, 3444, 338]]   # 加了 BOS（id=1），shape (1, 6)

attention_mask：

attention_mask = [[1, 1, 1, 1, 1, 1]]   # 所有位置都是有效 token，shape (1, 6)

注意——attention_mask 是 HuggingFace 用来标记 PAD 的——1 表示真实 token，0 表示 PAD。这里没有 PAD 所以全是 1。

Prefill 的数据流（详细版）

让我把数据从输入到输出的每一步都画出来。

Step 1：Token Embedding

input_ids: [1, 450, 7483, 310, 3444, 338]      shape (1, 6)
                ↓ 查 embedding 矩阵 W_E (32000, 4096)
embeddings:                                    shape (1, 6, 4096)
  位置 0: e_BOS      (4096 维向量)
  位置 1: e_The      (4096 维向量)
  位置 2: e_capital  (4096 维向量)
  位置 3: e_of       (4096 维向量)
  位置 4: e_France   (4096 维向量)
  位置 5: e_is       (4096 维向量)

注意，LLaMA 2 用 RoPE，所以这里不加 PE。embedding 直接进入第一层。位置信息通过 RoPE 在 attention 计算时注入。

Step 2：进入第 1 层 Decoder Block

LLaMA 2 7B 有 32 层——每层结构相同。让我们看第 1 层。

每一层的伪代码——

def llama_decoder_block(x, kv_cache_layer):
    # Sublayer 1: Attention
    h = rms_norm(x)                               # pre-norm
    h = self_attention(h, kv_cache_layer)         # 算 attention，更新 cache
    x = x + h                                     # 残差连接

    # Sublayer 2: FFN
    h = rms_norm(x)
    h = ffn(h)
    x = x + h                                     # 残差连接

    return x

让我把每一步拆开讲。

Step 2.1：第一次 RMSNorm

输入: embeddings shape (1, 6, 4096)
        ↓ RMSNorm（逐 token、逐 hidden_dim 归一化）
normalized_x: shape (1, 6, 4096)

RMSNorm 让每个 token 向量的"幅度"标准化——为后续 attention 做准备。

Step 2.2：算 Q、K、V

normalized_x: (1, 6, 4096)
       │
       ├── × W_Q (4096, 4096) → Q: (1, 6, 4096)
       ├── × W_K (4096, 4096) → K: (1, 6, 4096)
       └── × W_V (4096, 4096) → V: (1, 6, 4096)

接着 reshape 成 multi-head 形式——

Q: (1, 6, 4096) → reshape → (1, 32, 6, 128)   # 32 个头，每头 128 维
K: (1, 6, 4096) → reshape → (1, 32, 6, 128)
V: (1, 6, 4096) → reshape → (1, 32, 6, 128)

关键步骤：RoPE 旋转——

Q: (1, 32, 6, 128)
       ↓ 应用 RoPE（基于位置 0,1,2,3,4,5 旋转 Q）
Q_rotated: (1, 32, 6, 128)

K: (1, 32, 6, 128)
       ↓ 应用 RoPE（基于位置 0,1,2,3,4,5 旋转 K）
K_rotated: (1, 32, 6, 128)

V 不需要 RoPE——只有 Q 和 K 需要旋转（因为位置信息只影响"匹配过程"）。

Step 2.3：把 K, V 存进 KV cache

这是 prefill 阶段的关键动作——

kv_cache[layer_0].keys = K_rotated         # shape (1, 32, 6, 128)
kv_cache[layer_0].values = V               # shape (1, 32, 6, 128)

6 个位置的 K/V 全部存起来——decode 阶段就不用重算了。

Step 2.4：算 Attention

Q_rotated: (1, 32, 6, 128)
K_rotated: (1, 32, 6, 128)
V:         (1, 32, 6, 128)

算 attention scores——

scores = Q × K^T / √128

shape: (1, 32, 6, 6)    # 每个头有一个 6×6 attention 矩阵

应用 causal mask——

              k_BOS  k_The  k_cap  k_of   k_Fr   k_is
q_BOS     [  ?     -∞    -∞    -∞    -∞    -∞   ]
q_The     [  ?     ?     -∞    -∞    -∞    -∞   ]
q_cap     [  ?     ?     ?     -∞    -∞    -∞   ]
q_of      [  ?     ?     ?     ?     -∞    -∞   ]
q_Fr      [  ?     ?     ?     ?     ?     -∞   ]
q_is      [  ?     ?     ?     ?     ?     ?    ]

注意——上三角（未来位置）被设成 $- \infty$ 。

为什么 prefill 也要 causal mask？

因为 LLaMA 是 decoder-only 架构——它的 attention 永远是 causal 的。即使在 prefill 阶段（一次性处理整个 prompt），每个位置也只能看到自己之前的 token。

否则——位置 0 的 BOS 看到了未来的 "France" 和 "is"——它的输出向量就被"未来信息"污染了。但这违反了模型训练时的约定（模型训练时是 causal）——会导致输出分布偏离训练分布。

Softmax 后——

attention_weights: (1, 32, 6, 6)    # 每行加起来 = 1，PAD/未来位置为 0

乘以 V——

attention_output: attention_weights × V
shape: (1, 32, 6, 128)

reshape 回——

attention_output: (1, 32, 6, 128) → (1, 6, 4096)

最后一次线性投影——

attention_output × W_O (4096, 4096) → (1, 6, 4096)

Step 2.5：第一个残差连接

x = embeddings + attention_output       # shape (1, 6, 4096)

Step 2.6：FFN（SwiGLU）

x: (1, 6, 4096)
   ↓ RMSNorm
normalized_x: (1, 6, 4096)
   ↓ SwiGLU FFN
   |   gate = x × W_gate (4096, 11008)    # shape (1, 6, 11008)
   |   up = x × W_up (4096, 11008)        # shape (1, 6, 11008)
   |   activated = silu(gate) * up         # 逐元素相乘
   |   ffn_out = activated × W_down (11008, 4096)
   ↓
ffn_output: (1, 6, 4096)

注意——SwiGLU 比普通 ReLU FFN 多一个 gate 矩阵——这是为什么 FFN 中间维度是 11008 而不是标准的 16384（4×4096）。

Step 2.7：第二个残差连接

layer_1_output = x + ffn_output       # shape (1, 6, 4096)

Step 3：经过所有 32 层

layer_1_output → layer_2 → layer_2_output → layer_3 → ...
                                                    → layer_32_output

每一层都做相同的事——算 Q/K/V → 存 KV cache → attention → FFN → 残差。

最后一层输出的形状还是 (1, 6, 4096)——但每个位置的向量都已经"知道"了完整的 prompt 上下文。

Step 4：最后的 RMSNorm

layer_32_output: (1, 6, 4096)
       ↓ Final RMSNorm
final_hidden: (1, 6, 4096)

Step 5：取最后一个位置 + Linear + Softmax

final_hidden[:, -1, :]: (1, 4096)   # 只取位置 5（"is"）的向量
       ↓ × W_lm_head (4096, 32000)
logits: (1, 32000)                   # 词表大小的实数向量
       ↓ softmax
probs: (1, 32000)                    # 概率分布
       ↓ argmax 或采样
next_token_id = 7042                 # 假设对应 "Paris"

Prefill 完成——预测出第一个新 token "Paris"。

Prefill 之后的状态

Prefill 结束时，几个关键状态——

1. KV cache 已经填满 6 个位置

KV cache (32 层 × 32 头):
  位置 0: BOS 的 K, V
  位置 1: The 的 K, V
  位置 2: capital 的 K, V
  位置 3: of 的 K, V
  位置 4: France 的 K, V
  位置 5: is 的 K, V

2. 已生成的序列：[BOS, The, capital, of, France, is, Paris]——最后一个是 prefill 预测出来的。

3. 准备进入 decode 阶段：把 "Paris" 这一个 token 喂进模型，继续生成。

attention_mask 的作用

你提到 HuggingFace 的 input 里有 attention_mask——让我专门讲这个。

attention_mask 是一个 0/1 矩阵——

1 表示"这个位置是真实 token"
0 表示"这个位置是 PAD"

它的作用：在 attention 计算时，mask 掉 PAD 列——确保真实 token 不会从 PAD 那里吸收信息。

在你的例子里：

input_ids = [[1, 450, 7483, 310, 3444, 338]]      # 6 个真实 token
attention_mask = [[1, 1, 1, 1, 1, 1]]              # 全 1，没有 PAD

attention_mask 全是 1——所以实际上没有任何位置被屏蔽掉——只有 causal mask 生效。

如果有 PAD 怎么办？（batch 处理时会发生）

假设你 batch 里两个 prompt——

Prompt 1: "The capital is"        (4 个 token)
Prompt 2: "Hello"                  (2 个 token)

batch 后（左 padding）：
input_ids = [
    [1, 450, 7483, 338],          # Prompt 1
    [PAD, PAD, 1, 15043]           # Prompt 2（左边补 PAD）
]
attention_mask = [
    [1, 1, 1, 1],                  # Prompt 1 全有效
    [0, 0, 1, 1]                   # Prompt 2 前 2 个是 PAD
]

attention 计算时——

第二个样本的 PAD 列在 softmax 前设为 $- \infty$
真实 token 不会从 PAD 吸收信息

注意——LLaMA 在推理时用的是左 padding（PAD 在前面）——这是因为生成时只看最后一个位置的输出，左 padding 让所有 batch 内样本的"最后一个位置"对齐。

完整的可视化

让我把整个 prefill 流程用一张图概括——

input_ids: [1, 450, 7483, 310, 3444, 338]          (1, 6)
attention_mask: [1, 1, 1, 1, 1, 1]                 (1, 6)
       │
       ↓
┌─────────────────────────────────────────┐
│  Embedding (32000, 4096)                │
└─────────────────────────────────────────┘
       │
       ↓ embeddings: (1, 6, 4096)
┌─────────────────────────────────────────┐
│  Layer 1                                │
│  ┌────────────────────────────────────┐ │
│  │  RMSNorm                           │ │
│  ├────────────────────────────────────┤ │
│  │  Attention (with RoPE, causal mask)│ │
│  │  Q, K, V (1, 32, 6, 128)           │ │
│  │  K, V → 写入 KV cache              │  │
│  │  attention 矩阵 (1, 32, 6, 6)      │  │
│  │  output (1, 6, 4096)               │  │
│  ├────────────────────────────────────┤  │
│  │  + 残差                            │  │
│  ├────────────────────────────────────┤  │
│  │  RMSNorm                           │  │
│  ├────────────────────────────────────┤  │
│  │  SwiGLU FFN                        │  │
│  │  4096 → 11008 → 4096               │  │
│  ├────────────────────────────────────┤  │
│  │  + 残差                            │  │
│  └────────────────────────────────────┘  │
└──────────────────────────────────────────┘
       │
       ↓ layer_1_output: (1, 6, 4096)
        ...（重复 31 次）
       │
       ↓ layer_32_output: (1, 6, 4096)
┌─────────────────────────────────────────┐
│  Final RMSNorm                          │
└─────────────────────────────────────────┘
       │
       ↓ final_hidden: (1, 6, 4096)
       │
       ↓ 取最后一个位置 [-1]
       │
       ↓ last_hidden: (1, 4096)
┌─────────────────────────────────────────┐
│  Linear (4096, 32000) → logits          │
│  Softmax → probs                        │
└─────────────────────────────────────────┘
       │
       ↓ next_token_id (e.g., "Paris")

微妙但重要的点：为什么计算了 6 个位置的输出，但只用最后一个？

你可能会想——既然只用位置 5 的输出预测下一个 token，为什么不只算位置 5 的？

两个原因——

1. attention 的本质需要看前面所有位置

要算位置 5 的 attention 输出——必须有位置 0-4 的 K 和 V。所以前面的 K/V 必须算出来。

2. 顺便就把 KV cache 填了

一次性算所有位置的 K/V，正好填满 KV cache——decode 阶段直接复用，不用再算。

所以 prefill 看起来"算了 6 个位置但只用了 1 个"——但前面 5 个位置的 K/V 不是浪费，它们是 KV cache 的核心内容。

总结

Prefill 的意义有两层：

让最后一个位置（"is"）的输出向量"知道"整个 prompt 的上下文——用来预测第一个新 token
把 32 层 × 32 头 × 6 个位置的 K/V 全部填进 KV cache——为后续 decode 阶段铺好基础

数据流：input_ids → embedding → 32 层 (RMSNorm + Attention + FFN) → final RMSNorm → 取最后位置 → Linear + Softmax → 下一个 token

attention_mask 是 HuggingFace 用来标记 PAD 的——配合 batch 处理时屏蔽 PAD 位置——你单条样本的情况下它全是 1。

理解了 prefill 的数据流——decode 就简单了：decode = prefill 的"单 token 版本"——只算 1 个新 token，但要从 KV cache 读前面所有 K/V。

Decode

前提状态——prefill 刚结束，此时——

KV cache 里有 32 层 × 32 头 × 6 个位置的 K 和 V
已生成的序列：[BOS, The, capital, of, France, is, Paris]（最后一个是 prefill 预测的）

decode 的目标——把 "Paris" 喂回模型，预测下一个 token（可能是 "."、"," 或别的）。

数据流

Step 1：输入只有 1 个 token

input: "Paris" 这一个 token  (token_id, 比如 7042)
   ↓
input_ids: [[7042]]    shape (1, 1)

关键差别——不是把整个序列 [BOS, The, capital, of, France, is, Paris] 重新喂进来——只喂这一个新 token。

为什么不喂整个序列？因为前 6 个位置的 K/V 已经在 cache 里了——重新算一遍就是浪费。

Step 2：Embedding

input_ids: [[7042]]              shape (1, 1)
            ↓ 查 W_E
embeddings: shape (1, 1, 4096)   ← 只有 1 个向量

Step 3：进入第 1 层

3.1 RMSNorm + 算 Q/K/V

x:           (1, 1, 4096)
   ↓ RMSNorm
   ↓ × W_Q, W_K, W_V
Q, K, V:     (1, 1, 4096)
   ↓ reshape multi-head
Q, K, V:     (1, 32, 1, 128)     ← 注意 seq_len 维度只有 1

关键——只有 1 个新 token，所以只算出 1 个 q、1 个 k、1 个 v（每个 head）。

3.2 RoPE

q (位置 6) 应用 RoPE → q'   ← 注意位置是 6（不是 0）
k (位置 6) 应用 RoPE → k'

为什么位置是 6？因为前面已经有 6 个 token（位置 0-5）——新 token "Paris" 是第 7 个，position_id = 6。

RoPE 怎么知道这个位置？需要传入 position_ids = [6]——HuggingFace 的 generate 函数内部会自动维护这个状态。

3.3 把新的 K/V 拼到 cache 里

之前的 KV cache（这一层）：
  K_cache:  (1, 32, 6, 128)     ← 6 个旧位置的 K
  V_cache:  (1, 32, 6, 128)     ← 6 个旧位置的 V

新算出来的 k', v'（这一层）：
  k_new:    (1, 32, 1, 128)
  v_new:    (1, 32, 1, 128)

拼接：
  K_full = concat([K_cache, k_new], dim=seq) → (1, 32, 7, 128)
  V_full = concat([V_cache, v_new], dim=seq) → (1, 32, 7, 128)

更新 cache：
  KV cache（这一层）现在是 7 个位置

这是 KV cache 的核心动作——增量地往里加新 K/V。

3.4 算 Attention

TODO

q':       (1, 32, 1, 128)        ← 只有 1 个 query
K_full:   (1, 32, 7, 128)        ← 7 个 key
V_full:   (1, 32, 7, 128)        ← 7 个 value

scores = q' × K_full^T / √128
shape: (1, 32, 1, 7)              ← 1 行 7 列

关键差别——

Prefill 阶段：attention 矩阵是 6×6（每个 token 都要算一行）
Decode 阶段：attention 矩阵是 1×7（只有当前新 token 这一行）

这就是为什么 decode 远比 prefill 计算量小——只有 1 个 query，矩阵-向量乘法。

Causal mask 还需要吗？严格说不需要——因为新 token (位置 6) 本来就只关注前 7 个位置（0-6），没有"未来"可看。但实现上通常还是会画个 mask（全是 1，没有 $- \infty$ ）保持一致性。

softmax → attention weights: (1, 32, 1, 7)
        × V_full → output: (1, 32, 1, 128)
        → reshape → (1, 1, 4096)
        × W_O → (1, 1, 4096)

3.5 残差 + FFN

x = x + attention_output           # (1, 1, 4096)
   ↓ RMSNorm
   ↓ SwiGLU FFN
ffn_output: (1, 1, 4096)
x = x + ffn_output                 # (1, 1, 4096)

FFN 也只对这 1 个 token 算——比 prefill 少 6 倍计算量。

Step 4：经过所有 32 层

每层重复 Step 3——每层都从自己的 KV cache 里读旧 K/V，算 attention，更新 cache。

最后输出 (1, 1, 4096)——只有 1 个向量，对应新 token "Paris" 经过完整 32 层处理后的表示。

Step 5：Final RMSNorm + Linear + Softmax

final_hidden: (1, 1, 4096)
   ↓ Final RMSNorm
   ↓ × W_lm_head (4096, 32000)
logits: (1, 32000)
   ↓ softmax
probs: (1, 32000)
   ↓ argmax / sampling
next_token_id = ... (比如 "." 对应的 id)

预测出下一个 token——比如 "."。

Step 6：循环

while 还没生成完：
    1. 把上一步预测的 token 作为输入
    2. 走一遍 Step 2-5 的流程
    3. 得到下一个 token
    4. 如果是 EOS，停止；否则继续

每一步 decode：

KV cache 长度 +1
序列长度 +1
生成 1 个新 token

Decode 第 n 步与第 1 步的对比

只有计算 attention 的部分，第 n 步和第 1 步不一样，其他部分都是一样的。

attention 部分的计算量随累积长度线性增长：

Decode step	KV cache 位置数	Attention 矩阵	该步 attention 计算量
1	7	1×7	7 次向量（128d）点积算注意力，7 次标量乘向量，6 次向量相加算输出向量
100	106	1×106	106 次向量（128d）点积算注意力，106 次标量乘向量，105 次向量相加算输出向量
1000	1006	1×1006	1006 次向量（128d）点积算注意力，1006 次标量乘向量，1005 次向量相加算输出向量

softmax 那一步，因为 attention 矩阵尺寸的变化，计算量也有变化，我没有在表格里列出来。

decode 大部分时间花在 FFN 和 Q K V 投影上

但是，实际上 decode 大部分时间不是花在 attention 上，而是花在 FFN 和 Q/K/V 投影上，即使在 step 1000。为什么？

因为虽然 attention 计算量从 7×128 涨到 1006×128（涨了 144 倍），但 FFN 的计算量绝对值更大（涉及 4096 → 11008 → 4096 的两次大矩阵乘）。只有当序列长度极大（几万 token）时，attention 才会成为 decode 的主要瓶颈。

所以：
- 短序列（< 2K）：decode 的 GPU 时间主要花在 FFN 和投影
- 长序列（> 8K）：attention 占比开始上升，但 FFN 仍然不可忽略
- 极长序列（> 32K）：attention 成为主要瓶颈
真正的瓶颈是内存带宽

但要注意，计算量 ≠ 实际耗时。

decode 的真正瓶颈是内存带宽（HBM bandwidth），不是计算量。

每一步 decode 需要：
- 读取所有模型权重（约 14 GB，FP16）
- 读取整个 KV cache（随 step 增长）
- 算完写回输出
而每一步只算 1 个 token——计算单元大量空闲，瓶颈在"数据搬运"。

这就是为什么 decode 是 memory-bound——GPU 算力浪费，时间花在等数据从 HBM 到 SRAM。

Prefill vs Decode 对比表

维度	Prefill	Decode
输入 token 数	N（整个 prompt）	1（只有新 token）
Q 数量	N	1
K, V 计算	全部 N 个新算	只算 1 个新的，从 cache 读 N-1 个
Attention 矩阵	N×N	1×(N+1)
FFN 处理	N 个 token 并行	1 个 token
KV cache 操作	写入 N 个位置	追加 1 个位置
主要算子	矩阵-矩阵乘法	矩阵-向量乘法
GPU 利用率	高（compute-bound）	低（memory-bound）
并行性	序列内并行	必须串行

矩阵乘法的形状：

Q K V 投影

Prefill: (B, L, h) x WQ/WK/WV (h, h) = (B, L, h)
Decode:  (B, 1, h) x WQ/WK/WV (h, h) = (B, 1, h)

算 attention

Prefill:
Q (B, n_heads, L, d_k) × K^T (B, n_heads, d_k, L) → scores (B, n_heads, L, L)
scores (B, n_heads, L, L) × V (B, n_heads, L, d_k) = Output (B, n_heads, L, d_k)
                                                   = Output (B, L, h)
Decode:
Q (B, n_heads, 1, d_k) × K^T (B, n_heads, d_k, ?) → scores (B, n_heads, 1, ?)
scores (B, n_heads, 1, ?) × V (B, n_heads, ?, d_k) = Output (B, n_heads, 1, d_k)
                                                   = Output (B, 1, h)

attention output 投影

Prefill:
attention_output (B, L, h) × W_O (h, h) = (B, L, h)

Decode:
attention_output (B, 1, h) × W_O (h, h) = (B, 1, h)

FFN

Prefill:
W_gate: (B, L, h) × (h, h_ffn) → (B, L, h_ffn)   FLOPs: 2 × B × L × h × h_ffn
W_up:   (B, L, h) × (h, h_ffn) → (B, L, h_ffn)   FLOPs: 2 × B × L × h × h_ffn
W_down: (B, L, h_ffn) × (h_ffn, h) → (B, L, h)   FLOPs: 2 × B × L × h × h_ffn

Decode:
W_gate: (B, 1, h) × (h, h_ffn) → (B, 1, h_ffn)   FLOPs: 2 × B × 1 × h × h_ffn
W_up:   (B, 1, h) × (h, h_ffn) → (B, 1, h_ffn)   FLOPs: 2 × B × 1 × h × h_ffn
W_down: (B, 1, h_ffn) × (h_ffn, h) → (B, 1, h)   FLOPs: 2 × B × 1 × h × h_ffn