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 decode_req.req.hidden_states_tensor = output_hidden_states
if decode_req.req.return_logprob: 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() 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() output_token_logprobs_idx[0].item()
) )
decode_req.req.output_top_logprobs_val.append( decode_req.req.logprob.output_top_logprobs_val.append(
output_top_logprobs_val[: decode_req.req.top_logprobs_num].tolist() output_top_logprobs_val[
: decode_req.req.logprob.top_logprobs_num
].tolist()
) )
decode_req.req.output_top_logprobs_idx.append( decode_req.req.logprob.output_top_logprobs_idx.append(
output_top_logprobs_idx[: decode_req.req.top_logprobs_num].tolist() output_top_logprobs_idx[
: decode_req.req.logprob.top_logprobs_num
].tolist()
) )
decode_req.kv_receiver.clear() 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) self.seq_lens_sum = sum(seq_lens)
if self.return_logprob: if self.return_logprob:
self.top_logprobs_nums = [r.top_logprobs_num for r in reqs] self.top_logprobs_nums = [r.logprob.top_logprobs_num for r in reqs]
self.token_ids_logprobs = [r.token_ids_logprob 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.extend_num_tokens = extend_num_tokens
self.prefix_lens = [len(r.prefix_indices) for r in reqs] 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][2] = req.cached_tokens_host
self.cached_tokens[req.metadata_buffer_index][3] = req.cached_tokens_storage self.cached_tokens[req.metadata_buffer_index][3] = req.cached_tokens_storage
if req.return_logprob: 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] = ( 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] = ( 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][ 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( ] = 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][ 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( ] = 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 # For PD + spec decode
if req.hidden_states_tensor is not None: 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) req.finished_reason = FINISH_ABORT(error_message, status_code)
if req.return_logprob: if req.return_logprob:
req.input_token_logprobs_val = [] req.logprob.input_token_logprobs_val = []
req.input_token_logprobs_idx = [] req.logprob.input_token_logprobs_idx = []
req.input_top_logprobs_val = [] req.logprob.input_top_logprobs_val = []
req.input_top_logprobs_idx = [] req.logprob.input_top_logprobs_idx = []
req.input_token_ids_logprobs_val = [] req.logprob.input_token_ids_logprobs_val = []
req.input_token_ids_logprobs_idx = [] 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 req, num_tokens in zip(batch.reqs, num_tokens_per_req, strict=True):
for _ in range(num_tokens): for _ in range(num_tokens):
if req.return_logprob: if req.return_logprob:
req.output_token_logprobs_val.append(next_token_logprobs[pt]) req.logprob.output_token_logprobs_val.append(next_token_logprobs[pt])
req.output_token_logprobs_idx.append(accept_tokens_list[pt]) req.logprob.output_token_logprobs_idx.append(accept_tokens_list[pt])
if req.top_logprobs_num > 0: if req.logprob.top_logprobs_num > 0:
assert ( assert (
should_top_logprobs should_top_logprobs
), "Inconsistent state: should_top_logprobs is False" ), "Inconsistent state: should_top_logprobs is False"
req.output_top_logprobs_val.append(token_top_logprobs_val[pt]) req.logprob.output_top_logprobs_val.append(
req.output_top_logprobs_idx.append(token_top_logprobs_idx[pt]) token_top_logprobs_val[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.logprob.output_top_logprobs_idx.append(
req.output_token_ids_logprobs_idx.append(token_ids_logprobs_idx[pt]) 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 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 # 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): class Req(ReqDllmMixin):
"""The input and output status of a request.""" """The input and output status of a request."""
@@ -767,18 +788,14 @@ class Req(ReqDllmMixin):
self.return_logprob = return_logprob self.return_logprob = return_logprob
# Start index to compute logprob from. # Start index to compute logprob from.
self.logprob_start_len = 0 self.logprob_start_len = 0
self.top_logprobs_num = top_logprobs_num self.logprob = ReqLogprob(
self.token_ids_logprob = token_ids_logprob top_logprobs_num=top_logprobs_num,
token_ids_logprob=token_ids_logprob,
)
# Logprobs (return values) # Logprobs (return values)
# True means the input logprob has been already sent to detokenizer. # True means the input logprob has been already sent to detokenizer.
self.input_logprob_sent: bool = False 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. # Temporary holder to store input_token_logprobs.
self.input_token_logprobs: Optional[List[Tuple[int]]] = None self.input_token_logprobs: Optional[List[Tuple[int]]] = None
self.temp_input_top_logprobs_val: Optional[List[torch.Tensor]] = None self.temp_input_top_logprobs_val: Optional[List[torch.Tensor]] = None
@@ -788,22 +805,14 @@ class Req(ReqDllmMixin):
if return_logprob: if return_logprob:
# shape: (bs, 1) # shape: (bs, 1)
self.output_token_logprobs_val = [] self.logprob.output_token_logprobs_val = []
self.output_token_logprobs_idx = [] self.logprob.output_token_logprobs_idx = []
# shape: (bs, k) # shape: (bs, k)
self.output_top_logprobs_val = [] self.logprob.output_top_logprobs_val = []
self.output_top_logprobs_idx = [] self.logprob.output_top_logprobs_idx = []
# Can contain either lists or GPU tensors (delayed copy optimization for prefill-only scoring) # Can contain either lists or GPU tensors (delayed copy optimization for prefill-only scoring)
self.output_token_ids_logprobs_val: List[ self.logprob.output_token_ids_logprobs_val = []
Union[List[float], torch.Tensor] self.logprob.output_token_ids_logprobs_idx = []
] = []
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.hidden_states: List[List[float]] = [] 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.hidden_states_tensor = None # Note: use tensor instead of list to transfer hidden_states when PD + MTP
self.output_topk_p = None self.output_topk_p = None
@@ -1954,8 +1963,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.multi_item_delimiter_indices = None self.multi_item_delimiter_indices = None
if self.return_logprob: if self.return_logprob:
self.top_logprobs_nums = [r.top_logprobs_num for r in reqs] self.top_logprobs_nums = [r.logprob.top_logprobs_num for r in reqs]
self.token_ids_logprobs = [r.token_ids_logprob 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_logprob_start_lens = [r.extend_logprob_start_len for r in reqs]
self.extend_input_logprob_token_ids = extend_input_logprob_token_ids 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: 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.token_ids_logprobs = [list(r.origin_input_ids) for r in reqs]
batch.sampling_info = SamplingBatchInfo.from_schedule_batch( 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 max_accept = 1
for j, tok_id in enumerate(accepted_ids): for j, tok_id in enumerate(accepted_ids):
req.output_token_logprobs_val.append(accepted_logprobs[j]) req.logprob.output_token_logprobs_val.append(accepted_logprobs[j])
req.output_token_logprobs_idx.append(tok_id) req.logprob.output_token_logprobs_idx.append(tok_id)
if req.top_logprobs_num > 0: if req.logprob.top_logprobs_num > 0:
flat_idx = i * max_accept + j 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] 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] 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 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] 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] 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 # Process logprob values - handle multi-item scoring vs regular requests
if is_multi_item_scoring: if is_multi_item_scoring:
# Multi-item scoring: use all logprobs as-is # 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: else:
# Regular request: add None at start, remove last (sampling token) # 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 # Process logprob indices based on scoring type
if is_multi_item_scoring: if is_multi_item_scoring:
@@ -49,21 +49,21 @@ class SchedulerLogprobResultProcessor:
input_token_logprobs_idx = req.origin_input_ids[req.logprob_start_len :] input_token_logprobs_idx = req.origin_input_ids[req.logprob_start_len :]
# Clip padded hash values from image tokens to prevent detokenization errors # 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 x if x < self.model_config.vocab_size - 1 else 0
for x in input_token_logprobs_idx for x in input_token_logprobs_idx
] ]
def _process_input_top_logprobs(self, req: Req) -> None: def _process_input_top_logprobs(self, req: Req) -> None:
"""Process input top logprobs.""" """Process input top logprobs."""
if req.top_logprobs_num <= 0: if req.logprob.top_logprobs_num <= 0:
return return
is_multi_item_scoring = self._is_multi_item_scoring(req) is_multi_item_scoring = self._is_multi_item_scoring(req)
# Initialize arrays - multi-item scoring starts empty, others start with None # Initialize arrays - multi-item scoring starts empty, others start with None
req.input_top_logprobs_val = [] if is_multi_item_scoring else [None] req.logprob.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_idx = [] if is_multi_item_scoring else [None]
# Extend arrays with temp values # Extend arrays with temp values
for val, idx in zip( for val, idx in zip(
@@ -71,13 +71,13 @@ class SchedulerLogprobResultProcessor:
req.temp_input_top_logprobs_idx, req.temp_input_top_logprobs_idx,
strict=True, strict=True,
): ):
req.input_top_logprobs_val.extend(val) req.logprob.input_top_logprobs_val.extend(val)
req.input_top_logprobs_idx.extend(idx) req.logprob.input_top_logprobs_idx.extend(idx)
# Remove last token (sampling token) for non multi-item scoring requests # Remove last token (sampling token) for non multi-item scoring requests
if not is_multi_item_scoring: if not is_multi_item_scoring:
req.input_top_logprobs_val.pop() req.logprob.input_top_logprobs_val.pop()
req.input_top_logprobs_idx.pop() req.logprob.input_top_logprobs_idx.pop()
# Clean up temp storage # Clean up temp storage
req.temp_input_top_logprobs_idx = None req.temp_input_top_logprobs_idx = None
@@ -85,14 +85,18 @@ class SchedulerLogprobResultProcessor:
def _process_input_token_ids_logprobs(self, req: Req) -> None: def _process_input_token_ids_logprobs(self, req: Req) -> None:
"""Process input token IDs logprobs.""" """Process input token IDs logprobs."""
if req.token_ids_logprob is None: if req.logprob.token_ids_logprob is None:
return return
is_multi_item_scoring = self._is_multi_item_scoring(req) is_multi_item_scoring = self._is_multi_item_scoring(req)
# Initialize arrays - multi-item scoring starts empty, others start with None # 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.logprob.input_token_ids_logprobs_val = (
req.input_token_ids_logprobs_idx = [] if is_multi_item_scoring else [None] [] 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 # Process temp values - convert tensors to lists and extend arrays
for val, idx in zip( for val, idx in zip(
@@ -101,15 +105,15 @@ class SchedulerLogprobResultProcessor:
strict=True, strict=True,
): ):
val_list = val.tolist() if isinstance(val, torch.Tensor) else val 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] 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 # Remove last token (sampling token) for non multi-item scoring requests
if not is_multi_item_scoring: if not is_multi_item_scoring:
req.input_token_ids_logprobs_val.pop() req.logprob.input_token_ids_logprobs_val.pop()
req.input_token_ids_logprobs_idx.pop() req.logprob.input_token_ids_logprobs_idx.pop()
# Clean up temp storage # Clean up temp storage
req.temp_input_token_ids_logprobs_idx = None req.temp_input_token_ids_logprobs_idx = None
@@ -198,11 +202,11 @@ class SchedulerLogprobResultProcessor:
if req.temp_input_token_ids_logprobs_idx is None: if req.temp_input_token_ids_logprobs_idx is None:
req.temp_input_token_ids_logprobs_idx = [] 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 # The input logprob has been already computed. It only happens
# upon retract. # upon retract.
if req.top_logprobs_num > 0: if req.logprob.top_logprobs_num > 0:
assert req.input_token_logprobs_val is not None assert req.logprob.input_token_logprobs_val is not None
return return
# Important for the performance. # Important for the performance.
@@ -213,11 +217,11 @@ class SchedulerLogprobResultProcessor:
] ]
req.input_token_logprobs.extend(input_token_logprobs) 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_val.append(output.input_top_logprobs_val[i])
req.temp_input_top_logprobs_idx.append(output.input_top_logprobs_idx[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( req.temp_input_token_ids_logprobs_val.append(
output.input_token_ids_logprobs_val[i] output.input_token_ids_logprobs_val[i]
) )
@@ -228,10 +232,10 @@ class SchedulerLogprobResultProcessor:
if last_prefill_chunk: if last_prefill_chunk:
input_token_logprobs = req.input_token_logprobs input_token_logprobs = req.input_token_logprobs
req.input_token_logprobs = None req.input_token_logprobs = None
assert req.input_token_logprobs_val is None assert req.logprob.input_token_logprobs_val is None
assert req.input_token_logprobs_idx is None assert req.logprob.input_token_logprobs_idx is None
assert req.input_top_logprobs_val is None assert req.logprob.input_top_logprobs_val is None
assert req.input_top_logprobs_idx is None assert req.logprob.input_top_logprobs_idx is None
# Process all input logprob types using helper functions # Process all input logprob types using helper functions
self._process_input_token_logprobs(req, input_token_logprobs) self._process_input_token_logprobs(req, input_token_logprobs)
@@ -241,14 +245,24 @@ class SchedulerLogprobResultProcessor:
if req.return_logprob: if req.return_logprob:
relevant_tokens_len = self._calculate_relevant_tokens_len(req) relevant_tokens_len = self._calculate_relevant_tokens_len(req)
assert len(req.input_token_logprobs_val) == relevant_tokens_len assert len(req.logprob.input_token_logprobs_val) == relevant_tokens_len
assert len(req.input_token_logprobs_idx) == relevant_tokens_len assert len(req.logprob.input_token_logprobs_idx) == relevant_tokens_len
if req.top_logprobs_num > 0: if req.logprob.top_logprobs_num > 0:
assert len(req.input_top_logprobs_val) == relevant_tokens_len assert (
assert len(req.input_top_logprobs_idx) == relevant_tokens_len len(req.logprob.input_top_logprobs_val) == relevant_tokens_len
if req.token_ids_logprob is not None: )
assert len(req.input_token_ids_logprobs_val) == relevant_tokens_len assert (
assert len(req.input_token_ids_logprobs_idx) == relevant_tokens_len 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( def add_logprob_return_values(
self, self,
@@ -261,8 +275,8 @@ class SchedulerLogprobResultProcessor:
): ):
"""Attach logprobs to the return values.""" """Attach logprobs to the return values."""
if output.next_token_logprobs is not None: if output.next_token_logprobs is not None:
req.output_token_logprobs_val.append(output.next_token_logprobs[i]) req.logprob.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_idx.append(next_token_ids[i])
# Only add input logprobs if there are input tokens to process # 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, # Note: For prefill-only requests with default logprob_start_len, this will be 0,
@@ -279,20 +293,24 @@ class SchedulerLogprobResultProcessor:
else: else:
self._initialize_empty_logprob_containers(req) self._initialize_empty_logprob_containers(req)
if req.top_logprobs_num > 0: if req.logprob.top_logprobs_num > 0:
req.output_top_logprobs_val.append(output.next_token_top_logprobs_val[i]) req.logprob.output_top_logprobs_val.append(
req.output_top_logprobs_idx.append(output.next_token_top_logprobs_idx[i]) output.next_token_top_logprobs_val[i]
)
req.logprob.output_top_logprobs_idx.append(
output.next_token_top_logprobs_idx[i]
)
if ( 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 and output.next_token_token_ids_logprobs_val is not None
): ):
# Convert GPU tensor to list if needed # Convert GPU tensor to list if needed
logprobs_val = output.next_token_token_ids_logprobs_val[i] logprobs_val = output.next_token_token_ids_logprobs_val[i]
if isinstance(logprobs_val, torch.Tensor): if isinstance(logprobs_val, torch.Tensor):
logprobs_val = logprobs_val.tolist() logprobs_val = logprobs_val.tolist()
req.output_token_ids_logprobs_val.append(logprobs_val) req.logprob.output_token_ids_logprobs_val.append(logprobs_val)
req.output_token_ids_logprobs_idx.append( req.logprob.output_token_ids_logprobs_idx.append(
output.next_token_token_ids_logprobs_idx[i] 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 This is needed for prefill-only requests where the normal initialization
flow might be bypassed, but downstream code expects these fields to be lists. flow might be bypassed, but downstream code expects these fields to be lists.
""" """
if req.input_token_logprobs_val is None: if req.logprob.input_token_logprobs_val is None:
req.input_token_logprobs_val = [] req.logprob.input_token_logprobs_val = []
if req.input_token_logprobs_idx is None: if req.logprob.input_token_logprobs_idx is None:
req.input_token_logprobs_idx = [] req.logprob.input_token_logprobs_idx = []
if req.input_top_logprobs_val is None: if req.logprob.input_top_logprobs_val is None:
req.input_top_logprobs_val = [] req.logprob.input_top_logprobs_val = []
if req.input_top_logprobs_idx is None: if req.logprob.input_top_logprobs_idx is None:
req.input_top_logprobs_idx = [] req.logprob.input_top_logprobs_idx = []
if req.input_token_ids_logprobs_val is None: if req.logprob.input_token_ids_logprobs_val is None:
req.input_token_ids_logprobs_val = [] req.logprob.input_token_ids_logprobs_val = []
if req.input_token_ids_logprobs_idx is None: if req.logprob.input_token_ids_logprobs_idx is None:
req.input_token_ids_logprobs_idx = [] 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 # Decode server does not send input logprobs
and self.disaggregation_mode != DisaggregationMode.DECODE and self.disaggregation_mode != DisaggregationMode.DECODE
# Only send when input logprobs have been computed (after prefill) # 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_val.append(
self.input_token_logprobs_idx.append(req.input_token_logprobs_idx) req.logprob.input_token_logprobs_val
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_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( 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( 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 req.input_logprob_sent = True
else: else:
@@ -392,32 +396,32 @@ class _GenerationStreamAccumulator:
if req.return_logprob: if req.return_logprob:
logprob_end = max(len(output_ids_), 1) logprob_end = max(len(output_ids_), 1)
self.output_token_logprobs_val.append( self.output_token_logprobs_val.append(
req.output_token_logprobs_val[ req.logprob.output_token_logprobs_val[
send_output_token_logprobs_offset:logprob_end send_output_token_logprobs_offset:logprob_end
] ]
) )
self.output_token_logprobs_idx.append( self.output_token_logprobs_idx.append(
req.output_token_logprobs_idx[ req.logprob.output_token_logprobs_idx[
send_output_token_logprobs_offset:logprob_end send_output_token_logprobs_offset:logprob_end
] ]
) )
self.output_top_logprobs_val.append( self.output_top_logprobs_val.append(
req.output_top_logprobs_val[ req.logprob.output_top_logprobs_val[
send_output_token_logprobs_offset:logprob_end send_output_token_logprobs_offset:logprob_end
] ]
) )
self.output_top_logprobs_idx.append( self.output_top_logprobs_idx.append(
req.output_top_logprobs_idx[ req.logprob.output_top_logprobs_idx[
send_output_token_logprobs_offset:logprob_end send_output_token_logprobs_offset:logprob_end
] ]
) )
self.output_token_ids_logprobs_val.append( 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 send_output_token_logprobs_offset:logprob_end
] ]
) )
self.output_token_ids_logprobs_idx.append( 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 send_output_token_logprobs_offset:logprob_end
] ]
) )