Skip to content

Commit 7726146

Browse files
committed
Update vllm_ascend/worker/model_runner_v1.py
Co-authored-by: Angazenn <92204292+Angazenn@users.noreply.github.com> Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
1 parent e3a9cd5 commit 7726146

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
@@ -631,9 +631,9 @@ def _get_forward_metadata_across_dp(
631631
with_prefill: bool) -> tuple[torch.Tensor, bool]:
632632
local_forward_metadata = torch.tensor([num_tokens, with_prefill],
633633
device="npu",
634-
dtype=torch.int32)
634+
dtype=torch.int32).unsqueeze(0)
635635
global_forward_metadata = get_dp_group().all_gather(
636-
local_forward_metadata)
636+
local_forward_metadata, dim=0)
637637
num_tokens_across_dp = global_forward_metadata[:, 0].cpu()
638638
with_prefill = bool(global_forward_metadata[:, 1].any())
639639
return num_tokens_across_dp, with_prefill

0 commit comments

Comments
 (0)