inclusionAI/dInfer
inclusionAI/dInfer
inclusionAI团队发布了dInfer,一个专为扩散语言模型设计的高效推理框架。该框架旨在解决扩散模型在文本生成领域推理速度慢、资源消耗大的核心挑战。dInfer通过一系列底层优化技术,显著提升了推理效率,能够更快地生成文本,同时降低计算成本,为扩散模型在更广泛的实际应用场景中部署提供了关键技术支持。
蚂蚁把扩散语言模型的推理框架开源了,这类模型的推理效率一直是落地瓶颈,做端侧或低成本部署的团队值得看看能不能接上。
简介
dInfer 是一个高效且可扩展的 dLLM 推理框架。如下方架构图所示,它将推理模块化为四个组件:模型、扩散迭代管理器、解码器和 KV-cache 管理器。它为每个组件提供了精心设计的 API,以支持灵活的算法组合。目前它已支持批量推理,以提升吞吐量。
图:dInfer 的整体架构
dInfer 支持多种 dLLM 变体,包括 LLaDA、LLaDA-MoE 和 LLaDA2。
新闻
[2025/12/21] 发布 v0.2。本次发布的主要特性可参见此处。
[2025/12/10] 支持并加速块扩散 LLM 的正式版本(LLaDA2-mini 和 LLaDA2-flash)。支持 LLaDA2-mini 和 LLaDA2-flash 的量化版本。
[2025/11/15] 支持块扩散 LLM(LLaDA2-mini-preview 和 LLaDA2-flash-preview)的推理。
[2025/10/10] 发布 dInfer 框架的首个版本。
目录
支持的模型
dInfer 支持多种不同架构和规模的扩散语言模型变体。以下是 HuggingFace 模型链接及其对应的实现文件:
| 模型 | 规模 | 实现 | HuggingFace 链接 |
|---|---|---|---|
| LLaDA2.0-mini | 16B | LLaDA2MoeModelLM | inclusionAI/LLaDA2.0-mini |
| LLaDA2.0-flash | 100B | LLaDA2MoeModelLM | inclusionAI/LLaDA2.0-flash |
| LLaDA2.0-mini-preview | 16B | LLaDA2MoeModelLM | inclusionAI/LLaDA2.0-mini-preview |
| LLaDA2.0-flash-preview | 100B | LLaDA2MoeModelLM | inclusionAI/LLaDA2.0-flash-preview |
| LLaDA-MoE-7B-A1B-Base | 7B | LLaDAMoeModelLM | inclusionAI/LLaDA-MoE-7B-A1B-Base |
| LLaDA-MoE-7B-A1B-Instruct | 7B | LLaDAMoeModelLM | inclusionAI/LLaDA-MoE-7B-A1B-Instruct |
| LLaDA-8B-Base | 8B | LLaDAModelLM | GSAI-ML/LLaDA-8B-Base |
| LLaDA-8B-Instruct | 8B | LLaDAModelLM | GSAI-ML/LLaDA-8B-Instruct |
| LLaDA-1.5 | 8B | LLaDAModelLM | GSAI-ML/LLaDA-1.5 |
快速开始
安装 dInfer
git clone https://github.com/inclusionAI/dInfer.git
cd dInfer
pip install .
要配合 vLLM 后端使用(它适用于 LLaDA 和 LLaDA-MoE),请先安装 vLLM。
pip install vllm==0.10.2
要配合 SGLang 后端使用(它适用于 LLaDA2),请先安装 SGLang。
pip install sglang==0.5.3.post1
转换为 FusedMoE(仅限 LLaDA-MoE)
要运行从 HuggingFace 下载的 LLaDA-MoE 模型,我们需要先将其转换为 dInfer 支持的格式。dInfer 提供了一个脚本 tools/transfer.py 用于格式转换。
1) 下载并转换
pip install -U huggingface_hub hf_transfer
export HF_HUB_ENABLE_HF_TRANSFER=1
# Download Instruct checkpoint
hf download inclusionAI/LLaDA-MoE-7B-A1B-Instruct \
--repo-type model \
--local-dir /path/to/LLaDA-MoE-7B-A1B-Instruct
# Convert to FusedMoE
python -m tools.transfer \
--input /path/to/LLaDA-MoE-7B-A1B-Instruct \
--output /path/to/LLaDA-MoE-7B-A1B-Instruct-fused
2) 加载模型
from dinfer.model import AutoModelForCausalLM
from transformers import AutoTokenizer
m = "/path/to/LLaDA-MoE-7B-A1B-Instruct-fused"
tok = AutoTokenizer.from_pretrained(m, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(m, trust_remote_code=True, torch_dtype="bfloat16")
运行推理
基准测试(仅测速度)
仅测量吞吐量(TPS);预测结果保存在 --output_dir 下,不进行自动评分。
LLaDA2 模型
- LLaDA2-flash 数据集性能分析(阈值解码器,TP 跨 4 块 GPU):
python benchmarks/benchmark_dataset_sglang.py \ --model_name inclusionAI/LLaDA2.0-flash \ --dataset dataset_path \ --gen_len 2048 \ --block_length 32 \ --gpu 0,1,2,3 \ --output_dir runs/llada2_flash \ --use_tp \ --parallel_decoding threshold \ --threshold 0.9 \ --cache prefix \ --use_bd- LLaDA2-mini 数据集性能分析(阈值解码器,TP 跨 4 块 GPU):
python benchmarks/benchmark_dataset_sglang.py \ --model_name inclusionAI/LLaDA2.0-mini \ --dataset dataset_path \ --gen_len 2048 \ --block_length 32 \ --gpu 0,1,2,3 \ --output_dir runs/llada2_mini \ --use_tp \ --parallel_decoding threshold \ --threshold 0.9 \ --cache prefix \ --use_bdLLaDA、LLaDA1.5 和 LLaDA-MoE 模型
- LLaDA-MoE 数据集性能分析(阈值解码器,TP 跨 4 块 GPU):
python benchmarks/benchmark_dataset.py \ --model_name inclusionAI/LLaDA-MoE-7B-A1B-Instruct \ --model_type llada_moe \ --dataset dataset_path \ --gen_len 1024 \ --block_length 64 \ --gpu 0,1,2,3 \ --output_dir runs/llada_moe_threshold \ --use_tp \ --parallel_decoding threshold \ --threshold 0.8 \ --cache dual \ --prefix_look 16 \ --after_look 16 \ --warmup_times 4 \ --cont_weight 0.3- LLaDA 单样本性能分析(阈值解码器,TP 跨 4 块 GPU):
python benchmarks/benchmark.py \ --model_name GSAI-ML/LLaDA-8B-Instruct \ --model_type llada \ --gen_len 2048 \ --block_length 32 \ --gpu 0,1,2,3 \ --use_tp \ --parallel_decoding threshold \ --threshold 0.9 \ --cache prefix- LLaDA、LLaDA1.5、LLaDA-MoE 均可使用 benchmark_dataset.py 和 benchmark.py。
使用 lm-eval 进行评估(速度 + 准确率)
- 基于 HuggingFace
lm-eval-harness构建,用于计算 TPS 和基准测试分数。 - Tasks provided:
gsm8k_llada:数学推理。mbpp_sanitized_llada:经过清洗的 Python 代码生成。
- 更多示例和完整说明,请参阅我们的快速入门指南。
基准测试结果
LLaDA-MoE 上的推理速度(TPS)
dInfer 在 batch size 为 1 时于 HumanEval 上实现超过 1,100 TPS,并在单节点 8×H800 GPU 上跨六个基准测试平均达到 800+ TPS。
图:LLaDA-MoE 上的基准测试结果
加速比对比:
- 比 Fast-dLLM 快 10 倍,同时保持精度
- 在 vLLM(LLaDA-MoE)上比 Qwen2.5-3B 快 2-3 倍,质量相当
LLaDA2-flash-CAP 上的推理速度(TPS)
推理速度是在 LLaDA2-flash-CAP(100B 参数)上、使用 8 块 H20 GPU 测得的(并行解码阈值=0.95,生成长度=1000)。
| 基准测试 | batch size = 1 | batch size = 32 |
|---|---|---|
| openai_humaneval | 753.10 | 2558.51 |
| gsm8k | 591.90 | 2111.79 |
| IFEval | 222.60 | 931.89 |
| CruxEval-O | 562.90 | 1967.08 |
| mbpp | 773.00 | 2262.45 |
| AVG | 580.70 | 1966.34 |
局限性
- Block Diffusion:在 LLaDA Dense/MoE 模型上不受支持(仅 LLaDA2 可配合
--use_bd使用)
联系我们
- 微信群
引用
@article{dinfer,
title={dInfer: An Efficient Inference Framework for Diffusion Language Models},
author={Yuxin Ma, Lun Du, Lanning Wei, Kun Chen, Qian Xu, Kangyu Wang, Guofeng Feng, Guoshan Lu, Lin Liu, Xiaojing Qi, Xinyuan Zhang, Zhen Tao, Haibo Feng, Ziyun Jiang, Ying Xu, Zenan Huang, Yihong Zhuang, Haokai Xu, Jiaqi Hu, Zhenzhong Lan, Junbo Zhao, Jianguo Li, Da Zheng},
year={2025},
journal={arXiv preprint arXiv:2510.08666}
}
来源:蚂蚁 inclusionAI:GitHub 新仓库 · github.com