向量计算-第13章-实战4-Transformer内核-Self-Attention与KV Cache

清华大学出版社《向量计算》配套课件·课堂讲义第13章实战4：Transformer内核——Self-Attention与KVCache换个角度学AI，打通任督二脉刘鹏著·清华大学出版社用向量计算视角解构Self-Attention，掌握KVCache推理加速《向量计算》配套课件第13章实战4：Transformer内核——Self-Attention与KVCache第13章实战4：Transformer内核——Self-Attention与KVCache2/40教学目标1理解Self-Attention的动机：并行计算、长距离依赖、全局信息2掌握缩放点积注意力的五步推导与完整公式3理解正弦/余弦、可学习、相对偏置与RoPE位置编码的取舍4掌握KVCache的Prefill/Decode流程与显存估算公式5了解FlashAttention、PagedAttention、GQA等推理优化技术重点与难点1重点：Attention(Q,K,V)=softmax(QKᵀ/√dₖ)·V的维度推演与物理含义2难点：RoPE为何能通过旋转让点积只依赖相对位置3难点：KVCache的显存—计算权衡及其工程优化方案《向量计算》配套课件第13章实战4：Transformer内核——Self-Attention与KVCache第13章实战4：Transformer内核——Self-Attention与KVCache3/40本章内容导航13.1Self-Attention的QKV机制13.2位置编码13.3KVCache：推理加速的基石13.4Transformer推理优化全景13.5本章小结《向量计算》配套课件·第13章实战4：Transformer内核——Self-Attention与KVCache13.1Self-Attention的QKV机制为什么需要·QKV直观理解·数学原理·详细计算示例·多头注意力·可视化与比较《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache5/40为什么需要Self-Attention？●“苹果发布了新手机”——“苹果”是水果还是公司？答案依赖上下文●传统RNN/LSTM逐个处理单词，存在三大局限–顺序处理：必须先处理第1个词，无法并行计算–长距离依赖：句首信息传到句尾已经衰减–深层网络训练时梯度容易消失或爆炸●Self-Attention：每个单词直接与句中所有单词“对话”，不受距离限制《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache6/40QKV直观理解：图书馆检索系统●Query：你提出的问题——“找关于人工智能的书”●Key：每本书的索引标签——“人工智能”“计算机”“机器学习”●Value：书的具体内容——书名、作者、章节内容●检索：查询与所有索引匹配，匹配度高的书返回其内容●Self-Attention中每个单词都生成自己的Q、K、V向量–例“我爱程序设计”：“爱”的Q寻找主宾语，与“我”“程序设计”的K匹配《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache7/40步骤1：生成Q、K、V向量Q=X·W_QK=X·W_KV=X·W_VX：输入词嵌入向量（可以是原始嵌入，也可以经过位置编码）W_Q、W_K、W_V：三个可学习的权重矩阵（线性变换）Q、K、V的维度可以相同，也可以不同●代码中用三个nn.Linear(bias=False)实现，且要求head_dim·heads=embed_size●直观：Q表示“我想查找什么”，K表示“我能提供什么”，V是“我包含的内容”《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache8/40步骤2+3：注意力得分与缩放Attention_Score=Q·KᵀScaled_Score=Q·Kᵀ/√dₖQ、K形状均为[seq_len,dₖ]；Kᵀ为转置[dₖ,seq_len]；Score为[seq_len,seq_len]Score第i行第j列=第i个与第j个Token的相似度——全对全（All-to-All）计算dₖ：Key向量维度（如64或512）；√dₖ为缩放因子●为何用点积：方向一致时点积最大，且可用矩阵乘法并行计算●除以√dₖ：防止数值爆炸，避免Softmax进入饱和区导致梯度消失●统计意义：把点积结果的方差控制在1附近，保证数据分布稳定《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache9/40步骤4：Softmax归一化Attention_Weight=softmax(Scaled_Score)softmax(xᵢ)=exp(xᵢ)/Σⱼexp(xⱼ)xᵢ：输入向量（某一行缩放分数）中的每一个元素对每一行独立归一化：每个Query分配给所有Key的权重之和严格等于1●非负性：所有权重在(0,1)之间，表示“关注度”的百分比●归一性：每个Query的注意力总预算固定为100%●马太效应：exp放大高分项、抑制低分项，让模型更坚决地关注重要Token《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache10/40完整公式：缩放点积注意力Attention(Q,K,V)=softmax(Q·Kᵀ/√dₖ)·V权重矩阵[seq_len,seq_len]×值矩阵V[seq_len,dᵥ]→输出[seq_len,dᵥ]输出的每一行都是所有Value向量的线性组合——上下文信息在此融合●①Q·Kᵀ计算相似度②÷√dₖ数值缩放，保持梯度稳定●③softmax归一化为概率权重④·V按权重提取信息●例：TokenA对B的权重为0.9，则A的输出向量含90%的B信息《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache11/40详细计算示例①：生成Q（句子“猫追逐激光点”）Q("猫")=X[0]·W_Q=[0.8,0.2,0.5,0.1]·[[0.1,0.5,0.3,0.7],[0.9,0.2,0.4,0.1],[0.5,0.8,0.1,0.3],[0.2,0.6,0.7,0.4]]=[1.02,0.85,0.73,0.78]注：设词嵌入与Q、K、V维度均为4；X[0]=[0.8,0.2,0.5,0.1]是“猫”的嵌入，W_Q为可学习权重矩阵，K(“猫”)、V(“猫”)同理可得《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache12/40详细计算示例②：注意力得分（“追逐”的Query）Q("追逐")=[0.56,0.92,0.48,0.65]Attention_Score=Q("追逐")·K^T#与"猫"的Key的点积score1=[0.56,0.92,0.48,0.65]·[0.61,0.37,0.63,0.88]^T=0.56×0.61+0.92×0.37+0.48×0.63+0.65×0.88=0.34+0.34+0.30+0.57=1.55#与"追逐"的Key的点积score2=[0.56,0.92,0.48,0.65]·[0.42,0.75,0.28,0.53]^T=0.56×0.42+0.92×0.75+0.48×0.28+0.65×0.53=0.24+0.69+0.13+0.34注：四个点积全算完：与“激光”“点”分别为1.54、1.14，故Attention_Score(“追逐”)=[1.55,1.40,1.54,1.14]；d_k=4，缩放因子√4=2《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache13/40详细计算示例③：缩放后的Softmax归一化#以"追逐"为例：scores=[0.78,0.70,0.77,0.57]exp_scores=[e^0.78,e^0.70,e^0.77,e^0.57]=[2.18,2.01,2.16,1.77]sum_exp=2.18+2.01+2.16+1.77=8.12weights=[2.18/8.12,2.01/8.12,2.16/8.12,1.77/8.12]=[0.27,0.25,0.27,0.22]#归一化后的注意力权重矩阵：Attention_Weight=[[0.28,0.25,0.27,0.20],[0.27,0.25,0.27,0.22],[0.26,0.24,0.28,0.22],[0.25,0.23,0.26,0.26]]注：步骤6加权求和：Output(“追逐”)=0.27·V(“猫”)+0.25·V(“追逐”)+0.27·V(“激光”)+0.22·V(“点”)=[0.53,0.64,0.57,0.50]——“猫”与“激光”（主语、宾语）权重最高，输出向量融合了整句信息《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache14/40多头注意力（Multi-HeadAttention）headᵢ=Attention(Q·Wᵢ^Q,K·Wᵢ^K,V·Wᵢ^V)MultiHead=Concat(head₁,…,headₕ)·WᴼWᵢ^Q、Wᵢ^K∈R^(d_model×dₖ)，Wᵢ^V∈R^(d_model×dᵥ)：第i个头的独立投影Wᴼ∈R^(h·dᵥ×d_model)：输出变换矩阵，融合所有头的信息典型维度：d_model=512，h=8，每头dₖ=dᵥ=512/8=64●像同时雇多个“侦探”从不同角度调查案件，最后汇总发现●形状流：[seq,d_model]→分头[h,seq,dₖ]→各头独立算注意力→拼接→Wᴼ●优势：多视角（语法/语义/指代）、高表达能力、各头并行计算《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache15/40多头注意力：einsum核心计算#queries:(N,query_len,heads,head_dim)#keys:(N,key_len,heads,head_dim)#energy:(N,heads,query_len,key_len)energy=torch.einsum("nqhd,nkhd->nhqk",[queries,keys])#缩放energy=energy/(self.embed_size**(1/2))#应用mask（如果需要）ifmaskisnotNone:energy=energy.masked_fill(mask==0,float("-1e20"))#Softmaxattention=torch.softmax(energy,dim=3)#加权求和#attention:(N,heads,query_len,key_len)#values:(N,value_len,heads,head_dim)#out:(N,query_len,heads,head_dim)out=torch.einsum("nhql,nlhd->nqhd",[attention,values])#合并多头:(N,query_len,heads*head_dim)out=out.reshape(N,query_len,self.heads*self.head_dim)#线性变换out=self.fc_out(out)注：einsum一步完成Q·Kᵀ批量矩阵乘；mask位置填-1e20，经softmax后权重≈0，防止模型看到后续token《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache16/40Self-Attention的可视化与解释●句“那只追逐激光点的猫很可爱”：“猫”最关注“激光”(0.25)、“追逐”(0.20)–权重矩阵的每一行=当前token对所有token的关注度分布●底层：注意力集中在邻近词（第1层“猫”关注“追逐”“激光”）●中层：关注更长距离的依赖；高层：捕捉全局语义关系●BERT中发现句法头、指代消解头、位置头，以及权重均匀的“无关头”《向量计算》配套课件13.1Self-Attention的QKV机制第13章实战4：Transformer内核——Self-Attention与KVCache17/40表13-1Self-Attention与RNN、CNN的对比特性Self-AttentionRNN/LSTMCNN计算复杂度O(n²)O(n)O(n)并行计算支持不支持支持长距离依赖O(1)O(n)O(logn)感受野全局顺序递增局部内存占用高低中可解释性强弱中适用场景长序列建模时序数据局部特征注：劣势O(n²)复杂度与n×n矩阵显存的优化方向：稀疏注意力(Longformer)、线性注意力(Performer)、局部滑动窗口《向量计算》配套课件·第13章实战4：Transformer内核——Self-Attention与KVCache13.2位置编码置换不变性问题·正弦/余弦编码·可学习编码·相对位置偏置·RoPE旋转位置编码《向量计算》配套课件13.2位置编码第13章实战4：Transformer内核——Self-Attention与KVCache19/40为什么需要位置编码：置换不变性Zᵢ=Xᵢ+PᵢXᵢ：第i个词的词嵌入向量；Pᵢ：第i个位置的位向量Zᵢ：融合词义与位置信息的最终输入向量（也可采用拼接方式）●Self-Attention把输入当作“集合”而非“序列”，交换顺序结果不变●“我爱吃苹果”与“苹果爱吃我”词集相同、含义迥异●类比电影院座位：同一个词在不同位置扮演不同成分（主语/宾语）●即使词相同，位置不同→最终输入表示不同《向量计算》配套课件13.2位置编码第13章实战4：Transformer内核——Self-Attention与KVCache20/40绝对位置编码：正弦/余弦方案PE(pos,2i)=sin(pos/10000^(2i/d_model))PE(pos,2i+1)=cos(pos/10000^(2i/d_model))pos：词在序列中的位置索引(0,1,2,…)；i：维度索引(0,…,d_model/2−1)偶数维度用sin、奇数维度用cos；d_model为隐藏层维度（如512或768）10000为基数，控制频率衰减速度；出自《AttentionIsAllYouNeed》●多尺度如时钟：低维频率高变化快（秒针），高维变化慢（时针）●对任意固定偏移k，PE(pos+k)可表示为PE(pos)的线性函数→易学相对位置●外推性：确定性函数，理论可处理比训练更长的序列●例d_model=4：位置1≈[0.841,0.540,0.010,0.999]《向量计算》配套课件13.2位置编码第13章实战4：Transformer内核——Self-Attention与KVCache21/40正弦位置编码：PyTorch实现#位置索引:[0,1,2,...,max_len-1]position=torch.arange(0,max_len,dtype=torch.float).unsqueeze(1)#频率项:1/[10000^(2i/d_model)]div_term=torch.exp(torch.arange(0,d_model,2).float()*(-math.log(10000.0)/d_model))#计算位置编码#偶数维度使用sinpe[:,0::2]=torch.sin(position*div_term)#奇数维度使用cospe[:,1::2]=torch.cos(position*div_term)#增加batch维度:[1,max_len,d_model]pe=pe.unsqueeze(0)#注册为buffer（不是模型参数，但会保存到state_dict）self.register_buffer('pe',pe)defforward(self,x):#x:[batch_size,seq_len,d_model]#截取与输入序列长度匹配的位置编码returnx+self.pe[:,:x.size(1)]注：pe注册为buffer（不参与梯度但随state_dict保存），forward中与词嵌入相加。对比可学习方案nn.Embedding(max_len,d_model)：简单灵活，但外推性差、参数增加《向量计算》配套课件13.2位置编码第13章实战4：Transformer内核——Self-Attention与KVCache22/40相对位置编码：距离依赖的偏置项Attention_Score′(i,j)=Qᵢ·Kⱼᵀ+bᵢ₋ⱼQᵢ·Kⱼᵀ：原始注意力得分；bₖ是只依赖相对距离k=i−j的可学习偏置k取值如−5,−4,…,0,…,4,5；实现为一张[2·max_relative_position−1,num_heads]偏置表●优势1：直接编码相对位置，更符合语言建模的直觉●优势2：外推性好，可处理训练时未见过的相对距离●劣势：额外参数开销；须修改注意力计算的核心逻辑●代表方案：Transformer-XL《向量计算》配套课件13.2位置编码第13章实战4：Transformer内核——Self-Attention与KVCache23/40旋转位置编码RoPE：核心思想z′=z·e^(iθ)θᵢ=10000^(−2(i−1)/d)把Q、K分成d/2个二维向量；位置m处第i个二维向量旋转m·θᵢ角度e^(iθ)=cosθ+i·sinθ，等价于旋转矩阵[[cosθ,−sinθ],[sinθ,cosθ]]i：二维向量索引；d：向量总维度；LLaMA、GPT-NeoX、PaLM均采用●关键洞察：位置m的Q与位置n的K做点积，结果只依赖相对位置m−n●证明要点：旋转矩阵正交且可交换，R(m)ᵀ·R(n)化为只含n−m的旋转●优势：相对编码、外推性好、无参数、不改注意力核心逻辑●劣势：需复数/旋转运算；通常要求head_dim为偶数《向量计算》配套课件13.2位置编码第13章实战4：Transformer内核——Self-Attention与KVCache24/40RoPE实现：复数乘法即旋转classRotaryPositionEmbedding(nn.Module):def__init__(self,d_model,max_len=512):super(RotaryPositionEmbedding,self).__init__()#计算旋转频率self.d_model=d_modelinv_freq=1.0/(10000**(torch.arange(0,d_model,2).float()/d_model))#预计算位置编码self.register_buffer('inv_freq',inv_freq)self.max_len=max_lendefforward(self,x,position_ids):#x:[batch_size,seq_len,num_heads,head_dim]batch_size,seq_len,num_heads,head_dim=x.size()#计算旋转角度:[seq_len,head_dim//2]position_ids=position_ids.unsqueeze(-1).float()freqs=torch.matmul(position_ids,self.inv_freq.unsqueeze(0))freqs=torch.polar(torch.ones_like(freqs),freqs)#复数表示#将x分成实部和虚部x_complex=torch.view_as_complex(x.float())注：torch.polar构造e^(iθ)；续：x_rotated=x_complex*freqs完成旋转，再view_as_real转回实数《向量计算》配套课件13.2位置编码第13章实战4：Transformer内核——Self-Attention与KVCache25/40表13-3不同位置编码方案的对比方案类型参数外推性相对位置实现复杂度应用模型正弦/余弦绝对无好需学习低Transformer原始版可学习绝对有差需学习低BERT,GPT早期相对位置偏置相对有好直接中Transformer-XLRoPE相对无好直接高LLaMA,PaLM注：选择：长度固定→可学习；需外推长序列→RoPE/正弦；现代大模型中RoPE是事实上的标准《向量计算》配套课件·第13章实战4：Transformer内核——Self-Attention与KVCache13.3KVCache：推理加速的基石自回归生成的挑战·Prefill/Decode流程·显存开销·更新策略·优化技术《向量计算》配套课件13.3KVCache：推理加速的基石第13章实战4：Transformer内核——Self-Attention与KVCache27/40自回归生成的效率瓶颈●自回归：根据已生成token预测下一个，循环直到<EOS>或最大长度●生成“我喜欢吃苹果”：第3步输入“<BOS>我喜欢”→预测“吃”●挑战1：每步都要处理整个序列，新token与前面所有token算注意力，计算量O(n²)●挑战2：已处理token的K、V向量每一步都被重复计算●KVCache：缓存历史token的Key和Value向量，避免重复计算《向量计算》配套课件13.3KVCache：推理加速的基石第13章实战4：Transformer内核——Self-Attention与KVCache28/40KVCache的工作原理●Prefill：接收完整prompt，一次算出所有token的Q、K、V–把全部K、V存入缓存→生成第一个新token（如“苹果”）●Decode：每步只计算新token的Q、K、V–新Q与缓存K算注意力、加权缓存V，再把新K、V追加进缓存●复杂度：无缓存每步O(n²)→有缓存每步O(n)（总计算量≈O(n)）《向量计算》配套课件13.3KVCache：推理加速的基石第13章实战4：Transformer内核——Self-Attention与KVCache29/40KVCache推理伪代码#Prefill:处理初始promptoutput,kv_cache=model.forward_prefill(tokens)next_token=sample(output)tokens.append(next_token)#Decode:逐个生成tokenforiinrange(max_len-1):#只计算新token的Q、K、Voutput,kv_cache=model.forward_decode(next_token,kv_cache)#计算复杂度:O(n)next_token=sample(output)tokens.append(next_token)注：无缓存版本每步model(tokens)重算全部token的Q、K、V（复杂度O(n²)）；有缓存时forward_decode只处理新token《向量计算》配套课件13.3KVCache：推理加速的基石第13章实战4：Transformer内核——Self-Attention与KVCache30/40表13-4生成1000个token的计算复杂度对比方法第1步第2步…第1000步总计算量无KVCache1²2²…1000²≈334,000,000有KVCache12…1000≈500,000注：使用KVCache能将计算量减少约667倍《向量计算》配套课件13.3KVCache：推理加速的基石第13章实战4：Transformer内核——Self-Attention与KVCache31/40KVCache显存估算（式13-14）Memory=num_layers×2×batch_size×num_heads×seq_len×head_dim×sizeof(dtype)×2：Key与Value两部分；sizeof(dtype)：FP16为2字节，FP32为4字节seq_len随生成不断追加→显存占用随序列长度线性增长●例：32层×32头×128维×FP16，100个token时单层≈800KB，总量≈25MB；4096token≈1GB●LLaMA-7B（batch=1,FP16）：seq1024→约512MB，2048→1GB，4096→2GB，16K→8GB●A100实测（表13-6）：加速比从2.3×（512）升至10×（4096），序列越长收益越大《向量计算》配套课件13.3KVCache：推理加速的基石第13章实战4：Transformer内核——Self-Attention与KVCache32/40缓存更新策略与显存优化技术●追加（Append）：新K/V拼到缓存末尾，保留全部历史——显存线性增长●滑动窗口（SlidingWindow）：只留最近N个token——有上限但丢远期信息●量化：FP16→INT8省50%，→INT4省75%，需额外量化/反量化计算●PagedAttention：按页表管理非连续内存块，几乎消除碎片、支持多序列共享KV●GQA：32个Query头共享4组K/V头，缓存压缩至4/32=1/8，性能损失通常<1%●动态剪枝：按注意力权重/熵/时间衰减只保留top-k（如256）重要token的KV《向量计算》配套课件·第13章实战4：Transformer内核——Self-Attention与KVCache13.4Transformer推理优化全景Prefill与Decode·FlashAttention·算子融合与编译·内存优化·推理框架与调优《向量计算》配套课件13.4Transformer推理优化全景第13章实战4：Transformer内核——Self-Attention与KVCache34/40表13-7Prefill与Decode阶段的对比特性Prefill阶段Decode阶段计算复杂度O(n²)O(n)内存访问HBM读写频繁主要是KVCache读取计算模式并行串行主要瓶颈计算密集型内存带宽密集型优化重点计算效率内存管理注：时间占比：Prefill约占20%–50%（取决于prompt长度），Decode约占50%–80%（取决于生成长度）《向量计算》配套课件13.4Transformer推理优化全景第13章实战4：Transformer内核——Self-Attention与KVCache35/40计算与内存优化技术●FlashAttention：128×128分块入SRAM+反向重计算，避免n=4096时约67MB的注意力矩阵反复读写HBM–FlashAttention-2：速度提升2–4倍、显存降50%–70%；需手写CUDA或用Triton●算子融合：LayerNorm4个算子合为单次内核启动，典型加速10%–30%●编译优化：XLA、TVM、Triton——常量折叠、公共子表达式/死代码消除、内存复用，约1.5–2倍●梯度检查点：前向只存关键节点，降显存30%–50%，多花10%–20%计算●权重量化：INT8体积缩约75%、推理提速2–4倍；PTQ部署便捷，QAT训练时模拟量化更保精度《向量计算》配套课件13.4Transformer推理优化全景第13章实战4：Transformer内核——Self-Attention与KVCache36/40vLLM：高吞吐推理框架示例fromvllmimportLLM,SamplingParams#初始化模型llm=LLM(model="meta-llama/Llama-2-7b-chat-hf")#设置采样参数sampling_params=SamplingParams(temperature=0.8,top_p=0.95,max_tokens=100)#生成文本prompts=["请介绍一下人工智能。","什么是机器学习？"]outputs=llm.generate(prompts,sampling_params)