RNN-Based Stitching Mechanism in LLMs

• RNN-Based Stitching Mechanism is a technique that uses centroid summarization and token-level refinement to address long-range dependencies in large language models.
• It implements a two-stage retrieval process—centroid-level recall followed by token re-ranking—to drastically reduce memory use and computational load compared to dense full attention.
• Empirical results show 3–4× speedups and less than 1% accuracy loss, demonstrating its scalability for processing sequences with up to 10^5–10^6 tokens on limited hardware.

A Recurrent Neural Network (RNN)-Based Stitching Mechanism refers to a class of architectural and algorithmic techniques for efficiently handling long-range dependencies and massive contexts in sequence modeling, particularly in long-context LLMs. The principal objective is to overcome the memory and computational inefficiency endemic to dense attention in conventional autoregressive transformer inference by introducing a two-stage retrieval paradigm based on centroid summarization and token-level refinement. This mechanism deploys classical RNN-like ideas of state summarization via compact centroids, coupled with efficient index-based retrieval to selectively reconstruct relevant context at each decoding step.

1. Motivation and Problem Setting

In long-context LLMs, the naive retention and processing of all past Key-Value (KV) pairs for attention scaling linearly with context length $N$ yield prohibitive memory and latency costs, especially as sequence lengths approach $10^5$–$10^6$ tokens. The RNN-based stitching mechanism is motivated by three core observations:

• KV-Cache Memory Explosion: Storing all KV pairings on GPU quickly exhausts memory budget, bottlenecking throughput and restricting practical sequence lengths.
• Memory-Bound Full Attention: Every token generation at inference requires $O(Nd)$ computation and memory bandwidth, rendering full dense attention infeasible for long sequences.
• Query Redundancy After Positional Encoding: Query vectors $q_t$ and $q_{t+1}$, after Rotary Position Embedding (RoPE), remain highly similar (cosine similarity > 0.8 even over extended spans), indicating redundancy in the context each new query requires (Lu et al., 17 Dec 2025).

This context redundancy is reminiscent of RNN dynamics, where hidden states summarize temporal information without storing explicit token-wise history.

2. Centroid-Based Summarization

The mechanism initiates with centroid construction during a prefilling phase. Rather than tracking all queries or keys throughout the sequence, the last $C$ query vectors (post-RoPE) are extracted as centroids $$Q^{(c)} \in \mathbb{R}^{b \times g \times C \times d}$$, where $b$ is batch size, $g$ is the number of grouped-query attention (GQA) key-value head groups, $d$ is embedding dimensionality, and $C \ll N$ is the number of centroids.

For each centroid, an inverted file index ($qcIVF$) is built by recording the indices of the top-$\rho$ most similar Key vectors (according to the softmax similarity within each KV group). This process is both computationally lightweight and robust to changes in $C$, bypassing costly global approximate nearest-neighbor (ANN) indexing (Lu et al., 17 Dec 2025). The resulting centroid-key mapping compactly represents recent context, playing an analogous role to RNN hidden states but with explicit anchor points for content retrieval.

3. Two-Stage Retrieval and Stitching

At each decoding (generation) step, context "stitching" is performed in two hierarchically ordered phases:

Stage 1: Centroid-Level Recall

Given a new query $q_t$, cosine similarities to the set of centroids $Q^{(c)}$ are computed, yielding $Cos_q$. For each KV group, the $C'$ most similar centroids are identified. All unique (deduplicated) Key indices referenced by these centroids (via $qcIVF$) are then gathered, resulting in a candidate pool of $L^{Recall} \ll N$ keys.

Stage 2: Token-Level Re-Ranking

A finer attention step is then performed: softmaxed similarity scores between $q_t$ and the recalled Keys produce $A'_q$, followed by a max-over-head reduction to $A'_{kv}$. The top-$\rho'$ Keys are then selected for each group for final attention computation.

This two-stage filter-refine dynamic sharply decreases the read bandwidth and computational load, as only a small, adaptively chosen subset of context is reconstituted ("stitched") for each query—a procedure that functionally resembles RNN hidden state updates but with explicit data-driven context reconstruction. The full algorithm is detailed as Algorithm 1 in (Lu et al., 17 Dec 2025), with CPU–GPU co-execution offering additional efficiency.

Stage	Operation	Output (per step)
Centroid Recall	Cosine similarities, top-$C'$ select	$L^{Recall}$ candidate Key indices
Token Re-rank	Softmax + max-over-head, top-$\rho'$	$L^{Rerank}$ (final attention set)

4. System Implementation: Memory Layout and CPU–GPU Overlap

All centroid and inverted index data structures ($qcIVF$) are modest (typically $C \leq 2048$, $\rho \leq 1024$ per centroid), allowing their joint residency in both GPU and CPU memory. The full offloaded Key and Value caches occupy CPU DRAM, while only a "local" static $(L^{init} + L^{local})$ segment remains on the GPU for fast in-place attention.

During decoding, the GPU handles centroid similarity computations and static attention, whereas the CPU executes sparse gathering, softmax, and final attention among the recalled dynamic context. Overlapping these compute streams using multi-stream CUDA kernels amortizes data transfer and hides much of the end-to-end latency (Lu et al., 17 Dec 2025). Notably, this design accommodates hardware with limited GPU VRAM by streaming only relevant context on demand.

5. Empirical Results and Comparative Performance

Extensive benchmarking with Llama-3-8B (262K context) and Yi-9B (200K context) models at $96$K token context length shows that the RNN-based centroid stitching approach—specifically the Centroid Refinement Decoder (CRD) in CTKVR—delivers 3–4$\times$ throughput speedups compared to standard full-KV attention, with less than 1% decrement in accuracy across RULER, LongBench, and Needle-in-a-Haystack benchmarks.

Model	FullKV (A6000-48GB)	CRD (A6000-48GB)	Speedup
Llama-3-8B	16 KTokens/sec	50 KTokens/sec	$\times$3.1
Yi-9B	10 KTokens/sec	40 KTokens/sec	$\times$4.0

Even on 24GB GPUs, the method enables sequence lengths at 96K and beyond—contexts infeasible for dense bridging—which suggests practical scalability to future LLM deployments (Lu et al., 17 Dec 2025).

6. Extensions and Adaptations

The principle of summarizing long-range sequences through centroid-anchored, RNN-like stateful compact memories generalizes beyond token-level LLM KV retrieval. Applications include:

• Document retrieval + attention: Centroids per document chunk enable selective passage retrieval and stitched context formation in retrieval-augmented generation.
• Streaming dialogue: Maintains a FIFO centroid buffer summarizing recent conversational turns, supporting scalable, low-latency inference.
• Multimodal contexts: Extends centroid-based stitching to joint text-video embeddings for efficient long-context cross-modal attention.

A plausible implication is that the RNN-based stitching paradigm can be transferred to any setting where memory and compute scale with sequence length, and temporal state summarization is meaningful.

7. Significance and Outlook

RNN-based stitching mechanisms—exemplified by centroid-driven retrieval architectures—achieve a synergistic blend of computational and memory efficiency, near-lossless accuracy, and hardware flexibility. Their hybridization of classic RNN-style summarization dynamics with explicit attention-based retrieval constitutes a structural advance for long-context LLM inference. As the reliance on increasingly long context windows intensifies, such mechanisms will likely underpin next-generation neural architectures that operate on $10^5$–$10^6$ token contexts with bounded memory requirements (Lu et al., 17 Dec 2025).