GSAI-ML/LLaDA-MoE-v2-30B-A3B-Instruct

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])

@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},
}