From 2d40f45193b297824f8f245e3e7d6dfa5400fde5 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Tue, 19 May 2026 09:18:19 +0800 Subject: [PATCH] Pack scattered request logprob state into a dedicated container (#25712) --- python/sglang/srt/disaggregation/decode.py | 16 ++- .../decode_schedule_batch_mixin.py | 4 +- python/sglang/srt/disaggregation/utils.py | 36 ++--- python/sglang/srt/layers/utils/logprob.py | 27 ++-- python/sglang/srt/managers/schedule_batch.py | 57 ++++---- python/sglang/srt/managers/scheduler.py | 2 +- .../batch_result_processor.py | 16 +-- .../logprob_result_processor.py | 126 ++++++++++-------- .../scheduler_components/output_streamer.py | 30 +++-- 9 files changed, 182 insertions(+), 132 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 5be837eb3..5d79d1e90 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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() diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 3c4a7e3aa..43e58bfdf 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -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] diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 951fa5b7d..ad6904a0a 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -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 = [] diff --git a/python/sglang/srt/layers/utils/logprob.py b/python/sglang/srt/layers/utils/logprob.py index 3b14b2510..0fdd1b170 100644 --- a/python/sglang/srt/layers/utils/logprob.py +++ b/python/sglang/srt/layers/utils/logprob.py @@ -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 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 33da2b1e3..109fa0e64 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 67ac18d8c..72b574cfb 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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( diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 6a96fb4ef..d0e316ccc 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -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] ) diff --git a/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py b/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py index 6c814e4b6..f58429f0a 100644 --- a/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py @@ -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 = [] diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 6941fc357..90fdebb1f 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -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 ] )