transformers-qwen3-moe-fused/test_indexing.py at master · woct0rdho/transformers-qwen3-moe-fused · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
#!/usr/bin/env python3

import torch

from qwen3_moe_fused.kernels.indexing import (
    get_expert_offsets_and_idx_blocks,
    get_expert_offsets_and_idx_naive,
    get_expert_offsets_and_idx_parallel,
)


def main():
    N = 1024
    E = 128
    device = "cuda"

    s = torch.randint(0, E, (N,), device=device, dtype=torch.int32)

    offsets_naive, idx_naive, inv_idx_naive = get_expert_offsets_and_idx_naive(s, E)

    offsets_parallel, idx_parallel, inv_idx_parallel = get_expert_offsets_and_idx_parallel(s, E)
    torch.testing.assert_close(offsets_parallel, offsets_naive)
    torch.testing.assert_close(idx_parallel, idx_naive)
    torch.testing.assert_close(inv_idx_parallel, inv_idx_naive)

    offsets_blocks, idx_blocks, inv_idx_blocks = get_expert_offsets_and_idx_blocks(s, E)
    torch.testing.assert_close(offsets_blocks, offsets_naive)
    torch.testing.assert_close(idx_blocks, idx_naive)
    torch.testing.assert_close(inv_idx_blocks, inv_idx_naive)


if __name__ == "__main__":
    main()