468 lines
14 KiB
Python
468 lines
14 KiB
Python
"""
|
|
Benchmark: Streaming Session Inter-Turn Latency
|
|
|
|
Tests:
|
|
1. Latency (bs=8): regular vs streaming, assert speedup >= 2x
|
|
2. Correctness (bs=1): regular vs streaming, assert output equal + speedup
|
|
3. Random lengths (bs=8): streaming only, random input/output lens, no crash
|
|
|
|
Usage:
|
|
python -m pytest test_session_latency.py -s
|
|
python -m unittest test_session_latency.BenchSessionLatency
|
|
"""
|
|
|
|
import random
|
|
import time
|
|
import unittest
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, List, Optional
|
|
|
|
import requests
|
|
from tabulate import tabulate
|
|
|
|
from sglang.srt.utils import kill_process_tree
|
|
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
from sglang.test.test_utils import (
|
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
DEFAULT_URL_FOR_TEST,
|
|
CustomTestCase,
|
|
popen_launch_server,
|
|
)
|
|
|
|
register_cuda_ci(
|
|
est_time=122,
|
|
suite="stage-b-test-1-gpu-large",
|
|
)
|
|
|
|
NUM_TURNS = 150
|
|
INPUT_LEN = 16
|
|
GEN_LEN = 8
|
|
NUM_CONCURRENT = 8
|
|
TAIL_TURNS = 10
|
|
SAMPLE_TURNS = 8
|
|
|
|
NUM_TURNS_RANDOM = 50
|
|
RANDOM_INPUT_LEN_RANGE = (8, 64)
|
|
RANDOM_OUTPUT_LEN_RANGE = (4, 32)
|
|
|
|
FILLER_TEXT = (
|
|
"The quick brown fox jumps over the lazy dog. "
|
|
"Pack my box with five dozen liquor jugs. "
|
|
"How vexingly quick daft zebras jump. "
|
|
"Sphinx of black quartz, judge my vow. "
|
|
) * 200
|
|
|
|
SAMPLING_PARAMS = {
|
|
"temperature": 0,
|
|
"max_new_tokens": GEN_LEN,
|
|
"no_stop_trim": True,
|
|
"skip_special_tokens": False,
|
|
"ignore_eos": True,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class TurnResult:
|
|
turn: int
|
|
context_len: int
|
|
cached_tokens: int
|
|
prompt_tokens: int
|
|
completion_tokens: int
|
|
client_latency_ms: float
|
|
e2e_latency_ms: float
|
|
|
|
|
|
@dataclass
|
|
class ModeResult:
|
|
mode: str
|
|
turns: List[TurnResult] = field(default_factory=list)
|
|
outputs: List[str] = field(default_factory=list)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _generate_input_chunks(
|
|
tokenizer, num_turns: int, input_len: int, offset: int = 0
|
|
) -> List[List[int]]:
|
|
all_ids = tokenizer.encode(FILLER_TEXT)
|
|
if all_ids and all_ids[0] == tokenizer.bos_token_id:
|
|
all_ids = all_ids[1:]
|
|
|
|
start = offset * num_turns * input_len
|
|
needed = start + num_turns * input_len
|
|
while len(all_ids) < needed:
|
|
all_ids = all_ids + all_ids
|
|
chunks = [
|
|
all_ids[start + i * input_len : start + (i + 1) * input_len]
|
|
for i in range(num_turns)
|
|
]
|
|
|
|
if tokenizer.bos_token_id is not None:
|
|
chunks[0] = [tokenizer.bos_token_id] + chunks[0]
|
|
|
|
return chunks
|
|
|
|
|
|
def _generate_random_input_chunks(
|
|
tokenizer,
|
|
num_turns: int,
|
|
min_len: int,
|
|
max_len: int,
|
|
rng: random.Random,
|
|
offset: int = 0,
|
|
) -> List[List[int]]:
|
|
all_ids = tokenizer.encode(FILLER_TEXT)
|
|
if all_ids and all_ids[0] == tokenizer.bos_token_id:
|
|
all_ids = all_ids[1:]
|
|
|
|
total_max = offset * num_turns * max_len + num_turns * max_len
|
|
while len(all_ids) < total_max:
|
|
all_ids = all_ids + all_ids
|
|
|
|
chunks: List[List[int]] = []
|
|
pos = offset * num_turns * max_len
|
|
for i in range(num_turns):
|
|
length = rng.randint(min_len, max_len)
|
|
chunk = all_ids[pos : pos + length]
|
|
pos += length
|
|
chunks.append(chunk)
|
|
|
|
if tokenizer.bos_token_id is not None:
|
|
chunks[0] = [tokenizer.bos_token_id] + chunks[0]
|
|
|
|
return chunks
|
|
|
|
|
|
def _send_generate(base_url: str, payload: dict) -> dict:
|
|
resp = requests.post(base_url + "/generate", json=payload)
|
|
if resp.status_code != 200:
|
|
raise RuntimeError(f"Generate failed ({resp.status_code}): {resp.text}")
|
|
return resp.json()
|
|
|
|
|
|
def _record_turn(
|
|
turn_idx: int, context_len: int, meta: dict, client_latency_ms: float
|
|
) -> TurnResult:
|
|
return TurnResult(
|
|
turn=turn_idx + 1,
|
|
context_len=context_len,
|
|
cached_tokens=meta["cached_tokens"],
|
|
prompt_tokens=meta["prompt_tokens"],
|
|
completion_tokens=meta["completion_tokens"],
|
|
client_latency_ms=client_latency_ms,
|
|
e2e_latency_ms=meta.get("e2e_latency", 0) * 1000,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Single-session runner (called by worker threads)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _run_one_session(
|
|
base_url: str,
|
|
chunks: List[List[int]],
|
|
streaming: bool = False,
|
|
per_turn_gen_lens: Optional[List[int]] = None,
|
|
) -> ModeResult:
|
|
mode = "streaming_session" if streaming else "regular_session"
|
|
result = ModeResult(mode=mode)
|
|
|
|
default_gen = GEN_LEN
|
|
if per_turn_gen_lens is not None:
|
|
max_gen = max(per_turn_gen_lens)
|
|
else:
|
|
max_gen = default_gen
|
|
capacity = sum(len(c) for c in chunks) + len(chunks) * max_gen + 1024
|
|
|
|
open_payload: dict = {"capacity_of_str_len": capacity}
|
|
if streaming:
|
|
open_payload["streaming"] = True
|
|
session_id = requests.post(base_url + "/open_session", json=open_payload).json()
|
|
|
|
rid = None
|
|
context_len = 0
|
|
|
|
for turn_idx, chunk_ids in enumerate(chunks):
|
|
context_len += len(chunk_ids)
|
|
|
|
if per_turn_gen_lens is not None:
|
|
sp = {**SAMPLING_PARAMS, "max_new_tokens": per_turn_gen_lens[turn_idx]}
|
|
else:
|
|
sp = SAMPLING_PARAMS
|
|
|
|
t0 = time.perf_counter()
|
|
response = _send_generate(
|
|
base_url,
|
|
{
|
|
"input_ids": chunk_ids,
|
|
"session_params": {"id": session_id, "rid": rid},
|
|
"sampling_params": sp,
|
|
},
|
|
)
|
|
client_lat = (time.perf_counter() - t0) * 1000
|
|
|
|
meta = response["meta_info"]
|
|
rid = meta["id"]
|
|
context_len += meta["completion_tokens"]
|
|
|
|
result.turns.append(_record_turn(turn_idx, context_len, meta, client_lat))
|
|
result.outputs.append(response["text"])
|
|
|
|
requests.post(base_url + "/close_session", json={"session_id": session_id})
|
|
return result
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Stats & reporting
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _collect_latencies(
|
|
results: List[ModeResult], last_n: Optional[int] = None
|
|
) -> List[float]:
|
|
lats = []
|
|
for r in results:
|
|
turns = r.turns[1:] # skip turn 1
|
|
if last_n is not None:
|
|
turns = r.turns[-last_n:]
|
|
lats.extend(t.client_latency_ms for t in turns)
|
|
return lats
|
|
|
|
|
|
def _avg(values: List[float]) -> float:
|
|
return sum(values) / len(values) if values else 0.0
|
|
|
|
|
|
def _print_mode_table(result: ModeResult, label: str = ""):
|
|
tag = f"{result.mode} ({label})" if label else result.mode
|
|
print(f"\n [{tag}] {len(result.turns)} turns")
|
|
|
|
n = len(result.turns)
|
|
if n <= SAMPLE_TURNS * 2:
|
|
indices = list(range(n))
|
|
else:
|
|
indices = list(range(SAMPLE_TURNS)) + [-1] + list(range(n - SAMPLE_TURNS, n))
|
|
|
|
rows = []
|
|
for idx in indices:
|
|
if idx == -1:
|
|
rows.append(["..."] * 5)
|
|
continue
|
|
t = result.turns[idx]
|
|
rows.append(
|
|
[
|
|
t.turn,
|
|
t.context_len,
|
|
t.cached_tokens,
|
|
f"{t.client_latency_ms:.1f}ms",
|
|
f"{t.e2e_latency_ms:.1f}ms",
|
|
]
|
|
)
|
|
print(
|
|
tabulate(
|
|
rows,
|
|
headers=["Turn", "Context", "Cached", "Client Lat", "E2E Lat"],
|
|
colalign=("right",) * 5,
|
|
)
|
|
)
|
|
|
|
|
|
def _print_summary(all_results: Dict[str, List[ModeResult]]):
|
|
stats = [
|
|
(
|
|
mode,
|
|
_avg(_collect_latencies(rs)),
|
|
_avg(_collect_latencies(rs, last_n=TAIL_TURNS)),
|
|
)
|
|
for mode, rs in all_results.items()
|
|
]
|
|
base_all, base_tail = (stats[0][1] or 1.0), (stats[0][2] or 1.0)
|
|
tail_label = f"last {TAIL_TURNS}"
|
|
|
|
print(f"\n SUMMARY ({NUM_CONCURRENT} sessions x {NUM_TURNS} turns)")
|
|
rows = [
|
|
[
|
|
mode,
|
|
f"{a:.1f}ms",
|
|
f"{t:.1f}ms",
|
|
f"{base_all / a:.2f}x" if a else "inf",
|
|
f"{base_tail / t:.2f}x" if t else "inf",
|
|
]
|
|
for mode, a, t in stats
|
|
]
|
|
print(
|
|
tabulate(
|
|
rows,
|
|
headers=[
|
|
"Mode",
|
|
"Avg (all)",
|
|
f"Avg ({tail_label})",
|
|
"Speedup (all)",
|
|
f"Speedup ({tail_label})",
|
|
],
|
|
colalign=("left", "right", "right", "right", "right"),
|
|
)
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test class
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSessionLatency(CustomTestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.model = "openai/gpt-oss-20b"
|
|
cls.base_url = DEFAULT_URL_FOR_TEST
|
|
# NOTE: Overlap scheduling commits KV cache one step ahead,
|
|
# so the last decode token is cached (unlike non-overlap).
|
|
# Disable overlap to keep session cache behavior consistent.
|
|
cls.process = popen_launch_server(
|
|
cls.model,
|
|
cls.base_url,
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=[
|
|
"--disable-overlap-schedule",
|
|
"--enable-streaming-session",
|
|
"--mem-fraction-static",
|
|
"0.70",
|
|
"--disable-piecewise-cuda-graph",
|
|
"--page-size",
|
|
"4",
|
|
],
|
|
)
|
|
cls.tokenizer = get_tokenizer(cls.model)
|
|
|
|
requests.post(cls.base_url + "/flush_cache")
|
|
_send_generate(
|
|
cls.base_url,
|
|
{
|
|
"input_ids": cls.tokenizer.encode("Hello world"),
|
|
"sampling_params": {"temperature": 0, "max_new_tokens": 1},
|
|
},
|
|
)
|
|
|
|
cls.all_results: Dict[str, List[ModeResult]] = {}
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
if len(cls.all_results) > 1:
|
|
_print_summary(cls.all_results)
|
|
kill_process_tree(cls.process.pid)
|
|
|
|
def _run_concurrent_session(
|
|
self,
|
|
streaming: bool = False,
|
|
num_concurrent: int = NUM_CONCURRENT,
|
|
num_turns: int = NUM_TURNS,
|
|
input_len: int = INPUT_LEN,
|
|
per_turn_gen_lens: Optional[List[int]] = None,
|
|
random_input_chunks: bool = False,
|
|
rng: Optional[random.Random] = None,
|
|
) -> List[ModeResult]:
|
|
requests.post(self.base_url + "/flush_cache")
|
|
|
|
def run_one(session_idx):
|
|
if random_input_chunks and rng is not None:
|
|
per_session_rng = random.Random(rng.randint(0, 2**32) + session_idx)
|
|
chunks = _generate_random_input_chunks(
|
|
self.tokenizer,
|
|
num_turns,
|
|
RANDOM_INPUT_LEN_RANGE[0],
|
|
RANDOM_INPUT_LEN_RANGE[1],
|
|
per_session_rng,
|
|
offset=session_idx,
|
|
)
|
|
else:
|
|
chunks = _generate_input_chunks(
|
|
self.tokenizer, num_turns, input_len, offset=session_idx
|
|
)
|
|
return _run_one_session(
|
|
self.base_url,
|
|
chunks,
|
|
streaming=streaming,
|
|
per_turn_gen_lens=per_turn_gen_lens,
|
|
)
|
|
|
|
with ThreadPoolExecutor(max_workers=num_concurrent) as pool:
|
|
return list(pool.map(run_one, range(num_concurrent)))
|
|
|
|
# ------------------------------------------------------------------
|
|
# Test methods (alphabetical order matters for dependencies)
|
|
# ------------------------------------------------------------------
|
|
|
|
def test_regular_session(self):
|
|
"""Run regular (non-streaming) sessions for latency baseline."""
|
|
results = self._run_concurrent_session(streaming=False)
|
|
self.__class__.all_results["regular_session"] = results
|
|
_print_mode_table(results[0], label="session 0")
|
|
|
|
def test_streaming_session(self):
|
|
"""Latency test: bs=8, assert streaming >= 2x faster than regular."""
|
|
results = self._run_concurrent_session(streaming=True)
|
|
self.__class__.all_results["streaming_session"] = results
|
|
_print_mode_table(results[0], label="session 0")
|
|
|
|
reg_list = self.__class__.all_results.get("regular_session")
|
|
if reg_list:
|
|
reg_tail = _avg(_collect_latencies(reg_list, last_n=TAIL_TURNS))
|
|
stm_tail = _avg(_collect_latencies(results, last_n=TAIL_TURNS))
|
|
speedup = reg_tail / stm_tail if stm_tail > 0 else float("inf")
|
|
self.assertGreaterEqual(
|
|
speedup,
|
|
1.4,
|
|
f"streaming should be >=1.4x faster on last {TAIL_TURNS} turns "
|
|
f"(regular={reg_tail:.1f}ms, streaming={stm_tail:.1f}ms, speedup={speedup:.2f}x)",
|
|
)
|
|
|
|
def test_streaming_session_correctness(self):
|
|
"""Correctness test: bs=1, assert output equal + latency speedup."""
|
|
reg = self._run_concurrent_session(streaming=False, num_concurrent=1)
|
|
stm = self._run_concurrent_session(streaming=True, num_concurrent=1)
|
|
|
|
_print_mode_table(reg[0], label="correctness regular")
|
|
_print_mode_table(stm[0], label="correctness streaming")
|
|
|
|
reg_out = reg[0].outputs
|
|
stm_out = stm[0].outputs
|
|
mismatches = sum(1 for a, b in zip(reg_out, stm_out) if a != b)
|
|
self.assertEqual(
|
|
mismatches,
|
|
0,
|
|
f"regular vs streaming (bs=1): {mismatches}/{len(reg_out)} turns differ",
|
|
)
|
|
|
|
def test_streaming_session_random_lengths(self):
|
|
"""Stress test: bs=8, streaming only, random input/output lens."""
|
|
rng = random.Random(42)
|
|
gen_lens = [
|
|
rng.randint(*RANDOM_OUTPUT_LEN_RANGE) for _ in range(NUM_TURNS_RANDOM)
|
|
]
|
|
|
|
results = self._run_concurrent_session(
|
|
streaming=True,
|
|
num_turns=NUM_TURNS_RANDOM,
|
|
per_turn_gen_lens=gen_lens,
|
|
random_input_chunks=True,
|
|
rng=random.Random(42),
|
|
)
|
|
|
|
for i, r in enumerate(results):
|
|
self.assertEqual(
|
|
len(r.turns),
|
|
NUM_TURNS_RANDOM,
|
|
f"session {i}: expected {NUM_TURNS_RANDOM} turns, got {len(r.turns)}",
|
|
)
|
|
_print_mode_table(results[0], label="random streaming session 0")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|