A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 07:15 UTC
[FEAT][ROCm] Add AITER grouped topk for DeepSeekV2 (#18825)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
vllmellm committed
May 31, 2025 at 18:39 UTC
0f5e0d567e9949ffee4f1677bc9dd767804eadd0
3 files changed
+157
-16
tests/kernels/moe/test_rocm_aiter_topk.py
+93
@@ -35,6 +35,15 @@ def test_rocm_aiter_biased_grouped_topk_custom_op_registration():
35
assert callable(torch.ops.vllm.rocm_aiter_biased_grouped_topk)
36
37
38
+def test_rocm_aiter_grouped_topk_custom_op_registration():
39
+ """Test that the custom op is correctly registered."""
40
+ # Check if the op exists in torch.ops.vllm
41
+ assert hasattr(torch.ops.vllm, 'rocm_aiter_grouped_topk')
42
+
43
+ # Check if the op is callable
44
+ assert callable(torch.ops.vllm.rocm_aiter_grouped_topk)
45
+
46
+
47
def test_rocm_aiter_biased_grouped_topk_torch_compile_compatibility():
48
"""Test that the op can be used with torch.compile."""
49
# Create test tensors
@@ -120,3 +129,87 @@ def test_rocm_aiter_biased_grouped_topk_torch_compile_compatibility():
129
rtol=1e-2,
130
atol=1e-2)
131
assert torch.allclose(topk_ids_original, topk_ids_compiled)
132
+
133
+
134
+def test_rocm_aiter_grouped_topk_torch_compile_compatibility():
135
+ """Test that the op can be used with torch.compile."""
136
+ # Create test tensors
137
+ token = 64
138
+ expert = 256
139
+ num_expert_group = 8
140
+ topk = 8
141
+ topk_group = 4
142
+ renormalize = True
143
+ scoring_func = "softmax"
144
+ scale_factor = 1.0
145
+
146
+ gating_output = torch.randn((token, expert),
147
+ dtype=torch.bfloat16,
148
+ device="cuda")
149
+
150
+ device = gating_output.device
151
+ topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device)
152
+ topk_weights = torch.empty((token, topk),
153
+ dtype=torch.float32,
154
+ device=device)
155
+
156
+ # Define a function that uses the op
157
+ def grouped_topk_fn(gating_output, topk_weights, topk_ids, scoring_func):
158
+ return torch.ops.vllm.rocm_aiter_grouped_topk(
159
+ gating_output, topk_weights, topk_ids, num_expert_group,
160
+ topk_group, renormalize, scoring_func, scale_factor)
161
+
162
+ # Verify the op's fake implementation
163
+ torch.library.opcheck(torch.ops.vllm.rocm_aiter_grouped_topk,
164
+ (gating_output, topk_weights, topk_ids),
165
+ kwargs={
166
+ "num_expert_group": num_expert_group,
167
+ "topk_group": topk_group,
168
+ "need_renorm": renormalize,
169
+ "scoring_func": scoring_func,
170
+ "routed_scaling_factor": scale_factor
171
+ },
172
+ test_utils=("test_faketensor"))
173
+
174
+ # Compile the function with appropriate settings
175
+ compiled_fn = torch.compile(grouped_topk_fn,
176
+ fullgraph=True,
177
+ backend="inductor",
178
+ mode="reduce-overhead",
179
+ dynamic=False)
180
+
181
+ topk_weights_original = torch.empty((token, topk),
182
+ dtype=torch.float32,
183
+ device=device)
184
+ topk_ids_original = torch.empty((token, topk),
185
+ dtype=torch.int32,
186
+ device=device)
187
+
188
+ topk_weights_compiled = torch.empty((token, topk),
189
+ dtype=torch.float32,
190
+ device=device)
191
+ topk_ids_compiled = torch.empty((token, topk),
192
+ dtype=torch.int32,
193
+ device=device)
194
+
195
+ # Run both compiled (V1 graph mode) and uncompiled versions (V1 eager mode)
196
+ grouped_topk_fn(gating_output, topk_weights_original, topk_ids_original,
197
+ scoring_func)
198
+ compiled_fn(gating_output, topk_weights_compiled, topk_ids_compiled,
199
+ scoring_func)
200
+
201
+ # Sort the results for comparison since the order might not be deterministic
202
+ topk_ids_original, indices_original = torch.sort(topk_ids_original)
203
+ topk_weights_original = torch.gather(topk_weights_original, 1,
204
+ indices_original)
205
+
206
+ topk_ids_compiled, indices_compiled = torch.sort(topk_ids_compiled)
207
+ topk_weights_compiled = torch.gather(topk_weights_compiled, 1,
208
+ indices_compiled)
209
+
210
+ # Verify results match
211
+ assert torch.allclose(topk_weights_original,
212
+ topk_weights_compiled,
213
+ rtol=1e-2,
214
+ atol=1e-2)
215
+ assert torch.allclose(topk_ids_original, topk_ids_compiled)
vllm/model_executor/layers/fused_moe/layer.py
+1
-1
@@ -45,7 +45,7 @@ else:
45
FusedMoEPrepareAndFinalize = None # type: ignore
46
if is_rocm_aiter_moe_enabled():
47
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
48
- rocm_aiter_biased_group_topk as grouped_topk)
48
+ rocm_aiter_grouped_topk as grouped_topk)
49
else:
50
from vllm.model_executor.layers.fused_moe.fused_moe import grouped_topk
51
if current_platform.is_tpu():
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+63
-15
@@ -140,6 +140,36 @@ def rocm_aiter_biased_grouped_topk_fake(
140
pass
141
142
143
+def rocm_aiter_grouped_topk_impl(
144
+ gating_output: torch.Tensor,
145
+ topk_weights: torch.Tensor,
146
+ topk_ids: torch.Tensor,
147
+ num_expert_group: int,
148
+ topk_group: int,
149
+ need_renorm: bool,
150
+ scoring_func: str = "softmax",
151
+ routed_scaling_factor: float = 1.0 # mul to topk_weights
152
+) -> None:
153
+
154
+ from aiter import grouped_topk
155
+
156
+ grouped_topk(gating_output, topk_weights, topk_ids, num_expert_group,
157
+ topk_group, need_renorm, scoring_func, routed_scaling_factor)
158
+
159
+
160
+def rocm_aiter_grouped_topk_fake(
161
+ gating_output: torch.Tensor,
162
+ topk_weights: torch.Tensor,
163
+ topk_ids: torch.Tensor,
164
+ num_expert_group: int,
165
+ topk_group: int,
166
+ need_renorm: bool,
167
+ scoring_func: str = "softmax",
168
+ routed_scaling_factor: float = 1.0 # mul to topk_weights
169
+) -> None:
170
+ pass
171
+
172
+
173
def rocm_aiter_fused_moe_impl(
174
hidden_states: torch.Tensor,
175
w1: torch.Tensor,
@@ -218,36 +248,54 @@ if current_platform.is_rocm():
248
dispatch_key=current_platform.dispatch_key,
249
)
250
251
+ direct_register_custom_op(
252
+ op_name="rocm_aiter_grouped_topk",
253
+ op_func=rocm_aiter_grouped_topk_impl,
254
+ mutates_args=["topk_weights", "topk_ids"],
255
+ fake_impl=rocm_aiter_grouped_topk_fake,
256
+ dispatch_key=current_platform.dispatch_key,
257
+ )
258
+
259
222
-def rocm_aiter_biased_group_topk(
260
+def rocm_aiter_grouped_topk(
261
hidden_states: torch.Tensor,
262
gating_output: torch.Tensor,
263
topk: int,
264
renormalize: bool,
265
num_expert_group: int = 0,
266
topk_group: int = 0,
229
- scoring_func: str = "sigmoid",
267
+ scoring_func: str = "softmax",
268
e_score_correction_bias: Optional[torch.Tensor] = None
269
) -> tuple[torch.Tensor, torch.Tensor]:
232
- assert scoring_func == "sigmoid", (
233
- "rocm_aiter_biased_group_topk only supports 'sigmoid' scoring_func.")
234
- assert e_score_correction_bias is not None, (
235
- "'e_score_correction_bias' must not be None.")
270
token = hidden_states.shape[0]
271
device = hidden_states.device
272
topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device)
273
topk_weights = torch.empty((token, topk),
274
dtype=torch.float32,
275
device=device)
242
- torch.ops.vllm.rocm_aiter_biased_grouped_topk(
243
- gating_output,
244
- e_score_correction_bias,
245
- topk_weights,
246
- topk_ids,
247
- num_expert_group,
248
- topk_group,
249
- renormalize,
250
- )
276
+
277
+ if e_score_correction_bias is not None:
278
+ torch.ops.vllm.rocm_aiter_biased_grouped_topk(
279
+ gating_output,
280
+ e_score_correction_bias,
281
+ topk_weights,
282
+ topk_ids,
283
+ num_expert_group,
284
+ topk_group,
285
+ renormalize,
286
+ )
287
+ else:
288
+ assert (scoring_func == "softmax" or scoring_func == "sigmoid")
289
+ torch.ops.vllm.rocm_aiter_grouped_topk(
290
+ gating_output,
291
+ topk_weights,
292
+ topk_ids,
293
+ num_expert_group,
294
+ topk_group,
295
+ renormalize,
296
+ scoring_func,
297
+ )
298
+
299
return topk_weights, topk_ids
300
301