【自用】AI infra相关

一、推理两阶段1、预填充prefill对完整的输入prompt进行一次性并行计算为每一个token生成并缓存keyvalue向量。这一步只需要进行一次缓存内容将用于后续的生成步骤。token化输入批大小为batch_size的若干prompt每个prompt被处理为长度为seq_len的token序列。嵌入层这些token序列会被映射为隐藏向量X形状为[batch_size, seq_len, hidden_dim]。线性映射通过权重矩阵W_Q, W_K, W_V[hidden_dim, num_heads * head_dim])分别计算得到Q,K,V并重塑为[batch_size, num_heads, seq_len, head_dim]。kv cache计算得到的k和v矩阵会被缓存下来供后续解码阶段使用。注意力计算生成第一个token在工程实践中为减少一次显存读写和前向传播开销通常直接利用prompt最后一个位置的hidden state来生成第一个token。prefill阶段是高度并行的能够充分利用gpu因此属于compute-bound计算受限瓶颈在于算力而不是内存。2、解码decode模型以自回归方式逐步生成新token。每一轮仅需计算当前要生成的token并与此前缓存的kv包括用户查询和已经生成的token进行注意力计算从而避免与整个历史序列重复计算显著降低计算量。输入准备使用prefill阶段已经计算并存储的kv cache以及当前轮刚生成的token第一次decode时是prefill阶段生成的第一个token。嵌入层在第t步将新生成的token转哈为嵌入向量计算得到QtKtVt并将KtVt更新到kv cache中。计算注意力基于Qt和现有kv cache进行注意力计算计算复杂度为O(seq_len)相比prefill阶段大幅降低。生成下一个token将最后一层的注意力输出经过线性层和softmax得到下一个token的概率分布并根据解码策略如greedy search、beam search选出下一个token作为下一次decode的输入。decode阶段是典型的通信密集型任务主要受限于显存带宽因为生成一条回复通常需要多次decode每次都需要访问和更新kv cache。随着生成内容长度的增加kv cache也会越来越大对gpu内存带宽的压力也会随之增大。对比项prefilldecode主要功能一次性并行处理完整输入序列计算并缓存kv用于后续生成首个及所有后续token每次只生成一个新token利用缓存的kv cache生成新输出并持续更新缓存计算类型全序列自注意力计算高度并行、计算密集型增量式注意力计算串行生成、通信密集型GPU利用率充分利用算力资源能高度并行效率高并行度低通常是小批量或单token操作GPU利用率较低计算复杂度O(seq_len²)O(seq_len)kv cache管理不依赖缓存直接基于原始输入并行计算构建并缓存所有prompt的kv cache每步都查kv cache利用其计算注意力并将新token的k_tv_t持续追加到缓存中kv cache用空间换时间通过使用kv cache避免了对之前token的重复计算加速模型推理速度。kv cache的形状[batch_size, num_heads, seq_len, head_dim]cache显存2*层数*批大小*序列长*隐藏层维度*精度字节第一个2key和value两个缓存张量序列长度指需要缓存的长度等于prompt的长度加上已经生成的token数量。这种与序列长度成正比的显存增长是制约大模型走向更长上下文的核心瓶颈之一。传统推理系统通常采用预分配策略请求开始时就按最大可能生成的长度如2048 token为每个请求分配一整块连续内存空间。这种方式带来显著的内存浪费产生许多碎片。在parallel sampling等复杂推理模式中一个请求可能包含多个生成分支如多个候选答案这些分支共享相同的prompt因此理论上可以共享其kv cache。但传统系统每个分支的kv cache存放在独立的连续内存中物理上无法共享只能进行冗余复制进而显著增加了内存开销。因此催生了PagedAttentionkv cache量化滑动窗口注意力等一系列关键的推理优化技术。需要注意模型有n层例如7b的llama有32层所以kv cache的尺寸第一位就是第几层然后才是后面的[batch_size, num_heads, seq_len, head_dim]模型在prefill阶段生成整个prompt的k和v到了decode的时候对于每个层先加载第i个layer的k和v然后拼接起来随着decode越来越多kv cache只是在seq_len维度上增大了之前的内容都没有变。layer级别的kv cache transfer是比较好的方式可以避免一次性传输整个kv cache因为相比gpu的高速传输cpu的PCIE带宽会变成传输瓶颈。PD分离二、vLLMvLLM是一个用于LLM推理的高性能库。解决了什么问题答memory wastevllm解决了下列三种memory waste预留token传统推理会给每个request预留kv cache需要使用的memory因为llm本身无法确定这个请求会产生多少token所以会按固定长度预留。内部碎片系统通常也会以『块』为单位进行分配而不是单个token一个sequence长度300系统会分配2个块共512个token的容量第二块中就有212个token的位置wasted。外部碎片每个请求使用的块大小不一样memory在切割时会产生一些外部碎片这些碎片可能不足以再分给一个新的请求。1、解决方案vllm基于PagedAttention算法通过『分页』机制替代传统的『预先分配连续最大内存』策略来减少内存碎片。解决预留token不再为sequence预分配整个生命周期的内存。系统维护一个全局的物理块池。sequence只在需要时才申请新的物理块。解决内部碎片将kv cache存储在固定大小的块中一个sequence的kv cache可以分布在多个非连续的物理块中。最后一个块的未用空间就是内部碎片但由于块大小固定且相对较小如16个token这种浪费被控制在了很低水平。解决外部碎片由于sequence可以使用非连续的物理块系统只需维护一个全局的空闲物理块列表任何新请求都可以从列表中获得任意空闲块而无需寻找连续空间。这从根本上消除了外部碎片。注意vllm是把waste控制在很低的水平不是100% none waste。2、核心组件1Paged Attention在os中虚拟内存将程序的地址空间划分为固定大小的页这些页可以映射到物理内存中任意位置实现非连续分配从而有效解决了内存碎片和共享问题。pagedAttention被这一思路引入kv cache的管理中并带来了3项关键改进kv cache被切分为固定大小的block将每个序列的kv cache切分为固定大小的block(默认是16个token每个block存储若干个token的key和value向量。这种设计统一了内存分配力度使系统能够以更标准的方式管理kv cache的分配和回收从而提升内存复用效率并有效减少内存碎片。block可以存放在非连续的物理内存中与传统的attention不同pagedattention不在要求这些kv向量在内存中连续排列而是通过逻辑block和物理block的映射实现非连续存储。映射关系由block table维护它类似于操作系统中的页表用于记录每个逻辑block对应的物理内存位置确保模型在推理过程中可以正确读取所需的kv cache。支持灵活的分配与释放以及共享机制多个序列可以共享相同的kv cache block。如果某个共享block还没有被写满而不同序列需要向其中写入不同新token的kv则不能直接修改同一个block。此时系统通过copy-on-write复制该block为发生分叉的序列创建新的物理block从而既共享公共前缀又保证各自新增kv的隔离。每个逻辑块仅在前一个块被填满后才会分配新的物理块从而最大程度地减少内存浪费。在计算时我们操作的是逻辑块也就是说这些token在形式上是连续排列的与此同时vllm会通过block table映射关系在后台将逻辑块映射到实际的物理块从而完成数据的读取与计算。通过这种方式每个请求仿佛都能在一个连续且充足的内存空间中运行尽管这些数据在物理内存中实际上是非连续存储的。虚拟内存2continuous catchingcontinuous catching可以实现数倍乃至数十倍的系统吞吐提升被各大框架广泛使用。它的提出主要解决三类问题对early-finished requests处理一个batch中会有不同的请求每个请求所生成的文本长度不一致这就导致了在一个batch中生成短文本的请求得等生成长文本的请求结束后整个才结束造成了时间浪费。所以需要一个将已生成结束的请求从batch中移除并提前返回结果的机制一次来降低服务响应时间。对late-joining requests的处理同第一点有提前走的就有进来的就需要一个将新请求插入推理batch的机制。多个请求合并到一次attention计算每个请求对应的QKV Tensor的length维度各不相同在批量计算attention时需要处理此问题。3、框架关键模块为了解决通用深度学习框架中存在的不足vllm设计了几个关键模块①调度器scheduler用于解决多请求之间的调度协同问题②显存管理kv cache manager为请求分配kv cache内存资源③执行器model runner完成模型的计算上述3个模块放在引擎核engine core中。有了关键模块后再采用API服务的方式得到如图所示的改进方案4、高级特性主要聚焦于分块预填充chunked prefill前缀缓存prefix caching上面两个特性聚焦于prefill阶段会与continuous batching和kv cache深度结合。下面的特性聚焦于decoding阶段引导解码guided decoding推理解码speculative decodingPD分离prefill/decoding分离1chuncked prefill分块预填充是一种处理长提示词的技术主要在prefill阶段将一个长请求进行分块处理可以让其他prefill请求有机会插入执行。因为预填充的本质是对prompt的所有token执行一次完整前向传播属于计算密集型任务计算量随prompt长度线性增长。当遇到超长prompt如数百甚至数千token时单次预填充会占用GPU大量计算资源且执行时间极长在此期间引擎无法处理其他请求尤其是短prompt请求造成『资源闲置』与『请求排队拥堵』的双重问题短请求的等待时间latency会极速增长破坏服务的响应性如对话场景中用户等待多久将prefill拆分成多个子块并与其他请求的decode一起执行这种方式能够有效降低资源空跑提升整体的资源利用率。特性的实现在推理框架中主要改动点是调度器的逻辑保证多次分块的prefill请求能够衔接。调度器的输出格式{请求序号tokens位置}通过控制每个请求本轮需要计算的tokens数量实现任意chunk大小的组合下发。2prefix caching前缀缓存是一种避免重复计算多请求共享前缀token的技术。比如多个prompt在开头包含相同长前缀时只需在prefill阶段计算一次这个前缀的kv cache后续就无需重复执行前缀的prefill步骤。这里的『长前缀』定义为长度超过一个kv缓存块默认16个token的前缀。若前缀长度未对齐缓存块边界即long_prefix_len % block_size ! 0则未对齐的部分余数token仍需重新计算因为无法缓存不完整的kv块。在实际场景中多个请求共享长前缀的情况极为常见对话场景所有请求均包含相同的『系统提示』如『你是一个ai助手需遵循安全准则回答』批量处理场景多个任务的输入包含相同指令前缀如『总结以下文本』所以prefix的缓存极为重要。vllm中前缀缓存的实现依赖『哈希匹配』与『kv块关联』两大核心逻辑具体分为以下步骤步骤一第一次generate无缓存计算并存储前缀首先通过hash_request_token()计算前缀的哈希值。按kv块的大小将long_prefixprompt[0]的token id拆分为多个完整块。对每个块组合『前一个块的哈希值、当前块的token id、可选元数据』生成唯一的hash将每个块的hash值和token id封装成BlockHash对象存入哈希列表。通过find_longest_cache_hit()检查缓存命中。查询已经缓存哈希表判断是否有匹配的。因为是第一次生成的结果所以返回空如下图通过之前提到的allocate_slots()为当前请求分块kv cache块并通过coordinator.cache_blocks()将BlockHash对象与分配的KV块ID绑定记录到map映射中执行prefill步骤二第二次generate有缓存复用前缀的kv块对long_prefixprompt[1]执行相同的Token拆分和块哈希计算得到与第一次生成相同的『前缀块哈希列表』因long_prefix相同这时调用fine_longest_cache_hit()缓存命中匹配就直接返回匹配的kv块id列表所以就不用对long_prefix再匹配和计算kv再看第一次请求是否完成因为有continuous batching和chunk prefill如果完成就重新标记『已占用』如果还在运行第一次请求就将块的引用计数1确保块不被释放。仅prefill新的token并复用前缀的kv块进行前向传播。过程压缩成一张图第一次请求│↓long_prefix prompt[0]│↓切分 KV Block│↓计算 Block Hash│↓find_longest_cache_hit│MISS│↓allocate_slots│↓分配 Block 100/101/102│↓Prefill│↓┌─────────────────────┐│ H0 → Block 100 ││ H1 → Block 101 ││ H2 → Block 102 │└─────────────────────┘│││ 缓存↓第二次请求│↓long_prefix prompt[1]│↓切分 KV Block│↓计算 Block Hash│↓find_longest_cache_hit│HIT H0/H1│↓┌─────────────────────┐│ H0 → Block 100 │ ← 复用│ H1 → Block 101 │ ← 复用└─────────────────────┘│↓只 Prefill新 token│↓X Y Z3sepculative decoding投机采样通过同时利用小型和大型模型来加速token的生成。投机采样小模型/简单方法先猜一串token大模型一次性检查这一串token猜对的全部接受猜错的从那里截断。