검색 기반 희소 어텐션 (Retrieval-Based Sparse Attention)을 위한 범용 프로세싱-인-메모리
요약
검색 기반 희소 어텐션을 지원하기 위해 KV 캐시를 GPU 외부로 재배치하는 범용 PNM 설계인 KARAT을 제안합니다. GPU와 PNM 노드를 분리하여 연산 효율을 높이고, 스케줄링 최적화를 통해 파이프라인 버블을 최소화합니다.
핵심 포인트
- KV 캐시를 외부로 재배치하는 이기종 디코딩 서빙 시스템 제안
- 범용 PNM 설계인 KARAT을 통해 다양한 희소 어텐션 알고리즘 수용
- OFMS 스케줄링 및 CMR 재균형을 통한 파이프라인 효율 극대화
- GPU 전용 대비 TDP당 처리량을 최대 6.13배 향상
본 논문은 최근의 프런티어 LLM들이 백만 토큰 컨텍스트를 서빙하기 위해 채택하고 있는 검색 기반 희소 어텐션 (Retrieval-Based Sparse Attention)에 착안하여, KV 캐시를 GPU 메모리 외부로 재배치하는 이기종 디코딩 단계 (decode-phase) 서빙 시스템을 제시합니다. 이 시스템은 디코딩 단계를 연산 유형에 따라 분할합니다. GPU 노드는 모델 가중치를 보유하고 프로젝션 (projections) 및 MoE 레이어를 실행하는 반면, 프로세싱-인-메모리 (PNM) 노드는 KV 캐시와 키 인덱스 (index keys)를 보유하며 이를 읽는 모든 연산을 실행합니다. 우리는 먼저 기존의 PIM 및 PNM 설계의 전제 조건이 이러한 연산에는 더 이상 유효하지 않음을 보여주고, 해당 노드를 위한 네 가지 설계 요구사항을 도출합니다. 이러한 요구사항을 바탕으로, 우리는 네 가지 조건을 모두 충족하는 설계 지점인 범용 PNM 설계인 KARAT (KV-cache-resident Accelerator for Retrieval-based ATtention)를 제안합니다. KARAT 장치는 대용량 LPDDR 용량과 검색 인덱서 (retrieval indexer)에 적합한 크기의 범용 연산 능력을 결합하여, 저강도 GEMV를 목표로 구축된 기존 PIM/PNM 설계의 연산 강도 (operational intensity)를 넘어서는 성능을 제공하며, 고정 기능 유닛 (fixed-function units)이 진화에 따라 지원할 수 없는 다양한 희소 어텐션 (sparse attention) 알고리즘을 수용합니다. 두 장치 유형이 마이크로 배치 (micro-batches) 사이에서 교차할 때 발생하는 파이프라인 버블 (pipeline bubbles)을 줄이기 위해, 우리는 추가적으로 기회주의적 미세 조정 마이크로 배치 스케줄링 (OFMS, opportunistic, fine-grained micro-batch scheduling)을 제안합니다. 이는 다른 마이크로 배치의 GEMM 연산 뒤로 전문가 올투올 (expert all-to-all) 연산을 숨깁니다. 또한 컨텍스트 길이의 편차에도 불구하고 토큰 수를 균등하게 맞추는 컨텍스트 길이 인식 마이크로 배치 재균형 (CMR, context-length-aware micro-batch rebalancing)을 제안합니다. 세 가지 최첨단 모델과 실제 에이전트 트레이스 (agentic traces)를 통해 실험한 결과, 제안된 시스템은 서비스 수준 목표 (SLO) 하에서 GPU 전용 베이스라인 대비 TDP당 처리량을 2.09-6.13배 향상시켰으며, 훈련이 필요 없는 (training-free) 희소 어텐션 방법들을 1.36-3.21배 개선된 성능으로 실행합니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.AR의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기