【论文分享】ASCEND-CC: Confidential Computing on Heterogeneous NPU for Emerging Generative AI Workloads

华为昇腾910A NPU：ASCEND-CC架构 【目前还挂在2024 Arxiv上】

摘要

云工作负载主导了基于大型语言模型 (LLM) 的生成式 AI。 GPU、NPU 和 TPU 等专用硬件加速器由于其性能优于通用 CPU，因此在人工智能采用中发挥着关键作用。人工智能模型和数据通常高度敏感，并且来自互不信任的各方。现有的基于 CPU 的 TEE（例如 Intel SGX 或 AMD SEV）无法提供足够的保护。像 Nvidia-CC 这样以设备为中心的 TEE 仅适用于紧密耦合的 CPU-GPU 系统，其专有解决方案需要主机 CPU 端的 TEE。另一方面，现有的学术提案是针对特定的 CPU-TEE 平台量身定制的。

为了解决这一差距，我们提出了 ASCEND-CC，这是一种基于离散 NPU 设备的机密计算架构，无需信任主机系统。 ASCEND-CC 通过确保数据和模型加密来提供强大的安全性，不仅保护数据，还保护模型参数和算子二进制文件。 ASCEND-CC 使用基于委托的内存语义来确保与主机软件堆栈的隔离，任务证明提供了强大的模型完整性保证。我们对 Llama2 和 Llama3 等最先进的LLMs进行的 ASCEND-CC 实施和评估表明，ASCEND-CC 引入的开销最小，且人工智能软件堆栈没有任何变化。

引言

最近，生成式人工智能 (GenAI) 势头强劲，大型语言模型 (LLM) 被用于聊天机器人 [1]、图像和视频生成 [2, 3]、代码补全 [4] 等应用中。 GenAI 被认为是未来通用人工智能 (AGI) 的构建块 [5]。主要云提供商提供以 AI 为中心的服务 [6-9]，这些服务通常利用 GPU、NPU 和 TPU 等专用加速器。

Security concerns. 这些大型 GenAI 工作负载给数据中心环境带来了众多安全挑战。训练LLMs所需的大量数据和计算资源使得它们极其昂贵[10]，并且对模型提供商来说是主要的知识产权。其次，用户对LLMs的查询通常包含敏感信息，例如健康数据、个人信息甚至商业机密[11]。在数据中心部署中，存在三个互不信任的各方：模型提供商、数据提供商和云提供商。模型提供商开发并拥有人工智能模型。数据提供者将数据提供给AI模型进行处理。最后，云提供商拥有计算基础设施，人工智能模型在其中进行训练或运行推理。数据提供者与计算基础设施交互，将其数据传递给人工智能模型。因此，人工智能模型和数据需要云提供商和软件堆栈的保护。

Gap in prior works.现有的基于 CPU 的可信执行环境 (TEE) [12-15] 可实现与操作系统/虚拟机管理程序等恶意软件堆栈隔离的安全应用程序或飞地。它们通常可以通过采用内存加密来抵御不受信任的 DRAM 和总线。在 CPU 之外，GPU [16-18]、IPU [19]、FPGA [20-22] 等设备提供 TEE。一些现有的工作 [23-27] 还将 CPU TEE 的安全原语扩展到连接的设备。然而，除了极少数提案[19]之外，大多数现有提案都需要CPU TEE，这增加了可信计算基（TCB）。最近的侧通道攻击[28-32]表明攻击者可以破坏CPU TEE的安全保证。此外，需要机密 VM 或 C-VM（例如 CCA 和 SEV）的提案通过信任 C-VM 操作系统进一步增加了 TCB。一些提案仅考虑集成 GPU [24,26,27] 或 NPU [25]，通过不考虑 PCIe 通信以及将内存空间与 CPU 分开来简化问题。将驱动程序的一部分放置在高权限可信安全监视器 [25,26,33] 中可以减轻关键安全决策的负担，例如任务和二进制文件的内存分配、内存共享以及对可信实体的访问控制。这种设计扩展了 TCB，不太适合主机完全恶意的情况。虽然 GraphCore IPU [19] 的提案可以承受无 TEE 的主机，但它不支持现代人工智能软件堆栈与加速器进行交互式会话，需要大量的硬件修改。最后，几乎所有现有的 TEE 提案都考虑较旧的基于 CNN 的 AI 模型，或使用矩阵乘法或 SVD 等运算符进行评估，这些运算符占用的内存较小。因此，这些解决方案无法扩展到内存占用较大且需要低延迟响应的LLMs。

Our contribution. 我们设计了 ASCEND-CC，这是一种不依赖 CPU TEE 的离散 NPU 机密计算解决方案。 ASCEND-CC的TCB只是NPU本身，整个主机都是不可信的。 NPU 中的硬件信任根 (HW-RoT) 有助于密钥派生，以在模型和数据提供者之间建立安全通道并启用证明。测量启动可确保 NPU 使用硬件制造商签名的正确固件启动。 ASCEND-CC 接受来自数据和模型提供商的完全加密的数据和模型。在推理过程中，NPU 从主机的虚拟地址空间中删除所有 DMA 映射（通过删除 SMMU 条目），以防止恶意 DMA 操作进入其模型、数据和工作空间（算子执行空间）。只有在内存取消映射之后，数据和模型解密才会开始。在相应的内存被 DMA 映射到主机之前，结果会被加密。 NPU 运行时根据 AI 模型创建任务，这些任务规定了运算符执行（例如矩阵乘法或 ReLU）和内存操作（例如从主机到设备的 DMA 复制）的顺序。恶意主机可以将任务（例如，执行 DMA 复制）注入模型中，以损害数据和模型的机密性。 ASCEND-CC 在模型开始执行之前执行模型的任务和二进制证明，以确保保留任务和模型的完整性。提供隔离和端到端加密，无需对 PyTorch 等 AI 软件堆栈进行更改。因此，ASCEND-CC 不会给 AI 程序员带来负担，将机密计算引入现有代码库。

我们通过修改固件在最先进的 NPU 华为 Ascend 910A 上演示了 ASCEND-CC。我们使用最先进的基于 Transformer 的 LLM 来评估 ASCEND-CC，例如 GPT-neo-125M、Llama-2-7B-Base、Llama-2-7B-chat、Llama-2-13B-instruct、Llama -3-8B-Base、Llama3-ChatQA1.5-8B、Llama-3-8B-Instruct、CodeLlama-7B-Instruct，用于聊天、句子完成和代码完成等任务。我们的评估表明，ASCEND-CC 在推理过程中引入了最小的性能损失（GPT2-neo125M 中为 0.91%，Llama3-Chat-QA-1.5-8B 模型中为 0.028%，均为 2K 输入序列大小）和一次性 set-向上。尽管我们的提案专注于特定的 NPU 实现，但我们的设计理念扩展到其他 AI 加速器，例如展示基于任务的执行模型的 GPU 和 TPU。综上所述，本文做出以下贡献：

(1) Identify fundamental building blocks for NPU-based confidential computing.我们确定安全属性，以保护模型和数据免受不受信任的主机和云提供商的侵害。具体来说，如何阻止特权主机访问 NPU 上的数据和模型并更改其内存映射。我们在第 2 节中列出了这些观察结果和要求。

(2) System design and analysis.我们设计的 ASCEND-CC 可以在 NPU 中实现机密计算。

(3) End-to-end evaluation.我们使用最先进的LLMs实施和评估 ASCEND-CC，并表明 ASCEND-CC 引入的开销最小，且无需对 AI 软件堆栈进行任何更改。

背景

NPU Background

NPU hardware architecture

在本文中，我们重点关注基于 Ascend 910A SoC 的华为 Atlas 300T 卡 [34]。 Ascend 910A 是一款最先进的 AI SoC，主要为大型数据中心和云计算提供训练和推理加速器。图 1 显示了 Ascend 910 NPU SoC 架构的高级视图。 SoC 的所有组件均通过内部总线连接。 NPU SoC 有两种类型的计算核心来执行 AI 任务：AI CPU 和 AI Core。四个 AI CPU 核心为通用华为泰山（ARM A73 配置文件），具有硬件加密扩展。 32 个 AI 核心基于华为达芬奇 [35] 架构，针对执行神经网络运算进行了优化。控制CPU是泰山核心（类似于AI CPU核心），运行NPU固件并管理PCIe接口。 NPU上电后，控制CPU启动（测量启动）最小的Linux内核并初始化硬件组件。任务调度程序将专用硬件组件与在 Taishan 核心上运行的固件相结合。它将AI任务分配给NPU的计算资源、AI CPU和AI Core。

NPU driver, runtime, and AI stack

NPU 运行时堆栈转换来自更高级别 AI 软件堆栈 (PyTorch/TensorFlow) 的指令，并与 NPU 驱动程序通信。 NPU 驱动程序是一组 Linux 内核模块，用于管理通过 DMA 的通信、发出 MMIO 命令来发送指令以及监视 NPU 运行状况。 Ascend PyTorch 适配器（torch_npu [36]）提供了必要的接口，将高级 AI 特定 API 桥接到较低级 NPU 驱动程序调用。 PyTorch 提供两种类型的模型执行： eager和dynamo。 Eager 是一种交互式执行，任务一到达 NPU 就立即执行。这是通过compileAndExecute 指令完成的，该指令将单个操作编译为图形并在 NPU 上执行。 Dynamo，即图形模式，在执行任务之前等待整个图形被编译并部署到 NPU。 NPU运行时将数据和模型复制到NPU HBM。该模型包括与每个 AI 模型层相关的参数以及相应的算子二进制文件（例如矩阵乘法、ReLU 等）。接下来，NPU 运行时创建一组要在 AI CPU 或 AI Core 上执行的任务，例如矩阵乘法。任务包含运算符元数据，例如运算符二进制文件在 NPU 内存 (PC_START) 上的位置、数据参数的位置以及存储中间变量的工作区。运行时将任务发送到任务缓冲区，即保留的 NPU 内存位置。发送完所有任务后，NPU 发送executeModel 任务，或者对于Eager 模式，发送compileAndExecute。这会触发任务调度器按顺序读取任务缓冲区，选择第一个任务，并将其提交给相应的AI CPU或AI Core。执行任务后，调度程序将任务条目移动到完成队列（CQ）并继续执行，直到任务缓冲区为空。 NPU运行时可以读取CQ来了解当前的执行进度。上述描述是针对单个PCI流上下文的，并且多个这样的上下文可以同时存在。这些流由任务调度程序并行处理。

动机 / 相关工作

Motivation: Gap in the Prior Art

涉及敏感和专有数据的大型 ML/AI 工作负载需要保护云中的数据和计算 [37,40]。值得注意的是，LLMs的兴起需要三方的机密计算设置：数据提供商、云提供商和模型提供商。模型和数据都不能泄露给其他方。我们将这种设置称为多驻留 TEE，因为 TEE 需要访问来自不同互不信任方的代码和数据。这与涉及两方的传统云-TEE 设置存在明显偏差：飞地用户和不可信云。

先前的工作将机密计算范例移植到特定于 ML 的加速器（例如 NPU [25]、GPU [16-18]）。正如表 1 所示，大多数现有提案依赖于基于 CPU 的 TEE 解决方案，例如 Intel-SGX [17, 18]、AMD-SEV [38]、TrustZone [26, 27] 或 ARM CCA [23, 24]在主机上，在此基础上构建以跨设备扩展安全保证。这样做的优点是已经拥有可信组件，即主机上的飞地，以确保与（可信）设备的安全通信。这使得能够使用集成设备的经过身份验证的加密、签名或内存隔离来发送经过身份验证的命令。基于 CPU 的 TEE 方法显着扩大了 TCB，不仅需要信任加速器设备，还需要信任 CPU 以及各种相应的监视器和驱动程序。此外，它使兼容性变得复杂，因为该解决方案依赖于特定的 CPU主机，现有的攻击[28-30]甚至可能破坏CPU TEE的安全保证。

据我们所知，只有两项先前的工作完全消除了主机的信任，并仅使用设备作为硬件信任根（HRoT）：SheF [21] 和 Graphcore IPU [19]。 SheF 使用 FPGA 作为可信设备，通过经过身份验证和加密的比特流以及设备隔离来确保完整性和机密性。 Graphcore 依靠专门的编译器在洁净室环境中将模型转换为加密且经过身份验证的二进制文件，该二进制文件可以直接发送到设备，无需主机干预。它需要对硬件进行重大更改，并消除交互式人工智能软件堆栈。因此，这些系统对于现实世界的大规模模型来说是不切实际的。

我们的提案仅依赖 NPU 作为信任根，不需要 CPU TEE，可与当前的 AI 框架配合使用，无需更改硬件，扩展到其他基于任务的 AI 加速器（例如 TPU），并针对运行现代真实环境进行了优化-世界LLMs。

A case against spacial sharing

LLMs规模的不断扩大影响着资源共享和利用策略。早期的 ML 特定 TEE 侧重于空间共享以提高设备利用率。它们依靠复杂的技术（多页表、监视器、专用硬件支持）来促进资源和性能隔离。我们在商业机密计算解决方案中观察到这种空间共享（也称为多租户）的趋势，例如 H100 和 B100 上的 Nvidia-CC（使用 MIG [41]），以及一些学术提案 [17, 23–26]。多租户通常是内存利用率较低的工作负载的选择，如旧式 CNN（ResNet、VGG、AlexNet 和 MobileNet）等现有学术提案或矩阵乘法或 SVD 等隔离操作中所见。然而，此类工作负载并不代表生成式人工智能等现代人工智能工作负载。这一点在图 2 中很明显，我们评估了最先进的基于 Transformer 的 Llama2 和 Llama3（具有 32GB HBM 的华为 Ascend 910A NPU）的内存利用率，我们的 HBM 利用率超过 90%。此外，单个 NPU 的内存不足以加载 70B 参数模型或在 2K 输入序列长度上执行 13B 参数模型。内部数据结构（例如 KV 缓存）与输入序列长度呈二次方增长。因此，增加序列长度提出了记忆容量的挑战。例如，在输入序列长度超过2K的Llama-2-13B中，NPU耗尽了HMB。其次，聊天机器人 [1]、代码生成 [4]、搜索 [42, 43] 等生成式人工智能应用程序对延迟敏感。由于多个租户之间的空间共享而导致计算资源减少，从而导致更高的延迟响应。鉴于 LLM 的内存占用如此之大以及延迟敏感性，我们得出结论，多租户是无关紧要的。因此，我们明确致力于单租户解决方案。

Settings

人工智能工作负载涉及三方。模型提供者带来他的 IP 模型，而数据提供者使用她的秘密数据来运行推理工作负载。模型提供者还可以使用这些数据来训练或微调模型。硬件，即主机系统、NPU、网络基础设施、操作系统、虚拟机管理程序和人工智能软件堆栈，由云提供商部署，所有计算都在云提供商中进行。云提供商还可以为数据提供商提供运行推理服务的模型。

Trust assumption and attacker model

在典型的可信执行场景中，软件堆栈和云提供商是不可信的。另外，我们的设置涉及两类TEE用户：模型和数据提供者。模型和数据提供者相互不信任，也不信任云提供商：

(1) Cloud/infrastructure provider 云服务提供商 (CSP) 负责配置和维护所有运行所需的硬件和软件资源。 CSP控制所有节点（CPU、NPU），管理所有基础设施，如网络接口、交换机等，并维护所有软件，例如操作系统、虚拟机管理程序、设备驱动程序、固件、AI/ML 软件堆栈（例如 PyTorch 或 TensorFlow）。我们假设云提供商提供的所有硬件和软件都是不可信的，除了执行 AI 模型的特定 NPU 之外。

(2) Model provider 模型提供者开发和训练模型，并对模型的组成和参数保密。模型提供商可以使用专有数据集训练开源模型。在这种情况下，模型的参数是秘密的。模型提供者也可以是多方：基础模型所有者和微调者。微调器使用特定的数据集训练基础模型，以适应视频生成或聊天等应用场景。 CSP 还可以是机器学习即服务 (MLaaS) 场景中的模型提供商。

(3) Data owner 通常，数据所有者是想要使用特定模型和云基础设施（CPU、内存、NPU）进行训练或推理的客户。这是数据所有者可以带来她的模型的场景。通常，我们假设数据和模型提供者是两个独立的、互不信任的各方。

我们仅假设部署 AI/ML 工作负载的特定 NPU 是可信的。 NPU 具有片上硬件安全模块 (HSM)，充当硬件信任根。最后，我们假设拒绝服务 (DoS) 和旁道攻击不在本文的讨论范围内。

Security Challenges and Requirements

我们使用矩阵乘法作为运行示例，如图 3 所示。NPU 运行时将矩阵乘法内核 (torch.mm) 的 NPU 优化二进制文件以及 M1 和 M2 复制到 NPU 的内存中。内核执行后，NPU将M3从NPU内存复制到CPU的主内存。该示例对应的三个任务如图4所示，即从主机到NPU的张量（M1，M2）的内存复制、从NPU的结果张量（M3）的矩阵乘法和内存复制到主机内存。主机端 NPU 运行时在 NPU HBM 上为矩阵乘法运算符的二进制、张量：M1、M2、M3 和工作空间（运算符的工作空间，例如堆和堆栈）保留内存空间。基于此执行模型，假设攻击者控制主机和云提供商，我们观察到一些安全挑战。基于这些观察，我们制定了一组 NPU 必须提供的安全要求，以确保模型、数据和执行的安全。

Security Challenge 1:不受信任的主机运行操作系统、虚拟机管理程序和设备驱动程序等特权软件，以及处理数据和模型的 AI 软件堆栈（例如 PyTorch）。因此，在传统 AI/ML 场景中的任何时候，不受信任的主机都可以完全访问数据和模型。

→ Requirement 1:所有输入和输出张量、模型参数和模型二进制文件都需要端到端身份验证加密（例如 AES-GCM），以确保攻击者控制的主机无法观察或操纵数据和模型。主机在任何时候仅处理经过身份验证的加密 (AES-GCM) 数据、模型和结果。

Security Challenge 2: NPU 需要在执行之前对模型和数据进行解密。当 NPU 在执行前解密数据时，攻击者控制的主机可以提取数据和模型。因此，确保一旦数据和模型在 NPU 内部被解密，攻击者就无法提取或操纵 NPU 上的模型和数据，这一点至关重要。同样，执行结果的端到端认证加密也是不够的，因为主机可以在结果加密之前从 NPU 内存访问纯文本结果。

→ Requirement 2:原子执行不变性：在数据和模型被解密以使模型执行之前，主机必须失去对数据和模型的访问权限，因为数据和模型需要在执行之前解密。这可以通过从 NPU 侧删除数据和模型所在的 NPU 内存区域的 DMA 映射来实现。

Security Challenge 3: 主机定义模型输入和输出的内存映射。这包括指向输入数据和模型的指针以及模型的输出结果。恶意主机可以将输出指针声明为与输入数据或模型相同，从而损害模型的机密性。

→ Requirement 3:内存不变性：对于从 NPU 到主机的每次内存复制，我们必须确保 NPU 拒绝任何纯文本输入数据和模型、中间结果或输出所在的内存复制。

Security Challenge 4: 即使没有直接访问数据或模型，不受信任的主机也可以向 NPU 发送恶意命令，例如将部分模型参数复制到发送给数据提供者的结果中，从而损害模型的机密性。

→ Requirement 4:模型执行缺乏完整性会损害模型和数据的安全性。

Security Challenge 5: 只有存在验证 NPU 的机制，机密计算的安全原语才有效且安全。如果没有系统的方法来检查 NPU 固件（包括控制 CPU、任务调度程序、NPU 内存管理器等的二进制文件）的完整性，我们就无法断言 NPU 机密计算能力的可信度。

→ Requirement 5:我们需要为 NPU 提供一个经过测量的启动等效原语，其中 NPU 只接受制造商认证的固件，并且不允许攻击者控制的主机刷新其固件或更改运行时配置。

Security Challenge 6: NPU 和相应的软件堆栈提供了多种正确性和性能的调试方法，例如检查 NPU 内存、获取内存区域的快照、观察执行时间等。此类机制允许攻击者提取模型和数据。

→ Requirement 6:为了确保数据和模型的机密性，必须限制所有与调试相关的操作。

Related Work

有几项工作旨在通过在 CPU TEE（例如 Intel SGX）内运行敏感计算来保护来自模型提供商的敏感数据。虽然某些方法在 CPU TEE 上运行整个模型 [49]，但这削弱了专用 ML 任务加速器的优势。因此，其他方法试图部分缓解这一缺点，方法是仅在 CPU-TEE 内运行选定的部分工作负载，同时利用加速器执行更密集的任务 [50-52]。然而，这些方法仍然有很大的开销，并且容易受到隐私窃取攻击[53]。

与此同时（如第 3 节中详述），还有一长串工作旨在将 CPU-TEE 或 C-VM 直接扩展到特定的加速器设备 [16、18、23、54、55]。

与本文最接近的是，Graphcore [19] 和 SheF [21] 等工作将信任完全从主机转移到用于运行工作负载的设备。同样，GuardNN [56] 通过将 FPGA 重新设计为 ML 任务的全新安全加速器，消除了主机的信任。在更大范围内，一些方法 [39, 57] 尝试将机密计算范式扩展到数据中心架构，允许许多设备跨用户分配。作为上述技术的补充，一些工作纯粹专注于通过改进功能来改善内存隔离[58, 59]，或增强机密计算的 I/O 隔离机制[60]，这通常直接有利于 TEE 设备架构。

相比之下，一些提案旨在通过使用秘密共享[61, 62]或同态加密[63,64]的加速器支持的安全多方计算（MPC）来通过算法来保护隐私。尽管在这些领域取得了重大进展，但这两种解决方案在应用于许多实际的现实工作负载时仍然会产生大量开销[65]，这使得它们对于我们的设置来说不切实际。

虽然使用 TPM 进行远程认证是最常见的方法，但也有一些基于软件的认证示例 [66, 67]，主要针对物联网设备。即将推出的 PCIe 功能支持连接机制CPU-TEE 与 DSA-TEE。具体来说，PCIe-6 的 TEE 设备接口安全协议 (TDISP) 使处理器上的 TEE 能够连接到支持 TEE 的 PCIe 加速器 [68]。 PCIe-5 上的完整性和数据加密 (IDE) 可对处理器和设备上的 PCIe 流量进行加密和完整性保护 [69]。采用这些 PCIe 扩展作为用于 Intel TDX 的 TDX-Connect、用于 AMD SEV-SNP 的 SEV-TIO、用于 Arm CCA 的设备连接 (DA) 或用于 RISC-V 的 IOPMP，使这些支持 TEE 的设备能够从安全的直接访问中受益到处理器上的 TEE 内存 [70-73]。

虽然侧通道超出了范围，但一些工作提出了在使用 TEE 时保护工作负载免受此类攻击的解决方案。例如，Telekine [74] 可以保护 Graviton [17] 中提出的 TEE 的工作负载。

细节设计

• Basic Building Blocks for Confidential Computing on the Ascend NPU
  • Model and Data Encryption
    • AI CPU-based custom operator
  • Executing AI CPU operator with model
    • Preparing the model
    • Model execution on the NPU
  • Enforcing Memory Lock Invariants
  • Model and Task Attestation
  • Firmware and Runtime Integrity
• ASCEND-CC
  • Initial setup
  • ASCEND-CC: End-to-end system
  • Programming interface
  • ASCEND-CC lifecycle