Skip to content

Commit f071cc2

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

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
@@ -620,10 +620,12 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
620620
if batch_changed:
621621
self.input_batch.refresh_sampling_metadata()
622622

623-
def _get_forward_metadata_across_dp(self, num_tokens: int,
624-
with_prefill: bool) -> tuple[int, bool]:
623+
def _get_forward_metadata_across_dp(
624+
self, num_tokens: int,
625+
with_prefill: bool) -> tuple[torch.Tensor, bool]:
625626
local_forward_metadata = torch.tensor([num_tokens, with_prefill],
626-
device="npu", dtype=torch.int32)
627+
device="npu",
628+
dtype=torch.int32)
627629
global_forward_metadata = get_dp_group().all_gather(
628630
local_forward_metadata)
629631
num_tokens_across_dp = global_forward_metadata[:, 0].cpu()
@@ -1868,10 +1870,11 @@ def _dummy_run(
18681870
for k, v in self.intermediate_tensors.items()
18691871
})
18701872

1871-
with set_forward_context(None,
1872-
self.vllm_config,
1873-
num_tokens=num_tokens,
1874-
num_tokens_across_dp=num_tokens_across_dp):
1873+
with set_forward_context(
1874+
None,
1875+
self.vllm_config,
1876+
num_tokens=num_tokens,
1877+
num_tokens_across_dp=num_tokens_across_dp):
18751878
if self.torchair_graph_enabled and not with_prefill:
18761879
attn_metadata = self.attn_metadata_builder.build_dummy(
18771880
num_reqs=num_tokens, num_actual_tokens=1)

0 commit comments

Comments
 (0)