Day 8·1 自注意力机制:Q·K^T/softmax/V的在线计算

1. 知识点:三步公式

对第 t 个位置的 query,注意力要"回看"前面所有已缓存的 token:

text 复制代码
score(pos) = (Q_t · K_pos) / sqrt(d)    ① 点积打分(d = head_dim,防数值爆炸的缩放)
w(pos)     = softmax(score(0..t))[pos]  ② 归一化成权重
O_t        = Σ_pos w(pos) · V_pos        ③ 加权求和输出
  • ① 为什么除 √d:Q、K 各维度独立同分布时,点积的量级随 d 增长(约 √d)。不缩放,长维度时分数会越来越大,softmax 输入跑到 exp 的饱和/溢出区;
  • ② softmax 是"赢家通吃":把分数变成 0--1 权重,让高分 token 主导输出;
  • ③ 加权 V :输出 = 缓存的 V 按注意力权重混合。V 不参与打分------它只被"挑"。

2. 对应代码:引擎里最朴素的参考实现

vllm_transformer.c 有一段"教科书式"的因果注意力(第 77--128 行),三步与公式一一对应:

c 复制代码
/* Scaled Dot-Product Attention with Causal Mask & KV-Cache
 * For position `cur_pos` (the token we're generating):
 *   2. Compute scores = Q[cur_pos] @ K[0..cur_pos]^T / sqrt(head_dim) */
float scale = 1.0f / sqrtf((float)head_dim);
for (int h = 0; h < num_heads; h++) {
    const float *qh = q + h * head_dim;
    float max_score = -1e9f;
    for (int t = 0; t < seq_len; t++) {
        const float *kh = k_cache + t * kv_dim + h * head_dim;   /* ① K 缓存行 */
        float dot = 0.0f;
        for (int d = 0; d < head_dim; d++) dot += qh[d] * kh[d];
        scores[t] = dot * scale;                                  /* ÷√d */
        if (scores[t] > max_score) max_score = scores[t];
    }
    /* Softmax:先减 max 再 exp(数值稳定) */
    float sum_exp = 0.0f;
    for (int t = 0; t < seq_len; t++) { scores[t] = expf(scores[t] - max_score); sum_exp += scores[t]; }
    for (int t = 0; t < seq_len; t++) scores[t] /= sum_exp;      /* ② softmax */
    /* Weighted sum of V */
    for (int d = 0; d < head_dim; d++) {
        float val = 0.0f;
        for (int t = 0; t < seq_len; t++) {
            const float *vh = v_cache + t * kv_dim + h * head_dim;
            val += scores[t] * vh[d];                             /* ③ Σ w·V */
        }
        output[h * head_dim + d] = val;
    }
}

注意两个"又朴素又正确"的细节(它们就是后面所有优化版的"及格线"):

  1. 因果:只回看 t ≤ 当前(scores 数组长度 seq_len ≤ cur_pos+1),不偷看未来;
  2. 一次读一行 KV :k_cache + t*kv_dim + h*head_dim 表示"第 t 个 token、第 h 个头"的连续 head_dim 个元素------KV 是 token 优先(token-major)排布的

同样把这三步写在"paged 缓存"变体里的是 vllm_attention.c 的 exact_attention(第 44--94 行:逐块 kvcache_read → 打分 → softmax_normalize → 加权 V),数学完全一致,只是取 KV 走块表(block table)。

真实推理不会跑这个三重循环(太慢),会换成 q8 KV + NEON 单遍版。但参考实现的价值就是"慢而绝对对"------优化版全部以它为对拍对象。

3. 改动后果:去掉 √d 缩放,看分数漂哪去

数学上这一步最容易被"优化掉":把 dot * scale 的 scale 去掉,只留 dot。看起来"反正是线性缩放,softmax 里会被约掉"------对单个分数成立,对整条分数不成立:softmax 的分母是 Σ exp,去掉 scale 后 exp 的输入整体变大,逐步逼近 exp 的溢出阈值(fp32 里 exp(89) 已经 inf)。维度越大(d=128 时点积量级约 ±20--30 甚至更高),越危险。

纪律:1/√d 是公式的一部分,不是可选的。删掉它,短上下文可能看不出问题,长上下文/大数值一上来就 NaN------这种 bug 只在压力测试里显形,最阴险。

4. 学员调试任务

  • A 档(板端动手):跑引擎自检前,先用纸笔算一个"2 头 × 4 token"的示例(自己给 Q/K/V 数字):手算 score 三行、softmax 权重四列、加权 O 一行,再对照代码第 98--126 行逐步核对。
  • B 档(纯读源码):读 vllm_transformer.c 第 77--128 行与 vllm_attention.c 第 44--94 行,标出两版各自对应的公式 ①②③ 行号,指出"因果在哪一行体现、√d 在哪一行乘"。

预期输出:你能脱稿写出三步公式,并指出引擎参考实现里"打分/归一化/加权 V"与"因果/KV 读取"各在哪几行。

5. 收尾

  • 本篇源码点名:vllm_transformer.c(scaled dot-product attention 第 77--128 行)、vllm_attention.c(exact_attention 第 44--94 行)。
  • 开源仓库:Kestrel-LLM (Gitee)(源码可得双许可:学习 / 学术研究免费)。
  • 下篇预告:公式里那句"先减 max 再 exp"不是洁癖------下一篇 8-2 亲手把 e^x 算爆,看 naive softmax 怎么在大上下文里变 NaN,以及引擎的在线版本怎么绕开它。

上一篇: Day 7·3 预热与测量抖动:为什么性能基准报告要"重复 3 次取中位"

下一篇: Day 8·2 在线softmax数值稳定性:为什么不能先算e^x再除

相关推荐
来让爷抱一个1 小时前
2026 纹理一致性实战:贴图不许花脸,百智云精灵图把材质契约写进素材包
人工智能·机器学习·材质·贴图
幻影123!1 小时前
AlphaZero 五子棋实战(五):评估工具 —— 学会在错误的地方掉头
人工智能·机器学习·强化学习·alpha zero
jbk33111 小时前
PixiuCut「视频批量分割」支持智能镜头分割、按时长、按数量分割,可定义最小时长避免无效片段等多可选参数可设置
人工智能·音视频·剪辑软件·剪映自动化软件
诺伦1 小时前
RiseClaw玄策:GEO优化工程实战,从零搭建生成式引擎优化体系
人工智能·chatgpt
星禾元亨1 小时前
企业 AI 落地 5 步路线图:诊断、企业知识库与 GEO 获客的工程化实施步骤
大数据·人工智能·自动化·创业创新
苏打水com1 小时前
内容合规红线:大模型生成内容,哪些能做哪些会犯法?
人工智能·安全·大模型
南京兴帝文化传媒有限公司2 小时前
本地生活服务商户GEO优化技术实践:AI大模型收录机制与地图POI权重算法拆解
大数据·人工智能·算法·生活·geo优化实操·csdn运营技巧·ai内容收录
2601_962293532 小时前
人工智能 & 神经网络完整入门路线(零基础可走,分阶段)
人工智能·python·深度学习·神经网络·机器学习