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

[feat] Opt-in flat response format for prompt top logprobs (#32078)

Browse Source
This commit is contained in:
Sam Shleifer
2026-07-26 23:44:34 -07:00
committed by GitHub
parent 4ea17169b0
commit 5cc273a780
3 changed files with 407 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/io_struct.py
+13
View File
@@ -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=(
python/sglang/srt/managers/tokenizer_manager.py
+88 -9
View File
@@ -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
test/registered/unit/managers/test_flat_raw_top_logprobs.py
+306
View File
@@ -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="a<sep>b",
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)