LLM 추론을 위한 통신 효율적인 검증 가능한 어텐션

Communication-Efficient Verifiable Attention for LLM Inference

원격에서 제공되는 대규모 언어 모델(LLM)의 계산 무결성은 문제가 될 수 있습니다. 기존의 심층 신경망(DNN)의 경우, Trusted Execution Environment (TEE)로 보호된 DNN 분할(TSDP) 방식은 TEE를 사용하여 비선형 연산을 수행하고 GPU로 오프로드된 선형 연산의 무결성을 검증합니다. 그러나 트랜스포머 기반 LLM에 TSDP를 직접 적용하면 상당한 TEE 계산 및 TEE-GPU 통신 오버헤드가 발생합니다. 본 논문에서는 검증 가능한 LLM 추론 속도를 향상시키기 위한 통신 효율적인 TEE-GPU 어텐션( extsc{VeriAttn})을 제안합니다. extsc{VeriAttn}은 어텐션의 선형 및 비선형 연산을 모두 GPU로 오프로드하고, TEE는 검증 작업을 수행합니다. 또한, 프레필 과정에서 extsc{VeriAttn}은 데이터 이동, TEE 전/후 처리 및 GPU 계산을 병렬화하기 위한 2단계 파이프라인을 사용합니다. 디코딩 과정에서 키-값 캐시가 사용 가능한 GPU 메모리를 초과하는 경우, extsc{VeriAttn}은 어텐션을 TEE와 GPU로 분할하여 반복적인 키-값 전송을 줄입니다. Intel TDX 플랫폼에서의 평가 결과, extsc{VeriAttn}은 6k 토큰 프롬프트 및 10k 토큰 출력에 대해 프레필 및 디코딩 과정에서 각각 TSDP보다 2.60-3.38배 및 3.86-5.42배의 성능 향상을 달성했습니다.

Computation integrity of remote large language model (LLM) serving can be questionable. For conventional deep neural networks (DNNs), the existing TEE-shielded DNN partitioning (TSDP) approach uses Trusted Execution Environment (TEE) to compute non-linear components and verify the integrity of linear components offloaded to an untrusted GPU. However, directly applying TSDP to Transformer-based LLMs incurs significant TEE computation and TEE-GPU communication overhead. This paper presents Communication-efficient TEE-GPU Attention (\textsc{VeriAttn}) for accelerating verifiable LLM inference. \textsc{VeriAttn} offloads both linear and non-linear computations of attention to the GPU, while TEE performs verification. Moreover, for prefill, \textsc{VeriAttn} uses a two-level pipeline to overlap data movement, TEE pre-/post-processing, and GPU computation. For decoding, when the key-value cache exceeds available GPU memory, \textsc{VeriAttn} partitions attention across TEE and GPU to reduce repeated key-value transfers. Evaluation on an Intel TDX platform shows that \textsc{VeriAttn} achieves 2.60-3.38$\times$ and 3.86-5.42$\times$ acceleration over TSDP for 6k-token prompts and 10k-token outputs during prefill and decoding, respectively.