Skip to content

[BugFix] Fix aclgraph accu problem in A2. #3163

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wangxiyuan merged 4 commits into from
Sep 28, 2025
Merged

[BugFix] Fix aclgraph accu problem in A2. #3163

Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
58c14f1
fix a2 accu problem
whx-sjtu Sep 24, 2025
b256ff4
make maybe_all_reduce_tensor_model_parallel custom op
whx-sjtu Sep 25, 2025
02da01a
fix lint
whx-sjtu Sep 25, 2025
64d351d
change register position
whx-sjtu Sep 26, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 29 additions & 18 deletions vllm_ascend/ops/common_fused_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
from vllm.model_executor.layers.fused_moe.layer import (
FusedMoE, UnquantizedFusedMoEMethod, determine_expert_map)
from vllm.model_executor.layers.shared_fused_moe import SharedFusedMoE
from vllm.utils import direct_register_custom_op

from vllm_ascend.ascend_config import get_ascend_config
from vllm_ascend.ascend_forward_context import MoECommType
Expand Down Expand Up @@ -203,12 +204,8 @@ def maybe_all_reduce_tensor_model_parallel(
`finalize` function. In `allgathercommimpl`, we still need to all-reduce the
outputs since each rank only has partial outputs.
"""
forward_context = get_forward_context()
moe_comm_type = forward_context.moe_comm_type
if moe_comm_type in {MoECommType.ALLTOALL, MoECommType.MC2}:
return final_hidden_states
else:
return tensor_model_parallel_all_reduce(final_hidden_states)
return torch.ops.vllm.maybe_all_reduce_tensor_model_parallel(
final_hidden_states)

def forward_impl(self, hidden_states: torch.Tensor,
router_logits: torch.Tensor):
Expand Down Expand Up @@ -333,6 +330,15 @@ def forward(
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
shared_out, fused_out = AscendFusedMoE.forward(
self,
hidden_states=hidden_states,
router_logits=router_logits,
)
return shared_out, fused_out
Comment on lines +332 to +337

@gemini-code-assist gemini-code-assist Bot Sep 24, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The call to AscendFusedMoE.forward is incorrect. The MRO for AscendSharedFusedMoE will cause this to resolve to vllm.model_executor.layers.fused_moe.layer.FusedMoE.forward, which returns a single tensor. Attempting to unpack this single tensor into two variables, shared_out and fused_out, will raise a ValueError at runtime.

Given that the implementation logic has been moved into forward_impl, which correctly returns a tuple of two tensors, the forward method should likely just call self.forward_impl.

        return self.forward_impl(
            hidden_states=hidden_states,
            router_logits=router_logits,
        )


def forward_impl(self, hidden_states: torch.Tensor,
router_logits: torch.Tensor):
# Make sure the shared experts stream begins after hidden_states are ready.
if self.multistream_overlap_shared_expert:
self.shared_expert_stream.wait_stream( # type: ignore
Expand All @@ -347,27 +353,32 @@ def forward(
moe_comm_type = forward_context.moe_comm_type
if moe_comm_type in {MoECommType.ALLTOALL, MoECommType.MC2}:
shared_out = tensor_model_parallel_all_reduce(shared_out)

_, fused_out = AscendFusedMoE.forward(
fused_output = AscendFusedMoE.forward_impl(
self,
hidden_states=hidden_states,
router_logits=router_logits,
)
# Make sure the default stream waits for the shared experts stream to finish.
if self.multistream_overlap_shared_expert:
torch.npu.current_stream().wait_stream(self.shared_expert_stream)
return shared_out, fused_out
return shared_out, fused_output


def _maybe_all_reduce_tensor_model_parallel_impl(
Comment thread
whx-sjtu marked this conversation as resolved.
Outdated
final_hidden_states: torch.Tensor) -> torch.Tensor:
forward_context = get_forward_context()
moe_comm_type = forward_context.moe_comm_type
if moe_comm_type in {MoECommType.ALLTOALL, MoECommType.MC2}:
return final_hidden_states
else:
return tensor_model_parallel_all_reduce(final_hidden_states)

def forward_impl(self, hidden_states: torch.Tensor,
router_logits: torch.Tensor):
shared_output = torch.empty(1)
fused_output = AscendFusedMoE.forward_impl(
self,
hidden_states=hidden_states,
router_logits=router_logits,
)
return shared_output, fused_output

direct_register_custom_op(op_name="maybe_all_reduce_tensor_model_parallel",
op_func=_maybe_all_reduce_tensor_model_parallel_impl,
fake_impl=lambda x, label: x,
Comment thread
whx-sjtu marked this conversation as resolved.
Outdated
mutates_args=[],
dispatch_key="PrivateUse1")

UnquantizedFusedMoEMethod.__init__ = unquantized_fused_moe_init_func
UnquantizedFusedMoEMethod.process_weights_after_loading = process_weights_after_loading
Expand Down
Loading