[Feature] Beam search support (#31626)
Co-authored-by: cswuyg <cswuyg@gmail.com> Co-authored-by: cswuyg <496090217@qq.com> Co-authored-by: Vedant Jhaveri <vedantjh2@gmail.com> Co-authored-by: Vedant Jhaveri <vjhaveri@linkedin.com>
This commit is contained in:
Liangsheng Yin
co-authored by
cswuyg
cswuyg
Vedant Jhaveri
Vedant Jhaveri
cswuyg
cswuyg
Vedant Jhaveri
Vedant Jhaveri
parent
e5a1c5a423
commit
ec4bdbfa4a
@@ -0,0 +1,122 @@
|
||||
"""Beam search parity acceptance test (executable API spec).
|
||||
|
||||
- Trigger: sampling_params.beam_width = k (> 1); no server-level beam flag.
|
||||
- Response: one response per rid; meta_info.beam_results holds the top-n
|
||||
sequences (n <= beam_width, default 1 as in HF/OpenAI), best score first.
|
||||
- Acceptance: sequence-set overlap vs HF transformers >= 0.8 for k in {2, 10}.
|
||||
|
||||
Manual test (GPU host): python3 test_beam_parity.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
from typing import List
|
||||
|
||||
import requests
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
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,
|
||||
)
|
||||
|
||||
PROMPT = "Hello SGLang"
|
||||
MAX_NEW_TOKENS = 10
|
||||
OVERLAP_THRESHOLD = 0.8
|
||||
|
||||
|
||||
def get_transformers_beam_sequences(
|
||||
model_path: str, prompt: str, beam_width: int, max_new_tokens: int
|
||||
) -> List[str]:
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||||
model = AutoModelForCausalLM.from_pretrained(model_path, dtype="auto")
|
||||
model = model.to("cuda" if torch.cuda.is_available() else "cpu")
|
||||
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
|
||||
input_length = inputs["input_ids"].shape[1]
|
||||
|
||||
with torch.no_grad():
|
||||
generated = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=max_new_tokens,
|
||||
num_beams=beam_width,
|
||||
num_return_sequences=beam_width,
|
||||
do_sample=False,
|
||||
)
|
||||
sequences = [
|
||||
tokenizer.decode(seq[input_length:].cpu().tolist(), skip_special_tokens=True)
|
||||
for seq in generated
|
||||
]
|
||||
|
||||
del model
|
||||
torch.cuda.empty_cache()
|
||||
return sequences
|
||||
|
||||
|
||||
class TestBeamParity(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = os.environ.get(
|
||||
"SGLANG_TEST_BEAM_MODEL", DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
)
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
# No beam-specific server flag. Overlap is pinned off so a parity
|
||||
# mismatch can only come from the search, not from scheduling.
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=["--disable-overlap-schedule"],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def _generate_beams(self, beam_width, n=None):
|
||||
sampling_params = {"beam_width": beam_width, "max_new_tokens": MAX_NEW_TOKENS}
|
||||
if n is not None:
|
||||
sampling_params["n"] = n
|
||||
resp = requests.post(
|
||||
f"{self.base_url}/generate",
|
||||
json={"text": PROMPT, "sampling_params": sampling_params},
|
||||
timeout=120,
|
||||
)
|
||||
self.assertEqual(resp.status_code, 200, resp.text)
|
||||
beam_results = resp.json().get("meta_info", {}).get("beam_results")
|
||||
self.assertIsNotNone(beam_results, "response carries no beam_results")
|
||||
return beam_results
|
||||
|
||||
def test_parity_vs_transformers(self):
|
||||
for beam_width in [2, 10]:
|
||||
# n=beam_width to compare the whole beam set; the API default is 1.
|
||||
beam_results = self._generate_beams(beam_width, n=beam_width)
|
||||
self.assertEqual(len(beam_results), beam_width)
|
||||
|
||||
scores = [r["meta_info"]["sequence_score"] for r in beam_results]
|
||||
self.assertEqual(scores, sorted(scores, reverse=True))
|
||||
|
||||
sglang_sequences = {r["text"] for r in beam_results}
|
||||
hf_sequences = set(
|
||||
get_transformers_beam_sequences(
|
||||
self.model, PROMPT, beam_width, MAX_NEW_TOKENS
|
||||
)
|
||||
)
|
||||
overlap = len(sglang_sequences & hf_sequences) / beam_width
|
||||
print(f"beam_width={beam_width} overlap={overlap:.2%}")
|
||||
self.assertGreaterEqual(overlap, OVERLAP_THRESHOLD)
|
||||
|
||||
def test_return_top_n(self):
|
||||
beam_results = self._generate_beams(beam_width=10, n=3)
|
||||
self.assertEqual(len(beam_results), 3)
|
||||
|
||||
def test_default_returns_one_sequence(self):
|
||||
self.assertEqual(len(self._generate_beams(beam_width=10)), 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Beam search load and admission-saturation tests.
|
||||
|
||||
- Mixed-width load: 100 requests at 10 QPS, widths in [2, 100]; expect
|
||||
100/100 OK, report p50/p90/p99 latency.
|
||||
- Extreme fanout: beam_width=3200, so each request owns 3200 req-to-token
|
||||
slots and the admission gate serializes them. Phase 1 measures
|
||||
single-inflight service time (arrival-rate percentiles sit on the queueing
|
||||
knee and are not usable as an SLO); phase 2 drives 0.8x the measured
|
||||
capacity and expects a stable queue.
|
||||
|
||||
Manual test (GPU host): python3 test_beam_search_load.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
MAX_NEW_TOKENS = 10
|
||||
CLIENT_TIMEOUT_S = 600
|
||||
PROMPT = "Write a short story about a robot learning to paint."
|
||||
|
||||
|
||||
async def _generate(session, base_url, width):
|
||||
start = time.perf_counter()
|
||||
async with session.post(
|
||||
f"{base_url}/generate",
|
||||
json={
|
||||
"text": PROMPT,
|
||||
"sampling_params": {"beam_width": width, "max_new_tokens": MAX_NEW_TOKENS},
|
||||
},
|
||||
) as resp:
|
||||
payload = await resp.json()
|
||||
latency = time.perf_counter() - start
|
||||
beam_results = payload.get("meta_info", {}).get("beam_results") or []
|
||||
return resp.status, len(beam_results), latency
|
||||
|
||||
|
||||
async def _run_at_qps(base_url, widths, qps):
|
||||
"""Fire one request per width at a fixed rate; return per-request results."""
|
||||
timeout = aiohttp.ClientTimeout(total=CLIENT_TIMEOUT_S)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
|
||||
async def delayed(i, width):
|
||||
await asyncio.sleep(i / qps)
|
||||
return await _generate(session, base_url, width)
|
||||
|
||||
return await asyncio.gather(
|
||||
*[delayed(i, width) for i, width in enumerate(widths)]
|
||||
)
|
||||
|
||||
|
||||
def _report_latencies(name, results):
|
||||
latencies_ms = [lat * 1000 for _, _, lat in results]
|
||||
p50, p90, p99 = np.percentile(latencies_ms, [50, 90, 99])
|
||||
print(f"{name}: n={len(results)} p50/p90/p99 = {p50:.0f}/{p90:.0f}/{p99:.0f} ms")
|
||||
|
||||
|
||||
class _BeamLoadTestBase(CustomTestCase):
|
||||
extra_server_args = []
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = os.environ.get("SGLANG_TEST_BEAM_LOAD_MODEL", "Qwen/Qwen2.5-0.5B")
|
||||
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-overlap-schedule", "--disable-radix-cache"]
|
||||
+ cls.extra_server_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def _check_all_ok(self, results, widths):
|
||||
num_ok = sum(1 for status, _, _ in results if status == 200)
|
||||
self.assertEqual(num_ok, len(results), f"{len(results) - num_ok} failed")
|
||||
for (_, num_beams, _), width in zip(results, widths):
|
||||
self.assertGreaterEqual(num_beams, 1)
|
||||
self.assertLessEqual(num_beams, width)
|
||||
|
||||
|
||||
class TestBeamSearchMixedWidthLoad(_BeamLoadTestBase):
|
||||
"""100 requests at 10 QPS with beam widths mixed in [2, 100]."""
|
||||
|
||||
def test_mixed_width_load(self):
|
||||
rng = random.Random(42)
|
||||
widths = [rng.randint(2, 100) for _ in range(100)]
|
||||
results = asyncio.run(_run_at_qps(self.base_url, widths, qps=10))
|
||||
self._check_all_ok(results, widths)
|
||||
_report_latencies("mixed-width 100 reqs @ 10 QPS", results)
|
||||
|
||||
|
||||
class TestBeamMixedWithNormalTraffic(_BeamLoadTestBase):
|
||||
"""Beam and normal requests finishing in shared batches: the carrier must
|
||||
stay index-aligned across IPC and normal outputs must keep their text."""
|
||||
|
||||
def test_mixed_traffic(self):
|
||||
async def run():
|
||||
timeout = aiohttp.ClientTimeout(total=CLIENT_TIMEOUT_S)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
|
||||
async def normal():
|
||||
async with session.post(
|
||||
f"{self.base_url}/generate",
|
||||
json={
|
||||
"text": PROMPT,
|
||||
"sampling_params": {"max_new_tokens": MAX_NEW_TOKENS},
|
||||
},
|
||||
) as resp:
|
||||
return resp.status, await resp.json()
|
||||
|
||||
return await asyncio.gather(
|
||||
*[_generate(session, self.base_url, 4) for _ in range(10)],
|
||||
*[normal() for _ in range(10)],
|
||||
)
|
||||
|
||||
results = asyncio.run(run())
|
||||
beam_results, normal_results = results[:10], results[10:]
|
||||
self._check_all_ok(beam_results, [4] * 10)
|
||||
for status, payload in normal_results:
|
||||
self.assertEqual(status, 200)
|
||||
self.assertTrue(
|
||||
payload["text"], "normal request lost its text in a mixed batch"
|
||||
)
|
||||
self.assertNotIn("beam_results", payload["meta_info"])
|
||||
|
||||
|
||||
class TestBeamSearchExtremeFanout(_BeamLoadTestBase):
|
||||
"""beam_width=3200: measure single-group service time, then 0.8x-capacity load."""
|
||||
|
||||
# Each beam_width=3200 request owns 3200 req-to-token slots; make the pool
|
||||
# size deterministic so exactly one group fits at a time.
|
||||
extra_server_args = ["--max-running-requests", "4000"]
|
||||
|
||||
WIDTH = 3200
|
||||
|
||||
async def _run_sequential(self, num_requests):
|
||||
timeout = aiohttp.ClientTimeout(total=CLIENT_TIMEOUT_S)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
return [
|
||||
await _generate(session, self.base_url, self.WIDTH)
|
||||
for _ in range(num_requests)
|
||||
]
|
||||
|
||||
def test_extreme_fanout(self):
|
||||
# Phase 1: single-inflight service time (first request dropped as warmup).
|
||||
results = asyncio.run(self._run_sequential(6))
|
||||
self._check_all_ok(results, [self.WIDTH] * 6)
|
||||
service_samples = [lat for _, _, lat in results[1:]]
|
||||
service_s = sum(service_samples) / len(service_samples)
|
||||
print(
|
||||
f"extreme fanout n={self.WIDTH} single-inflight service: "
|
||||
f"{service_s * 1000:.0f} ms/group"
|
||||
)
|
||||
|
||||
# Phase 2: stable-queue load at 0.8x the measured serial capacity.
|
||||
qps = 0.8 / service_s
|
||||
num_requests = max(10, int(30 * qps))
|
||||
widths = [self.WIDTH] * num_requests
|
||||
results = asyncio.run(_run_at_qps(self.base_url, widths, qps=qps))
|
||||
self._check_all_ok(results, widths)
|
||||
_report_latencies(
|
||||
f"extreme fanout n={self.WIDTH} @ 0.8x capacity ({qps:.2f} QPS)", results
|
||||
)
|
||||
|
||||
# Server must still be alive and serving after the burst.
|
||||
final = asyncio.run(_run_at_qps(self.base_url, [2], qps=1))
|
||||
self._check_all_ok(final, [2])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,149 @@
|
||||
"""Beam width sweep benchmark: 100 concurrent ShareGPT prompts
|
||||
(prompt_len < 100), max_new_tokens=10, widths 10/50/100/200/400.
|
||||
Primary metric is aggregate beam tok/s.
|
||||
|
||||
Manual test (GPU host): python3 test_beam_search_perf_sweep.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import aiohttp
|
||||
|
||||
from sglang.benchmark.datasets.sharegpt import sample_sharegpt_requests
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
BEAM_WIDTHS = [10, 50, 100, 200, 400]
|
||||
NUM_PROMPTS = 100
|
||||
MAX_PROMPT_LEN = 100
|
||||
MAX_NEW_TOKENS = 10
|
||||
CLIENT_TIMEOUT_S = 1200
|
||||
|
||||
|
||||
async def _generate(session, base_url, prompt, width):
|
||||
start = time.perf_counter()
|
||||
async with session.post(
|
||||
f"{base_url}/generate",
|
||||
json={
|
||||
"text": prompt,
|
||||
"sampling_params": {"beam_width": width, "max_new_tokens": MAX_NEW_TOKENS},
|
||||
},
|
||||
) as resp:
|
||||
payload = await resp.json()
|
||||
latency = time.perf_counter() - start
|
||||
beam_results = payload.get("meta_info", {}).get("beam_results") or []
|
||||
return resp.status, len(beam_results), latency
|
||||
|
||||
|
||||
class _BeamSweepBase(CustomTestCase):
|
||||
# Primary metric is aggregate beam tok/s (reqs x width x new_tokens /
|
||||
# elapsed); QPS is secondary since it conflates width.
|
||||
|
||||
extra_server_args = []
|
||||
pool_label = "default pool"
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = os.environ.get("SGLANG_TEST_BEAM_MODEL", "Qwen/Qwen3-1.7B")
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
# 0.7 leaves headroom for the full-vocab [num_beam_rows, vocab]
|
||||
# logprobs tensor, which OOMs at large width x concurrency.
|
||||
other_args=["--disable-overlap-schedule", "--mem-fraction-static", "0.7"]
|
||||
+ cls.extra_server_args,
|
||||
)
|
||||
|
||||
tokenizer = get_tokenizer(cls.model)
|
||||
rows = sample_sharegpt_requests(
|
||||
dataset_path="", num_requests=4000, tokenizer=tokenizer
|
||||
)
|
||||
cls.prompts = [r.prompt for r in rows if r.prompt_len < MAX_PROMPT_LEN][
|
||||
:NUM_PROMPTS
|
||||
]
|
||||
assert (
|
||||
len(cls.prompts) == NUM_PROMPTS
|
||||
), f"only {len(cls.prompts)} short prompts sampled"
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
async def _run_one_width(self, width):
|
||||
timeout = aiohttp.ClientTimeout(total=CLIENT_TIMEOUT_S)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
start = time.perf_counter()
|
||||
results = await asyncio.gather(
|
||||
*[
|
||||
_generate(session, self.base_url, prompt, width)
|
||||
for prompt in self.prompts
|
||||
]
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
return results, elapsed
|
||||
|
||||
def _run_sweep(self):
|
||||
report = []
|
||||
for width in BEAM_WIDTHS:
|
||||
results, elapsed = asyncio.run(self._run_one_width(width))
|
||||
|
||||
num_ok = sum(1 for status, _, _ in results if status == 200)
|
||||
self.assertEqual(
|
||||
num_ok, NUM_PROMPTS, f"width={width}: {NUM_PROMPTS - num_ok} failed"
|
||||
)
|
||||
for status, num_beams, _ in results:
|
||||
self.assertGreaterEqual(num_beams, 1)
|
||||
self.assertLessEqual(num_beams, width)
|
||||
|
||||
beam_tok_s = NUM_PROMPTS * width * MAX_NEW_TOKENS / elapsed
|
||||
qps = NUM_PROMPTS / elapsed
|
||||
report.append((width, beam_tok_s, qps, elapsed))
|
||||
print(
|
||||
f"width={width:4d} beam_tok/s={beam_tok_s:9.0f} "
|
||||
f"qps={qps:6.2f} elapsed={elapsed:6.2f}s"
|
||||
)
|
||||
|
||||
print(f"\nBeam width sweep ({self.pool_label}):")
|
||||
print("| beam width | beam tok/s | qps | elapsed (s) |")
|
||||
print("|---|---|---|---|")
|
||||
for width, beam_tok_s, qps, elapsed in report:
|
||||
print(f"| {width} | {beam_tok_s:.0f} | {qps:.2f} | {elapsed:.2f} |")
|
||||
|
||||
|
||||
class TestBeamSweepDefaultPool(_BeamSweepBase):
|
||||
"""Default req-slot pool (4096): beam rows pin at ~4000 for width >= 50, so
|
||||
this curve saturates at the pool, not the engine."""
|
||||
|
||||
def test_beam_width_sweep(self):
|
||||
self._run_sweep()
|
||||
|
||||
|
||||
class TestBeamSweepLargePool(_BeamSweepBase):
|
||||
"""Enlarged pool (16384) for the engine ceiling: ~40 width-400 groups run
|
||||
concurrently, affordable because --context-length is short."""
|
||||
|
||||
extra_server_args = [
|
||||
"--max-running-requests",
|
||||
"16384",
|
||||
"--context-length",
|
||||
"2048",
|
||||
]
|
||||
pool_label = "large pool (16384 slots)"
|
||||
|
||||
def test_beam_width_sweep(self):
|
||||
self._run_sweep()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user