[논문리뷰] TPLA: Tensor Parallel Latent Attention for Efficient Disaggregated Prefill & Decode Inference

링크: 논문 PDF로 바로 열기

저자: Xiaojuan Tang, Fanxu Meng, Pingzhi Tang, Yuxuan Wang, Di Yin, Xing Sun, Muhan Zhang

핵심 연구 목표

본 논문은 DeepSeek-V2 에서 도입된 Multi-Head Latent Attention (MLA) 이 Tensor Parallelism (TP) 환경에서 KV 캐시 메모리 절감 효과를 잃는 문제를 해결하고자 합니다. 특히, TP 환경에서 각 디바이스가 전체 latent vector (cKV) 를 로드해야 하는 비효율성을 개선하여, MLA의 압축 이점과 TP 효율성을 동시에 달성하면서도 표현 능력(representational capacity) 을 유지하는 것을 목표로 합니다.

핵심 방법론

제안하는 Tensor-Parallel Latent Attention (TPLA) 은 latent representation 과 각 헤드의 입력 차원을 디바이스 간에 분할하고, 각 샤드에서 독립적으로 어텐션을 수행한 후 all-reduce 로 결과를 결합합니다. TPLA 는 각 어텐션 헤드가 전체 latent representation 을 활용하게 하여 표현 능력을 유지하며, 디바이스는 KV 캐시의 파티션만 로드 합니다. 또한, Hadamard transform 또는 PCA 와 같은 직교 변환 을 RMSNorm 및 softmax 연산에 적용하여 크로스-샤드 간섭을 완화하고 정확도 저하를 최소화합니다. Prefill 단계 에서는 MLA 방식을 사용하고 디코딩 단계 에서는 TPLA 를 사용하는 Prefill/Decode Separation 전략을 채택하여 각 단계의 효율성을 최적화합니다.

주요 결과

DeepSeek-V3 및 Kimi-K2 모델에서 32K 토큰 컨텍스트 길이 기준으로 디바이스당 KV 캐시를 감소시켜 각각 1.79배 및 1.93배의 속도 향상 을 달성했습니다. 이러한 성능 향상은 LongBench 및 commonsense benchmarks 에서 성능 저하 없이 이루어졌으며, FlashAttention-3 와 호환되어 실용적인 구현이 가능함을 보였습니다. 특히, PCA 기반 재매개변수화(reparameterization) 는 RMSNorm과 softmax를 동시에 병렬화할 때 최상의 성능 을 일관되게 제공했습니다.

AI 실무자를 위한 시사점

TPLA 는 MLA 기반 LLM 의 Tensor Parallelism 추론 효율성을 크게 향상시켜 장문 컨텍스트 추론 비용 을 절감할 수 있는 실용적인 솔루션을 제공합니다. 기존 사전 훈련된 MLA 모델 에 재훈련 없이 적용 가능하여 도입 장벽이 낮으며, FlashAttention-3 와 같은 최적화 라이브러리와의 호환성으로 end-to-end 성능 향상 을 기대할 수 있습니다. Prefill/Decode 단계 분리 전략은 각 단계의 컴퓨팅 및 메모리 특성을 고려한 최적화를 가능하게 하여 전체 추론 파이프라인의 효율성을 높입니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.