Skip to content

Commit bfec114

Browse files
committed
Update vllm_ascend/worker/model_runner_v1.py
Co-authored-by: Angazenn <92204292+Angazenn@users.noreply.github.com> Signed-off-by: jianzs <zhengshoujian.zsj@outlook.com>
1 parent 6100e0d commit bfec114

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

vllm_ascend/worker/model_runner_v1.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -625,9 +625,9 @@ def _get_forward_metadata_across_dp(
625625
with_prefill: bool) -> tuple[torch.Tensor, bool]:
626626
local_forward_metadata = torch.tensor([num_tokens, with_prefill],
627627
device="npu",
628-
dtype=torch.int32)
628+
dtype=torch.int32).unsqueeze(0)
629629
global_forward_metadata = get_dp_group().all_gather(
630-
local_forward_metadata)
630+
local_forward_metadata, dim=0)
631631
num_tokens_across_dp = global_forward_metadata[:, 0].cpu()
632632
with_prefill = bool(global_forward_metadata[:, 1].any())
633633
return num_tokens_across_dp, with_prefill

0 commit comments

Comments
 (0)