Reconstructing KV Caches with Cross-Layer Fusion for Enhanced Transformer

⚠️ Notice

The goal of this release is to provide the essential code for researchers to understand and reproduce the core methodology described in our paper.

Please note that this is not a standalone project. This repository only contains the core modifications and key components built upon the NVIDIA Megatron-LM framework. For internal policy reasons, we are unable to release the full training infrastructure. We appreciate your understanding. 🙏

🎯 Core Contributions

This codebase includes the implementation of the following key contributions from our paper:

• FusedKV/FusedKV-Lite: The main implementation can be found in CrossLayerAttention class within attention.py.
• Layout of KV Share: We provide a cross-layer-attention configuration for different methods in cross_layer_attention_layout.py.
• Fusion Coefficient: The implementation of fusion coefficient can be found in layer_kv_coeff.py.
• Inference: The implementation of fusion coefficient can be found in inference directory.

🧩 Intergration

To use this code, you must first set up the Megatron-LM environment and then integrate our modifications. The file structure should look like this after copying from this repository into your cloned Megatron-LM:

Megatron-LM/megatron/core/transformer/
    ├── 🔄 attention.py                           <-- REPLACED by this repo
    ├── 🆕 cross_layer_attention_layout.py        <-- ADDED by this repo
    └── 🆕 layer_kv_coeff.py                      <-- ADDED from this repo
... (other Megatron-LM files)

Beyond integrating the core files listed above, some minor modifications to the existing Megatron-LM codebase may be necessary to support our method. These changes are typically small and context-dependent, varying based on the specific version of Megatron-LM you are using and your own environment.

🚀 Inference

We provide a Triton-based implementation of FusedKV, designed for efficient LLM inference.

Benchmarking

For MHA/GQA:

python benchmark_gqa.py --config-path model.json --attn gqa

For FusedKV:

python benchmark_fusedkv.py --config-path model.json --attn fusedkv

📖 Citation

@misc{lin2026reconstructingkvcachescrosslayer,
      title={Reconstructing KV Caches with Cross-layer Fusion For Enhanced Transformers}, 
      author={Hongzhan Lin and Zhiqi Bai and Xinmiao Zhang and Sen Yang and Xiang Li and Siran Yang and Yunlong Xu and Jiaheng Liu and Yongchi Zhao and Jiamang Wang and Yuchi Xu and Wenbo Su and Bo Zheng},
      year={2026},
      eprint={2512.03870},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2512.03870}, 
}