LLM推理框架之vLLM

大概分为以下几个部分:

  • 介绍vllm的特性
  • 代码分析
  • 服务化

PagedAttention:如何解决 GPU 显存瓶颈

在 self-attention 中,计算速度比内存速度快得多,因此进程(操作)越来越多地受到内存(HBM)访问的瓶颈。

研究发现,在 vLLM 库中 LLM 服务的性能受到内存瓶颈的影响。在自回归 decoder 中,所有输入到 LLM 的 token 会产生注意力 key 和 value 的张量,这些张量保存在 GPU 显存中以生成下一个 token。这些缓存 key 和 value 的张量通常被称为 KV cache,其具有以下特点:

  • 显存占用大:在 LLaMA-13B 中,缓存单个序列最多需要 1.7GB 显存;
  • 动态变化:KV 缓存的大小取决于序列长度,这是高度可变和不可预测的。因此,这对有效管理 KV cache 挑战较大。该研究发现,由于碎片化和过度保留,现有系统浪费了 60% - 80% 的显存。

为了解决这个问题,提出了PagedAttention,这是一种受操作系统中虚拟内存和分页经典思想启发的注意力算法。与传统的注意力算法不同,PagedAttention 允许在非连续的内存空间中存储连续的key 和value 。具体而言,PagedAttention 将每个序列的 KV cache 划分为块,每个块包含固定数量tokens 的keys 和 values。在注意力计算过程中,PagedAttention 内核能够高效地识别并提取这些块。

PagedAttention: KV Cache are partitioned into blocks. Blocks do not need to be contiguous in memory space.

由于blocks不需要在内存中连续,我们可以像在操作系统的虚拟内存中那样以更灵活的方式管理 keys 和values:可以把blocks看作pages,tokens看作bytes,sequences看作processes。序列的连续逻辑块通过块表映射到非连续的物理块。随着新tokens的生成,物理块会按需分配。

Example generation process for a request with PagedAttention.

在分页注意力中,内存浪费只发生在序列的最后一个块。实际上,这导致了接近最优的内存使用,仅浪费不到4%。这种内存效率的提高非常有利:它使系统能够将更多的序列批量处理,提高GPU利用率,从而显著提高吞吐量.

分页注意力还有另一个关键优势:有效的内存共享。例如,在并行采样中,从同一提示生成多个输出序列。在这种情况下,提示的计算和内存可以在输出序列之间共享。

并行采样示例

PagedAttention 通过其 block table 自然地实现了内存共享。与进程共享物理页类似,分页注意力中的不同序列可以通过将其逻辑块映射到同一物理块来共享块。为了确保安全共享,分页注意力跟踪物理块的引用计数,并实现 Copy-on-Write mechanism。

对于请求生成多个输出的示例生成过程

PageAttention的内存共享大大减少了复杂采样算法的内存开销,如parallel sampling and beam search,减少了高达55%的内存使用。这可以转化为高达2.2倍的吞吐量提升。这使得这种采样方法在LLM服务中变得实用。

与triton的结合

参考这篇文章:https://github.com/vllm-project/vllm/issues/541

优缺点

优点:

  • 吞吐量大,多batch下真正有效果的应该还是continuous batching。单batch下,除非输出序列特别长,就是提前申请max_len的kv cache显存空间,推理速度也没有比huggingface快多少
  • batch=1的时候不如hf
  • 其采样参数支持的还不是很完备(与huggingface相比),参数对齐后跑humaneval指标比hf低一些

代码细节

Engine


离线推理(LLM)和在线推理(AsyncLLMEngine)的Engine不一样,但是核心都是使用的LLMEngine:

class LLMEngine:
    """An LLM engine that receives requests and generates texts.

    This is the main class for the vLLM engine. It receives requests
    from clients and generates texts from the LLM. It includes a tokenizer, a
    language model (possibly distributed across multiple GPUs), and GPU memory
    space allocated for intermediate states (aka KV cache). This class utilizes
    iteration-level scheduling and efficient memory management to maximize the
    serving throughput.

    The `LLM` class wraps this class for offline batched inference and the
    `AsyncLLMEngine` class wraps this class for online serving.

    NOTE: The config arguments are derived from the `EngineArgs` class. For the
    comprehensive list of arguments, see `EngineArgs`.
    """

设置终止符的地方

CompletionOutput-vllm

# vllm/engine/llm_engine.py
    def _stop_sequences(self, seq_groups: List[SequenceGroup]) -> None:
        """Stop the finished sequences."""
        for seq_group in seq_groups:
            sampling_params = seq_group.sampling_params
            for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
                # Check if the sequence has generated a stop string.
                stopped = False
                for stop_str in sampling_params.stop:
                    if seq.output_text.endswith(stop_str):
                        # Truncate the output text so that the stop string is
                        # not included in the output.
                        seq.output_text = seq.output_text[:-len(stop_str)]
                        self.scheduler.free_seq(
                            seq, SequenceStatus.FINISHED_STOPPED)
                        stopped = True
                        break
                if stopped:
                    continue

                # Check if the sequence has reached max_seq_len.
                if seq.get_len() > self.scheduler_config.max_model_len:
                    self.scheduler.free_seq(
                        seq, SequenceStatus.FINISHED_LENGTH_CAPPED)
                    continue
                # Check if the sequence has reached max_tokens.
                if seq.get_output_len() == sampling_params.max_tokens:
                    self.scheduler.free_seq(
                        seq, SequenceStatus.FINISHED_LENGTH_CAPPED)
                    continue
                # Check if the sequence has generated the EOS token.
                if not sampling_params.ignore_eos:
                    if seq.get_last_token_id() == self.tokenizer.eos_token_id:
                        self.scheduler.free_seq(
                            seq, SequenceStatus.FINISHED_STOPPED)
                        continue

请求细节

batch

vllm/benchmarks/benchmark_throughput.py测试的是离线推理,参数--num-prompts指的是有多少个请求被传到engine中处理,根据数量for循环,add_request到engine中进行schedule。

目前不确定这个和batch的关系

目前缺陷

参考