Table of Contents
Fetching ...

SubGCache: Accelerating Graph-based RAG with Subgraph-level KV Cache

Qiuyu Zhu, Liang Zhang, Qianxiong Xu, Cheng Long, Jie Zhang

TL;DR

This work addresses the inefficiency of graph-based retrieval-augmented generation when processing in-batch queries that retrieve overlapping subgraphs. It introduces SubGCache, a subgraph-level KV caching framework that clusters queries by subgraph embeddings, builds a representative subgraph per cluster, and reuses a cluster-wide KV cache during inference. The authors demonstrate substantial latency reductions across two datasets and multiple backbones while preserving or improving generation quality, with up to about 6.68× TTFT speedups and notable ACC gains on Scene Graph. The approach is plug-and-play, model-agnostic, and scalable to large batch workloads, offering practical speedups for graph-based RAG in real-world settings.

Abstract

Graph-based retrieval-augmented generation (RAG) enables large language models (LLMs) to incorporate structured knowledge via graph retrieval as contextual input, enhancing more accurate and context-aware reasoning. We observe that for different queries, it could retrieve similar subgraphs as prompts, and thus we propose SubGCache, which aims to reduce inference latency by reusing computation across queries with similar structural prompts (i.e., subgraphs). Specifically, SubGCache clusters queries based on subgraph embeddings, constructs a representative subgraph for each cluster, and pre-computes the key-value (KV) cache of the representative subgraph. For each query with its retrieved subgraph within a cluster, it reuses the pre-computed KV cache of the representative subgraph of the cluster without computing the KV tensors again for saving computation. Experiments on two new datasets across multiple LLM backbones and graph-based RAG frameworks demonstrate that SubGCache consistently reduces inference latency with comparable and even improved generation quality, achieving up to 6.68$\times$ reduction in time-to-first-token (TTFT).

SubGCache: Accelerating Graph-based RAG with Subgraph-level KV Cache

TL;DR

This work addresses the inefficiency of graph-based retrieval-augmented generation when processing in-batch queries that retrieve overlapping subgraphs. It introduces SubGCache, a subgraph-level KV caching framework that clusters queries by subgraph embeddings, builds a representative subgraph per cluster, and reuses a cluster-wide KV cache during inference. The authors demonstrate substantial latency reductions across two datasets and multiple backbones while preserving or improving generation quality, with up to about 6.68× TTFT speedups and notable ACC gains on Scene Graph. The approach is plug-and-play, model-agnostic, and scalable to large batch workloads, offering practical speedups for graph-based RAG in real-world settings.

Abstract

Graph-based retrieval-augmented generation (RAG) enables large language models (LLMs) to incorporate structured knowledge via graph retrieval as contextual input, enhancing more accurate and context-aware reasoning. We observe that for different queries, it could retrieve similar subgraphs as prompts, and thus we propose SubGCache, which aims to reduce inference latency by reusing computation across queries with similar structural prompts (i.e., subgraphs). Specifically, SubGCache clusters queries based on subgraph embeddings, constructs a representative subgraph for each cluster, and pre-computes the key-value (KV) cache of the representative subgraph. For each query with its retrieved subgraph within a cluster, it reuses the pre-computed KV cache of the representative subgraph of the cluster without computing the KV tensors again for saving computation. Experiments on two new datasets across multiple LLM backbones and graph-based RAG frameworks demonstrate that SubGCache consistently reduces inference latency with comparable and even improved generation quality, achieving up to 6.68 reduction in time-to-first-token (TTFT).
Paper Structure (22 sections, 5 figures, 8 tables)

This paper contains 22 sections, 5 figures, 8 tables.

Figures (5)

  • Figure 1: Overview of graph-based RAG without and with caching.
  • Figure 2: Overview of SubGCache and its integration into the standard graph-based RAG pipeline.
  • Figure 3: Impact of cluster number on ACC (%) and TTFT (s).
  • Figure 4: Cluster processing time vs. LLM response time by varying cluster numbers.
  • Figure 5: Case study.