Skip to content

Commit 98f9bc9

Browse files
committed
chore: lint
Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
1 parent dc2a7ec commit 98f9bc9

File tree

1 file changed

+10
-7
lines changed

1 file changed

+10
-7
lines changed

vllm_ascend/worker/model_runner_v1.py

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -612,10 +612,12 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
612612
if batch_changed:
613613
self.input_batch.refresh_sampling_metadata()
614614

615-
def _get_forward_metadata_across_dp(self, num_tokens: int,
616-
with_prefill: bool) -> tuple[int, bool]:
615+
def _get_forward_metadata_across_dp(
616+
self, num_tokens: int,
617+
with_prefill: bool) -> tuple[torch.Tensor, bool]:
617618
local_forward_metadata = torch.tensor([num_tokens, with_prefill],
618-
device="npu", dtype=torch.int32)
619+
device="npu",
620+
dtype=torch.int32)
619621
global_forward_metadata = get_dp_group().all_gather(
620622
local_forward_metadata)
621623
num_tokens_across_dp = global_forward_metadata[:, 0].cpu()
@@ -1830,10 +1832,11 @@ def _dummy_run(
18301832
for k, v in self.intermediate_tensors.items()
18311833
})
18321834

1833-
with set_forward_context(None,
1834-
self.vllm_config,
1835-
num_tokens=num_tokens,
1836-
num_tokens_across_dp=num_tokens_across_dp):
1835+
with set_forward_context(
1836+
None,
1837+
self.vllm_config,
1838+
num_tokens=num_tokens,
1839+
num_tokens_across_dp=num_tokens_across_dp):
18371840
if self.torchair_graph_enabled and not with_prefill:
18381841
attn_metadata = self.attn_metadata_builder.build_dummy(
18391842
num_reqs=num_tokens, num_actual_tokens=1)

0 commit comments

Comments
 (0)