vTensor: Flexible Virtual Tensor Management for Efficient LLM Serving

发表时间: 2024-07 · arXiv:2407.15309 (preprint)

原文: https://arxiv.org/abs/2407.15309

Jiale Xu, Rui Zhang, Cong Guo, Weiming Hu, Zihan Liu, Feiyang Wu, Yu Feng, Shixuan Sun, Changxu Shao, Yuhong Guo, Junping Zhao, Ke Zhang, Minyi Guo, Jingwen Leng
Shanghai Jiao Tong University, Shanghai Qi Zhi Institute, Ant Group

速读

一句话结论 本文提出了基于 GPU 虚拟内存管理的 vTensor 抽象，将大语言模型推理中的 KV Cache 内存碎片整理与计算 Kernel 解耦，在释放 71.25% 显存的同时实现了端到端平均 1.86 倍的吞吐量提升。

要解决什么问题 大语言模型推理时为了避免重复计算会缓存历史的 Key 和 Value（KV Cache），但由于不同请求的序列长度不同，传统的连续内存分配会导致严重的显存碎片化。现有的主流系统（如 vLLM）引入了 PagedAttention 机制，通过页表来实现内存碎片整理。然而，这种做法将页表管理和计算 Kernel 深度耦合，带来了两个致命卡点。首先是计算效率受限：由于页表地址转换逻辑复杂，GPU 的 Tensor Core 无法直接支持，导致 vLLM 只能依赖计算能力较弱的 CUDA Core 来执行 PagedAttention。在面对多查询注意力（MQA）或分组查询注意力（GQA）等计算访存比极高的场景时，算力瓶颈尤为明显。其次是显存灵活性极差：PagedAttention 需要在初始化时静态预分配几乎所有的 GPU 显存来构建页表，这些显存被独占后即使处于空闲状态也无法动态释放给激活值或其他模型实例使用。此外，这种耦合设计使得开发新特性的工程成本极高。

怎么做的 核心思路是将内存碎片整理逻辑从 GPU 计算 Kernel 中剥离，交由 CPU 异步处理，从而让 GPU 能够使用原生、高度优化的 Kernel（如 FlashAttention）在 Tensor Core 上满血运行。作者基于 GPU 底层的虚拟内存管理（VMM）API，设计了 vTensor 这一全新的张量抽象。对于计算 Kernel 而言，vTensor 表现为一个标准的、地址连续的 CUDA 张量指针，但其底层物理内存是按需映射的离散块。关键设计由 vTensor Manager（VTM）和 FlexInfer 调度器构成。VTM 运行在 CPU 上，包含三个核心数据结构：用于管理连续虚拟地址空间的 vSet、用于管理 2MB 物理内存块句柄的 pSet，以及用于多轮对话前缀匹配的基数树 rTree。其核心不变量可以表示为预留的虚拟地址空间大小与实际分配的物理块总容量之间的解耦关系：

$$L_{virtual} \ge N_{physical} \times S_{chunk}$$

这意味着系统可以为新请求预留最大序列长度的虚拟地址空间 $L_{virtual}$（例如 4096 个 Token），但初始只分配极少量的物理块 $N_{physical}$。在实际运行中，FlexInfer 调度器采用 CPU-GPU 异构调度策略。当模型处于自回归解码阶段需要生成新 Token 时，CPU 会提前触发 Extend 操作，调用底层 API 分配新的物理块并将其映射到已有的虚拟地址上。由于物理内存的分配和映射完全在 CPU 端异步完成，这一过程与 GPU 上的矩阵计算完美重叠，彻底隐藏了内存操作的延迟。同时，得益于基数树 rTree 的设计，当遇到共享前缀的请求时，系统只需分配新的虚拟地址并将其映射到已有的物理块句柄上，无需发生任何显存拷贝。

效果如何 实验在单机 8 卡 NVIDIA A100 (80GB) 硬件上进行，测试了 Yi-6B、Yi-9B 和 Yi-34B 模型。对比基线包括代表静态分页内存路线的 vLLM（PagedAttention）、代表前缀预填充优化路线的 SGLang（Triton Kernel），以及不带内存优化的原生 FlashAttention。在 Kernel 层面的微观测试中，vTensor 在解码阶段相比 vLLM 的 PagedAttention 最高实现了 3.27 倍的加速；在前缀预填充阶段，相比 SGLang 的 Triton Kernel 最高实现了 3.92 倍的加速。在端到端服务场景（包含单次生成、多轮对话和前缀共享）下，搭载 vTensor 的 FlexInfer 系统相比 vLLM 实现了平均 1.86 倍的吞吐量提升，在多轮对话场景中最高达到 2.42 倍。在显存代价方面，由于打破了静态预分配机制，vTensor 在 A100 上相比 vLLM 平均释放了 71.25%（约 57GB）的空闲显存，使得系统能够支持更高密度的任务混部。作者承认的局限在于，为了维护页表和支持动态扩展，系统在大 Batch Size（如 64）时会产生约 4.99% 的显存保留开销，但这些开销在请求结束后会被立即释放。

主要贡献

大型语言模型（LLMs）在各个领域得到广泛应用，每天处理数百万个请求。这种需求的激增给优化吞吐量和延迟同时控制成本带来了重大挑战。键值（KV）缓存作为保留先前计算的标准方法，使得LLM推理受到严重的内存限制。虽然批处理策略可以提高性能，但它们经常导致严重的内存碎片。尽管像vLLM这样最先进的系统使用分页注意力（paged Attention）机制来缓解KV缓存碎片，但由于页面管理和计算内核紧密耦合，它们仍然受到低效的内存和计算操作的困扰。

本研究介绍了一种基于GPU虚拟内存管理（VMM）的创新张量结构vTensor，用于LLM推理。vTensor通过将计算与内存碎片整理分离并提供动态可扩展性，解决了现有的局限性。我们的框架采用CPU-GPU异构方法，确保高效、无碎片的内存管理，同时适应不同LLM架构的各种计算内核。实验结果表明，vTensor在不同模型上实现了平均$1.86\times$的加速，在多轮对话场景中最高可达$2.42\times$。此外，vTensor在内核评估中提供了平均$2.12\times$和$3.15\times$的加速，与SGLang Triton前缀预填充内核和vLLM paged Attention内核相比，最高分别达到$3.92\times$和$3.27\times$。此外，与vLLM相比，它在NVIDIA A100 GPU上释放了约$71.25\%$（57GB）的内存，从而支持更多内存密集型工作负载。

图1：三种KV缓存内存管理策略：(a) 具有原生GPU分配的原生KV缓存存在大量碎片；(b) vLLM采用分页内存管理以消除大部分碎片，但内核紧密耦合；(c) vTensor解耦了计算和内存分配，管理更加灵活。
图1：三种KV缓存内存管理策略：(a) 具有原生GPU分配的原生KV缓存存在大量碎片；(b) vLLM采用分页内存管理以消除大部分碎片，但内核紧密耦合；(c) vTensor解耦了计算和内存分配，管理更加灵活。

背景知识与设计原则

LLM推理

生成式LLM的自回归模式。生成式大型语言模型以自回归模式运行，按顺序生成token。LLM通常执行两个基本操作：线性投影操作和注意力操作。线性投影操作首先使用各自的权重$W_Q$、$W_K$和$W_V$计算查询（$Q$）、键（$K$）和值（$V$）张量。随后，注意力机制计算$Q$和$K$之间的相似度，应用此相似度对相应的$V$值进行加权以计算注意力输出。最后，被称为前馈网络的线性投影操作处理注意力输出。

推理过程的两个阶段。面向用户的生成式LLM推理过程通常包括两个不同的阶段：预填充（prefill）阶段和解码（decode）阶段。在预填充阶段，模型接收一个提示序列来为后续文本生成设定上下文，并产生初始token。此外，预填充阶段会缓存$K$和$V$张量，以防止在随后的解码阶段进行冗余的重新计算。在预填充阶段之后，解码阶段使用最后一个token和KV缓存迭代地生成下一个token并更新KV缓存。换句话说，KV缓存在推理过程结束之前会不断扩展。

LLM服务优化

连续批处理策略。连续批处理是一种动态策略，它在一个批次中的请求完成之后立即用新请求替换它。它应用迭代级别的调度，通过确定每次迭代的批次大小，而不是等待所有请求完成。这种方法允许在另一个请求完成时立即插入新请求，从而显著提高了GPU利用率。

前缀缓存优化。在典型的基于LLM的应用中，系统提示在各种请求中通常保持不变。在多轮对话中，每条新消息都建立在所有先前消息的上下文之上。这两种情况都可能导致对相同提示重复计算键值（KV）缓存，这在计算上成本高昂且耗时。因此，先前的研究提出了一种称为前缀缓存的优化技术【55, Automatic prefix caching + vLLM + URL】，它存储共享提示前缀的KV缓存以消除冗余计算并减少初始token生成延迟。当新请求到达时，LLM可以绕过对缓存部分重新计算KV缓存，从而减少初始token计算时间。这对于具有长且重复提示或频繁交互的应用（如聊天机器人或虚拟助手）特别有利。然而，前缀缓存增加了GPU内存使用，并在将前缀缓存优化与高效计算内核结合时带来了挑战。

GQA和MQA优化。注意力操作中的KV缓存是受内存限制场景下LLM推理的一个关键问题。多头注意力（MHA）中每个查询头对应一个不同的键和值头，导致显著的KV缓存内存消耗。为了解决这个问题，开发了几种KV缓存优化技术。在多查询注意力（MQA）中，所有查询头共享一个键和值头。分组查询注意力（GQA）将查询头分成$G$组，其中每组共享一个单独的键头和值头。

LLM服务系统

现有系统的局限性。LLM服务系统如vLLM、TensorRT-LLM、LMDeploy和TGI，是支持LLM服务过程的框架。现有的LLM服务系统集成了包括上述优化和分页KV缓存（paged KV-cache）在内的一系列优化，以提高LLM推理的效率。然而，大多数这些优化都与内存相关，以更多的内存换取更少的计算，或者减少KV缓存的内存占用以提高请求批次大小并增加吞吐量。很少有系统关注LLM服务的计算效率。一个明显的原因是现有的分页KV缓存机制将内存管理与计算内核耦合在一起。这种耦合需要大量努力才能将现有的高效实现迁移到支持分页KV缓存的系统中，换句话说，缺乏灵活性。

GPU虚拟内存管理

CUDA底层虚拟内存管理。由于应用程序越来越需要以极低的延迟和高效率管理内存，CUDA引入了底层虚拟内存管理（VMM）【11, Cuda driver api documentation + 2024 + NVIDIA + URL】【12, Introducing lowlevel gpu virtual memory management + 2024 + NVIDIA + URL】【29, Gmlake: Efficient and transparent gpu memory defragmentation for large-scale dnn training with virtual memory stitching + 2024 + arXiv】。VMM提供了如保留（reserve）和映射（map）等原生操作来操作虚拟地址空间，提供了比传统方法（如cudaMalloc）更细的粒度。新的API，包括cuMemCreate、cuMemAddressReserve、cuMemMap和cuMemSetAccess，使得创建更高效的动态数据结构成为可能，并提供了对GPU内存使用的更好控制。VMM减少了内部碎片并消除了昂贵的内存操作。

内存灵活性（Motivation）

现有方法的内存占用分析。为了进一步探索现有的局限性，我们对未进行内存优化的FlashAttention（Native）和vLLM进行了初步的内存占用分析。我们将它们与我们的设计FlexInfer进行了比较。如图2所示，原生方法存在严重的碎片问题，这些碎片随着批次大小（BS）的增加而增加，并在BS为64时导致内存不足（OOM）错误。vLLM试图通过将KV碎片转换为保留内存（图2中的黄色条）来缓解这种情况。vLLM需要静态预分配所有可用的GPU设备内存，以构建能够寻址所有KV token的页表。一旦页表建立，预分配的内存就专门为KV缓存保留，不能用于其他分配，包括激活和额外的LLM推理实例。虽然vLLM解决了单个LLM服务实例内的碎片问题，但这种内存限制仍然严重限制了LLM服务系统的整体功能。

FlexInfer的内存灵活性。我们的解决方案FlexInfer引入了两个关键的灵活性。第一个灵活性是可用内存。FlexInfer可以动态释放未使用的内存，在不同批次大小下平均释放57GB的内存，这占NVIDIA A100 GPU上$80GB$总内存的$71.25\%$，如图2所示。第二个灵活性是动态扩展。虽然我们的方法需要一些保留开销（当批次大小BS为64时约为$4.99\%$），但这是存储页表和支持动态扩展所必需的。此外，此保留会在请求结束时立即释放。

图2：在具有$80GB$内存的GPU A100上，使用FlashAttention（Native）、vLLM和FlexInfer的GPU内存使用细分。
图2：在具有$80GB$内存的GPU A100上，使用FlashAttention（Native）、vLLM和FlexInfer的GPU内存使用细分。

计算灵活性（Motivation）

现有系统的性能问题。原生系统受到碎片和低效的困扰，由于KV缓存碎片而经历了巨大的开销。vLLM具有更好的整体性能，并且现在是具有KV缓存碎片整理的LLM服务的事实标准。然而，它的计算效率仍然很低。分页注意力（paged Attention）会产生与地址转换相关的CUDA内核开销，这些开销无法重叠或优化。此外，vLLM更深层次的原因是它严重限制了编程模型。vLLM目前的实现仅依赖于CUDA核心，这极大地限制了其计算能力。具体来说，vLLM只能在解码阶段使用CUDA核心运行分页注意力，由于GPU张量核心（tensor cores）不支持复杂的页表转换，导致效率较低。

Roofline模型对比。图3显示了比较vLLM和FlexInfer的Roofline模型。vLLM的性能随着批次大小的增加而增加。vLLM和FlexInfer在算术强度为0.99的MHA中表现出相似的性能，这是一个极度受内存限制的应用，仅使用CUDA核心就足以满足计算要求。然而，从MHA到共享头数减少的GQA再到MQA，算术强度不断增加，这意味着它们具有更高的计算/内存比和更高的计算能力要求。因此，从GQA-H16（具有16个共享头）到GQA-H2和MQA，由于基于CUDA内核的分页机制，vLLM几乎保持相同的性能，最终被限制在3.6 TFLOP/s。对于具有GQA或MQA的LLM，现有的分页KV缓存管理要么使用低效的CUDA核心，要么需要大量努力来优化内核，但效率仍然很低。相比之下，FlexInfer的性能随着算术强度的增加而提高。FlexInfer可以轻松地与高效的CUDA内核集成，基于GPU张量核心的计算和CPU的内存管理显著提高了性能。对于具有MQA的LLM，FlexInfer达到27.3 TFLOP/s，比vLLM高出$7.58\times$。

FlexInfer的计算优势。现有系统在计算效率方面表现出两个主要缺点。首先，虽然vLLM在某些场景下由于原始注意力内核的算术强度较低而表现出色，但当应用新的内存高效优化时，这一优势成为显着的瓶颈。其次，在实现新的LLM功能时，vLLM中手工制作的内核会产生巨大的开发成本。我们的解决方案FlexInfer有效地解决了这些限制。首先，解耦机制确保我们可以避免以前的性能问题。其次，基于VMM的KV内存管理及其服务支持显著提高了计算和内存的灵活性。

图3：GPU A100上LLM注意力的Roofline模型。
图3：GPU A100上LLM注意力的Roofline模型。

方法细节

4 Overview

FlexInfer Scheduler。正如前一节所述，内存碎片整理对于在LLM服务中实现高性能变得越来越关键。然而，以前的大多数工作很少关注与内存相关的操作，经常将它们与计算内核耦合，导致严重的性能损失。为了解决这个问题，我们提出了FlexInfer，这是一个CPU和GPU异构框架，它将大多数内存操作解耦并卸载到CPU，并通过重叠GPU计算来隐藏它们。与以前在GPU上运行内存操作（例如，页表映射和转换）的工作相比，CPU更擅长与内存相关的动作。当请求发送到FlexInfer时，它根据请求的配置（如批次大小和序列长度）解耦内存和计算操作。对于计算，FlexInfer调度器采用原始的高度优化的内核在GPU张量核心上运行LLM，保持高计算灵活性和效率，而不受算术强度限制。对于内存管理，FlexInfer专注于调度LLM服务不同阶段中与内存相关的行为。具体来说，FlexInfer提供了定制的调度计划，以在开始、预填充、解码和结束阶段重叠并隐藏内存分配和释放与计算。该策略显著减轻了LLM服务系统中内存操作的开销。

图4：FlexInfer服务框架概述。
图4：FlexInfer服务框架概述。

vTensor Manager。vTensor Manager (VTM) 是用于内存碎片整理的核心设计。它包含三个组件：vTensor Pool (VTP)、vTensor Operation (VTO) 和 vTensor Scheduler (VTS)。vTensor Scheduler是vTensor Manager中的关键组件。VTS接收来自FlexInfer调度器的指令，然后基于元信息为每个指令创建特定的策略。这些元信息包括多个数据结构（例如，哈希表）来记录和跟踪与内存相关的细节，例如未分配的物理块（PCs）和vTensor状态。基于这些策略，VTS调度特定操作以通过vTensor Operation (VTO)分配或释放内存。VTO充当调度策略和CUDA底层VMM API之间的“翻译器”。它包括所有vTensor内存相关动作，将策略翻译为原始API，并在GPU上异步执行它们。执行结果然后返回到vTensor Pool (VTP)，VTP存储所有虚拟张量，包括物理块和虚拟内存之间的虚拟内存地址映射信息。VTP还更新VTS中的元信息，然后为CUDA内核生成vTensor或释放vTensor。

5 vTensor and Management

5.1 vTensor。vTensor是为CUDA内核设计的虚拟内存的高效抽象。从CUDA内核的角度来看，vTensor是一个指向数组的简单张量指针，它与标准的CUDA分配的指针相同。然而，vTensor比其简单的抽象要复杂得多。如图5所示，vTensor指针由vTensor Manager (VTM) 生成。当向vTensor Scheduler (VTS) 发送请求时，VTS将创建分配策略，然后让vTensor Operation (VTO) 分配虚拟内存和相应的物理块。vTensor指针*A指向一个GPU虚拟地址（VA），该地址必须是连续的，以兼容标准的CUDA分配张量。在虚拟地址之下，GPU虚拟内存管理（VMM）维护由VTM在CUDA内核分派到GPU之前注册的物理块（PCs）的完整映射信息。VMM允许GPU张量核心通过虚拟地址访问GPU全局内存中所需的数据。请注意，物理块分配在GPU内存中，但物理块句柄（PHs）和虚拟内存可以由CPU访问和管理。句柄和虚拟内存仅具有物理块的索引信息，而无需设备-主机内存传输。具体来说，我们采用每个大小为2MB的物理块并将其存储在GPU全局内存中。其对应的句柄大小只有几个字节，由vTensor Pool (VTP) 记录并存储在CPU的主内存中。vTensor的一个关键特征是物理块、虚拟地址和请求之间灵活的映射关系。vTensor映射有一些基本特征：（1）一个请求可以放置在多个非连续的物理块中（例如，图5中显示的PC 9、PC 6和PC 8），但虚拟地址必须是连续的；（2）一个物理块可以被多个虚拟地址引用；（3）虚拟地址可能具有比实际物理块（例如，6）更大的内存容量（例如，8），这意味着该虚拟地址的某些部分未映射到物理块。因此，VTM可以为新请求设置具有最大序列长度的vTensor虚拟内存地址，但不分配真实的物理块，自然地消除了KV缓存中的碎片。

图5：vTensor的设计。
图5：vTensor的设计。

5.2 vTensor Pool。vTensor Pool包含用于有效管理KV缓存vTensor的基本数据结构。有两种主要类型的数据结构：有序集合和基数树（即前缀树）。如图5所示，vTensor的两种类型的信息返回给VTP，即一个虚拟地址和对应的一组物理句柄。由于vTensor灵活的映射机制，我们将虚拟地址和物理句柄分别存储在两个有序集合中，分别记为vSet和pSet。我们还构建了一个基数树rTree，以解决多轮对话场景。
- vSet：vSet由一个排序集合组成，用于存储vTensor的虚拟地址。虚拟地址包含一个记录所有页面对应物理句柄的页表。因此，虚拟地址可以被视为vTensor。如图5所示，我们可以超额分配虚拟地址大小，但只将几个页面映射到真实的物理句柄。在实践中，当请求传入时，我们分配一个具有最大序列长度（例如，4096个tokens）的虚拟地址，但只分配几个物理块（例如，256个tokens）。
- pSet：pSet也基于排序集合来存储物理句柄，它与物理块有严格的一对一对应关系。pSet维护每个物理句柄的必要属性以记录其状态，例如活动状态和它映射到的vTensor。因为多个vTensor可以引用同一个物理句柄，所以我们像Linux系统中的“硬链接”一样管理物理块。我们为每个物理句柄分配一个引用计数器。一旦在GPU上分配了物理块，其相应的物理句柄就会插入到pSet中。当引用计数器为0时，相应的物理句柄将从pSet中删除。
- rTree：rTree表示一种基数（前缀）树数据结构，将vTensor作为树节点存储，并促进不同LLM服务请求之间的前缀匹配。一旦前一个请求完成并生成了一个完整的vTensor，这个vTensor将成为下一个请求的前缀KV缓存。然后，vSet中vTensor的指针将根据请求的前缀模式插入到rTree中。当下一个请求需要其前缀KV缓存时，rTree可以有效地搜索前缀vTensor。rTree是LLM多轮对话场景中前缀匹配功能的基础设施。

5.3.1 Allocation (分配)。GPU虚拟内存管理（VMM）API是vTensor和GPU之间的基本接口。VMM分配发生在三个关键的原始操作中。首先，地址保留（cuMemReserve）：这一初始步骤涉及确定所需的内存大小并确保适当的虚拟地址空间。接下来，物理创建（cuMemCreate）：在这里，系统在GPU物理内存中生成实际的数据段。最后，映射（cuMemMap）：此阶段将物理数据段链接到保留的虚拟地址，允许张量有效地访问内存。这三个步骤（保留、创建和映射）构成了分配原始虚拟内存块的综合机制。VTO通过操作调用这些原始API以支持各种KV缓存管理。
- pAlloc(N)：物理分配pAlloc(N)首先将检查pSet以找到未存储tokens的可用物理块。假设pSet中剩余的可用块为$M$，且小于$N$。在这种情况下，pAlloc将继续利用cuMemCreate API从GPU全局内存中分配新的$N - M$个物理块，并将返回的句柄插入pSet。在本研究中，我们使用$2MB$的块大小。这是唯一能增加GPU内存使用的vTensor操作。物理分配是独立于虚拟分配的基本内存操作。
- vAlloc(S)：虚拟分配函数vAlloc(S)同样检查vSet以寻找可用的虚拟地址空间，否则通过cuMemAddressReserve API分配并保留具有S字节虚拟地址空间的虚拟地址。在LLM任务中，我们分配虚拟地址，其中$S = $最大序列长度，并且所有虚拟地址共享相同的长度。这是一个轻量级操作，仅为vTensor保留虚拟地址空间，不分配任何物理块。虚拟地址将被记录在vSet中。虚拟分配不分配新的物理内存。在由CUDA VMM API引起的大量开销下，解耦虚拟和物理分配设计正是高效KV缓存碎片整理的关键设计。
- Map(VA, PCs)：在分配了虚拟地址（VA）和物理块（PCs）之后，我们需要将虚拟地址映射到物理块的一个子集，并使用cuMemMap API使虚拟地址可供真实数据访问。映射函数是确保我们的设计对内存管理灵活的关键。

5.3.2 Deallocation (释放)。释放操作是对应于分配的逆操作。未映射的物理块和虚拟地址的状态将被修改，使它们可供将来重用。调用Unmap()操作以从物理块中取消映射虚拟地址，vFree()和pFree()用于释放GPU上的资源。释放模块引入了延迟释放（lazy deallocation）机制，而不是主动释放vTensor和物理块，以减少内存操作开销。在实践中，我们通常使用取消映射操作，而不是释放它们在GPU上的资源，直到服务任务结束。

5.3.3 Tree Operations (树操作)。树操作用于操作前缀树。rPush(vTensor)操作用于将vTensor指针插入rTree以实现前缀匹配功能。在我们的rTree设计中，我们将每个vTensor存储为一个不同的节点。首先，rTree将检查是否有任何节点与正在插入的vTensor共享公共前缀模式。如果有，我们将继续遍历其子节点，直到没有节点匹配剩余的前缀，此时vTensor将作为当前节点的子节点插入。如果没有这样的节点，那么这个vTensor将作为rTree中新基数树的根插入。rPrefixMatch(vTensor)操作用于在rTree中搜索与vTensor前缀模式匹配的节点。rTree的遍历过程类似于rPush操作。rPrefixMatch将返回存储的vTensor指针，并仍然为后续请求保留此前缀记录。

5.4 vTensor Scheduler。本节介绍了vTensor Scheduler (VTS)，它利用vTensor Pool和vTensor Operation来实现高效的内存调度。vTensor Scheduler被设计为接收来自FLEXINFER调度器的内存指令并执行vTensor操作。图6通过一个例子显示了VTS的调度策略。vTensor调度器有五个基本操作，包括创建、扩展、释放以及两个前缀相关操作：前缀记录和前缀匹配。
- Create (创建)：Create用于初始化新请求。它将采用pAlloc和vAlloc来分配所需的物理块和虚拟内存地址。然后，VTS将使用Map操作映射虚拟内存的物理块，组合成一个完整的vTensor，该vTensor将返回给CUDA内核。
- Extend (扩展)：如图6 (2) 所示，Extend操作是FlexInfer LLM服务的基础。VTS在解码期间动态扩展序列长度以迭代生成新tokens。VTS使用pAlloc分配一个新的物理块，并将其映射到在Create操作中分配的虚拟地址。这种动态扩展策略对vTensor的设计至关重要。该过程涉及两个关键优势。首先，KV缓存中新物理块的扩展是在与现有块计算之前抢占式分派的。这种异步方法允许内存分配与计算重叠，从而最大限度地减少开销。值得注意的是，pAlloc可能会重用pSet中可用的物理块，从而可能避免新的GPU分配。其次，vTensor始终保持紧凑的GPU内存使用。
- Prefix Record and Match (前缀记录和匹配)：前缀记录和匹配操作支持前缀KV缓存，减少了多轮对话中的重新计算。如图6(3)所示，当对话被标记为继续时，VTS将vTensor记录在rTree中作为前缀候选。图6(4)演示了当后续请求需要前缀时，VTS如何使用基数树有效地检索vTensor。在获得前缀vTensor (VA-0) 后，VTS分配一个新的虚拟地址 (VA-1) 并复制VA-0的页表和元数据。重要的是，物理块不会被重新分配；相反，它们被映射到新虚拟地址 (VA-1) 的现有块。这种方法显著减少了内存分配开销和使用，从而有效支持前缀KV缓存。随后，共享前缀的请求转换为常规请求，利用Extend机制生成额外的tokens，如图6(5)所示。
- Release (释放)：从根本上说，我们的VTM设计保留了GPU资源，直到LLM服务结束或调用了显式的清空内存操作。这种方法允许VTM保留分配的资源，最大限度地减少内存开销。然而，仍然需要一个Release策略将虚拟地址和物理块返回到其可用状态，使其他请求能够使用它们，如图6(6)所示。

图6：vTensor调度示例。
图6：vTensor调度示例。

6 FlexInfer Scheduler

FlexInfer请求调度策略。FlexInfer通过调度高度优化的内核而不受内存寻址限制，实现了在GPU张量核心上的高效LLM执行。FlexInfer调度器将不同的请求类型转换为不同的内存指令，然后由VTM处理。算法1总结了FlexInfer的请求调度策略，该策略异步分派计算和内存动作。对于新请求，FlexInfer调用VTM中的Create或Prefix Match操作来初始化请求内存并构建vTensor。随后，它将vTensor指针提供给计算内核，并将其异步分派到GPU张量核心。在解码期间，它动态扩展vTensor的物理块以满足内存需求。如果资源不足，调度器抢占优先级最低的请求并重新分配其资源。调度器采用一系列vTensor调度策略与VTM交互，简化了调度逻辑。在程序终止时，调度器释放所有vTensor和相关资源。这种异步调度确保扩展操作先于计算内核，有效地重叠并隐藏了LLM服务期间的内存开销。

算法1：FlexInfer调度策略
算法1：FlexInfer调度策略

实验环境

实验结果

解码内核评估 (Decoding Kernel Evaluation)。在图7左侧，我们在KV缓存序列长度固定为16K的情况下，评估了不同批次大小的不同注意力内核。FlexInfer内核与vLLM PagedAttention相比平均加速$2.78\times$，与原生FlashAttention相比性能为$1.02\times$。当批次大小达到32时，加速达到峰值，比PagedAttention快$3.08\times$。在图7中间，我们在批次大小为16的情况下对KV缓存的不同序列长度进行了内核评估。FlexInfer内核与PageAttention相比平均性能提高$2.67\times$，与Paged FlashAttention相比提高$1.59\times$。当序列长度等于1024时出现峰值加速，比PageAttention高$3.27\times$。在图7右侧，我们探讨了具有不同KV头数的解码内核的性能。FlexInfer内核的加速随着KV头数的减少而增加，从KV头数为32时的MHA的$1.03\times$增加到KV头数为1时的MQA的$8.0\times$。

图7：从不同角度对解码内核的评估。
图7：从不同角度对解码内核的评估。

前缀预填充内核评估 (Prefix-prefilling Kernel Evaluation)。在图8左侧，我们在序列长度固定为16K且前缀比例为0.5的情况下，评估了具有不同批次大小的前缀预填充内核。FlexInfer的前缀预填充内核比SGLang提出并被vLLM采用的triton内核实现了$3.40\times$的加速。当批次大小为16时达到峰值，与triton内核相比为$3.49\times$。在图8右侧，随着前缀与提示比例的下降，加速从$2.9\times$上升到$3.92\times$。

图8：对前缀预填充内核实现的评估。
图8：对前缀预填充内核实现的评估。

端到端单次生成场景 (End-to-End Single-generation Scenario)。我们使用张量并行度分别为1、2和4的Yi-6B-200K、Yi-9B-32K和Yi-34B-32K模型评估了单次生成场景。如图9所示，随着批次大小的增长，3个模型的吞吐量平均分别提高了$1.8\times$、$1.3\times$和$1.4\times$。当批次大小为64时，这些模型达到峰值吞吐量，与vLLM相比，性能分别为$2.02\times$、$1.5\times$和$1.53\times$。

图9：端到端单次生成中对不同模型和并行度的计算灵活性评估。
图9：端到端单次生成中对不同模型和并行度的计算灵活性评估。

端到端前缀缓存场景 (End-to-End Prefix-caching Scenarios)。如图10右侧所示，当计算需求较低时，前缀缓存为FlexInfer和vLLM都实现了加速。然而，当批次大小增长时，FlexInfer的吞吐量比vLLM高出高达$2.06\times$。如图10左侧的多轮对话场景所示，FlexInfer的吞吐量最高可比vLLM高出$2.42\times$。

图10：使用Yi-6B进行的端到端前缀缓存计算灵活性评估。
图10：使用Yi-6B进行的端到端前缀缓存计算灵活性评估。

FlexInfer的内存效率 (Memory Efficiency of FlexInfer)。图11a描述了运行批次大小与FlexInfer和vLLM内存使用之间的关系。当批次大小很小时，与以前的工作相比，FlexInfer几乎可以节省所有的KV缓存GPU内存。图11b表明，现有的解决方案（如vLLM）静态划分GPU内存用于KV缓存内存管理，导致请求率波动时GPU内存浪费。FlexInfer提供了内存预留的灵活性。

图11：Yi-6B模型的内存灵活性评估。(b) 显示了不同请求率下的内存轨迹。
图11：Yi-6B模型的内存灵活性评估。(b) 显示了不同请求率下的内存轨迹。

结论

通过引入vTensor抽象并将内存管理与计算解耦，FlexInfer解决了现有LLM服务系统中的关键低效问题。FlexInfer的计算灵活性促使内核性能比LLM服务系统中广泛采用的内核实现提高了最高$3.92\times$。在端到端实验中，我们的框架在各种模型和应用中实现了最高$2.4\times$和平均$1.86\times$的加速。此外，FlexInfer平均释放$71.25\%$（57GB）GPU内存的能力为内存密集型任务开辟了新的可能性。这些在计算效率和内存利用率方面的改进标志着LLM部署向前迈出了关键一步。

补充细节

相关工作 (Related Works)。虚拟内存管理：长期以来，机器学习系统一直考虑使用虚拟内存管理来提高内存利用率并实现统一的内存管理。NVIDIA的统一虚拟内存（UVM）技术允许应用程序分配CPU或GPU上执行的代码都可以读写的内存。vDNN提出了一种运行时内存管理解决方案，使DNN工作负载的内存使用虚拟化。在LLM时代，GMLake利用CUDA底层虚拟内存管理来管理大规模模型微调工作负载的内存并减少内存碎片。对于服务场景，vAttention是一项并发工作，侧重于使用类似的CUDA底层VMM API动态管理KV缓存并促进高效的LLM服务。