MInference：面向长上下文大模型的百万令牌提示推理

| 项目页面 | 论文 | HF 演示 | SCBench | MMInference |

https://github.com/microsoft/MInference/assets/30883354/52613efc-738f-4081-8367-7123c81d6b19

现在，借助 MInference 1.0，您可以在单张 A100 上以 10 倍的速度处理 100 万令牌的上下文，同时保持更高的准确性！使用 LLaMA-3-8B-1M、GLM-4-1M 等长上下文大模型，立即体验吧！

📰 新闻

• 🐝 [25/05/02] MMInference 已被 ICML'25 接受。
• 👨‍💻‍ [25/04/14] SGLang 和 vLLM 已合并 MInference 的稀疏注意力核。MInference 已经支持这些优化后的核。 只需运行 pip install sglang 即可。您将获得高达 1.64× (64K), 2.4× (96K), 2.9× (128K), 5.2× (256K), 8× (512K), 以及 15× (1M) 的加速效果。值得一提的是，SGLang 还将其适配到了 FlashAttention-3 上。特别感谢 @zhyncs 和 @yinfan98 的贡献！
• 👾 [25/04/23] 我们很高兴地宣布推出多模态工作 MMInference，该工作采用 模态感知的置换稀疏注意力 来加速长上下文视觉语言模型的推理。我们将在 Microsoft 展位 和 FW-Wild at ICLR'25 上展示 MMInference。新加坡见！
• 🤗 [25/01/27] MInference 已集成到 Qwen2.5-1M 及其在线服务中。详情请参阅 论文 和 vLLM 实现。
• 🪸 [25/01/23] SCBench 已被 ICLR'25 接受。

更多新闻

• 🍩 [24/12/13] 我们很高兴地宣布发布以 KV 缓存为中心的分析工作，SCBench,它从 KV 缓存的角度评估长上下文方法。
• 🧤 [24/09/26] MInference 被选为 NeurIPS'24 的 亮点。温哥华见！
• 👘 [24/09/16] 我们很荣幸地宣布发布 KV 缓存卸载工作，RetrievalAttention,它通过向量检索加速长上下文大模型的推理。
• 🥤 [24/07/24] MInference 现在支持 meta-llama/Meta-Llama-3.1-8B-Instruct。
• 🪗 [24/07/07] 感谢 @AK 的赞助。您现在可以在 HF 演示 中使用 MInference，无需 GPU。
• 📃 [24/07/03] 由于 arXiv 出现问题，目前无法在其上获取 PDF 文件。您可以通过此链接查看论文：链接。
• 🧩 [24/07/03] 我们将在 ICML'24 的 Microsoft 展位 和 ES-FoMo 上展示 MInference 1.0。维也纳见！

TL;DR

MInference 1.0 利用大模型注意力的动态稀疏特性——尽管存在一些静态模式——来加速长上下文大模型的预填充过程。它首先离线确定每个注意力头所属的稀疏模式，然后在线近似稀疏索引，并使用最优的自定义核动态计算注意力。这种方法能够在 A100 上实现高达 10 倍的预填充加速，同时保持准确性。

• MInference 1.0：通过动态稀疏注意力加速长上下文大模型的预填充（NeurIPS'24 亮点，ICML'24 ES-FoMo）
  Huiqiang Jiang†, Yucheng Li†, Chengruidong Zhang†, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H. Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang 和 Lili Qiu

SCBench 从 KV 缓存为核心视角 分析长上下文方法，覆盖整个 KV 缓存生命周期（如 KV 缓存的生成、压缩、检索和加载）。它在两种共享上下文模式下评估了 12 项任务，涵盖长上下文能力的四个类别：字符串检索、语义检索、全局信息以及多任务场景。

• SCBench：长上下文方法的 KV 缓存中心分析（ICLR'25，ENLSP @ NeurIPS'24）
  Yucheng Li, Huiqiang Jiang, Qianhui Wu, Xufang Luo, Surin Ahn, Chengruidong Zhang, Amir H. Abdi, Dongsheng Li, Jianfeng Gao, Yuqing Yang 和 Lili Qiu

MMInference 使用 模态感知的置换稀疏注意力 来加速长上下文视觉语言模型的预填充阶段。具体而言，我们实现了三种基于置换的稀疏注意力机制——FlashAttention、FlashDecoding 和 PIT——以应对视觉输入中的网格模式以及混合模态场景中的模态边界问题。

• MMInference：通过模态感知的置换稀疏注意力加速长上下文视觉语言模型的预填充（ICML'25，FM-Wild @ ICLR'25）
  Yucheng Li, Huiqiang Jiang, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Amir H. Abdi, Dongsheng Li, Jianfeng Gao, Yuqing Yang 和 Lili Qiu

🎥 概览

🎯 快速开始

需求

• Torch
• FlashAttention-2（可选）
• Triton
• Transformers >= 4.46.0

要开始使用 MInference，只需通过 pip 安装即可：

pip install minference

支持的高效方法

您可以通过运行以下代码获取支持的所有高效方法的完整列表：

from minference import MInferenceConfig
supported_attn_types = MInferenceConfig.get_available_attn_types()
supported_kv_types = MInferenceConfig.get_available_kv_types()

目前，我们支持以下长上下文方法：

• [① KV缓存生成]： MInference、xAttention、FlexPrefill、A-shape、Tri-shape、MInference w/ static、Dilated、Strided
• [② KV缓存压缩]： StreamingLLM、SnapKV、PyramidKV、KIVI
• [③ KV缓存检索]： CacheBlend
• [④ KV缓存加载]： Quest、RetrievalAttention

有关KV缓存生命周期的更多详细信息，请参阅SCBench。请注意，某些模式由vLLM支持，而所有模式均由HF支持。

支持的模型

通用MInference 支持任何解码型LLM，包括LLaMA风格的模型和Phi系列模型。 我们已经适配了市面上几乎所有的开源长上下文LLM。 如果您的模型不在支持列表中，请随时在问题区告知我们，或者您可以按照指南手动生成稀疏注意力头配置。

您可以通过运行以下代码获取支持的LLM完整列表：

from minference import get_support_models
get_support_models()

目前，我们支持以下LLM：

• Qwen2.5：Qwen/Qwen2.5-7B-Instruct、Qwen/Qwen2.5-32B-Instruct、Qwen/Qwen2.5-72B-Instruct、Qwen/Qwen2.5-7B-Instruct-1M、Qwen/Qwen2.5-14B-Instruct-1M
• LLaMA-3.1：meta-llama/Meta-Llama-3.1-8B-Instruct、meta-llama/Meta-Llama-3.1-70B-Instruct
• LLaMA-3：gradientai/Llama-3-8B-Instruct-262k、gradientai/Llama-3-8B-Instruct-Gradient-1048k、gradientai/Llama-3-8B-Instruct-Gradient-4194k、gradientai/Llama-3-70B-Instruct-Gradient-262k、gradientai/Llama-3-70B-Instruct-Gradient-1048k
• GLM-4：THUDM/glm-4-9b-chat-1m
• Yi：01-ai/Yi-9B-200K
• Phi-3：microsoft/Phi-3-mini-128k-instruct
• Qwen2：Qwen/Qwen2-7B-Instruct

如何使用MInference

[!TIP] 为了享受快速的内核实现，我们建议安装SGLang或vLLM。 对于sglang：

uv pip install "sglang[all]>=0.4.6.post4"

对于vLLM：

uv pip install "vllm>=0.9.0"
uv pip install git+https://github.com/vllm-project/flash-attention

对于HF：

from transformers import pipeline
+from minference import MInference

pipe = pipeline("text-generation", model=model_name, torch_dtype="auto", device_map="auto")

# 替换MInference模块，
# 如果您使用本地路径，请在初始化MInference时使用HF中的model_name。
+minference_patch = MInference("minference", model_name)
+pipe.model = minference_patch(pipe.model)

pipe(prompt, max_length=10)

# 使用稀疏KV方法，例如snapkv、quest、retr_attn、kivi
+minference_patch = MInference(attn_type="minference", model_name=model_name，kv_type="quest")
+pipe.model = minference_patch(pipe.model)

pipe(prompt, max_length=10)

对于vLLM：

目前，请使用vllm>=0.4.1版本。

from vllm import LLM、SamplingParams
+ from minference import MInference

llm = LLM(model_name，enforce_eager=True，max_model_len=128_000，enable_chunked_prefill=False)

# 替换MInference模块，
# 如果您使用本地路径，请在初始化MInference时使用HF中的model_name。
+minference_patch = MInference("vllm", model_name)
+llm = minference_patch(llm)

outputs = llm.generate(prompts，sampling_params)

对于带有张量并行的vLLM：

1. 将minference_patch_vllm_tp和minference_patch_vllm_executor从minference/patch.py复制到vllm/worker/worker.py中Worker类的末尾。确保缩进minference_patch_vllm_tp。
2. 调用VLLM时，务必设置enable_chunked_prefill=False。
3. 参考位于https://github.com/microsoft/MInference/blob/main/experiments/benchmarks/run_e2e_vllm_tp.sh的脚本。

from vllm import LLM、SamplingParams
+ from minference import MInference

llm = LLM(model_name，enforce_eager=True，max_model_len=128_000，enable_chunked_prefill=False，tensor_parallel_size=2)

# 替换MInference模块，
# 如果您使用本地路径，请在初始化MInference时使用HF中的model_name。
+minference_patch = MInference("vllm", model_name)
+llm = minference_patch(llm)

outputs = llm.generate(prompts，sampling_params)

仅使用内核：

from minference import vertical_slash_sparse_attention、block_sparse_attention、streaming_forward

attn_output = vertical_slash_sparse_attention(q、k、v、vertical_topk、slash)
attn_output = block_sparse_attention(q、k、v、topk)
attn_output = streaming_forward(q、k、v、init_num、local_window_num)

对于本地Gradio演示

git clone https://huggingface.co/spaces/microsoft/MInference
cd MInference
pip install -r requirments.txt
pip install flash_attn
python app.py

如需更多详情，请参阅我们的示例和实验。您可以在这篇论文以及GitHub上找到关于动态编译器PIT的更多信息。

SCBench

[!注意]

• datasets >= 2.15.0

加载数据

您可以通过 Hugging Face 数据集下载并加载 SCBench 数据（🤗 HF 仓库）：

from datasets import load_dataset

datasets = ["scbench_kv", "scbench_prefix_suffix", "scbench_vt", "scbench_repoqa", "scbench_qa_eng", "scbench_qa_chn", "scbench_choice_eng", "scbench_many_shot", "scbench_summary", "scbench_mf", "scbench_summary_with_needles", "scbench_repoqa_and_kv"]

for dataset in datasets:
    data = load_dataset("microsoft/SCBench", dataset, split="test")

数据格式

SCBench 中的所有数据都统一为以下格式：

{
    "id": "每条数据的随机 ID。",
    "context": "任务所需的长上下文，例如代码库、长文档或多轮示例。",
    "multi_turns": [{"input": "多轮问题", "answer": "多轮参考答案"}],
}

实验

我们在 GreedySearch 和 GreedySearch_vllm 两个类中，分别使用 HF 和 vLLM 实现了 多轮模式 和 多请求模式。请参考以下脚本运行实验。

对于所有方法：

cd scbench
# 单 GPU，多轮模式
VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 CUDA_VISIBLE_DEVICES=0 VLLM_WORKER_MULTIPROC_METHOD=spawn bash scripts/run_all_tasks.sh meta-llama/Llama-3.1-8B-Instruct 1 multi-turn
# 多 GPU，多轮模式
VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 CUDA_VISIBLE_DEVICES=0,1 VLLM_WORKER_MULTIPROC_METHOD=spawn bash scripts/run_all_tasks.sh meta-llama/Llama-3.1-8B-Instruct 2 multi-turn
# 多 GPU，多请求模式
VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 CUDA_VISIBLE_DEVICES=0,1 VLLM_WORKER_MULTIPROC_METHOD=spawn bash scripts/run_all_tasks.sh meta-llama/Llama-3.1-8B-Instruct 2 scdq

对于单个方法：

cd scbench
# 单 GPU，多轮模式，使用 attn_type: vllm，kv_type: dense
VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 CUDA_VISIBLE_DEVICES=0 VLLM_WORKER_MULTIPROC_METHOD=spawn bash scripts/run_single_method.sh meta-llama/Llama-3.1-8B-Instruct 1 multi-turn vllm dense
# 多 GPU，多轮模式，使用 attn_type: vllm，kv_type: dense
VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 CUDA_VISIBLE_DEVICES=0,1 VLLM_WORKER_MULTIPROC_METHOD=spawn bash scripts/run_single_method.sh meta-llama/Llama-3.1-8B-Instruct 2 multi-turn vllm dense
# 多 GPU，多请求模式，使用 attn_type: vllm，kv_type: dense
VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 CUDA_VISIBLE_DEVICES=0,1 VLLM_WORKER_MULTIPROC_METHOD=spawn bash scripts/run_single_method.sh meta-llama/Llama-3.1-8B-Instruct 2 scdq vllm dense

关于 attn_type 和 kv_type 的更多细节，请参阅此部分：支持的高效方法。

常见问题解答

如需更多见解和解答，请访问我们的 FAQ 部分。

Q1：如何有效评估动态稀疏注意力对长上下文大模型能力的影响？

为了评估像 LLaMA-3-8B-Instruct-1M 和 GLM-4-9B-1M 这样的长上下文大模型的能力，我们测试了：1) 使用 RULER 测量上下文窗口；2) 使用 InfiniteBench 测试通用任务；3) 使用 Needle in a Haystack 测试检索任务；4) 使用 PG-19 测试语言模型预测。
我们发现传统方法在检索任务中的表现较差，难度排序如下：KV 检索 > 北针任务 > Retrieval.Number > Retrieval PassKey。主要挑战在于针与 haystack 之间的语义差异。当这种差异较大时，比如 passkey 任务，传统方法表现较好。而 KV 检索则需要更高的检索能力，因为任何键都可能成为目标，多针任务更是复杂。
我们将在后续版本中继续用更多模型和数据集更新结果。

Q2：这种动态稀疏注意力模式是否只存在于未完全训练的长上下文大模型中？

首先，注意力本身就是动态稀疏的，这是其机制 inherent 的特性。我们选择了最先进的长上下文大模型，GLM-4-9B-1M 和 LLaMA-3-8B-Instruct-1M，它们的有效上下文窗口分别为 64K 和 16K。借助 MInference，这些窗口可以分别扩展到 64K 和 32K。我们将继续将该方法应用于其他先进的长上下文大模型，并更新结果，同时探索这种动态稀疏注意力模式的理论基础。

Q3：这种动态稀疏注意力模式是否只存在于自回归语言模型或基于 RoPE 的语言模型中？

类似的垂直和斜线稀疏模式已在 BERT[1] 和多模态大模型[2] 中被发现。我们对 T5 注意力模式的分析显示，即使在双向注意力中，这些模式也会出现在不同的头中。
[1] SparseBERT：重新思考自注意力中的重要性分析，ICML 2021。
[2] LOOK-M：KV 缓存中的一次性优化，用于高效的多模态长上下文推理，2024年。

图 1. T5 编码器中的稀疏模式。

Q4：MInference、SSM、线性注意力和稀疏注意力之间有什么关系？

这四种方法（MInference、SSM、线性注意力和稀疏注意力）都能高效地优化 Transformer 中的注意力复杂度，但它们引入归纳偏置的方式各不相同。后三种方法需要从头开始训练。最近的一些工作，如 Mamba-2 和 Unified Implicit Attention Representation，将 SSM 和 线性注意力统一为静态稀疏注意力，而 Mamba-2 本身则是一种分块稀疏方法。尽管这些方法由于注意力中的稀疏冗余而显示出潜力，但在处理复杂任务中的动态语义关联时，静态稀疏注意力可能会遇到困难。相比之下，动态稀疏注意力更适合管理这些关系。

Q5：在 _prepare_4d_causal_attention_mask_with_cache_position 中出现 CUDA 内存不足。

解决方案：在模型初始化时添加以下参数，将 Hugging Face 模型的注意力后端设置为 FlashAttention-2：_attn_implementation="flash_attention_2"。

Q6：在 logits = self.lm_head(hidden_states[:, -logits_to_keep:, :]) 中出现 CUDA 内存不足。

解决方案：在模型前向传播中将 logits_to_keep 设置为 1。

引用

如果您认为 MInference 对您的项目和研究有用或相关，请您引用我们的论文：

@inproceedings{jiang2024minference,
  author = {Huiqiang Jiang and Yucheng Li and Chengruidong Zhang and Qianhui Wu and Xufang Luo and Surin Ahn and Zhenhua Han and Amir H. Abdi and Dongsheng Li and Chin-Yew Lin and Yuqing Yang and Lili Qiu},
  booktitle = {The Thirty-eighth Annual Conference on Neural Information Processing Systems},
  title = {{MI}nference 1.0: Accelerating Pre-filling for Long-Context {LLM}s via Dynamic Sparse Attention},
  url = {https://openreview.net/forum?id=fPBACAbqSN},
  year = {2024}
}

@inproceedings{li2025scbench,
  title={{SCB}ench: A {KV} Cache-Centric Analysis of Long-Context Methods},
  author={Yucheng Li and Huiqiang Jiang and Qianhui Wu and Xufang Luo and Surin Ahn and Chengruidong Zhang and Amir H. Abdi and Dongsheng Li and Jianfeng Gao and Yuqing Yang and Lili Qiu},
  booktitle={The Thirteenth International Conference on Learning Representations},
  year={2025},
  url={https://openreview.net/forum?id=gkUyYcY1W9}
}

@inproceedings{li2025mminference,
    title={{MMI}ference: Accelerating Pre-filling for Long-Context VLMs via Modality-Aware Permutation Sparse Attention},
    author={Li, Yucheng and Jiang, Huiqiang and Zhang, Chengruidong and Wu, Qianhui and Luo, Xufang and Ahn, Surin and Abdi, Amir H and Li, Dongsheng and Gao, Jianfeng and Yang, Yuqing and Qiu, Lili},
    booktitle={Forty-second International Conference on Machine Learning},
    year={2025},
    url={https://openreview.net/forum?id=me6PfbATWM}
}

贡献

本项目欢迎各种贡献和建议。大多数贡献都需要您同意贡献者许可协议（CLA），以声明您有权并将您的贡献权利授予我们使用。有关详细信息，请访问 https://cla.opensource.microsoft.com。

当您提交拉取请求时，CLA 机器人会自动判断您是否需要提供 CLA，并相应地标记 PR（例如状态检查、评论）。请按照机器人提供的指示操作即可。对于所有使用我们 CLA 的仓库，您只需完成一次此步骤。

本项目已采用 微软开源行为准则。更多信息请参阅 行为准则常见问题解答 或通过电子邮件 opencode@microsoft.com 联系我们，以获取更多问题或意见。

商标

本项目可能包含项目、产品或服务的商标或徽标。未经授权使用微软商标或徽标必须遵守并遵循 微软商标与品牌指南。在本项目的修改版本中使用微软商标或徽标不得造成混淆或暗示微软的赞助。任何第三方商标或徽标的使用均受该第三方政策的约束。

MInference 快速上手指南

MInference 是一款专为长上下文大语言模型（LLM）设计的加速工具。它利用注意力机制的动态稀疏特性，在单张 A100 显卡上可实现高达 10 倍 的预填充（prefilling）速度提升，同时保持甚至提高模型准确率。支持处理百万级（1M）token 上下文。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐)
• Python: 3.8+
• 核心依赖:
  • PyTorch
  • Triton
  • Transformers >= 4.46.0
• 可选优化依赖 (强烈推荐安装以获得最佳性能):
  • FlashAttention-2
  • SGLang 或 vLLM (用于启用优化的底层内核)

提示: 为了获得最快的推理速度，建议优先安装 sglang 或 vllm。

安装步骤

1. 安装基础包

使用 pip 直接安装 MInference：

pip install minference

2. 安装高性能后端（推荐）

为了启用自定义稀疏注意力内核并实现最大加速比，请选择以下任一方案进行安装：

方案 A：安装 SGLang (推荐)

uv pip install "sglang[all]>=0.4.6.post4"

方案 B：安装 vLLM

uv pip install "vllm>=0.9.0"
uv pip install git+https://github.com/vllm-project/flash-attention

基本使用

MInference 支持 Hugging Face (transformers)、vLLM 和 SGLang 等多种框架。以下是基于 Hugging Face 的最简使用示例。

Hugging Face 集成示例

只需几行代码即可将现有模型转换为支持 MInference 的模式：

from transformers import pipeline
from minference import MInference

# 1. 初始化标准的 HF pipeline
model_name = "Qwen/Qwen2.5-7B-Instruct-1M"  # 或其他支持的长上下文模型
pipe = pipeline("text-generation", model=model_name, torch_dtype="auto", device_map="auto")

# 2. 应用 MInference 补丁
# 注意：如果使用本地路径加载模型，初始化 MInference 时请仍使用 HF 上的模型名称
minference_patch = MInference("minference", model_name)
pipe.model = minference_patch(pipe.model)

# 3. 正常进行推理
prompt = "这是一段非常长的文本..."
outputs = pipe(prompt, max_length=10)

print(outputs)

进阶：结合稀疏 KV Cache 方法

MInference 还支持结合多种 KV Cache 压缩与检索策略（如 SnapKV, Quest, RetrievalAttention 等）：

from transformers import pipeline
from minference import MInference

model_name = "Qwen/Qwen2.5-7B-Instruct-1M"
pipe = pipeline("text-generation", model=model_name, torch_dtype="auto", device_map="auto")

# 指定注意力类型和 KV Cache 策略 (例如使用 Quest)
minference_patch = MInference(attn_type="minference", model_name=model_name, kv_type="quest")
pipe.model = minference_patch(pipe.model)

outputs = pipe(prompt, max_length=10)

查看支持列表

您可以运行以下代码查看当前版本支持的所有高效方法和模型：

from minference import MInferenceConfig, get_support_models

# 查看支持的注意力类型和 KV 类型
print(MInferenceConfig.get_available_attn_types())
print(MInferenceConfig.get_available_kv_types())

# 查看支持的模型列表
print(get_support_models())

注意: 目前完美支持 Qwen2.5-1M, LLaMA-3-8B/70B (Gradient 版本), GLM-4-1M 等主流长上下文开源模型。