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

TL;DR

This paper tackles secure, efficient collaborative tensor computing across CPU and NPU by addressing granularity mismatches in existing TEEs. It introduces TensorTEE, a unified tensor-granularity architecture comprising a CPU-side TenAnalyzer for on-chip tensor VN management, a tensor-wise MAC regime with delayed verification on the NPU, and a direct data transfer protocol between secure enclaves to avoid re-encryption. The approach yields substantial end-to-end improvements for LLM training—averaging 4.0x faster than prior heterogeneous TEEs with only about 2% overhead relative to non-secure training—and scales with model size, making secure collaboration practical for large models. The work demonstrates that unified tensor granularity, combined with delayed integrity verification and direct secure transfers, can dramatically reduce memory-bandwidth bottlenecks while preserving security guarantees in heterogeneous CPU-NPU systems.

Abstract

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.

Figures (21)

• Figure 1: Zero-Offload dataflow. Light computation like weight update and high precision data like weights and optimization states are offloaded to CPU.
• Figure 2: Two types of memory protection schemes. Tensor-wise protection eliminates the off-chip memory access for VN since the per-tensor VNs could be stored on-chip.
• Figure 3: CPU TEE incurs large performance overhead and leads to the transition of the Adam workload from computation-intensive to memory-intensive due to additional memory access for metadata.
• Figure 4: Access characteristics. The data is stored and accessed in tensor format, with small numbers and large sizes.
• Figure 5: Breakdown of collaborative computing that GPT2-M model training with ZeRO-Offload. The communication occurs unacceptable overhead with TEE.