HedraRAG: Coordinating LLM Generation and Database Retrieval in Heterogeneous RAG Serving
https://dl.acm.org/doi/10.1145/3731569.3764806
https://sigops.org/s/conferences/sosp/2025/schedule.html
summary
LLM이 답변을 제공하기 위해서는 검색과 생성의 단계를 거치게 되는데 검색과 생성은 각각 CPU와 GPU를 사용하기 때문에 작업 단계마다 사용하는 하드웨어 자원이 다르다. 작업이 단순했던 과거와 달리 점점 더 다단계 추론이 복잡해지고 목적에 따라 워크플로우의 구조가 다양해지기 때문에 리소스를 효율적으로 사용하기 어렵다. 또한 LLM은 토큰을 조금씩 생성해서 빠르게 진행되지만 벡터 검색은 한 번에 많은 쿼리를 배치 처리하기 때문에 LLM은 갑자기 벡터 검색이 필요해져도 검색이 끝날 때까지 stall 되어야 하고 LLM이 조금씩 보내는 요청 때문에 벡터 검색의 처리량은 떨어질 수 있다. 그래서 검색이 필요한 요청이 밀리거나 GPU가 under utilization 될 수 있다는 단점도 한계도 있다.
논문은 리소스를 효율적으로 사용하기 위해 각기 다른 작업에서 수행되는 검색과 생성은 병렬 처리가 가능하다는 단계 간 병렬성, 비슷한 작업이 반복되고 생성되는 결과 또한 비슷한 내용이 반복되는 요청 내 유사성, 유독 자주 검색되는 영역인 핫스팟이 생기는 요청 간 편향성이라는 세 가지 찬스를 발견했다.
hedraRAG는 리소스를 비효율적으로 활용하게 되는 기존의 RAG 시스템의 한계를 해결하기 위해 검색과 생성 단계를 하나로 설계하고 런타임에 동적으로 조정할 수 있게 했고 복잡한 워크플로우는 그래프 기반의 RAGraph로 표현했다. 노드와 엣지로 구성된 RAGraph를 이용해서 기존에 각각 큰 하나의 작업으로 처리되던 검색과 생성을 여러 서브 스테이지로 파이프라이닝하여 단계 간 병렬 처리를 이루고, semantic aware 재정렬과 추측 실행을 통해서 작업을 오버랩 시키고 ANNS 검색의 early termination을 유도해서 전체적인 레이턴시를 줄였다. 요청 간 편향성을 활용해 핫 스팟 클러스터만 GPU에 캐싱하는 부분 GPU 인덱스 캐싱으로 GPU 메모리 제약 문제를 해결한다.
strengths
- 미세하게 조정되는 서브 스테이지 파이프라이닝과 다이나믹한 배치로 기존에 가변 길이로 인해서 생기던 파이프라인의 stall 문제를 해결하고 토큰 단위로 디코딩하는 LLM과 배치 단위로 처리하는 벡터 검색의 불일치 문제를 해결할 수 있었다.
- 요청의 의미들이 비슷하다는 특징을 활용해서 기존의 정적 추측 방식보다 복잡한 워크플로우에서의 성능을 높였고 GPU 메모리 한계를 극복하기 위해 편향성을 보이는 핫 스팟 클러스터만 캐싱하고 비동기적인 업데이트를 통해 검색 효율을 극대화 시키는 등 heterogenous한 워크로드를 실행할 때 발생하는 효율성 문제를 해결했다.
- 기존 프레임워크 대비 성능 개선이 있는 것 뿐만 아니라 multistep, IRG처럼 더 복잡한 워크플로우에서 더 큰 성능 개선 효과를 보이고 특히 서로 다른 RAG 워크플로우가 동시에 실행되는 환경에서도 최대 5.5배의 레이턴시 감소와 3.3배 이상 처리량을 높였다.
weakness
- 서브 스테이지로 분할할 때 서브 스테이지의 최대 실행 시간의 예산을 결정하는 것이 성능에 큰 영향을 주는데 레이턴시를 줄이는 것과 스케쥴링이나 분할로 인한 오버헤드 간의 트레이드 오프를 잘 고려해야한다. 그런데 최적값을 찾는 것은 경험적인 측정과 복잡한 모델링이 필요하기 때문에 워크로드가 조금만 달라져도 성능이 크게 나빠질 수 있다는 문제가 있다.
- GPU 메모리는 LLM 가중치, key value 캐시, 부분 인덱스 캐시 공간으로 나뉘어서 할당되어야 하는데 검색 속도를 높이기 위한 부분 인덱스 캐시 공간을 너무 작게 할당하면 성능 개선이 적고 너무 크게 할당하면 LLM 추론을 위한 kv 캐시 공간도 부족해지기 때문에 스와핑을 일으켜서 LLM 생성의 속도를 낮추는 문제가 있다. 서버를 시작하기 전에 예상되는 요청 속도와 검색 처리량, 생성 처리량을 기반으로 최적화된 KV 캐시 크기를 찾아야 한다. 그러나 이렇게 정적으로 할당하는 방식은 실시간으로 요청 패턴이 달라지면 비효율적이다.
- 레이턴시를 줄이기 위해서 제안된 추측 실행은 예측 불가능한 LLM의 아웃풋과 동적인 워크플로우 패턴으로 인해서 실행 시점을 결정하기 어렵다는 단점이 있고 hedraRAG은 시스템 처리량 활용도가 낮을 때 추측을 시작하는 적응형 전략을 사용하게 되는데 이때 Tcurr 값을 경험적으로 추정하고 시스템 최대 처리량인 Tmax와 임계값 또한 경험적으로 추정하기 때문에 예측 불가능한 시나리오에서 일관적으로 높은 추측 정확도를 유지하며 일반화하기에는 어려움이 있을 수 있다.
detailed comments
RAGraph라는 그래프 기반의 방법을 도입해서 복잡하고 heterogenous한 RAG 워크플로우를 통일된 방식으로 표현하고 동적으로 최적화할 수 있게 만들었다는 점이 가장 핵심적인 기여라고 생각한다. 튜닝하기만 한 것이 아니라 시스템 구조 자체를 바꿈으로써 기존의 시스템이 단순한 2단계 파이프라인이었던 것과 다르게 복잡한 멀티 스테이지 워크플로우를 유연하게 정의할 수 있다. 오픈소스 프레임워크와 호환되는 그래프 구성 API를 제공해서 쉽게 hedraRAG을 적용하고 통합할 수 있게 지원하며 노드 분할, 재정렬, 엣지 추가, 의존성 rewiring처럼 다이나믹한 그래프 변환을 가능하게 해서 기존 스케쥴링 프레임워크가 스테이지 중심이었던 것보다 더 다양하고 세분화된 최적화를 할 수 있었다.