LIA: 结合协作式AMX CPU-GPU计算和CXL卸载的单GPU LLM推理加速

作者/机构:
- Hyungyo Kim, Nachuan Wang, Qirong Xia, Jinghan Huang, Nam Sung Kim (伊利诺伊大学厄巴纳-香槟分校)
- Amir Yazdanbakhsh (Google DeepMind)

A1 主要贡献

核心问题:
大型语言模型（LLM）的参数规模持续增长，导致在单块GPU上进行推理变得不可行，因为模型参数和中间值（如KV缓存和激活值）的存储需求超过了高端GPU（如H100）的显存容量。现有解决方案，如多GPU部署，成本高昂；而将模型参数卸载到CPU内存并通过PCIe传输的系统级卸载方法，则受到PCIe带宽的限制，导致推理延迟显著增加。此外，以往的CPU-GPU协作框架因CPU计算吞吐量远低于GPU而效果有限，并且未能根据不同的批量大小（batch size）和序列长度动态调整计算卸载策略，导致资源利用不佳。

研究目标:
本文旨在解决单GPU在LLM推理中的内存瓶颈和性能问题，提出一个名为LIA（LLM Inference Acceleration）的框架，该框架利用英特尔CPU的先进矩阵扩展（AMX）技术和CXL内存技术，在小批量和高吞吐量两种场景下加速单GPU LLM推理。

创新点:
1. 对AMX矩阵乘法性能的全面分析：本文对集成了AMX（内置矩阵乘法加速器）的英特尔第四代（Sapphire Rapids, SPR）和第六代（Granite Rapids, GNR）至强可扩展处理器进行了微基准测试。测试结果显示，搭载AMX的SPR和GNR处理器分别能提供高达4.5倍和9倍于AVX512的GEMM（通用矩阵乘法）吞吐量，其性能分别达到A100 GPU的11%和22%，H100 GPU的5%和10%。在GEMV（通用矩阵向量乘法）方面，其吞吐量分别达到A100的38%和44%，H100的35%和41%。这些发现证明了配备AMX的CPU在计算卸载方面的巨大潜力，远超以往的CPU卸载框架。

  1. 基于AMX驱动的CPU-GPU协同LLM推理加速：基于上述分析，LIA框架被设计用于系统性地协调CPU和GPU的协同工作。LIA包含两个主要部分：

    • (C1) 算法前端: 该前端根据批量大小、输入序列长度、各子层的计算/访存比、CPU与GPU间的数据传输量以及两者的计算和内存带宽，系统性地决策应将哪些子层卸载到CPU计算。这使得LIA能够根据不同负载动态选择最优卸载策略，而不仅限于卸载计算最不密集的子层。
    • (C2) 执行后端: 该后端扩展了Intel Extension for PyTorch (IPEX)，使其能够支持AMX CPU与NVIDIA GPU之间的无缝协作推理。后端还引入了优化措施，以进一步提高GPU显存和CPU/GPU计算资源的利用率。实验表明，在SPR-H100（GNR-H100）系统上，LIA相比最新的单GPU卸载框架，延迟降低了4.0–7.0倍（10–19倍），吞吐量提高了1.2–3.7倍（1.6–5.1倍）。

  2. 利用CXL进行吞吐量驱动LLM推理的内存卸载策略：为满足高吞吐量应用中大批量推理对内存容量的巨大需求（例如，OPT-175B在批量为256时需1.6TB内存），LIA提出了一种使用低成本CXL内存（由退役服务器的DDR4模块构成）来扩展内存容量的策略。该策略将所有模型参数存储在CXL内存中，同时将中间值保留在DDR内存中，以减轻CXL内存相对较高的延迟和较低的带宽对性能的影响。这种方法既避免了CPU和CXL内存之间的低延迟PCIe传输开销，又能在给定批量大小下将DDR内存使用量减少高达43%，或者在DDR内存容量限制内支持更大的批量（最多增加1.76倍），从而将吞吐量额外提升高达1.45倍。


图 1: OPT模型架构及推理数据流。热力图展示了OPT-175B在输入令牌长度L为512、批量大小B为180时，预填充（prefill）和解码（decoding）阶段各子层的“操作数/字节”比。

A3 背景知识

2.1 LLM 推理

LLM模型结构与推理流程。一个典型的LLM由嵌入/编码层、N个解码器层和一个语言模型（LM）头（线性层和softmax）组成。本文重点关注解码器层，因为它们在推理时间和内存消耗中占主导地位。一个解码器层包含多个子层，主要由矩阵乘法和其他操作（如层归一化、残差连接和softmax）构成。所有N个解码器层结构相同，但各自拥有独立的参数。当模型接收到一个输入令牌序列时，会通过所有层生成一个输出令牌，这个过程称为预填充（prefill）或Sum阶段。在此阶段，与输入序列对应的K和V矩阵会被计算并作为KV缓存存储在内存中。随后，模型将生成的输出令牌作为新的输入，再次通过所有层来产生下一个输出令牌，这个过程称为解码（decoding）或Gen阶段，并递归进行直到生成完整的输出序列。图1展示了OPT-175B的网络架构，并突出了预填充阶段和第一个解码阶段。

KV缓存与注意力计算。解码阶段的各层可以对单个输入令牌或一批输入令牌进行计算，但注意力得分（attention scoring）子层除外。在计算注意力得分（$Q \times K^T$ 和 $Attention \times V$ 子层）时，来自KV映射子层的KV向量会与KV缓存拼接。KV缓存首先在预填充阶段生成，并在后续输入令牌的注意力计算中被重复使用。然而，KV缓存在批处理中是不共享的，因为批次内的每个输入令牌都有其独特的KV缓存。

操作密度（Operations/Byte）的差异。根据推理阶段（预填充或解码）、输入令牌长度（L）和批量大小（B）的不同，一个层内各子层之间以及不同层中相同子层的“操作数/字节”比（operations/byte）差异巨大。例如，对于OPT-175B模型，当L=512且B=180时，这个比值范围可以从1到50,000，如图1中的热力图所示。通常，解码器层中的softmax、层归一化和残差连接子层由于其操作密度较低，常与相邻的子层进行融合。

2.2 英特尔先进矩阵扩展 (AMX)

AMX架构与功能。为了显著提升CPU在机器学习应用中的吞吐量，英特尔从2023年发布的第四代至强CPU（代号Sapphire Rapids, SPR）开始，集成了AMX，一个片上矩阵乘法加速器，并提供了相应的指令集架构（ISA）支持。如图2所示，AMX加速器架构包含两个核心组件：(1) 一个二维寄存器阵列（称为tiles）和 (2) 一个瓦片矩阵乘法单元（Tile matrix multiply unit, TMUL）【索引6，Intel® Architecture Instruction Set Extensions and Future Features】。这些组件设计用于支持INT8和BF16数据格式。Tiles用于存储矩阵的子块，而TMUL作为一个二维的乘加单元阵列，对这些tiles进行操作。TMUL的二维结构通过大规模基于tile的矩阵计算，实现了比一维计算引擎（如AVX引擎）更高的每周期操作数和并行度。CPU通过分派AMX指令（如tile加载/存储和加速器命令）来控制多周期的AMX单元。该加速器的内存访问与CPU的内存访问保持一致性。


图 2: 英特尔先进矩阵扩展 (AMX)。

2.3 CXL 内存

CXL作为内存扩展方案。当前DRAM及其（并行）DDR接口技术在满足数据中心服务器对容量和带宽需求方面变得越来越昂贵，并受到CPU封装和PCB板的物理限制。为了应对这一挑战，CXL（Compute Express Link）【索引42，Compute Express Link (CXL): Enabling Heterogeneous Data-Centric Computing With Heterogeneous Memory Hierarchy, 2022, IEEE Micro】已成为一种有前途的内存接口，能够以高成本效益的方式扩展内存系统的容量和带宽，作为传统DDR接口的补充。例如，基于（串行）PCIe接口构建的CXL与DDR5 DIMM插槽相比，所需的引脚数减少了3倍【索引48，Demystifying CXL Memory with Genuine CXL-Ready Systems and Devices, 2023, MICRO】。此外，CXL协议可以将CXL内存暴露为远程NUMA节点中的内存，同时将内存技术与CPU特定的内存接口解耦。因此，它可以将任何过去几代的DDR DRAM（例如DDR4）连接到CPU，即使CPU本身不再支持该DDR版本。尽管如此，与DDR相比，CXL会使内存访问延迟增加140-170纳秒，并且在多模块系统中，其补充的带宽最多只能达到DDR内存所提供带宽的50%【索引48，Demystifying CXL Memory with Genuine CXL-Ready Systems and Devices, 2023, MICRO】。

A3 卸载框架在LLM推理中的性能瓶颈

最新的框架使得在GPU显存不足以容纳整个模型时，也能使用单个GPU进行LLM推理【索引13，DeepSpeed-Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale, 2022, SC; 索引22，Accelerate: Training and inference at scale made simple, efficient and adaptable, 2022, GitHub; 索引43，FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU, 2023, ICML】。具体来说，内存卸载（memory-offloading）将所有参数和中间值存储在CPU内存甚至存储设备中，然后通过PCIe将计算给定层所需的参数和中间值传输到GPU显存。除了内存卸载，还提出了计算卸载（compute-offloading），将注意力得分（attention scoring）子层的计算交给CPU，以减少CPU-GPU的数据传输。本节使用最新的LLM推理卸载框架FlexGen【索引43，FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU, 2023, ICML】，在一个SPR-A100系统上运行OPT-175B模型，并分析其性能瓶颈。模型参数、KV缓存和激活值的数据传输时间使用NVIDIA Nsight Systems【索引1，NVIDIA Nsight Systems, 2025】和NVTX标记【索引2，NVIDIANVIDIA Nsight VSE Documentation, 2025】进行测量，而CPU注意力计算延迟则使用Python时间模块测量。在B=1时，所有参数都卸载到CPU内存，而KV缓存和激活值存储在GPU显存中。在B=32时，由于KV缓存和激活值随B增大而需要更多内存，它们也被卸载到CPU内存。

洞见-1: 内存卸载因缓慢的CPU-GPU传输导致推理延迟过长。
洞见-2: 计算卸载虽避免了KV缓存的CPU-GPU传输，但当CPU计算时间超过CPU-GPU传输时间时，反而会增加推理延迟。

3.1 内存卸载：缓慢的数据传输

CPU-GPU传输是主要瓶颈。如图3所示，当批量大小B为1时，对于短序列长度L，CPU-GPU的参数传输占用了预填充和解码阶段计算延迟的98%以上。对于长序列L，预填充阶段由CPU-GPU传输贡献的延迟百分比下降到87%，这是因为预填充阶段的计算量与L成比例增加，而传输的参数量不变。然而，解码阶段的计算量几乎与L无关（对于相同的B），因此该百分比保持不变。

大批量下的内存与传输开销。当B增加到32时，KV缓存和激活值的大小会随B×L增加，可能需要高达145GB的内存。因此，它们也必须存储在CPU内存中并按需传输到GPU显存。同时，虽然计算量也随B×L增加，但其增长速度快于CPU-GPU传输量。例如，当L从64增加到1024时，计算量增加了16.2倍，而CPU-GPU传输量仅增加了1.8倍。因此，预填充阶段计算延迟中由CPU-GPU传输贡献的百分比随L显著下降。与预填充阶段相反，解码阶段计算延迟中由CPU-GPU传输贡献的百分比仍然保持在80%以上，无论L如何变化。这是因为解码阶段的计算量与B成比例，而不是B×L。


图 3: OPT-175B预填充和解码阶段的延迟构成，包括参数、KV缓存和激活值的CPU-GPU传输时间，以及相应的传输数据量。

3.2 计算卸载：受限的CPU计算吞吐量

计算卸载的权衡。KV缓存首先在预填充阶段生成，然后在多个解码阶段与新生成的KV向量持续拼接。为了避免在解码阶段进行KV缓存的CPU-GPU传输，FlexGen提出将使用KV缓存的注意力得分（attention scoring）子层的计算卸载到CPU。如图4所示，在B=32，L=1024时，将注意力得分子层计算卸载到CPU，解码阶段的延迟仅降低了10.2%。这有两个原因：(1) 参数的CPU-GPU传输仍然占延迟的很大一部分；(2) CPU计算注意力得分子层的延迟相当可观，与CPU-GPU传输KV缓存的延迟相比并不占优（例如，图4中CPU计算耗时1毫秒，而传输耗时0.4毫秒）。需要注意的是，使用AVX512的SPR处理器的实测最大矩阵乘法吞吐量不到A100的1%。这导致CPU计算注意力得分子层耗时很长，当L较短时（如图4中的L=64和128），甚至会增加推理延迟。


图 4: 在B=32时，CPU计算卸载子层的延迟与传输KV缓存到GPU的延迟对比，以及通过计算卸载实现的延迟降低。

A3 AMX的矩阵乘法吞吐量

本节使用源自OPT-175B的微基准测试来展示AMX的矩阵乘法吞吐量。吞吐量在SPR（40核）和GNR（128核）CPU上进行测量，并与SPR上的AVX512以及四代Nvidia GPU（P100、V100、A100和H100）进行比较。评估使用GEMM和GEMV微基准测试来反映LLM推理中的实际工作负载。GEMM基准测试模拟预填充阶段计算密集的FC1子层，涉及维度为$(B \times L, d_{model})$和$(d_{model}, 4 \times d_{model})$的矩阵乘法，其中$d_{model}$是模型维度（例如，OPT-175B为12,288）。GEMV基准测试则捕捉解码阶段内存密集的$Q \times K^T$子层，涉及维度为$(B \times n_h, 1, d_h)$的向量和$(B \times n_h, d_h, L)$的矩阵的分批乘法，其中$n_h$和$d_h$分别代表头的数量和头维度。基准测试在不同范围的$B \times L$（对于GEMM）和$L/d_h$（对于GEMV）值下进行，使用各架构支持的半精度浮点格式：AMX、A100和H100使用BF16，AVX512、P100和V100使用FP16。选择FC1和$Q \times K^T$子层是因为它们分别代表了OPT-175B所有子层中操作数/字节比的两个极端——FC1是计算最密集的，而$Q \times K^T$是内存最密集的。

洞见-3: AMX提供的吞吐量比AVX高出4.4倍（理论上可达8倍），并与前几代GPU性能相当，这使得计算卸载成为比仅支持AVX的CPU远为吸引人的选项。

4.1 GEMM 吞吐量：计算密集型

AMX的计算性能。图5（左）显示了不同架构的GEMM吞吐量。总体性能排名从高到低为：H100, A100, V100, GNR-AMX, SPR-AMX, P100, AVX512。SPR-AMX的理论最大吞吐量为90.1 TFLOPS，分别比AVX512和P100高8倍和4.7倍。在评估的$B \times L$范围内，SPR-AMX的实测最大吞吐量分别比AVX512和P100高4.5倍和2.4倍。SPR-AMX的峰值性能利用率较低，是因为新引入的AMX库相比成熟的AVX和P100库优化程度较低。与近期的GPU相比，在$B \times L = 64–36,684$的范围内，SPR-AMX的吞吐量分别达到V100的14-28%，A100的7-15%，以及H100的4-11%。SPR-AMX的吞吐量可扩展性受其40核数量的限制，因为AMX性能与核心数成正比。GNR-AMX通过128个核心解决了这个问题，其吞吐量比SPR-AMX高出2.4倍。一个双插槽的GNR系统可进一步将吞吐量提高1.8倍，分别达到V100、A100和H100吞吐量的68%、30%和16%。尽管原始吞吐量较低，但AMX巨大的内存容量使其能够通过消除缓慢的CPU-GPU传输来提供具有竞争力的端到端性能，特别是在超出GPU显存容量的LLM推理任务中。


图 5: AVX512, SPR-AMX, GNR-AMX, P100, V100, A100, 和 H100 的GEMM和（分批）GEMV吞吐量。

4.2 GEMV 吞吐量：内存密集型

AMX的访存性能。图5（右）展示了不同架构的GEMV吞吐量，性能排名略有变化：H100, A100, V100, P100, GNR-AMX, SPR-AMX, AVX512。排名的变化归因于这些系统中计算与内存带宽比率的不同。SPR-AMX拥有8个DDR5-4800通道，其峰值吞吐量为199 GFLOPS，与AVX512几乎相同，差异不到10%。这种相似性是由于GEMV工作负载是内存密集型的，并且AMX和AVX512共享相同的内存带宽（SPR系统上为260 GB/s）。SPR-AMX的GEMV吞吐量分别达到了P100、V100、A100和H100的54%、31%、19%和15%，这与其相对内存带宽的41%、34%、20%和15%是一致的。在较小的矩阵和向量尺寸下（即较小的$L$和/或$d_h$），SPR-AMX相对于GPU实现了更高的吞吐量，分别达到P100、V100、A100和H100吞吐量的70%、57%、38%和35%。这种提升是由于GPU内核调用的开销减少，这在小计算量时变得显著。最后，GNR-AMX通过利用其12个DDR5-5600通道，将可用内存带宽有效翻倍，从而比SPR-AMX的GEMV吞吐量提高了70%。

A2 方法细节

最新的计算卸载框架FlexGen仅将注意力得分（attention scoring）子层卸载到CPU。由于旧一代CPU（SPR之前）上AVX的吞吐量有限，这种经验性的选择只能带来适度的延迟降低。利用现代CPU中AMX的高吞吐量，我们提出了LIA，一个用于CPU-GPU协同LLM推理加速的框架，它能有效利用CPU的计算资源。该框架无缝集成了Intel IPEX【索引7，Intel® Extension for PyTorch】。图6提供了LIA的概览。该框架假设CPU在其大内存中存储所有模型参数和中间值，而GPU则像朴素的数据卸载框架一样逐层计算模型。LIA构建了一个优化问题，以系统地确定应将哪些子层卸载到CPU，从而在给定的批量大小（B）和输入序列长度（L）下最小化推理延迟。该优化主要关注解码器层，因为它们主导了LLM的推理延迟和内存使用。


图 6: 提出的用于LLM推理的CPU-GPU协同计算框架。

5.1 计算卸载算法

优化问题定义。LIA优化的核心在于确定一个卸载策略，该策略由一个向量p = (p_1, p_2, ..., p_6)表示，其中每个元素$p_i \in \{0, 1\}$指示解码器层的第i个子层（如图6所示，共六个子层）是在CPU上计算（$p_i=1$）还是在GPU上计算（$p_i=0$）。最优的计算卸载策略旨在最小化LLM推理的延迟（或最大化吞吐量），表示为：


其中$L(\mathbf{p})$表示对应卸载向量的单个解码器层计算延迟，而$\mathbf{p}_{opt}$是最优卸载向量。

延迟模型构成。单个解码器层的延迟被分解为每个子层的三个组成部分：


其中$L_{i,load}(\mathbf{p})$, $L_{i,compute}(\mathbf{p})$, 和 $L_{i,store}(\mathbf{p})$分别代表将两个操作数矩阵加载到给定计算设备、执行子层计算以及将中间结果存储到CPU的延迟。

数据加载延迟。首先，子层i的数据加载延迟表示如下：


其中$L_{i,load,act}(\mathbf{p})$, $L_{i,load,p/kv}(\mathbf{p})$, 和 $L_{i,load,res}(\mathbf{p})$分别考虑了激活值（也称为隐藏状态）、模型参数或KV矩阵（或缓存），以及用于残差操作的前一个子层的激活值的数据移动。

数据移动量计算。$L_{i,load,act}(\mathbf{p})$, $L_{i,load,p/kv}(\mathbf{p})$, 和 $L_{i,load,res}(\mathbf{p})$的计算方式如下：
表 1: OPT模型解码器层中GEMM/GEMV子层的数据大小和计算量（BF16数据类型）。





其中$S_i^1$和$S_i^2$分别表示子层i的第一个和第二个矩阵操作数的数据大小，这些值总结在表1中，$B_{PCIe}$表示PCIe带宽。我们设定$p_0 = p_6$，因为子层1的激活值将被放置在前一个解码器层的子层6计算所在的设备上。

预填充阶段的特殊情况。公式(5)不适用于预填充阶段的i=2, 3的情况，在这些情况下：


这是因为在预填充阶段，子层2（Key）和子层3（Value）是由子层1生成的。

计算延迟。接下来，计算延迟为：


其中$B_{GPU}$和$B_{CPU}$分别是GPU和CPU的内存带宽，$C_i$是子层i的计算量，$T_{GPU}$和$T_{CPU}$分别是GPU和CPU的BF16计算吞吐量。

存储延迟。最后，存储延迟为：


其中$S_{KV}$是在子层1中生成的KV缓存大小。

模型参数依赖性总结。需要注意的是，计算量仅是B和L的函数，而CPU-GPU传输量不仅是B和L的函数，还取决于相邻子层的卸载选择。

5.2 性能优化

优化策略概述。为了进一步降低延迟，我们部署了两种优化措施，以更好地利用CPU和GPU的计算与内存资源。

优化-1：高效利用GPU显存。LIA首先将所有参数和中间值存储在CPU内存中，仅将计算当前子层所需的数据传输到GPU。这个过程通常会使大部分GPU显存处于未使用状态，特别是在小批量大小的情况下，此时总容量远超单个解码器层的需求。例如，在B=1的情况下进行OPT-30B推理时，超过90%的A100显存处于闲置状态。FlexGen利用这部分未使用的GPU显存来尽可能多地存储所有解码器层的各个子层，而LIA采取了不同的方法。LIA的做法是，在未使用的GPU显存中存储尽可能多的完整解码器层（包含所有子层）。这种方法使LIA能够更高效、更细粒度地利用GPU显存容量，同时也减少了CPU-GPU数据传输的需求。例如，在OPT-30B的情况下，LIA每存储一个解码器层增加约1.2GB的GPU显存使用，而FlexGen每存储一个所有解码器层的子层则增加4.7GB。当B=1且L=2016（即最大输入令牌序列长度）时，LIA在A100 GPU上使用35GB显存存储了62%的解码器层，而FlexGen使用32GB显存仅存储了58%的子层。

优化-2：重叠CPU/GPU计算与CPU-GPU传输。重叠计算和数据传输可以最大限度地减少CPU和/或GPU在传输期间的空闲时间。图7展示了LIA的重叠执行时序图。在预填充阶段，LIA遵循FlexGen的方法，将一个批次分成两个或多个小批次（mini-batches），将下一个解码器层的模型参数和相邻小批次的中间值的PCIe传输与当前小批次的计算重叠，以减少流水线气泡【索引43，FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU, 2023, ICML】。然而，在解码阶段，LIA将整个批次的计算与参数传输重叠，而FlexGen仍使用小批次重叠。最近的研究表明【索引37，AttAcc! Unleashing the Power of PIM for Batched Transformer-based Generative Model Inference, 2024, ASPLOS; 索引51，Duplex: A Device for Large Language Models with Mixture of Experts, Grouped Query Attention, and Continuous Batching, 2024, arXiv】，这种方法对解码有负面影响，因为计算时间不会随着小批次的减小而线性缩放，从而削弱了减少流水线气泡的好处。通过在解码阶段避免小批次重叠，LIA在B=900的OPT-175B推理中比FlexGen实现了1.1–1.3倍的延迟降低。


图 7: 单个解码器层的重叠时序图。示例为预填充策略 p = (0, 0, 0, 0, 0, 0) 和解码策略 p = (0, 1, 1, 0, 0, 0)。S{i}和O{i}分别表示子层i的计算和输出数据。

5.3 全系统框架实现

IPEX扩展与集成。我们通过扩展最新的IPEX库来实现完整的LIA框架。原始的IPEX库是基于pytorch-cpu和pytorch-CXX发行版构建的，用于英特尔CPU和GPU，与NVIDIA GPU不兼容。因此，我们重建了IPEX库，将其绑定到pytorch-cuda发行版。我们修改了IPEX Transformers【索引7，Intel® Extension for PyTorch】和HuggingFace Transformers【索引5，HuggingFace Transformers Library】库，以实现LIA的计算卸载算法和性能优化技术。

框架执行流程。首先，LIA框架根据输入特性确定最优卸载策略。这些特性包括每个子层的模型参数和中间值、批量大小（B）和输入长度（L）。框架还考虑系统性能特性，例如CPU-GPU接口的带宽以及CPU和GPU的计算吞吐量和内存带宽。利用这些信息来调度计算和CPU-GPU数据传输。最后，它使用推导出的最优卸载策略执行模型，确保高效、有效的推理。

A2 方法细节 (续)

6 利用CXL内存进行LLM推理

内存容量挑战与CXL方案。随着模型规模的增长，即使是大型CPU内存也不足以经济高效地存储所有参数和中间值，尤其是在批量大小（B）和输入长度（L）都很大时。例如，OPT-175B在L=1024且B=256时需要约1.4TB内存。为了应对这一挑战，我们提倡使用CXL内存，它由退役数据中心服务器中廉价的DDR4模块构成【索引54，Managing Memory Tiers with CXL in Virtualized Environments, 2024】，用于具有大批量（>860）的离线LLM推理。尽管CXL内存具有成本优势，但它也存在局限性：其带宽通常只有DDR内存的50%，延迟则高出2-3倍【索引48，Demystifying CXL Memory with Genuine CXL-Ready Systems and Devices, 2023, MICRO】。为了克服这些挑战，我们提出了一种内存卸载策略，当B很大时，将所有模型参数存储在CXL内存中，这一策略基于以下关键观察。

观察-1：CXL不影响参数的CPU-GPU传输带宽。虽然DDR内存子系统提供更高的带宽，但CPU-GPU传输的带宽最终受到PCIe接口的限制。因此，只要CXL内存能提供高于CPU-GPU之间PCIe带宽的带宽，CXL内存在CPU-GPU数据传输方面就能达到与DDR内存相同的带宽。当单个CXL内存扩展器提供的带宽低于CPU和GPU之间的PCIe连接时，可以通过页粒度的NUMA内存分配策略将多个CXL内存交错使用。这种交错能有效扩展CXL到GPU的数据传输带宽。图8(a)表明，通过交错使用两个CXL内存扩展器（每个提供17GB/s的内存带宽），在传输大数据块时（例如，OPT-175B模型每个子层≥300MB），组合带宽接近于DDR内存所能达到的最大值。这种带宽上的对等性确保了将参数存储在CXL内存中不会在数据传输过程中降低GPU性能。


图 8: (a) 不同数据大小时，DDR-GPU和CXL-GPU传输的带宽。(b) CPU从CXL内存获取模型参数和KV缓存的吞吐量，已归一化到DDR内存。

观察-2：盲目地将所有数据存入CXL内存会导致显著的性能下降。虽然CXL内存不影响GPU数据传输，但将数据存储在CXL内存中会显著降低CPU的计算吞吐量。图8(b)比较了当参数和KV缓存存储在CXL内存时，CPU AMX在子层1和子层2上的吞吐量（已归一化到DDR内存）。子层1（激活值与模型参数相乘）的吞吐量下降了11-70%，而子层2（解码阶段激活值与KV缓存相乘）的下降幅度更大，为10-82%。注意，子层2的计算不涉及模型参数。这种性能下降的差异源于这些子层的操作数/字节比。对于子层1，操作数/字节比在预填充阶段随B×L扩展，在解码阶段随B扩展；而对于子层2，无论B或L如何，该比值始终为1。这种差异使得子层2对内存带宽和延迟更为敏感。

为吞吐量驱动的推理设计的内存卸载策略。对于大批量B的吞DDR内存中，以避免对依赖KV缓存的子层（即子层2和3）的CPU计算性能造成下降。因此，即使所有参数都存储在CXL内存中，LIA也能确保推理吞吐量不受影响，同时优化内存成本。

A4 实验环境

  • 硬件配置:

    • 服务器: Supermicro X13DDW-A
    • CPU: 40核 Intel 第四代 Xeon 可扩展处理器 (Sapphire Rapids, SPR)。扩展性研究中使用了128核的Granite Rapids (GNR) CPU。
    • GPU: NVIDIA A100 (80GB) 和 H100 (80GB) GPU。评估了SPR-A100和SPR-H100两种配置。
    • 内存: 512GB DDR内存。CXL实验使用了两个三星CXL Type-3内存扩展器。
    • 互联: PCIe 5.0

  • 软件配置:

    • 实现: LIA框架通过扩展最新的Intel Extension for PyTorch (IPEX)【索引7，Intel® Extension for PyTorch】库实现。
    • 依赖库: PyTorch-CUDA, IPEX Transformers【索引7，Intel® Extension for PyTorch】, HuggingFace Transformers【索引5，HuggingFace Transformers Library】。
    • 基线框架: IPEX【索引7，Intel® Extension for PyTorch】（仅CPU，利用AMX）和 FlexGen【索引43，FlexGen: High-Throughput Generative Inference of Large Language Models with a single GPU, 2023, ICML】（最新的GPU+AVX卸载框架）。
    • 操作系统: Ubuntu 22.04.4 LTS, Linux 内核 6.8。

  • 模型与数据集:

    • 模型架构:
      • 主要评估模型: OPT-30B, OPT-66B, OPT-175B (BF16精度)。
      • 泛化性测试模型: Llama2-70B, Chinchilla-70B, Bloom-176B。
    • 数据集/负载:
      • 输入/输出序列长度: 基于Azure LLM推理追踪统计【索引38，Splitwise: Efficient Generative LLM Inference Using Phase Splitting, 2024, arXiv】设定。输入长度$L_{in}$从32到2048不等，输出长度$L_{out}$设为32（代码生成场景）和256（对话场景）。

  • 实验场景与指标:

    • 在线推理 (Latency-driven): 批量大小B=1，衡量指标为每次查询的秒数(s/query)。
    • 离线推理 (Throughput-driven): 批量大小B=64和B=900，衡量指标为每秒处理的令牌数(tokens/s)。
    • 能效: 使用ipmitool测量系统平均功耗，计算每令牌能耗(energy/token)。
    • 成本效益: 与多GPU系统（DGX-A100）进行每GPU吞吐量和每百万令牌成本的比较。
    • 延迟模型: 对于超出评估系统内存容量的配置，使用一个分析模型（基于公式(2)）进行性能估计，该模型平均误差为12%。模型推导的结果在图中用星号(⋆)标记。

A4 实验结果

7.1 计算卸载的最优策略

  • 实验内容: 分析OPT-175B在不同批量大小（B）和输入长度（L）下的最优计算卸载策略。
  • 实验结果:
    • 预填充阶段: 卸载策略取决于 B 和 L 的乘积。当 B×L 较小（对于OPT-175B，大约<850）时，将所有子层卸载到CPU以最小化延迟；当 B×L 较大时，所有子层都在GPU上计算以利用其高吞吐量。
    • 解码阶段: 卸载策略仅取决于 B。当 B 较小（对于OPT-175B，<858）时，所有子层卸载到CPU；当 B 较大时，GPU计算KV映射、输出投影、FC1和FC2子层，CPU处理其他子层。
  • 分析结论: 跨所有评估的OPT模型，LIA主要确定了三种策略：部分CPU卸载、完全CPU卸载、完全GPU计算。策略的选择受模型大小和GPU能力（H100比A100更倾向于GPU计算）的影响。这些策略对于类似OPT架构的模型（如GPT, Llama2）具有普适性（见图9）。


图 9: OPT-175B在SPR-A100和SPR-H100系统上，针对不同$L_{in}$和B组合的最优计算卸载策略。

7.2 在线推理延迟

  • 实验内容: 比较LIA、IPEX和FlexGen在B=1时的在线推理延迟。
  • 实验结果:
    • SPR-A100系统: 对于OPT-30B (OPT-175B)，LIA的延迟分别比IPEX低1.8–2.1倍 (1.1–1.3倍)，比FlexGen低5.3–7.3倍 (8.5–12倍)。
    • SPR-H100系统: 对于OPT-66B (OPT-175B)，LIA的延迟分别比IPEX低2.1–2.5倍 (1.1–1.5倍)，比FlexGen低4.9–7.0倍 (4.0–5.1倍)。
  • 分析结论: LIA通过大幅减少CPU-GPU数据传输（相对于FlexGen）和利用GPU处理计算密集型任务（相对于IPEX）来显著降低延迟。H100的更强性能放大了LIA相对于IPEX的优势，但缩小了与FlexGen的相对差距，因为FlexGen也能从更强的GPU和PCIe中受益（见图10）。


图 10: LIA, IPEX, 和 FlexGen 在不同模型和系统配置下的推理延迟比较。

7.3 离线推理吞吐量

  • 实验内容: 比较LIA、IPEX和FlexGen在B=64和B=900时的离线推理吞吐量。
  • 实验结果:
    • SPR-A100系统: 对于OPT-30B (OPT-175B)，LIA的吞吐量分别比IPEX高1.5–6.0倍 (1.1–6.1倍)，比FlexGen高2.0–5.9倍 (1.3–6.0倍)。
    • SPR-H100系统: 对于OPT-66B (OPT-175B)，LIA的吞吐量分别比IPEX高1.3–8.3倍 (1.2–10倍)，比FlexGen高1.2–3.3倍 (1.5–3.7倍)。
    • CXL卸载: LIA通过CXL卸载参数，可将DDR内存使用减少高达43.1%，而性能损失在1%以内；或者在相同DDR内存占用下，通过支持更大的批量，实现高达1.45倍的吞吐量提升（见表3）。
  • 分析结论: LIA通过智能卸载策略和利用AMX的高计算吞吐量，在各种配置下均实现了更高的吞吐量。CXL卸载策略被证明是提高大批量场景下吞吐量和成本效益的有效手段（见图11）。


图 11: LIA, IPEX, 和 FlexGen 在不同模型和系统配置下的推理吞吐量（tokens/s）比较。

7.4 消融研究与运行时分解

  • 实验内容: 分析LIA各项优化技术（高效GPU显存利用、计算传输重叠）和卸载策略的贡献，并分解LIA、IPEX和FlexGen的运行时。
  • 实验结果:
    • 消融研究(表4): 优化-1（GPU显存）对小批量(B=1)延迟降低贡献显著；优化-2（重叠）对大批量(B=900)效果更佳。LIA的卸载策略本身带来了巨大提升，在B=1和B=64时比FlexGen策略分别提升6.2倍和3.5倍。
    • 运行时分解(表5): LIA相比FlexGen显著减少了通信开销，相比IPEX显著减少了总计算时间。
  • 分析结论: LIA的性能提升是其智能卸载策略和多项优化技术协同作用的结果。

7.5 能效

  • 实验内容: 比较LIA、IPEX和FlexGen在SPR-A100系统上的每令牌能耗。
  • 实验结果: LIA的能效分别比IPEX和FlexGen高1.1–5.8倍和1.6–10.3倍。
  • 分析结论: LIA通过缩短推理延迟减少了系统静态功耗的影响，并通过将计算密集型任务卸载到能效更高的GPU，实现了更高的整体能效（见图12）。


图 12: 在SPR-A100系统上，IPEX和FlexGen相对于LIA的归一化每令牌能耗。

7.6 使用Granite Rapids (GNR) CPU进行扩展

  • 实验内容: 评估LIA在配备128核GNR CPU系统上的性能，并比较升级CPU（GNR-A100）与升级GPU（SPR-H100）的效果。
  • 实验结果:
    • 性能差距变化(表6): 从SPR升级到GNR后，LIA与IPEX的性能差距平均缩小14%，而与FlexGen的差距平均扩大1.7倍。
    • CPU vs GPU扩展(图13): GNR-A100系统在在线推理上比SPR-H100系统延迟低1.4–2.0倍，在离线推理（B=64）上吞吐量高1.9倍，但在B=900时稍逊一筹。
  • 分析结论: LIA能有效利用更强的CPU。GNR-A100系统在提供优越性能的同时，系统成本和能耗更低，展示了平衡CPU-GPU选择的重要性。

7.7 模型普适性、成本效益及与其他工作对比

  • 模型普适性: LIA在Llama2-70B、Chinchilla-70B和Bloom-176B等模型上同样表现出色，相比FlexGen延迟降低6.1-11倍，吞吐量提升1.2-5.0倍。
  • 成本效益 (vs. 多GPU, 图14): 在GNR-A100上，LIA在B=1时比8xA100的DGX系统具有更高的单卡吞吐量和更低的成本。虽然在B=64时稍逊，但潜力巨大，且支持DGX系统因内存不足而无法运行的B=900场景。LIA系统成本仅为DGX的10%。
  • 与PowerInfer对比 (图15): 在GNR-A100上，LIA对Llama2-70B的推理延迟比PowerInfer低1.4–9.0倍，吞吐量高1.5–15倍。PowerInfer在AMX-enabled系统上性能不佳，且在大批量时会内存溢出。

A5 结论

本文提出了LIA，一个利用英特尔AMX和CXL内存技术的单GPU与CPU协同的大型语言模型加速框架。我们首先通过基准测试证明，AMX的计算吞吐量与部分GPU相当，为高效的CPU计算卸载奠定了基础。基于此，我们开发了一种计算卸载算法，并将其在真实硬件上实现。实验结果表明，与基于英特尔AMX的CPU-only框架（IPEX）和最新的CPU-GPU协同LLM推理框架（FlexGen）相比，LIA分别实现了高达2.1倍和24倍的延迟降低，10倍和6.0倍的吞吐量提升，以及5.8倍和10.3倍的能效提升。这些成果展示了通过协同利用现代CPU的专用加速器和新型内存技术来加速LLM推理的巨大潜力。

A6 附录

A.1 摘要

附录目的。本附录为使用LIA进行LLM推理提供了指南，并说明如何复现本文的两个关键结果：1) LIA与IPEX和FlexGen在在线/离线推理中的性能对比；2) LIA针对大批量推理的CXL卸载功能。以下小节概述了访问、构建和设置软件环境以在支持AMX的CPU系统和NVIDIA GPU上部署LIA的步骤。

A.2 组件清单（元信息）

  • 运行环境：Ubuntu 22.04.4 LTS, Linux内核版本6.8。
  • 硬件：Intel Xeon Platinum 8460H处理器，NVIDIA A100 GPU，以及两个三星CXL Type-3内存扩展器。
  • 所需磁盘空间：380 GB，其中20 GB用于Docker镜像，其余用于OPT-30B和OPT-175B的模型权重。
  • 输出：来自LIA、IPEX和FlexGen程序的stdout.log文件，包含推理延迟。
  • 指标：在线推理的端到端推理延迟（秒）和离线推理的推理吞吐量（令牌/秒）。
  • 安装时间：1小时。
  • 实验时长：13小时。
  • 复现的实验：图10和图11中的SPR-A100结果，以及表3。
  • 公开可用性：是。
  • 代码许可证：Apache-2.0许可证。
  • 存档地址：https://zenodo.org/records/15104666

A.3 描述

A.3.1 如何访问。LIA的源代码、脚本和指南已在GitHub (https://github.com/Hyungyo1/LIA_AMXGPU) 和Zenodo (https://zenodo.org/records/15104666) 上公开。我们比较的其他框架可在以下地址找到：
- Intel Extension for PyTorch: https://github.com/intel/intel-extension-for-pytorch.git
- FlexGen: https://github.com/FMInference/FlexLLMGen
- PowerInfer: https://github.com/SJTU-IPADS/PowerInfer

A.3.1 硬件依赖。运行LIA进行LLM推理需要一个支持AMX的CPU（第四代Xeon CPU或更新版本）系统，配备一个NVIDIA GPU和两个或更多CXL type-3内存扩展器。我们使用特定命令将CXL内存从设备DAX模式转换为系统RAM模式。

A.4 安装

代码获取与初始化。首先，下载GitHub仓库，初始化并更新子模块：

$ git clone https://github.com/Hyungyo1/LIA_AMXGPU.git
$ cd LIA_AMXGPU
$ git submodule sync
$ git submodule update --init --recursive

Docker镜像构建。接下来，构建一个docker镜像并从源码编译LIA：

$ DOCKER_BUILDKIT=1 docker build -f \
examples/cpu/inference/python/llm/Dockerfile \
--build-arg COMPILE=ON -t lia-amxgpu:main .

或者，直接下载Docker镜像。

容器运行与环境激活。使用GPU并挂载存储来运行docker镜像，然后激活环境。

A.5 实验流程

模型权重准备。首先，使用以下命令为OPT-175B模型创建虚拟权重。

实验脚本执行。然后，运行以下脚本来收集LIA和IPEX在在线和离线推理中的数据。对于LIA的CXL卸载离线推理实验，运行另一个特定脚本，它会将结果保存到你挂载的存储目录。FlexGen的在线和离线推理性能可以通过运行Zenodo存档组件中的另一个脚本来收集。

A.6 评估与预期结果

复现内容。本文中的关键实验将被复现：(1) LIA、IPEX和FlexGen的在线（图10）和离线（图11）推理的延迟和吞吐量，以及(2) LIA使用CXL卸载的性能（表3）。每个实验结束后，可以使用以下命令自动收集结果数据以生成每个图表。

方法细节中的引用汇总

  • 引用文献 [5] HuggingFace Transformers Library

    • 引用段落: 5.3 全系统框架实现
    • 原文描述: 提及修改了HuggingFace Transformers库以实现LIA的计算卸载算法和性能优化技术。

  • 引用文献 [7] Intel® Extension for PyTorch

    • 引用段落: 5. LIA: AMX-aware Framework... (引言), 5.3 全系统框架实现
    • 原文描述: 在第5节引言中提到，LIA框架无缝集成了Intel IPEX。在5.3节中详细说明，LIA的实现是通过扩展最新的IPEX库完成的，并修改了其中的IPEX Transformers部分。

  • 引用文献 [37] AttAcc! Unleashing the Power of PIM for Batched Transformer-based Generative Model Inference, 2024, ASPLOS

    • 引用段落: 5.2 性能优化
    • 原文描述: 引用该研究以支持“在解码阶段使用小批次（mini-batch）重叠会对性能产生负面影响”的观点，因为计算时间不会随着小批次的减小而线性缩放。

  • 引用文献 [43] FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU, 2023, ICML

    • 引用段落: 5.2 性能优化
    • 原文描述: 引用FlexGen在预填充阶段采用的将批次拆分为小批次并重叠PCIe传输与计算的方法，LIA在预填充阶段也遵循了此方法。

  • 引用文献 [48] Demystifying CXL Memory with Genuine CXL-Ready Systems and Devices, 2023, MICRO

    • 引用段落: 6. 利用CXL内存进行LLM推理
    • 原文描述: 引用该文献来说明CXL内存的局限性，即其带宽通常只有DDR内存的50%，延迟高出2-3倍。

  • 引用文献 [51] Duplex: A Device for Large Language Models with Mixture of Experts, Grouped Query Attention, and Continuous Batching, 2024, arXiv

    • 引用段落: 5.2 性能优化
    • 原文描述: 与[37]一同被引用，以支持“在解码阶段使用小批次重叠对性能有负面影响”的观点。

  • 引用文献 [54] Managing Memory Tiers with CXL in Virtualized Environments, 2024

    • 引用段落: 6. 利用CXL内存进行LLM推理
    • 原文描述: 引用该文献以支持使用由退役数据中心服务器的廉价DDR4模块构成的CXL内存作为LLM推理的经济高效方案。