EngineCode#_initialize_kv_caches: kv cache 的分配按照 attention head 的数量分配比如 gpt2 有12 个 head那么 kv cache 的规格如下profile分析模型的峰值内存占用从而决定可用于 KV Cache的内存大小。model_executor.determine_available_memory()可以通过指定 kv_cache_memory_bytes 参数直接指定 KV Cache 的容量其不受 gpu_memory_utilization 参数的影响但存在 OOM 的风险。profile_run使用虚拟输入执行一次前向传播以分析模型的内存使用情况memory_profiling - 内存性能分析上下文管理器contextlib.contextmanagerdefmemory_profiling(baseline_snapshot:MemorySnapshot,weights_memory:int0,)-Generator[MemoryProfilingResult,None,None]baseline_snapshot当前 vLLM 实例创建之前的内存快照。weights_memory: PyTorch 加载模型权重时使用的内存。请注意在加载模型权重之前我们还会初始化设备和分布式环境这可能会消耗一些内存。这部分内存不包含在 weights_memory 中因为 PyTorch 并不控制它。单张 GPU 上的内存可以分为 3 类1. 被当前 vLLM 实例之外的其他内容占用的内存。2. 被当前 vLLM 实例中的 torch 占用的内存。3. 在当前 vLLM 实例中但不由 torch 控制的内存如 NCCL、注意力后端缓冲区等。量化示例 创建当前 vLLM 实例之前 * 类别 11 GiB * 类别 20 GiB * 类别 30 GiB 创建当前 vLLM 实例并加载模型之后即性能分析之前 * 类别 11 GiB * 类别 22 GiB模型权重占 2 GiB * 类别 30.5 GiBNCCL 使用的内存 性能分析期间峰值 * 类别 11 GiB * 类别 24 GiB峰值激活张量占 2 GiB * 类别 31 GiBNCCL 某些注意力后端的缓冲区 性能分析之后 * 类别 11 GiB * 类别 23 GiB激活张量被垃圾回收后 * 类别 31 GiBNCCL 某些注意力后端的缓冲区 在这种情况下非 KV Cache 总共占用 5 GiB包括 * a. 模型权重占用的 2 GiB类别 2 * b. 为峰值激活张量预留的 2 GiB类别 2 * c. 非 torch 组件占用的 1 GiB类别 3 加载权重所使用的内存a.直接从参数 weights_memory 中获取。 性能分析期间 torch.accelerator.memory_stats()[allocated_bytes.all.peak] 的增量即为b.。 从创建当前 vLLM 实例开始到性能分析结束non_torch_memory 的增量即为c.。_dummy_run涉及的参数max_num_seqsmax_num_tokens组装 forward 批次每个 seq 所需的 token 数量比如num_scheduled_tokens_list [2, 2, 2, 2, 2, 6]num_reqsmin(num_tokens,max_num_reqs)min_tokens_per_reqnum_tokens//num_reqs num_scheduled_tokens_list[min_tokens_per_req]*num_reqs num_scheduled_tokens_list[-1]num_tokens%num_reqsavailable_kv_cache_memory_bytesself.available_kv_cache_memory_bytes(self.requested_memory# vLLM 允许使用的总显存预算-profile_result.non_kv_cache_memory# 非 KV cache 占用的显存-cudagraph_memory_estimate_applied# CUDA Graph 额外预估占用)这几行是在计算当前 GPU 上还能留给 KV cache 用的显存有多少。self.requested_memoryvLLM 根据gpu_memory_utilization算出来的可用显存预算。比如启动时设置--gpu-memory-utilization 0.95。大致表示vLLM 最多可以使用当前 GPU 空闲显存中的 95%。这个预算就是requested_memory。profile_result.non_kv_cache_memory这是 profiling 阶段测出来的非KV cache 显存占用通常包括模型权重显存activation 峰值显存torch 分配器额外占用非 torch 显存增长其他运行时开销cudagraph_memory_estimate_applied这是 CUDA Graph 可能额外需要的显存预估值cudagraph_memory_estimate_applied(cudagraph_memory_estimateifenvs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHSelse0)所以只有开启VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS时才会把 CUDA Graph 的预估显存从 KV cache 预算里扣掉。举个例子requested_memory 20 GBnon_kv_cache_memory 6 GBcudagraph_memory_estimate_applied 1 GB。那么available_kv_cache_memory_bytes 20 - 6 - 1 13 GB。即vLLM 最多可以拿 13GB 显存来分配 KV cache blocks。这一步很关键因为 KV cache 大小直接决定了 vLLM 能同时容纳多少 token也就是最大并发、最大上下文长度、以及调度时能不能接更多请求。get_kv_cache_configsholdinitialize_from_config (gpuworker.py)gpu_model_runnerinitialize kv_cache根据 kv_cache_config每个 head 分配一个 tensor911523840,forkv_cache_tensorinkv_cache_config.kv_cache_tensors:tensortorch.zeros(kv_cache_tensor.size,dtypetorch.int8,deviceself.device)forlayer_nameinkv_cache_tensor.shared_by:kv_cache_raw_tensors[layer_name]tensor地址划分GPU 底层分配出来的确是一整段连续地址。按 kv_cache_tensor.size 分配一块连续的 raw buffer。所谓 “block” 不是 CUDA allocator 分配了很多小块也不是物理地址天然不连续而是 vLLM 在这块连续 buffer 上人为切出来的固定大小 page/block。一整块连续 raw_tensor 会 reshape 成 attention backend 需要的 KV cache 形状。所以 “block” 是一个逻辑/索引层面的概念raw_tensor 连续内存:[block0][block1][block2][block3]...[block N-1]每个 block 在地址上通常也是连续排布的但对于某个 request 来说它拿到的 block id 可以是不连续的block_table[seq_A][7,2,19,...]seq A 的逻辑块:logical block0-physical block7logical block1-physical block2logical block2-physical block19PagedAttention 的好处也在这里KV cache 总池子是连续大 buffer但每个 sequence 不需要占用一段连续的 KV cache。它只需要拿若干物理 block靠 block_table 拼成逻辑连续的上下文。这样可以减少因为不同序列长度导致的碎片和搬移。reshape kv_cache把 raw buffer reshape 成 attn backend 需要的 KV cache 形状2,18545,16,12,64K/V, kernel_num_blocks, kernel_block_size, num_kv_heads, head_sizenum_blocksraw_tensor.numel()//kv_cache_spec.page_size_bytes num_blocks_per_kv_block(kv_cache_spec.block_size//kernel_block_size)kernel_num_blocksnum_blocks*num_blocks_per_kv_blockbind kv_cache将已分配的 KV Cache 同时绑定到 ModelRunner 和 forward context使得 KV Cache 可以在前向传播中被使用。该函数的作用将 kv_caches 填充到ModelRunner的 KV Cache 列表runner_kv_caches中。(self.kv_caches)将 forward_context 中的每个注意力层与 kv_caches 中对应的 KV Cache 关联起来。(self.compilation_config.static_forward_context)逻辑到物理地址映射slot_mapping_prepare_input (GPUModelRunner)query_start_loc: 表示每个 request 在当前 batch token 数组里的起止位置。比如num_scheduled_tokens [2, 3, 1]那么 query_start_loc [0, 2, 5, 6]其中req0 对应 token_idx [0, 2)req1 对应 token_idx [2, 5)req2 对应 token_idx [5, 6)positions表示每个 token 在自己 request 序列里的绝对位置。比如 decode 阶段某个 request 已经算了 100 个 token本轮新 token 的 position 就是 100。block_table.compute_slot_mapping(num_reqs,self.query_start_loc.gpu[:num_reqs1],self.positions[:total_num_scheduled_tokens],)_compute_slot_mapping_kernel(block_table.py)为当前 batch 的每个 token 准备 KV cache 写入位置映射。prefill 时会把一批 token 映射到一串 slotdecode 时通常每个 request 映射一个新 slot。# per-token 查表操作# 1. 从 token position 算出它落在第几个逻辑 blockblock_indicespos//block_size# 2. 查 block_table拿到这个逻辑 block 对应的物理 block idblock_numberstl.load(block_table_ptrrow_offsetblock_indices)# physical_block_id block_table[row_offset][block_indices]# 3. 物理 block id * block_size 块内偏移 物理 slot idslot_idsblock_numbers*block_sizelocal_block_offsets# 这里从连续地址 kv_cache_blocks 查找# 形状block_table:shape[max_num_reqs,max_num_blocks_per_req]dtypeint32 CPU 一份GPU 一份 slot_mapping:shape[max_num_batched_tokens]dtypeint64 CPU 一份GPU 一份 num_blocks_per_row:shape[max_num_reqs]dtypenumpy int32 只有 CPU为什么需要 slot_mapping不能直接用 block_tableblock_table 是 per-request 的逻辑→物理 block 映射二维req × max_blocks。slot_mapping 是 per-token 的最终物理位置一维num_tokens已经把 batch 中所有 request 摊平、把 block 偏移加好了。attention kernel 在 token 维度并行per-token 拿到一个 slot_id 直接就能 scatter 写入省去 kernel 内再做 block_table[req][block_idx] 间接寻址。所以 slot_mapping 本质是预先在 CPU/GPU 上算好的写入索引是 prepare_inputs 阶段的产物。CPU/GPU 上构造 slot_mapping逻辑块 → 物理 slot核心 kernelblock_table.py:319-373是个 Triton kernel_compute_slot_mapping_kernel。每个 token 的计算# 输入positions[i]# 第 i 个 token 在它 request 里的绝对位置含已生成的block_table[req,*]# 这个 request 的逻辑块号 → 物理块号映射block_size# 一个物理块装多少 token (例如 16)# 计算block_table.py:356-371block_indicespos//block_size# 逻辑块下标这个 token 落在 request 的第几个块block_numbersblock_table[req,block_indices]# 物理块号block manager 分配的block_offsetspos%block_size# 这个 token 在块内的位置0..block_size-1slot_idblock_numbers*block_sizeblock_offsets# 最终扁平槽位reshape_and_cache kernel 按 slot 写入const int64_t slot_idxslot_mapping[token_idx];//一次访存拿到目标 slotif(slot_idx0)return;//padding token 跳过 const int64_t block_idxslot_idx/block_size;//反算物理块号 const int64_t block_offsetslot_idx%block_size;//反算块内偏移用 block_idx 和 block_offset 直接索引到key_cache[num_blocks, num_heads, head_size/x, block_size, x]目标位置key_dstkey_cacheblock_idx*(num_heads*h_block_count*block_size*x)//跳到目标块head_idx*(h_block_count*block_size*x)//跳到目标 headh_block*(block_size*x)//head_size 内的子块block_offset*x;//块内 token 偏移gpu layoutflash attentionintnum_tokensslot_mapping.size(0);intnum_headskey.size(1);inthead_sizekey.size(2);dim3 grid(num_tokens);dim3 block(std::min(num_heads*head_size,512));Flash Attentionkernel_num_blocks, kernel_block_size, num_kv_heads, head_size# FA 内存布局token0├─ head0:d0 d1 d2...d63 ← head_size 一次性连续 └─ head1:d0 d1 d2...d63 ← head_size 一次性连续 token1├─ head0:d0 d1 d2...d63 └─ head1:d0 d1 d2...d63 token2├─ head0:d0 d1 d2...d63 └─ head1:d0 d1 d2...d63 token3├─ head0:d0 d1 d2...d63 └─ head1:d0 d1 d2...d63为什么这种布局核心目标让一个 block 的 K/V 是一段连续大内存可以用 cp.async / TMA 一次性异步装载到 shared memory然后直接喂给Tensor Core MMA。// FA attention kernel// FlashAttention 把 K block 整块搬到 shared memory__shared__ half K_smem[block_size][head_size];// [4][16]// 用 cp.async / TMA 一次性异步搬运 256Bcp_async_bulk(K_smem[0][0],key_cache[block_id,0,head,0],256);__syncthreads();// 然后用 Tensor Core MMA 计算 Q K^Tmma.sync(...);// m16n8k16 等指令 【vLLM 源码解析】KV Cache
EngineCode#_initialize_kv_caches: kv cache 的分配按照 attention head 的数量分配比如 gpt2 有12 个 head那么 kv cache 的规格如下profile分析模型的峰值内存占用从而决定可用于 KV Cache的内存大小。model_executor.determine_available_memory()可以通过指定 kv_cache_memory_bytes 参数直接指定 KV Cache 的容量其不受 gpu_memory_utilization 参数的影响但存在 OOM 的风险。profile_run使用虚拟输入执行一次前向传播以分析模型的内存使用情况memory_profiling - 内存性能分析上下文管理器contextlib.contextmanagerdefmemory_profiling(baseline_snapshot:MemorySnapshot,weights_memory:int0,)-Generator[MemoryProfilingResult,None,None]baseline_snapshot当前 vLLM 实例创建之前的内存快照。weights_memory: PyTorch 加载模型权重时使用的内存。请注意在加载模型权重之前我们还会初始化设备和分布式环境这可能会消耗一些内存。这部分内存不包含在 weights_memory 中因为 PyTorch 并不控制它。单张 GPU 上的内存可以分为 3 类1. 被当前 vLLM 实例之外的其他内容占用的内存。2. 被当前 vLLM 实例中的 torch 占用的内存。3. 在当前 vLLM 实例中但不由 torch 控制的内存如 NCCL、注意力后端缓冲区等。量化示例 创建当前 vLLM 实例之前 * 类别 11 GiB * 类别 20 GiB * 类别 30 GiB 创建当前 vLLM 实例并加载模型之后即性能分析之前 * 类别 11 GiB * 类别 22 GiB模型权重占 2 GiB * 类别 30.5 GiBNCCL 使用的内存 性能分析期间峰值 * 类别 11 GiB * 类别 24 GiB峰值激活张量占 2 GiB * 类别 31 GiBNCCL 某些注意力后端的缓冲区 性能分析之后 * 类别 11 GiB * 类别 23 GiB激活张量被垃圾回收后 * 类别 31 GiBNCCL 某些注意力后端的缓冲区 在这种情况下非 KV Cache 总共占用 5 GiB包括 * a. 模型权重占用的 2 GiB类别 2 * b. 为峰值激活张量预留的 2 GiB类别 2 * c. 非 torch 组件占用的 1 GiB类别 3 加载权重所使用的内存a.直接从参数 weights_memory 中获取。 性能分析期间 torch.accelerator.memory_stats()[allocated_bytes.all.peak] 的增量即为b.。 从创建当前 vLLM 实例开始到性能分析结束non_torch_memory 的增量即为c.。_dummy_run涉及的参数max_num_seqsmax_num_tokens组装 forward 批次每个 seq 所需的 token 数量比如num_scheduled_tokens_list [2, 2, 2, 2, 2, 6]num_reqsmin(num_tokens,max_num_reqs)min_tokens_per_reqnum_tokens//num_reqs num_scheduled_tokens_list[min_tokens_per_req]*num_reqs num_scheduled_tokens_list[-1]num_tokens%num_reqsavailable_kv_cache_memory_bytesself.available_kv_cache_memory_bytes(self.requested_memory# vLLM 允许使用的总显存预算-profile_result.non_kv_cache_memory# 非 KV cache 占用的显存-cudagraph_memory_estimate_applied# CUDA Graph 额外预估占用)这几行是在计算当前 GPU 上还能留给 KV cache 用的显存有多少。self.requested_memoryvLLM 根据gpu_memory_utilization算出来的可用显存预算。比如启动时设置--gpu-memory-utilization 0.95。大致表示vLLM 最多可以使用当前 GPU 空闲显存中的 95%。这个预算就是requested_memory。profile_result.non_kv_cache_memory这是 profiling 阶段测出来的非KV cache 显存占用通常包括模型权重显存activation 峰值显存torch 分配器额外占用非 torch 显存增长其他运行时开销cudagraph_memory_estimate_applied这是 CUDA Graph 可能额外需要的显存预估值cudagraph_memory_estimate_applied(cudagraph_memory_estimateifenvs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHSelse0)所以只有开启VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS时才会把 CUDA Graph 的预估显存从 KV cache 预算里扣掉。举个例子requested_memory 20 GBnon_kv_cache_memory 6 GBcudagraph_memory_estimate_applied 1 GB。那么available_kv_cache_memory_bytes 20 - 6 - 1 13 GB。即vLLM 最多可以拿 13GB 显存来分配 KV cache blocks。这一步很关键因为 KV cache 大小直接决定了 vLLM 能同时容纳多少 token也就是最大并发、最大上下文长度、以及调度时能不能接更多请求。get_kv_cache_configsholdinitialize_from_config (gpuworker.py)gpu_model_runnerinitialize kv_cache根据 kv_cache_config每个 head 分配一个 tensor911523840,forkv_cache_tensorinkv_cache_config.kv_cache_tensors:tensortorch.zeros(kv_cache_tensor.size,dtypetorch.int8,deviceself.device)forlayer_nameinkv_cache_tensor.shared_by:kv_cache_raw_tensors[layer_name]tensor地址划分GPU 底层分配出来的确是一整段连续地址。按 kv_cache_tensor.size 分配一块连续的 raw buffer。所谓 “block” 不是 CUDA allocator 分配了很多小块也不是物理地址天然不连续而是 vLLM 在这块连续 buffer 上人为切出来的固定大小 page/block。一整块连续 raw_tensor 会 reshape 成 attention backend 需要的 KV cache 形状。所以 “block” 是一个逻辑/索引层面的概念raw_tensor 连续内存:[block0][block1][block2][block3]...[block N-1]每个 block 在地址上通常也是连续排布的但对于某个 request 来说它拿到的 block id 可以是不连续的block_table[seq_A][7,2,19,...]seq A 的逻辑块:logical block0-physical block7logical block1-physical block2logical block2-physical block19PagedAttention 的好处也在这里KV cache 总池子是连续大 buffer但每个 sequence 不需要占用一段连续的 KV cache。它只需要拿若干物理 block靠 block_table 拼成逻辑连续的上下文。这样可以减少因为不同序列长度导致的碎片和搬移。reshape kv_cache把 raw buffer reshape 成 attn backend 需要的 KV cache 形状2,18545,16,12,64K/V, kernel_num_blocks, kernel_block_size, num_kv_heads, head_sizenum_blocksraw_tensor.numel()//kv_cache_spec.page_size_bytes num_blocks_per_kv_block(kv_cache_spec.block_size//kernel_block_size)kernel_num_blocksnum_blocks*num_blocks_per_kv_blockbind kv_cache将已分配的 KV Cache 同时绑定到 ModelRunner 和 forward context使得 KV Cache 可以在前向传播中被使用。该函数的作用将 kv_caches 填充到ModelRunner的 KV Cache 列表runner_kv_caches中。(self.kv_caches)将 forward_context 中的每个注意力层与 kv_caches 中对应的 KV Cache 关联起来。(self.compilation_config.static_forward_context)逻辑到物理地址映射slot_mapping_prepare_input (GPUModelRunner)query_start_loc: 表示每个 request 在当前 batch token 数组里的起止位置。比如num_scheduled_tokens [2, 3, 1]那么 query_start_loc [0, 2, 5, 6]其中req0 对应 token_idx [0, 2)req1 对应 token_idx [2, 5)req2 对应 token_idx [5, 6)positions表示每个 token 在自己 request 序列里的绝对位置。比如 decode 阶段某个 request 已经算了 100 个 token本轮新 token 的 position 就是 100。block_table.compute_slot_mapping(num_reqs,self.query_start_loc.gpu[:num_reqs1],self.positions[:total_num_scheduled_tokens],)_compute_slot_mapping_kernel(block_table.py)为当前 batch 的每个 token 准备 KV cache 写入位置映射。prefill 时会把一批 token 映射到一串 slotdecode 时通常每个 request 映射一个新 slot。# per-token 查表操作# 1. 从 token position 算出它落在第几个逻辑 blockblock_indicespos//block_size# 2. 查 block_table拿到这个逻辑 block 对应的物理 block idblock_numberstl.load(block_table_ptrrow_offsetblock_indices)# physical_block_id block_table[row_offset][block_indices]# 3. 物理 block id * block_size 块内偏移 物理 slot idslot_idsblock_numbers*block_sizelocal_block_offsets# 这里从连续地址 kv_cache_blocks 查找# 形状block_table:shape[max_num_reqs,max_num_blocks_per_req]dtypeint32 CPU 一份GPU 一份 slot_mapping:shape[max_num_batched_tokens]dtypeint64 CPU 一份GPU 一份 num_blocks_per_row:shape[max_num_reqs]dtypenumpy int32 只有 CPU为什么需要 slot_mapping不能直接用 block_tableblock_table 是 per-request 的逻辑→物理 block 映射二维req × max_blocks。slot_mapping 是 per-token 的最终物理位置一维num_tokens已经把 batch 中所有 request 摊平、把 block 偏移加好了。attention kernel 在 token 维度并行per-token 拿到一个 slot_id 直接就能 scatter 写入省去 kernel 内再做 block_table[req][block_idx] 间接寻址。所以 slot_mapping 本质是预先在 CPU/GPU 上算好的写入索引是 prepare_inputs 阶段的产物。CPU/GPU 上构造 slot_mapping逻辑块 → 物理 slot核心 kernelblock_table.py:319-373是个 Triton kernel_compute_slot_mapping_kernel。每个 token 的计算# 输入positions[i]# 第 i 个 token 在它 request 里的绝对位置含已生成的block_table[req,*]# 这个 request 的逻辑块号 → 物理块号映射block_size# 一个物理块装多少 token (例如 16)# 计算block_table.py:356-371block_indicespos//block_size# 逻辑块下标这个 token 落在 request 的第几个块block_numbersblock_table[req,block_indices]# 物理块号block manager 分配的block_offsetspos%block_size# 这个 token 在块内的位置0..block_size-1slot_idblock_numbers*block_sizeblock_offsets# 最终扁平槽位reshape_and_cache kernel 按 slot 写入const int64_t slot_idxslot_mapping[token_idx];//一次访存拿到目标 slotif(slot_idx0)return;//padding token 跳过 const int64_t block_idxslot_idx/block_size;//反算物理块号 const int64_t block_offsetslot_idx%block_size;//反算块内偏移用 block_idx 和 block_offset 直接索引到key_cache[num_blocks, num_heads, head_size/x, block_size, x]目标位置key_dstkey_cacheblock_idx*(num_heads*h_block_count*block_size*x)//跳到目标块head_idx*(h_block_count*block_size*x)//跳到目标 headh_block*(block_size*x)//head_size 内的子块block_offset*x;//块内 token 偏移gpu layoutflash attentionintnum_tokensslot_mapping.size(0);intnum_headskey.size(1);inthead_sizekey.size(2);dim3 grid(num_tokens);dim3 block(std::min(num_heads*head_size,512));Flash Attentionkernel_num_blocks, kernel_block_size, num_kv_heads, head_size# FA 内存布局token0├─ head0:d0 d1 d2...d63 ← head_size 一次性连续 └─ head1:d0 d1 d2...d63 ← head_size 一次性连续 token1├─ head0:d0 d1 d2...d63 └─ head1:d0 d1 d2...d63 token2├─ head0:d0 d1 d2...d63 └─ head1:d0 d1 d2...d63 token3├─ head0:d0 d1 d2...d63 └─ head1:d0 d1 d2...d63为什么这种布局核心目标让一个 block 的 K/V 是一段连续大内存可以用 cp.async / TMA 一次性异步装载到 shared memory然后直接喂给Tensor Core MMA。// FA attention kernel// FlashAttention 把 K block 整块搬到 shared memory__shared__ half K_smem[block_size][head_size];// [4][16]// 用 cp.async / TMA 一次性异步搬运 256Bcp_async_bulk(K_smem[0][0],key_cache[block_id,0,head,0],256);__syncthreads();// 然后用 Tensor Core MMA 计算 Q K^Tmma.sync(...);// m16n8k16 等指令