29 lines
1.6 KiB
Python
29 lines
1.6 KiB
Python
|
"""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)
|