跳到正文
原文
蚂蚁 inclusionAI:GitHub 新仓库· inclusionAI·· 2025-09-29精选AI 评分58

inclusionAI/dInfer

inclusionAI/dInfer

AI 导读

inclusionAI团队发布了dInfer,一个专为扩散语言模型设计的高效推理框架。该框架旨在解决扩散模型在文本生成领域推理速度慢、资源消耗大的核心挑战。dInfer通过一系列底层优化技术,显著提升了推理效率,能够更快地生成文本,同时降低计算成本,为扩散模型在更广泛的实际应用场景中部署提供了关键技术支持。

推荐理由

蚂蚁把扩散语言模型的推理框架开源了,这类模型的推理效率一直是落地瓶颈,做端侧或低成本部署的团队值得看看能不能接上。

正文 · AI 翻译

dInfer

简介

dInfer 是一个高效且可扩展的 dLLM 推理框架。如下方架构图所示,它将推理模块化为四个组件:模型、扩散迭代管理器、解码器和 KV-cache 管理器。它为每个组件提供了精心设计的 API,以支持灵活的算法组合。目前它已支持批量推理,以提升吞吐量。

dInfer v0.1 architecture
图: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_bd
    
  • LLaDA、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。

dInfer v0.1 speedup
图: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 使用)

联系我们

  • 微信群

Wechat Group

引用

@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