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

Pack scattered request logprob state into a dedicated container (#25712)

Browse Source
This commit is contained in:
fzyzcjy
2026-05-19 09:18:19 +08:00
committed by GitHub
parent fb7e49d4eb
commit 2d40f45193
9 changed files with 182 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+10 -6
View File
@@ -1450,17 +1450,21 @@ class DecodeTransferQueue:
decode_req.req.hidden_states_tensor = output_hidden_states
if decode_req.req.return_logprob:
decode_req.req.output_token_logprobs_val.append(
decode_req.req.logprob.output_token_logprobs_val.append(
output_token_logprobs_val[0].item()
)
decode_req.req.output_token_logprobs_idx.append(
decode_req.req.logprob.output_token_logprobs_idx.append(
output_token_logprobs_idx[0].item()
)
decode_req.req.output_top_logprobs_val.append(
output_top_logprobs_val[: decode_req.req.top_logprobs_num].tolist()
decode_req.req.logprob.output_top_logprobs_val.append(
output_top_logprobs_val[
: decode_req.req.logprob.top_logprobs_num
].tolist()
)
decode_req.req.output_top_logprobs_idx.append(
output_top_logprobs_idx[: decode_req.req.top_logprobs_num].tolist()
decode_req.req.logprob.output_top_logprobs_idx.append(
output_top_logprobs_idx[
: decode_req.req.logprob.top_logprobs_num
].tolist()
)
decode_req.kv_receiver.clear()
python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py
+2 -2
View File
@@ -85,8 +85,8 @@ class ScheduleBatchDisaggregationDecodeMixin:
self.seq_lens_sum = sum(seq_lens)
if self.return_logprob:
self.top_logprobs_nums = [r.top_logprobs_num for r in reqs]
self.token_ids_logprobs = [r.token_ids_logprob for r in reqs]
self.top_logprobs_nums = [r.logprob.top_logprobs_num for r in reqs]
self.token_ids_logprobs = [r.logprob.token_ids_logprob for r in reqs]
self.extend_num_tokens = extend_num_tokens
self.prefix_lens = [len(r.prefix_indices) for r in reqs]
python/sglang/srt/disaggregation/utils.py
+20 -16
View File
@@ -263,26 +263,30 @@ class MetadataBuffers:
self.cached_tokens[req.metadata_buffer_index][2] = req.cached_tokens_host
self.cached_tokens[req.metadata_buffer_index][3] = req.cached_tokens_storage
if req.return_logprob:
if req.output_token_logprobs_val: # not none or empty list
if req.logprob.output_token_logprobs_val: # not none or empty list
self.output_token_logprobs_val[req.metadata_buffer_index][0] = (
req.output_token_logprobs_val[0]
req.logprob.output_token_logprobs_val[0]
)
if req.output_token_logprobs_idx: # not none or empty list
if req.logprob.output_token_logprobs_idx: # not none or empty list
self.output_token_logprobs_idx[req.metadata_buffer_index][0] = (
req.output_token_logprobs_idx[0]
req.logprob.output_token_logprobs_idx[0]
)
if req.output_top_logprobs_val: # not none or empty list
if req.logprob.output_top_logprobs_val: # not none or empty list
self.output_top_logprobs_val[req.metadata_buffer_index][
: len(req.output_top_logprobs_val[0])
: len(req.logprob.output_top_logprobs_val[0])
] = torch.tensor(
req.output_top_logprobs_val[0], dtype=torch.float32, device="cpu"
req.logprob.output_top_logprobs_val[0],
dtype=torch.float32,
device="cpu",
)
if req.output_top_logprobs_idx: # not none or empty list
if req.logprob.output_top_logprobs_idx: # not none or empty list
self.output_top_logprobs_idx[req.metadata_buffer_index][
: len(req.output_top_logprobs_idx[0])
: len(req.logprob.output_top_logprobs_idx[0])
] = torch.tensor(
req.output_top_logprobs_idx[0], dtype=torch.int32, device="cpu"
req.logprob.output_top_logprobs_idx[0],
dtype=torch.int32,
device="cpu",
)
# For PD + spec decode
if req.hidden_states_tensor is not None:
@@ -635,9 +639,9 @@ def prepare_abort(req: Req, error_message: str, status_code=None):
req.finished_reason = FINISH_ABORT(error_message, status_code)
if req.return_logprob:
req.input_token_logprobs_val = []
req.input_token_logprobs_idx = []
req.input_top_logprobs_val = []
req.input_top_logprobs_idx = []
req.input_token_ids_logprobs_val = []
req.input_token_ids_logprobs_idx = []
req.logprob.input_token_logprobs_val = []
req.logprob.input_token_logprobs_idx = []
req.logprob.input_top_logprobs_val = []
req.logprob.input_top_logprobs_idx = []
req.logprob.input_token_ids_logprobs_val = []
req.logprob.input_token_ids_logprobs_idx = []
python/sglang/srt/layers/utils/logprob.py
+19 -8
View File
@@ -415,17 +415,28 @@ def add_output_logprobs_for_spec_v1(
for req, num_tokens in zip(batch.reqs, num_tokens_per_req, strict=True):
for _ in range(num_tokens):
if req.return_logprob:
req.output_token_logprobs_val.append(next_token_logprobs[pt])
req.output_token_logprobs_idx.append(accept_tokens_list[pt])
if req.top_logprobs_num > 0:
req.logprob.output_token_logprobs_val.append(next_token_logprobs[pt])
req.logprob.output_token_logprobs_idx.append(accept_tokens_list[pt])
if req.logprob.top_logprobs_num > 0:
assert (
should_top_logprobs
), "Inconsistent state: should_top_logprobs is False"
req.output_top_logprobs_val.append(token_top_logprobs_val[pt])
req.output_top_logprobs_idx.append(token_top_logprobs_idx[pt])
if req.token_ids_logprob is not None and should_token_ids_logprobs:
req.output_token_ids_logprobs_val.append(token_ids_logprobs_val[pt])
req.output_token_ids_logprobs_idx.append(token_ids_logprobs_idx[pt])
req.logprob.output_top_logprobs_val.append(
token_top_logprobs_val[pt]
)
req.logprob.output_top_logprobs_idx.append(
token_top_logprobs_idx[pt]
)
if (
req.logprob.token_ids_logprob is not None
and should_token_ids_logprobs
):
req.logprob.output_token_ids_logprobs_val.append(
token_ids_logprobs_val[pt]
)
req.logprob.output_token_ids_logprobs_idx.append(
token_ids_logprobs_idx[pt]
)
pt += 1
python/sglang/srt/managers/schedule_batch.py
+33 -24
View File
@@ -568,6 +568,27 @@ class MultimodalInputs:
# other args would be kept intact
@dataclasses.dataclass(slots=True, kw_only=True)
class ReqLogprob:
top_logprobs_num: int
token_ids_logprob: Optional[List[int]]
input_token_logprobs_val: Optional[List[float]] = None
input_token_logprobs_idx: Optional[List[int]] = None
input_top_logprobs_val: Optional[List[List[float]]] = None
input_top_logprobs_idx: Optional[List[List[int]]] = None
input_token_ids_logprobs_val: Optional[List[List[float]]] = None
input_token_ids_logprobs_idx: Optional[List[List[int]]] = None
output_token_logprobs_val: Optional[list] = None
output_token_logprobs_idx: Optional[list] = None
output_top_logprobs_val: Optional[list] = None
output_top_logprobs_idx: Optional[list] = None
# Can contain either lists or GPU tensors (delayed copy optimization for prefill-only scoring)
output_token_ids_logprobs_val: Optional[List[Union[List[float], torch.Tensor]]] = (
None
)
output_token_ids_logprobs_idx: Optional[list] = None
class Req(ReqDllmMixin):
"""The input and output status of a request."""
@@ -767,18 +788,14 @@ class Req(ReqDllmMixin):
self.return_logprob = return_logprob
# Start index to compute logprob from.
self.logprob_start_len = 0
self.top_logprobs_num = top_logprobs_num
self.token_ids_logprob = token_ids_logprob
self.logprob = ReqLogprob(
top_logprobs_num=top_logprobs_num,
token_ids_logprob=token_ids_logprob,
)
# Logprobs (return values)
# True means the input logprob has been already sent to detokenizer.
self.input_logprob_sent: bool = False
self.input_token_logprobs_val: Optional[List[float]] = None
self.input_token_logprobs_idx: Optional[List[int]] = None
self.input_top_logprobs_val: Optional[List[float]] = None
self.input_top_logprobs_idx: Optional[List[int]] = None
self.input_token_ids_logprobs_val: Optional[List[float]] = None
self.input_token_ids_logprobs_idx: Optional[List[int]] = None
# Temporary holder to store input_token_logprobs.
self.input_token_logprobs: Optional[List[Tuple[int]]] = None
self.temp_input_top_logprobs_val: Optional[List[torch.Tensor]] = None
@@ -788,22 +805,14 @@ class Req(ReqDllmMixin):
if return_logprob:
# shape: (bs, 1)
self.output_token_logprobs_val = []
self.output_token_logprobs_idx = []
self.logprob.output_token_logprobs_val = []
self.logprob.output_token_logprobs_idx = []
# shape: (bs, k)
self.output_top_logprobs_val = []
self.output_top_logprobs_idx = []
self.logprob.output_top_logprobs_val = []
self.logprob.output_top_logprobs_idx = []
# Can contain either lists or GPU tensors (delayed copy optimization for prefill-only scoring)
self.output_token_ids_logprobs_val: List[
Union[List[float], torch.Tensor]
] = []
self.output_token_ids_logprobs_idx = []
else:
self.output_token_logprobs_val = self.output_token_logprobs_idx = (
self.output_top_logprobs_val
) = self.output_top_logprobs_idx = self.output_token_ids_logprobs_val = (
self.output_token_ids_logprobs_idx
) = None
self.logprob.output_token_ids_logprobs_val = []
self.logprob.output_token_ids_logprobs_idx = []
self.hidden_states: List[List[float]] = []
self.hidden_states_tensor = None # Note: use tensor instead of list to transfer hidden_states when PD + MTP
self.output_topk_p = None
@@ -1954,8 +1963,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.multi_item_delimiter_indices = None
if self.return_logprob:
self.top_logprobs_nums = [r.top_logprobs_num for r in reqs]
self.token_ids_logprobs = [r.token_ids_logprob for r in reqs]
self.top_logprobs_nums = [r.logprob.top_logprobs_num for r in reqs]
self.token_ids_logprobs = [r.logprob.token_ids_logprob for r in reqs]
self.extend_logprob_start_lens = [r.extend_logprob_start_len for r in reqs]
self.extend_input_logprob_token_ids = extend_input_logprob_token_ids
python/sglang/srt/managers/scheduler.py
+1 -1
View File
@@ -2272,7 +2272,7 @@ class Scheduler(
)
if batch.return_logprob:
batch.top_logprobs_nums = [r.top_logprobs_num for r in reqs]
batch.top_logprobs_nums = [r.logprob.top_logprobs_num for r in reqs]
batch.token_ids_logprobs = [list(r.origin_input_ids) for r in reqs]
batch.sampling_info = SamplingBatchInfo.from_schedule_batch(
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+8 -8
View File
@@ -712,22 +712,22 @@ class SchedulerBatchResultProcessor:
max_accept = 1
for j, tok_id in enumerate(accepted_ids):
req.output_token_logprobs_val.append(accepted_logprobs[j])
req.output_token_logprobs_idx.append(tok_id)
if req.top_logprobs_num > 0:
req.logprob.output_token_logprobs_val.append(accepted_logprobs[j])
req.logprob.output_token_logprobs_idx.append(tok_id)
if req.logprob.top_logprobs_num > 0:
flat_idx = i * max_accept + j
req.output_top_logprobs_val.append(
req.logprob.output_top_logprobs_val.append(
logits_output.next_token_top_logprobs_val[flat_idx]
)
req.output_top_logprobs_idx.append(
req.logprob.output_top_logprobs_idx.append(
logits_output.next_token_top_logprobs_idx[flat_idx]
)
if req.token_ids_logprob is not None:
if req.logprob.token_ids_logprob is not None:
flat_idx = i * max_accept + j
req.output_token_ids_logprobs_val.append(
req.logprob.output_token_ids_logprobs_val.append(
logits_output.next_token_token_ids_logprobs_val[flat_idx]
)
req.output_token_ids_logprobs_idx.append(
req.logprob.output_token_ids_logprobs_idx.append(
logits_output.next_token_token_ids_logprobs_idx[flat_idx]
)
python/sglang/srt/managers/scheduler_components/logprob_result_processor.py
+72 -54
View File
@@ -31,10 +31,10 @@ class SchedulerLogprobResultProcessor:
# Process logprob values - handle multi-item scoring vs regular requests
if is_multi_item_scoring:
# Multi-item scoring: use all logprobs as-is
req.input_token_logprobs_val = input_token_logprobs
req.logprob.input_token_logprobs_val = input_token_logprobs
else:
# Regular request: add None at start, remove last (sampling token)
req.input_token_logprobs_val = [None] + input_token_logprobs[:-1]
req.logprob.input_token_logprobs_val = [None] + input_token_logprobs[:-1]
# Process logprob indices based on scoring type
if is_multi_item_scoring:
@@ -49,21 +49,21 @@ class SchedulerLogprobResultProcessor:
input_token_logprobs_idx = req.origin_input_ids[req.logprob_start_len :]
# Clip padded hash values from image tokens to prevent detokenization errors
req.input_token_logprobs_idx = [
req.logprob.input_token_logprobs_idx = [
x if x < self.model_config.vocab_size - 1 else 0
for x in input_token_logprobs_idx
]
def _process_input_top_logprobs(self, req: Req) -> None:
"""Process input top logprobs."""
if req.top_logprobs_num <= 0:
if req.logprob.top_logprobs_num <= 0:
return
is_multi_item_scoring = self._is_multi_item_scoring(req)
# Initialize arrays - multi-item scoring starts empty, others start with None
req.input_top_logprobs_val = [] if is_multi_item_scoring else [None]
req.input_top_logprobs_idx = [] if is_multi_item_scoring else [None]
req.logprob.input_top_logprobs_val = [] if is_multi_item_scoring else [None]
req.logprob.input_top_logprobs_idx = [] if is_multi_item_scoring else [None]
# Extend arrays with temp values
for val, idx in zip(
@@ -71,13 +71,13 @@ class SchedulerLogprobResultProcessor:
req.temp_input_top_logprobs_idx,
strict=True,
):
req.input_top_logprobs_val.extend(val)
req.input_top_logprobs_idx.extend(idx)
req.logprob.input_top_logprobs_val.extend(val)
req.logprob.input_top_logprobs_idx.extend(idx)
# Remove last token (sampling token) for non multi-item scoring requests
if not is_multi_item_scoring:
req.input_top_logprobs_val.pop()
req.input_top_logprobs_idx.pop()
req.logprob.input_top_logprobs_val.pop()
req.logprob.input_top_logprobs_idx.pop()
# Clean up temp storage
req.temp_input_top_logprobs_idx = None
@@ -85,14 +85,18 @@ class SchedulerLogprobResultProcessor:
def _process_input_token_ids_logprobs(self, req: Req) -> None:
"""Process input token IDs logprobs."""
if req.token_ids_logprob is None:
if req.logprob.token_ids_logprob is None:
return
is_multi_item_scoring = self._is_multi_item_scoring(req)
# Initialize arrays - multi-item scoring starts empty, others start with None
req.input_token_ids_logprobs_val = [] if is_multi_item_scoring else [None]
req.input_token_ids_logprobs_idx = [] if is_multi_item_scoring else [None]
req.logprob.input_token_ids_logprobs_val = (
[] if is_multi_item_scoring else [None]
)
req.logprob.input_token_ids_logprobs_idx = (
[] if is_multi_item_scoring else [None]
)
# Process temp values - convert tensors to lists and extend arrays
for val, idx in zip(
@@ -101,15 +105,15 @@ class SchedulerLogprobResultProcessor:
strict=True,
):
val_list = val.tolist() if isinstance(val, torch.Tensor) else val
req.input_token_ids_logprobs_val.extend(
req.logprob.input_token_ids_logprobs_val.extend(
val_list if isinstance(val_list, list) else [val_list]
)
req.input_token_ids_logprobs_idx.extend(idx)
req.logprob.input_token_ids_logprobs_idx.extend(idx)
# Remove last token (sampling token) for non multi-item scoring requests
if not is_multi_item_scoring:
req.input_token_ids_logprobs_val.pop()
req.input_token_ids_logprobs_idx.pop()
req.logprob.input_token_ids_logprobs_val.pop()
req.logprob.input_token_ids_logprobs_idx.pop()
# Clean up temp storage
req.temp_input_token_ids_logprobs_idx = None
@@ -198,11 +202,11 @@ class SchedulerLogprobResultProcessor:
if req.temp_input_token_ids_logprobs_idx is None:
req.temp_input_token_ids_logprobs_idx = []
if req.input_token_logprobs_val is not None:
if req.logprob.input_token_logprobs_val is not None:
# The input logprob has been already computed. It only happens
# upon retract.
if req.top_logprobs_num > 0:
assert req.input_token_logprobs_val is not None
if req.logprob.top_logprobs_num > 0:
assert req.logprob.input_token_logprobs_val is not None
return
# Important for the performance.
@@ -213,11 +217,11 @@ class SchedulerLogprobResultProcessor:
]
req.input_token_logprobs.extend(input_token_logprobs)
if req.top_logprobs_num > 0:
if req.logprob.top_logprobs_num > 0:
req.temp_input_top_logprobs_val.append(output.input_top_logprobs_val[i])
req.temp_input_top_logprobs_idx.append(output.input_top_logprobs_idx[i])
if req.token_ids_logprob is not None:
if req.logprob.token_ids_logprob is not None:
req.temp_input_token_ids_logprobs_val.append(
output.input_token_ids_logprobs_val[i]
)
@@ -228,10 +232,10 @@ class SchedulerLogprobResultProcessor:
if last_prefill_chunk:
input_token_logprobs = req.input_token_logprobs
req.input_token_logprobs = None
assert req.input_token_logprobs_val is None
assert req.input_token_logprobs_idx is None
assert req.input_top_logprobs_val is None
assert req.input_top_logprobs_idx is None
assert req.logprob.input_token_logprobs_val is None
assert req.logprob.input_token_logprobs_idx is None
assert req.logprob.input_top_logprobs_val is None
assert req.logprob.input_top_logprobs_idx is None
# Process all input logprob types using helper functions
self._process_input_token_logprobs(req, input_token_logprobs)
@@ -241,14 +245,24 @@ class SchedulerLogprobResultProcessor:
if req.return_logprob:
relevant_tokens_len = self._calculate_relevant_tokens_len(req)
assert len(req.input_token_logprobs_val) == relevant_tokens_len
assert len(req.input_token_logprobs_idx) == relevant_tokens_len
if req.top_logprobs_num > 0:
assert len(req.input_top_logprobs_val) == relevant_tokens_len
assert len(req.input_top_logprobs_idx) == relevant_tokens_len
if req.token_ids_logprob is not None:
assert len(req.input_token_ids_logprobs_val) == relevant_tokens_len
assert len(req.input_token_ids_logprobs_idx) == relevant_tokens_len
assert len(req.logprob.input_token_logprobs_val) == relevant_tokens_len
assert len(req.logprob.input_token_logprobs_idx) == relevant_tokens_len
if req.logprob.top_logprobs_num > 0:
assert (
len(req.logprob.input_top_logprobs_val) == relevant_tokens_len
)
assert (
len(req.logprob.input_top_logprobs_idx) == relevant_tokens_len
)
if req.logprob.token_ids_logprob is not None:
assert (
len(req.logprob.input_token_ids_logprobs_val)
== relevant_tokens_len
)
assert (
len(req.logprob.input_token_ids_logprobs_idx)
== relevant_tokens_len
)
def add_logprob_return_values(
self,
@@ -261,8 +275,8 @@ class SchedulerLogprobResultProcessor:
):
"""Attach logprobs to the return values."""
if output.next_token_logprobs is not None:
req.output_token_logprobs_val.append(output.next_token_logprobs[i])
req.output_token_logprobs_idx.append(next_token_ids[i])
req.logprob.output_token_logprobs_val.append(output.next_token_logprobs[i])
req.logprob.output_token_logprobs_idx.append(next_token_ids[i])
# Only add input logprobs if there are input tokens to process
# Note: For prefill-only requests with default logprob_start_len, this will be 0,
@@ -279,20 +293,24 @@ class SchedulerLogprobResultProcessor:
else:
self._initialize_empty_logprob_containers(req)
if req.top_logprobs_num > 0:
req.output_top_logprobs_val.append(output.next_token_top_logprobs_val[i])
req.output_top_logprobs_idx.append(output.next_token_top_logprobs_idx[i])
if req.logprob.top_logprobs_num > 0:
req.logprob.output_top_logprobs_val.append(
output.next_token_top_logprobs_val[i]
)
req.logprob.output_top_logprobs_idx.append(
output.next_token_top_logprobs_idx[i]
)
if (
req.token_ids_logprob is not None
req.logprob.token_ids_logprob is not None
and output.next_token_token_ids_logprobs_val is not None
):
# Convert GPU tensor to list if needed
logprobs_val = output.next_token_token_ids_logprobs_val[i]
if isinstance(logprobs_val, torch.Tensor):
logprobs_val = logprobs_val.tolist()
req.output_token_ids_logprobs_val.append(logprobs_val)
req.output_token_ids_logprobs_idx.append(
req.logprob.output_token_ids_logprobs_val.append(logprobs_val)
req.logprob.output_token_ids_logprobs_idx.append(
output.next_token_token_ids_logprobs_idx[i]
)
@@ -305,15 +323,15 @@ class SchedulerLogprobResultProcessor:
This is needed for prefill-only requests where the normal initialization
flow might be bypassed, but downstream code expects these fields to be lists.
"""
if req.input_token_logprobs_val is None:
req.input_token_logprobs_val = []
if req.input_token_logprobs_idx is None:
req.input_token_logprobs_idx = []
if req.input_top_logprobs_val is None:
req.input_top_logprobs_val = []
if req.input_top_logprobs_idx is None:
req.input_top_logprobs_idx = []
if req.input_token_ids_logprobs_val is None:
req.input_token_ids_logprobs_val = []
if req.input_token_ids_logprobs_idx is None:
req.input_token_ids_logprobs_idx = []
if req.logprob.input_token_logprobs_val is None:
req.logprob.input_token_logprobs_val = []
if req.logprob.input_token_logprobs_idx is None:
req.logprob.input_token_logprobs_idx = []
if req.logprob.input_top_logprobs_val is None:
req.logprob.input_top_logprobs_val = []
if req.logprob.input_top_logprobs_idx is None:
req.logprob.input_top_logprobs_idx = []
if req.logprob.input_token_ids_logprobs_val is None:
req.logprob.input_token_ids_logprobs_val = []
if req.logprob.input_token_ids_logprobs_idx is None:
req.logprob.input_token_ids_logprobs_idx = []
python/sglang/srt/managers/scheduler_components/output_streamer.py
+17 -13
View File
@@ -368,17 +368,21 @@ class _GenerationStreamAccumulator:
# Decode server does not send input logprobs
and self.disaggregation_mode != DisaggregationMode.DECODE
# Only send when input logprobs have been computed (after prefill)
and req.input_token_logprobs_val is not None
and req.logprob.input_token_logprobs_val is not None
):
self.input_token_logprobs_val.append(req.input_token_logprobs_val)
self.input_token_logprobs_idx.append(req.input_token_logprobs_idx)
self.input_top_logprobs_val.append(req.input_top_logprobs_val)
self.input_top_logprobs_idx.append(req.input_top_logprobs_idx)
self.input_token_logprobs_val.append(
req.logprob.input_token_logprobs_val
)
self.input_token_logprobs_idx.append(
req.logprob.input_token_logprobs_idx
)
self.input_top_logprobs_val.append(req.logprob.input_top_logprobs_val)
self.input_top_logprobs_idx.append(req.logprob.input_top_logprobs_idx)
self.input_token_ids_logprobs_val.append(
req.input_token_ids_logprobs_val
req.logprob.input_token_ids_logprobs_val
)
self.input_token_ids_logprobs_idx.append(
req.input_token_ids_logprobs_idx
req.logprob.input_token_ids_logprobs_idx
)
req.input_logprob_sent = True
else:
@@ -392,32 +396,32 @@ class _GenerationStreamAccumulator:
if req.return_logprob:
logprob_end = max(len(output_ids_), 1)
self.output_token_logprobs_val.append(
req.output_token_logprobs_val[
req.logprob.output_token_logprobs_val[
send_output_token_logprobs_offset:logprob_end
]
)
self.output_token_logprobs_idx.append(
req.output_token_logprobs_idx[
req.logprob.output_token_logprobs_idx[
send_output_token_logprobs_offset:logprob_end
]
)
self.output_top_logprobs_val.append(
req.output_top_logprobs_val[
req.logprob.output_top_logprobs_val[
send_output_token_logprobs_offset:logprob_end
]
)
self.output_top_logprobs_idx.append(
req.output_top_logprobs_idx[
req.logprob.output_top_logprobs_idx[
send_output_token_logprobs_offset:logprob_end
]
)
self.output_token_ids_logprobs_val.append(
req.output_token_ids_logprobs_val[
req.logprob.output_token_ids_logprobs_val[
send_output_token_logprobs_offset:logprob_end
]
)
self.output_token_ids_logprobs_idx.append(
req.output_token_ids_logprobs_idx[
req.logprob.output_token_ids_logprobs_idx[
send_output_token_logprobs_offset:logprob_end
]
)