@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 149 lines 4.73 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests compute_expert_num_tokens kernels."""
4
5 import dataclasses
6
7 import pytest
8 import torch
9
10 from vllm.model_executor.layers.fused_moe.utils import count_expert_num_tokens
11 from vllm.platforms import current_platform
12
13 DEVICE = current_platform.device_type
14
15 pytestmark = pytest.mark.skipif(
16 not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
17 reason="Triton MoE kernels require CUDA/ROCm/XPU.",
18 )
19
20
21 @dataclasses.dataclass
22 class TestTensors:
23 topk_ids: torch.Tensor
24 expert_map: torch.Tensor | None = None
25
26 def to_device(self, device: str):
27 self.topk_ids = self.topk_ids.to(device=device)
28 if self.expert_map is not None:
29 self.expert_map = self.expert_map.to(device=device)
30
31 @staticmethod
32 def make(
33 num_tokens: int,
34 num_topk: int,
35 num_experts: int,
36 device: str,
37 topk_ids_dtype: torch.dtype,
38 ) -> "TestTensors":
39 # make topk ids
40 topk_ids = torch.empty((num_tokens, num_topk), device=device, dtype=torch.int64)
41 for x in range(num_tokens):
42 topk_ids[x] = torch.randperm(num_experts)[:num_topk]
43 topk_ids = topk_ids.to(dtype=torch.int64)
44 return TestTensors(topk_ids=topk_ids)
45
46 def with_ep_rank(
47 self, ep_rank: int, num_global_experts: int, num_local_experts: int, device: str
48 ):
49 # make an expert map
50 expert_map = torch.empty((num_global_experts), device=device, dtype=torch.int32)
51 expert_map.fill_(-1)
52 s = ep_rank * num_local_experts
53 e = s + num_local_experts
54 expert_map[s:e] = torch.tensor(list(range(num_local_experts)), device=device)
55
56 return TestTensors(topk_ids=self.topk_ids.clone(), expert_map=expert_map)
57
58
59 def ref_impl(tt: TestTensors, expert_num_tokens: torch.Tensor):
60 # do the reference in cpu
61 tt.to_device("cpu")
62 expert_ids, counts = tt.topk_ids.unique(return_counts=True)
63
64 for eid, count in zip(expert_ids, counts):
65 if eid != -1 and tt.expert_map is not None:
66 eid = tt.expert_map[eid]
67
68 if eid == -1:
69 continue
70
71 expert_num_tokens[eid] += count
72
73
74 def do_test_compute_expert_num_tokens(
75 num_tokens: int,
76 num_topk: int,
77 num_experts: int,
78 ep_size: int,
79 topk_ids_dtype: torch.dtype,
80 ):
81 assert num_topk <= num_experts
82
83 tt = TestTensors.make(
84 num_tokens, num_topk, num_experts, topk_ids_dtype=topk_ids_dtype, device="cpu"
85 )
86
87 num_global_experts = num_experts
88 assert num_global_experts % ep_size == 0
89 num_local_experts = num_global_experts // ep_size
90 for ep_rank in range(ep_size):
91 tt_rank = tt.with_ep_rank(ep_rank, num_global_experts, num_local_experts, "cpu")
92
93 ref_expert_num_tokens = torch.zeros(
94 (num_local_experts), device="cpu", dtype=torch.int32
95 )
96 ref_impl(tt_rank, ref_expert_num_tokens)
97 ref_expert_num_tokens = ref_expert_num_tokens.to(DEVICE)
98
99 tt_rank.to_device(DEVICE)
100 # Test with expert_map
101 triton_expert_num_tokens_w_emap = count_expert_num_tokens(
102 tt_rank.topk_ids, num_local_experts, tt_rank.expert_map
103 )
104
105 # Test without expert map
106 topk_ids = tt_rank.expert_map[tt_rank.topk_ids].to(topk_ids_dtype)
107 triton_expert_num_tokens_wo_emap = count_expert_num_tokens(
108 topk_ids, num_local_experts, expert_map=None
109 )
110
111 torch.testing.assert_close(
112 ref_expert_num_tokens, triton_expert_num_tokens_w_emap, atol=0, rtol=0
113 )
114 torch.testing.assert_close(
115 ref_expert_num_tokens, triton_expert_num_tokens_wo_emap, atol=0, rtol=0
116 )
117
118
119 @pytest.mark.parametrize("num_tokens", [1, 4, 8, 11, 127, 128, 3333, 7317])
120 @pytest.mark.parametrize("num_topk", [2, 6, 8])
121 @pytest.mark.parametrize("num_experts", [64])
122 @pytest.mark.parametrize("ep_size", [1, 2, 4])
123 @pytest.mark.parametrize("topk_ids_dtype", [torch.int64])
124 def test_compute_expert_num_tokens(
125 num_tokens: int,
126 num_topk: int,
127 num_experts: int,
128 ep_size: int,
129 topk_ids_dtype: torch.dtype,
130 ):
131 do_test_compute_expert_num_tokens(
132 num_tokens, num_topk, num_experts, ep_size, topk_ids_dtype
133 )
134
135
136 @pytest.mark.parametrize("numel", list(range(1, 8192, 111)))
137 @pytest.mark.parametrize("num_experts", [32])
138 @pytest.mark.parametrize("ep_size", [2])
139 @pytest.mark.parametrize("topk_ids_dtype", [torch.int64])
140 def test_compute_expert_num_tokens_from_numel(
141 numel: int, num_experts: int, ep_size: int, topk_ids_dtype: torch.dtype
142 ):
143 do_test_compute_expert_num_tokens(
144 num_tokens=numel,
145 num_topk=1,
146 num_experts=num_experts,
147 ep_size=ep_size,
148 topk_ids_dtype=topk_ids_dtype,
149 )