我要投稿

蚂蚁集团：大模型推理显存优化的深度探索与实践

发布日期：2024-12-28 21:54:43 浏览次数： 1556 来源：DataFunSummit

导读赵军平老师团队主要专注于大模型的推理优化，同时也涉及异构算力的优化。本次分享的重点是我们团队从去年（特别是今年）在大模型推理方面的主要工作，具体集中在显存优化这一领域。本次分享题目为大模型推理-显存优化探索蚂蚁集团。

主要介绍：

1. 大模型推理显存挑战

2. 蚂蚁显存优化探索

3. 结语

4. Q&A

分享嘉宾｜赵军平蚂蚁集团技术总监

张锐蚂蚁集团高级研发工程师

编辑整理｜向隆

内容校对｜李瑶

出品社区｜DataFun

大模型推理显存挑战

在探讨大模型领域，我们不得不关注显存的需求问题。显存对于大模型来说至关重要，因为它直接关系到模型能否顺利运行。以 Llama 65B 模型为例，这是一个相对较大且复杂的模型。它采用了 MHA（多头注意力机制）来加载并运行首字模型，同时支持不同长度的上下文，例如在 100 长度时，需要 8K 的显存。

显存需求可以通过一个基本的公式来估算。这个公式涉及到模型的精度（如 FP16 或 BF16 等 16 位的精度）、乘数以及头的数量等因素。显存需求基本上是随着这些参数的增加而线性增长的。如果我们想要运行如此大规模的模型和长上下文，就需要估算所需的显存量。这通常意味着需要使用多张显卡，甚至可能需要跨多台机器来实现。

幸运的是，业界已经进行了许多优化工作，以减少显存的需求。例如，GQA 技术就是一种量化压缩技术，它通过减少模型运行时所需的显存量，使得在有限的硬件资源下也能运行大型模型。这些技术的发展对于推动大模型的应用和普及具有重要意义。

左边的图表展示了英伟达显卡显存容量的发展历程，其中特别提到了 H100 显卡。同时，红色线条表示的是大模型权重的增长情况。总体来看，显存的需求增长速度更快，这一点无需赘述。

接下来，我们将分析显存和访存带宽的增长情况与算力增长的对比。我们将使用 roofline 模型来进行分析，以 V100、A100 和最新的 B200 显卡为例。虽然这些数据难以获取，但我们关注的是趋势。实际上，算力的增长速度更快，而访存带宽虽然也在增长，尤其是旗舰型显卡，但仍然跟不上算力的增长速度。

这种算力增长意味着在很多情况下，尤其是在推理场景中，访存成为了一个关键因素。首先，我们需要确保模型能够被放入显存中；其次，模型的加载速度需要足够快，这些都是实际问题。特别是在线上场景中，由于需要提供良好的用户体验，通常采用较小的 batch size，以确保响应时间（RT）足够短。当然，在离线压测场景中，可能会使用较大的 batch size，这时显存的利用率可能会接近满载。

我们主要围绕显存开展了一些工作，希望与大家分享一些思路和解决方案。

蚂蚁显存优化探索

在 GPU 显存管理方面，我们有一些新的观察。传统上，开发者在 GPU 上进行显存分配时，通常使用 CUDA Malloc。然而，从 CUDA 10.2 版本开始，引入了一个新的 API——虚拟内存管理（Virtual Memory Management，简称 VMM）。这个 API 已经存在一段时间了，它的作用与传统的显存分配方式有所不同。

传统的 CUDA Malloc 是一次性分配显存，分配后程序员或框架无法进行更多的操作，只能直接使用。而 VMM 提供了一个分两步进行显存申请的机会。首先，它允许我们申请一个虚拟地址，这个地址可以非常大，但实际上在没有创建 handle 之前，这个地址是廉价的，几乎不需要成本。这个虚拟地址是连续的，但有一个最小单位，称为 chunk，规定为两兆。

接下来，我们可以通过另一个 API 申请一个 handle，这个 handle 可以简单理解为代表一个物理显存。然后，我们可以通过第三个 API 将虚拟地址和 handle 关联起来，进行映射。映射之后，程序才能正常访问显存。这样，虚拟地址是连续的，而物理地址可以不连续。

这种机制为我们提供了新的创新机会。例如，在典型的场景中，我们可以解决显存碎片问题。如果只有一层地址空间，我们很难进行有效的管理。有了 VMM API，我们可以进行创新，解决训练中的显存碎片问题，以及推理场景中的 KV 缓存碎片问题。

我们的工作就是基于这个 API 进行的。当然，仅仅有一个简单的 malloc 操作是不够的，就像我们不能仅用一个 malloc 就写出一个数据库或 Redis 一样。在实际应用中，我们还需要做很多工作，比如内存池（pooling）和缓存（caching）等。

为了说明虚拟内存管理（VMM）API的实际应用，让我们通过一个简单的例子来阐述。假设我们有一段虚拟地址空间，其中存放了一部分数据，而其余部分则是空白的，已经被标记为空闲并释放掉了，里面没有存放任何真实的数据。

现在，我们面临一个需求：我们能否将这些空白的物理显存回收掉，同时保留有数据的部分？这实际上是存储领域中经常研究的碎片整理或低碎片存储问题的一个典型例子。

通过使用 VMM API，我们可以达成这样的效果。具体来说，我们可以将数据最小对齐，保留斜线部分的数据，而对应的物理块（比如一两兆的范围）也会被保留。在这个过程中，我们不需要搬运任何数据，因此开销相对较小。而那些已经变成灰色的物理块，即没有数据的部分，可以被归还给系统。这样，即使最初申请的是一个连续的地址空间，现在我们也可以将其中没有数据的部分释放回系统，从而实现显存碎片整理的效果。

简而言之，通过 VMM API，我们能够在不移动数据的情况下，将有数据的部分保留，将没有数据的部分释放，有效地减少了显存碎片，提高了显存的利用率。

我们目前的工作主要围绕 Virtual Tensor 展开，其核心目标与 Page Attention 解决的问题完全一致，因此我们的工作是对标 Page Attention 的。

在 Page Attention 出现之前，由于碎片化问题，早期的 Fast Transformer 等方案难以有效解决这一挑战。虽然这些方案可以通过牺牲显存或吞吐量来提升实时性能（RT），但在吞吐量方面表现不佳。

vLLM 提出的 Page Attention 在这一领域产生了重大影响，效果显著。然而，我们也注意到，PagedAttention 后续的升级迭代较为缓慢。最初，Page Attention 仅支持 CUDA core，随后逐步扩展到 Tensor core，并引入了 Flash Attention 等技术。这一过程从 2023 年五六月份开始，经历了较长时间才逐步完善。

Page Attention 技术从最初支持 CUDA 核心到扩展至 Tensor 核心，并引入 Flash Attention 等技术，整个过程耗时较长，可能需要四到七个月。这段时间本足以训练出一个非常大的模型，但底层优化和推理优化的迭代进度却相对较慢，这是一个亟待解决的问题。

具体来说，在训练场景中，一个 kernel 可能包含前向和后向计算。业界存在多种算法，例如 Flash Attention 等，尽管这些算法在不断发展，但目前仍存在一些限制。例如，Flash Attention 可能尚未完全支持官方版本，8 比特（FP8、int8）以及各种稀疏注意力机制（如 Sage Attention、Seer Attention、Dou Attention、Re Attention 等）也在不断涌现。这些技术的影响力较大，例如 PyTorch 官方的 Flex Tensor 和百度的 Flash Mask 等，我们也在密切关注这些领域的进展。

尽管注意力机制在业界不断创新和引入，但很多情况下，这些技术首先关注的是训练效果，而在从训练到推理的迁移过程中，仍然存在较大的差距。其中一个重要原因是这些技术本身较为复杂，涉及多种 mask 机制。能否快速将这些技术迁移到 vLLM 上，仍然是一个挑战。

我们也在尝试进行量化和定制化优化，但发现修改 vLLM 的难度较大。例如，将一个 kernel 适配到 vLLM Attention 中，需要仔细调整 shared memory、寄存器、block table 等底层资源。这不仅要求开发人员具备算法和工程知识，还需要深入理解寄存器和 shared memory 等底层机制。最终，适配和调优的过程可能需要 2 到 3 周甚至更长时间，且有时调优后的效果并不理想。

我们希望解决这些问题，避免当前这种耗费大量时间和精力却难以调优的情况。同时，我们希望在管理显存和优化碎片方面达到与 vLLM 当前水平相当的性能。

另一个问题是至少在 vLLM 的早期版本中，实现上存在一些问题。最初的版本可能不支持 Tensor core，这限制了其性能表现。随后，在 GQA（可能是指一个特定的应用或框架）上的性能也不够理想。

为了解决这些问题，vLLM 的开发团队开始引入 Flash Attention 的一些 kernel，以期提升性能。然而，这个过程耗费了相当长的时间。从最初的不支持 Tensor core，到后来的性能问题，再到引入新的技术以改善性能，vLLM 的发展经历了一段漫长的旅程。

我们的工作目标可以概括为三个小目标。首先，我们希望能够像 vLLM 的 Page Attention 一样解决 KV 缓存的碎片问题，这是我们着手这项工作的初衷，也是我们的基础目标。其次，我们追求的性能目标是达到原生性能，例如 Flash Infer 的性能，这显然需要利用 Tensor core 来实现。第三，我们认为在现实场景中非常关键且影响迭代效率的一点是，各种 attention 算法可能会不断演进，我们希望能够在部署到推理阶段时，实现分钟级别的适配到 VMM 上，这样我们就不需要改动 kernel。这意味着在一开始编写好量化或 Spaas 的初始 kernel 后，不需要为了适应 vLLM 而进行调整。

我们的最终目标是，编写一次 kernel 后，就不需要区分训练或推理。我们希望能够将 attention kernel 计算和显存管理这两层解耦。这样做的最大好处是迭代会非常快。例如，官方可能还没有集成好 FlashAttention3 ，而我们可能只需要一个下午的调试就能将 Flash Attention 3 集成上去。包括各种前缀的稀疏 attention，如 PyTorch 官方的 Flex Attention，基本上都可以通过三行代码进行修改实现。

我们的团队发表了一篇论文，介绍了我们的工作，这项工作在今年六七月份对外发布。我们的目标是提供一个分化的 Tensor 结构，这个 Tensor 是基于 PyTorch 的。您可以将这个 Tensor 视为 PyTorch 的一个派生类，它专门用于管理 KV 缓存。所有的复杂细节都被封装在这个 PyTorch 对象中，用户无需关心这些细节。

假设您之前使用的是 Flash Attention，现在您想将 Flash Attention 3 集成到您的 Hopper 系列 GPU 上，不论其性能如何，您只需对 Tensor 进行简单的修改。在我们的 Tensor 内部，我们的库会管理 KV 缓存的虚拟地址分配和物理地址分配，这是最基本的功能。除此之外，我们还需要执行许多额外的工作来降低开销，包括进行内存池化（pooling）、预先分配（pre-allocation）、异步操作（async operations）等，例如使用 MMAP 重新映射内存。这些逻辑相对复杂，但我们的用户不需要关心这些内部细节。我们提供了一些简单的 API，例如允许用户重新分配（re-allocate）或扩展（extend）Tensor 的大小。

在 NVIDIA GPU 上进行调整是我们工作的重点，我们主要利用了 NVIDIA 的 VMM 接口。据我们了解，AMD 的显卡也有类似的接口，华为的显卡也具备类似的功能，但功能和产品可能会有细微差别。因此，我们的工作主要以 NVIDIA 的显卡为例。

vTensor 内部支持分配、释放和 map 操作，并在 prefetch cache 场景中集成了一些功能，它可以在这种 promote cache 中自动调用我们的 Tensor。具体细节可以参考我们的论文，论文中详细介绍了如何预先分配、异步 map 或 unmap 缓存类对象。

为了降低开销，我们进行了 pooling、预先分配，并采用了异步操作，例如 MMAP 和重新 map，这些逻辑相对复杂，但内部用户不需要关心这些细节。我们提供了一些简单的 API，例如可以进行 re-allocate 或扩展 Tensor 的大小。

让我们来看一下用户最终如何使用我们的系统。实际上，用户只需要引入我们提供的 Python 包。这个包中会包含我们自定义的一个对象，可以认为它是一个 Tensor 对象。用户可以使用这个对象来创建 K（键）和 V（值）对象，并且可以对这些对象进行扩展（extend）或者缩减（shrink）操作。

在调用 Attention kernel 时，原先需要传递 K 和 V 对象，现在只需替换为我们的 K 和 V 对象即可。这意味着，如果你原本有一个可以运行的 kernel，在 vLLM 中可以直接使用你的 kernel，无需进行任何修改。以前可能需要传递一个 block table 并进行查找等操作，但现在这些步骤完全不需要了。

virtualTensor 实现的效果达到了我们设定的三个小目标。首先，我们成功解决了 KV 缓存的碎片问题，实际上，我们的效果可以达到与 Page Attention 相似的水平。此外，virtualTensor 还支持按需增长，这意味着可以根据策略动态地增长 table cache，而不需要立即预分配一个很大的缓存空间，这个特性在其他场景中也非常有用。

第二点是关于我们工作的最大优势之一——快速迁移和高效迭代。

如果你有一些现有的 kernel，例如业界的 FlashAtten3、PyTorch 的 Flex Attention，或者 ICML 中提到的 Double-sparsity 等优化算法，这些通常在 Hugging Face 或 PyTorch 上实现。然而，这些算法在迁移到 vLLM 时，可能并不会被算法同学特别关注。

使用我们的方法，迁移过程非常简单。以下是两个示例：

只需导入相关包。
在 Attention kernel 中，替换两行代码的参数即可。

整个过程基本上是分钟级别的，仅需几行代码。这种快速迭代的效率非常高，能够显著提升开发和优化的速度。

我们自己进行了实验，使用我们的方法将 Flash Attention 3 集成到 vLLM 的 0.5.5 版本中，并测试了其性能。结果显示：

与直接使用 Flash Attention 3 的 kernel 相比，性能基本持平。
与 vLLM 自带的 Flash Attention 2 相比，在某些场景下，性能提升了两倍以上。

这种显著的性能提升主要归功于 Flash Attention 3 的功能优势。我们通过构建一个中间层（bridge），能够快速地将现有的 kernel 引入到 vLLM 中，而无需深入了解底层的硬件细节。这种无脑集成的方式极大地提升了日常迭代的效率，我们认为这在实际开发中非常重要。

以上是第一个工作的主要内容，其核心目标是：

消除显存碎片。
提高研发效率，快速适配各种 kernel。

希望这个思路能为大家提供参考。

第二个工作是关于优化首字生成的问题，我们将其命名为 LayerKV，主要在 vLLM 上进行了优化。

首字生成的耗时通常包含两部分，这两部分都需要优化：

Prefill 阶段：这是前向计算阶段，通常是计算密集型的，耗时较长。
排队阶段：排队通常是由于等待显存或等待前一个请求完成（例如前一个 prefill 阶段）。我们主要关注的是显存等待问题。

什么情况下会等待显存？通常是在压力较大且上下文较长的情况下，显存可能不足。即使启用了类似 chunk prefill 等功能，显存仍然可能不够用。在压力较大的场景下，这种等待的耗时占比会变得很长。

以 Llama 7B 为例，假设在单卡上运行一个 7B 模型，观察到红色部分的占比会随着上下文长度的增加而增加。此外，我们还发现负载排队时间占比也会逐渐增加。

由于显存问题导致的排队现象，主要原因如下：像 vLLM 以及其他一些引擎通常的做法是，当一个请求（例如 16K 的上下文）进来时，系统需要检查实时计算所需的 KV CACHE 总量。例如，假设总量为 60M ~ 64M，系统需要一次性确保显存足够，以保证在 prefill 过程中不会中断。如果显存不够，当前的实现通常会采取等待策略，直到前一个请求的 decode 阶段结束并释放显存后，才能继续处理当前请求。即使 prefill 阶段结束，显存也不会立即释放，而是需要等到足够多的请求完成 decode 并释放显存后，当前请求才能被排入队列。

上面这个示意图：在负载压力逐渐增加的情况下（例如每秒请求数从低到高），我们运行一个超大的模型（如 80B 或 70B 模型，使用多卡），观察到首次延迟和 P99 延迟在某个点之后会显著上升。这种情况通常是由于排队导致的。就像日常生活中的排队现象一样，虽然任务本身并不复杂，但排队时间却很长。

因此，我们希望重点优化这一排队阶段，以减少显存等待带来的延迟。

我们对这一问题进行了分析，实际情况是：

在大模型中（如 80 层模型），计算是逐层进行的。在 prefill 阶段，虽然 KV cache 显存需要一次性检查是否足够，但实际消耗是逐层进行的。这意味着在时间和空间上存在一个 trade-off 窗口。

我们的优化思路是：

分层申请显存：在 prefill 阶段，逐层申请 KV cache 显存，而不是一次性申请。
分层 offloading：在不同卡和不同模型长度下，推算出计算时间和传输时间，尝试将两者互相隐藏。

具体来说，我们可以计算在不同卡上（如 L20、H20）以及不同模型长度下，计算 KV cache 的时间和将其传输到 CPU 的时间。经过推算，这些时间基本上是可以互相隐藏的，尤其是在算力受限的场景下（如 L20、H20 的算力可能只有正常的一半甚至 1/3），总会有时间和空间窗口来完成以下操作：

计算完 KV cache 后，立即传输到 CPU。
异步地将 KV cache 从 CPU 拉回到 GPU。

这种优化方式可以有效隐藏计算和传输的时间，提升整体效率。

我们在显存管理和 Offloading 层面上做了一些优化工作，这些工作是在 vLLM 的基础上进行修改的。具体改动如下：

1. 新请求到达时的显存检查优化：

我们放松了对显存检查的要求。原本需要一次性检查整个请求所需的显存是否足够，现在只需检查部分显存（例如 4 层或 8 层）是否足够即可。
通过配置参数，系统可以在显存足够的情况下继续推进请求。
有两种机会来处理显存不足的情况：
一种是将显存 Offloading 到其他设备。
另一种是等待其他请求完成并释放显存后继续推进。

2. 显存分配优化：

我们采用了逐层分配显存的方式，而不是一次性分配。
在计算完一层的 KV cache 后，立即开启异步的 Uploading 操作，将数据传输回 GPU。

3. 调度优化：

为了平衡首次延迟和后续 token 生成的效率，我们设计了调度策略，动态调整显存分配和 Offloading 的优先级。
通过配置不同的策略，系统可以在首字生成和后续生成阶段之间进行动态 trade-off。

4. 与现有 vLLM 实现的对比：

当前 vLLM 的实现是请求级别的显存管理，即在 prefill 阶段，需要一次性为整个请求的上下文长度预留显存。
如果显存不足，vLLM 支持 Offloading 或重新计算（recompute）。
首字生成阶段仍然需要确保显存足够，这可能导致排队等待。

和 DistServe 等相关工作对比。我们的优化主要体现在 KV cache 管理层的细粒度控制上，包括显存的 reserve、分配、Offloading 以及调度策略的动态调整。总体而言，我们的优化使得显存管理更加灵活和动态，能够更好地平衡首字生成和后续 token 生成的显存消耗。

我们后续还有一篇论文工作，稍后会提供链接。我们测试了一些典型的模型，从小到大包括 7B、30B 和 70B 的 LLAMA 模型。对于小模型，我们使用单卡运行；对于大模型，我们使用多卡运行。测试中，我们调整了不同的上下文长度和负载压力，观察了优化前后的对比效果。

以下是部分测试结果的摘录（更夸张的结果未展示，供参考）：

以 L20 为例，我们测试了三个不同规模的模型：34B、7B 和 70B。基线结果用蓝色表示，优化后的结果用橙色表示。

1. ShareGPT 负载测试：

随着负载压力的增加，优化后的首次延迟（橙色）虽然也会增长，但增长非常缓慢。
而基线（蓝色）在负载增加到一定程度后，首次延迟会显著上升，优化后的效果比基线好十几倍。

2. 7B 模型测试：

在计算压力较大的情况下，优化后的效果仍然优于基线，大约提升了 8 倍。
在 QPS（每秒查询数）方面，优化后相比基线最多降低了 5%（70B 模型），而对于 7B 模型，QPS 仅降低了 1%。

总体来看，优化后的系统在首次延迟和 QPS 方面表现更好，尤其是在高负载和长上下文的情况下，优化效果显著。

实际下降情况通常会根据业务需求来评估。我们内部通常不会只看最右边的极端情况（即已经无法使用的场景），而是会约定一个合理的首次延迟指标。例如，对于 7B 模型，我们要求首次延迟小于 1 秒，同时生成速度需要满足人眼可读的体验，这会约束一个体验指标。

在这种实际场景下，我们主要关注的是：

QPS（吞吐量）：在首次延迟小于 1 秒且生成速度满足体验要求的情况下，优化后的 QPS 相比基线大约提升了 30%。
放宽条件：如果放宽一些条件（例如允许首次延迟稍长），QPS 的提升幅度会更大，具体提升多少取决于业务场景的需求。

从另一个角度来看，如果我们希望在不影响生成速度的情况下优化首次延迟，优化后的效果大约可以提升 1.4 倍。

因此，优化效果可以根据不同的业务场景和需求进行调整。你可以根据业务的具体诉求，选择合适的指标进行优化。

刚才的测试是基于 7B 模型的结果。如果我们改变输入和输出的长度，可以得到更多的测试结果。趋势如下：

1. 给定模型和硬件条件下：

如果生成的长度越长，优化后的加速效果会越好。
主要原因是，在负载较大的情况下，生成长度越长，基线系统需要等待前一个请求的 decode 阶段结束并释放显存后，才能开始执行新的请求。
这种等待会导致新的请求被阻塞，因为显存不足。

2. 测试场景：

我们测试了不同的输入输出长度组合，例如 1K 输入到 300 输出，以及 1K 输入到 1K 输出。
在不同的场景下，你可以根据业务需求约束首字延迟和生成速度，观察 QPS 的变化。
通常情况下，首字延迟会有非常明显的提升。

3. 极限情况：

在论文中，我们展示了一些极限情况下的测试结果，优化后的效果可能会有 60 多倍的提升。

结语

最后，快速总结一下我们工作的思路，供大家参考：

我们的优化主要集中在显存管理和 Offloading 层面，通过分层申请显存、动态调度以及异步传输等技术，显著提升了首次延迟和 QPS。
优化效果与输入输出长度、负载压力等因素密切相关，生成长度越长，优化效果越显著。
可以根据不同的业务场景和需求，灵活调整优化策略，以达到最佳性能。

我们的显存优化工作从训练场景开始，最早可以追溯到 2023 年的一篇相关工作。该工作主要优化了训练场景中的显存碎片问题，核心技术是使用了 CUDA 的 VMM API 接口。

在 2024 年初，我们开始重点关注推理场景的优化。推理场景的优化主要分为两个工作：

1. Virtual Tensor：

目标是消除显存碎片，提升适配各种 kernel 的效率。
最终效果是仅需 3 行代码即可适配 attention kernel，例如将 Flash Attention 3 集成到 vLLM 中，在 H100 上测试，性能提升了两倍。

2. LayerKV：

专注于优化首次生成阶段的显存管理，特别是 prefill 阶段的 KV Cache 管理。
通过显存 Offloading 等技术，在 QPS 持平的情况下，优化了平均延迟和 P99 延迟，极限情况下首次延迟提升了 60 多倍。

目前，我们进一步优化了 KV Cache 和其他显存开销的动态管理，目标是提高吞吐量。特别是在新型模型架构（如 GQA、Mamba 或多模态模型）上，优化效果显著。

我们已开源了部分工作：

GMLake：已开源。
Virtual Tensor：部分开源，代码仓库中有相关实现。
LayerKV：尚未开源，计划完成后整理报告并开源。

欢迎大家关注和交流。

我们也有一个微信群，但由于微信的时间限制，我先提供一个个人微信号，感兴趣的朋友可以先加我。

最后简单介绍一下我们团队的工作：

我们主要关注大模型的推理优化，重点包括：

语言模型：例如 LLAMA 和公司内部的百灵大模型，是我们重点优化的目标。
多模态模型：例如文生图、文生视频等。
搜索、推荐、广告类模型：这也是我们需要关注的领域。

大家如果有兴趣，欢迎交流，甚至可以一起探讨如何共建开源项目。

我的分享内容到此结束，谢谢大家！

Q&A

Q1：请问 virtualTensor 在训练中也有收益吗呃？

A1：针对训练场景，我们有一个单独的工作，名为 GMLake。你可以先快速了解一下这个工作。我们在实际测试中验证了它的收益。

其核心思路是：

使用虚拟显存和物理显存的动态管理和映射。
通过动态映射技术，优化显存的使用效率，从而带来显著的收益。

Q2：请问 vtensor 怎么解决碎片问题的？比如头尾各有个碎片，怎么合到一起？

A2：我再完整说明一下这个例子。假设数据在中间，前后都是空闲的显存空间。我们希望达到的效果是：将数据保留在中间，同时将前后的空白显存释放掉。

实现效果如下：

物理显存释放：
物理显存的头和尾部分被释放回系统，而中间有数据的部分仍然保留，用户可以通过指针继续访问数据。
我们没有移动数据，数据仍然保持在原来的位置。
实现过程：
假设当前虚拟显存和物理显存对应关系如下：虚拟显存对应四个物理块（四个 block）。
我们需要通过 mmap 操作将虚拟显存和物理显存解绑（unmap）。
解绑后，我们只保留中间有数据的物理块，而将前后的物理块通过 mmap 操作解绑。
解绑后，可以通过 release 操作将物理显存释放回系统，从而让硬件能够看到更多的可用物理显存。
后续操作：
释放后的物理显存已经归还给硬件，用户可以继续通过虚拟显存的方式分配显存。
数据仍然保留在中间的物理块中，不会被移动或影响。
虚拟地址连续性：
虚拟地址是连续的，用户可以一次性分配一个足够长的虚拟地址空间（例如48 比特，即 2^48），具体长度取决于硬件支持。

这种实现方式的优势在于：

碎片整理：通过动态映射和释放，实现了显存碎片的整理，而不会影响用户数据的访问。
数据不动：在实现过程中，数据不会被移动，保证了数据的完整性和一致性。

当然，如果需要，也可以通过将数据重新搬移到头尾位置，然后一次性释放的方式来实现，但这是一种不同的实现方式。

Q3：请问 LayerKV 有开源计划吗？

A3：我们的代码其实已经放在了 GitHub 的某个目录下，但由于我们目前正在撰写论文，暂时不太适合进行宣传。大家可以加入微信群，获取更多信息。

关于代码的情况：

当前放出的代码可能有些混乱，因为我们正在进行重构。
老代码是基于 vLLM 0.4.2 版本的，而我们现在最新的代码是基于 vLLM 0.5.5 版本的。
我们会尽快放出新的代码，请大家稍作等待。

关于 LayerKV 在不同硬件上的效果：

我们主要测试了 L20 硬件，而在 A100 和 H 卡上的效果可能会有所不同。
具体效果可以参考我们的论文，论文中可能包含了 A100 和 L40 的测试结果。
虽然我们没有在 H 卡上进行大量测试，但 A100 和 L40 上仍然会有一定的优化效果。

以上就是本次分享的内容，谢谢大家。

分享嘉宾

INTRODUCTION

赵军平 蚂蚁集团 技术总监

张锐 蚂蚁集团高级研发工程师

赵军平，蚂蚁集团技术总监，负责大模型异构算力和推理。CCF HPC 和存储专委委员， ~200 中美技术专利。"数据密集型应用系统设计"译者。

往期推荐

人工智能在汽车制造上的落地应用探讨

实时语音交互的游戏队友——网易伏羲 AI Agent 创新应用

货拉拉技术应用：LaLaEval大模型应用评测框架及实践

1688AI 助手“源宝”的产品演变之路

中小银行大数据应用建设实践

抖音集团指标管理与消费体系建设实践

下一代 RAG：tidb.ai 使用知识图谱增强 RAG 能力

百度飞桨：多模态大模型技术进展与产业应用实践

数据指标实战：从0到1构建工业数据指标体系！

阿里云出海数据合规挑战与解决方案

点个在看你最好看

SPRING HAS ARRIVED

53AI，企业落地应用大模型首选服务商

产品：大模型应用平台+智能体定制开发+落地咨询服务

承诺：先做场景POC验证，看到效果再签署服务协议。零风险落地应用大模型，已交付160+中大型企业

160+中大型企业正在使用53AI

立即咨询预约演示

百度智能云邀53AI：共创AI新纪元，启航智能新时代

2024-05-27

钉钉恒星计划：53AI与百余位企业家及钉钉生态伙伴，共议“AI 浪潮下的新机遇”

2024-05-22

热点资讯

太爽了！o1 现在可以天天享用了

2024-09-18

FP8 低精度训练：Transformer Engine 简析

2024-07-11

万字综述 10+ 种 LLM 投机采样推理加速方案

2024-07-11

微调神器LLaMA-Factory官方保姆级教程来了，从环境搭建到模型训练评估全覆盖

2024-07-09

通义千问超越GPT4了？

2024-06-11

深度剖析：字节跳动大模型训练被实习生“投毒”事件

2024-10-20

彻底理解系列之：FP32、FP16、TF32、BF16、混合精度

2024-07-26

Qwen2大模型微调入门实战-命名实体识别（NER）任务

2024-07-23

基于 Qwen2 大模型微调技术详细教程（LoRA 参数高效微调和 SwanLab 可视化监控）

2024-07-20

DifySandbox 的构建背景和实现机制

2024-07-12

大家都在问

AI模型训练到底在训练什么？

2024-12-26

微软变脸OpenAI，模型价值之争压不住了？

2024-12-25

大模型时代的软件工程教育，路在何方？

2024-12-24

字节ReFT技术：OpenAI强化微调的幕后功臣？

2024-12-09

OpenAI最新推出的RFT是什么？

2024-12-07

LLM Quantization ：理解 LLM 中模型大小？

2024-12-06

企业部署生成式 AI，五个最大的挑战是什么？

2024-12-05

如何生成Function Calling微调数据？

2024-12-02

开箱即用的企业大模型应用平台

工作+AI

业务+AI

AIx业务

大模型咨询

大模型定制

相关资讯