分享好友 Global News首页 Global News分类 切换频道

Peking University’s MapFlow Reuses Attention Across Models without Recomputing the KV Cache

2026-11-08 15:0010RoboInd Editorial TeamWeChat Official Account

Model A is responsible for planning, Model B is responsible for refining and summarizing; complex problems are given to strong models, simple problems are given to lightweight models, and multi-model collaboration provides a pathway for AI Agent applications that is both cost-effective and cost effective.

But the interaction between the agents, or models, hides a cost of repeated calculations.

In the previous round, the big model had just completed a long document; in the next round, the smaller model took over, and often had to re-calculate the same content: databases, structure tool descriptions, dialogue history, task background... these recurring contexts became a new prefill burden after switching models.

图片

If one model has already done the calculation, can it help another model process the same context faster?

MapFlow, developed by a team of assistant professors at PKU-DAIR's Data and Intelligent Laboratory in Beijing, offers a new approach to this problem: using attention graphs as an intermediate state for cross-model replication, through light mapping, scarcity compression and integration of GPU cores, allowing models of different sizes to share already calculated token dependency relationships.

In paper evaluation, MapFlow reduces prefill latency by 15% to 20% while maintaining the quality of the task generated; in multi-round dialog routing scenarios, the same P90 token increases throughput by up to 25% under delay constraints. The relevant paper has been accessed by SIGMOD 2027 and the code is open source.

图片

Why is it so hard to keep using caches instead of changing models?

For large models, it is necessary to process the input in context before starting to respond, which is called prefill. The longer the input, the higher the calculation cost, the easier it is to extend the user's waiting time for the first token.

Forecast caching has been able to reduce part of the repetition of computations: when the same model encounters the same forecast again, the KV Cache that was previously generated can be used. Historical messages in multi-round conversations, as well as system prompts shared by multiple requests, can benefit from it.

However, the cross-model scenario encounters new obstacles. KV Cache is closely related to the parameters and internals of the model itself. Even if the two models process exactly the same text, they may have different layers, number of attention heads and hidden dimensions, and the resulting KV states are not naturally compatible.

Therefore, although multi-modeling systems can work at the task level, it is still difficult to share the underlying computational results. Particularly when small models take over large model accumulation in the long context, the repeated prefill of model switching will weaken the efficiency gains that were originally intended to be achieved.

MapFlow has shifted focus to another category of intermediate states: Attention Map.

From the inside, the dependency relationship between going to tokens

The attention chart describes how much attention a token allocates to other tokens when computing. Its two main dimensions correspond to the location of the token in the input.

This provides a natural alignment basis for cross-model replication: when two models use the same tokenizer to process the same contexts, the positions in the attention graph can match. Even if vectors within the model are represented differently, there may still be usable similarities to the engravings depending on the relationship in the text.

It can be understood that two readers read the same material, forming different internal understandings, but may focus on the same key paragraphs and contextual connections. The latter reader can use the connections marked by the previous reader to speed up their own processing process.

MapFlow replicates the attention structure between these tokens. The receiving model still uses its own Value representation and subsequent calculations to generate its own output.

However, there are also differences in the focus charts of different models. In order to keep duplication effective, MapFlow combines partial recalculation and lightweight mapping.

The receiving model first normally calculates several layers before establishing its own hidden state; in subsequent layers suitable for replication, the attention map of the sending model is mapped to the receiving model.

The process only trains mapping weights, and the original parameters of the two large models remain frozen. In the Ministral 3-8B → 3B configuration given in the paper, the new weight is less than 0.001% of the model parameters received. The mapping relationship is obtained by offline calibration, which can then be used for continuous online requests.

When it comes to real savings in reuse, there are also storage and execution issues that need to be addressed

The attention chart is portable, but that doesn't mean it needs to be cached.

A direct problem is volume: the size of the complete attention matrix increases by the square of the length of the sequence. If a tight matrix is stored for each layer, each head, the context will soon bring enormous storage pressure. Reading and converting these states may also offset the computing time saved.

图片

MapFlow designed three interconnected system components around this problem.

First, the importance of retaining the sparse blocks. Attention is not uniformly distributed in all locations. MapFlow selects important areas based on the accumulated attention probability, and focuses on compressing the sparse blocks in a compact format. The adoption of two-dimensional blocks with rules also makes subsequent access and computation more suitable for GPU execution.

Second, the cache is managed by a shared cache. The number of sparse blocks stored in different layers and locations is not the same. MapFlow uses the cache hash to index these states that contain context information in the preceding sequence, so that the receiving model finds the corresponding cache blocks, and the process of retrieving the cache overlaps with the previous layers. The same sending end state can also serve multiple downstream models.

Third, execute the mapping with the integration of Value. If the complete attention matrix of the receiving model is first restored and then multiplied by Value, a large amount of intermediate data is reintroduced. The fusion core of the MapFlow directly reads the sparse blocks, completes the linear combination between the attention heads on the film, and immediately multiplies and adds the output with the Value of the receiving model, thus reducing the display read-write and intermediate volume expense.

The system is integrated into vLLM, integrates retrieval kernel with Triton implementation, and supports mechanisms such as PagedAttention, Chunked Prefill, and variable sequence packing.

Experiments: From faster prefill to higher online throughput

In an environment equipped with the NVIDIA A800 80GB GPU, the team evaluated the two Qwen3 and Ministral 3 model families, covering four groups of different sized model pairs, as well as queries, conversations, and summaries.

图片

On prefill delay and quality balance, MapFlow reduces the prefill delay by 15% to 20% compared to the full prefill and maintains the generating quality in the corresponding evaluation tasks; compared to FlexPrefill, it further reduces the prefill delay by 7% to 12% compared to the equivalent or better configuration of the quality.

In multi-round conversation routes, the team used LMSYS and ShareGPT4 to construct analog routing loads for conversation data. Under the same P90 token delay constraint, MapFlow's request throughput increased by up to 25%. This benefit corresponds to online service capacity: the system can handle more requests on the condition that it meets the response time requirement.

In the three model-series workflows, the 14B model was responsible for planning, and the 8B and 3B models were responsible for refining and summarizing, respectively, while the two downstream models shared the attentional states produced by the 14B model. In the experiment, the workflow decreased by 0.48 to 0.73 seconds per shift, with end-to-end acceleration reaching up to 5.4%.

Adding a layer of computational reuse for multi-model collaboration

From model routing to Agent workflows, multi-model systems need to decide by whom to handle tasks, and they also need to handle repeated computations when models are communicated.

MapFlow demonstrates the value of attention graphs as an intermediate state between models: in addition to passing text, dependency relationships can be replicated between different models that have already been calculated for the same context. This replication can be converted into actual delay and throughput gains through mapping, compression, cache management and integration.

Let the context calculation be used for multiple subsequent models. This is one possibility that MapFlow opens up for efficient multi-model reasoning.


点赞 0
举报
收藏 0
分享 1