Accelerating Cross-Model LLM Inference via Attention Map Reuse

summary

Video file (mp4)

The gist

The gist The gist MapFlow, a system for efficient heterogeneous LLM serving via cross-model attention map reuse, reduces prefilling latency and improves end-to-end serving efficiency while preserving

In short

MapFlow is a system that reuses attention maps between different, heterogeneous LLMs to speed up inference. It addresses the inefficiency of standard prefix caching when multiple models share context but differ in architecture. By transferring and adapting these attention patterns, MapFlow reduces prefilling latency by 15-20% and boosts throughput by up to 25% while maintaining generation quality.

Key concepts

Attention Maps
These are internal representations within a Transformer model that encode the token-to-token dependency structure. The paper suggests they are more structurally similar across different models than standard Key-Value (KV) tensors, making them a viable state for cross-model reuse.
Partial Recomputation
This technique involves running the receiver model on the initial layers of a workflow to get accurate hidden states. For later layers, instead of performing expensive attention calculations, the system skips them and reuses previously cached attention maps from the sender model.
Layer Mapping and Head Mapping
These components manage how to translate attention maps from one model (sender) to another (receiver). Layer mapping links receiver layers to specific sender layers, while head mapping uses linear transformations to determine how individual attention heads are computed in the receiver model.

Terminology used across episodes

This episode discusses

The paper

Accelerating Cross-Model LLM Inference via Attention Map Reuse · Read on arXiv

Zhan Jin, Yuze Luo, Bin Cui, Xupeng Miao

Peking University

LLM-augmented data workflows often involve multiple heterogeneous models operating over substantial shared context, such as schemas, tool specifications, interaction history, and intermediate results. Although prefix caching can reduce prefilling overhead, existing approaches are largely limited to reuse within the same model because KV states are tightly coupled to model-specific architectures. We present MapFlow, a system for efficient heterogeneous LLM serving via cross-model attention map reuse. Instead of reusing KV states, MapFlow reuses attention maps, which capture token dependency patterns and are more transferable across models. To enable efficient and accurate reuse, we combine partial recomputation with lightweight linear translation across layers and heads. To make attention reuse practical for long contexts, we further introduce block-level sparse compression and a fused retrieval kernel. Experiments show that MapFlow reduces prefilling latency and improves end-to-end serving efficiency while preserving generation quality across heterogeneous model pairs and multi-turn routing workloads. These results suggest that attention maps provide an effective reusable state for heterogeneous LLM serving, especially in prefix-heavy data workflows. Our code is available: https://github.com/PKU-DAIR/MapFlow.

DOI: 10.1145/3856312

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Accelerating Cross-Model LLM Inference via Attention Map Reuse".

Jane: The gist The gist MapFlow, a system for efficient heterogeneous LLM serving via cross-model attention map reuse,

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So we just talked about how this paper, "Accelerating Cross-Model LLM Inference via Attention Map Reuse," proposes using attention maps instead of KV states to speed up serving when you use different models in one workflow. The main thesis is that attention maps capture the token dependency structure better than KV states for cross-model transfer.

Jane: Right, so the authors point out that existing prefix caching methods are mostly limited to reusing KV tensors because they are too model-specific, and MapFlow tries to solve this by reusing attention maps across different models.

Lu: The paper claims that by combining partial recomputation with lightweight linear translation across layers and heads, you can enable efficient and accurate reuse of these attention maps even when the models have different internal structures.

Meng: So it’s not just a direct copy-paste of maps, but there's this mapping process to handle the differences between layers and heads, right? That sounds like a lot of bookkeeping on top of the core idea.

Tom: Exactly. They introduce Layer Mapping to map each receiver layer j to a sender layer m(j), and they get that attention map by sourcing it from A m(j)(S) where S is the sender model.

Jane: And then there’s Head Mapping, which figures out how each reuse receiver head Aˆj i(R) is computed using a linear transformation applied to the source attention map A m(j)(S).

Lu: This entire design addresses the issue of raw attention maps having quadratic size in context length and the overhead of directly reconstructing them on the receiver side.

Meng: That sounds like they're tackling the memory bloat problem head-on, which is huge when you're dealing with long contexts. How do they handle storing those maps compactly?

Tom: They tackle that with an Attention State Compressor that filters the attention map before it’s cached, using a hardware-friendly block-level Top-p filtering scheme to store the attention map in a compact Block Sparse Row format.

Jane: And for retrieval, they have this custom GPU kernel that fuses the reconstruction of the receiver’s attention map with the matrix multiplication of its value vectors, which avoids creating any intermediate dense material.

Conclusion: Tom: So looking at "Accelerating Cross-Model LLM Inference via Attention Map Reuse," the authors are arguing that cross-model prefix reuse is a bottleneck in complex AI data workflows, and they provide a system called MapFlow to tackle it with attention map reuse.

Jane: In simple terms, what this means for the listener is that when you have multiple different AI models working together on a long task—say, one model plans and another refines—this system can speed up the initial response time by fifteen to twenty percent.

Lu: It suggests that attention maps are a viable intermediate execution state for heterogeneous LLM serving, moving away from the tight architectural coupling of reusing KV states and offering more flexibility across models.

Meng: For someone who actually runs this stuff in production, it means you get a tangible performance boost on prefilling latency and throughput when you’re running those multi-model routing assistants.

Tom: It’s about finding a way to make that reuse practical for long contexts, and the results suggest it can serve as an effective execution primitive for those kinds of workflows.

Jane: So, the core idea is that attention maps are a promising way to achieve prefix reuse when you need to switch between different models in a sequence.

More episodes

← Home