vLLM源码之框架执行

原文：https://zhuanlan.zhihu.com/p/710927181

本文主要介绍vLLM推理引擎的框架执行流程 (v0.1.2)，相关文章：

vLLM源码之PagedAttention（持续更新）

引用

本文内容主要源于：

1. https://tech.scatterlab.co.kr/vllm-implementation-details/ 。

2. （https://tech.scatterlab.co.kr/vllm-implementation-details/）

3. 
   

4. https://docs.google.com/presentation/d/1QL-XPFXiFpDBh86DbEegFXBXFXjix4v032GhShbKf3s/edit#slide=id.g24ad94a0065_0_209，vLLM的meetup slides

5. （https://docs.google.com/presentation/d/1QL-XPFXiFpDBh86DbEegFXBXFXjix4v032GhShbKf3s/edit#slide=id.g24ad94a0065_0_209）
   

感谢 @lipi 提供该技术博客的原网址。

随着 LLM 时代的到来，为服务开发和研究了多种优化方法。今天，我们将分析 vLLM，其性能比 Hugging Face 提高了 24 倍。请注意，本文分析的内容基于 vLLM 刚发布时的实现版本 (v0.1.2)，因此可能存在一些已更改的部分。其中包含了大量深入乃至代码级别的知识，因此建议有深入了解愿望的人阅读！

vLLM

vLLM 是一种利用 PagedAttention 技术大幅提升句子生成速度的方法。除此之外，它还包含了许多可用于实际服务的要素。例如，为了在多集群环境中实现稳定服务，它使用了 Ray Cluster；为了能够并行处理大型模型和数据，它借用了 Megatron LM 的并行性（张量并行）。Ray Cluster 或 Megatron LM 已有诸多相关文章，本文中不再赘述。在本篇博文中，我们重点了解一下 vLLM 的核心技术：PagedAttention 和连续批处理技术，并对它们进行代码级别的分析。

整体结构

vLLM 的整体架构。LLMEngine 中生成用于分布式处理的 Worker、管理 PagedAttention 区块的区块管理器、管理 KV Cache 等组件，并通过调度器按顺序更改每个请求中请求的提示的生成顺序。为了在 LM 中生成句子，需要在模型中反复前馈，直到出现最后一个令牌，调度器会根据内存和优先级，通过有效的方式，实现 GPU 实用程序。为了进一步提升内存效率，使用了 PagedAttention 技术，并在 GPU 内存不足时，以交换到 CPU 内存的方式，稳定管理中间计算过程。

现在逐一分析上述组件之间的关系和流程。

使用方式

目前vLLM支持2种方式，分别是在线方式和离线的方式。

在线方式使用如下：

离线方式使用如下：

组件关系

如上图所示，LLMEngine是核心的类，离线和在线均会使用，提供了add_request()和step 2个方法，LLM类是离线推理使用的类，另外，作为服务提供的时候，会有api_server相关的类，底层调用AsyncLLMEngine，下面详细分析LLMEngine和LLM。

LLM类

如上文vLLM离线demo，vLLM最基本的使用方法有2个。

1. 创建LLM()类，即调用LLM的__init__()方法；

2. 使用generate()函数生成句子。

LLM.__init__()

我们不妨想一下初始化主要需要做什么。

首先，肯定是需要创建并且加载模型的，如下图所示。

LLM初始化的时候会调用LLMEngine的初始化方法，整体调用流程如上图所示。

其次，初始化时需要创建kv cache，这部分内容在后文LLMEngine类那一节会介绍。

generate

def generate(
    self,
    prompts: Optional[Union[str, List[str]]] = None,
    sampling_params: Optional[SamplingParams] = None,
    prompt_token_ids: Optional[List[List[int]]] = None,
    use_tqdm: bool = True,
) -> List[RequestOutput]:

本代码是 generate() 函数接受的参数的列表。prompts 变量接收字符串组成的提示列表，并接收 SamplingParams 参数进行生成。

SamplingParams 选项列表

class SamplingParams:
    # 根据生成的best_of个句子中选择n个最好的句子.
    n: int = 1
    # 针对每个提示生成多少个句子. 默认值与n相同，并且该值必须大于或等于n.
    best_of: Optional[int] = None
    # 已出现令牌的惩罚(包括提示中的令牌)
    # 大于0时，将倾向于生成尚未出现的令牌，
    # 小于0时，将倾向于重复生成令牌.
    presence_penalty: float = 0.0
    # 令牌出现频率的惩罚(包括提示中的令牌)
    # presence_penalty会对出现一次的令牌给予相同的惩罚，
    # 而frequency_penalty会根据频率区别对待惩罚.
    # 类似地，大于0时，倾向于新令牌，而小于0时，则倾向于已出现的令牌.
    frequency_penalty: float = 0.0
    temperature: float = 1.0
    top_p: float = 1.0
    top_k: int = -1
    use_beam_search: bool = False
    # 指定特定字符串或列表, 如果生成以该字符串结尾，则停止生成.
    stop: Union[None, str, List[str]] = None
    # 如果为True，即使出现EOS，也会继续生成.
    ignore_eos: bool = False
    # 要生成的令牌的最大数目.
    max_tokens: int = 16
    # 除了生成的令牌之外，还可以获取具有最高概率的令牌. 将返回logprobs中指定数目具有高概率的令牌.
    logprobs: Optional[int] = None

SamplingParams 类可以指定生成所需选项。

# 如果给定提示 (由字符串组成的提示列表)，则为其设置该值，
if prompts is not None:
    num_requests = len(prompts)
# 否则，使用提示的令牌。
else:
    num_requests = len(prompt_token_ids)
for i in range(num_requests):
    prompt = prompts[i] if prompts is not None else None
    if prompt_token_ids is None:
        token_ids = None
    else:
        token_ids = prompt_token_ids[i]
    self._add_request(prompt, sampling_params, token_ids)
return self._run_engine(use_tqdm)

generate() 函数的逻辑为：

1. 如果没有给出提示，但是给出了 prompt_token_ids，则使用已经编码好的 token 进行生成；

2. 然后为每个提示分别通过 LLMEngine.add_request() 将提示添加到生成中；

3. 当所有内容都添加完毕后，调用 _run_engine()；

LLMEngine 执行

当请求通过_add_request加入到LLMEngine之后，会调用_run_engine，_run_engine代码如下。

def _run_engine(self, use_tqdm: bool) -> List[RequestOutput]:
    # 运行 LLMEngine 以生成句子。
    outputs: List[RequestOutput] = []

    while self.llm_engine.has_unfinished_requests():
        step_outputs = self.llm_engine.step()

        for output in step_outputs:
            if output.finished:
                outputs.append(output)

    return outputs

在 _run_engine() 函数中，对于使用 add_request() 添加的所有提示，它会调用 LLMEngine.step()，直到它们完成。step() 函数对模型执行一次前向传播 ，在该批次中生成提示的标记，一个接一个。

LLMEngine类

LLMEngine 是负责实际生成的引擎。LLMEngine底层核心依赖Scheduler和Worker。在初始化时，它将生成生成所需的所有组件（如 tokenizer、worker 和 scheduler）并对其进行初始化。此时，将为每个等级生成一个 worker 以进行并行处理。并行化使用 MegatronLM 的并行性，分布环境使用 Ray 集群。

在引擎初始化之后，它将执行 KV Cache 的初始操作。

缓存初始化

通过 _init_cache 函数进行初始化以存储 KV 缓存。

# 函数 profile_num_available_blocks() 内部

# 首先运行一次模型前向传播
self.model(
    input_ids=input_tokens,
    positions=input_positions,
    kv_caches=[(None, None)] * num_layers,
    input_metadata=input_metadata,
    cache_events=None,
)

torch.cuda.synchronize()

# 计算可用块数
peak_memory = torch.cuda.max_memory_allocated()
# 获得全部 GPU 内存大小
total_gpu_memory = get_gpu_memory()
# 根据参数计算缓存块大小
cache_block_size = CacheEngine.get_cache_block_size(
    block_size, self.model_config, self.parallel_config)
# 根据可用内存大小计算最大块数
num_gpu_blocks = int((total_gpu_memory * gpu_memory_utilization
                      - peak_memory) // cache_block_size)
num_cpu_blocks = int(cpu_swap_space // cache_block_size)
num_gpu_blocks = max(num_gpu_blocks, 0)
num_cpu_blocks = max(num_cpu_blocks, 0)

_init_cache 函数通过 Worker 类的 profile_num_available_blocks 函数计算块数。块是 PagedAttention 中使用的概念，类似于操作系统内存管理方法之一页，一个块中存储多个令牌。

profile_num_available_blocks方法通过以下过程计算：

1. 一次使用给定的参数转发模型。

2. 使用 PyTorch 的 max_memory_allocated 函数获取最大已使用 GPU 内存。

3. 从给定的最大 GPU 使用限制值（gpu_memory_utilization）中减去第 2 步中获得的大小，以计算可用于存储块的内存大小。

4. 使用该大小计算最大可用块数。换句话说，除了用于转发的 GPU 之外，其余的用作缓存内存。

def get_cache_block_size(block_size, model_config, parallel_config):
    # 每个头的维度
    head_size = model_config.get_head_size()
    # 将头的数量除以Tensor并行数（并行是Megatron LM的参数）
    num_heads = model_config.get_num_heads(parallel_config)
    # 将总层数除以Pipeline并行数（并行是Megatron LM的参数）
    num_layers = model_config.get_num_layers(parallel_config)
    
    # 包含在块中的元素数量（block_size）乘以每个元素包含的参数数量
    key_cache_block = block_size * num_heads * head_size
    value_cache_block = key_cache_block
    total = num_layers * (key_cache_block + value_cache_block)
    dtype_size = _get_dtype_size(model_config.dtype)
    
    # cache_block_size = dtype_size * total
    return dtype_size * total

cache_block_size表示PagedAttention中使用的块所占用的实际内存大小。它接收的参数是每个块最多可以容纳多少个令牌信息，因此实际占用多少内存需要单独计算。

队列

将句子生成请求提示添加到队列中。在这个过程中，将根据SamplingParams中指定的best_of数量复制序列并将其存储在队列中。每个提示将被创建为一个Sequence对象，并将一个请求(=提示)复制best_of次以将SequenceGroup对象添加到等待队列中。

step

step是进行一次前向执行，step的功能如下图解：

上图执行了2次step函数，详细介绍如下（为简单作图，这里的一个字母如a、b、x、y等均代表一个token，可能与实际情况不一定相同）：

1. 初始2条数据，分别为 abc、xyz；

2. 经过一次step，2个句子各自产生一个token；

3. 在执行下一个step之前，队列中又多了一个句子，hi；

4. 第二次step，这3条句子各自都生成了一个次。

step执行总共有2个步骤。

1. 调度程序使用Scheduler.schedule()函数从当前保存的请求队列中获取要转发到模型的序列，并准备缓存块以应用PagedAttention。

2. 然通过Worker.execute_model将模型转发，生成并保存下一个标记。

seq_group_metadata_list, scheduler_outputs = self.scheduler.schedule()
if (not seq_group_metadata_list) and scheduler_outputs.is_empty():
    # 无需执行任何操作。
    return []

# 指示模型的前向传播。
output = self._run_workers(
    "execute_model",
    seq_group_metadata_list=seq_group_metadata_list,
    blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in,
    blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out,
    blocks_to_copy=scheduler_outputs.blocks_to_copy,
)
# 将前向传播后的结果更新到调度器。
seq_groups = self.scheduler.update(output)

# 解码前向传播中生成的令牌。
self._decode_sequences(seq_groups)
# 调整达到生成中断标准的序列。
self._stop_sequences(seq_groups)
# 从队列中移除生成已完成的序列。
self.scheduler.free_finished_seq_groups()

# 储存生成结果。
request_outputs: List[RequestOutput] = []
for seq_group in seq_groups:
    request_output = RequestOutput.from_seq_group(seq_group)
    request_outputs.append(request_output)
return request_outputs

Schedule

然后执行一个调度任务，以便根据优先级对需要一次性处理的大量请求进行整理，因为无法一次性处理。然后，让我们了解如何分配内存以应用 PagedAttention。

在开始之前，我们先整理一下术语：

• Sequence：这是转发到实际模型的顺序。

• SequenceGroup：可以指定为每个提示生成句子的数量，这是按每个提示分组的。如果想要使用总共 3 个提示为每个提示生成 4 个句子，则总顺序数为 3 x 4 = 12，顺序组为 3。

• Block：在 PagedAttention 中，顺序中令牌的键值对会被分成一个固定数量并进行管理，类似操作系统中的页面。

• Slot：块中分配或可用于令牌的键或值的空间。如果块大小为 8，则某个块的槽位数为 8。

状态

状态已分为总共3类。

Waiting

• 这是正在等待SequenceGroup队列的SequenceGroup。包含KVCache未分配的SequenceGroup。

Running

• 已转发到模型的队列，此处正在生成句子。因此，已分配了提示和生成标记的KVCache。

Swapped

• SequenceGroup处于运行阶段，但由于逐个生成标记而可能内存不足，因此无法进入正在运行的状态，而是在等待。

• KVCache已交换到CPU内存。

Slot Swap Out

调度的第一步是检查当前正在运行（状态为 Running）的 SequenceGroup 中是否有可以按照优先级分配槽的 SequenceGroup。由于在生成令牌时所需的内存会增加，因此如果优先级较高的 SequenceGroup 没有足够的内存，则需要将优先级较低的 SequenceGroup 的槽从 GPU 内存中移除。

本过程通过重复以下顺序实现。

1. 从处于运行状态的序列组中弹出优先级最高的组（=GroupA）。

2. 如果此组无法分配一个高速缓存块槽，则重复以下过程：

2.1 从运行状态的组中选择一个优先级最低的组（=GroupB）。

如果没有，则预占当前组（GroupA）。

2.2 预占选定的组（GroupB）。

3. 将当前组（GroupA）添加到运行状态。

抢占策略如下：

1. 在 SequenceGroup 中获取尚未完成生成的序列数。

2. 如果该值为 1，则释放该序列组的所有 KVCache 并将其更改为等待状态。

当 SequenceGroup 更改为正在运行状态时，将重新计算 KV 缓存。（RECOMPUTE）

3. 如果不这样做，则将交换 KVCache 内存。（SWAP）

将要交换出的组的 KV Cache 块存储在 blocks_to_swap_out 变量中。

等待状态中的 SequenceGroup 分配

如果存在已交换组和处于 Waiting 状态的组，请执行以下步骤，将处于 Waiting 状态的序列组更改为 Running 状态。此时，无需按优先级排序，原因是未交换且已抢占的组（即，只有 KV 缓存飞走的 SequenceGroup）保持在 Waiting 队列的最前面，不会排序。

逐个获取最前面的序列组并按照以下顺序执行：

1. 如果当前调度认为该组是一个已抢占组，则跳过。

2. 如果当前无法分配到内存，则跳过。

3. 如果该组超过了特定的基准（最大批处理大小、最大令牌数），则跳过。

4. 为该序列组分配 KV 缓存空间，并更改为 Running 状态。

此过程结束后，调度器的输出产品共有 3 种。

• blocks_to_swap_in：如果已交换状态的组需要重新变为 Running 状态并推进流程，则需要 Swap-in CPU↔GPU 之间的块映射表。

• blocks_to_swap_out：如果处于 Running 状态的组已更改为 Swapped 状态，则需要 Swap-out GPU↔CPU 之间的块映射表。

• blocks_to_copy：用于将 GPU 内的块复制到其他位置的块对列表。我们将在块管理器部分中再次了解。

并针对处于 Running 状态的序列组生成 SequenceGroupMetadata，其中包含以下项目：

• request_id：每个提示都会获得一个的 request_id 值。

• is_prompt：它表示该组是否已从 Waiting → Running 更改。如果没有 KV 缓存，则需要使用提示计算 KV 缓存，此值用于确定是否向模型中添加提示令牌。

• seq_data：对于组内所有序列，都存在一个以 ID 为键、以对应令牌为值的映射表。

• sampling_params：用于采样。其中包含 top_k、top_p 等。

• block_tables：映射该序列组内每个序列的物理块列表的表。我们将在块管理器部分中再次了解。

Block Manager

PagedAttention类似于操作系统的虚拟内存，将逻辑 ↔ 物理部分划分，并将每个标记的KV缓存值分区存储为块。负责执行这两项任务的组件是BlockManager。

根据之前的计算，在GPU和CPU内初始化管理器的最大可用块数。在此过程中使用的类是BlockAllocator，其执行以下功能：

• 未分配的块管理

• 块分配和释放

• 引用计数管理

在vLLM中，可以使用相同的提示生成多个句子，此时初始提示标记是共享的，正在生成的标记则各自拥有。

• 管理逻辑块和物理块之间的映射表（块表）

引用计数表示特定块正在使用的序列数量，块表包含序列的块在GPU上时存储GPU块，在换出后块在CPU上时存储CPU块。

Allocation

def allocate(self, seq_group: SequenceGroup) -> None:
    seq = seq_group.get_seqs()[0]

    # 为了保存提示令牌，分配一个新的物理块。
    block_table: BlockTable = []
    for _ in range(len(seq.logical_token_blocks)):
        block = self.gpu_allocator.allocate()
        # 由于提示词令牌都是共享的，Ref Count 被初始化为序列组中所有序列的数量。
        # 用 Sequence Group 中所有“序列”的计数来初始化。
        block.ref_count = seq_group.num_seqs()
        block_table.append(block)

    # 存储每个序列的块表。
    for seq in seq_group.get_seqs():
        self.block_tables[seq.seq_id] = block_table.copy()

Block Manager接收SequenceGroup或Sequence作为参数，然后分配块或槽（=块内令牌所在的空间）。SequenceGroup会像处理Virtual Memory一样自管理Logical block。因此，与实际的物理内存大小无关，都可以创建用于存储令牌的block，并且当每个block内的槽达到block_size大小时，就会发布并分配一个新的Logical block。Block Manager负责将SequenceGroup拥有的Logical block分配到CPU/GPU内存中分配的Physical block。

SequenceGroup 的块分配由 allocate 函数分配，实际分配的块存储在块表中。

而且，如果要生成新标记，则必须添加新槽。allocate 函数是将 SequenceGroup 拥有的逻辑块分配给物理块的函数，而 append_slot 是 SequenceGroup 内的每个序列为标记的新槽重新分配块的函数。

按以下顺序进行：

1. 如果 SequenceGroup 中的逻辑块数量大于块管理器管理的 SequenceGroup 的块表中包含的块数量，则表示为新标记重新分配了逻辑块，因此分配新的物理块并存储在块表中。

2. 如果不是，则表示在现有块中分配了槽。vLLM 是序列在生成时共享块，对应于 SequenceGroup 中的提示，除了每个块用于生成标记的块之外，还报告最后一个块的参考计数，并确认是 1。

3. 如果为 1，则表示每个序列都有自己的新块已经被分配并使用，因此不会分配新块。

4. 如果它不是 1 而是表示块也被其他序列使用，则删除该块的引用（free）并分配新块。此时，返回的值返回现有的块号和新分配的块号，以便随后将其放入 blocks_to_copy 映射表中。这样，以前共享的块就会被复制，每个序列都会拥有自己的块。

Swap

def swap_out(self, seq_group: SequenceGroup) -> Dict[int, int]:
    # GPU block -> CPU block.
    mapping: Dict[PhysicalTokenBlock, PhysicalTokenBlock] = {}
    for seq in seq_group.get_seqs():
        if seq.is_finished():
            continue
        new_block_table: BlockTable = []
        block_table = self.block_tables[seq.seq_id]
  
        for gpu_block in block_table:
            if gpu_block in mapping:
                cpu_block = mapping[gpu_block]
                cpu_block.ref_count += 1
            else:
                cpu_block = self.cpu_allocator.allocate()
                mapping[gpu_block] = cpu_block
            new_block_table.append(cpu_block)
            # Free the GPU block swapped out to CPU.
            self.gpu_allocator.free(gpu_block)
        self.block_tables[seq.seq_id] = new_block_table
  
    block_number_mapping = {
        gpu_block.block_number: cpu_block.block_number
        for gpu_block, cpu_block in mapping.items()
    }
    return block_number_mapping

BlockManager 负责修改 block_table 以进行 Swap-In/Out。当发生 Swap-Out 时，需要将 GPU 内存上的块移至 CPU。

1. 使用 CPU Allocator 为 GPU 上的所有块分配对应的块。

2. Free GPU 块，然后使 SequenceGroup 的块与 1 中分配的块对应。

3. 返回以前存在于 GPU 上的块号与现在分配到 CPU 内存上的新的块号之间的映射表。

如果发生 Swap-In，则执行以上过程的逆过程。

Worker

if blocks_to_swap_in:
    self.cache_engine.swap_in(blocks_to_swap_in)
    issued_cache_op = True
if blocks_to_swap_out:
    self.cache_engine.swap_out(blocks_to_swap_out)
    issued_cache_op = True
if blocks_to_copy:
    self.cache_engine.copy(blocks_to_copy)
    issued_cache_op = True

调度程序已经决定了应将哪些块交换进/出，应复制哪些块（blocks_to_copy），以及应转发哪些 SequenceGroup 到模型。但到目前为止，数据实际上并未在 GPU↔CPU 内存之间移动。这仅是哪些数据应移动的信息，使用 CacheEngine 组件即可实际移动数据。

# Prepare input tensors.
input_tokens, input_positions, input_metadata = self._prepare_inputs(
    seq_group_metadata_list)

# Execute the model.
output = self.model(
    input_ids=input_tokens,
    positions=input_positions,
    kv_caches=self.gpu_cache,
    input_metadata=input_metadata,
    cache_events=cache_events,
)

最后，在生成令牌之前，所有必需的基础作业都已完成！现在准备序列组的令牌 ID（_prepare_inputs）以生成并在模型中转发。

Prepare Inputs

模型输入正由_prepare_inputs函数转换。采用Continuous Batching进行分批处理，如上图所示。因此，最大限度高效地利用空白空间，节约内存，快速推理。有关更详细的说明，请参阅vLLM博文！

首先为提示令牌创建输入 ID。将属于 SequenceGroup 内的所有 Sequence 令牌扩展为 input_tokens 变量。然后将提示令牌的大小放入 input_positions 中的位置信息数组。此值充当 position_ids 的角色。最后生成每个令牌 ID 分配到的插槽的插槽映射表。(slot_mapping)

提示令牌仅进入从 Waiting 过渡为 Running 状态的 SequenceGroup。也就是说，仅将提示令牌放入尚未计算出 KV 缓存的 SequenceGroup 中，而对于已交换分组或以前处于 Running 状态（已计算出 KV 缓存）的分组不放入提示令牌。

接下来，为每个序列添加一个生成的令牌。令牌 ID 是提示中最末尾的令牌 ID，该令牌对应的块和时隙已经预先分配给调度器中所有处于运行状态的 SequenceGroup。

这样创建的 input_tokens、input_positions、slot_mapping 和 SequenceGroup 的块表以及其他元信息（提示长度、上下文长度）将输入到模型中。Attention

我们来了解下作为 Transformer 模型核心功能的 Attention（本文中仅介绍 GPT2 的多头 Attention）。在 vLLM 的实现中，主要根据上述结构对 Attention 进行了变更，其余部分（除并行相关的实现外）与上述结构相同。

# Reshape the query, key, and value tensors.
query = query.view(-1, self.num_heads, self.head_size)
key = key.view(-1, self.num_heads, self.head_size)
value = value.view(-1, self.num_heads, self.head_size)

# Pre-allocate the output tensor.
output = torch.empty_like(query)

# Compute the attention op for prompts.
num_prompt_tokens = input_metadata.num_prompt_tokens
if num_prompt_tokens > 0:
    self.set_attn_bias(input_metadata)
    self.multi_query_kv_attention(
        output[:num_prompt_tokens],
        query[:num_prompt_tokens],
        key[:num_prompt_tokens],
        value[:num_prompt_tokens],
        input_metadata,
    )

注意力层的输入接收以下三个张量：

• Query：形状为 [num_tokens, num_heads * head_size] 的张量

• Key：形状为 [num_tokens, num_heads * head_size] 的张量

• Value：形状为 [num_tokens, num_heads * head_size] 的张量

换句话说，按顺序排列所有标记，并且根据多头注意力，头数存在维度。将每个张量重新分割为 [num_tokens, num_heads, head_size] 的大小。

之后，为与提示部分相对应的标记（即除了最后 N（=序列数） 个之外的其余标记），对 Q、K、V 应用多头注意力。

也就是说，对于整个 QKV，不对没有高速缓存的 SequenceGroup 执行常规注意力计算。注意力使用 xformers 的内核。（xops.memory_efficient_attention_forward）

# 将键值存储到缓存中
cache_ops.reshape_and_cache(
    key[:num_valid_tokens],
    value[:num_valid_tokens],
    key_cache,
    value_cache,
    input_metadata.slot_mapping,
)

# Single Query Attention
self.single_query_cached_kv_attention(
    output[num_prompt_tokens:num_valid_tokens],
    query[num_prompt_tokens:num_valid_tokens], key_cache,
    value_cache, input_metadata)

在对 Key，Value 值进行注意力计算之后，将这些值存储在缓存中，然后对照缓存和生成标记（语境中的最后标记）应用单一查询注意力。

至此，本文介绍完vLLM生成的整体流程，直到模型执行到attention部分。后续attention得分析见：

手抓饼熊：vLLM源码之PagedAttention（持续更新）