diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index d8d2023fbbf5..368e25bfa9e1 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -2929,13 +2929,38 @@ def _commit_transfer_to_req(self, decode_req: DecodeRequest): ) else: committed_output_id = output_id[0].item() - decode_req.req.output_ids.append(committed_output_id) - # The handoff token is generated on the prefill worker, so it does - # not pass through the decode worker's normal batch-result path. - # Account for it here using the same request-selected reasoning - # terminator matcher as subsequent decode tokens. - self.scheduler.batch_result_processor._maybe_update_reasoning_tokens( - decode_req.req, committed_output_id + req = decode_req.req + logprob_token_id = logprob = top_logprobs_val = top_logprobs_idx = None + if req.return_logprob: + top_logprobs_num = req.logprob.top_logprobs_num + logprob_token_id = output_token_logprobs_idx[0].item() + logprob = output_token_logprobs_val[0].item() + top_logprobs_val = output_top_logprobs_val[:top_logprobs_num].tolist() + top_logprobs_idx = output_top_logprobs_idx[:top_logprobs_num].tolist() + sampling_mask_idx = sampling_mask_logprobs = None + if req.return_sampling_mask: + assert output_token_sampling_mask_idx is not None, ( + "sampling mask buffer disabled on decode side" + ) + sampling_mask_len = int(output_token_sampling_mask_len[0].item()) + num_logprobs = ( + sampling_mask_len if req.sampling_logprobs_mode == "support" else 1 + ) + sampling_mask_idx = ( + output_token_sampling_mask_idx[:sampling_mask_len].cpu().numpy() + ) + sampling_mask_logprobs = ( + output_token_sampling_logprobs[:num_logprobs].cpu().numpy() + ) + self.scheduler.batch_result_processor.commit_handoff_token( + req, + committed_output_id, + logprob_token_id=logprob_token_id, + logprob=logprob, + top_logprobs_val=top_logprobs_val, + top_logprobs_idx=top_logprobs_idx, + sampling_mask_idx=sampling_mask_idx, + sampling_mask_logprobs=sampling_mask_logprobs, ) decode_req.req.cached_tokens = cached_tokens[0].item() # The prefill node already reported its prefix-cache hit in @@ -2965,42 +2990,6 @@ def _commit_transfer_to_req(self, decode_req: DecodeRequest): output_dsa_topk_indices = None decode_req.req.output_dsa_topk_indices = output_dsa_topk_indices - if decode_req.req.return_logprob and not resumes_output: - decode_req.req.logprob.output_token_logprobs_val.append( - output_token_logprobs_val[0].item() - ) - decode_req.req.logprob.output_token_logprobs_idx.append( - output_token_logprobs_idx[0].item() - ) - decode_req.req.logprob.output_top_logprobs_val.append( - output_top_logprobs_val[ - : decode_req.req.logprob.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() - ) - if decode_req.req.logprob.token_ids_logprob is not None: - # Prefill omits these scores; reserve a row for the streaming offset. - decode_req.req.logprob.output_token_ids_logprobs_val.append([]) - decode_req.req.logprob.output_token_ids_logprobs_idx.append([]) - if decode_req.req.return_sampling_mask and not resumes_output: - assert output_token_sampling_mask_idx is not None, ( - "sampling mask buffer disabled on decode side" - ) - sampling_mask_len = int(output_token_sampling_mask_len[0].item()) - num_logprobs = ( - sampling_mask_len - if decode_req.req.sampling_logprobs_mode == "support" - else 1 - ) - decode_req.req.sampling_mask_rows.append( - output_token_sampling_mask_idx[:sampling_mask_len].cpu().numpy(), - output_token_sampling_logprobs[:num_logprobs].cpu().numpy(), - ) - decode_req.kv_receiver.clear() decode_req.kv_receiver = None decode_req.req.time_stats.set_wait_queue_entry_time() 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 90708e4461f0..a37975ad4097 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -1261,24 +1261,16 @@ def _apply_decode_logprobs( max_accept = 1 for j, tok_id in enumerate(accepted_ids): - 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.logprob.output_top_logprobs_val.append( - logits_output.next_token_top_logprobs_val[flat_idx] - ) - req.logprob.output_top_logprobs_idx.append( - logits_output.next_token_top_logprobs_idx[flat_idx] - ) - if req.logprob.token_ids_logprob is not None: - flat_idx = i * max_accept + j - req.logprob.output_token_ids_logprobs_val.append( - logits_output.next_token_token_ids_logprobs_val[flat_idx] - ) - req.logprob.output_token_ids_logprobs_idx.append( - logits_output.next_token_token_ids_logprobs_idx[flat_idx] - ) + self.logprob_result_processor.append_output_token_logprobs( + req, + tok_id, + accepted_logprobs[j], + row=i * max_accept + j, + top_logprobs_val=logits_output.next_token_top_logprobs_val, + top_logprobs_idx=logits_output.next_token_top_logprobs_idx, + token_ids_logprobs_val=logits_output.next_token_token_ids_logprobs_val, + token_ids_logprobs_idx=logits_output.next_token_token_ids_logprobs_idx, + ) def add_sampling_mask_return_values( self, @@ -1587,6 +1579,35 @@ def _mamba_can_cache_finished_req( ) return len(retained[i]) == num_verified + def commit_handoff_token( + self, + req: Req, + token_id: int, + *, + logprob_token_id: Optional[int] = None, + logprob: Optional[float] = None, + top_logprobs_val: Optional[List[float]] = None, + top_logprobs_idx: Optional[List[int]] = None, + sampling_mask_idx: Optional[np.ndarray] = None, + sampling_mask_logprobs: Optional[np.ndarray] = None, + ) -> None: + """Commit the PD prefill worker's sampled token to a decode request.""" + req.output_ids.append(token_id) + self._maybe_update_reasoning_tokens(req, token_id) + if req.return_logprob: + # Prefill omits token-ids scores; reserve a row for the streaming offset. + self.logprob_result_processor.append_output_token_logprobs( + req, + logprob_token_id, + logprob, + top_logprobs_val=[top_logprobs_val], + top_logprobs_idx=[top_logprobs_idx], + token_ids_logprobs_val=[[]], + token_ids_logprobs_idx=[[]], + ) + if req.return_sampling_mask: + req.sampling_mask_rows.append(sampling_mask_idx, sampling_mask_logprobs) + def _maybe_update_reasoning_tokens( self, req: Req, 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 86e1a7bfe70f..dc36dfd96d20 100644 --- a/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import ( List, + Optional, Tuple, ) @@ -326,9 +327,20 @@ def add_logprob_return_values( output: LogitsProcessorOutput, ): """Attach logprobs to the return values.""" - if output.next_token_logprobs is not None: - req.logprob.output_token_logprobs_val.append(output.next_token_logprobs[i]) - req.logprob.output_token_logprobs_idx.append(next_token_ids[i]) + self.append_output_token_logprobs( + req, + next_token_ids[i], + ( + output.next_token_logprobs[i] + if output.next_token_logprobs is not None + else None + ), + row=i, + top_logprobs_val=output.next_token_top_logprobs_val, + top_logprobs_idx=output.next_token_top_logprobs_idx, + token_ids_logprobs_val=output.next_token_token_ids_logprobs_val, + token_ids_logprobs_idx=output.next_token_token_ids_logprobs_idx, + ) # 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, @@ -345,35 +357,40 @@ def add_logprob_return_values( else: self._initialize_empty_logprob_containers(req) - if ( - req.logprob.top_logprobs_num > 0 - and output.next_token_top_logprobs_val is not None - ): - # Guarded like next_token_logprobs above: a backend may leave the - # top-logprob fields unset even for a request that asked for them - # (indexing None raises TypeError). - 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] - ) + return num_input_logprobs + @staticmethod + def append_output_token_logprobs( + req: Req, + token_id: int, + logprob: Optional[float], + *, + row: int = 0, + top_logprobs_val: Optional[list] = None, + top_logprobs_idx: Optional[list] = None, + token_ids_logprobs_val: Optional[list] = None, + token_ids_logprobs_idx: Optional[list] = None, + ) -> None: + """Append one output token's logprobs; list arguments are read at `row`, + and None (a backend may leave a field unset) appends nothing.""" + if logprob is not None: + req.logprob.output_token_logprobs_val.append(logprob) + req.logprob.output_token_logprobs_idx.append(token_id) + if req.logprob.top_logprobs_num > 0 and top_logprobs_val is not None: + req.logprob.output_top_logprobs_val.append(top_logprobs_val[row]) + req.logprob.output_top_logprobs_idx.append(top_logprobs_idx[row]) if ( req.logprob.token_ids_logprob is not None - and output.next_token_token_ids_logprobs_val is not None + and token_ids_logprobs_val is not None ): - # Convert GPU tensor to list if needed - logprobs_val = output.next_token_token_ids_logprobs_val[i] + logprobs_val = token_ids_logprobs_val[row] if isinstance(logprobs_val, torch.Tensor): logprobs_val = logprobs_val.tolist() 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] + token_ids_logprobs_idx[row] ) - return num_input_logprobs - def _initialize_empty_logprob_containers(self, req: Req) -> None: """ Initialize logprob fields to empty lists if unset. diff --git a/test/registered/core/test_srt_endpoint.py b/test/registered/core/test_srt_endpoint.py index 072ead572630..b19ee90a9512 100644 --- a/test/registered/core/test_srt_endpoint.py +++ b/test/registered/core/test_srt_endpoint.py @@ -887,7 +887,7 @@ def test_detokenize_invalid_tokens(self): r2 = self.session.post( self.detokenize_url, json={"model": self.model, "tokens": [1, -1, 2]} ) - self.assertEqual(r2.status_code, 500) + self.assertEqual(r2.status_code, 400) # --------------------------------------------------------------------------- diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index ed737f6ce712..d39322b20da9 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -1,4 +1,5 @@ import concurrent.futures +import dataclasses import threading import unittest from collections import defaultdict @@ -49,6 +50,9 @@ from sglang.srt.layers.attention.dsa.utils import should_use_dsa_fused_topk from sglang.srt.managers.overlap_utils import FutureMap, RelayPayload from sglang.srt.managers.schedule_batch import ReqKvInfo +from sglang.srt.managers.scheduler_components.batch_result_processor import ( + SchedulerBatchResultProcessor, +) from sglang.srt.mem_cache.deepseek_v4_compress_state import ( CompressStatePool, c4_state_transfer_indices, @@ -1845,12 +1849,16 @@ def test_sampling_mask_row_reaches_decode_unchanged(self): sampling_logprobs_mode=mode, ) ) + fields = dict.fromkeys( + f.name for f in dataclasses.fields(SchedulerBatchResultProcessor) + ) + fields["model_config"] = SimpleNamespace( + self_opened_think_start_ids=None + ) queue = DecodeTransferQueue.__new__(DecodeTransferQueue) queue.scheduler = SimpleNamespace( kv_checksum_computer=None, - batch_result_processor=SimpleNamespace( - _maybe_update_reasoning_tokens=lambda req, token_id: None - ), + batch_result_processor=SchedulerBatchResultProcessor(**fields), ) queue.spec_algorithm = SimpleNamespace(is_none=lambda: True) queue.metadata_buffers = buffers @@ -1859,6 +1867,7 @@ def test_sampling_mask_row_reaches_decode_unchanged(self): bootstrap_host="127.0.0.1", bootstrap_room=9, output_ids=[], + require_reasoning=False, return_logprob=False, return_sampling_mask=True, sampling_logprobs_mode=mode,