[ROCm] Pass pointers to FlyDSL MoE kernels (#50378)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 30, 2026 at 14:49 UTC
3f90c7e9e6d3f8e874756ada2dbcd353bf0803d9
1 file changed
+46
-69
vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py
+46
-69
@@ -7,6 +7,7 @@ import json
7
import os
8
9
import flydsl.compiler as flyc
10
+import flydsl.expr as fx
11
import torch
12
from aiter.fused_moe import moe_sorting as aiter_moe_sorting
13
from aiter.ops.flydsl.kernels.moe_gemm_2stage import (
@@ -43,6 +44,12 @@ _FLYDSL_MOE_DEFAULT_CONFIG = {
44
}
45
46
47
+def _as_flydsl_ptr(tensor: torch.Tensor | None):
48
+ """Adapt a torch tensor to the pointer ABI expected by AITER FlyDSL."""
49
+ data_ptr = None if tensor is None else tensor.data_ptr()
50
+ return flyc.from_c_void_p(fx.Uint8, data_ptr)
51
+
52
+
53
def moe_sorting(
54
topk_ids: torch.Tensor,
55
topk_weights: torch.Tensor,
@@ -219,6 +226,22 @@ def fused_flydsl_moe_impl(
226
)
227
228
compiled_exe1 = _FLYDSL_MOE_GEMM1_CACHE.get(key1)
229
+ args1 = (
230
+ _as_flydsl_ptr(out_stage1),
231
+ _as_flydsl_ptr(hidden_states),
232
+ _as_flydsl_ptr(w1),
233
+ _as_flydsl_ptr(scale_x_1d),
234
+ _as_flydsl_ptr(w1_scale),
235
+ _as_flydsl_ptr(sorted_token_ids),
236
+ _as_flydsl_ptr(sorted_expert_ids),
237
+ _as_flydsl_ptr(sorted_weights_1d),
238
+ _as_flydsl_ptr(num_valid_ids),
239
+ tokens,
240
+ inter_dim,
241
+ model_dim,
242
+ int(blocks),
243
+ stream,
244
+ )
245
if compiled_exe1 is None:
246
exe1 = compile_moe_gemm1(
247
model_dim=model_dim,
@@ -235,41 +258,10 @@ def fused_flydsl_moe_impl(
258
use_cshuffle_epilog=False,
259
scale_is_bf16=scale_is_bf16,
260
)
238
- compiled_exe1 = flyc.compile(
239
- exe1,
240
- out_stage1,
241
- hidden_states,
242
- w1,
243
- scale_x_1d,
244
- w1_scale,
245
- sorted_token_ids,
246
- sorted_expert_ids,
247
- sorted_weights_1d,
248
- num_valid_ids,
249
- tokens,
250
- inter_dim,
251
- model_dim,
252
- int(blocks),
253
- stream,
254
- )
261
+ compiled_exe1 = flyc.compile(exe1, *args1)
262
_FLYDSL_MOE_GEMM1_CACHE[key1] = compiled_exe1
256
-
257
- compiled_exe1(
258
- out_stage1,
259
- hidden_states,
260
- w1,
261
- scale_x_1d,
262
- w1_scale,
263
- sorted_token_ids,
264
- sorted_expert_ids,
265
- sorted_weights_1d,
266
- num_valid_ids,
267
- tokens,
268
- inter_dim,
269
- model_dim,
270
- int(blocks),
271
- stream,
272
- )
263
+ else:
264
+ compiled_exe1(*args1)
265
266
a2_1d = out_stage1.view(-1).contiguous()
267
a2_scale_1d = torch.empty((0,), device=device, dtype=torch.float32)
@@ -291,6 +283,23 @@ def fused_flydsl_moe_impl(
283
)
284
285
compiled_exe2 = _FLYDSL_MOE_GEMM2_CACHE.get(key2)
286
+ args2 = (
287
+ _as_flydsl_ptr(out_stage2),
288
+ _as_flydsl_ptr(a2_1d),
289
+ _as_flydsl_ptr(w2),
290
+ _as_flydsl_ptr(a2_scale_1d),
291
+ _as_flydsl_ptr(w2_scale),
292
+ _as_flydsl_ptr(sorted_token_ids),
293
+ _as_flydsl_ptr(sorted_expert_ids),
294
+ _as_flydsl_ptr(sorted_weights_1d),
295
+ _as_flydsl_ptr(num_valid_ids),
296
+ tokens,
297
+ model_dim,
298
+ inter_dim,
299
+ int(blocks),
300
+ stream,
301
+ )
302
+ out_stage2.zero_()
303
if compiled_exe2 is None:
304
exe2 = compile_moe_gemm2(
305
model_dim=model_dim,
@@ -306,42 +315,10 @@ def fused_flydsl_moe_impl(
315
doweight_stage2=bool(doweight_stage2),
316
scale_is_bf16=scale_is_bf16,
317
)
309
- compiled_exe2 = flyc.compile(
310
- exe2,
311
- out_stage2,
312
- a2_1d,
313
- w2,
314
- a2_scale_1d,
315
- w2_scale,
316
- sorted_token_ids,
317
- sorted_expert_ids,
318
- sorted_weights_1d,
319
- num_valid_ids,
320
- tokens,
321
- model_dim,
322
- inter_dim,
323
- int(blocks),
324
- stream,
325
- )
318
+ compiled_exe2 = flyc.compile(exe2, *args2)
319
_FLYDSL_MOE_GEMM2_CACHE[key2] = compiled_exe2
327
-
328
- out_stage2.zero_()
329
- compiled_exe2(
330
- out_stage2,
331
- a2_1d,
332
- w2,
333
- a2_scale_1d,
334
- w2_scale,
335
- sorted_token_ids,
336
- sorted_expert_ids,
337
- sorted_weights_1d,
338
- num_valid_ids,
339
- tokens,
340
- model_dim,
341
- inter_dim,
342
- int(blocks),
343
- stream,
344
- )
320
+ else:
321
+ compiled_exe2(*args2)
322
return out_stage2
323
324