"""Tiny real P800 routing test, run after docker build with an exposed XPU."""
import torch
from vllm.platforms import current_platform
current_platform.pre_register_and_update()
for count in (2, 128):
    for scale in (1.0, 2.5):
        torch.manual_seed(17)
        logits = torch.randn(count, 256, dtype=torch.float32)
        bias = torch.randn(256, dtype=torch.float32) * 0.1
        scores = logits.sigmoid()
        adjusted = scores + bias
        groups = adjusted.view(count, 8, 32).topk(2, dim=-1).values.sum(-1).topk(4, dim=-1).indices
        mask = torch.zeros(count, 8, dtype=torch.bool).scatter_(1, groups, True).repeat_interleave(32, dim=1)
        reference_ids = adjusted.masked_fill(~mask, float('-inf')).topk(8, dim=-1).indices
        ids = torch.empty(count, 8, dtype=torch.int32, device='cuda')
        weights = torch.empty(count, 8, dtype=torch.float32, device='cuda')
        stats = torch.zeros(12, 256, dtype=torch.int32, device='cuda')
        torch.ops._C.moe_sigmoid_group_topk_norm(
            x=logits.cuda(), topk_index=ids, norm_score=weights, block_static=stats,
            bias=bias.cuda(), scale=scale, n_group=8, topk_group=4)
        torch.cuda.synchronize()
        actual_ids = ids.cpu().long()
        assert torch.equal(actual_ids.sort(-1).values, reference_ids.sort(-1).values)
        expected = scores.gather(1, actual_ids)
        expected = expected / expected.sum(-1, keepdim=True) * scale
        error = (weights.cpu() - expected).abs().max().item()
        assert error < 1e-5, error
        print('P800_ROUTE_PASS', count, scale, error, flush=True)