TensorTEE: Unifying Heterogeneous TEE Granularity for Efficient Secure Collaborative Tensor Computing

📄 arXiv: 2407.08903v1 📥 PDF

作者: Husheng Han, Xinyao Zheng, Yuanbo Wen, Yifan Hao, Erhu Feng, Ling Liang, Jianan Mu, Xiaqing Li, Tianyun Ma, Pengwei Jin, Xinkai Song, Zidong Du, Qi Guo, Xing Hu

分类: cs.CR, cs.AI, cs.AR

发布日期: 2024-07-12

备注: Accepted by ASPLOS 2024

DOI: 10.1145/3622781.3674168

💡 一句话要点

TensorTEE：统一异构TEE粒度，实现高效安全协同张量计算

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 可信执行环境 异构计算 张量计算 安全计算 NPU CPU 大型语言模型 TEE

📋 核心要点

  1. 现有异构TEE设计在CPU和NPU协同计算中效率低，主要原因是CPU和NPU之间存在精细且不同的内存粒度，导致额外的内存访问和存储压力。
  2. TensorTEE的核心思想是统一CPU和NPU的TEE粒度，通过虚拟张量粒度支持、张量粒度MAC管理和预测执行等技术，优化内存访问和数据传输。
  3. 实验结果表明，TensorTEE相比现有工作在LLM训练工作负载上性能提升4.0倍，与非安全训练相比仅增加2.1%的开销，具有实际应用价值。

📝 摘要（中文）

由于NPU和CPU的异构协同计算具有显著的性能优势，因此受到了广泛关注。为了确保计算过程中的数据保密性和完整性，可信执行环境（TEE）被认为是一种有前景的解决方案，因为它开销相对较低。然而，现有的异构TEE设计对于协同计算效率低下，这是由于CPU和NPU之间存在精细且不同的内存粒度。具体来说，CPU TEE的缓存行粒度因额外的内存访问而加剧了内存压力，NPU的缓存行粒度MAC增加了有限内存存储的压力。此外，跨异构enclave的数据传输依赖于非安全区域的中转，导致繁琐的重新加密和调度。为了解决这些问题，我们提出了TensorTEE，一种统一的张量粒度异构TEE，用于高效的安全协同张量计算。首先，我们在CPU TEE中虚拟支持张量粒度，通过检测和维护片上张量结构来消除片外元数据访问。其次，我们提出了具有预测执行的张量粒度MAC管理，以避免计算停顿，同时消除片外MAC存储和访问。此外，基于统一的粒度，我们实现了直接数据传输，避免了重新加密和调度困境。我们的评估建立在增强的Gem5和周期精确的NPU模拟器上。结果表明，与现有工作相比，TensorTEE将大型语言模型（LLM）训练工作负载的性能提高了4.0倍，与非安全训练相比，仅产生2.1%的开销，为LLM训练提供了实际的安全保障。

🔬 方法详解

问题定义：现有的异构计算系统中，CPU和NPU之间的协同计算面临安全性和效率的双重挑战。使用TEE可以保证安全性，但CPU和NPU的内存访问粒度不同（CPU通常是cacheline，NPU可能是更小的粒度），导致频繁的跨TEE边界的数据传输和重新加密，增加了计算开销，降低了效率。此外，NPU的片上存储空间有限，传统的MAC（Message Authentication Code）存储方式会占用大量空间，进一步限制了性能。

核心思路：TensorTEE的核心思路是统一CPU和NPU的TEE粒度，使其都以张量为单位进行操作。通过在CPU TEE中虚拟化张量粒度，减少了片外元数据访问。同时，采用张量粒度的MAC管理和预测执行，避免了计算停顿，并消除了片外MAC存储和访问。这种统一粒度的设计简化了跨异构enclave的数据传输，避免了不必要的重新加密和调度。

技术框架：TensorTEE的整体框架包括以下几个主要模块：1) CPU TEE中的张量粒度虚拟化模块，负责检测和维护片上张量结构，实现张量粒度的内存访问。2) NPU中的张量粒度MAC管理模块，负责生成和验证张量级别的MAC，确保数据的完整性。3) 预测执行模块，用于提前计算MAC，避免计算停顿。4) 跨异构enclave的直接数据传输机制，允许在CPU和NPU之间直接传输张量数据，无需经过非安全区域。

关键创新：TensorTEE的关键创新在于统一了异构TEE的粒度，使其都以张量为单位进行操作。这与传统的基于cacheline的TEE设计有本质区别。通过张量粒度的虚拟化和MAC管理，减少了内存访问和数据传输的开销，提高了整体的计算效率。此外，预测执行的引入进一步优化了性能，避免了计算停顿。

关键设计：TensorTEE的关键设计包括：1) 张量结构检测算法，用于在CPU TEE中自动识别和维护张量结构。2) 张量粒度MAC生成和验证算法，确保张量数据的完整性。3) 预测执行的调度策略，用于提前计算MAC，避免计算停顿。4) 跨异构enclave的直接数据传输协议，保证数据传输的安全性和效率。具体的参数设置和网络结构取决于具体的应用场景和硬件平台。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，TensorTEE在大型语言模型（LLM）训练工作负载上相比现有工作性能提升了4.0倍。更重要的是，与非安全训练相比，TensorTEE仅引入了2.1%的开销，这表明该方案在保证安全性的同时，能够实现接近原生性能的计算效率。这些结果验证了TensorTEE在实际应用中的可行性和优越性。

🎯 应用场景

TensorTEE适用于需要高性能和高安全性的异构计算场景，例如大型语言模型的训练和推理、联邦学习、隐私保护的机器学习等。该技术可以显著提高计算效率，同时保证数据的保密性和完整性，具有重要的实际应用价值。未来，TensorTEE可以扩展到更多的异构计算平台和应用领域，为安全计算提供更强大的支持。

📄 摘要（原文）

Heterogeneous collaborative computing with NPU and CPU has received widespread attention due to its substantial performance benefits. To ensure data confidentiality and integrity during computing, Trusted Execution Environments (TEE) is considered a promising solution because of its comparatively lower overhead. However, existing heterogeneous TEE designs are inefficient for collaborative computing due to fine and different memory granularities between CPU and NPU. 1) The cacheline granularity of CPU TEE intensifies memory pressure due to its extra memory access, and 2) the cacheline granularity MAC of NPU escalates the pressure on the limited memory storage. 3) Data transfer across heterogeneous enclaves relies on the transit of non-secure regions, resulting in cumbersome re-encryption and scheduling. To address these issues, we propose TensorTEE, a unified tensor-granularity heterogeneous TEE for efficient secure collaborative tensor computing. First, we virtually support tensor granularity in CPU TEE to eliminate the off-chip metadata access by detecting and maintaining tensor structures on-chip. Second, we propose tensor-granularity MAC management with predictive execution to avoid computational stalls while eliminating off-chip MAC storage and access. Moreover, based on the unified granularity, we enable direct data transfer without re-encryption and scheduling dilemmas. Our evaluation is built on enhanced Gem5 and a cycle-accurate NPU simulator. The results show that TensorTEE improves the performance of Large Language Model (LLM) training workloads by 4.0x compared to existing work and incurs only 2.1% overhead compared to non-secure training, offering a practical security assurance for LLM training.