Accelerating Cross-Model LLM Inference via Attention Map Reuse

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 generation quality across heterogeneous model pairs and multi-turn routing workloads.

Motivation

LLM-augmented data workflows often involve multiple heterogeneous models operating over substantial shared context, such as schemas, tool specifications, interaction history, and intermediate results to long, structured contexts, including database schemas, system constraints, tool specifications, prior interaction history, and intermediate results. When multiple requests share the same prompt prefix and are served by the same model Key-Value (KV) tensors can be reused to avoid redundant prefilling and reduce Time-to-First-Token (TTFT), as demonstrated by recent LLM serving systems such as vLLM [33] and SGLang [84]. However, this assumption does not hold in many LLM-augmented data systems, where a single data workflow often involves multiple heterogeneous LLMs or agents. Existing prefix caching mechanisms are largely restricted to identical model instances.

Attention Maps as Reuse States

The key insight is that while KV states are model-specific latent representations, attention maps encode token-to-token dependency structure in Transformer models. For models that share the same tokenizer and process the same context, these dependency patterns can exhibit stronger structural similarity across models than KV tensors do. MapFlow transfers attention maps from a sender model to a receiver model, uses lightweight linear translation to bridge differences in layers and heads, and combines this with partial recomputation to preserve semantic fidelity. This makes attention maps a more suitable reusable state for heterogeneous execution.

MapFlow System Design

MapFlow addresses the challenges of raw attention maps having quadratic size in context length and the overhead of directly reconstructing dense receiver-side attention. It is designed with three core components:

1. Partial Recomputation, where for the first r layers, the receiver model R performs standard computation to obtain accurate hidden states, and for subsequent layers (j > r), R skips the heavy attention computation and instead reuses the cached attention maps from the sender model S.

2. Layer Mapping, which maps each reuse receiver layer j to a sender layer m(j) and Aˆj(R) is sourced from A m(j)(S).

3. Head Mapping, which determines how each reuse receiver head Aˆj i(R) is computed via A m(j)(S), utilizing linear transformation as head mapping.

Practical Implementation Details

To manage memory bloat, MapFlow employs an Attention State Compressor that filters the attention map before it is cached, using a hardware-friendly block-level Top-p filtering scheme to store the attention map in a compact Block Sparse Row (BSR) format. It also implements a Block Manager that organizes the cache by maintaining a prefix-hash-based indexing system for efficient retrieval. The Attention Retrieval uses a custom GPU kernel that fuses the reconstruction of the receiver’s attention map with the matrix multiplication of the receiver’s value vectors, directly producing the attention output without intermediate dense materialization.

Experimental Results

Experiments show that MapFlow reduces TTFT by 15%–20% while maintaining generation quality, and improves end-to-end throughput by up to 25% under the same latency target. These results suggest that cross-model attention reuse can serve as an effective execution primitive for LLM-augmented data systems. The benefits are especially pronounced in long-context, multi-turn settings, where repeated workflow-level context amplifies the cost of standard prefilling. MapFlow achieves lower P90 TTFT across all arrival rates and achieves at most 25% higher throughput (QPS) under same P90 TTFT constraint across different model pairs.

Conclusion

Ultimately, MapFlow provides a scalable, memory-efficient solution that significantly accelerates serving times for complex multi-model workflows. It demonstrates that attention maps provide a promising systems abstraction for heterogeneous prefix reuse. The system achieves a 15% to 20% reduction in Time-toFirst-Token and up to a 25% increase in end-to-end throughput for multi-turn routing, all while maintaining generation quality. It is particularly suitable for TTFT-sensitive, routing-heavy, and long-context workflows where repeated prefill constitutes a substantial fraction of execution time.

Acknowledgement

We thank the anonymous SIGMOD reviewers for their comments and helpful feedback. This work is supported by National Natural Science Foundation of China (U23B2048, U22B2037), Beijing Municipal Natural Science Foundation (Grant No. L2603008), and the Fundamental Research Funds for the Central Universities, Peking University. The corresponding authors are Bin Cui and Xupeng Miao.

References

[1] Amey Agrawal, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav S. Gulavani, and Ramachandran Ramjee. 2024. SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24). USENIX Association, 537–552.

[2] Anysphere Inc. 2025. Cursor: The AI Code Editor. https://www.cursor.com/. Accessed: 2025-04-23.

[3] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. 2024. LongBench: A Bilingual Multitask Benchmark for Long Context Understanding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). Association for Computational Linguistics.

[4] Butterfly Effect AI. 2025. Manus: AI Agent. https://manus.im/. Accessed: 2025-04-23.

[5] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. 2021. Scatterbrain: Unifying sparse and low-rank attention. In Advances in Neural Information Processing Systems, Vol. 34.

[6] Liana Chen et al. 2024. Punica: Multi-Tenant LoRA Serving. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS).

[7] Lingjiao Chen, Matei Zaharia, and James Zou. 2023. FrugalGPT: How to Use Large Language Models While Reducing Cost and Improving Performance. In Advances in Neural Information Processing Systems (NeurIPS).

[8] Zui Chen, Lei Cao, Sam Madden, Tim Kraska, Zeyuan Shang, Ju Fan, Nan Tang, Zihui Gu, Chunwei Liu, and Michael Cafarella. 2023. SEED: Domain-specific data curation with large language models. arXiv preprint arXiv:2310.00749 (2023).

[9] Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, Ion Stoica, and Eric P. Xing. 2023. Vicuna:

1. Bold header: Cross-Model Prefix Reuse via Attention Map Transfer

This improvement allows heterogeneous LLM serving systems to reuse attention maps instead of tightly coupled KV states, addressing the limitation that KV tensors are tightly coupled to model-specific architectures by transferring attention maps from a sender model to a receiver model.

2. Bold header: Lightweight Linear Translation for Head Mapping

To bridge architectural differences between models, the system uses lightweight linear translation across layers and heads via learned weights, formally defined as applying the transformation where Aˆji(R) = (Aⱼi(R), 1 ≤ j ≤ r Ali(S) · wⱼi to approximate the receiver's attention map.

3. Bold header: Hardware-Friendly Block-Level Sparse Compression

The system employs a hardware-friendly block-level top−p sparse compression scheme to store attention maps in a compact block-sparse layout, which reduces the quadratic space complexity of attention maps and ensures minimal loss in precision by retaining only tokens necessary to satisfy a cumulative probability threshold p.

4. Bold header: Fused Retrieval Kernel for Online Reconstruction

MapFlow implements a custom fused GPU kernel (Algorithm 2) that reconstructs the receiver’s attention and computes the output on the fly, without allocating global memory for intermediate dense states, effectively overcoming the bottleneck where explicitly materializing An(R) requires O (nL2) memory.

5. Bold header: Accelerated Prefix Caching in Data Workflows

The resulting system can serve prefix-heavy data workflows by reducing prefilling latency by 15%–20% and improving end-to-end throughput by up to 25%, which is especially beneficial in long-context, multi-turn settings where context accumulation amplifies standard prefilling costs.

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.