Derivative-work notice: This repository builds on inclusionAI/dInfer and references the expert-offloading evaluation architecture from TIDE, “Efficient and Lossless MoE Diffusion LLM Inference with I/O-aware Expert Offload” (arXiv:2605.20179). It adds reproducible LLaDA 2.0 Mini runs for expert caps 32/64/128, fixed sample generation, metrics aggregation, HumanEval validation, and a self-contained environment setup workflow. See NOTICE.

What this repository does

This repository provides a reproducible data-collection pipeline for studying TIDE predictive expert offloading on LLaDA 2.0 Mini. It runs the same evaluation workload with different GPU expert capacities and records generation results, throughput, latency, expert-cache hit/miss rates, CPU/GPU compute time, and expert-migration overhead.

The three supported expert caps are 32, 64, and 128. A cap is the maximum number of MoE experts kept on the GPU for each sparse layer. Running all three caps makes it possible to compare memory pressure and offloading behavior under the same prompts and decoding settings.

The main workflow is:

fixed dataset samples (optional)
          |
          v
run_caps.sh -> cap32 / cap64 / cap128 generations
          |
          +-> tide_metrics_summary.json / .csv
          |
          +-> HumanEval pass@1 validation (optional)

Requirements

• Linux with an NVIDIA CUDA GPU
• Python 3.10–3.12
• A local copy of inclusionAI/LLaDA2.0-mini
• Enough CPU RAM and GPU memory for the model and selected expert cap
• Internet access for the initial dependency and dataset download

Model weights and generated datasets are intentionally not committed to Git.

Installation

Clone this repository and run the provided installer:

git clone https://github.com/jy2247/TIDE_DATA_COLLECTION.git
cd TIDE_DATA_COLLECTION
./script/setup_env.sh

The installer creates .venv inside the repository and installs PyTorch, vLLM, Transformers, lm-eval, datasets, TIDE dependencies, and this dInfer checkout. It is safe to run again to update an existing environment.

To discard only the selected virtual environment and rebuild it:

./script/setup_env.sh --recreate

setup_env.sh options:

The same values can be set with VENV_DIR, PYTHON_BIN, and PIP_CACHE_DIR.

Quick start

First set the model and output paths. The default model path in the script is specific to the original development machine, so other users must override it:

export MODEL_PATH=/absolute/path/to/LLaDA2.0-mini
export OUTPUT_ROOT="$PWD/outputs"

Run one sample for all three caps to verify the installation:

./script/run_caps.sh --limit 1

Run the full HumanEval dataset:

./script/run_caps.sh \
  --caps "32 64 128" \
  --tasks "humaneval" \
  --gen-length 256 \
  --gpu 0

Caps and tasks run sequentially on the selected GPU. This avoids loading multiple model copies at the same time and makes measurements easier to compare.

run_caps.sh parameters

./script/run_caps.sh [options]

Important parameter notes:

• Lower caps keep fewer experts on GPU and normally cause more CPU execution and expert migration.
• tau controls how frequently TIDE refreshes its expert-placement prediction.
• threshold affects how aggressively tokens are accepted during parallel diffusion decoding.
• Use identical gen-length, block-length, tau, threshold, and samples when comparing caps.
• --samples and --limit may be used together, but --limit further reduces the selected sample set.

Example with two tasks and a custom output directory:

./script/run_caps.sh \
  --model /models/LLaDA2.0-mini \
  --output-root "$PWD/outputs" \
  --caps "32 64 128" \
  --tasks "gsm8k_llada_mini mbpp_sanitized_llada_mini" \
  --gen-length 256 \
  --block-length 32 \
  --tau 6 \
  --threshold 0.95 \
  --gpu 0

Output files

Each command creates one timestamped directory:

<output-root>/<run-name>/
├── base_out.txt
├── run_config.txt
├── status.log
├── summarize.log
├── tide_metrics_summary.json
├── tide_metrics_summary.csv
├── cap32_gen256/
│   ├── logs/<task>.log
│   └── tasks/<task>/
│       ├── rank_0.jsonl
│       └── tide_metrics.json
├── cap64_gen256/
└── cap128_gen256/

Key files:

• run_config.txt: exact model, tasks, caps, decoding settings, and paths.
• status.log: start/end time and return code for every cap/task pair.
• rank_0.jsonl: one generated response plus per-sample measurements per line.
• tide_metrics.json: task-level aggregate TIDE measurements.
• tide_metrics_summary.csv: compact cap/task comparison table.
• tide_metrics_summary.json: detailed machine-readable metric summary.

The summary includes available measurements such as generated tokens, sample latency, tokens per second, decoding steps, hit/miss rates, CPU/GPU compute time, and migration time.

Generate reproducible sample indices

Use a fixed sample set when comparing caps:

./.venv/bin/python script/generate_sample_manifest.py \
  --output-dir "$PWD/outputs/fixed_samples" \
  --samples-per-task 100 \
  --tasks \
    mbpp_sanitized_llada_mini \
    gsm8k_llada_mini \
    hendrycks_math500 \
    humaneval \
    ifeval

Generator parameters:

The command writes:

• *.jsonl: selected dataset documents with local and global indices.
• *.indices.json: task-to-index mapping accepted by run_caps.sh --samples.
• *.summary.json: sample counts and output paths.

Run all caps with that fixed index file:

./script/run_caps.sh \
  --tasks "mbpp_sanitized_llada_mini gsm8k_llada_mini hendrycks_math500 humaneval ifeval" \
  --samples "$PWD/outputs/fixed_samples/llada2_mini_5tasks_100each.indices.json"

Rebuild metric summaries

The run script summarizes successful outputs automatically. To rebuild a summary later:

./.venv/bin/python script/summarize_tide_metrics.py /path/to/run-directory

The summarizer accepts both the current cap*/tasks/<task>/rank_0.jsonl layout and the older cap*/rank_0.jsonl HumanEval layout.

Validate HumanEval correctness

HumanEval validation executes model-generated Python code. Run it only on outputs that you trust:

HF_ALLOW_CODE_EVAL=1 \
./.venv/bin/python script/validate_humaneval.py \
  /path/to/run-directory \
  --caps 32 64 128 \
  --gen-length 256 \
  --workers 8 \
  --timeout 3

Validator parameters:

It writes humaneval_validation.json inside each HumanEval task directory and humaneval_accuracy_summary.json at the run root.

All tools under script/ support --help. A shorter command reference is also available in script/README.md.

Introduction

dInfer is an efficient and extensible inference framework for dLLMs. As illustrated in the following architecture, it modularizes inference into four components: model, diffusion iteration manager, decoder and KV-cache manager. It provides well-designed APIs for flexible algorithms combinations in each component. It now supports batched inference for improved throughput.

Figure: Overall Architecture of dInfer

dInfer supports multiple dLLM variants, including LLaDA, LLaDA-MoE and LLaDA2.

News

[2025/12/21] release v0.2. The major features of this release can be found here.

[2025/12/10] Support and speed up the formal version of block diffusion LLMs (LLaDA2-mini and LLaDA2-flash). Support quant versions of LLaDA2-mini and LLaDA2-flash.

[2025/11/15] Support the inference on block diffusion LLMs (LLaDA2-mini-preview and LLaDA2-flash-preview).

[2025/10/10] Release the first version of the dInfer framework.

Contents

• Supported Models
• Quick Start
• Benchmark Results

Supported Models

dInfer supports multiple diffusion language model variants with different architectures and sizes. Below are the HuggingFace model links and their corresponding implementation files:

Quick Start

Install dInfer

git clone https://github.com/inclusionAI/dInfer.git
cd dInfer
pip install .

To use it with vLLM backend (it works with LLaDA and LLaDA-MoE), please install vLLM.

pip install vllm==0.10.2

To use it with SGLang backend (it works with LLaDA2), please install SGLang.

pip install sglang==0.5.3.post1

Convert to FusedMoE (LLaDA-MoE only)

To run LLaDA-MoE model downloaded from HuggingFace, we need to first convert it to a format supported by dInfer. dInfer provides a script tools/transfer.py for the format conversion.

1) Download and Convert

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) Load the model

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")

Run Inference

Benchmark (speed only)

Measure throughput (TPS) only; predictions are saved under --output_dir with no automatic scoring.

• LLaDA2 model

  • LLaDA2-flash Dataset profiling (threshold decoder, TP across 4 GPUs):

  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 Dataset profiling (threshold decoder, TP across 4 GPUs):

  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 and LLaDA-MoE model

  • LLaDA-MoE Dataset profiling (threshold decoder, TP across 4 GPUs):

  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 Single-sample profiling (threshold decoder, TP across 4 GPUs):

  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 can use benchmark_dataset.py and benchmark.py.

Evaluation with lm-eval (speed + accuracy)

• Built on HuggingFace lm-eval-harness to compute TPS and benchmark scores.
• Tasks provided:
  • gsm8k_llada: math reasoning.
  • mbpp_sanitized_llada: sanitized Python code generation.
• For more examples and comprehensive instructions, see our quickstart guide.

Benchmark Results

The inference speed (TPS) on LLaDA-MoE

dInfer delivers over 1,100 TPS at batch size 1 on HumanEval and on average 800+ TPS across six benchmarks on a single node with 8×H800 GPUs.

Figure: Benchmark results on LLaDA-MoE

Speedup comparisons:

• 10× faster than Fast-dLLM while maintaining accuracy
• 2-3× faster than Qwen2.5-3B on vLLM (LLaDA-MoE) with comparable quality

The inference speed (TPS) on LLaDA2-flash-CAP

The inference speed is measured on LLaDA2-flash-CAP (with 100B parameters) on 8 H20 GPUs (parallel decoding threshold=0.95, generation length=1000).

Limitations

• Block Diffusion: Not supported on LLaDA Dense/MoE models (use --use_bd with LLaDA2 only)

Contact us

• Wechat Group

Citation

@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}
}