Skip to content

Commit 263af3b

Browse files
authored
Fix a bug of ascend_forward_context (#1449)
### What this PR does / why we need it? Fix a bug of ascend_forward_context, when dp metadata is None ### Does this PR introduce _any_ user-facing change? ### How was this patch tested? Signed-off-by: liziyu <liziyu16@huawei.com>
1 parent bf17152 commit 263af3b

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

vllm_ascend/ascend_forward_context.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,11 +56,13 @@ def set_ascend_forward_context(
5656
forward_context.in_profile_run = in_profile_run
5757

5858
dp_world_size = get_dp_group().world_size
59-
if dp_world_size > 1:
59+
if dp_world_size > 1 and forward_context.dp_metadata is not None:
6060
forward_context.max_tokens_across_dp = forward_context.dp_metadata.max_tokens_across_dp_cpu.item(
6161
)
62+
elif attn_metadata is not None:
63+
forward_context.max_tokens_across_dp = num_tokens or attn_metadata.num_actual_tokens
6264
else:
63-
forward_context.max_tokens_across_dp = num_tokens
65+
forward_context.max_tokens_across_dp = None
6466

6567
try:
6668
yield

0 commit comments

Comments
 (0)