diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index f2e6e2560e70..b7b652273728 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -10,6 +10,7 @@ import torch import vllm.envs as envs +from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import ( MoEActivation, @@ -1394,6 +1395,15 @@ def apply( expert_tokens_meta=expert_tokens_meta, ) + # Aiter don't use temp buffer, so can skip copy if output is contiguous + if ( + rocm_aiter_ops.is_fused_moe_enabled() + and not self.inplace + and fused_out.shape == output.shape + and fused_out.is_contiguous() + ): + output = fused_out + return self._finalize( output, fused_out, diff --git a/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py b/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py index d24bda101ffa..bf75b80b339e 100644 --- a/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py @@ -437,4 +437,4 @@ def apply( num_local_tokens=num_local_tokens, output_dtype=output.dtype, ) - output.copy_(result) + output.data = result diff --git a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py index 4cebe608a6b4..5c5a923648f0 100644 --- a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py +++ b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py @@ -69,7 +69,8 @@ def apply( f"But got output={output.size()}, " f"used_expert_output={fused_expert_output.size()}" ) - output.copy_(fused_expert_output, non_blocking=True) + if output.data_ptr() != fused_expert_output.data_ptr(): + output.copy_(fused_expert_output, non_blocking=True) return output