404 lines
12 KiB
Python
404 lines
12 KiB
Python
"""
|
|
Benchmark: Streaming Session Inter-Turn Latency
|
|
|
|
Measures per-turn latency across three modes as context grows:
|
|
- no_session: re-send full context each turn (radix tree prefix match)
|
|
- regular_session: session append (radix tree insert + match)
|
|
- streaming_session: session append (O(1) KV direct transfer)
|
|
|
|
Each mode runs NUM_CONCURRENT parallel sessions, each doing NUM_TURNS sequential
|
|
requests (16 input / 8 output per turn).
|
|
|
|
Usage:
|
|
python -m pytest bench_session_latency.py -s
|
|
python -m unittest bench_session_latency.BenchSessionLatency.test_streaming_session
|
|
python -m unittest bench_session_latency.BenchSessionLatency
|
|
"""
|
|
|
|
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_MODEL_NAME_FOR_TEST,
|
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
DEFAULT_URL_FOR_TEST,
|
|
CustomTestCase,
|
|
popen_launch_server,
|
|
)
|
|
|
|
register_cuda_ci(est_time=100, suite="stage-b-test-large-1-gpu")
|
|
|
|
NUM_TURNS = 300
|
|
INPUT_LEN = 16
|
|
GEN_LEN = 8
|
|
NUM_CONCURRENT = 4
|
|
TAIL_TURNS = 10
|
|
SAMPLE_TURNS = 8
|
|
|
|
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 _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 runners (called by worker threads)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _run_one_no_session(
|
|
base_url: str, tokenizer, chunks: List[List[int]]
|
|
) -> ModeResult:
|
|
result = ModeResult(mode="no_session")
|
|
accumulated_ids: List[int] = []
|
|
|
|
for turn_idx, chunk_ids in enumerate(chunks):
|
|
accumulated_ids.extend(chunk_ids)
|
|
|
|
t0 = time.perf_counter()
|
|
response = _send_generate(
|
|
base_url,
|
|
{"input_ids": accumulated_ids, "sampling_params": SAMPLING_PARAMS},
|
|
)
|
|
client_lat = (time.perf_counter() - t0) * 1000
|
|
|
|
meta = response["meta_info"]
|
|
result.turns.append(
|
|
_record_turn(turn_idx, len(accumulated_ids), meta, client_lat)
|
|
)
|
|
result.outputs.append(response["text"])
|
|
|
|
output_ids = tokenizer.encode(response["text"])
|
|
if output_ids and output_ids[0] == tokenizer.bos_token_id:
|
|
output_ids = output_ids[1:]
|
|
accumulated_ids.extend(output_ids)
|
|
|
|
return result
|
|
|
|
|
|
def _run_one_session(
|
|
base_url: str, chunks: List[List[int]], streaming: bool = False
|
|
) -> ModeResult:
|
|
mode = "streaming_session" if streaming else "regular_session"
|
|
result = ModeResult(mode=mode)
|
|
|
|
capacity = sum(len(c) for c in chunks) + len(chunks) * GEN_LEN + 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)
|
|
|
|
t0 = time.perf_counter()
|
|
response = _send_generate(
|
|
base_url,
|
|
{
|
|
"input_ids": chunk_ids,
|
|
"session_params": {"id": session_id, "rid": rid},
|
|
"sampling_params": SAMPLING_PARAMS,
|
|
},
|
|
)
|
|
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 BenchSessionLatency(CustomTestCase):
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
|
|
cls.base_url = DEFAULT_URL_FOR_TEST
|
|
cls.process = popen_launch_server(
|
|
cls.model,
|
|
cls.base_url,
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=[
|
|
"--attention-backend",
|
|
"flashinfer",
|
|
"--enable-streaming-session",
|
|
],
|
|
)
|
|
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_no_session(self) -> List[ModeResult]:
|
|
requests.post(self.base_url + "/flush_cache")
|
|
|
|
def run_one(session_idx):
|
|
chunks = _generate_input_chunks(
|
|
self.tokenizer, NUM_TURNS, INPUT_LEN, offset=session_idx
|
|
)
|
|
return _run_one_no_session(self.base_url, self.tokenizer, chunks)
|
|
|
|
with ThreadPoolExecutor(max_workers=NUM_CONCURRENT) as pool:
|
|
return list(pool.map(run_one, range(NUM_CONCURRENT)))
|
|
|
|
def _run_concurrent_session(self, streaming: bool = False) -> List[ModeResult]:
|
|
requests.post(self.base_url + "/flush_cache")
|
|
|
|
def run_one(session_idx):
|
|
chunks = _generate_input_chunks(
|
|
self.tokenizer, NUM_TURNS, INPUT_LEN, offset=session_idx
|
|
)
|
|
return _run_one_session(self.base_url, chunks, streaming=streaming)
|
|
|
|
with ThreadPoolExecutor(max_workers=NUM_CONCURRENT) as pool:
|
|
return list(pool.map(run_one, range(NUM_CONCURRENT)))
|
|
|
|
# ------------------------------------------------------------------
|
|
# Test methods
|
|
# ------------------------------------------------------------------
|
|
|
|
def test_no_session(self):
|
|
results = self._run_concurrent_no_session()
|
|
self.__class__.all_results["no_session"] = results
|
|
_print_mode_table(results[0], label="session 0")
|
|
|
|
def test_regular_session(self):
|
|
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):
|
|
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_out = reg_list[0].outputs
|
|
stm_out = results[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 (session 0): {mismatches}/{len(reg_out)} turns differ",
|
|
)
|
|
|
|
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,
|
|
2.0,
|
|
f"streaming should be >=2x faster on last {TAIL_TURNS} turns "
|
|
f"(regular={reg_tail:.1f}ms, streaming={stm_tail:.1f}ms, speedup={speedup:.2f}x)",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|