Views
No views yet
1import torch
2import torch.nn.functional as F
3from transformers import AutoModel, AutoTokenizer
4
5
6def add_gumbel_noise(logits, temperature):
7 if temperature == 0:
8 return logits
9 logits = logits.to(torch.float64)
10 noise = torch.rand_like(logits, dtype=torch.float64)
11 return logits.exp() / (-torch.log(noise)) ** temperature
12
13
14def get_num_transfer_tokens(mask_index, steps):
15 mask_num = mask_index.sum(dim=1, keepdim=True)
16 base = mask_num // steps
17 remainder = mask_num % steps
18 num_transfer_tokens = torch.zeros(
19 mask_num.size(0), steps, device=mask_index.device, dtype=torch.long
20 ) + base
21 for i in range(mask_num.size(0)):
22 num_transfer_tokens[i, : remainder[i].item()] += 1
23 return num_transfer_tokens
24
25
26@torch.no_grad()
27def generate(
28 model,
29 prompt,
30 mask_id,
31 eos_id,
32 steps=64,
33 gen_length=1024,
34 block_length=64,
35 temperature=0.0,
36 threshold=1.0,
37 minimal_topk=1,
38 opt_softmax=True,
39 eos_early_stop=True,
40):
41 if prompt.ndim != 2 or prompt.size(0) != 1:
42 raise ValueError("This compact example supports batch size 1 only.")
43 if gen_length % block_length != 0:
44 raise ValueError("gen_length must be divisible by block_length.")
45 if steps <= 0 or block_length <= 0:
46 raise ValueError("steps and block_length must be positive.")
47 if threshold is not None and steps * minimal_topk < block_length:
48 raise ValueError("steps * minimal_topk must cover one block.")
49
50 prompt_length = prompt.size(1)
51 x = torch.full(
52 (1, prompt_length + block_length),
53 mask_id,
54 dtype=torch.long,
55 device=prompt.device,
56 )
57 x[:, :prompt_length] = prompt
58
59 generated = 0
60 while generated < gen_length:
61 block_start = x.size(1) - block_length
62 block_end = x.size(1)
63 block_mask = x[:, block_start:block_end].eq(mask_id)
64 scheduled_transfers = get_num_transfer_tokens(block_mask, steps)
65
66 for step in range(steps):
67 mask_index = x.eq(mask_id)
68 mask_index[:, :block_start] = False
69 mask_index[:, block_end:] = False
70 if not mask_index.any().item():
71 break
72
73 logits = model(input_ids=x, use_cache=False).logits
74 x0 = add_gumbel_noise(logits, temperature).argmax(dim=-1)
75
76 if opt_softmax:
77 masked_probs = F.softmax(
78 logits[mask_index].float(), dim=-1
79 ).to(logits.dtype)
80 else:
81 masked_probs = F.softmax(logits[mask_index], dim=-1)
82 masked_confidence = masked_probs.gather(
83 dim=-1, index=x0[mask_index].unsqueeze(-1)
84 ).squeeze(-1)
85
86 confidence = torch.full(
87 x.shape, -torch.inf, device=x.device, dtype=logits.dtype
88 )
89 confidence[mask_index] = masked_confidence
90 transfer_index = torch.zeros_like(mask_index)
91
92 for batch_idx in range(x.size(0)):
93 if threshold is None:
94 k = scheduled_transfers[batch_idx, step].item()
95 else:
96 k = mask_index[batch_idx].sum().item()
97 if k == 0:
98 continue
99
100 selected = torch.topk(confidence[batch_idx], k=k).indices
101 if threshold is not None:
102 keep = confidence[batch_idx, selected] >= threshold
103 keep[: min(minimal_topk, k)] = True
104 selected = selected[keep]
105 transfer_index[batch_idx, selected] = True
106
107 x[transfer_index] = x0[transfer_index]
108
109 if x[:, block_start:block_end].eq(mask_id).any().item():
110 raise RuntimeError("A block was not completed; increase steps.")
111
112 if eos_early_stop:
113 eos_offsets = x[:, block_start:block_end].eq(eos_id).nonzero(as_tuple=True)[1]
114 if eos_offsets.numel() > 0:
115 return x[:, : block_start + eos_offsets[0].item() + 1]
116
117 generated += block_length
118 if generated < gen_length:
119 next_block = torch.full(
120 (1, block_length),
121 mask_id,
122 dtype=torch.long,
123 device=x.device,
124 )
125 x = torch.cat([x, next_block], dim=1)
126
127 return x
128
129
130device = "cuda"
131model_id = "/path/to/your/model"
132tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
133model = AutoModel.from_pretrained(
134 model_id,
135 trust_remote_code=True,
136 torch_dtype=torch.bfloat16,
137 low_cpu_mem_usage=True,
138).to(device).eval()
139mask_id = model.config.mask_token_id
140model = torch.compile(model)
141
142prompt = "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?"
143input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
144output_ids = generate(
145 model,
146 input_ids,
147 mask_id=mask_id,
148 eos_id=tokenizer.eos_token_id,
149 steps=64,
150 gen_length=1024,
151 block_length=64,
152 temperature=0.0,
153 threshold=1.0,
154 minimal_topk=1,
155 opt_softmax=True,
156 eos_early_stop=True,
157)
158generated_ids = output_ids[:, input_ids.size(1):]
159print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])1@misc{zhu2026lladamoev2scaling,
2 title={LLaDA MoE v2: Scaling Mixture-of-Experts Diffusion Language Models},
3 author={Fengqi Zhu and Shaoxuan Xu and Jingyang Ou and Zebin You and Yipeng Xing and Huabin Liu and Xiaolu Zhang and Jun Zhou and Zhenzhong Lan and Yankai Lin and Wayne Xin Zhao and Jianguo Li and Chongxuan Li and Ji-Rong Wen},
4 year={2026},
5 eprint={2608.03457},
6 archivePrefix={arXiv},
7 primaryClass={cs.AI},
8 url={https://arxiv.org/abs/2608.03457},
9}