Abstract

NPU和CPU的异构协同计算因其显着的性能优势而受到广泛关注。为了确保计算过程中数据的机密性和完整性，可信执行环境（TEE）因其相对较低的开销而被认为是一种有前途的解决方案。然而，由于CPU和NPU之间的内存粒度细且不同，现有的异构TEE设计对于协作计算来说效率低下。 1）CPU TEE的cacheline粒度由于额外的内存访问而加剧了内存压力；2）NPU的cacheline粒度MAC加剧了有限内存存储的压力。 3）跨异构飞地的数据传输依赖于非安全区域的传输，导致繁琐的重新加密和调度。

为了解决这些问题，我们提出了 TensorTEE，一种用于高效安全协作张量计算的统一张量粒度异构 TEE。首先，我们实际上支持 CPU TEE 中的张量粒度，通过检测和维护片上张量结构来消除片外元数据访问。其次，我们提出具有预测执行的张量粒度 MAC 管理，以避免计算停顿，同时消除片外 MAC 存储和访问。此外，基于统一的粒度，我们可以实现直接数据传输，而无需重新加密和调度困境。我们的评估基于增强的 Gem5 和周期精确的 NPU 模拟器。结果表明TensorTEE提高了大型语言模型的性能。与现有工作相比，LLM训练工作量提高了4.0倍，与非安全训练相比，开销仅为2.1%，为LLM训练提供了切实可行的安全保障。

Introduction

NPU和CPU的异构协同计算在各个领域都表现出了显着的性能和能耗提升[12,30,36,44,64,69-71]，特别是在处理NPU的DRAM容量有限带来的挑战时，例如如 LLM 训练 [70]、自动驾驶 [82] 和 AlphaFold2 [46]，因此异构协作架构成为最先进的趋势：Nvidia Grace Hopper [20] 和 AMD Instinct MI300 [5]。这种改进源于异构协同计算系统可以有效地利用不同架构的特点。例如，DeepSpeed [68] 的 ZeRO-offload [70] 将前向和后向计算卸载到 NPU，而 CPU 管理高精度权重和优化器状态。这种异构协作方式有效缓解了 DRAM 存储压力，从而能够在单个 V100 GPU 上训练 10B 参数的 GPT-2 [70]。

鉴于协同计算在计算效率和数据隐私方面同样重要，异构协同计算的突出表现日益激发人们对安全的追求。这种趋势使得可信执行环境（TEE）获得了极大的关注，它是为CPU [13,16,53]开发并扩展到NPU架构[13,43,56,78,84]，因为它可以提供安全的飞地与替代解决方案（例如完全同态加密和多方计算）相比，具有强大的隐私保证和更低的开销。具体来说，TEE只认为片上硬件是可信的，因此CPU/NPU片外存储器和片间通信都需要额外的保护机制，以防止来自操作系统或物理攻击的攻击。在本文中，我们重点关注具有片外 GDDR 存储器的分立 NPU，其配置与最近的工作相同 [2,34,62]。

然而，现有的异构TEE设计[34, 78]，在CPU和NPU之间具有不同的内存保护粒度，会产生显着的性能开销，并且对于协作计算来说不可行。它们面临三个方面的效率问题：1）CPU TEE 中采用的版本号（VN）的缓存行粒度会为 VN 和 Merkle Tree 遍历产生大量额外的内存访问，加剧内存密集型工作负载中的内存压力，并导致显着的开销。 2）虽然VN在NPU中是张量管理的，但它们的缓存行粒度MAC管理带来了额外的存储开销，进一步加剧了本已受限的内存存储的压力。 3）由于异构TEE的粒度差异，跨异构Enclave的数据传输依赖于非安全区域的传输，需要重新加密和解密以维护数据安全。这种重新加密和解密过程引入了大量的额外内存访问，这反过来又导致数据传输和计算之间对内存带宽的竞争，从而阻止并行执行。

为了解决这些问题，我们的工作创建了一个统一的TEE来有效支持安全异构协作计算，消除CPU中密集批处理数据访问的VN开销、NPU中大粒度MAC引起的计算停顿以及安全通信通道的重新加密开销。 1）对于CPU来说，TensorTEE虚拟地支持片上张量粒度的VN管理（兼容cacheline粒度），从而消除了片外VN访问和Merkle Tree遍历开销。 TensorTEE 通过利用维度流访问模式和灵活的条目合并操作，有效地检测和管理片上张量结构。 2）对于NPU，TensorTEE提出了具有延迟验证的张量式MAC管理，以减少存储开销并消除计算停顿。它通过提出的张量毒物追踪和验证屏障机制来维护数据完整性，以确保 NPU enclave 中的数据篡改受到限制，并且可以在通信之前很快被检测到。为了确保代码完整性，我们限制指令内存请求遵循正常的非延迟验证数据流。 3）对于异构通信，TensorTEE的统一粒度带来了跨异构enclave的密文兼容性，避免了数据传输中繁琐的重新加密和调度。

Contribution

我们分析了现有异构 TEE 设计中 CPU 和 NPU 之间不同粒度带来的开销，这启发了我们统一张量粒度的异构 TEE 设计。

我们提出了TensorTEE，一种统一粒度的异构TEE架构，包括基于硬件的张量粒度CPU TEE、具有延迟验证的张量MAC管理以及安全飞地内存之间无需重新加密的直接数据传输协议。

我们使用增强型 Gem5 和周期精确的 NPU 模拟器评估了 TensorTEE。我们的评估结果表明，与现有工作相比，TensorTEE 将 LLM 训练的性能提高了 4.0 倍，并且与非安全训练相比仅产生 2.1% 的开销。

Background

Collaborative Computing

由于CPU和NPU独特的计算模型和局限性，协作计算变得越来越流行[46,70,82]。

LLM CPU-NPU collaborative computing

由于NPU DRAM容量有限，大型语言模型训练趋向于NPU和CPU的异构协同计算[12,30,70]。本文重点关注用于评估研究的ZeRO-Offload [70]，这是LLM协作训练中流行的框架。主要原因是ZeRO-Offload可以将计算负载较低但内存使用率较高的阶段从NPU卸载到CPU，有效减轻NPU的内存存储压力，从而实现更大的模型训练。例如，通过将梯度和优化器状态卸载到 CPU，同时在 NPU 上保持前向和后向计算，ZeRO-Offload 可以将训练模型大小增加 10 倍。如图1所示，ZeRO-Offload的计算流程主要涉及三个阶段。首先，NPU 执行前向和后向计算，在反向传播期间将更新的梯度传递给 CPU。然后，CPU 进行优化器迭代以更新优化器状态和权重。最后，CPU 将更新后的权重传回 NPU。

Memory Protection

在 TEE 中，内存保护包括通过加密确保数据机密性、通过消息身份验证代码 (MAC) 验证确保数据完整性以及使用单调计数器确保数据新鲜度。

Memory encryption

TEE 通常采用计数器模式 AES 进行内存加密，它引入了由物理地址 (PA) 和版本号 (VN) 组成的计数器。 VN 随着每次写回而增加，以防止重放攻击。设 KAES、P 和 C 分别表示 AES 加密密钥、明文和密文。那么，加密可以表示为 C = AES (KAES, (PA, V N )) ⊕P，其中 ⊕ 表示 XOR。由于 XOR 性质，解密遵循相同的过程。

粒度不匹配：图 2 说明了 VN 管理的两种不同粒度。首先，CPU 为每个缓存行保留一个专用的 VN，以支持随机内存访问（图 2 (a)）。然而，细粒度，例如 64B 数据块的 56 位 VN，会引入较大的存储开销（11%）和内存访问开销[13]。此外，片外 VN 依靠 Merkle Tree 来确保其完整性。其次，最近的工作 [34, 78] 采用了 NPU 的张量 VN，利用了 NPU 工作负载中普遍存在的常规内存访问模式（图 2 (b)）。对于张量粒度，单个 VN 用于张量内的所有缓存行。这种方法显着降低了存储压力，并实现了 VN 的片上存储，从而减少了 VN 的内存访问并消除了对 Merkle Tree 的需求。

现有工作的局限性：SoftVN [83] 通过在软件代码中显式声明来保留片上 VN 表中的张量结构。它以最小的硬件开销在简单的静态场景中高效执行。然而，SoftVN有一定的局限性：1）不适用于复杂和动态的场景：SoftVN[83]假设程序中的数据流是静态的和简单的。然而，在LLM训练中，存在各种运行时并行策略和更复杂的操作，例如张量分裂和合并[67]，这使得跟踪张量的数据流并在顶层用户代码中的每次写回时指定VN变得具有挑战性。 2）提高实用性的困境：由于SoftVN中VN的获取发生在关键路径（缓存访问）上，因此增加可管理的张量条目数量会导致性能下降。这使得SoftVN在复杂应用下陷入实用性和高性能之间的困境。 3）条目浪费：在多个核并行使用单个张量的场景中，SoftVN 导致每个子张量占用相应核中的一个 VN Table 条目。这加剧了容量压力，凸显了 SoftVN 的一个关键限制。

Integrity verification

为了确保数据完整性，在写回期间计算数据的MAC并将其存储在DRAM中。在随后的读取中，MAC 会被重新计算和验证，以确保它与片外存储中的 MAC 相同。 MAC计算可以表示为MAC = Hash(KMAC, (C, PA, V N ))。为了防止重放攻击，MAC 应存储在片内。但由于片上存储有限，采用VN和Merkle Tree来减少片上存储开销。 BMT[72]进一步提出Merkle Tree只需要保护VN，减少Merkle Tree的宽度。最终，Merkle Tree 的根节点安全地存储在片内，而 VN 和 MAC 则存储在片外。每次内存访问的完整性验证都依赖于 Merkle Tree 的递归逐层验证，导致开销很大。对于 NPU，VN 安全地存储在片上，无需迭代访问 Merkle 树。

Heterogeneous NPU TEE

最近的工作已将 CPU TEE 扩展到异构 TEE，用于 NPU 内的关键任务计算。异构NPU TEE根据与CPU的连接情况可以分为集成型和分立型。

Integrated NPU TEE

集成架构将CPU和NPU放置在同一芯片上并共享相同的DRAM，因此CPU和NPU之间的数据传输没有问题[15, 56]。然而，由于功耗和面积的限制，集成的NPU TEE的性能受到限制，这使得它不足以满足诸如LLM训练等涉及大量计算的高要求应用。

Discrete NPU TEE

离散异构架构具有独立的CPU和NPU芯片，并通过GDDR等数据总线将它们连接起来。分立 NPU 具有专用控制器和内存，可提供增强的计算能力来支持 LLM 训练。现有的 NPU 内存保护工作更喜欢张量方式的 VN 管理来提高性能：Common Counter [62] 在内核切换时扫描内存，将具有相同 VN 的数据区域的元数据保存在芯片上，从而减少元数据的额外内存访问开销。另一方面，MGX [34]和Securator [78]根据内核执行期间的片上执行状态为每次内存访问生成相应的VN，从而最大限度地减少VN存储和访问开销。由于离散异构架构在产品环境中应用更为广泛，因此我们将重点关注离散异构TEE。

Threat Model

作为一个异构系统，我们的可信计算基（TCB）仅由片上架构（如CPU、NPU和缓存）组成。片外主机内存和 NPU GDDR 内存驻留在 TCB 外部。对于 NPU，遵循与最近的工作[2,34,62]相同的配置，我们考虑具有片外连接 GDDR 存储器的 NPU。片外内存很容易受到物理攻击[88]，例如总线窥探攻击[6,35,37,77]，已经开发出成熟的工具包[10]，以及Intel TDX[42]、AMD SEV[48]假设的冷启动攻击。

我们假设对手完全控制操作系统或特权软件，并且可以进行物理攻击。攻击者可以窥探总线（内存总线或 PCIe 总线）[35, 37] 或采用冷启动攻击 [1, 42, 56] 绕过安全隔离来窃取数据。他们还可以操纵总线信号进行数据损坏或中继攻击。我们不考虑其他侧信道攻击，例如基于时间、基于功率和基于电磁的攻击[39,47,49,52,63]，以及拒绝服务攻击和对抗性攻击[9 、27、28]。模型提取攻击可以通过正交方法来缓解[80]。

Motivation

Inefficient CPU TEE Computation with Cacheline-Wise Metadata

在异构协作 LLM 训练中，CPU 负责优化器更新，如 Adam 优化器步骤。图 3 显示引入 TEE 时会导致高达 3.7 倍的速度下降。在SGX环境中，随着线程数量的增加，性能优势逐渐减弱，表明从计算密集型任务过渡到内存密集型任务。内存访问量的增加主要归因于 VN 和 MAC 等元数据。

对于 Adam 更新，张量数据是按元素访问的。常规访问模式可以实现张量粒度的内存保护。图 4 说明了不同模型的优化器更新期间张量的数量和大小。我们观察到张量大小增长到 MByte，但张量数量增长速度很慢，仅达到几百。张量数量少、规模大的数据特性意味着张量粒度的内存保护将获得巨大的性能增益一个小的开销。然而，CPU总是以缓存行的粒度访问内存，并且内存组织或指令没有张量类型。

Cacheline-Wise MAC of NPU Leads to Large Storage and Performance Overhead

NPU上的内存资源有限，相比于服务器CPU上高达512GB的CPU内存，NPU内存容量通常只有40GB。 enclave元数据的存储进一步加剧了内存资源的短缺。最近的 NPU 与张量粒度 VN [34, 79] 一起工作，减少了 VN 的存储开销，但来自 MAC 的接近 10% 的存储开销和 12% 的性能开销仍未解决。 Plutus [2] 采用预测执行来减少 MAC 的额外内存访问，但它没有解决存储开销问题。 MGX [34]和GuardNN [33]将MAC的粒度从64B提升到512B，但他们没有考虑粒度增加导致的计算停顿问题（见4.3节）。因此，迫切需要一种能够减少存储和性能开销的完整性验证方法。

为了解决上述挑战，我们建议利用带有延迟验证的张量MAC，并通过延迟验证为完整性提供以下保证：1）篡改只能发生在数据上，而不能发生在代码上，并且可以很快检测到篡改的数据。 2) 篡改数据NPU不能离开NPU enclave，包括NPU芯片和GDDR内存。首先，使用NPU代码的正常非延迟验证来防止代码篡改并呈现基于延迟验证的攻击[76, 77]。其次，我们确保通信数据的完整性，防止被篡改的数据离开NPU enclave。第三，虽然由于验证延迟，NPU数据篡改可能会在短时间内发生，但由于NPU作为纯粹的加速器缺乏文件I/O能力，因此无法利用它进行更强大的攻击（例如数据窃取）。

Inefficient Data Transfer Protocols

图 5 显示，在非安全模式下，通信仅占用 12% 的开销，但在 CPU 具有类 SGX TEE 和 NPU 具有类 MGX TEE 的基准系统中，通信开销显着增加到 53%。原因如下：

CPU和NPU的保护粒度不匹配会导致额外的重新加密。 CPU和NPU的粒度不匹配，CPU是在cacheline粒度上进行保护，而NPU是在张量粒度上进行保护。由于粒度不兼容，TEE 即使具有相同的密钥，CPU（或NPU）TEE的加密数据无法被NPU（或CPU）正确解密，因为它们的硬件引擎固定且元数据格式不同。因此，通信依赖于非安全存储器作为中继（图6（a））。发送方需要解密安全​​内存区域的数据并将其保存到非安全区域。通信后，接收方将数据加载到安全区域中。为了保护数据的隐私，需要对非安全区域的数据进行重新加密。

即使加密引擎足够，通信（重新加密）和计算（IO 读/写）引起的带宽争用也会导致两个任务被迫按顺序执行（图 7）。对于资源受限的NPU，硬件成本限制了多个加密引擎的部署，导致加密带宽的竞争[90]。最近的研究 SecureLoop [55] 报告称，全流水线加密引擎占 Eyeriss 类 NPU 中逻辑门的近 35%。在我们的模拟中，即使是NPU计算，一个AES引擎也无法提供足够的带宽（提供8GB/s，但至少需要20GB/s），因此通信数据必须通过NPU计算进行延迟和串行化。即使对于拥有足够 AES 引擎的大型 NPU，通信过程中的重新加密也会加剧内存带宽的压力，并阻塞与计算相关的 IO。重新加密过程非常消耗内存，并且在我们的模拟中占用了大部分 DRAM 带宽（84%），阻碍了其与 NPU 计算的并行化。上述两者都表明计算和通信的顺序执行导致数据流效率低下。在本文中，我们假设每个通道都有其专用的加密引擎来平衡硬件开销和带宽要求之间的权衡。

Optimization Goal

这些观察启发了统一的张量粒度异构 TEE 架构，它具有以下几个优点：1) 具有片上 VN 的张量 CPU TEE 减少了元数据的内存访问。 2）NPU的Tensor-wise MAC管理减轻了存储和性能开销，同时保证了NPU数据的安全。 3）统一的异构张量TEE实现安全内存之间的直接数据传输。

TensorTEE Design

Overview

我们的整体架构如图 8 所示，引入了三项优化：

Tensor-wise memory protection on CPU

为了统一保护粒度并提高协同计算上的CPU TEE性能，我们提出了基于硬件的Tensor-Analyzer（TenAnalyzer）组件，用于张量内存保护。 TenAnalyzer接收并分析来自CPU内核的内存访问请求，在运行时构建和维护Meta Table数据结构。元表的每个条目都保存张量内所有缓存行的共享元数据（如 VN），从而消除了后续内存请求中元数据内存访问的开销。我们的设计是可选的，是对 SGX [13] 等传统保护方法的补充，通过启用张量管理标志（EnTMF）启用或禁用张量管理，而不影响软件编程。在非张量应用中，张量管理被禁用，开销可以忽略不计。

Tensor-wise MAC with delayed verification on NPU

为了提高性能和存储效率，TensorTEE提出了张量式MAC管理，仅对NPU张量数据进行延迟验证。延迟验证可实现验证和计算之间的并行化，消除粗粒度 MAC 验证造成的停顿。为了确保完整性，TensorTEE 1）按照正常的非延迟验证数据流限制代码访问，以维护代码完整性并防止任何潜在的攻击，2）通过提出的张量中毒跟踪和验证屏障机制确保通信数据离开 NPU enclave 时的完整性。

Implementation of direct data transfer protocol

利用统一的张量粒度，我们实现了高效的数据传输协议，消除了通过非安全区域的传输，并实现了数据传输和计算之间的并行。该协议包含两个通道：可信通道用于片上元数据的加密传输，直接通道用于异构安全区域之间密文张量的传输。

CPU TEE with Tensor Granularity

为了增加实用性并减轻程序员的负担，需要一种基于硬件的张量CPU内存管理机制。它应该包含张量管理功能，同时保持较低的开销。

Challenges of hardware-based tensor data structure detection:

为了解决条目浪费问题，张量管理结构应位于内核外部，例如内存控制器（MC）中。然而，MC 中的张量检测和维护存在挑战： 1）缓存干扰：如图 9 所示，进程切换会导致随机缓存行驱逐，扰乱内存访问模式。此外，当发出错误的预取请求时，缓存预取器（如跨步预取器）会打乱访问顺序。 2）MC接收物理地址，它们的连续性受到页面大小的限制。 3）延迟较长：请求到达MC之前对大容量缓存的访问会带来显着的内存访问延迟，导致维护和访问张量数据结构的延迟较长。

Tensor structure management in CPU TEE

为了解决上述挑战，我们提出了位于MC内的TenAnalyzer，它直接接收来自内核的内存访问请求，如图8所示。TenAnalyzer具有以下优点： 1）它直接接收来自内核的内存访问请求，避免了逐出和预取带来的缓存干扰如图 9 (a) 所示。 2) 利用核心的虚拟地址最大化张量连续性的利用率，如图 9 (b) 所示。为了避免同样的 VA 问题，每个 enclave 使用其专用密钥，并保存和恢复元表以供上下文切换情况使用。 3）张量数据结构的检测、维护和访问可以通过多级缓存并行化，有效隐藏其访问延迟。因此，TenAnalyzer克服了SoftVN [83]的局限性，并为基于张量的内存管理提供了高效且可扩展的解决方案。

TenAnalyzer主要包含两个组件：Meta Table和Tensor Filter，如图10所示。 1）Meta Table维护张量结构，并且可以通过边界命中不断更新。具体来说，Meta Table接收Core的访问请求，命中时将地址和VN传递给MEE进行加密，这样就避免了对内存中元数据的额外访问。 2）Tensor Filter 处理 Meta Table 缺失，收集并分析请求地址。当张量过滤器中的特定条目达到其集合地址限制（在我们的设置中为 4）时，它会检查该条目是否满足张量条件：具有相同的 VN 以及地址之间的一致模式。然后，成功检查的条目将填充到元数据表中作为张量的初始化结构。此外，还使用了其他组件，即位图和两个用于指示缓存行更新状态的标志以及用于启用或禁用 TenAnalyzer 的 EnTMF 标志。

Dataflow for reading (or detection)

如图10所示，TenAnalyzer在请求到达时分为三种情况：命中、命中边界和未命中。 1）Hit in是指请求命中Meta Table中某个条目的地址范围，然后将对应的VN传输到MEE开始加解密。 2）命中边界是指请求命中了Meta Table中VN无法保证正确的条目的地址范围边界（请求地址==最后地址+步幅）。此时，假设这个VN是正确的，则仍然将VN发送到MEE，同时开始DRAM中的VN访问。如果DRAM中的VN与假设的VN一致，则更新相关Meta Table条目并扩大其地址范围，从而导致张量检测逐渐覆盖。 3）Miss表示Meta Table中没有匹配，则通过DRAM访问获取VN，并将该请求发送到Tensor Filter进行后续的模式检测和过滤。对于读取，片上张量方式 VN 始终保持等于片外缓存行方式 VN，同时高效地提供 VN 以最大限度地减少片外内存访问和 Merkle Tree 遍历。

TenAnalyzer 通过条目合并有效检测具有复杂访问模式的张量。对于流操作，例如 Adam 的更新，单个条目足以检测和管理大张量（图 11 (a)）。然而，对于需要平铺的张量运算（例如矩阵乘法），使用单个条目检测完整的张量信息变得具有挑战性。尽管如此，TenAnalyzer 仍然可以像管理图块的一行一样管理每个短流数据，并根据维度流访问模式推断图块暗淡（例如，大小为 d1 的第一个蓝色行）。接下来，TenAnalyzer 在创建新条目时尝试合并一些最近更新的条目。合并操作不限于地址相邻的图块，而是允许在多个方向上合并，从而实现更及时的合并和更高效的片上存储（1D 张量 2 个方向，2D 张量 4 个方向，3D 张量 6 个方向）。要合并具有非相邻地址的图块，需要图块变暗、步幅和 VN 匹配。成功的合并允许推断张量维度（例如，合并两个蓝色行推断张量维度 D1）。推断出的维度可以作为后续合并的约束，以提高合并的准确性。通过迭代这个过程，可以实现对平铺张量的统一管理，如图11（b）所示。

TenAnalyzer 逐渐完成张量结构，但需要额外的内存访问来验证其正确性。由于来自 NPU 的数据传输指令通常包括张量结构信息，例如地址、大小和步长，TensorTEE 利用它们来更新具有较大重叠地址范围的元表条目，以加快 CPU 中张量结构的创建速度。

Dataflow for writing (or update)

在张量写入期间，TenAnalyzer 通过使用位图 (BM)、更新标志 (UF) 和位状态 (BS) 确认张量内的每个缓存行只能在张量更新内更新一次，从而有效地维护张量 VN 的正确更新。张量的位图和地址范围日志缓存线更新，翻转写请求时地址上的位。 UF表示更新张量，BS表示更新前的位图位值。 TenAnalyzer 在更新所有张量缓存行后递增 VN（所有相关位图位翻转）。硬件隔离机制可保护 DRAM 中的完整位图，小型片上缓存可提高位图访问效率[21]。

TenAnalyzer 在张量写入中包含 3 种场景，如图 12 所示：命中边缘、命中和未命中。 1）命中边缘是指请求命中与常见的基于张量的第一个地址（开始更新）或最后一个地址（完成更新）匹配。应用程序并允许复杂的内存访问，例如平铺。在与 BS 检查相等后，位图位翻转，标记缓存行已更新。当最后一个cacheline地址到达时，​​检查地址范围内的所有位图位是否等于翻转的BS，以指示张量更新完成；然后片上张量递增，BS 翻转，UF 重置。 2) 命中指请求命中地址范围但未命中边缘。检查 UF 和位图状态后，我们翻转位图位。 3) Miss表示地址超出范围，只需要片外VN更新。请注意，来自内核的写入地址由 LLC 过滤以匹配片外数据更新。写入是在后台执行的，而不是在关键路径上执行。

断言（Assert1、2 和 3）确保在张量更新完成之前每个缓存行只能更新一次，以保证片上张量 VN 的正确更新。 Assert1 确保在张量更新开始之前每个缓存行不能更新，而 Assert2 确保每个缓存行必须在张量更新完成时更新一次。访问模式在基于张量的应用程序中很常见。当断言被违反时，元表条目无效，从而防止在以下极端情况下张量 VN 和片外 VN 之间出现不一致： 1）检测到的张量与非张量变量混合。2）一个条目包含多个具有不同更新频率的张量。

Delayed Verification for Tensor-Wise MAC

MAC 的粒度权衡了存储开销和验证开销。 Cacheline-wise MAC 导致 NPU 中存储开销较大，由于内存容量有限，这进一步阻碍了 LLM 工作负载在 NPU 上的部署。最近的研究 [33, 34] 选择了 512 字节的 MAC 粒度，从而在存储开销和验证成本之间建立了平衡。然而，使用大粒度 MAC 会导致稍后的验证，从而导致已解密的缓存行的计算停顿，如图 13 (b) 所示。具体来说，尽管在两个粒度上具有相同的 MAC 重新生成时间，但后续 MAC 验证阶段的阻塞计算会导致许多管道气泡，从而导致管道效率低下（4KB 粒度下的开销为 13%）。

为了减少存储开销，同时避免计算停顿，利用LLM协作计算的特点，我们提出了一种具有延迟完整性验证的张量MAC管理方法。我们对解密数据和重新生成张量 MAC 进行并行内核计算，如图 13 © 所示。完成tensor-wise MAC后，我们对整个张量进行完整性验证。通过这种方式，我们显着减少了存储开销和内存流量，同时避免了长停顿造成的大气泡，如图 13 (b) 所示。对于张量 MAC 计算，我们采用一种简单而高效的算法，该算法涉及对所有缓存行的 MAC 值进行异或，类似于 MEE [24]。该算法可以表述如下，其中 MACi 表示张量中第 i 个缓存行的 MAC，n 是缓存行的数量： MACtensor = MAC0 ⊕ MAC1 ⊕ … ⊕ MACn−1。基于 XOR 的算法对顺序不敏感，允许在 NPU 计算中进行各种优化，例如张量平铺。

Integrity guarantee

我们通过对代码访问和通信数据严格的完整性保证来确保 NPU 计算上张量数据的完整性。这允许在计算期间延迟验证中间张量，从而在不影响安全性的情况下提高性能。延迟验证允许临时篡改中间张量，但很快就能被 MAC 检测到。 NPU 中间张量的临时篡改是可以接受的，因为它不能被用来进行更强大的攻击，例如数据盗窃。这是因为NPU作为纯粹的加速器，缺乏特权软件（如操作系统）和文件I/O能力。相比之下，离开NPU和GDDR内存的NPU代码和通信数据需要严格的完整性保证。

TensorTEE 通过限制遵循正常非延迟 MAC 验证数据流的代码访问请求来保证代码完整性。 NPU代码的非延迟验证可以禁用基于代码篡改的延迟验证攻击，例如指针转换、二分搜索或泄露内核[76]。我们在内存控制器中区分代码和数据请求，并禁用代码请求的延迟验证。通过在取指组件或末级指令高速缓存发出期间在指令请求包中添加 isInst 标志，可以轻松实现所需的区分。具体来说，在本文中，指示符标志由指令缓冲区提供（假设类似于 TPU 架构 [45]）。

TensorTEE通过确保在通信之前完成对涉及的张量的验证来保证NPU通信数据的完整性。这是通过两个步骤完成的：张量毒害追踪和验证屏障。首先，从[57]中汲取灵感，我们为张量引入了一个毒物位来跟踪毒物传播，如图 14 © 所示。未经验证的张量的毒害位设置为 1，潜在的毒害效应传播到输出张量。验证完成后，毒害位被清除。其次，我们介绍一下在通信之前的代码中插入pragma verification_barrier，如图14(a)所示。该编译指示在编译时生成完整性同步指令，该指令会阻止后续通信指令，直到清除相关张量的有害位为止。为了自动插入同步指令并减少手动开销，我们在编译时识别特定函数，例如 .cpu()。与图14（b）相比，这种张量毒害追踪+验证屏障机制成功有效地防止了张量被篡改导致的通信数据损坏。此外，通过计数器限制未验证张量的数量，以避免检测到验证失败后进行无意义的计算。

Security analysis

鉴于XOR运算不会减少MAC的输出空间（56位），张量粒度上的防伪能力并没有明显下降[25, 85]。反复“盲猜”仍需要近256次尝试才能伪造成功。

Direct Data Transfer Protocol Implementation

在本节中，我们提出了一种高效的异构 TEE 数据传输协议的实现，该协议消除了重新加密和解密的需要，并实现了计算和数据传输之间的并行。

Optimizations Based on Unified Tensor Granularity.

现有的异构CPU和NPU TEE具有不同的内存保护粒度，这导致Enclave之间的数据传输产生巨大的开销：首先，数据传输依赖于非安全的内存区域中继，这会产生重新加密的开销，如图6所示（a）。其次，由于计算和数据传输（重新加密）都是内存密集型的，有限的 AES/内存带宽使得数据传输与计算不并行，如图 7 所示。

在这项工作中，我们根据提出的统一张量管理粒度减少了数据传输开销。一方面，CPU和NPU enclave具有统一的张量数据结构，这意味着片外加密存储器可以在双方上解密。这样，两者之间的数据传输不需要依赖于非安全区域的传输，这使得安全区域之间的直接数据传输成为可能，如图6(b)所示。另一方面，在消除了 AES 数据传输需求后，我们打破了 AES/内存带宽限制，这阻碍了并行调度优化。如图15所示，基于统一的管理粒度，我们将数据传输隐藏在计算中，大大减少了数据传输的开销。

Implementation

该协议主要由两个阶段组成：身份验证和数据传输。

Authentication phase

CPU和NPU通过远程认证建立信任。首先，在工作负载执行开始时，CPU启动enclave创建过程，将代码和数据副本从非安全内存复制到安全内存，并计算enclave的报告。其次，CPU enclave向NPU发送enclave创建请求，创建NPU enclave，流程类似。第三，两个飞地根据报告相互授权。为了实现低成本的数据传输，CPU和NPU TEE持有相同的加密/解密密钥。证明后，两个飞地执行类似 Diffie-Hellman 的密钥交换协议，该协议在两个飞地中启用相同的密钥，而不会在通信过程中泄漏密钥。交换后，密钥始终保存在片内存储器中，以确保安全。

Communication phase

有两种类型的传输：可信元数据传输和直接张量传输。我们以CPU到NPU的数据传输为例来描述这个过程。如图8所示，在元数据传输阶段，NPU发送包含张量地址范围的数据传输请求。 CPU内存控制器查询元表以获取元数据。获得的张量VN、MAC、地址通过可信加密通道传输到NPU进行后续解密和验证。同时，该地址范围对应的张量数据通过直接通道直接从CPU DRAM传输到NPU的DRAM，而不需要CPU和NPU的参与。值得注意的是，这两个阶段可以并行执行，只需要在传输完成后使用同步机制。

Experimental Methodology

Simulation Infrastructures

我们使用增强型 CPU 模拟器 Gem5 [8]（与 Ramulator [51] 集成的实现的周期精确 NPU 模拟器）以及具有测量带宽的建模系统通信协议来执行详细模拟。此外，我们分别与SGX CPU和A100 GPU进行性能比对，以确保我们定制的Gem5和NPU模拟器的时间模型的正确性。

CPU Simulator

为了评估基线 SGX 和 TensorTEE 上 CPU 工作负载的性能，我们使用类基线 SGX 和我们提出的张量内存保护机制扩展了 Gem5 [8] 模拟器。CPU 配置如表 1 所示。Gem5 是一款周期精确的事件驱动 CPU 性能模拟器，支持多核、全系统和内存模拟。我们实现了一个类似于 Intel SGX MEE [24] 的方案作为基准，其中包括每 64 字节缓存线一个 56 位 VN 和 MAC 以及一个 8 元 Merkle 树。对于SGX，到达内存控制的每个内存请求都会触发VN和MAC的额外访问。 MEE对读请求进行解密，对写请求进行加密。 MEE增加元数据缓存，缓解访问VN的访问压力。对于 TensorTEE，我们实现并模拟了元表和张量过滤器扩展。为了满足协同计算的要求，我们选择了多核乱序CPU和多通道大容量DRAM。

NPU Simulator

为了评估 NPU 部分计算的性能，我们实现了与 DRAM 模拟器 Ramulator [51] 集成的周期精确模拟器，并通过内存保护机制对其进行了扩展。 NPU配置如表1所示。所实现的NPU模拟器采用TPUv3架构，采用输出固定数据流，并结合自动平铺和层间优化。通过仿真比较，我们的模拟器实现了与A100 GPU相似的计算性能。我们使用类似于 MGX [34] 的片上 VN 管理和 NPU 与 DRAM 之间的内存保护组件来扩展 NPU 模拟器。

Communication

我们对基本的类 Graviton [84] 协议和我们的直接数据传输协议进行建模以进行性能评估，并将通信模拟与 CPU 和 NPU 模拟器连接起来。我们采用常用的PCIe 4.0*16作为CPU和NPU之间的通信总线。

Workload

TensorTEE 适用于任何基于张量的应用程序。在本文中，我们以最具代表性和挑战性的工作负载之一——LLM训练来进行评估研究。对于LLM训练，我们选择了几个流行的模型作为我们的工作负载，参数数量从100M到7B不等。模型及其配置如表 2 所示。我们采用 DeepSpeed 框架的调度策略，并针对协作异构计算进行卸载优化。在本文中，我们重点关注具有主机多核CPU和单个NPU的LLM训练场景。为了适应NPU的内存容量，我们为不同的模型选择不同的batch size。此外，我们选择带有平铺优化的 2D 矩阵乘法（GEMM）作为 CPU 工作负载，以展示 TensorTEE 在具有复杂访问模式的工作负载上的有效性。

我们在实验中比较了三种类型的配置： 1）非安全：在此设置中，CPU 和 NPU 均不利用 TEE 等隔离和内存保护机制。此配置用作性能参考。 2）CPU SGX + NPU MGX 配置：在此设置中，CPU 采用类似 SGX 的缓存行粒度保护机制，而 NPU 采用类似 MGX 的张量粒度保护机制。 CPU和NPU之间的通信由于粒度不匹配，需要额外的加密和解密。 3）TensorTEE（我们的）：在此设置中，我们利用论文中提出的基于张量的统一管理方法和优化技术来实现增强的性能，同时保持安全性。

Experiment Results

System Performance Improvement

Improvement of Tensor-Wise CPU TEE

Improvement of Delayed Verification

Improvement of Data Transfer Protocol

Hardware Overhead

Related Works

Privacy-preserving computing on LLM

由于对训练数据和模型隐私的担忧，人们提出了各种安全技术。具有乱码电路和秘密共享的 MPC [3,17,18,26,32,60]，擅长 ReLU 计算等布尔逻辑电路，但不擅长矩阵乘法等算术逻辑。差分隐私（DP）[19,31,58,59,65]在训练数据或梯度中添加高斯噪声等噪声，以提高其隐形性，使攻击者无法从差分推理中获取数据信息。尽管提高了性能，但 DP 由于添加了噪声而降低了模型精度。 FHE [29,38,61,66,74]旨在使用加密数据或模型进行计算，由于密文扩展和繁重的操作而导致不可接受的开销。 TEE [13,79,87]是一个安全区域，旨在保护数据的机密性和完整性加载到受保护内存区域的数据、模型和代码，通过内存加密、消息认证码 (MAC) 和Merkel Tree实现。与其他隐私保护技术相比，TEE 可以提供强大的隐私计算保证，同时产生较低的开销，因此受到了广泛的关注。与基于 VM 的 TEE 相比，SGX 等基于 Merkle-Tree 的 TEE [13, 22] 提供增强的安全性，并在区块链 [11, 14]、ML [50, 89] 和金融 [40, 41] 等高安全场景中保持竞争力。

Memory encryption and integrity

人们提出了许多工作来改进CPU TEE，包括反模式加密[81]、减少VN大小的优化[73, 86]、元数据缓存[23, 54]以及预测VN或推测性地使用未经验证的VN[ 57, 75, 76]，可扩展性增强[21, 22]。但它们都需要片外存储器中的版本号，并且具有大量额外的存储器访问开销。 SoftVN[83]中的VN由软件显式指定并编译成指令，从而消除了片外访问。然而，其静态数据流假设限制了其适应动态和复杂场景的能力。最近的工作已将CPU TEE扩展到NPU TEE，涉及建立可信通信通道[4,43,84]和基于张量的VN管理[33,34,56,62,78]，预测MAC[2]。然而，这些研究未能解决与CPU粒度级别的兼容性问题，导致协作计算中的通信开销巨大。

Conclusion

针对高效、安全的协作异构计算，本文提出了一种针对 CPU 和 NPU 的统一张量安全内存管理方法，具有对 CPU TEE 的虚拟支持、张量 MAC 的延迟验证和直接数据传输协议。这些解决方案可在 LLM 协作计算中实现高效的片上内存访问和异构数据传输，最终增强端到端性能。