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

[Session] Add streaming mode with SessionAwareCache fast path (#19171)

Browse Source 
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
This commit is contained in:
Aurick Qiao
2026-02-27 16:31:08 -08:00
committed by GitHub
co-authored by hnyls2002
parent 2bd2b60b5c
commit c6cb0c9649
12 changed files with 999 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/sessions/test_session_control.py
+947
View File
@@ -0,0 +1,947 @@
"""
Usage:
python3 -m unittest test_session_control.TestSessionControl.test_session_control
python3 -m unittest test_session_control.TestSessionControl.test_session_control_with_branching
python3 -m unittest test_session_control.TestSessionControl.test_session_control_backtrack_with_abort
python3 -m unittest test_session_control.TestSessionControl.test_streaming_session
python3 -m unittest test_session_control.TestSessionControlVision.test_session_control
"""
import asyncio
import json
import time
import unittest
import aiohttp
import requests
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_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
register_cuda_ci(est_time=60, suite="stage-b-test-large-1-gpu")
def remove_prefix(text: str, prefix: str) -> str:
return text[len(prefix) :] if text.startswith(prefix) else text
class TestSessionControl(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_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",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_session_control(self, gen_len=12):
chunks = [
"Let me tell you something about France.",
"The capital of France is",
"The population of the city is",
"A brief history about that city is",
]
tokenizer = get_tokenizer(self.model)
chunks_ids = [tokenizer.encode(x) for x in chunks]
for i in range(1, len(chunks_ids)):
if chunks_ids[i][0] == tokenizer.bos_token_id:
chunks_ids[i] = chunks_ids[i][1:]
# 1. using session control
requests.post(self.base_url + "/flush_cache")
session_id = requests.post(
self.base_url + "/open_session",
json={"capacity_of_str_len": 1000},
).json()
rid = None
# open an existing session, should get session_id as None
ret = requests.post(
self.base_url + "/open_session",
json={"capacity_of_str_len": 1000, "session_id": session_id},
)
self.assertNotEqual(ret.status_code, 200)
first_rid = None
outputs_from_session = []
logprobs_from_session = []
cur_logprob_start_len = 0
for i, chunk_ids in enumerate(chunks_ids):
max_new_tokens = gen_len if i > 0 else 1 # prefill only for the first chunk
response = requests.post(
self.base_url + "/generate",
json={
"input_ids": chunk_ids,
"session_params": {
"id": session_id,
"rid": rid,
"offset": -1,
"replace": True,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": max_new_tokens,
"no_stop_trim": True,
"skip_special_tokens": False,
},
"return_logprob": True,
"logprob_start_len": cur_logprob_start_len - 1,
},
).json()
rid = response["meta_info"]["id"]
if i == 0:
first_rid = rid
if i > 0:
outputs_from_session.append(response["text"])
logprobs_from_session.extend(
[
round(sublist[0], 2)
for sublist in response["meta_info"]["output_token_logprobs"]
]
)
cur_logprob_start_len += len(chunk_ids) + max_new_tokens
# query with a logprob_start_len longer than the request, should see error
ret = requests.post(
self.base_url + "/generate",
json={
"input_ids": chunk_ids,
"session_params": {
"id": session_id,
"rid": rid,
"offset": -1,
"replace": True,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": max_new_tokens,
"no_stop_trim": True,
"skip_special_tokens": False,
},
"return_logprob": True,
"logprob_start_len": cur_logprob_start_len + len(chunk_ids),
},
)
self.assertNotEqual(ret.status_code, 200)
# backtrack to the first request and regenerate
cur_logprob_start_len = 0
response = requests.post(
self.base_url + "/generate",
json={
"input_ids": chunks_ids[-1],
"session_params": {
"id": session_id,
"rid": first_rid,
"offset": -1,
"replace": True,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
"return_logprob": True,
"logprob_start_len": cur_logprob_start_len,
},
).json()
outputs_from_session.append(response["text"])
logprobs_from_session.extend(
[
round(sublist[0], 2)
for sublist in response["meta_info"]["output_token_logprobs"]
]
)
# query with a non-existing rid (the last one should be disappeared because of backtrack), should see abort
ret = requests.post(
self.base_url + "/generate",
json={
"input_ids": chunks_ids[-1],
"session_params": {
"id": session_id,
"rid": rid,
"offset": -1,
"replace": True,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
"return_logprob": True,
},
)
self.assertNotEqual(ret.status_code, 200)
ret = requests.post(
self.base_url + "/close_session",
json={"session_id": session_id},
)
self.assertEqual(ret.status_code, 200)
# send a request to a closed session, should see abort
ret = requests.post(
self.base_url + "/generate",
json={
"input_ids": chunks_ids[-1],
"session_params": {
"id": session_id,
"rid": first_rid,
"offset": -1,
"replace": True,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
"return_logprob": True,
},
)
self.assertNotEqual(ret.status_code, 200)
# 2. not use session control
requests.post(self.base_url + "/flush_cache")
input_ids_first_req = None
input_ids = []
outputs_normal = []
logprobs_normal = []
for i, chunk_ids in enumerate(chunks_ids):
input_ids += chunk_ids
response = requests.post(
self.base_url + "/generate",
json={
"input_ids": input_ids,
"sampling_params": {
"temperature": 0,
"max_new_tokens": (
gen_len if i > 0 else 1
), # prefill only for the first chunk
"no_stop_trim": True,
"skip_special_tokens": False,
},
"return_logprob": True,
},
).json()
if i > 0:
output_ids = tokenizer.encode(response["text"])
if output_ids[0] == tokenizer.bos_token_id:
output_ids = output_ids[1:]
input_ids += output_ids[:-1]
outputs_normal.append(response["text"])
logprobs_normal.extend(
[
round(sublist[0], 2)
for sublist in response["meta_info"]["output_token_logprobs"]
]
)
if i == 0:
input_ids_first_req = input_ids.copy()
input_ids_first_req += chunks_ids[-1]
response = requests.post(
self.base_url + "/generate",
json={
"input_ids": input_ids_first_req,
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
"return_logprob": True,
},
).json()
outputs_normal.append(response["text"])
logprobs_normal.extend(
[
round(sublist[0], 2)
for sublist in response["meta_info"]["output_token_logprobs"]
]
)
print("outputs from chunked queries with session control:")
print(outputs_from_session)
print("outputs from normal queries:")
print(outputs_normal)
self.assertEqual(outputs_from_session, outputs_normal)
print("logprobs from chunked queries with session control:")
print(logprobs_from_session)
print("logprobs from normal queries:")
print(logprobs_normal)
assert len(logprobs_from_session) == len(
logprobs_normal
), "logprobs must have equal length"
for a, b in zip(logprobs_from_session, logprobs_normal):
assert abs(a - b) <= 0.15, f"logprobs {a} and {b} differ by more than 0.15"
async def async_generate(self, payload):
url = self.base_url + "/generate"
async with aiohttp.ClientSession() as session:
async with session.post(url=url, json=payload) as response:
assert response.status == 200
async for chunk_bytes in response.content:
chunk_bytes = chunk_bytes.strip()
if not chunk_bytes:
continue
chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
if chunk == "[DONE]":
yield "", None, ""
else:
data = json.loads(chunk)
finish_reason = (
data["meta_info"]["finish_reason"]["type"]
if data["meta_info"]["finish_reason"]
else ""
)
yield data["text"], data["meta_info"]["id"], finish_reason
async def run_session_control_backtrack_with_abort(self, replace):
chunks = [
"Let me tell you something about France.",
"The capital of France is",
]
tokenizer = get_tokenizer(self.model)
chunks_ids = [tokenizer.encode(x) for x in chunks]
for i in range(1, len(chunks_ids)):
if chunks_ids[i][0] == tokenizer.bos_token_id:
chunks_ids[i] = chunks_ids[i][1:]
# 1. using session control
requests.post(self.base_url + "/flush_cache")
session_id = requests.post(
self.base_url + "/open_session",
json={"capacity_of_str_len": 1000},
).json()
rid = None
payload = {
"input_ids": chunks_ids[0],
"session_params": {
"id": session_id,
"rid": rid,
"offset": -1,
"replace": True,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": 100,
"no_stop_trim": True,
"skip_special_tokens": False,
"ignore_eos": True,
},
"stream": True,
}
gen_so_far = ""
finish_reason = ""
second_output = ""
async for chunk, rid, finish_reason_chunk in self.async_generate(payload):
gen_so_far += chunk
if finish_reason == "":
finish_reason = finish_reason_chunk
if len(gen_so_far) > 50 and second_output == "":
payload2 = {
"input_ids": chunks_ids[1],
"session_params": {
"id": session_id,
"rid": rid,
"offset": 50,
"replace": replace,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": 32,
"no_stop_trim": True,
"skip_special_tokens": False,
},
"stream": False,
"stream_output": True,
}
response = requests.post(
url=self.base_url + "/generate", json=payload2
).json()
second_output = response["text"]
if replace:
assert finish_reason == "abort"
print("first request output:")
print(gen_so_far)
print("second request output:")
print(second_output)
# close the session
ret = requests.post(
self.base_url + "/close_session",
json={"session_id": session_id},
)
assert ret.status_code == 200
if not replace:
assert response["meta_info"]["finish_reason"]["type"] == "abort"
else:
# 2. not using session control
requests.post(self.base_url + "/flush_cache")
output_ids = tokenizer.encode(gen_so_far)
if output_ids[0] == tokenizer.bos_token_id:
output_ids = output_ids[1:]
input_ids = chunks_ids[0] + output_ids
input_ids = input_ids[:50] + chunks_ids[1]
payload = {
"input_ids": input_ids,
"sampling_params": {
"temperature": 0,
"max_new_tokens": 32,
"no_stop_trim": True,
"skip_special_tokens": False,
},
"stream": False,
"stream_output": True,
}
response = requests.post(
url=self.base_url + "/generate", json=payload
).json()
output_no_session = response["text"]
print("second request output without session:")
print(output_no_session)
assert (
second_output == output_no_session
), f"second_output: {second_output}, output_no_session: {output_no_session}"
@unittest.skip("broken")
def test_session_control_backtrack_with_abort(self):
asyncio.run(self.run_session_control_backtrack_with_abort(replace=True))
asyncio.run(self.run_session_control_backtrack_with_abort(replace=False))
def test_streaming_session(self, gen_len=12):
chunks = [
"Let me tell you something about France.",
"The capital of France is",
"The population of the city is",
]
tokenizer = get_tokenizer(self.model)
chunks_ids = [tokenizer.encode(x) for x in chunks]
for i in range(1, len(chunks_ids)):
if chunks_ids[i][0] == tokenizer.bos_token_id:
chunks_ids[i] = chunks_ids[i][1:]
# === Part 1: streaming session ===
requests.post(self.base_url + "/flush_cache")
session_id = requests.post(
self.base_url + "/open_session",
json={"capacity_of_str_len": 1000, "streaming": True},
).json()
rid = None
outputs_from_session = []
prev_kv_len = 0
for turn_idx, chunk_ids in enumerate(chunks_ids):
response = requests.post(
self.base_url + "/generate",
json={
"input_ids": chunk_ids,
"session_params": {"id": session_id, "rid": rid},
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
},
).json()
rid = response["meta_info"]["id"]
outputs_from_session.append(response["text"])
cached = response["meta_info"]["cached_tokens"]
prompt_tokens = response["meta_info"]["prompt_tokens"]
completion_tokens = response["meta_info"]["completion_tokens"]
if turn_idx == 0:
# Turn 1 should have no cache hit (cache was flushed).
self.assertEqual(
cached, 0, "Turn 1 should have 0 cached tokens (clean start)"
)
else:
# Turns 2+ inherit KV from the previous turn (via inherit_kv_states,
# not radix tree matching). cached_tokens reflects the inherited prefix.
self.assertEqual(
cached,
prev_kv_len,
f"Turn {turn_idx + 1}: should inherit {prev_kv_len} KV tokens from previous turn",
)
prev_kv_len = prompt_tokens + completion_tokens
# Close the session before checking cache/memory state.
ret = requests.post(
self.base_url + "/close_session",
json={"session_id": session_id},
)
self.assertEqual(ret.status_code, 200)
# === Cache verification (after close, before flush) ===
# Assertion 2: turn 1's prompt was inserted to the cache.
verify_resp = requests.post(
self.base_url + "/generate",
json={
"input_ids": chunks_ids[0],
"sampling_params": {"temperature": 0, "max_new_tokens": 1},
},
).json()
self.assertGreater(
verify_resp["meta_info"]["cached_tokens"],
0,
"Turn 1's prompt should be cached in the radix tree",
)
# Assertion 3 (insertion): turn 2's prompt tokens should NOT be in cache.
# The tree should only contain turn 1's extent (prompt + output from
# cache_unfinished_req during decode). Turn 2's prompt starts fresh tokens
# that were never inserted.
verify_resp2 = requests.post(
self.base_url + "/generate",
json={
"input_ids": chunks_ids[1],
"sampling_params": {"temperature": 0, "max_new_tokens": 1},
},
).json()
self.assertEqual(
verify_resp2["meta_info"]["cached_tokens"],
0,
"Turn 2's prompt should not be in cache (no insertion for turns 2+)",
)
# === Memory verification ===
# Assertion 4 & 5: KV is released properly and no memory leak.
# SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE is True by default;
# the scheduler will crash if it detects a leak during idle.
time.sleep(2)
health_resp = requests.get(self.base_url + "/health")
self.assertEqual(
health_resp.status_code,
200,
"Server should be healthy after session close (no memory leak)",
)
# After flush, all cache should be reclaimed.
requests.post(self.base_url + "/flush_cache")
verify_resp3 = requests.post(
self.base_url + "/generate",
json={
"input_ids": chunks_ids[0],
"sampling_params": {"temperature": 0, "max_new_tokens": 1},
},
).json()
self.assertEqual(
verify_resp3["meta_info"]["cached_tokens"],
0,
"After session close + flush, cache should be fully reclaimed",
)
# === Part 2: non-session baseline for output comparison ===
requests.post(self.base_url + "/flush_cache")
outputs_normal = []
input_ids = chunks_ids[0][:]
for i in range(len(chunks_ids)):
response = requests.post(
self.base_url + "/generate",
json={
"input_ids": input_ids,
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
},
).json()
outputs_normal.append(response["text"])
if i + 1 < len(chunks_ids):
out_ids = tokenizer.encode(response["text"])
if out_ids and out_ids[0] == tokenizer.bos_token_id:
out_ids = out_ids[1:]
input_ids = input_ids + out_ids + chunks_ids[i + 1]
print("outputs from streaming session:")
print(outputs_from_session)
print("outputs from normal queries:")
print(outputs_normal)
self.assertEqual(outputs_from_session, outputs_normal)
def run_session_control_with_branching(
self, root_prompt, chunks_per_step, gen_len=16
):
for x in chunks_per_step:
assert len(x) == len(chunks_per_step[0])
# 1. using session control
requests.post(self.base_url + "/flush_cache")
session_id = requests.post(
self.base_url + "/open_session",
json={"capacity_of_str_len": 1000},
).json()
outputs_from_session = []
# send the root prompt
response = requests.post(
self.base_url + "/generate",
json={
"text": root_prompt,
"session_params": {
"id": session_id,
"rid": None,
"offset": 0,
"replace": False,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
},
).json()
rid_per_branch = [response["meta_info"]["id"]] * len(chunks_per_step[0])
outputs_from_session.append(response["text"])
# send the prompts in branches
for chunks_for_branches in chunks_per_step:
for j, chunk in enumerate(chunks_for_branches):
response = requests.post(
self.base_url + "/generate",
json={
"text": chunk,
"session_params": {
"id": session_id,
"rid": rid_per_branch[j],
"offset": 0,
"replace": False,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
},
).json()
rid = response["meta_info"]["id"]
rid_per_branch[j] = rid
outputs_from_session.append(response["text"])
# close the session
ret = requests.post(
self.base_url + "/close_session",
json={"session_id": session_id},
)
assert ret.status_code == 200
# 2. not use session control
requests.post(self.base_url + "/flush_cache")
outputs_normal = []
input_texts = [root_prompt] * len(chunks_per_step[0])
# send the root prompt
response = requests.post(
self.base_url + "/generate",
json={
"text": root_prompt,
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
},
).json()
outputs_normal.append(response["text"])
input_texts = [x + response["text"] for x in input_texts]
# send the prompts in branches
for chunks_for_branches in chunks_per_step:
for j, chunk in enumerate(chunks_for_branches):
input_texts[j] += chunk
response = requests.post(
self.base_url + "/generate",
json={
"text": input_texts[j],
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
},
).json()
outputs_normal.append(response["text"])
input_texts[j] += response["text"]
print("====== outputs from chunked queries with session control: =======")
print(outputs_from_session)
print("====== outputs from normal queries: =======")
print(outputs_normal)
assert (
outputs_from_session == outputs_normal
), f"outputs_from_session: {outputs_from_session}, outputs_normal: {outputs_normal}"
def test_session_control_with_branching(self):
root_prompt = "First, let me explain in one sentence about AI"
chunks_per_step = [
[
"Then, briefly, the positive side of AI is",
"But, briefly, AI could be harmful to human",
],
["For example", "For example"],
]
self.run_session_control_with_branching(
root_prompt=root_prompt, chunks_per_step=chunks_per_step, gen_len=8
)
root_prompt = "I have three apples."
chunks_per_step = [
["I then give one apple to my friend", "My friend give me another apple."],
["I still have", "I now have"],
]
self.run_session_control_with_branching(
root_prompt=root_prompt, chunks_per_step=chunks_per_step, gen_len=8
)
@unittest.skip("broken")
class TestSessionControlVision(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = "lmms-lab/llava-onevision-qwen2-7b-ov"
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={"--disable-radix"},
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_session_control(self):
text_chunks = [
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n",
"<|im_start|>user\n<image>\nDescribe this image in a very short sentence.<|im_end|>\n<|im_start|>assistant\n",
"<|im_start|>user\n<image>\nIs this image same with one of the previous images?<|im_end|>\n<|im_start|>assistant\n",
"<|im_start|>user\n<image>\nIs this image same with one of the previous images?<|im_end|>\n<|im_start|>assistant\n",
"<|im_start|>user\nDescribe this image in a very short sentence.<|im_end|>\nassistant:",
]
image_chunks = [
"https://raw.githubusercontent.com/sgl-project/sglang/main/examples/assets/example_image.png",
"https://raw.githubusercontent.com/sgl-project/sglang/main/examples/assets/example_image.png",
"https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png",
]
self.assertEqual(
len(text_chunks), len(image_chunks) + 2
) # the first and the last prompt does not contain images
tokenizer = get_tokenizer(self.model)
text_input_ids = [tokenizer.encode(x) for x in text_chunks]
for i in range(1, len(text_input_ids)):
if text_input_ids[i][0] == tokenizer.bos_token_id:
text_input_ids[i] = text_input_ids[i][1:]
gen_len = 32
# 1. using session control
requests.post(self.base_url + "/flush_cache")
session_id = requests.post(
self.base_url + "/open_session",
json={"capacity_of_str_len": 1000},
).json()
rid = None
# open an existing session, should get session_id as None
ret = requests.post(
self.base_url + "/open_session",
json={"capacity_of_str_len": 1000, "session_id": session_id},
)
self.assertNotEqual(ret.status_code, 200)
first_rid = None
outputs_from_session = []
for i in range(len(text_input_ids[:-1])):
response = requests.post(
self.base_url + "/generate",
json={
"input_ids": text_input_ids[i],
"image_data": image_chunks[i - 1] if i > 0 else None,
"modalities": ["multi-images"],
"session_params": {
"id": session_id,
"rid": rid,
"offset": 0,
"replace": True,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": (
gen_len if i > 0 else 0
), # prefill only for the first chunk
"no_stop_trim": True,
"skip_special_tokens": False,
},
},
).json()
rid = response["meta_info"]["id"]
if i == 0:
first_rid = rid
if i > 0:
outputs_from_session.append(response["text"])
# backtrack to the first request and regenerate
response = requests.post(
self.base_url + "/generate",
json={
"input_ids": text_input_ids[-1],
"session_params": {
"id": session_id,
"rid": first_rid,
"offset": 0,
"replace": True,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
},
).json()
outputs_from_session.append(response["text"])
# query with a non-existing rid (the last one should be disappeared because of backtrack), should see abort
ret = requests.post(
self.base_url + "/generate",
json={
"input_ids": text_input_ids[-1],
"session_params": {
"id": session_id,
"rid": rid,
"offset": 0,
"replace": True,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
},
)
self.assertNotEqual(ret.status_code, 200)
ret = requests.post(
self.base_url + "/close_session",
json={"session_id": session_id},
)
self.assertEqual(ret.status_code, 200)
# send a request to a closed session, should see abort
ret = requests.post(
self.base_url + "/generate",
json={
"input_ids": text_input_ids[-1],
"session_params": {
"id": session_id,
"rid": first_rid,
"offset": 0,
"replace": True,
},
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
},
)
self.assertNotEqual(ret.status_code, 200)
# 2. not use session control
requests.post(self.base_url + "/flush_cache")
input_ids_first_req = None
input_ids = []
outputs_normal = []
for i in range(len(text_input_ids[:-1])):
input_ids += text_input_ids[i]
image_data = image_chunks[:i] if i > 0 else None
response = requests.post(
self.base_url + "/generate",
json={
"input_ids": input_ids,
"image_data": image_data,
"modalities": ["multi-images"],
"sampling_params": {
"temperature": 0,
"max_new_tokens": (
gen_len if i > 0 else 0
), # prefill only for the first chunk
"no_stop_trim": True,
"skip_special_tokens": False,
},
},
).json()
if i > 0:
output_ids = tokenizer.encode(response["text"])
if output_ids[0] == tokenizer.bos_token_id:
output_ids = output_ids[1:]
input_ids += output_ids
outputs_normal.append(response["text"])
if i == 0:
input_ids_first_req = input_ids.copy()
input_ids_first_req += text_input_ids[-1]
response = requests.post(
self.base_url + "/generate",
json={
"input_ids": input_ids_first_req,
"sampling_params": {
"temperature": 0,
"max_new_tokens": gen_len,
"no_stop_trim": True,
"skip_special_tokens": False,
},
},
).json()
outputs_normal.append(response["text"])
print("outputs from chunked queries with session control:")
print(outputs_from_session)
print("outputs from normal queries:")
print(outputs_normal)
assert (
outputs_from_session == outputs_normal
), f"outputs_from_session: {outputs_from_session}, outputs_normal: {outputs_normal}"
if __name__ == "__main__":
unittest.main()
test/registered/sessions/test_session_latency.py
+399
View File
@@ -0,0 +1,399 @@
"""
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"],
)
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()