Skip to content

[Fix] Adjust use_aclgraph logic #2156

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
merged 1 commit into from
Aug 4, 2025
+11 −2
Merged

[Fix] Adjust use_aclgraph logic #2156

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
11 changes: 10 additions & 1 deletion vllm_ascend/ops/common_fused_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
from vllm.model_executor.layers.fused_moe.layer import \
UnquantizedFusedMoEMethod

from vllm_ascend.ascend_config import get_ascend_config
from vllm_ascend.ops.fused_moe import (fused_experts, fused_experts_moge,
select_experts)
from vllm_ascend.utils import is_310p
Expand All @@ -33,7 +34,15 @@ def unquantized_fused_moe_init_func(self, *args, **kwargs):
original_unquantized_fused_moe_init_func(self, *args, **kwargs)
vllm_config = get_current_vllm_config()
self.max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
self.use_aclgraph = vllm_config.compilation_config.level == CompilationLevel.PIECEWISE and not vllm_config.model_config.enforce_eager

ascend_config = get_ascend_config()
Copy link
Collaborator

Choose a reason for hiding this comment

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

I suggest referring to #1841 and adding a new class called AscendUnquantizedFusedMoEMethod, incorporating the ACLGraph switch into this class's __init__ method, while also placing the new forward_oot method within this class.

Copy link
Collaborator

Choose a reason for hiding this comment

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

@ApsarasX it's on the todo list. we'll first combine common_fused_moe and fused_moe into one. Then refactor to Custom OP register way.

Copy link
Collaborator

Choose a reason for hiding this comment

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

@ApsarasX it's on the todo list. we'll first combine common_fused_moe and fused_moe into one. Then refactor to Custom OP register way.

OK


if ascend_config.torchair_graph_config.enabled:
self.use_aclgraph = False
else:
self.use_aclgraph = (vllm_config.compilation_config.level
== CompilationLevel.PIECEWISE
and not vllm_config.model_config.enforce_eager)


def forward_oot(
Expand Down
2 changes: 1 addition & 1 deletion vllm_ascend/ops/fused_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1084,7 +1084,7 @@ def apply(
# this is a naive implementation for experts load balance so as
# to avoid accumulating too much tokens on a single rank.
# currently it is only activated when doing profile runs.
if enable_force_load_balance:
if enable_force_load_balance and not self.use_aclgraph:
topk_ids = torch.randint_like(topk_ids, 0, global_num_experts)

fused_moe_state = get_forward_context().fused_moe_state
Expand Down
Loading