Fix _GenerationStreamAccumulator logprob_end off-by-one under retract (#26510)

Co-authored-by: Qiaolin Yu <liin1211@outlook.com>

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

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()

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()