reference/code/DMax/dInfer/tests/test.py

import os
import logging
from multiprocessing import Process
import random
import pytest

import torch
import torch.distributed as dist
from transformers import AutoTokenizer, AutoModel, AutoConfig
from vllm.config import CompilationConfig, ParallelConfig
from vllm.config import VllmConfig, set_current_vllm_config, get_current_vllm_config

from dinfer.model import LLaDAMoeModelLM, LLaDAModelLM
from dinfer import BlockWiseDiffusionLLM, VicinityCacheDiffusionLLM, IterSmoothDiffusionLLM, IterSmoothWithVicinityCacheDiffusionLLM, BlockWiseDiffusionLLMWithSP
from dinfer import ThresholdParallelDecoder, HierarchyDecoder
from dinfer import DiffusionLLMServing, SamplingParams

from dinfer.model.modeling_llada_fastdllm import LLaDAModelLM as LLaDAModelLM_fastdllm
from dinfer.decoding.generate_fastdllm import generate, generate_with_prefix_cache, generate_with_dual_cache
from dinfer.decoding.generate_dist import generate as generate_sp
from dinfer.decoding.generate_uniform import BaseDiffusionIteration
from dinfer.decoding.generate_hierarchy import generate_hierarchy
from dinfer.decoding.utils import TokenArray, DistAlignedTokenArray, BlockIterator, BlockIteratorFactory, KVCacheFactory, gather_sequence_block, BlockLoc
from dinfer.decoding.utils import DiffusionKVCacheManager
from dinfer.decoding.generate_merge import generate_merge
os.environ['TOKENIZERS_PARALLELISM'] = 'false'

from test_generate import IterSmoothDiffusionLLM as IterSmoothDiffusionLLM_test
from test_generate import IterSmoothWithVicinityCacheDiffusionLLM as IterSmoothWithVicinityCacheDiffusionLLM_test

model_path = "/mnt/dllm/model_hub/LLaDA-1.5/"
# model_path = "/data/myx/llm/vllm/model/LLaDA-1_5/"
moe_model_path = '/mnt/dllm/fengling/moe/workdir/7bA1b_anneal_15t_0827_500B_further_8k_enneal_train_4k_ep3_v7_1e-5/step45567_converted_hf_fusemoe'
# moe_model_path = '/data/dulun/models/llada-moe-sft/llada-moe-sft-model/7bA1b_anneal_19t_500B_further_8k_anneal_train_4k_ep3_v8p5/step45567_converted_hf_fusemoe/'

def test_block_iterator():
    prompt = torch.tensor([1, 2, 3, 4, 5, 6, 7]).view(1, 7)
    x = TokenArray(prompt, gen_length=10, mask_id=17, eos_id=18, device='cpu')
    it = BlockIterator(x, block_length=5)
    num_iters = 0
    for block_id, (block_loc, block) in enumerate(it):
        num_iters += 1
        assert block_loc.start == block_id * 5 + prompt.shape[1]
        assert block_loc.end == min((block_id + 1) * 5 + prompt.shape[1], prompt.shape[1] + 10)
    assert num_iters == 2

def test_token_array():
    prompt = torch.tensor([1, 2, 3, 4, 5, 6, 7]).view(1, 7)
    arr = TokenArray(prompt, gen_length=20, mask_id=17, eos_id=18, device='cpu')
    assert arr.total_length == prompt.shape[1] + 20
    assert torch.all(arr[:, 0:5] == prompt[:, 0:5])
    arr[:, 8:10] = torch.tensor([9, 10]).view(1, 2)

    arr = DistAlignedTokenArray(prompt, gen_length=20, mask_id=17, eos_id=18, device='cpu', rank=0, world_size=4)
    assert arr.total_length == prompt.shape[1] + 20 + 1
    assert torch.all(arr[:, 0:5] == prompt[:, 0:5])
    arr[:, 8:10] = torch.tensor([9, 10]).view(1, 2)

class SimulateBlockIterator:
    """ This class simulates the block iterator in VicinityCacheDiffusionLLM.
    """
    def __init__(self, x, block_length, mask_id):
        self.x = x
        self.iter = 0
        self.block_length = block_length
        self.mask_id = mask_id

    def __iter__(self):
        self.iter = 0
        return self

    def move_next(self):
        current_block_start = self.x.prompt.shape[1] + self.iter * self.block_length
        current_block_end = current_block_start + self.block_length
        current_block_end = min(current_block_end, self.x.total_length)
        # If all tokens have been decoded, move to the next block.
        if torch.all(self.x[:, current_block_start:current_block_end] != self.mask_id):
            self.iter += 1

    def __next__(self):
        self.move_next()
        current_block_start = self.x.prompt.shape[1] + self.iter * self.block_length
        if current_block_start >= self.x.total_length:
            raise StopIteration
        current_block_end = current_block_start + self.block_length
        current_block_end = min(current_block_end, self.x.total_length)
        return BlockLoc(current_block_start, current_block_end), self.x[:, current_block_start:current_block_end]

class SimulateBlockIteratorFactory:
    def create(self, x, block_length):
        return SimulateBlockIterator(x, block_length, 126336)

def get_prompts(tokenizer, mask_id, device, num=1):
    prompt = "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? "
    m = [{"role": "user", "content": prompt}, ]
    prompt = tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
    input_ids1 = torch.tensor(tokenizer(prompt)['input_ids']).to(device).unsqueeze(0)
    len1 = input_ids1.shape[1]

    if num == 2:
        prompt = "Lily can run 12 kilometers per hour for 4 hours. How many kilometers can she run in 4 hours? "
        m = [{"role": "user", "content": prompt}, ]
        prompt = tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
        input_ids2 = torch.tensor(tokenizer(prompt)['input_ids']).to(device).unsqueeze(0)
        len2 = input_ids2.shape[1]
        ret = torch.zeros(2, max(len1, len2), dtype=input_ids1.dtype)
        ret[0, 0:len1] = input_ids1
        ret[1, 0:len2] = input_ids2
    else:
        ret = input_ids1

    return ret

@ torch.no_grad()
def check_iteration(model, decoder, input_ids):
    print('check_iteration start')
    x1 = TokenArray(input_ids, 256, decoder.mask_id, decoder.eos_id, model.device)
    x2 = TokenArray(torch.cat([input_ids, input_ids]), 256, decoder.mask_id, decoder.eos_id, model.device)
    kv_cache1 = DiffusionKVCacheManager(cache_type='dual')
    kv_cache2 = DiffusionKVCacheManager(cache_type='dual')
    it1 = BlockIterator(x1, 32)
    it2 = BlockIterator(x2, 32)
    diff_iter1 = BaseDiffusionIteration()
    diff_iter2 = BaseDiffusionIteration()
    model = model.to(torch.float32)
    for block_id, ((block_loc1, block1), (block_loc2, block2)) in enumerate(zip(it1, it2)):
        output1 = model(x1.data, use_cache=True, output_hidden_states=True)
        output2 = model(x2.data, use_cache=True, output_hidden_states=True)
        for i in range(16):
            assert torch.all(output1.past_key_values._data[i] == output2.past_key_values._data[i][0])
            assert torch.all(output1.hidden_states[i][0] == output2.hidden_states[i][0]), output1.hidden_states[i][0] - output2.hidden_states[i][0]

        updated_cache1, logits1 = diff_iter1.forward(model, decoder, x1, kv_cache1, block1, block_loc1, block_id)
        updated_cache2, logits2 = diff_iter2.forward(model, decoder, x2, kv_cache2, block2, block_loc2, block_id)
        assert torch.all(logits1 == logits2)
        for layer_id, (layer1, layer2) in enumerate(zip(kv_cache1.past_key_values._data, kv_cache2.past_key_values._data)):
            assert torch.all(layer1 == layer2)
        assert torch.all(x1.data[0] == x2.data[0])
    print('check_iteration end')

def test_moe_diffusion():
    torch.cuda.set_device(0)
    device = torch.device(0)

    decoder = ThresholdParallelDecoder(0, threshold=0.9, mask_id=156895, eos_id=156892, use_float64=True)
    h_decoder = HierarchyDecoder(0, threshold=0.9, mask_id=156895, eos_id=156892, low_threshold=0.4)
    tokenizer = AutoTokenizer.from_pretrained(moe_model_path, trust_remote_code=True)
    input_ids = get_prompts(tokenizer, mask_id=156895, device=device)

    from vllm import distributed
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = random.randint(50000, 60000).__str__()
    distributed.init_distributed_environment(1, 0, 'env://', 0, 'nccl')
    distributed.initialize_model_parallel(1, backend='nccl')
    print("[Loading model]")
    # setup EP
    parallel_config = ParallelConfig(enable_expert_parallel = True)
    with set_current_vllm_config(VllmConfig(parallel_config = parallel_config)):
        model_config = AutoConfig.from_pretrained(moe_model_path, trust_remote_code=True)
        model = LLaDAMoeModelLM(config=model_config).eval()
        model.load_weights(moe_model_path, torch_dtype=torch.bfloat16)
        tokenizer = AutoTokenizer.from_pretrained(moe_model_path, trust_remote_code=True)
        model = model.to(device)

        # # Test diffusion iteration.
        # check_iteration(model, decoder, input_ids)

        # Test generation without cache.
        print('Test block-wise diffusion MOE-LLM without KV-cache')
        dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True)
        res = dllm.generate(input_ids, gen_length=128, block_length=32)
        res1, nfe = generate(model, input_ids, gen_length=128, block_length=32, threshold=0.9, mask_id=156895, eos_id=156892)
        res2, nfe = generate_merge(model, input_ids, None, gen_length=128, block_length=32, threshold=0.9, mask_id=156895, eos_id=156892, parallel_decoding='threshold', early_stop=False,)
        res1 = res1[res1 != 156892]
        res2 = res2[res2 != 156892]
        assert res.shape[1] == len(res1)
        assert res.shape[1] == len(res2)
        res1 = res1.to(res.device)
        res2 = res2.to(res.device)
        assert torch.all(res == res1)
        assert torch.all(res == res2)

        # Test generation without cache with batch size == 2.
        #print('Test block-wise diffusion MOE-LLM without KV-cache and batch size == 2')
        #input_ids2 = get_prompts(tokenizer, mask_id=156895, device=device, num=2)
        #res11 = dllm.generate(input_ids2[0].unsqueeze(0), gen_length=128, block_length=32)
        #res12 = dllm.generate(input_ids2[1].unsqueeze(0), gen_length=128, block_length=32)
        #res2 = dllm.generate(input_ids2, gen_length=128, block_length=32)
        #assert res2.shape[0] == 2
        #res21 = res2[0]
        #res22 = res2[1]
        #res21 = res21[res21 != 156892]
        #res22 = res22[res22 != 156892]
        #assert res11.shape[1] == len(res21)
        #assert res12.shape[1] == len(res22)
        # assert torch.all(res11[0] == res21)
        # assert torch.all(res12[0] == res22)

        # Test generation with iteration smooth without kv-cache.
        print('Test block-wise diffusion MOE-LLM with iteration smooth without kv-cache')
        dllm = IterSmoothDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True)
        dllm1 = IterSmoothDiffusionLLM_test(model, decoder, BlockIteratorFactory(), early_stop=True)
        res = dllm.generate(input_ids, gen_length=128, block_length=32)
        res1 = dllm1.generate(input_ids, gen_length=128, block_length=32)
        assert dllm.num_forwards == dllm1.num_forwards
        assert dllm.cache_updates == 0
        assert res.shape[1] == res1.shape[1]
        res1 = res1.to(res.device)
        assert torch.all(res == res1)

        # Test generation with dual cache
        print('Test block-wise diffusion MOE-LLM with dual KV-cache')
        dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('dual'))
        res = dllm.generate(input_ids, gen_length=256, block_length=32)
        res1, nfe = generate_with_dual_cache(model, input_ids, gen_length=256, block_length=32, threshold=0.9, mask_id=156895, eos_id=156892)
        res1 = res1[res1 != 156892]
        assert res.shape[1] == len(res1)
        res1 = res1.to(res.device)
        assert torch.all(res == res1)

        # Test generation with dual cache with batch size == 2
        #print('Test block-wise diffusion MOE-LLM with dual KV-cache and batch size == 2')
        #dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('dual'))
        #input_ids2 = get_prompts(tokenizer, mask_id=156895, device=device, num=2)
        #res2 = dllm.generate(input_ids2, gen_length=256, block_length=32)
        #assert res2.shape[0] == 2
        #res21 = res2[0]
        #res22 = res2[1]
        #res21 = res21[res21 != 156892]
        #res22 = res22[res22 != 156892]
        #assert res1.shape[0] == len(res21)
        #assert res1.shape[0] == len(res22)
        # assert torch.all(res1 == res21)
        # assert torch.all(res1 == res22)

        # Test generation with iteration smooth with kv-cache.
        print('Test block-wise diffusion MOE-LLM with iteration smooth with kv-cache')
        dllm = IterSmoothDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('dual'))
        dllm1 = IterSmoothDiffusionLLM_test(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('dual'))
        res = dllm.generate(input_ids, gen_length=128, block_length=32)
        res1 = dllm1.generate(input_ids, gen_length=128, block_length=32)
        assert dllm.num_forwards == dllm1.num_forwards
        assert dllm.cache_updates > 0
        assert dllm.cache_updates == dllm1.cache_updates
        assert res.shape[1] == res1.shape[1]
        res1 = res1.to(res.device)
        assert torch.all(res == res1)

        # Test generation with iteration smooth and vicinity cache update.
        print('Test block-wise diffusion MOE-LLM with iteration smooth with vicinity cache update')
        dllm = IterSmoothWithVicinityCacheDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('dual'))
        dllm1 = IterSmoothWithVicinityCacheDiffusionLLM_test(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('dual'))
        res = dllm.generate(input_ids, gen_length=128, block_length=32)
        res1 = dllm1.generate(input_ids, gen_length=128, block_length=32)
        assert dllm.num_forwards == dllm1.num_forwards
        assert dllm.cache_updates > 0
        assert dllm.cache_updates == dllm1.cache_updates
        assert res.shape[1] == res1.shape[1]
        res1 = res1.to(res.device)
        assert torch.all(res == res1)

        # Test generation without cache.
        print('Test block-wise hierarchical diffusion MOE-LLM without KV-cache')
        dllm = BlockWiseDiffusionLLM(model, h_decoder, BlockIteratorFactory(), early_stop=True)
        res = dllm.generate(input_ids, gen_length=128, block_length=32)
        res1, nfe = generate_hierarchy(model, input_ids, gen_length=128, block_length=32, threshold=0.9, mask_id=156895, eos_id=156892,decoding='hierarchy_fast_v2',
                                        low_threshold=0.4, remask_threshold=0.4)
        res1 = res1[res1 != 156892]
        assert res.shape[1] == len(res1)
        res1 = res1.to(res.device)
        assert torch.all(res == res1)

    distributed.destroy_model_parallel()
    distributed.destroy_distributed_environment()
    

def test_diffusion():
    torch.cuda.set_device(0)
    device = torch.device(0)
    config = AutoConfig.from_pretrained(model_path)
    config.flash_attention = True
    model = LLaDAModelLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, config=config).eval()
    model = model.to(device)
    fastdllm_model = LLaDAModelLM_fastdllm.from_pretrained(model_path, torch_dtype=torch.bfloat16, config=config).eval()
    fastdllm_model = fastdllm_model.to(device)
    decoder = ThresholdParallelDecoder(0, threshold=0.9, use_float64=True)

    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    input_ids = get_prompts(tokenizer, mask_id=126336, device=device)
    batch_size = 1
    input_ids = torch.tensor(input_ids).to(device).repeat(batch_size, 1)

    print('Test sliding-window diffusion LLM with dual KV-cache')
    dllm = VicinityCacheDiffusionLLM(model, decoder, SimulateBlockIteratorFactory(), KVCacheFactory('dual'))
    res = dllm.generate(input_ids, gen_length=128, block_length=32)
    res1, nfe = generate_with_dual_cache(fastdllm_model, input_ids, gen_length=128, block_length=32, threshold=0.9)
    res1 = res1[res1 != 126081]
    assert res.shape[1] == len(res1)
    res1 = res1.to(res.device)
    assert torch.all(res == res1)

    # Test generation without cache.
    print('Test block-wise diffusion LLM without KV-cache')
    dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True)
    res = dllm.generate(input_ids, gen_length=128, block_length=32)
    res1, nfe = generate(fastdllm_model, input_ids, gen_length=128, block_length=32, threshold=0.9)
    res1 = res1[res1 != 126081]
    assert res.shape[1] == len(res1)
    res1 = res1.to(res.device)
    assert torch.all(res == res1)

    # Test generation with prefix cache
    print('Test block-wise diffusion LLM with prefix KV-cache')
    dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True, cache_factory=KVCacheFactory('prefix'))
    res = dllm.generate(input_ids, gen_length=128, block_length=32)
    res1, nfe = generate_with_prefix_cache(fastdllm_model, input_ids, gen_length=128, block_length=32, threshold=0.9)
    res1 = res1[res1 != 126081]
    assert res.shape[1] == len(res1)
    res1 = res1.to(res.device)
    assert torch.all(res == res1)

    # Test generation with dual cache
    print('Test block-wise diffusion LLM with dual KV-cache')
    dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), cache_factory=KVCacheFactory('dual'), early_stop=True)
    res = dllm.generate(input_ids, gen_length=128, block_length=32)
    res1, nfe = generate_with_dual_cache(fastdllm_model, input_ids, gen_length=128, block_length=32, threshold=0.9)
    res1 = res1[res1 != 126081]
    assert res.shape[1] == len(res1)
    res1 = res1.to(res.device)
    assert torch.all(res == res1)

def setup_distributed(rank, world_size):
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = random.randint(30000, 40000).__str__()
    print(f'rank={rank}, world size={world_size}')
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)

def check_worker(rank, world_size, gpu):
    setup_distributed(rank, world_size)
    torch.cuda.set_device(gpu)
    device = torch.device(gpu)
    full_data = torch.arange(100).view(4, 25).to(device)

    # Partition size is smaller than block size.
    block_size = 6
    part_size = 4
    first_part_start = 1
    last_part_end = first_part_start + part_size * world_size
    assert last_part_end <= full_data.shape[1]
    partial_start = first_part_start + part_size * rank
    partial_end = partial_start + part_size
    part_data = full_data[:, partial_start:partial_end]
    # The accessed block must be covered by all parts.
    for block_start in range(first_part_start, last_part_end - block_size):
        block_end = block_start + block_size
        block_data = gather_sequence_block(part_data, partial_start, partial_end, block_start, block_end, rank, world_size)
        assert torch.all(block_data == full_data[:, block_start:block_end])

    # Partition size is larger than block size.
    block_size = 4
    part_size = 6
    first_part_start = 1
    last_part_end = first_part_start + part_size * world_size
    assert last_part_end <= full_data.shape[1]
    partial_start = first_part_start + part_size * rank
    partial_end = partial_start + part_size
    part_data = full_data[:, partial_start:partial_end]
    # The accessed block must be covered by all parts.
    for block_start in range(first_part_start, last_part_end - block_size):
        block_end = block_start + block_size
        block_data = gather_sequence_block(part_data, partial_start, partial_end, block_start, block_end, rank, world_size)
        assert torch.all(block_data == full_data[:, block_start:block_end])

    # Partition size is equal to block size.
    block_size = 4
    part_size = 4
    first_part_start = 1
    last_part_end = first_part_start + part_size * world_size
    assert last_part_end <= full_data.shape[1]
    partial_start = first_part_start + part_size * rank
    partial_end = partial_start + part_size
    part_data = full_data[:, partial_start:partial_end]
    # The accessed block must be covered by all parts.
    for block_start in range(first_part_start, last_part_end - block_size):
        block_end = block_start + block_size
        block_data = gather_sequence_block(part_data, partial_start, partial_end, block_start, block_end, rank, world_size)
        assert torch.all(block_data == full_data[:, block_start:block_end])

    dist.destroy_process_group()

def test_dist():
    num_gpus = 4
    procs = []
    for i, gpu in enumerate(range(num_gpus)):
        p = Process(target=check_worker, args=(i, num_gpus, i))
        procs.append(p)
        p.start()
    for p in procs:
        p.join()

def check_diffusion_worker(rank, world_size, gpu):
    setup_distributed(rank, world_size)
    torch.cuda.set_device(gpu)
    device = torch.device(gpu)

    config = AutoConfig.from_pretrained(model_path)
    config.flash_attention = True
    model = LLaDAModelLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, config=config).eval()
    model = model.to(device)

    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    input_ids = get_prompts(tokenizer, mask_id=126336, device=device)
    batch_size = 1
    input_ids = torch.tensor(input_ids).to(device).repeat(batch_size, 1)

    # Test generation without cache.
    print('Test diffusion LLM without KV-cache')
    decoder = ThresholdParallelDecoder(0, threshold=0.9, use_float64=True)
    dllm = BlockWiseDiffusionLLMWithSP(rank, world_size, model, decoder, BlockIteratorFactory())
    res = dllm.generate(input_ids, gen_length=128, block_length=32)
    res1, nfe = generate_sp(model, input_ids, rank=rank, world_size=world_size, gen_length=128, block_length=32, threshold=0.9)
    res1 = res1[res1 != 126081]
    assert res.shape[1] == len(res1)
    res1 = res1.to(res.device)
    assert torch.all(res == res1)

    dist.destroy_process_group()

def test_diffusion_sp():
    num_gpus = 4
    procs = []
    for i, gpu in enumerate(range(num_gpus)):
        p = Process(target=check_diffusion_worker, args=(i, num_gpus, i))
        procs.append(p)
        p.start()
    for p in procs:
        p.join()

@pytest.mark.skip(reason="Produces errors during certain runs")
def test_moe_server(require_init=True):
    print('test serving of diffusion-MOE')
    params = SamplingParams(temperature=0, threshold=0.9, mask_id=156895, eos_id=156892, early_stop=True, cache='', cont_weight=0, enable_torch_compile=True)

    device = torch.device(0)
    tokenizer = AutoTokenizer.from_pretrained(moe_model_path, trust_remote_code=True)
    input_ids = get_prompts(tokenizer, mask_id=156895, device=device)
    input_ids = torch.tensor(input_ids)

    from vllm import distributed
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = random.randint(30000, 40000).__str__()
    distributed.init_distributed_environment(1, 0, 'env://', 0, 'nccl')
    distributed.initialize_model_parallel(1, backend='nccl')

    batch_size = 1
    decoder = ThresholdParallelDecoder(0, threshold=0.9, mask_id=156895, eos_id=156892)
    # setup EP
    parallel_config = ParallelConfig(enable_expert_parallel = True)
    with set_current_vllm_config(VllmConfig(parallel_config = parallel_config)):
        model_config = AutoConfig.from_pretrained(moe_model_path, trust_remote_code=True)
        model = LLaDAMoeModelLM(config=model_config).eval()
        model.load_weights(moe_model_path, torch_dtype=torch.bfloat16)
        model.forward = torch.compile(model.forward, mode='reduce-overhead', fullgraph=False, dynamic=True)
        tokenizer = AutoTokenizer.from_pretrained(moe_model_path, trust_remote_code=True)
        model = model.to(device)

        dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True)
        res1 = dllm.generate(input_ids, gen_length=256, block_length=32).cpu()

    # Test DP == 1 and TPEP == 1
    print('Test serving: DP == 1 and TPEP == 1')
    llm = DiffusionLLMServing(model=moe_model_path, is_moe=True, sample_params=params, num_gpus=1, server_port=random.randint(50000, 60000))
    res = llm.generate(input_ids, gen_length=256, block_length=32)
    assert res.shape == res1.shape
    res1 = res1.to(res.device)
    assert torch.all(res == res1)
    llm.stop_serving()

    input_ids2 = torch.cat([input_ids, input_ids])
    # Test DP == 2 and TPEP == 1
    print('Test serving: DP == 2 and TPEP == 1')
    llm = DiffusionLLMServing(model=moe_model_path, is_moe=True, sample_params=params, num_gpus=2, dp_size=2, tpep_size=1, server_port=random.randint(50000, 60000))
    res2 = llm.generate(input_ids2, gen_length=256, block_length=32)

    # Remove EOS and padding tokens before comparison
    assert torch.all(res2[0][res2[0] != 156892] == res[0][res[0] != 156892])
    assert torch.all(res2[1][res2[1] != 156892] == res[0][res[0] != 156892])
    llm.stop_serving()

    # Test DP == 2 and TPEP == 2
    print('Test serving: DP == 2 and TPEP == 2')
    llm = DiffusionLLMServing(model=moe_model_path, is_moe=True, sample_params=params, num_gpus=2, dp_size=1, tpep_size=2, server_port=random.randint(50000, 60000))
    res = llm.generate(input_ids, gen_length=256, block_length=32)
    llm.stop_serving()

    input_ids2 = torch.cat([input_ids, input_ids])
    llm = DiffusionLLMServing(model=moe_model_path, is_moe=True, sample_params=params, num_gpus=4, dp_size=2, tpep_size=2, server_port=random.randint(40000, 50000))
    res2 = llm.generate(input_ids2, gen_length=256, block_length=32)

    assert torch.all(res2[0][res2[0] != 156892] == res[0][res[0] != 156892])
    llm.stop_serving()

    distributed.destroy_model_parallel()
    distributed.destroy_distributed_environment()

@pytest.mark.skip(reason="Produces errors during certain runs")
def test_server():
    print('test serving of diffusion')
    params = SamplingParams(temperature=0, threshold=0.9, mask_id=126336, eos_id=126081, early_stop=True, cache='', cont_weight=0, enable_torch_compile=True)

    torch.cuda.set_device(0)
    device = torch.device(0)

    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    input_ids = get_prompts(tokenizer, mask_id=126336, device=device)
    input_ids = torch.tensor(input_ids)
    config = AutoConfig.from_pretrained(model_path)
    config.flash_attention = True
    model = LLaDAModelLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, config=config).eval()
    model.forward = torch.compile(model.forward, mode='reduce-overhead', fullgraph=False, dynamic=True)
    model = model.to(device)
    # Test generation without cache.
    decoder = ThresholdParallelDecoder(0, threshold=0.9)
    dllm = BlockWiseDiffusionLLM(model, decoder, BlockIteratorFactory(), early_stop=True)
    res1 = dllm.generate(input_ids, gen_length=256, block_length=32).cpu()

    # Test DP == 1 and TPEP == 1
    print('Test serving: DP == 1 and TPEP == 1')
    llm = DiffusionLLMServing(model=model_path, is_moe=False, sample_params=params, num_gpus=1, server_port=random.randint(40000, 50000))
    res = llm.generate(input_ids, gen_length=256, block_length=32)
    llm.stop_serving()
    assert res.shape == res1.shape
    res1 = res1.to(res.device)
    assert torch.all(res == res1)


if __name__ == '__main__':
    torch.multiprocessing.set_start_method('spawn')
    logging.basicConfig(level=logging.INFO)

    test_token_array()
    test_block_iterator()

    test_moe_diffusion()
    test_diffusion()
    test_server()
    test_moe_server(False)

    test_dist()
    test_diffusion_sp()