762 lines
29 KiB
Python
762 lines
29 KiB
Python
"""Unit tests for the flat raw prompt top logprob response format
|
|
(`return_flat_raw_top_logprobs` / `return_flat_raw_top_logprobs_b64`).
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import os
|
|
import pickle
|
|
import time
|
|
import unittest
|
|
from array import array
|
|
from types import SimpleNamespace
|
|
|
|
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 import runtime_context as rc
|
|
from sglang.srt.managers.io_struct import (
|
|
BatchTokenIDOutput,
|
|
GenerateReqInput,
|
|
build_flat_input_top_logprobs_arrays,
|
|
msgpack_decode,
|
|
msgpack_encode,
|
|
)
|
|
from sglang.srt.managers.schedule_batch import Req
|
|
from sglang.srt.managers.scheduler_components.logprob_result_processor import (
|
|
SchedulerLogprobResultProcessor,
|
|
)
|
|
from sglang.srt.managers.tokenizer_manager import (
|
|
ReqState,
|
|
TokenizerManager,
|
|
_build_flat_input_top_logprobs_fields,
|
|
_build_flat_input_top_logprobs_fields_from_arrays,
|
|
)
|
|
from sglang.srt.observability.req_time_stats import APIServerReqTimeStats
|
|
from sglang.srt.sampling.sampling_params import SamplingParams
|
|
|
|
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]]
|
|
|
|
# Float32-exact values for scheduler-flat equivalence tests: the scheduler
|
|
# ships float32 arrays, so equality against the python-float rows needs values
|
|
# that survive the float64 -> float32 round trip (production logprobs do, being
|
|
# computed in float32).
|
|
_EXACT_VAL_ROWS = [None, [-0.5, -2.5], [-0.25, -1.5], [-0.125, -4.0]]
|
|
|
|
|
|
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)
|
|
|
|
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):
|
|
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_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
|
|
)
|
|
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"],
|
|
)
|
|
|
|
|
|
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"],
|
|
)
|
|
|
|
|
|
def _make_logprob_processor() -> SchedulerLogprobResultProcessor:
|
|
# enable_mis comes from the published exec bag, not from here; see setUp.
|
|
return SchedulerLogprobResultProcessor(
|
|
model_config=SimpleNamespace(vocab_size=1_000_000),
|
|
)
|
|
|
|
|
|
# Per-position rows as computed during prefill: one row per prompt position
|
|
# from logprob_start_len on, the last row being the sampling position that
|
|
# scheduler-side assembly pops.
|
|
_SCHED_VAL_ROWS = [
|
|
[-0.5, -2.5],
|
|
[-0.25, -1.5],
|
|
[-0.125, -4.0],
|
|
[-1.0, -3.0],
|
|
[-2.0, -5.0],
|
|
]
|
|
_SCHED_IDX_ROWS = [[11, 22], [33, 44], [55, 66], [77, 88], [99, 100]]
|
|
|
|
|
|
class TestSchedulerFlatAssembly(CustomTestCase):
|
|
"""Scheduler-side flat assembly in the logprob result processor."""
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
self._server_args_override = rc.get_context().override_server_args()
|
|
self._server_args_override.install()
|
|
|
|
def tearDown(self):
|
|
self._server_args_override.restore()
|
|
|
|
def _make_req(self, flat: bool, num_tokens: int = 5) -> Req:
|
|
return Req(
|
|
"r0",
|
|
"",
|
|
array("q", range(1, num_tokens + 1)),
|
|
SamplingParams(),
|
|
return_logprob=True,
|
|
top_logprobs_num=2,
|
|
return_flat_raw_top_logprobs=flat,
|
|
)
|
|
|
|
def _run_prefill(self, req: Req, chunk_sizes) -> None:
|
|
processor = _make_logprob_processor()
|
|
token_logprobs = [row[0] for row in _SCHED_VAL_ROWS]
|
|
pt = 0
|
|
for chunk_idx, size in enumerate(chunk_sizes):
|
|
output = SimpleNamespace(
|
|
input_token_logprobs=tuple(token_logprobs[pt : pt + size]),
|
|
input_top_logprobs_val=[_SCHED_VAL_ROWS[pt : pt + size]],
|
|
input_top_logprobs_idx=[_SCHED_IDX_ROWS[pt : pt + size]],
|
|
)
|
|
processor.add_input_logprob_return_values(
|
|
0,
|
|
req,
|
|
output,
|
|
0,
|
|
size,
|
|
last_prefill_chunk=chunk_idx == len(chunk_sizes) - 1,
|
|
)
|
|
pt += size
|
|
|
|
def test_flat_arrays_replace_nested_rows(self):
|
|
flag_off = self._make_req(flat=False)
|
|
self._run_prefill(flag_off, [3, 2])
|
|
flag_on = self._make_req(flat=True)
|
|
self._run_prefill(flag_on, [3, 2])
|
|
|
|
# Flag off: nested rows as today, no arrays.
|
|
self.assertIsNone(flag_off.logprob.input_top_logprobs_val_flat)
|
|
self.assertIsNone(flag_off.logprob.input_top_logprobs_flat_null_prefix)
|
|
self.assertEqual(
|
|
flag_off.logprob.input_top_logprobs_val, [None] + _SCHED_VAL_ROWS[:-1]
|
|
)
|
|
|
|
# Flag on: arrays carrying the nested rows' content, nested emptied.
|
|
val_arr = flag_on.logprob.input_top_logprobs_val_flat
|
|
idx_arr = flag_on.logprob.input_top_logprobs_idx_flat
|
|
self.assertEqual(val_arr.dtype, np.float32)
|
|
self.assertEqual(idx_arr.dtype, np.int32)
|
|
self.assertEqual(flag_on.logprob.input_top_logprobs_flat_null_prefix, 1)
|
|
np.testing.assert_array_equal(
|
|
val_arr, np.asarray(_SCHED_VAL_ROWS[:-1], dtype=np.float32)
|
|
)
|
|
np.testing.assert_array_equal(
|
|
idx_arr, np.asarray(_SCHED_IDX_ROWS[:-1], dtype=np.int32)
|
|
)
|
|
self.assertEqual(flag_on.logprob.input_top_logprobs_val, [])
|
|
self.assertEqual(flag_on.logprob.input_top_logprobs_idx, [])
|
|
# The non-top logprob results are untouched.
|
|
self.assertEqual(
|
|
flag_on.logprob.input_token_logprobs_val,
|
|
flag_off.logprob.input_token_logprobs_val,
|
|
)
|
|
self.assertEqual(
|
|
flag_on.logprob.input_token_logprobs_idx,
|
|
flag_off.logprob.input_token_logprobs_idx,
|
|
)
|
|
|
|
def test_chunked_matches_one_shot(self):
|
|
one_shot = self._make_req(flat=True)
|
|
self._run_prefill(one_shot, [5])
|
|
chunked = self._make_req(flat=True)
|
|
self._run_prefill(chunked, [2, 2, 1])
|
|
np.testing.assert_array_equal(
|
|
one_shot.logprob.input_top_logprobs_val_flat,
|
|
chunked.logprob.input_top_logprobs_val_flat,
|
|
)
|
|
np.testing.assert_array_equal(
|
|
one_shot.logprob.input_top_logprobs_idx_flat,
|
|
chunked.logprob.input_top_logprobs_idx_flat,
|
|
)
|
|
self.assertEqual(
|
|
one_shot.logprob.input_top_logprobs_flat_null_prefix,
|
|
chunked.logprob.input_top_logprobs_flat_null_prefix,
|
|
)
|
|
|
|
def test_unrepresentable_rows_fall_back_to_nested(self):
|
|
req = self._make_req(flat=True, num_tokens=3)
|
|
processor = _make_logprob_processor()
|
|
val_rows = [[-0.5, -2.5], [-0.25], [-0.125, -4.0]]
|
|
idx_rows = [[11, 22], [33], [55, 66]]
|
|
output = SimpleNamespace(
|
|
input_token_logprobs=(-0.5, -0.25, -0.125),
|
|
input_top_logprobs_val=[val_rows],
|
|
input_top_logprobs_idx=[idx_rows],
|
|
)
|
|
with self.assertLogs(
|
|
"sglang.srt.managers.scheduler_components.logprob_result_processor",
|
|
level="WARNING",
|
|
):
|
|
processor.add_input_logprob_return_values(
|
|
0, req, output, 0, 3, last_prefill_chunk=True
|
|
)
|
|
self.assertIsNone(req.logprob.input_top_logprobs_val_flat)
|
|
self.assertIsNone(req.logprob.input_top_logprobs_flat_null_prefix)
|
|
self.assertEqual(req.logprob.input_top_logprobs_val, [None] + val_rows[:-1])
|
|
self.assertEqual(req.logprob.input_top_logprobs_idx, [None] + idx_rows[:-1])
|
|
|
|
|
|
class TestFromArraysMatchesFromRows(CustomTestCase):
|
|
"""The tokenizer-manager from-arrays builder must produce the same
|
|
response fields as the rows-based builder."""
|
|
|
|
def _both(self, return_b64: bool):
|
|
from_rows = _build_flat_input_top_logprobs_fields(
|
|
_EXACT_VAL_ROWS, _IDX_ROWS, top_logprobs_num=2, return_b64=return_b64
|
|
)
|
|
val_arr, idx_arr, null_prefix = build_flat_input_top_logprobs_arrays(
|
|
_EXACT_VAL_ROWS, _IDX_ROWS, top_logprobs_num=2
|
|
)
|
|
from_arrays = _build_flat_input_top_logprobs_fields_from_arrays(
|
|
val_arr, idx_arr, null_prefix, return_b64=return_b64
|
|
)
|
|
return from_rows, from_arrays
|
|
|
|
def test_non_b64(self):
|
|
from_rows, from_arrays = self._both(return_b64=False)
|
|
self.assertEqual(from_rows, from_arrays)
|
|
|
|
def test_b64(self):
|
|
from_rows, from_arrays = self._both(return_b64=True)
|
|
self.assertEqual(from_rows, from_arrays)
|
|
|
|
def test_all_null_rows(self):
|
|
val_arr, idx_arr, null_prefix = build_flat_input_top_logprobs_arrays(
|
|
[None], [None], top_logprobs_num=2
|
|
)
|
|
self.assertEqual(val_arr.shape, (0, 2))
|
|
self.assertEqual(null_prefix, 1)
|
|
fields = _build_flat_input_top_logprobs_fields_from_arrays(
|
|
val_arr, idx_arr, null_prefix
|
|
)
|
|
self.assertEqual(
|
|
fields,
|
|
_build_flat_input_top_logprobs_fields([None], [None], top_logprobs_num=2),
|
|
)
|
|
|
|
|
|
class TestMetaInfoFromSchedulerArrays(CustomTestCase):
|
|
"""add_logprob_to_meta_info consumes scheduler-flat arrays directly."""
|
|
|
|
def _rows_meta(self, **state_kwargs) -> dict:
|
|
state = _make_state(
|
|
return_logprob=True,
|
|
top_logprobs_num=2,
|
|
return_flat_raw_top_logprobs=True,
|
|
**state_kwargs,
|
|
)
|
|
state.input_top_logprobs_val.extend(_EXACT_VAL_ROWS)
|
|
state.input_top_logprobs_idx.extend(_IDX_ROWS)
|
|
return _add_logprob_meta_info(state)
|
|
|
|
def _arrays_state(self, **state_kwargs) -> ReqState:
|
|
state = _make_state(
|
|
return_logprob=True,
|
|
top_logprobs_num=2,
|
|
return_flat_raw_top_logprobs=True,
|
|
**state_kwargs,
|
|
)
|
|
# Scheduler-flat requests arrive with empty nested rows and the arrays.
|
|
state.input_top_logprobs_scheduler_flat = build_flat_input_top_logprobs_arrays(
|
|
_EXACT_VAL_ROWS, _IDX_ROWS, top_logprobs_num=2
|
|
)
|
|
return state
|
|
|
|
def test_matches_rows_path_field_for_field(self):
|
|
got = _add_logprob_meta_info(self._arrays_state())
|
|
self.assertEqual(got, self._rows_meta())
|
|
|
|
def test_b64_matches_rows_path_field_for_field(self):
|
|
got = _add_logprob_meta_info(
|
|
self._arrays_state(return_flat_raw_top_logprobs_b64=True)
|
|
)
|
|
self.assertEqual(got, self._rows_meta(return_flat_raw_top_logprobs_b64=True))
|
|
|
|
def test_fields_cached_across_chunks(self):
|
|
state = self._arrays_state()
|
|
first = _add_logprob_meta_info(state)
|
|
again = _add_logprob_meta_info(state)
|
|
self.assertIs(
|
|
again["input_top_logprobs_val_flat"], first["input_top_logprobs_val_flat"]
|
|
)
|
|
|
|
|
|
def _make_batch_token_id_output(**overrides) -> BatchTokenIDOutput:
|
|
"""A two-request BatchTokenIDOutput with the required fields stubbed."""
|
|
n = 2
|
|
fields = dict(
|
|
rids=["r0", "r1"],
|
|
finished_reasons=[None] * n,
|
|
decoded_texts=["", ""],
|
|
decode_ids=[array("q", [1]), array("q", [2])],
|
|
read_offsets=[0] * n,
|
|
output_ids=None,
|
|
skip_special_tokens=[True] * n,
|
|
spaces_between_special_tokens=[True] * n,
|
|
no_stop_trim=[False] * n,
|
|
prompt_tokens=[5] * n,
|
|
reasoning_tokens=[0] * n,
|
|
completion_tokens=[1] * n,
|
|
cached_tokens=[0] * n,
|
|
input_token_logprobs_val=[[], []],
|
|
input_token_logprobs_idx=[[], []],
|
|
output_token_logprobs_val=[[], []],
|
|
output_token_logprobs_idx=[[], []],
|
|
input_top_logprobs_val=[[], []],
|
|
input_top_logprobs_idx=[[], []],
|
|
output_top_logprobs_val=[[], []],
|
|
output_top_logprobs_idx=[[], []],
|
|
input_token_ids_logprobs_val=[[], []],
|
|
input_token_ids_logprobs_idx=[[], []],
|
|
output_token_ids_logprobs_val=[[], []],
|
|
output_token_ids_logprobs_idx=[[], []],
|
|
output_token_entropy_val=None,
|
|
output_token_sampling_mask=None,
|
|
output_token_sampling_logprobs=None,
|
|
output_hidden_states=None,
|
|
routed_experts=None,
|
|
indexer_topk=None,
|
|
placeholder_tokens_idx=None,
|
|
placeholder_tokens_val=None,
|
|
)
|
|
fields.update(overrides)
|
|
return BatchTokenIDOutput(**fields)
|
|
|
|
|
|
class TestBatchOutputTransport(CustomTestCase):
|
|
"""The flat array fields must survive both IPC transports: pickle
|
|
(SGLANG_USE_PICKLE_IPC, the default) and msgpack (enc/dec hooks)."""
|
|
|
|
def _flat_output(self) -> BatchTokenIDOutput:
|
|
val_arr, idx_arr, null_prefix = build_flat_input_top_logprobs_arrays(
|
|
_EXACT_VAL_ROWS, _IDX_ROWS, top_logprobs_num=2
|
|
)
|
|
return _make_batch_token_id_output(
|
|
input_top_logprobs_val_flat=[None, val_arr],
|
|
input_top_logprobs_idx_flat=[None, idx_arr],
|
|
input_top_logprobs_flat_null_prefix=[None, null_prefix],
|
|
)
|
|
|
|
def _check_roundtrip(self, decoded, original):
|
|
self.assertIsNone(decoded.input_top_logprobs_val_flat[0])
|
|
self.assertIsNone(decoded.input_top_logprobs_idx_flat[0])
|
|
self.assertIsNone(decoded.input_top_logprobs_flat_null_prefix[0])
|
|
for got, sent in (
|
|
(
|
|
decoded.input_top_logprobs_val_flat[1],
|
|
original.input_top_logprobs_val_flat[1],
|
|
),
|
|
(
|
|
decoded.input_top_logprobs_idx_flat[1],
|
|
original.input_top_logprobs_idx_flat[1],
|
|
),
|
|
):
|
|
self.assertIsInstance(got, np.ndarray)
|
|
self.assertEqual(got.dtype, sent.dtype)
|
|
np.testing.assert_array_equal(got, sent)
|
|
self.assertEqual(decoded.input_top_logprobs_flat_null_prefix[1], 1)
|
|
|
|
def test_pickle_roundtrip(self):
|
|
output = self._flat_output()
|
|
decoded = pickle.loads(pickle.dumps(output, protocol=pickle.HIGHEST_PROTOCOL))
|
|
self._check_roundtrip(decoded, output)
|
|
|
|
def test_msgpack_roundtrip(self):
|
|
output = self._flat_output()
|
|
decoded = msgpack_decode(msgpack_encode(output))
|
|
self._check_roundtrip(decoded, output)
|
|
|
|
def test_fields_default_none(self):
|
|
output = _make_batch_token_id_output()
|
|
self.assertIsNone(output.input_top_logprobs_val_flat)
|
|
self.assertIsNone(output.input_top_logprobs_idx_flat)
|
|
self.assertIsNone(output.input_top_logprobs_flat_null_prefix)
|
|
|
|
|
|
@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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
def test_bench_ipc_pickle(self):
|
|
"""Inter-process cost of BatchTokenIDOutput input-top fields: nested
|
|
per-position rows vs scheduler-flat arrays (two ZMQ pickle hops each
|
|
pay dumps + loads)."""
|
|
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)
|
|
|
|
def best_of(fn, iters=10):
|
|
return min(
|
|
(lambda s=time.perf_counter(): (fn(), time.perf_counter() - s)[1])()
|
|
for _ in range(iters)
|
|
)
|
|
|
|
nested = _make_batch_token_id_output(
|
|
input_top_logprobs_val=[[None] + vals[1:].tolist(), []],
|
|
input_top_logprobs_idx=[[None] + idxs[1:].tolist(), []],
|
|
)
|
|
flat = _make_batch_token_id_output(
|
|
input_top_logprobs_val_flat=[vals[1:], None],
|
|
input_top_logprobs_idx_flat=[idxs[1:], None],
|
|
input_top_logprobs_flat_null_prefix=[1, None],
|
|
)
|
|
for name, obj in (("nested rows", nested), ("flat arrays", flat)):
|
|
payload = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
|
|
dumps_ms = best_of(
|
|
lambda o=obj: pickle.dumps(o, protocol=pickle.HIGHEST_PROTOCOL)
|
|
)
|
|
loads_ms = best_of(lambda p=payload: pickle.loads(p))
|
|
print(
|
|
f"{name}: pickle.dumps {dumps_ms * 1e3:.2f} ms, "
|
|
f"pickle.loads {loads_ms * 1e3:.2f} ms, {len(payload) / 1e6:.2f} MB"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2)
|