【声明】本博客所有内容均为个人业余时间创作所述技术案例均来自公开开源项目如GithubApache基金会不涉及任何企业机密或未公开技术如有侵权请联系删除标题232、【AI】【模型部署】基座模型研究一次前向、KV cache 与采样背景上篇 blog【AI】【模型部署】从零手写基座(二)训练、loss 与生成跑通了预训练闭环x与y错开一位做 next-token 预测、交叉熵算误差、反向传播更新权重135 万参数的微型模型在 CPU 上训 1200 步3 分 8 秒loss 从 7.3883≈ln1432纯乱猜降到 0.7821能续写出格式正确的博客文本——小模型的会写是统计模仿而 Qwen2 用的是同一套训练目标。本篇回到推理侧拿真正的 Qwen2.5-0.5B把一次前向的形状、KV cacheKey-Value cache键值缓存 到底省什么和采样怎么选下一个字这三件事用实验看清模型部署训练解决模型怎么变聪明推理解决用起来多快、输出怎么选。前面几篇一直在说 KV cache、temperature、top-k本篇用真模型把它们一次测清楚。一次前向形状从哪来、到哪去给模型一个单词用 hook 抓词嵌入的输出实测hmodel.model.embed_tokens.register_forward_hook(lambdam,i,o:seen.update(embedtuple(o.shape)))withtorch.no_grad():outmodel(**tok(hello,return_tensorspt))h.remove()print(embedding:,seen[embed])print(logits :,tuple(out.logits.shape))print(层数 :,len(model.model.layers))实测embedding (1,1,896)、logits (1,1,151936)、24 层。与 224 的数据流一致hello是 1 个 token → 查嵌入得 896 维 → 过 24 层形状始终(1,1,896)→ LM 头映射到词表 151936 维。所有层输入输出同形才能一路残差相加。形状里三个维度分别是batch同时处理几条、seq序列长度、hidden896。训练时batch与seq都大于 1、可并行推理逐字生成时seq每次只增 1。KV cache为什么生成要缓存自回归生成是逐字进行的生成第 n 个字时要把前面 n-1 个字的 K、V 重新算一遍——太浪费。KV cache 把算过的 K/V 存下来复用。用开关对比实测inputstok(用一句话介绍什么是 Transformer。,return_tensorspt)forcachein[True,False]:ttime.time()model.generate(**inputs,max_new_tokens64,do_sampleFalse,use_cachecache)print(fuse_cache{cache}:{time.time()-t:.2f}s)实测use_cacheTrue6.83suse_cacheFalse10.91s——关掉缓存慢了约60%而且生成长度越长差距越大无缓存时每步都要重算全部历史。这就是 228 里 GQA 省 KV cache 的意义缓存越小显存越省、长上下文越可行显存越大上下文越长但缓存减小后时间会增加。把账算具体每个 token 的 KV cache 2(K/V) × 24(层) × 2(KV 头) × 64(head_dim) × 2 字节(fp16)12 KB到 32768 上下文约384 MB。若换成 MHAMulti-Head Attention多头注意力14 个 KV 头同一上下文要 ~2.7 GB——GQA 省下的正是这一大块。为什么只缓存 K/V、不缓存 Q生成第 n 个字时历史位置的 Q 已经用完、不会再被查询而它们的 K/V 要被当前位置的 Q 反复查阅所以只需缓存 K/V。这也是KV cache名字的由来。KV cache 的取舍很直白开缓存省时间但占显存12KB/token关缓存省显存但慢。长上下文场景必须开缓存不然时间太慢、并靠 GQA 把缓存压小只有显存极度紧张时才考虑关。这是时间换空间在推理里的典型体现。生成是逐字循环一次前向只算一个位置要生成一句话得把上一次的输出接回输入、反复前向自回归for_inrange(max_new):logitsmodel(ids)[:,-1]# 只关心最后一个位置nxtsample(logits)# 采样得到下一个 tokenidstorch.cat([ids,nxt],dim1)# 拼回输入进入下一轮每生成一个字都要跑一次完整前向。没有 KV cache 时每次都要把前面所有字重算一遍——这正是前面 6.83s 与 10.91s 差距的来源。采样贪心与随机差在哪拿到 logits每个候选字的分数后怎么选下一个字两种常见策略实测gmodel.generate(**inp,max_new_tokens32,do_sampleFalse)smodel.generate(**inp,max_new_tokens32,do_sampleTrue,temperature1.2,top_k40)print(贪心:,tok.decode(g[0][n:],skip_special_tokensTrue))print(采样:,tok.decode(s[0][n:],skip_special_tokensTrue))实测同一句提示下的回答策略输出节选贪心do_sampleFalse“学习就像攀登一座高峰只有不断努力才能到达顶峰。”…采样temperature1.2, top_k40“成功往往来源于不断地尝试与学习而不是等待机遇…”贪心每步取概率最大的字结果稳定可复现采样按概率随机抽temperature调大更发散、top_k限制只从概率最高的 k 个里抽避免抽到离谱的字。产品里要稳定就用低温度/贪心要多样就开采样。再对比不同温度实测同一提示temperature输出风格0.3更收敛、句式平稳“……激励着无数人勇往直前……”1.5更发散、句式跳脱“……不断挑战……看到成功的光芒……”完整的采样链路是logits → 除以 temperature → top-k 过滤 → softmax → 按概率 multinomial 抽样。每一步都在改变候选字的概率分布最终抽出一个字追加到序列再进入下一轮。为什么需要采样logits 不是答案模型最后输出的是每个候选字的分数logits不是现成的答案。要变成文字必须选一个——这也解释了为什么同一个模型、同一句话输出可以每次不同随机性来自怎么选不来自模型本身。策略行为适用贪心每步取最大概率稳定、可复现采样按概率随机多样、有创意temperature调节分布陡峭度控制发散程度top-k截断到前 k 个防止离谱输出推理与训练的一个关键差别训练时一次前向处理整段(B, S)文本、并行算所有位置推理时逐字生成、每次只处理新位置——前者能并行、后者必须串行。KV cache 正是为缓解串行带来的重复计算这也是训练吃算力、推理吃显存带宽这一说法的来源。三个概念一句话对照概念作用实测一次前向输入 → logits(1,1,896) → (1,1,151936)24 层KV cache复用历史 K/V开 6.83s vs 关 10.91s采样从 logits 选下一个字贪心稳定、采样多样到这里模型怎么推理这条线也闭环了结构224~229→ 训练230~231→ 推理232。下一篇用下载好的基座版做对比看预训练与指令微调到底差在哪。一句话记忆一次前向把(1,1)的 token 变成(1,1,151936)的 logits24 层始终维持(1,1,896)同形KV cache 复用历史 K/V实测开 6.83s、关 10.91s慢约 60%越长越明显这正是 GQA 省缓存的动机采样用do_sample/temperature/top_k控制稳定 vs 多样——结构、训练、推理三条线到此闭环。OK本篇先到这里如有疑问欢迎评论区留言讨论祝各位功力大涨技术更上一层楼更多内容见下篇 blog【AI】【模型部署】基座模型研究基座与 Instruct 的差别 232、【AI】【模型部署】基座模型研究：一次前向、KV cache 与采样
【声明】本博客所有内容均为个人业余时间创作所述技术案例均来自公开开源项目如GithubApache基金会不涉及任何企业机密或未公开技术如有侵权请联系删除标题232、【AI】【模型部署】基座模型研究一次前向、KV cache 与采样背景上篇 blog【AI】【模型部署】从零手写基座(二)训练、loss 与生成跑通了预训练闭环x与y错开一位做 next-token 预测、交叉熵算误差、反向传播更新权重135 万参数的微型模型在 CPU 上训 1200 步3 分 8 秒loss 从 7.3883≈ln1432纯乱猜降到 0.7821能续写出格式正确的博客文本——小模型的会写是统计模仿而 Qwen2 用的是同一套训练目标。本篇回到推理侧拿真正的 Qwen2.5-0.5B把一次前向的形状、KV cacheKey-Value cache键值缓存 到底省什么和采样怎么选下一个字这三件事用实验看清模型部署训练解决模型怎么变聪明推理解决用起来多快、输出怎么选。前面几篇一直在说 KV cache、temperature、top-k本篇用真模型把它们一次测清楚。一次前向形状从哪来、到哪去给模型一个单词用 hook 抓词嵌入的输出实测hmodel.model.embed_tokens.register_forward_hook(lambdam,i,o:seen.update(embedtuple(o.shape)))withtorch.no_grad():outmodel(**tok(hello,return_tensorspt))h.remove()print(embedding:,seen[embed])print(logits :,tuple(out.logits.shape))print(层数 :,len(model.model.layers))实测embedding (1,1,896)、logits (1,1,151936)、24 层。与 224 的数据流一致hello是 1 个 token → 查嵌入得 896 维 → 过 24 层形状始终(1,1,896)→ LM 头映射到词表 151936 维。所有层输入输出同形才能一路残差相加。形状里三个维度分别是batch同时处理几条、seq序列长度、hidden896。训练时batch与seq都大于 1、可并行推理逐字生成时seq每次只增 1。KV cache为什么生成要缓存自回归生成是逐字进行的生成第 n 个字时要把前面 n-1 个字的 K、V 重新算一遍——太浪费。KV cache 把算过的 K/V 存下来复用。用开关对比实测inputstok(用一句话介绍什么是 Transformer。,return_tensorspt)forcachein[True,False]:ttime.time()model.generate(**inputs,max_new_tokens64,do_sampleFalse,use_cachecache)print(fuse_cache{cache}:{time.time()-t:.2f}s)实测use_cacheTrue6.83suse_cacheFalse10.91s——关掉缓存慢了约60%而且生成长度越长差距越大无缓存时每步都要重算全部历史。这就是 228 里 GQA 省 KV cache 的意义缓存越小显存越省、长上下文越可行显存越大上下文越长但缓存减小后时间会增加。把账算具体每个 token 的 KV cache 2(K/V) × 24(层) × 2(KV 头) × 64(head_dim) × 2 字节(fp16)12 KB到 32768 上下文约384 MB。若换成 MHAMulti-Head Attention多头注意力14 个 KV 头同一上下文要 ~2.7 GB——GQA 省下的正是这一大块。为什么只缓存 K/V、不缓存 Q生成第 n 个字时历史位置的 Q 已经用完、不会再被查询而它们的 K/V 要被当前位置的 Q 反复查阅所以只需缓存 K/V。这也是KV cache名字的由来。KV cache 的取舍很直白开缓存省时间但占显存12KB/token关缓存省显存但慢。长上下文场景必须开缓存不然时间太慢、并靠 GQA 把缓存压小只有显存极度紧张时才考虑关。这是时间换空间在推理里的典型体现。生成是逐字循环一次前向只算一个位置要生成一句话得把上一次的输出接回输入、反复前向自回归for_inrange(max_new):logitsmodel(ids)[:,-1]# 只关心最后一个位置nxtsample(logits)# 采样得到下一个 tokenidstorch.cat([ids,nxt],dim1)# 拼回输入进入下一轮每生成一个字都要跑一次完整前向。没有 KV cache 时每次都要把前面所有字重算一遍——这正是前面 6.83s 与 10.91s 差距的来源。采样贪心与随机差在哪拿到 logits每个候选字的分数后怎么选下一个字两种常见策略实测gmodel.generate(**inp,max_new_tokens32,do_sampleFalse)smodel.generate(**inp,max_new_tokens32,do_sampleTrue,temperature1.2,top_k40)print(贪心:,tok.decode(g[0][n:],skip_special_tokensTrue))print(采样:,tok.decode(s[0][n:],skip_special_tokensTrue))实测同一句提示下的回答策略输出节选贪心do_sampleFalse“学习就像攀登一座高峰只有不断努力才能到达顶峰。”…采样temperature1.2, top_k40“成功往往来源于不断地尝试与学习而不是等待机遇…”贪心每步取概率最大的字结果稳定可复现采样按概率随机抽temperature调大更发散、top_k限制只从概率最高的 k 个里抽避免抽到离谱的字。产品里要稳定就用低温度/贪心要多样就开采样。再对比不同温度实测同一提示temperature输出风格0.3更收敛、句式平稳“……激励着无数人勇往直前……”1.5更发散、句式跳脱“……不断挑战……看到成功的光芒……”完整的采样链路是logits → 除以 temperature → top-k 过滤 → softmax → 按概率 multinomial 抽样。每一步都在改变候选字的概率分布最终抽出一个字追加到序列再进入下一轮。为什么需要采样logits 不是答案模型最后输出的是每个候选字的分数logits不是现成的答案。要变成文字必须选一个——这也解释了为什么同一个模型、同一句话输出可以每次不同随机性来自怎么选不来自模型本身。策略行为适用贪心每步取最大概率稳定、可复现采样按概率随机多样、有创意temperature调节分布陡峭度控制发散程度top-k截断到前 k 个防止离谱输出推理与训练的一个关键差别训练时一次前向处理整段(B, S)文本、并行算所有位置推理时逐字生成、每次只处理新位置——前者能并行、后者必须串行。KV cache 正是为缓解串行带来的重复计算这也是训练吃算力、推理吃显存带宽这一说法的来源。三个概念一句话对照概念作用实测一次前向输入 → logits(1,1,896) → (1,1,151936)24 层KV cache复用历史 K/V开 6.83s vs 关 10.91s采样从 logits 选下一个字贪心稳定、采样多样到这里模型怎么推理这条线也闭环了结构224~229→ 训练230~231→ 推理232。下一篇用下载好的基座版做对比看预训练与指令微调到底差在哪。一句话记忆一次前向把(1,1)的 token 变成(1,1,151936)的 logits24 层始终维持(1,1,896)同形KV cache 复用历史 K/V实测开 6.83s、关 10.91s慢约 60%越长越明显这正是 GQA 省缓存的动机采样用do_sample/temperature/top_k控制稳定 vs 多样——结构、训练、推理三条线到此闭环。OK本篇先到这里如有疑问欢迎评论区留言讨论祝各位功力大涨技术更上一层楼更多内容见下篇 blog【AI】【模型部署】基座模型研究基座与 Instruct 的差别