diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index fa359c548..3ce179755 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -208,6 +208,9 @@ class GenerateReqInput: return_sampling_mask: Optional[Union[List[bool], bool]] = None # Whether to detokenize tokens in text in the returned logprobs. return_text_in_logprobs: bool = False + # Return prompt top logprobs as flat arrays plus shape metadata instead of + # the nested per-position [logprob, token_id, text] lists. + return_flat_raw_top_logprobs: bool = False # Whether to stream output. stream: bool = False # Whether to log metrics for this request (e.g. health_generate calls do not log metrics) @@ -369,6 +372,15 @@ class GenerateReqInput: raise ValueError( "Either text, input_ids or input_embeds should be provided." ) + if ( + self.return_flat_raw_top_logprobs + and self.multi_item_delimiter_indices is not None + ): + raise ValueError( + "return_flat_raw_top_logprobs does not support multi-item " + "scoring: delimiter-sparse top logprob rows have no contiguous " + "position mapping." + ) def _determine_batch_size(self): """Determine if this is a single example or a batch and the batch size.""" @@ -724,6 +736,7 @@ class GenerateReqInput: token_ids_logprob=self.token_ids_logprob[i], return_sampling_mask=self.return_sampling_mask[i], return_text_in_logprobs=self.return_text_in_logprobs, + return_flat_raw_top_logprobs=self.return_flat_raw_top_logprobs, stream=self.stream, log_metrics=self.log_metrics, return_hidden_states=( diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 7fe6e9491..ddec08f41 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -226,6 +226,11 @@ class ReqState: output_token_sampling_mask: List = dataclasses.field(default_factory=list) output_token_sampling_logprobs: List = dataclasses.field(default_factory=list) + # Cached flat-format prompt top logprob fields; rebuilt only when more + # prefill chunks arrive, so streaming decode chunks reuse the payload. + input_top_logprobs_flat_fields: Optional[Dict[str, Any]] = None + input_top_logprobs_flat_num_rows: int = -1 + # For detokenized logprobs input_token_logprobs: List[Any] = dataclasses.field(default_factory=list) output_token_logprobs: List[Any] = dataclasses.field(default_factory=list) @@ -254,6 +259,43 @@ def _slice_streaming_output_meta_info( meta_info[key] = meta_info[key][last_output_offset:] +def _build_flat_input_top_logprobs_fields( + input_top_logprobs_val: List[Optional[List[float]]], + input_top_logprobs_idx: List[Optional[List[int]]], + top_logprobs_num: int, +) -> Dict[str, Any]: + """Build the flat raw prompt top logprob response fields. + + `input_top_logprobs_shape` is the literal [rows, k] shape of the flat + arrays. The leading null positions (counted by + `input_top_logprobs_null_prefix`) precede the arrays, so covered position + i, entry j lives at flat[(i - null_prefix) * k + j] and the covered range + spans null_prefix + rows positions. + """ + num_rows = len(input_top_logprobs_val) + null_prefix = 0 + while null_prefix < num_rows and not input_top_logprobs_val[null_prefix]: + null_prefix += 1 + val_rows = input_top_logprobs_val[null_prefix:] + idx_rows = input_top_logprobs_idx[null_prefix:] + k = len(val_rows[0]) if val_rows else top_logprobs_num + for offset, row in enumerate(val_rows): + if row is None or len(row) != k: + # Not representable by (shape, null_prefix); e.g. multi-item scoring. + raise ValueError( + "return_flat_raw_top_logprobs requires rectangular top logprob " + f"rows with nulls only in the leading prefix; row {null_prefix + offset} " + f"has {None if row is None else len(row)} entries (expected {k})." + ) + + fields: Dict[str, Any] = {} + fields["input_top_logprobs_val_flat"] = [v for row in val_rows for v in row] + fields["input_top_logprobs_idx_flat"] = [i for row in idx_rows for i in row] + fields["input_top_logprobs_shape"] = [len(val_rows), k] + fields["input_top_logprobs_null_prefix"] = null_prefix + return fields + + class InputFormat(Enum): """Input format types for tokenization handling.""" @@ -2228,14 +2270,53 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): # 2. Handle top logprobs if top_logprobs_num > 0: - if len(state.input_top_logprobs_val) > len(state.input_top_logprobs): - state.input_top_logprobs.extend( - self.detokenize_top_logprobs_tokens( - state.input_top_logprobs_val[len(state.input_top_logprobs) :], - state.input_top_logprobs_idx[len(state.input_top_logprobs) :], - return_text_in_logprobs, + # Guarded by the caller's return_logprob check, so obj is a + # GenerateReqInput here. + use_flat = state.obj.return_flat_raw_top_logprobs + if use_flat: + # Flat replaces nested for the input side only. + if state.input_top_logprobs_flat_num_rows != len( + state.input_top_logprobs_val + ): + try: + state.input_top_logprobs_flat_fields = ( + _build_flat_input_top_logprobs_fields( + state.input_top_logprobs_val, + state.input_top_logprobs_idx, + top_logprobs_num, + ) + ) + except ValueError as e: + # A raise here would disrupt unrelated requests in the + # shared batch-output loop; degrade to nested instead. + state.input_top_logprobs_flat_fields = None + logger.error( + "Falling back to nested input top logprobs for " + "rid=%s: %s", + meta_info.get("id"), + e, + ) + state.input_top_logprobs_flat_num_rows = len( + state.input_top_logprobs_val ) - ) + if state.input_top_logprobs_flat_fields is not None: + meta_info.update(state.input_top_logprobs_flat_fields) + else: + use_flat = False + if not use_flat: + if len(state.input_top_logprobs_val) > len(state.input_top_logprobs): + state.input_top_logprobs.extend( + self.detokenize_top_logprobs_tokens( + state.input_top_logprobs_val[ + len(state.input_top_logprobs) : + ], + state.input_top_logprobs_idx[ + len(state.input_top_logprobs) : + ], + return_text_in_logprobs, + ) + ) + meta_info["input_top_logprobs"] = state.input_top_logprobs if len(state.output_top_logprobs_val) > len(state.output_top_logprobs): state.output_top_logprobs.extend( self.detokenize_top_logprobs_tokens( @@ -2244,8 +2325,6 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): return_text_in_logprobs, ) ) - - meta_info["input_top_logprobs"] = state.input_top_logprobs meta_info["output_top_logprobs"] = state.output_top_logprobs # 3. Handle token_ids_logprob diff --git a/test/registered/unit/managers/test_flat_raw_top_logprobs.py b/test/registered/unit/managers/test_flat_raw_top_logprobs.py new file mode 100644 index 000000000..77be10f00 --- /dev/null +++ b/test/registered/unit/managers/test_flat_raw_top_logprobs.py @@ -0,0 +1,306 @@ +"""Unit tests for the flat raw prompt top logprob response format +(`return_flat_raw_top_logprobs`). +""" + +import asyncio +import json +import os +import time +import unittest + +import numpy as np + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.managers.io_struct import GenerateReqInput +from sglang.srt.managers.tokenizer_manager import ( + ReqState, + TokenizerManager, + _build_flat_input_top_logprobs_fields, +) +from sglang.srt.observability.req_time_stats import APIServerReqTimeStats + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +# Synthetic per-position top-k rows (k=2). The leading None mirrors the first +# prompt position, which has no top logprobs. +_VAL_ROWS = [None, [-0.1, -2.5], [-0.3, -1.5], [-0.05, -4.0]] +_IDX_ROWS = [None, [11, 22], [33, 44], [55, 66]] + + +class _TokenizerManagerStub: + """Borrow the real logprob meta_info methods without a full manager.""" + + add_logprob_to_meta_info = TokenizerManager.add_logprob_to_meta_info + detokenize_logprob_tokens = TokenizerManager.detokenize_logprob_tokens + detokenize_top_logprobs_tokens = TokenizerManager.detokenize_top_logprobs_tokens + + +def _make_state(**obj_kwargs) -> ReqState: + obj = GenerateReqInput(text="hello", **obj_kwargs) + obj.normalize_batch_and_arguments() + return ReqState( + out_list=[], + finished=False, + event=asyncio.Event(), + obj=obj, + time_stats=APIServerReqTimeStats(), + ) + + +def _add_logprob_meta_info(state: ReqState, top_logprobs_num: int = 2) -> dict: + meta_info = {} + _TokenizerManagerStub().add_logprob_to_meta_info( + meta_info, + state, + top_logprobs_num=top_logprobs_num, + token_ids_logprob=None, + return_text_in_logprobs=False, + ) + return meta_info + + +class TestFlatRawTopLogprobsValidation(CustomTestCase): + def test_flag_defaults_off_and_valid(self): + for kwargs in ( + {}, + {"return_flat_raw_top_logprobs": True}, + ): + req = GenerateReqInput(text="hello", **kwargs) + req.normalize_batch_and_arguments() + + def test_flat_rejects_multi_item_scoring(self): + req = GenerateReqInput( + text="ab", + return_flat_raw_top_logprobs=True, + multi_item_delimiter_indices=[1], + ) + with self.assertRaisesRegex(ValueError, "multi-item"): + req.normalize_batch_and_arguments() + + def test_flag_propagates_to_batch_items(self): + req = GenerateReqInput( + text=["a", "b"], + return_flat_raw_top_logprobs=True, + ) + req.normalize_batch_and_arguments() + for i in range(2): + self.assertTrue(req[i].return_flat_raw_top_logprobs) + + +class TestFlatAssembly(CustomTestCase): + def test_flat_matches_nested_rows(self): + fields = _build_flat_input_top_logprobs_fields( + _VAL_ROWS, _IDX_ROWS, top_logprobs_num=2 + ) + self.assertEqual(fields["input_top_logprobs_shape"], [3, 2]) + self.assertEqual(fields["input_top_logprobs_null_prefix"], 1) + self.assertEqual( + fields["input_top_logprobs_val_flat"], + [v for row in _VAL_ROWS[1:] for v in row], + ) + self.assertEqual( + fields["input_top_logprobs_idx_flat"], + [i for row in _IDX_ROWS[1:] for i in row], + ) + # Reconstruct the nested rows: covered position i, entry j lives at + # flat[(i - null_prefix) * k + j]. + rows, k = fields["input_top_logprobs_shape"] + null_prefix = fields["input_top_logprobs_null_prefix"] + flat_val = fields["input_top_logprobs_val_flat"] + self.assertEqual(len(flat_val), rows * k) + for i in range(null_prefix, null_prefix + rows): + start = (i - null_prefix) * k + self.assertEqual(flat_val[start : start + k], _VAL_ROWS[i]) + + def test_all_null_rows(self): + fields = _build_flat_input_top_logprobs_fields( + [None], [None], top_logprobs_num=2 + ) + self.assertEqual(fields["input_top_logprobs_shape"], [0, 2]) + self.assertEqual(fields["input_top_logprobs_null_prefix"], 1) + self.assertEqual(fields["input_top_logprobs_val_flat"], []) + self.assertEqual(fields["input_top_logprobs_idx_flat"], []) + + def test_rejects_null_row_after_prefix(self): + with self.assertRaisesRegex(ValueError, "leading prefix"): + _build_flat_input_top_logprobs_fields( + [None, [-0.1, -2.5], None, [-0.3, -1.5]], + [None, [11, 22], None, [33, 44]], + top_logprobs_num=2, + ) + + def test_rejects_ragged_rows(self): + with self.assertRaisesRegex(ValueError, "rectangular"): + _build_flat_input_top_logprobs_fields( + [None, [-0.1, -2.5], [-0.3]], + [None, [11, 22], [33]], + top_logprobs_num=2, + ) + + +class TestAddLogprobToMetaInfo(CustomTestCase): + def _extend_input_top(self, state: ReqState, val_rows, idx_rows): + state.input_top_logprobs_val.extend(val_rows) + state.input_top_logprobs_idx.extend(idx_rows) + + def test_nested_path_unchanged_when_flags_unset(self): + state = _make_state(return_logprob=True, top_logprobs_num=2) + self._extend_input_top(state, _VAL_ROWS, _IDX_ROWS) + meta_info = _add_logprob_meta_info(state) + self.assertEqual( + meta_info["input_top_logprobs"], + [ + None if row is None else [(v, i, None) for v, i in zip(row, idx_row)] + for row, idx_row in zip(_VAL_ROWS, _IDX_ROWS) + ], + ) + self.assertIn("output_top_logprobs", meta_info) + for key in meta_info: + self.assertNotIn("_flat", key) + self.assertNotIn("input_top_logprobs_shape", meta_info) + self.assertNotIn("input_top_logprobs_null_prefix", meta_info) + + def test_flat_path_replaces_nested_input(self): + state = _make_state( + return_logprob=True, + top_logprobs_num=2, + return_flat_raw_top_logprobs=True, + ) + self._extend_input_top(state, _VAL_ROWS, _IDX_ROWS) + meta_info = _add_logprob_meta_info(state) + self.assertNotIn("input_top_logprobs", meta_info) + self.assertIn("output_top_logprobs", meta_info) + self.assertEqual(meta_info["input_top_logprobs_shape"], [3, 2]) + self.assertEqual(meta_info["input_top_logprobs_null_prefix"], 1) + self.assertEqual( + meta_info["input_top_logprobs_val_flat"], + [v for row in _VAL_ROWS[1:] for v in row], + ) + + def test_unrepresentable_rows_fall_back_to_nested(self): + # The shared batch-output loop must not raise on unrepresentable + # rows; the request degrades to the nested format. + state = _make_state( + return_logprob=True, + top_logprobs_num=2, + return_flat_raw_top_logprobs=True, + ) + val_rows = [_VAL_ROWS[1], None, _VAL_ROWS[2]] + idx_rows = [_IDX_ROWS[1], None, _IDX_ROWS[2]] + self._extend_input_top(state, val_rows, idx_rows) + meta_info = _add_logprob_meta_info(state) + self.assertIn("input_top_logprobs", meta_info) + self.assertNotIn("input_top_logprobs_val_flat", meta_info) + self.assertNotIn("input_top_logprobs_shape", meta_info) + self.assertEqual(len(meta_info["input_top_logprobs"]), 3) + + def test_chunked_accumulation_matches_one_shot(self): + # One shot. + one_shot = _make_state( + return_logprob=True, + top_logprobs_num=2, + return_flat_raw_top_logprobs=True, + ) + self._extend_input_top(one_shot, _VAL_ROWS, _IDX_ROWS) + expected = _add_logprob_meta_info(one_shot) + + # Rows arriving across two chunks, with meta_info assembled after each + # (as happens for streaming requests). + chunked = _make_state( + return_logprob=True, + top_logprobs_num=2, + return_flat_raw_top_logprobs=True, + ) + self._extend_input_top(chunked, _VAL_ROWS[:2], _IDX_ROWS[:2]) + _add_logprob_meta_info(chunked) + self._extend_input_top(chunked, _VAL_ROWS[2:], _IDX_ROWS[2:]) + got = _add_logprob_meta_info(chunked) + + flat_keys = [ + "input_top_logprobs_val_flat", + "input_top_logprobs_idx_flat", + "input_top_logprobs_shape", + "input_top_logprobs_null_prefix", + ] + for key in flat_keys: + self.assertEqual(got[key], expected[key]) + + # No new rows -> the encoded payload is reused, not rebuilt. + again = _add_logprob_meta_info(chunked) + self.assertIs( + again["input_top_logprobs_val_flat"], + got["input_top_logprobs_val_flat"], + ) + + +@unittest.skipUnless( + os.environ.get("SGLANG_BENCH_FLAT_RAW_TOP_LOGPROBS"), + "Serialization microbenchmark; set SGLANG_BENCH_FLAT_RAW_TOP_LOGPROBS=1 to run.", +) +class BenchFlatRawTopLogprobsSerialization(CustomTestCase): + """Round-trip cost of the formats: server assembly + json.dumps, then + client json.loads + reconstruction into [rows, k] arrays.""" + + def test_bench(self): + num_positions, k = 32768, 2 + rng = np.random.default_rng(0) + vals = rng.standard_normal((num_positions, k)).astype(np.float32) + idxs = rng.integers(0, 150000, size=(num_positions, k), dtype=np.int32) + val_rows = [None] + vals[1:].tolist() + idx_rows = [None] + idxs[1:].tolist() + + def best_of(fn, iters=5): + result = fn() + elapsed = min( + (lambda s=time.perf_counter(): (fn(), time.perf_counter() - s)[1])() + for _ in range(iters) + ) + return elapsed * 1e3, result + + def bench(name, build, decode): + encode_ms, payload = best_of(lambda: json.dumps(build())) + decode_ms, arrays = best_of(lambda: decode(payload)) + self.assertEqual(arrays[0].shape, (num_positions - 1, k)) + print( + f"{name}: encode {encode_ms:.1f} ms, decode {decode_ms:.1f} ms, " + f"{len(payload)} bytes" + ) + + def decode_nested(payload): + rows = [r for r in json.loads(payload) if r is not None] + return ( + np.array([[e[0] for e in r] for r in rows], dtype=np.float32), + np.array([[e[1] for e in r] for r in rows], dtype=np.int32), + ) + + def decode_flat(payload): + d = json.loads(payload) + shape = d["input_top_logprobs_shape"] + return ( + np.asarray(d["input_top_logprobs_val_flat"], np.float32).reshape(shape), + np.asarray(d["input_top_logprobs_idx_flat"], np.int32).reshape(shape), + ) + + bench( + "nested triples", + lambda: [ + (None if row is None else [(v, i, None) for v, i in zip(row, idx_row)]) + for row, idx_row in zip(val_rows, idx_rows) + ], + decode_nested, + ) + bench( + "flat lists", + lambda: _build_flat_input_top_logprobs_fields( + val_rows, idx_rows, top_logprobs_num=k + ), + decode_flat, + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2)