[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,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()
|
||||
Reference in New Issue
Block a user