dMoE-16B: dLLMs with Learnable Block Experts

dMoE is a block-level Mixture-of-Experts (MoE) framework designed for Diffusion Large Language Models (dLLMs). By aggregating token-level expert distributions within each block into a unified block-level distribution, dMoE substantially reduces the number of uniquely activated experts during inference, mitigating memory-bound bottlenecks without sacrificing performance.

• Paper: dMoE: dLLMs with Learnable Block Experts
• Project Page: https://fscdc.github.io/dMoE/
• Repository: https://github.com/fscdc/dMoE

Highlights

• Learnable Block Experts: Introduces block-level MoE routing into dLLMs, drastically compressing the number of activated unique experts across diffusion steps.
• Reduced MoE Bandwidth: Significantly reduces memory bandwidth consumed by expert weight loading during the block diffusion process.
• Improved Efficiency-Accuracy Trade-off: Achieves 1.14x to 1.66x end-to-end latency speedup while maintaining competitive performance on benchmarks.
• Plug-and-play on LLaDA-2.0: Built directly on top of LLaDA-2.0-mini without architectural changes.

Sample Usage

The model can be used with the Transformers library. Note that it requires trust_remote_code=True to load the custom architecture.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "FSCCS/dMoE-16B"
device = "cuda:0"

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME, trust_remote_code=True, torch_dtype=torch.bfloat16
).to(device).eval()

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)

prompt = "A robe takes 2 bolts of blue fiber and half that much white fiber. How many bolts in total does it take?" + "
Let's think step by step
"

messages = [[{"role": "user", "content": prompt}]]
input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)

inputs = tokenizer(input_text, return_tensors="pt", padding_side="left")
input_ids = inputs["input_ids"].to(device)

with torch.no_grad():
    out, unique_experts_count = model.generate(
        input_ids,
        steps=32,
        gen_length=2048,
        block_length=32,
        temperature=0.0,
        eos_early_stop=True,
    )

generated = out[:, input_ids.shape[1]:]
result = tokenizer.batch_decode(generated, skip_special_tokens=True)

print("Output:", result[0])
print("Unique experts count:", unique_experts_count)

Citation

@article{feng2026dmoe,
  title={dMoE: dLLMs with Learnable Block Experts},
  author={Feng, Sicheng and Chen, Zigeng and Fang, Gongfan and Ma, Xinyin and Wang, Xinchao},
  journal={arXiv preprint arXiv:2605.30876},
  year={2026}
}