minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Refactor Req.fill_ids into full_untruncated_fill_ids + fill_len with equivalence (#26637)

Browse Source
This commit is contained in:
fzyzcjy
2026-06-08 14:52:18 +08:00
committed by GitHub
parent 4201de11de
commit 259a2da3e0
29 changed files with 136 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/bench_one_batch.py
+10 -7
View File
@@ -371,9 +371,10 @@ def prepare_inputs_for_correctness_test(bench_args, tokenizer, custom_prompts):
origin_input_ids=array("q", tmp_input_ids),
sampling_params=sampling_params,
)
req.fill_ids = req.origin_input_ids
req.full_untruncated_fill_ids = req.origin_input_ids
req.fill_len = len(req.full_untruncated_fill_ids)
req.logprob_start_len = -1
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
reqs.append(req)
return input_ids, reqs
@@ -384,14 +385,15 @@ def prepare_extend_inputs_for_correctness_test(
):
for i in range(len(reqs)):
req: Req = reqs[i]
req.fill_ids.extend(input_ids[i][bench_args.cut_len :])
req.full_untruncated_fill_ids += input_ids[i][bench_args.cut_len :]
req.fill_len = len(req.full_untruncated_fill_ids)
if model_runner is not None:
# Use req.req_pool_idx instead of i to handle slot 0 padding correctly
req.prefix_indices = model_runner.req_to_token_pool.req_to_token[
req.req_pool_idx, : bench_args.cut_len
].to(req.prefix_indices.dtype)
req.logprob_start_len = -1
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
return reqs
@@ -416,9 +418,10 @@ def prepare_synthetic_inputs_for_latency_test(
origin_input_ids=array("q", input_ids[i]),
sampling_params=sampling_params,
)
req.fill_ids = req.origin_input_ids
req.full_untruncated_fill_ids = req.origin_input_ids
req.fill_len = len(req.full_untruncated_fill_ids)
req.logprob_start_len = -1
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
reqs.append(req)
return reqs
@@ -557,7 +560,7 @@ class _MlxBenchRunner:
req_ids = [str(req.rid) for req in reqs]
results = []
for rid, req in zip(req_ids, reqs):
token_ids = [int(t) for t in req.fill_ids]
token_ids = [int(t) for t in req.get_fill_ids()]
next_token = self.mlx_runner.prefill(
req_id=rid,
new_token_ids=token_ids,