ds4.c 学习笔记

Appearance

将原始分数（logits）转换为概率分布，是注意力和采样的核心归一化操作。

为什么需要 Softmax ​

注意力分数是 Q 和 K 的点积，值域没有约束（可以是任意实数）。我们需要把这些分数变成"权重"——非负且总和为 1，才能对 V 做加权平均。同理，模型输出 129280 个 logits，也需要归一化为概率才能采样。

核心原理 ​

输入: 任意实数 [x₁, x₂, ..., xₙ]
输出: 概率分布 [p₁, p₂, ..., pₙ]，每个 pᵢ ∈ (0,1)，Σpᵢ = 1

公式: pᵢ = exp(xᵢ - max) / Σexp(xⱼ - max)
              ↑ 减最大值是数值稳定技巧，防止 exp 溢出

不减最大值的问题：如果最大值是 1000，exp(1000) = ∞（float 溢出），整个计算产生 NaN。减去最大值后最大的 exp 项 = exp(0) = 1，其余都 < 1，安全。

在 ds4.c 中的实现 ​

c
// 注意力中的内联 softmax（ds4.c 约 4751 行）
// 1. 找最大值
for (uint32_t r = 0; r < n_rows; r++)
    if (score[r] > max) max = score[r];
// 2. 减最大值 → exp → 累加
for (uint32_t r = 0; r < n_rows; r++) {
    score[r] = expf(score[r] - max);
    sum += score[r];
}
// 3. 归一化
for (uint32_t r = 0; r < n_rows; r++)
    score[r] /= sum;

采样阶段使用 double 精度累加，减少 129280 个浮点数的累加误差。

ds4.c 还引入了 Attention Sink 机制：每个注意力头有一个可学习的 sink logit，参与 softmax 分母但不贡献 value，让模型学会"不关注任何位置"的能力。详见 Part 3 — 注意力机制。

相关概念 ​

  • kv-cache — softmax 在缓存的 KV 行上计算注意力权重
  • rope — RoPE 旋转后的 Q·K 点积作为 softmax 输入
  • moe — Router 使用 sqrt(softplus(x)) 而非标准 softmax 归一化专家分数

详见 Part 2 — 采样策略。

出链