WeDLM

Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference

WeChat AI, Tencent

     

🏆 The fastest diffusion language model, capable of outperforming vLLM-optimized autoregressive baselines in wall-clock speed

⬆️ Real-time comparison: Qwen3-8B-Instruct with vLLM (left) vs WeDLM-8B-Instruct (right) on the same prompt

💡 Why WeDLM?

Most diffusion language models use bidirectional attention, which breaks KV cache compatibility and fails to translate parallel prediction into actual speedups over optimized AR engines like vLLM.

WeDLM solves this by performing parallel mask recovery under standard causal attention, enabling:

• ✅ Native KV cache compatibility (FlashAttention, PagedAttention, CUDA Graphs)
• ✅ Direct initialization from pre-trained AR models (Qwen2.5, Qwen3)
• ✅ Real speedups measured against production-grade vLLM baselines

GSM8K benchmark: WeDLM achieves 3-6× speedup over vLLM while maintaining competitive accuracy.

🚀 Quick Start

Installation

git clone https://github.com/tencent/WeDLM.git
cd WeDLM && bash install.sh

# Step 1: PyTorch
pip install torch==2.8.0+cu129 --index-url https://download.pytorch.org/whl/cu129

# Step 2: flash-attn build dependencies
pip install psutil ninja packaging

# Step 3: flash-attn (requires torch first)
pip install flash-attn==2.7.4.post1 --no-build-isolation

# Step 4: WeDLM
git clone https://github.com/tencent/WeDLM.git
cd WeDLM && pip install -e .

# Pull the Docker image
docker pull aiweiliu/wedlm:v3

# Run the container with GPU support
docker run -it --gpus all -p 8080:8080 --name wedlm aiweiliu/wedlm:v3 /bin/bash

# Inside the container, run inference directly
python example.py --model tencent/WeDLM-8B-Instruct

# Inside the container, run web demo
python web_demo.py --model tencent/WeDLM-8B-Instruct

Run Inference

# Run simple generation
python example.py --model tencent/WeDLM-8B-Instruct

Example Output (NVIDIA H20):

Prompt: A store sells apples for $2 each and oranges for $3 each...

Response: To determine the total amount Tom spent...
Therefore, the total amount Tom spent is $22.

==================================================
  Generated tokens: 218
  Time elapsed: 0.32 seconds
  Speed: 689.18 tokens/s ⚡
==================================================

💻 Interactive Demo & API

Web Demo:

python web_demo.py --model tencent/WeDLM-8B-Instruct

⬆️ Interactive web interface for real-time generation

Python API:

from transformers import AutoTokenizer
from wedlm import LLM, SamplingParams

# Initialize
llm = LLM(model="tencent/WeDLM-8B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("tencent/WeDLM-8B-Instruct", trust_remote_code=True)

# Prepare Prompt
prompt = "Solve: 2x + 5 = 13"
messages = [{"role": "user", "content": prompt}]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

# Generate
outputs = llm.generate([text], SamplingParams(temperature=0.0, max_tokens=512))
print(outputs[0]["text"])

📊 Performance

Speed-Quality Tradeoff

WeDLM’s speedup varies by task characteristics. Structured, low-entropy tasks (math, code) see the largest gains.

[!NOTE] Acceleration comes with a quality-speed tradeoff. Conservative settings preserve accuracy; aggressive settings maximize speed. See our paper for detailed analysis.

Generation Quality

WeDLM preserves and often improves upon its base AR model capabilities.

🦁 Model Zoo

Model	Base	Context	Download
WeDLM-7B	Qwen2.5-7B	32k
WeDLM-7B-Instruct	Qwen2.5-7B	32k
WeDLM-8B	Qwen3-8B	32k
WeDLM-8B-Instruct ⭐	Qwen3-8B	32k

🛠️ Advanced Usage

Installation from Source

Requirements: Python 3.9+, PyTorch 2.1+, CUDA 11.8+

git clone https://github.com/tencent/WeDLM.git
cd WeDLM && bash install.sh

[!NOTE] flash-attn requires compilation and must be installed after PyTorch. The install.sh script handles this automatically (default: CUDA 12.9). For other CUDA versions: CUDA_VERSION=cu124 bash install.sh

Evaluation

Reproduce our results using the provided scripts:

# 1. Download datasets
python -m evaluation.download_datasets --all

# 2. Run WeDLM evaluation (e.g., GSM8K)
bash evaluation/evaluation_base.sh \
    --model_path "tencent/WeDLM-8B-Instruct" \
    --output_dir "output/" \
    --datasets "gsm8k" \
    --num_gpus 8

Evaluate AR Baseline Models:

To benchmark against autoregressive models using vLLM:

bash evaluation/ar_baseline_script.sh

[!TIP] Edit ar_baseline_script.sh to specify your model path (MODEL_PATH) and select datasets from the list.

See evaluation/demo.sh for all benchmark commands.

Fine-tuning

Fine-tune WeDLM on your own data using our training framework:

cd finetune

# Download sample dataset (Alpaca)
python scripts/download_alpaca_sft.py

# Launch training (multi-GPU with DeepSpeed)
accelerate launch --multi_gpu --num_processes 4 --mixed_precision bf16 \
    train.py --config configs/example.yaml

Attention Backends:

• dense — PyTorch native SDPA (works out-of-the-box)
• magi — MagiAttention for optimized training (requires separate installation)

See finetune/README.md for detailed configuration and data format.

HuggingFace Compatibility

For training or simple forward passes, we provide a standard HF interface.

[!WARNING] For fast inference, use the wedlm engine (shown in Quick Start). The HF interface below is for training/forward pass convenience only.

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("tencent/WeDLM-8B-Instruct", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("tencent/WeDLM-8B-Instruct", trust_remote_code=True)

inputs = tokenizer("Hello", return_tensors="pt")
out = model(**inputs)

🧠 How It Works

WeDLM introduces Topological Reordering to perform parallel mask recovery under standard causal attention, combined with Streaming Parallel Decoding for continuous prefix commitment.

• 👉 Project Page — Interactive explanations and visualizations
• 👉 Paper — Technical details and full experimental results

📜 Citation

If you find WeDLM useful for your research, please cite:

@article{liu2025wedlm,
  title={WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference},
  author={Liu, Aiwei and He, Minghua and Zeng, Shaoxun and Zhang, Linhao and Wu, Chuhan and Jia, Wei and Liu, Yuan and Yu, Yang and Zhou, Xiao and Zhou, Jie},
  journal={arXiv preprint arXiv:2512.22737},
  year={2025}
}

Acknowledgement

Built upon nano-vllm, Qwen, and vLLM.

License

Apache 2.0