Refactor: Moving extend_logprob_start_len calculation out of prepare_for_extend (#16105)

This commit is contained in:
Cheng Wan
2025-12-29 20:38:33 -08:00
committed by GitHub
parent 9e263c2162
commit 60f1ca6925
7 changed files with 25 additions and 36 deletions

View File

@@ -301,8 +301,8 @@ def prepare_inputs_for_correctness_test(bench_args, tokenizer, custom_prompts):
sampling_params=sampling_params,
)
req.fill_ids = req.origin_input_ids
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
req.logprob_start_len = -1
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
reqs.append(req)
return input_ids, reqs
@@ -312,13 +312,13 @@ def prepare_extend_inputs_for_correctness_test(
bench_args, input_ids, reqs, model_runner
):
for i in range(len(reqs)):
req = reqs[i]
req: Req = reqs[i]
req.fill_ids += input_ids[i][bench_args.cut_len :]
req.prefix_indices = model_runner.req_to_token_pool.req_to_token[
i, : bench_args.cut_len
]
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
req.logprob_start_len = -1
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
return reqs
@@ -344,8 +344,8 @@ def prepare_synthetic_inputs_for_latency_test(
sampling_params=sampling_params,
)
req.fill_ids = req.origin_input_ids
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
req.logprob_start_len = -1
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
reqs.append(req)
return reqs