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 4, 03:45 UTC
Eliminate redundant MoE buffer copies in AITER fused experts (without dependency on AITER changes) (#41713)
Signed-off-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Mehdi Ghanimifard committed
May 7, 2026 at 12:46 UTC
d4b00484040c9a0bbd4ee2d55983df7a50ab1fd3
3 files changed
+37
-1
vllm/model_executor/layers/fused_moe/modular_kernel.py
+19
@@ -1215,6 +1215,7 @@ class FusedMoEKernelModularImpl:
1215
expert_map: torch.Tensor | None,
1216
apply_router_weight_on_input: bool,
1217
expert_tokens_meta: ExpertTokensMetadata | None,
1218
+ output_alias: torch.Tensor | None = None,
1219
) -> torch.Tensor:
1220
_, M_full, N, K, top_k = self.fused_experts.moe_problem_size(
1221
a1q, w1, w2, topk_ids
@@ -1243,6 +1244,23 @@ class FusedMoEKernelModularImpl:
1244
activation,
1245
)
1246
1247
+ # If caller's output buffer already matches fused_out shape/dtype, alias
1248
+ # to skip the redundant copy in TopKWeightAndReduceNoOP.apply downstream.
1249
+ # This eliminates ~94% of __amd_rocclr_copyBuffer events (Copy 2 of the
1250
+ # double-copy MoE write-back path).
1251
+ if current_platform.is_rocm():
1252
+ from vllm._aiter_ops import rocm_aiter_ops
1253
+
1254
+ if (
1255
+ rocm_aiter_ops.is_fused_moe_enabled()
1256
+ and output_alias is not None
1257
+ and output_alias.shape == fused_out.shape
1258
+ and output_alias.dtype == fused_out.dtype
1259
+ and output_alias.device == fused_out.device
1260
+ and output_alias.is_contiguous()
1261
+ ):
1262
+ fused_out = output_alias
1263
+
1264
self.fused_experts.apply(
1265
output=fused_out,
1266
hidden_states=a1q,
@@ -1403,6 +1421,7 @@ class FusedMoEKernelModularImpl:
1421
expert_map=expert_map,
1422
apply_router_weight_on_input=apply_router_weight_on_input,
1423
expert_tokens_meta=expert_tokens_meta,
1424
+ output_alias=output,
1425
)
1426
1427
return self._finalize(
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+14
-1
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.fused_moe.config import (
17
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
18
TopKWeightAndReduceNoOP,
19
)
20
+from vllm.model_executor.layers.fused_moe.utils import disable_inplace
21
from vllm.model_executor.layers.quantization.utils.quant_utils import (
22
QuantKey,
23
kFp8Dynamic128Sym,
@@ -438,4 +439,16 @@ class AiterExperts(mk.FusedMoEExpertsModular):
439
num_local_tokens=num_local_tokens,
440
output_dtype=output.dtype,
441
)
441
- output.copy_(result)
442
+ # avoid redundant copy when output is a view of the result
443
+ if (
444
+ output.shape == result.shape
445
+ and output.dtype == result.dtype
446
+ and output.device == result.device
447
+ and output.is_contiguous()
448
+ and result.is_contiguous()
449
+ and output._base is None
450
+ and disable_inplace()
451
+ ):
452
+ output.set_(result)
453
+ else:
454
+ output.copy_(result)
vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py
+4
@@ -62,6 +62,10 @@ class TopKWeightAndReduceNoOP(mk.TopKWeightAndReduce):
62
if output is None:
63
return fused_expert_output
64
65
+ # Skip self-copy when caller aliased fused_out to output upstream.
66
+ if output is fused_expert_output:
67
+ return output
68
+
69
# MoEPrepareAndFinalizeNoDPEPModular needs the output to be in the `output`
70
# tensor.
71
assert output.size() == fused_expert_output.size(), (