@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 221 lines 7.55 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # This is a test for the AITER ops.
4 # It tests if the AITER ops are
5 # 1. correctly registered as custom ops
6 # 2. correctly defined the relationship between
7 # implementation and fake function
8 # 3. can be used with torch.compile
9 # This file will be skipped if AITER is not installed
10 # and the platform is not ROCm.
11
12 import importlib.util
13
14 import pytest
15 import torch
16
17 from vllm.platforms import current_platform
18
19 if not current_platform.is_rocm():
20 pytest.skip("This test can only run on ROCm.", allow_module_level=True)
21
22 # this import statement is needed to ensure the ops are registered
23 import vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe # noqa: F401
24
25 # need to import once to ensure the ops are registered
26 # Check if aiter package is installed
27 aiter_available = importlib.util.find_spec("aiter") is not None
28
29 if not aiter_available:
30 pytest.skip("These tests require AITER to run.", allow_module_level=True)
31
32
33 def test_rocm_aiter_biased_grouped_topk_custom_op_registration():
34 """Test that the custom op is correctly registered."""
35 # Check if the op exists in torch.ops.vllm
36 assert hasattr(torch.ops.vllm, "rocm_aiter_biased_grouped_topk")
37
38 # Check if the op is callable
39 assert callable(torch.ops.vllm.rocm_aiter_biased_grouped_topk)
40
41
42 def test_rocm_aiter_grouped_topk_custom_op_registration():
43 """Test that the custom op is correctly registered."""
44 # Check if the op exists in torch.ops.vllm
45 assert hasattr(torch.ops.vllm, "rocm_aiter_grouped_topk")
46
47 # Check if the op is callable
48 assert callable(torch.ops.vllm.rocm_aiter_grouped_topk)
49
50
51 def test_rocm_aiter_biased_grouped_topk_torch_compile_compatibility():
52 """Test that the op can be used with torch.compile."""
53 # Create test tensors
54 token = 64
55 expert = 256
56 num_expert_group = 8
57 topk = 8
58 topk_group = 4
59 renormalize = True
60 scale_factor = 1.0
61
62 gating_output = torch.randn((token, expert), dtype=torch.bfloat16, device="cuda")
63 e_score_correction_bias = torch.randn(
64 (expert,), dtype=torch.bfloat16, device="cuda"
65 )
66
67 device = gating_output.device
68 topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device)
69 topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device)
70
71 # Define a function that uses the op
72 def biased_grouped_topk_fn(
73 gating_output, e_score_correction_bias, topk_weights, topk_ids
74 ):
75 return torch.ops.vllm.rocm_aiter_biased_grouped_topk(
76 gating_output,
77 e_score_correction_bias,
78 topk_weights,
79 topk_ids,
80 num_expert_group,
81 topk_group,
82 renormalize,
83 scale_factor,
84 )
85
86 # Verify the op's fake implementation
87 torch.library.opcheck(
88 torch.ops.vllm.rocm_aiter_biased_grouped_topk,
89 (gating_output, e_score_correction_bias, topk_weights, topk_ids),
90 kwargs={
91 "num_expert_group": num_expert_group,
92 "topk_group": topk_group,
93 "need_renorm": renormalize,
94 "routed_scaling_factor": scale_factor,
95 },
96 test_utils=("test_faketensor"),
97 )
98
99 # Compile the function with appropriate settings
100 compiled_fn = torch.compile(
101 biased_grouped_topk_fn,
102 fullgraph=True,
103 backend="inductor",
104 mode="reduce-overhead",
105 dynamic=False,
106 )
107
108 topk_weights_original = torch.empty(
109 (token, topk), dtype=torch.float32, device=device
110 )
111 topk_ids_original = torch.empty((token, topk), dtype=torch.int32, device=device)
112
113 topk_weights_compiled = torch.empty(
114 (token, topk), dtype=torch.float32, device=device
115 )
116 topk_ids_compiled = torch.empty((token, topk), dtype=torch.int32, device=device)
117
118 # Run both compiled (V1 graph mode) and uncompiled versions (V1 eager mode)
119 biased_grouped_topk_fn(
120 gating_output, e_score_correction_bias, topk_weights_original, topk_ids_original
121 )
122 compiled_fn(
123 gating_output, e_score_correction_bias, topk_weights_compiled, topk_ids_compiled
124 )
125
126 # Sort the results for comparison since the order might not be deterministic
127 topk_ids_original, indices_original = torch.sort(topk_ids_original)
128 topk_weights_original = torch.gather(topk_weights_original, 1, indices_original)
129
130 topk_ids_compiled, indices_compiled = torch.sort(topk_ids_compiled)
131 topk_weights_compiled = torch.gather(topk_weights_compiled, 1, indices_compiled)
132
133 # Verify results match
134 assert torch.allclose(
135 topk_weights_original, topk_weights_compiled, rtol=1e-2, atol=1e-2
136 )
137 assert torch.allclose(topk_ids_original, topk_ids_compiled)
138
139
140 def test_rocm_aiter_grouped_topk_torch_compile_compatibility():
141 """Test that the op can be used with torch.compile."""
142 # Create test tensors
143 token = 64
144 expert = 256
145 num_expert_group = 8
146 topk = 8
147 topk_group = 4
148 renormalize = True
149 scoring_func = "softmax"
150 scale_factor = 1.0
151
152 gating_output = torch.randn((token, expert), dtype=torch.bfloat16, device="cuda")
153
154 device = gating_output.device
155 topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device)
156 topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device)
157
158 # Define a function that uses the op
159 def grouped_topk_fn(gating_output, topk_weights, topk_ids, scoring_func):
160 return torch.ops.vllm.rocm_aiter_grouped_topk(
161 gating_output,
162 topk_weights,
163 topk_ids,
164 num_expert_group,
165 topk_group,
166 renormalize,
167 scoring_func,
168 scale_factor,
169 )
170
171 # Verify the op's fake implementation
172 torch.library.opcheck(
173 torch.ops.vllm.rocm_aiter_grouped_topk,
174 (gating_output, topk_weights, topk_ids),
175 kwargs={
176 "num_expert_group": num_expert_group,
177 "topk_group": topk_group,
178 "need_renorm": renormalize,
179 "scoring_func": scoring_func,
180 "routed_scaling_factor": scale_factor,
181 },
182 test_utils=("test_faketensor"),
183 )
184
185 # Compile the function with appropriate settings
186 compiled_fn = torch.compile(
187 grouped_topk_fn,
188 fullgraph=True,
189 backend="inductor",
190 mode="reduce-overhead",
191 dynamic=False,
192 )
193
194 topk_weights_original = torch.empty(
195 (token, topk), dtype=torch.float32, device=device
196 )
197 topk_ids_original = torch.empty((token, topk), dtype=torch.int32, device=device)
198
199 topk_weights_compiled = torch.empty(
200 (token, topk), dtype=torch.float32, device=device
201 )
202 topk_ids_compiled = torch.empty((token, topk), dtype=torch.int32, device=device)
203
204 # Run both compiled (V1 graph mode) and uncompiled versions (V1 eager mode)
205 grouped_topk_fn(
206 gating_output, topk_weights_original, topk_ids_original, scoring_func
207 )
208 compiled_fn(gating_output, topk_weights_compiled, topk_ids_compiled, scoring_func)
209
210 # Sort the results for comparison since the order might not be deterministic
211 topk_ids_original, indices_original = torch.sort(topk_ids_original)
212 topk_weights_original = torch.gather(topk_weights_original, 1, indices_original)
213
214 topk_ids_compiled, indices_compiled = torch.sort(topk_ids_compiled)
215 topk_weights_compiled = torch.gather(topk_weights_compiled, 1, indices_compiled)
216
217 # Verify results match
218 assert torch.allclose(
219 topk_weights_original, topk_weights_compiled, rtol=1e-2, atol=1e-2
220 )
221 assert torch.allclose(topk_ids_original, topk_ids_compiled)