Remove normalized_prompt_logprobs from the engine to make code easier to maintain (#2902)

This commit is contained in:
Lianmin Zheng
2025-01-15 04:54:14 -08:00
parent b803b395b7
commit f65c13b559
12 changed files with 11 additions and 153 deletions
@@ -191,7 +191,6 @@ class DetokenizerManager:
input_top_logprobs_idx=recv_obj.input_top_logprobs_idx,
output_top_logprobs_val=recv_obj.output_top_logprobs_val,
output_top_logprobs_idx=recv_obj.output_top_logprobs_idx,
normalized_prompt_logprob=recv_obj.normalized_prompt_logprob,
)
)
-2
View File
@@ -340,7 +340,6 @@ class BatchTokenIDOut:
input_top_logprobs_idx: List[List]
output_top_logprobs_val: List[List]
output_top_logprobs_idx: List[List]
normalized_prompt_logprob: List[float]
@dataclass
@@ -366,7 +365,6 @@ class BatchStrOut:
input_top_logprobs_idx: List[List]
output_top_logprobs_val: List[List]
output_top_logprobs_idx: List[List]
normalized_prompt_logprob: List[float]
@dataclass
@@ -280,7 +280,6 @@ class Req:
self.top_logprobs_num = top_logprobs_num
# Logprobs (return value)
self.normalized_prompt_logprob = None
self.input_token_logprobs_val = None
self.input_token_logprobs_idx = None
self.input_top_logprobs_val = None
@@ -344,9 +343,6 @@ class Req:
max_prefix_len = min(max_prefix_len, input_len - 1)
if self.return_logprob:
if self.normalized_prompt_logprob is None:
# Need at least two tokens to compute normalized logprob
max_prefix_len = min(max_prefix_len, input_len - 2)
max_prefix_len = min(max_prefix_len, self.logprob_start_len)
max_prefix_len = max(max_prefix_len, 0)
@@ -433,7 +433,6 @@ class PrefillAdder:
or input_tokens <= self.rem_chunk_tokens
or (
req.return_logprob
and req.normalized_prompt_logprob is None
and req.logprob_start_len != len(req.origin_input_ids) - 1
)
):
+1 -12
View File
@@ -1038,9 +1038,6 @@ class Scheduler:
logits_output.input_token_logprobs = (
logits_output.input_token_logprobs.tolist()
)
logits_output.normalized_prompt_logprobs = (
logits_output.normalized_prompt_logprobs.tolist()
)
# Check finish conditions
logprob_pt = 0
@@ -1188,9 +1185,6 @@ class Scheduler:
# If logprob_start_len > 0, then first logprob_start_len prompt tokens will be ignored.
num_input_logprobs = req.extend_input_len - req.extend_logprob_start_len
if req.normalized_prompt_logprob is None:
req.normalized_prompt_logprob = output.normalized_prompt_logprobs[i]
if req.input_token_logprobs_val is None:
input_token_logprobs_val = output.input_token_logprobs[
pt : pt + num_input_logprobs - 1 - req.last_update_decode_tokens
@@ -1288,15 +1282,12 @@ class Scheduler:
input_top_logprobs_idx = []
output_top_logprobs_val = []
output_top_logprobs_idx = []
normalized_prompt_logprob = []
else:
input_token_logprobs_val = input_token_logprobs_idx = (
output_token_logprobs_val
) = output_token_logprobs_idx = input_top_logprobs_val = (
input_top_logprobs_idx
) = output_top_logprobs_val = output_top_logprobs_idx = (
normalized_prompt_logprob
) = None
) = output_top_logprobs_val = output_top_logprobs_idx = None
for req in reqs:
if req is skip_req:
@@ -1343,7 +1334,6 @@ class Scheduler:
input_top_logprobs_idx.append(req.input_top_logprobs_idx)
output_top_logprobs_val.append(req.output_top_logprobs_val)
output_top_logprobs_idx.append(req.output_top_logprobs_idx)
normalized_prompt_logprob.append(req.normalized_prompt_logprob)
# Send to detokenizer
if rids:
@@ -1370,7 +1360,6 @@ class Scheduler:
input_top_logprobs_idx,
output_top_logprobs_val,
output_top_logprobs_idx,
normalized_prompt_logprob,
)
)
else: # embedding or reward model
@@ -796,9 +796,6 @@ class TokenizerManager:
recv_obj.output_token_logprobs_idx[recv_obj_index],
return_text_in_logprobs,
)
meta_info["normalized_prompt_logprob"] = recv_obj.normalized_prompt_logprob[
recv_obj_index
]
if top_logprobs_num > 0:
meta_info["input_top_logprobs"] = self.detokenize_top_logprobs_tokens(
@@ -151,11 +151,6 @@ class TpModelWorkerClient:
logits_output.input_token_logprobs = (
logits_output.input_token_logprobs.to("cpu", non_blocking=True)
)
logits_output.normalized_prompt_logprobs = (
logits_output.normalized_prompt_logprobs.to(
"cpu", non_blocking=True
)
)
next_token_ids = next_token_ids.to("cpu", non_blocking=True)
copy_done.record()
@@ -174,9 +169,6 @@ class TpModelWorkerClient:
logits_output.input_token_logprobs = (
logits_output.input_token_logprobs.tolist()
)
logits_output.normalized_prompt_logprobs = (
logits_output.normalized_prompt_logprobs.tolist()
)
next_token_ids = next_token_ids.tolist()
return logits_output, next_token_ids