Views
No views yet
LLaDA-MoE-7B-A1B-Base: A base pre-trained model designed for research and secondary development.LLaDA-MoE-7B-A1B-Instruct: An instruction-tuned model optimized for practical applications.LLaDA-MoE-7B-A1B-Instruct-TD: A specialized instruction-tuned model, further optimized for accelerated inference using Trajectory Distillation.

| Model ID | Description | Hugging Face Link |
|---|---|---|
inclusionAI/LLaDA-MoE-7B-A1B-Base | Base pre-trained model for research and fine-tuning. | 🤗 Model Card |
inclusionAI/LLaDA-MoE-7B-A1B-Instruct | Instruction-tuned model, ready for downstream applications. | 🤗 Model Card |
inclusionAI/LLaDA-MoE-7B-A1B-Instruct-TD | An instruction-tuned model further optimized with Trajectory Distillation (TD) for accelerated inference. Decodes multiple tokens per forward pass. | 🤗 Model Card |

git clone https://github.com/inclusionAI/dInfer.git
cd dInfer
pip install .1# From repo root
2python tools/transfer.py \
3 --input /path/to/LLaDA-MoE-7B-A1B-Instruct \
4 --output /path/to/LLaDA-MoE-7B-A1B-Instruct-fused1import torch
2from transformers import AutoTokenizer
3
4from dinfer.model import AutoModelForCausalLM
5from dinfer.model import FusedOlmoeForCausalLM
6from dinfer import BlockIteratorFactory, KVCacheFactory
7from dinfer import ThresholdParallelDecoder, BlockWiseDiffusionLLM
8
9m = "/path/to/LLaDA-MoE-7B-A1B-Instruct-fused"
10tok = AutoTokenizer.from_pretrained(m, trust_remote_code=True)
11model = AutoModelForCausalLM.from_pretrained(m, trust_remote_code=True, torch_dtype="bfloat16")
12
13decoder = ThresholdParallelDecoder(0, threshold=0.9)
14dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(True), cache_factory=KVCacheFactory('dual'))
15
16prompt = "Lily can run 12 kilometers per hour for 4 hours. After that, she can run 6 kilometers per hour. How many kilometers can she run in 8 hours?"
17input_ids = tokenizer(prompt)['input_ids']
18input_ids = torch.tensor(input_ids).to(device).unsqueeze(0)
19res = dllm.generate(input_ids, gen_length=gen_len, block_length=block_len)transformers and its dependencies installed:1import torch
2import numpy as np
3import torch.nn.functional as F
4
5from transformers import AutoTokenizer, AutoModel
6
7
8def add_gumbel_noise(logits, temperature):
9 if temperature == 0:
10 return logits
11 logits = logits.to(torch.float64)
12 noise = torch.rand_like(logits, dtype=torch.float64)
13 gumbel_noise = (- torch.log(noise)) ** temperature
14 return logits.exp() / gumbel_noise
15
16
17def get_num_transfer_tokens(mask_index, steps):
18 mask_num = mask_index.sum(dim=1, keepdim=True)
19
20 base = mask_num // steps
21 remainder = mask_num % steps
22
23 num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base
24
25 for i in range(mask_num.size(0)):
26 num_transfer_tokens[i, :remainder[i]] += 1
27
28 return num_transfer_tokens
29
30
31@ torch.no_grad()
32def generate(model, prompt, steps=128, gen_length=128, block_length=128, temperature=0.,
33 cfg_scale=0., remasking='low_confidence', mask_id=156895):
34 x = torch.full((1, prompt.shape[1] + gen_length), mask_id, dtype=torch.long).to(model.device)
35 x[:, :prompt.shape[1]] = prompt.clone()
36 prompt_index = (x != mask_id)
37
38 assert gen_length % block_length == 0
39 num_blocks = gen_length // block_length
40 assert steps % num_blocks == 0
41 steps = steps // num_blocks
42
43 for num_block in range(num_blocks):
44 block_mask_index = (x[:, prompt.shape[1] + num_block * block_length: prompt.shape[1] + (num_block + 1) * block_length:] == mask_id)
45 num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps)
46 for i in range(steps):
47 mask_index = (x == mask_id)
48 if cfg_scale > 0.:
49 un_x = x.clone()
50 un_x[prompt_index] = mask_id
51 x_ = torch.cat([x, un_x], dim=0)
52 logits = model(x_).logits
53 logits, un_logits = torch.chunk(logits, 2, dim=0)
54 logits = un_logits + (cfg_scale + 1) * (logits - un_logits)
55 else:
56 logits = model(x).logits
57
58 logits_with_noise = add_gumbel_noise(logits, temperature=temperature)
59 x0 = torch.argmax(logits_with_noise, dim=-1) # b, l
60
61 if remasking == 'low_confidence':
62 p = F.softmax(logits, dim=-1)
63 x0_p = torch.squeeze(
64 torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1) # b, l
65 elif remasking == 'random':
66 x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)
67 else:
68 raise NotImplementedError(remasking)
69
70 x0_p[:, prompt.shape[1] + (num_block + 1) * block_length:] = -np.inf
71
72 x0 = torch.where(mask_index, x0, x)
73 confidence = torch.where(mask_index, x0_p, -np.inf)
74
75 transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)
76 for j in range(confidence.shape[0]):
77 _, select_index = torch.topk(confidence[j], k=num_transfer_tokens[j, i])
78 transfer_index[j, select_index] = True
79 x[transfer_index] = x0[transfer_index]
80
81 return x
82
83
84device = 'cuda'
85model = AutoModel.from_pretrained('inclusionAI/LLaDA-MoE-7B-A1B-Instruct', trust_remote_code=True, torch_dtype=torch.bfloat16).to(device).eval()
86tokenizer = AutoTokenizer.from_pretrained('inclusionAI/LLaDA-MoE-7B-A1B-Instruct', trust_remote_code=True)
87
88prompt = "Lily can run 12 kilometers per hour for 4 hours. After that, she runs 6 kilometers per hour. How many kilometers can she run in 8 hours?"
89m = [
90 {"role": "system", "content": "You are a helpful AI assistant."},
91 {"role": "user", "content": prompt}
92]
93prompt = tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
94
95input_ids = tokenizer(prompt)['input_ids']
96input_ids = torch.tensor(input_ids).to(device).unsqueeze(0)
97
98text = generate(model, input_ids, steps=128, gen_length=128, block_length=32, temperature=0., cfg_scale=0., remasking='low_confidence')
99print(tokenizer.batch_decode(text[:, input_ids.shape[1]:], skip_special_tokens=False)[0])@article{zhu2025llada,
title={LLaDA-MoE: A Sparse MoE Diffusion Language Model},
author={Fengqi Zhu and Zebin You and Yipeng Xing and Zenan Huang and Lin Liu and Yihong Zhuang and Guoshan Lu and Kangyu Wang and Xudong Wang and Lanning Wei and Hongrui Guo and Jiaqi Hu and Wentao Ye and Tieyuan Chen and Chenchen Li and Chengfu Tang and Haibo Feng and Jun Hu and Jun Zhou and Xiaolu Zhang and Zhenzhong Lan and Junbo Zhao and Da Zheng and Chongxuan Li and Jianguo Li and Ji-Rong Wen},
journal={arXiv preprint arXiv:2509.24389},
year={2025}
}