RAG Module-Aware Distributed Inference System for Heterogeneous AI Accelerators

Retrieval-Augmented Generation (RAG) is widely adopted to mitigate hallucinations in Large Language Models (LLMs) by leveraging external knowledge. More recently, a growing body of work has explored modularization of a RAG pipeline to enable distributed inference across heterogeneous AI accelerators. Such efforts, however, have given limited consideration to the alignment between module characteristics and accelerator architectures. As a result, a RAG-module placement may incur unnecessarily high inference costs. To address this limitation, this paper proposes a RAG module-aware distributed inference system that determines an energy-efficient RAG-module placement by considering two factors: 1) the execution characteristics of each module and 2) the placement strategies available to each module across heterogeneous AI accelerators. For the proposed system, RAG Module-Aware Taxonomy (RMAT) combines a module taxonomy with an associated profiling table to support RAG-module placement across heterogeneous AI accelerators, and RAG Module-Aware Machine Learning (RMML) is devised to determine candidate RAG-module placements. Experimental results show that the proposed system reduces energy consumption by 27.8%–69.1% and 42.7%–90.8% across request arrival rate and scaled latency ranges, respectively, compared with the dynamic placement strategy, while the static strategy does not satisfy every configuration in several ranges. Compared with exhaustive search, RMML has lower placement-selection cost for search spaces larger than approximately 15 and, as the search space expands, reduces the increase in placement-selection cost by $24.9\times$.