--- library_name: transformers license: apache-2.0 base_model: GSAI-ML/LLaDA-MoE-v2-30B-A3B-Base tags: - llada - dllm - diffusion - mixture-of-experts - moe - instruction-tuned - custom_code --- # LLaDA-MoE-v2-Instruct LLaDA MoE v2 is a family of mixture-of-experts diffusion language models. LLaDA-MoE-v2-Instruct is the instruction-tuned checkpoint in this model family. It uses the same public architecture implementation as LLaDA-MoE-v2-Base and different instruction-tuned weights. For more details, please see our paper: [LLaDA MoE v2: Scaling Mixture-of-Experts Diffusion Language Models](https://arxiv.org/abs/2608.03457). ## ⚡ Sampling with Transformers You can use the following example code to sample from LLaDA MoE v2. ```python import torch import torch.nn.functional as F from transformers import AutoModel, AutoTokenizer def add_gumbel_noise(logits, temperature): if temperature == 0: return logits logits = logits.to(torch.float64) noise = torch.rand_like(logits, dtype=torch.float64) return logits.exp() / (-torch.log(noise)) ** temperature def get_num_transfer_tokens(mask_index, steps): mask_num = mask_index.sum(dim=1, keepdim=True) base = mask_num // steps remainder = mask_num % steps num_transfer_tokens = torch.zeros( mask_num.size(0), steps, device=mask_index.device, dtype=torch.long ) + base for i in range(mask_num.size(0)): num_transfer_tokens[i, : remainder[i].item()] += 1 return num_transfer_tokens @torch.no_grad() def generate( model, prompt, mask_id, eos_id, steps=64, gen_length=1024, block_length=64, temperature=0.0, threshold=1.0, minimal_topk=1, opt_softmax=True, eos_early_stop=True, ): if prompt.ndim != 2 or prompt.size(0) != 1: raise ValueError("This compact example supports batch size 1 only.") if gen_length % block_length != 0: raise ValueError("gen_length must be divisible by block_length.") if steps <= 0 or block_length <= 0: raise ValueError("steps and block_length must be positive.") if threshold is not None and steps * minimal_topk < block_length: raise ValueError("steps * minimal_topk must cover one block.") prompt_length = prompt.size(1) x = torch.full( (1, prompt_length + block_length), mask_id, dtype=torch.long, device=prompt.device, ) x[:, :prompt_length] = prompt generated = 0 while generated < gen_length: block_start = x.size(1) - block_length block_end = x.size(1) block_mask = x[:, block_start:block_end].eq(mask_id) scheduled_transfers = get_num_transfer_tokens(block_mask, steps) for step in range(steps): mask_index = x.eq(mask_id) mask_index[:, :block_start] = False mask_index[:, block_end:] = False if not mask_index.any().item(): break logits = model(input_ids=x, use_cache=False).logits x0 = add_gumbel_noise(logits, temperature).argmax(dim=-1) if opt_softmax: masked_probs = F.softmax( logits[mask_index].float(), dim=-1 ).to(logits.dtype) else: masked_probs = F.softmax(logits[mask_index], dim=-1) masked_confidence = masked_probs.gather( dim=-1, index=x0[mask_index].unsqueeze(-1) ).squeeze(-1) confidence = torch.full( x.shape, -torch.inf, device=x.device, dtype=logits.dtype ) confidence[mask_index] = masked_confidence transfer_index = torch.zeros_like(mask_index) for batch_idx in range(x.size(0)): if threshold is None: k = scheduled_transfers[batch_idx, step].item() else: k = mask_index[batch_idx].sum().item() if k == 0: continue selected = torch.topk(confidence[batch_idx], k=k).indices if threshold is not None: keep = confidence[batch_idx, selected] >= threshold keep[: min(minimal_topk, k)] = True selected = selected[keep] transfer_index[batch_idx, selected] = True x[transfer_index] = x0[transfer_index] if x[:, block_start:block_end].eq(mask_id).any().item(): raise RuntimeError("A block was not completed; increase steps.") if eos_early_stop: eos_offsets = x[:, block_start:block_end].eq(eos_id).nonzero(as_tuple=True)[1] if eos_offsets.numel() > 0: return x[:, : block_start + eos_offsets[0].item() + 1] generated += block_length if generated < gen_length: next_block = torch.full( (1, block_length), mask_id, dtype=torch.long, device=x.device, ) x = torch.cat([x, next_block], dim=1) return x device = "cuda" model_id = "/path/to/your/model" tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) model = AutoModel.from_pretrained( model_id, trust_remote_code=True, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, ).to(device).eval() mask_id = model.config.mask_token_id model = torch.compile(model) messages = [ {"role": "system", "content": "You are a helpful AI assistant."}, {"role": "user", "content": "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?"}, ] input_ids = tokenizer.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_tensors="pt", ).to(device) output_ids = generate( model, input_ids, mask_id=mask_id, eos_id=tokenizer.convert_tokens_to_ids("<|role_end|>"), steps=64, gen_length=1024, block_length=64, temperature=0.0, threshold=1.0, minimal_topk=1, opt_softmax=True, eos_early_stop=True, ) generated_ids = output_ids[:, input_ids.size(1):] print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]) ``` ## 📚 Citation If you find our model useful, please cite: ```bibtex @misc{zhu2026lladamoev2scaling, title={LLaDA MoE v2: Scaling Mixture-of-Experts Diffusion Language Models}, 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}, year={2026}, eprint={2608.03457}, archivePrefix={arXiv}, primaryClass={cs.AI}, url={https://arxiv.org/abs/2608.03457}, } ``` ## 🌐 License This project is licensed under the terms of the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0). ## 🤝 Contact If you have any questions while using the model, feel free to contact [fengqizhu@ruc.edu.cn](mailto:fengqizhu@ruc.edu.cn).