diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index ef1617c8e..32486141c 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -211,6 +211,8 @@ class GenerateReqInput: # 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 + # Base64-encode the flat arrays. Requires return_flat_raw_top_logprobs. + return_flat_raw_top_logprobs_b64: 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) @@ -381,6 +383,13 @@ class GenerateReqInput: "scoring: delimiter-sparse top logprob rows have no contiguous " "position mapping." ) + if ( + self.return_flat_raw_top_logprobs_b64 + and not self.return_flat_raw_top_logprobs + ): + raise ValueError( + "return_flat_raw_top_logprobs_b64 requires return_flat_raw_top_logprobs." + ) def _determine_batch_size(self): """Determine if this is a single example or a batch and the batch size.""" @@ -737,6 +746,7 @@ class GenerateReqInput: 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, + return_flat_raw_top_logprobs_b64=self.return_flat_raw_top_logprobs_b64, 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 1ab68ead2..c8934b27e 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -37,6 +37,7 @@ from http import HTTPStatus from typing import Any, Awaitable, Dict, Iterable, List, Optional, Tuple, Union import fastapi +import numpy as np import pybase64 import torch import uvloop @@ -263,6 +264,7 @@ 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, + return_b64: bool = False, ) -> Dict[str, Any]: """Build the flat raw prompt top logprob response fields. @@ -270,7 +272,9 @@ def _build_flat_input_top_logprobs_fields( 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. + spans null_prefix + rows positions. With ``return_b64``, the arrays are + base64 contiguous little-endian binary; the dtype marker fields let the + widths change later without a wire break. """ num_rows = len(input_top_logprobs_val) null_prefix = 0 @@ -289,8 +293,20 @@ def _build_flat_input_top_logprobs_fields( ) 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] + if return_b64: + val_arr = np.asarray(val_rows, dtype=np.float32) + idx_arr = np.asarray(idx_rows, dtype=np.int32) + fields["input_top_logprobs_val_flat_b64"] = pybase64.b64encode( + val_arr.tobytes() + ).decode("utf-8") + fields["input_top_logprobs_idx_flat_b64"] = pybase64.b64encode( + idx_arr.tobytes() + ).decode("utf-8") + fields["input_top_logprobs_val_flat_b64_dtype"] = "float32" + fields["input_top_logprobs_idx_flat_b64_dtype"] = "int32" + else: + 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 @@ -2301,6 +2317,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): state.input_top_logprobs_val, state.input_top_logprobs_idx, top_logprobs_num, + return_b64=state.obj.return_flat_raw_top_logprobs_b64, ) ) except ValueError as e: diff --git a/test/registered/unit/managers/test_flat_raw_top_logprobs.py b/test/registered/unit/managers/test_flat_raw_top_logprobs.py index 77be10f00..3dab14a17 100644 --- a/test/registered/unit/managers/test_flat_raw_top_logprobs.py +++ b/test/registered/unit/managers/test_flat_raw_top_logprobs.py @@ -1,8 +1,9 @@ """Unit tests for the flat raw prompt top logprob response format -(`return_flat_raw_top_logprobs`). +(`return_flat_raw_top_logprobs` / `return_flat_raw_top_logprobs_b64`). """ import asyncio +import base64 import json import os import time @@ -90,6 +91,21 @@ class TestFlatRawTopLogprobsValidation(CustomTestCase): for i in range(2): self.assertTrue(req[i].return_flat_raw_top_logprobs) + def test_b64_requires_flat_flag(self): + req = GenerateReqInput(text="hello", return_flat_raw_top_logprobs_b64=True) + with self.assertRaisesRegex(ValueError, "return_flat_raw_top_logprobs"): + req.normalize_batch_and_arguments() + + def test_b64_flag_propagates_to_batch_items(self): + req = GenerateReqInput( + text=["a", "b"], + return_flat_raw_top_logprobs=True, + return_flat_raw_top_logprobs_b64=True, + ) + req.normalize_batch_and_arguments() + for i in range(2): + self.assertTrue(req[i].return_flat_raw_top_logprobs_b64) + class TestFlatAssembly(CustomTestCase): def test_flat_matches_nested_rows(self): @@ -116,6 +132,26 @@ class TestFlatAssembly(CustomTestCase): start = (i - null_prefix) * k self.assertEqual(flat_val[start : start + k], _VAL_ROWS[i]) + def test_b64_roundtrip(self): + fields = _build_flat_input_top_logprobs_fields( + _VAL_ROWS, _IDX_ROWS, top_logprobs_num=2, return_b64=True + ) + 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_b64_dtype"], "float32") + self.assertEqual(fields["input_top_logprobs_idx_flat_b64_dtype"], "int32") + # shape is the literal array shape, so the b64 buffer reshapes with it. + val = np.frombuffer( + base64.b64decode(fields["input_top_logprobs_val_flat_b64"]), + dtype=np.dtype(fields["input_top_logprobs_val_flat_b64_dtype"]), + ).reshape(fields["input_top_logprobs_shape"]) + idx = np.frombuffer( + base64.b64decode(fields["input_top_logprobs_idx_flat_b64"]), + dtype=np.dtype(fields["input_top_logprobs_idx_flat_b64_dtype"]), + ).reshape(fields["input_top_logprobs_shape"]) + np.testing.assert_array_equal(val, np.asarray(_VAL_ROWS[1:], dtype=np.float32)) + np.testing.assert_array_equal(idx, np.asarray(_IDX_ROWS[1:], dtype=np.int32)) + def test_all_null_rows(self): fields = _build_flat_input_top_logprobs_fields( [None], [None], top_logprobs_num=2 @@ -237,6 +273,29 @@ class TestAddLogprobToMetaInfo(CustomTestCase): ) +class TestB64MetaInfo(CustomTestCase): + def test_b64_fields_replace_flat_and_cache_reused(self): + state = _make_state( + return_logprob=True, + top_logprobs_num=2, + return_flat_raw_top_logprobs=True, + return_flat_raw_top_logprobs_b64=True, + ) + state.input_top_logprobs_val.extend(_VAL_ROWS) + state.input_top_logprobs_idx.extend(_IDX_ROWS) + meta_info = _add_logprob_meta_info(state) + self.assertNotIn("input_top_logprobs", meta_info) + self.assertNotIn("input_top_logprobs_val_flat", meta_info) + self.assertIn("input_top_logprobs_val_flat_b64", meta_info) + self.assertEqual(meta_info["input_top_logprobs_shape"], [3, 2]) + # No new rows -> the encoded payload is reused, not rebuilt. + again = _add_logprob_meta_info(state) + self.assertIs( + again["input_top_logprobs_val_flat_b64"], + meta_info["input_top_logprobs_val_flat_b64"], + ) + + @unittest.skipUnless( os.environ.get("SGLANG_BENCH_FLAT_RAW_TOP_LOGPROBS"), "Serialization microbenchmark; set SGLANG_BENCH_FLAT_RAW_TOP_LOGPROBS=1 to run.", @@ -301,6 +360,28 @@ class BenchFlatRawTopLogprobsSerialization(CustomTestCase): decode_flat, ) + def decode_b64(payload): + d = json.loads(payload) + shape = d["input_top_logprobs_shape"] + return ( + np.frombuffer( + base64.b64decode(d["input_top_logprobs_val_flat_b64"]), + np.dtype(d["input_top_logprobs_val_flat_b64_dtype"]), + ).reshape(shape), + np.frombuffer( + base64.b64decode(d["input_top_logprobs_idx_flat_b64"]), + np.dtype(d["input_top_logprobs_idx_flat_b64_dtype"]), + ).reshape(shape), + ) + + bench( + "flat b64", + lambda: _build_flat_input_top_logprobs_fields( + val_rows, idx_rows, top_logprobs_num=k, return_b64=True + ), + decode_b64, + ) + if __name__ == "__main__": unittest.main(verbosity=2)