minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[feat] Optional base64 encoding for the flat prompt top logprob arrays (#31960)

Browse Source
This commit is contained in:
Sam Shleifer
2026-07-29 12:15:56 -07:00
committed by GitHub
parent e4f7f7b380
commit d0e69d3881
3 changed files with 112 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/io_struct.py
+10
View File
@@ -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=(
python/sglang/srt/managers/tokenizer_manager.py
+20 -3
View File
@@ -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:
test/registered/unit/managers/test_flat_raw_top_logprobs.py
+82 -1
View File
@@ -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)