vLLM源码之框架执行

主要观点总结

本文介绍了vLLM推理引擎的框架执行流程（v0.1.2），包括其核心技术如PagedAttention和连续批处理技术，并进行了代码级别的分析。vLLM使用PagedAttention技术提升句子生成速度，并包含多个用于实际服务的要素，例如使用Ray Cluster实现多集群环境中的稳定服务，以及利用Megatron LM的并行性处理大型模型和数据。文章详细描述了vLLM的整体架构、使用方式、组件关系、LLMEngine、Worker和Scheduler的执行流程，并解释了PagedAttention的分配和交换策略，以及单查询注意力等关键部分。

关键观点总结

关键观点1: vLLM的框架执行流程

vLLM利用PagedAttention技术提升句子生成速度，并包含多个用于实际服务的要素，如使用Ray Cluster和Megatron LM的并行性。

关键观点2: vLLM的架构与组件

vLLM的架构包括LLMEngine、Worker、Scheduler等组件，用于分布式处理、管理PagedAttention区块和KV Cache，并通过调度器改变生成顺序。

关键观点3: 使用方式

vLLM支持在线和离线两种方式，通过LLM类初始化模型、创建kv cache，并使用generate()函数生成句子。

关键观点4: PagedAttention技术

PagedAttention技术用于提升内存效率，当GPU内存不足时，以交换到CPU内存的方式，稳定管理中间计算过程。

关键观点5: 注意力计算

vLLM对Key和Value值进行注意力计算后，将其存储在缓存中，并对照缓存和生成标记应用单查询注意力。

正文

请到「今天看啥」查看全文

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）乘以每个元素包含的参数数量

请到「今天看啥」查看全文