diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 18a42407a..b5b303720 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -503,7 +503,9 @@ class _GenerationStreamAccumulator: self.input_token_ids_logprobs_idx.append([]) if req.return_logprob: - logprob_end = max(len(output_ids_), 1) + logprob_end = ( + len(output_ids_) if req.is_retracted else max(len(output_ids_), 1) + ) self.output_token_logprobs_val.append( req.logprob.output_token_logprobs_val[ send_output_token_logprobs_offset:logprob_end diff --git a/test/registered/scheduler/test_retract_decode_logprob.py b/test/registered/scheduler/test_retract_decode_logprob.py new file mode 100644 index 000000000..688e844e3 --- /dev/null +++ b/test/registered/scheduler/test_retract_decode_logprob.py @@ -0,0 +1,150 @@ +"""Regression test for stream_output_generation logprob off-by-one under retract. + +Bug: under overlap scheduling, if a request is retracted between the moment +its prefill batch is launched and the moment that batch's result is +processed, the in-flight batch comes back through +`process_batch_result_prefill` with the now-retracted request still in +`batch.reqs`. The `if req.is_retracted: continue` guard at the top of the +for-loop skips `req.output_ids.append(next_token_id)`, so the subsequent +`self.stream_output(batch.reqs, ...)` sees a req with `len(output_ids) == 0`. +For a non-streaming `return_logprob=True` req, +`should_output = (0 % DEFAULT_FORCE_STREAM_INTERVAL == 0) == True`, so the +slice math in `_GenerationStreamAccumulator.handle_req` (formerly +`stream_output_generation`) fires: + + output_ids_ = req.output_ids_through_stop # empty + output_ids.append(output_ids_[send_token_offset:]) # 0 tokens + req.send_token_offset = len(output_ids_) # 0 + logprob_end = max(len(output_ids_), 1) # 1 <-- BUG + output_token_logprobs_val.append(... [send_lp_off:logprob_end]) # 1 entry + req.send_output_token_logprobs_offset = logprob_end # 1 + +The two send-offsets diverge by 1. Every subsequent stream tick for this +req ships N tokens and N-1 logprobs. The final response delivered to the +client has `len(meta_info["output_token_logprobs"]) == len(output_ids) - 1`. + +Fix: drop the `max(..., 1)` floor only for retracted requests: + + logprob_end = ( + len(output_ids_) if req.is_retracted else max(len(output_ids_), 1) + ) + +The non-retracted branch preserves the first logprob for legitimate +prefill-only requests with `max_new_tokens=0`. + +This test forces the trigger reliably via `SGLANG_TEST_RETRACT=True` +(retract every two forward steps) and asserts the 1:1 invariant across +many concurrent `return_logprob=True` requests. + +Run: + python -m unittest test_retract_decode_logprob.TestRetractDecodeLogprob +""" + +import unittest +from concurrent.futures import ThreadPoolExecutor, as_completed + +import requests + +from sglang.srt.environ import envs +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=300, stage="base-b", runner_config="1-gpu-small") +register_amd_ci(est_time=360, suite="stage-b-test-1-gpu-small-amd") + +N_REQUESTS = 32 +MAX_NEW_TOKENS = 256 + + +class TestRetractDecodeLogprob(CustomTestCase): + """python -m unittest test_retract_decode_logprob.TestRetractDecodeLogprob""" + + other_args = [] + + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST + cls.base_url = DEFAULT_URL_FOR_TEST + launch_args = [ + "--chunked-prefill-size", + "128", + "--max-running-requests", + "8", + "--mem-fraction-static", + "0.7", + ] + cls.other_args + with envs.SGLANG_TEST_RETRACT.override(True): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=launch_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def _one_request(self, idx: int) -> dict: + # NOTE: no `stream` field -> non-streaming. The bug affects the + # non-streaming path too because _stream_output_generation runs for + # both (non-streaming reqs get internally force-flushed every + # DEFAULT_FORCE_STREAM_INTERVAL decoded tokens). + payload = { + "text": f"Once upon a time #{idx},", + "sampling_params": { + "max_new_tokens": MAX_NEW_TOKENS, + "temperature": 0.0, + "ignore_eos": True, + }, + "return_logprob": True, + "logprob_start_len": -1, + } + r = requests.post(f"{self.base_url}/generate", json=payload, timeout=600) + r.raise_for_status() + data = r.json() + meta = data.get("meta_info", {}) + return { + "idx": idx, + "n_tokens": len(data.get("output_ids") or []), + "n_logprobs": len(meta.get("output_token_logprobs") or []), + } + + def test_output_logprobs_aligned_under_test_retract(self): + """Every non-streaming return_logprob=True response must have + len(output_ids) == len(meta_info["output_token_logprobs"]). + + Without the fix, with SGLANG_TEST_RETRACT=True forcing retraction + every two forward steps, ~6% of responses come back with one fewer + logprob than tokens. With the fix, all responses are 1:1.""" + + with ThreadPoolExecutor(max_workers=N_REQUESTS) as pool: + futs = [pool.submit(self._one_request, i) for i in range(N_REQUESTS)] + results = [f.result() for f in as_completed(futs)] + + mismatches = [r for r in results if r["n_tokens"] != r["n_logprobs"]] + self.assertEqual( + mismatches, + [], + msg=( + f"{len(mismatches)}/{N_REQUESTS} responses have " + f"len(output_ids) != len(output_token_logprobs). " + f"Sample: {mismatches[:5]}. " + "This is the _GenerationStreamAccumulator logprob_end " + "off-by-one (max(len(output_ids_), 1) in " + "scheduler_components/output_streamer.py)." + ), + ) + + assert self.process.poll() is None, "Server crashed during test" + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_output_streamer_logprobs.py b/test/registered/unit/managers/test_output_streamer_logprobs.py new file mode 100644 index 000000000..da3bd0643 --- /dev/null +++ b/test/registered/unit/managers/test_output_streamer_logprobs.py @@ -0,0 +1,102 @@ +import unittest +from types import SimpleNamespace + +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.managers.scheduler_components.output_streamer import ( + _GenerationStreamAccumulator, +) +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class _FakeReq: + def __init__(self, *, is_retracted: bool, max_new_tokens: int): + self.rid = "req" + self.http_worker_ipc = None + self.finished_reason = None + self.finished_output = False + self.finished_len = None + self.stream = False + self.sampling_params = SimpleNamespace( + max_new_tokens=max_new_tokens, + stream_interval=None, + skip_special_tokens=True, + spaces_between_special_tokens=True, + no_stop_trim=False, + ) + self.output_ids = [] + self.output_ids_through_stop = [] + self.send_token_offset = 0 + self.send_output_token_logprobs_offset = 0 + self.send_decode_id_offset = 0 + self.decoded_text = "" + self.origin_input_ids = [] + self.reasoning_tokens = 0 + self.cached_tokens = 0 + self.retraction_count = 0 + self.time_stats = None + self.mm_image_tokens = 0 + self.mm_audio_tokens = 0 + self.mm_video_tokens = 0 + self.multimodal_inputs = None + self.customized_info = None + self.is_retracted = is_retracted + + self.return_logprob = True + self.input_logprob_sent = True + self.logprob = SimpleNamespace( + output_token_logprobs_val=[-0.5], + output_token_logprobs_idx=[42], + output_top_logprobs_val=[[(-0.5, 42)]], + output_top_logprobs_idx=[[42]], + output_token_ids_logprobs_val=[[-0.5]], + output_token_ids_logprobs_idx=[[42]], + ) + + def finished(self): + return False + + def init_incremental_detokenize(self): + return self.output_ids_through_stop, 0 + + +def _make_accumulator() -> _GenerationStreamAccumulator: + return _GenerationStreamAccumulator( + return_logprob=True, + return_hidden_states=False, + return_routed_experts=False, + return_indexer_topk=False, + spec_algorithm=SpeculativeAlgorithm.NONE, + disaggregation_mode=DisaggregationMode.NULL, + default_stream_interval=1, + default_force_stream_interval=1, + get_cached_tokens_details=lambda req: None, + ) + + +class TestOutputStreamerLogprobs(unittest.TestCase): + def test_retracted_empty_output_does_not_advance_logprob_offset(self): + req = _FakeReq(is_retracted=True, max_new_tokens=16) + accumulator = _make_accumulator() + + accumulator.accept(req=req) + + self.assertEqual(req.send_token_offset, 0) + self.assertEqual(req.send_output_token_logprobs_offset, 0) + self.assertEqual(accumulator.output_token_logprobs_val, [[]]) + + def test_prefill_only_request_preserves_first_logprob(self): + req = _FakeReq(is_retracted=False, max_new_tokens=0) + accumulator = _make_accumulator() + + accumulator.accept(req=req) + + self.assertEqual(req.send_token_offset, 0) + self.assertEqual(req.send_output_token_logprobs_offset, 1) + self.assertEqual(accumulator.output_token_logprobs_val, [[-0.5]]) + + +if __name__ == "__main__": + unittest.main()