first commit

image_build/test_route_p800.py

@@ -0,0 +1,28 @@
+"""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)