Test reorganization: Move tests to manual/ (#13610)
This commit is contained in:
@@ -0,0 +1,104 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest test_ascend_w8a8_quantization.TestAscendW8A8.test_gsm8k
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
if "ASCEND_RT_VISIBLE_DEVICES" not in os.environ:
|
||||
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = "0,1"
|
||||
DEFAULT_PORT_FOR_SRT_TEST_RUNNER = (
|
||||
7000 + int(os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "0")[0]) * 100
|
||||
)
|
||||
DEFAULT_URL_FOR_TEST = f"http://127.0.0.1:{DEFAULT_PORT_FOR_SRT_TEST_RUNNER + 1000}"
|
||||
|
||||
|
||||
class TestAscendW8A8(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "vllm-ascend/Qwen2.5-0.5B-Instruct-w8a8"
|
||||
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=[
|
||||
"--trust-remote-code",
|
||||
"--disable-cuda-graph",
|
||||
"--device",
|
||||
"npu",
|
||||
"--attention-backend",
|
||||
"ascend",
|
||||
"--quantization",
|
||||
"w8a8_int8",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
url = urlparse(base_url)
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host=f"http://{url.hostname}",
|
||||
port=int(url.port),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(metrics)
|
||||
|
||||
self.assertGreaterEqual(metrics["accuracy"], 0.25)
|
||||
self.assertGreaterEqual(metrics["output_throughput"], 1000)
|
||||
|
||||
def run_decode(self, max_new_tokens):
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": "The capital of France is",
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": max_new_tokens,
|
||||
},
|
||||
"ignore_eos": True,
|
||||
},
|
||||
)
|
||||
return response.json()
|
||||
|
||||
def test_throughput(self):
|
||||
max_tokens = 256
|
||||
|
||||
tic = time.perf_counter()
|
||||
res = self.run_decode(max_tokens)
|
||||
tok = time.perf_counter()
|
||||
print(res["text"])
|
||||
throughput = max_tokens / (tok - tic)
|
||||
print(f"Throughput: {throughput} tokens/s")
|
||||
|
||||
if is_in_ci():
|
||||
self.assertGreaterEqual(throughput, 25)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest test_mindspore_models.TestMindSporeQwen3.test_gsm8k
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestMindSporeQwen3(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "Qwen/Qwen3-8B"
|
||||
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=[
|
||||
"--device",
|
||||
"npu",
|
||||
"--model-impl",
|
||||
"mindspore",
|
||||
"--attention-backend",
|
||||
"ascend",
|
||||
"--tp-size",
|
||||
"1",
|
||||
"--dp-size",
|
||||
"1",
|
||||
"--mem-fraction-static",
|
||||
0.8,
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.78)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,116 @@
|
||||
import copy
|
||||
import multiprocessing
|
||||
import os
|
||||
import traceback
|
||||
import unittest
|
||||
from multiprocessing import Process
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
from sglang.test.test_utils import CustomTestCase, find_available_port
|
||||
|
||||
|
||||
def run_distributed_test(rank, world_size, master_port, output_writer, fn):
|
||||
try:
|
||||
os.environ["RANK"] = str(rank)
|
||||
os.environ["WORLD_SIZE"] = str(world_size)
|
||||
os.environ["MASTER_ADDR"] = "localhost"
|
||||
os.environ["MASTER_PORT"] = str(master_port)
|
||||
os.environ["LOCAL_SIZE"] = str(world_size)
|
||||
|
||||
dist.init_process_group("gloo", rank=rank, world_size=world_size)
|
||||
torch.ops.sgl_kernel.initialize(world_size, rank)
|
||||
|
||||
fn(rank, world_size)
|
||||
|
||||
execution_ok = True
|
||||
except Exception as e:
|
||||
print(f"subprocess[{rank=}] has error: {e}", flush=True)
|
||||
traceback.print_exc()
|
||||
execution_ok = False
|
||||
|
||||
output_writer.send(execution_ok)
|
||||
output_writer.close()
|
||||
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def all_reduce_fn(rank, world_size):
|
||||
op = dist.ReduceOp.SUM
|
||||
for dtype in [torch.float32, torch.bfloat16, torch.float16]:
|
||||
tensor = torch.randn(2, 10, dtype=dtype)
|
||||
tensor_shm = copy.deepcopy(tensor)
|
||||
|
||||
dist.all_reduce(tensor, op=op)
|
||||
torch.ops.sgl_kernel.shm_allreduce(tensor_shm, op)
|
||||
|
||||
torch.testing.assert_close(tensor, tensor_shm)
|
||||
|
||||
|
||||
def all_gather_fn(rank, world_size):
|
||||
dim = -1
|
||||
|
||||
for dtype in [torch.float32, torch.bfloat16, torch.float16]:
|
||||
tensor = torch.randn(2, 10, dtype=dtype)
|
||||
|
||||
if dim < 0:
|
||||
# Convert negative dim to positive.
|
||||
dim += tensor.dim()
|
||||
|
||||
input_size = tensor.size()
|
||||
output_size = (input_size[0] * world_size,) + input_size[1:]
|
||||
output_tensor = torch.empty(
|
||||
output_size, dtype=tensor.dtype, device=tensor.device
|
||||
)
|
||||
dist.all_gather_into_tensor(output_tensor, tensor)
|
||||
output_tensor = output_tensor.reshape((world_size,) + input_size)
|
||||
output_tensor = output_tensor.movedim(0, dim)
|
||||
output_tensor = output_tensor.reshape(
|
||||
input_size[:dim] + (world_size * input_size[dim],) + input_size[dim + 1 :]
|
||||
)
|
||||
|
||||
output_shm = torch.ops.sgl_kernel.shm_allgather(tensor, dim)
|
||||
|
||||
torch.testing.assert_close(output_tensor, output_shm)
|
||||
|
||||
|
||||
class TestComm(CustomTestCase):
|
||||
def _spawn_and_check(self, fn, world_size=2):
|
||||
mp.set_start_method("spawn", force=True)
|
||||
master_port = find_available_port(23456)
|
||||
|
||||
processes = []
|
||||
output_reader, output_writer = multiprocessing.Pipe(duplex=False)
|
||||
|
||||
for rank in range(world_size):
|
||||
p = Process(
|
||||
target=run_distributed_test,
|
||||
kwargs=dict(
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
master_port=master_port,
|
||||
output_writer=output_writer,
|
||||
fn=fn,
|
||||
),
|
||||
)
|
||||
p.start()
|
||||
processes.append(p)
|
||||
|
||||
for _ in range(world_size):
|
||||
self.assertTrue(output_reader.recv(), "Subprocess fail. Check logs above.")
|
||||
|
||||
for p in processes:
|
||||
p.join()
|
||||
|
||||
def test_all_reduce(self):
|
||||
self._spawn_and_check(all_reduce_fn)
|
||||
|
||||
def test_all_gather(self):
|
||||
self._spawn_and_check(all_gather_fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,28 @@
|
||||
import json
|
||||
import unittest
|
||||
|
||||
from sglang.srt.debug_utils import log_parser
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestLogParser(CustomTestCase):
|
||||
def test_log_parser(self):
|
||||
lines = """
|
||||
(SGLangEngine pid=35555) [2025-10-31 03:45:20 TP0] Decode batch [51341], #running-req: 317, #token: 1094261, token usage: 0.67, cuda graph: True, gen throughput (token/s): 14806.57, #queue-req: 0,
|
||||
(SGLangEngine pid=111711, ip=10.15.36.1) [2025-10-31 03:45:20 TP0] Decode batch [39913], #running-req: 78, #token: 432100, token usage: 0.27, cuda graph: True, gen throughput (token/s): 7269.16, #queue-req: 0,
|
||||
[2025-11-03 14:31:10 DP6 TP6 EP6] Decode batch, #running-req: 251, #token: 2811200, token usage: 1.00, cuda graph: True, gen throughput (token/s): 2055.94, #queue-req: 655,
|
||||
"""
|
||||
expect_rows = json.loads(
|
||||
"""[{"line":"(SGLangEngine pid=35555) [2025-10-31 03:45:20 TP0] Decode batch [51341], #running-req: 317, #token: 1094261, token usage: 0.67, cuda graph: True, gen throughput (token/s): 14806.57, #queue-req: 0,","1":"(SGLangEngine pid=35555)","pid":35555,"ip":null,"time":"2025-10-31 03:45:20","dp_rank":null,"tp_rank":0,"ep_rank":null,"pp_rank":null,"9":" [51341]","num_running_req":317,"num_token":1094261,"token_usage":0.67,"gen_throughput":14806.57,"queue_req":0},{"line":"(SGLangEngine pid=111711, ip=10.15.36.1) [2025-10-31 03:45:20 TP0] Decode batch [39913], #running-req: 78, #token: 432100, token usage: 0.27, cuda graph: True, gen throughput (token/s): 7269.16, #queue-req: 0,","1":"(SGLangEngine pid=111711, ip=10.15.36.1)","pid":111711,"ip":"10.15.36.1","time":"2025-10-31 03:45:20","dp_rank":null,"tp_rank":0,"ep_rank":null,"pp_rank":null,"9":" [39913]","num_running_req":78,"num_token":432100,"token_usage":0.27,"gen_throughput":7269.16,"queue_req":0},{"line":"[2025-11-03 14:31:10 DP6 TP6 EP6] Decode batch, #running-req: 251, #token: 2811200, token usage: 1.00, cuda graph: True, gen throughput (token/s): 2055.94, #queue-req: 655,","1":null,"pid":null,"ip":null,"time":"2025-11-03 14:31:10","dp_rank":6,"tp_rank":6,"ep_rank":6,"pp_rank":null,"9":null,"num_running_req":251,"num_token":2811200,"token_usage":1.0,"gen_throughput":2055.94,"queue_req":655}]""",
|
||||
)
|
||||
|
||||
df = log_parser.parse(lines)
|
||||
print(df)
|
||||
print(df.write_json())
|
||||
|
||||
assert len(df) == len(lines.strip().splitlines()), f"{len(df)=}"
|
||||
self.assertEqual(json.loads(df.write_json()), expect_rows)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,205 @@
|
||||
"""
|
||||
Integration test for abort_request functionality with a SGLang server.
|
||||
|
||||
Run with:
|
||||
python -m unittest sglang.test.srt.entrypoints.http_server.test_abort_request -v
|
||||
"""
|
||||
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class TestAbortRequest(CustomTestCase):
|
||||
"""Integration test class for abort request functionality."""
|
||||
|
||||
model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Launch the server."""
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=["--disable-cuda-graph"],
|
||||
)
|
||||
|
||||
cls.completion_url = f"{cls.base_url}/generate"
|
||||
cls.abort_url = f"{cls.base_url}/abort_request"
|
||||
cls.health_url = f"{cls.base_url}/health"
|
||||
|
||||
print(f"Server started at {cls.base_url}")
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""Clean up the server."""
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def _send_completion_request(
|
||||
self,
|
||||
text: str,
|
||||
request_id: str,
|
||||
max_tokens: int = 50,
|
||||
temperature: float = 0.8,
|
||||
stream: bool = True,
|
||||
) -> requests.Response:
|
||||
"""Send a completion request to the server."""
|
||||
payload = {
|
||||
"text": text,
|
||||
"sampling_params": {
|
||||
"max_new_tokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
},
|
||||
"stream": stream,
|
||||
"rid": request_id,
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
self.completion_url,
|
||||
json=payload,
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=30,
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
def _send_abort_request(self, request_id: str) -> requests.Response:
|
||||
"""Send an abort request."""
|
||||
payload = {"rid": request_id}
|
||||
return requests.post(self.abort_url, json=payload, timeout=10)
|
||||
|
||||
def _check_server_health(self) -> bool:
|
||||
"""Check if server is healthy."""
|
||||
try:
|
||||
response = requests.get(self.health_url, timeout=5)
|
||||
return response.status_code == 200
|
||||
except:
|
||||
return False
|
||||
|
||||
def test_abort_during_non_streaming_generation(self):
|
||||
"""Test aborting a non-streaming request during generation."""
|
||||
self.assertTrue(self._check_server_health(), "Server should be healthy")
|
||||
|
||||
request_id = "test_abort_non_streaming"
|
||||
completion_result = {}
|
||||
|
||||
def run_completion():
|
||||
response = self._send_completion_request(
|
||||
"Write a detailed essay about artificial intelligence",
|
||||
max_tokens=500,
|
||||
temperature=1,
|
||||
request_id=request_id,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
completion_result["text"] = result.get("text", "")
|
||||
completion_result["finish_reason"] = result.get("meta_info", {}).get(
|
||||
"finish_reason"
|
||||
)
|
||||
|
||||
completion_thread = threading.Thread(target=run_completion)
|
||||
completion_thread.start()
|
||||
time.sleep(0.1)
|
||||
|
||||
abort_response = self._send_abort_request(request_id)
|
||||
completion_thread.join()
|
||||
|
||||
self.assertEqual(abort_response.status_code, 200)
|
||||
self.assertIsNotNone(completion_result, "Should have completion result")
|
||||
if completion_result:
|
||||
finish_reason_obj = completion_result.get("finish_reason")
|
||||
self.assertIsNotNone(finish_reason_obj, "Should have finish_reason")
|
||||
if finish_reason_obj:
|
||||
self.assertEqual(
|
||||
finish_reason_obj.get("type"), "abort", "Should be aborted"
|
||||
)
|
||||
|
||||
def test_batch_requests_with_selective_abort(self):
|
||||
"""Test multiple concurrent requests with selective abort of one request."""
|
||||
self.assertTrue(self._check_server_health(), "Server should be healthy")
|
||||
|
||||
request_ids = ["batch_test_0", "batch_test_1", "batch_test_2"]
|
||||
abort_target_id = "batch_test_1"
|
||||
completion_results = {}
|
||||
threads = []
|
||||
|
||||
def run_completion(req_id, prompt):
|
||||
response = self._send_completion_request(
|
||||
f"Write a story about {prompt}",
|
||||
max_tokens=100,
|
||||
temperature=0.8,
|
||||
request_id=req_id,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
completion_results[req_id] = {
|
||||
"text": result.get("text", ""),
|
||||
"finish_reason": result.get("meta_info", {}).get("finish_reason"),
|
||||
}
|
||||
|
||||
# Start all requests
|
||||
prompts = ["a knight's adventure", "a space discovery", "a chef's restaurant"]
|
||||
for i, req_id in enumerate(request_ids):
|
||||
thread = threading.Thread(target=run_completion, args=(req_id, prompts[i]))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
# Abort one request
|
||||
time.sleep(0.1)
|
||||
abort_response = self._send_abort_request(abort_target_id)
|
||||
|
||||
# Wait for completion
|
||||
for thread in threads:
|
||||
thread.join(timeout=30)
|
||||
|
||||
# Verify results
|
||||
self.assertEqual(abort_response.status_code, 200)
|
||||
|
||||
# Check aborted request
|
||||
aborted_result = completion_results.get(abort_target_id)
|
||||
self.assertIsNotNone(
|
||||
aborted_result, f"Aborted request {abort_target_id} should have result"
|
||||
)
|
||||
if aborted_result:
|
||||
aborted_finish_reason = aborted_result.get("finish_reason")
|
||||
self.assertIsNotNone(
|
||||
aborted_finish_reason, "Aborted request should have finish_reason"
|
||||
)
|
||||
if aborted_finish_reason:
|
||||
self.assertEqual(aborted_finish_reason.get("type"), "abort")
|
||||
|
||||
# Check other requests completed normally
|
||||
normal_completions = 0
|
||||
for req_id in request_ids:
|
||||
if req_id != abort_target_id and req_id in completion_results:
|
||||
result = completion_results[req_id]
|
||||
if result:
|
||||
finish_reason = result.get("finish_reason")
|
||||
if finish_reason and finish_reason.get("type") == "length":
|
||||
normal_completions += 1
|
||||
|
||||
self.assertEqual(
|
||||
normal_completions, 2, "Other 2 requests should complete normally"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2, warnings="ignore")
|
||||
@@ -0,0 +1,262 @@
|
||||
import os
|
||||
import random
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from typing import Dict
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.bench_serving import get_tokenizer
|
||||
from sglang.test.test_disaggregation_utils import TestDisaggregationBase
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
popen_launch_pd_server,
|
||||
)
|
||||
|
||||
|
||||
class DisaggregationHiCacheBase(TestDisaggregationBase):
|
||||
"""Base class for disaggregation with HiCache tests"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
parsed_url = urlparse(DEFAULT_URL_FOR_TEST)
|
||||
cls.base_host = parsed_url.hostname
|
||||
base_port = str(parsed_url.port)
|
||||
cls.lb_port = base_port
|
||||
cls.prefill_port = f"{int(base_port) + 100}"
|
||||
cls.decode_port = f"{int(base_port) + 200}"
|
||||
cls.prefill_url = f"http://{cls.base_host}:{cls.prefill_port}"
|
||||
cls.decode_url = f"http://{cls.base_host}:{cls.decode_port}"
|
||||
cls.lb_url = f"http://{cls.base_host}:{cls.lb_port}"
|
||||
print(f"{cls.base_host=} {cls.lb_port=} {cls.prefill_port=} {cls.decode_port=}")
|
||||
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
cls.temp_dir = tempfile.mkdtemp()
|
||||
cls.start_prefill()
|
||||
cls.start_decode()
|
||||
|
||||
# Block until both
|
||||
cls.wait_server_ready(cls.prefill_url + "/health")
|
||||
cls.wait_server_ready(cls.decode_url + "/health")
|
||||
|
||||
cls.launch_lb()
|
||||
|
||||
@classmethod
|
||||
def start_prefill(cls):
|
||||
# Prefill with HiCache enabled
|
||||
prefill_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"prefill",
|
||||
"--tp-size",
|
||||
"1",
|
||||
"--page-size",
|
||||
"64",
|
||||
"--enable-hierarchical-cache",
|
||||
"--hicache-ratio",
|
||||
"1.2",
|
||||
"--hicache-size",
|
||||
"0",
|
||||
"--hicache-write-policy",
|
||||
"write_through",
|
||||
"--hicache-storage-backend",
|
||||
"file",
|
||||
"--hicache-storage-prefetch-policy",
|
||||
"wait_complete",
|
||||
"--mem-fraction-static",
|
||||
"0.8",
|
||||
]
|
||||
prefill_args += cls.transfer_backend + cls.rdma_devices
|
||||
env = {
|
||||
**os.environ,
|
||||
"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
|
||||
}
|
||||
cls.process_prefill = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.prefill_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=prefill_args,
|
||||
env=env,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
pass
|
||||
|
||||
def gen_prompt(self, token_num: int) -> str:
|
||||
all_available_tokens = list(self.tokenizer.get_vocab().values())
|
||||
selected_tokens = random.choices(all_available_tokens, k=token_num)
|
||||
return self.tokenizer.decode(selected_tokens)
|
||||
|
||||
def send_request(
|
||||
self, prompt: str, max_tokens: int = 100, temperature: float = 0.0
|
||||
) -> Dict:
|
||||
"""Send a generate request and return response"""
|
||||
response = requests.post(
|
||||
f"{self.lb_url}/generate",
|
||||
json={
|
||||
"text": prompt,
|
||||
"sampling_params": {
|
||||
"temperature": temperature,
|
||||
"max_new_tokens": max_tokens,
|
||||
"ignore_eos": True,
|
||||
},
|
||||
},
|
||||
timeout=60,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
response.status_code,
|
||||
200,
|
||||
f"Request failed: {response.status_code} - {response.text}",
|
||||
)
|
||||
return response.json()
|
||||
|
||||
def trigger_offloading_and_flush(self):
|
||||
"""Helper method to trigger offloading and flush cache"""
|
||||
# Trigger offloading
|
||||
self.send_request(self.gen_prompt(1), max_tokens=150)
|
||||
|
||||
# Flush device cache to force remote storage access
|
||||
time.sleep(2)
|
||||
requests.post(self.prefill_url + "/flush_cache")
|
||||
|
||||
|
||||
class TestDisaggregationPrefillWithHiCache(DisaggregationHiCacheBase):
|
||||
"""Test disaggregation with HiCache enabled only on Prefill side"""
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
# Decode without HiCache offload
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp-size",
|
||||
"1",
|
||||
"--page-size",
|
||||
"64",
|
||||
"--mem-fraction-static",
|
||||
"0.8",
|
||||
"--base-gpu-id",
|
||||
"1",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
env = {
|
||||
**os.environ,
|
||||
"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
|
||||
}
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
env=env,
|
||||
)
|
||||
|
||||
def test_prefill_cache_hit(self):
|
||||
"""Test that prefill cache works with repeated queries"""
|
||||
|
||||
repeated_prompt = self.gen_prompt(800)
|
||||
|
||||
# First request - should miss cache
|
||||
self.send_request(repeated_prompt, max_tokens=100)
|
||||
|
||||
# Flush cache
|
||||
self.trigger_offloading_and_flush()
|
||||
|
||||
# Second request - should hit cache (faster)
|
||||
response2 = self.send_request(repeated_prompt, max_tokens=100)
|
||||
|
||||
# Assert cached tokens cnt
|
||||
self.assertGreater(response2["meta_info"]["cached_tokens"], 700)
|
||||
|
||||
|
||||
class TestDisaggregationDecodeWithHiCache(DisaggregationHiCacheBase):
|
||||
"""Test disaggregation with HiCache enabled on both Prefill and Decode sides"""
|
||||
|
||||
@classmethod
|
||||
def start_decode(cls):
|
||||
# Decode with HiCache offload enabled
|
||||
decode_args = [
|
||||
"--trust-remote-code",
|
||||
"--disaggregation-mode",
|
||||
"decode",
|
||||
"--tp-size",
|
||||
"1",
|
||||
"--page-size",
|
||||
"64",
|
||||
"--mem-fraction-static",
|
||||
"0.8",
|
||||
"--base-gpu-id",
|
||||
"1",
|
||||
"--disaggregation-decode-enable-offload-kvcache",
|
||||
"--hicache-ratio",
|
||||
"1.2",
|
||||
"--hicache-size",
|
||||
"0",
|
||||
"--hicache-storage-backend",
|
||||
"file",
|
||||
"--hicache-storage-prefetch-policy",
|
||||
"wait_complete",
|
||||
]
|
||||
decode_args += cls.transfer_backend + cls.rdma_devices
|
||||
env = {
|
||||
**os.environ,
|
||||
"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
|
||||
}
|
||||
cls.process_decode = popen_launch_pd_server(
|
||||
cls.model,
|
||||
cls.decode_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=decode_args,
|
||||
env=env,
|
||||
)
|
||||
|
||||
def test_multi_turn_conversation_cache(self):
|
||||
"""Test multi-turn conversation scenario with cache hit improvement"""
|
||||
|
||||
print("=== Multi-turn Conversation Cache Test ===")
|
||||
|
||||
# Turn 1
|
||||
initial_prompt = self.gen_prompt(300)
|
||||
|
||||
response1 = self.send_request(initial_prompt, max_tokens=200, temperature=0.1)
|
||||
current_context = initial_prompt + response1["text"]
|
||||
|
||||
# Turns 2-4: Continue generation based on previous context
|
||||
previous_cached_tokens = 0
|
||||
|
||||
for turn in range(2, 5):
|
||||
print(f"\nTurn {turn}: Continuing from previous context")
|
||||
|
||||
response = self.send_request(
|
||||
current_context, max_tokens=200, temperature=0.1
|
||||
)
|
||||
cached_tokens = response["meta_info"]["cached_tokens"]
|
||||
|
||||
print(f"Turn {turn} cached tokens: {cached_tokens}")
|
||||
print(f"Improvement: {cached_tokens - previous_cached_tokens} tokens")
|
||||
|
||||
# Assert cache improvement
|
||||
self.assertGreater(
|
||||
cached_tokens,
|
||||
previous_cached_tokens,
|
||||
f"Turn {turn} should have more cached tokens than turn {turn-1}",
|
||||
)
|
||||
|
||||
# Update context and cached tokens for next iteration
|
||||
current_context += response["text"]
|
||||
previous_cached_tokens = cached_tokens
|
||||
|
||||
# Flush prefill cache
|
||||
self.trigger_offloading_and_flush()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,281 @@
|
||||
"""
|
||||
Unit tests comparing TileLang and Triton implementations of activation quantization.
|
||||
Tests both accuracy and performance.
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import Tuple
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.nsa.tilelang_kernel import act_quant
|
||||
from sglang.srt.layers.attention.nsa.triton_kernel import act_quant as act_quant_triton
|
||||
|
||||
|
||||
def benchmark_kernel(
|
||||
fn,
|
||||
x: torch.Tensor,
|
||||
block_size: int,
|
||||
scale_fmt,
|
||||
warmup: int = 10,
|
||||
repeat: int = 100,
|
||||
use_cuda_graph: bool = True,
|
||||
) -> Tuple[float, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Benchmark a kernel function.
|
||||
|
||||
Args:
|
||||
fn: Function to benchmark
|
||||
x: Input tensor
|
||||
block_size: Block size for quantization
|
||||
scale_fmt: Scale format
|
||||
warmup: Number of warmup iterations
|
||||
repeat: Number of repeat iterations
|
||||
use_cuda_graph: Whether to use CUDA graphs for more accurate timing
|
||||
|
||||
Returns:
|
||||
Tuple of (avg_time_ms, quantized_output, scales)
|
||||
"""
|
||||
# Warmup
|
||||
for _ in range(warmup):
|
||||
y, s = fn(x, block_size=block_size, scale_fmt=scale_fmt)
|
||||
|
||||
if not x.is_cuda or not use_cuda_graph:
|
||||
# Fallback to regular timing
|
||||
if x.is_cuda:
|
||||
torch.cuda.synchronize()
|
||||
|
||||
start = time.perf_counter()
|
||||
for _ in range(repeat):
|
||||
y, s = fn(x, block_size=block_size, scale_fmt=scale_fmt)
|
||||
|
||||
if x.is_cuda:
|
||||
torch.cuda.synchronize()
|
||||
|
||||
end = time.perf_counter()
|
||||
avg_time_ms = (end - start) / repeat * 1000
|
||||
|
||||
return avg_time_ms, y, s
|
||||
|
||||
# Use CUDA graph for more accurate timing
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Allocate output buffers
|
||||
N = x.size(-1)
|
||||
y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
|
||||
s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32)
|
||||
|
||||
# Capture CUDA graph
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
y_cap, s_cap = fn(x, block_size=block_size, scale_fmt=scale_fmt)
|
||||
|
||||
# Warmup with graph
|
||||
for _ in range(warmup):
|
||||
graph.replay()
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Timing with CUDA graph
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
|
||||
start_event.record()
|
||||
for _ in range(repeat):
|
||||
graph.replay()
|
||||
end_event.record()
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
avg_time_ms = start_event.elapsed_time(end_event) / repeat
|
||||
|
||||
return avg_time_ms, y_cap, s_cap
|
||||
|
||||
|
||||
def check_accuracy(
|
||||
y_ref: torch.Tensor,
|
||||
s_ref: torch.Tensor,
|
||||
y_test: torch.Tensor,
|
||||
s_test: torch.Tensor,
|
||||
rtol: float = 1e-2,
|
||||
atol: float = 1e-2,
|
||||
) -> Tuple[bool, dict]:
|
||||
"""
|
||||
Check accuracy between reference and test outputs.
|
||||
|
||||
Args:
|
||||
y_ref: Reference quantized output
|
||||
s_ref: Reference scales
|
||||
y_test: Test quantized output
|
||||
s_test: Test scales
|
||||
rtol: Relative tolerance
|
||||
atol: Absolute tolerance
|
||||
|
||||
Returns:
|
||||
Tuple of (passed, metrics_dict)
|
||||
"""
|
||||
# Convert FP8 to float for comparison
|
||||
y_ref_float = y_ref.float()
|
||||
y_test_float = y_test.float()
|
||||
|
||||
# Compute differences
|
||||
y_diff = torch.abs(y_ref_float - y_test_float)
|
||||
s_diff = torch.abs(s_ref - s_test)
|
||||
|
||||
# Compute metrics
|
||||
y_max_diff = y_diff.max().item()
|
||||
y_mean_diff = y_diff.mean().item()
|
||||
s_max_diff = s_diff.max().item()
|
||||
s_mean_diff = s_diff.mean().item()
|
||||
|
||||
# Check relative and absolute tolerance
|
||||
y_close = torch.allclose(y_ref_float, y_test_float, rtol=rtol, atol=atol)
|
||||
s_close = torch.allclose(s_ref, s_test, rtol=rtol, atol=atol)
|
||||
|
||||
# Compute percentage of matching elements
|
||||
y_match_pct = (y_ref_float == y_test_float).float().mean().item() * 100
|
||||
|
||||
metrics = {
|
||||
"y_max_diff": y_max_diff,
|
||||
"y_mean_diff": y_mean_diff,
|
||||
"y_match_pct": y_match_pct,
|
||||
"s_max_diff": s_max_diff,
|
||||
"s_mean_diff": s_mean_diff,
|
||||
"y_close": y_close,
|
||||
"s_close": s_close,
|
||||
}
|
||||
|
||||
passed = y_close and s_close
|
||||
|
||||
return passed, metrics
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
def test_act_quant_comprehensive_benchmark(scale_fmt=None):
|
||||
"""Comprehensive benchmark across multiple sizes with CUDA graphs."""
|
||||
device = torch.device("cuda")
|
||||
dtype = torch.bfloat16
|
||||
block_size = 128
|
||||
|
||||
shapes = [
|
||||
(128, 512),
|
||||
(256, 1024),
|
||||
(512, 2048),
|
||||
(1024, 4096),
|
||||
(2048, 8192),
|
||||
(4096, 16384),
|
||||
]
|
||||
|
||||
print("\n" + "=" * 100)
|
||||
print("Comprehensive Performance Benchmark with CUDA Graphs")
|
||||
print("=" * 100)
|
||||
print(
|
||||
f"{'Shape':<20} {'TileLang (ms)':<15} {'Triton (ms)':<15} {'Speedup':<10} {'Status'}"
|
||||
)
|
||||
print("-" * 100)
|
||||
|
||||
for shape in shapes:
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(shape, dtype=dtype, device=device)
|
||||
|
||||
try:
|
||||
# Benchmark both with CUDA graphs
|
||||
time_tilelang, y_ref, s_ref = benchmark_kernel(
|
||||
act_quant,
|
||||
x,
|
||||
block_size,
|
||||
scale_fmt,
|
||||
warmup=5,
|
||||
repeat=50,
|
||||
use_cuda_graph=True,
|
||||
)
|
||||
time_triton, y_triton, s_triton = benchmark_kernel(
|
||||
act_quant_triton,
|
||||
x,
|
||||
block_size,
|
||||
scale_fmt,
|
||||
warmup=5,
|
||||
repeat=50,
|
||||
use_cuda_graph=True,
|
||||
)
|
||||
|
||||
# Check accuracy
|
||||
passed, _ = check_accuracy(y_ref, s_ref, y_triton, s_triton)
|
||||
|
||||
speedup = time_tilelang / time_triton if time_triton > 0 else 0
|
||||
status = "✓ PASS" if passed else "✗ FAIL"
|
||||
|
||||
print(
|
||||
f"{str(shape):<20} {time_tilelang:<15.4f} {time_triton:<15.4f} "
|
||||
f"{speedup:<10.2f} {status}"
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"{str(shape):<20} ERROR: {str(e)}")
|
||||
|
||||
print("=" * 100)
|
||||
|
||||
# Also run without CUDA graphs for comparison
|
||||
print("\n" + "=" * 100)
|
||||
print("Performance Benchmark WITHOUT CUDA Graphs (for comparison)")
|
||||
print("=" * 100)
|
||||
print(
|
||||
f"{'Shape':<20} {'TileLang (ms)':<15} {'Triton (ms)':<15} {'Speedup':<10} {'Status'}"
|
||||
)
|
||||
print("-" * 100)
|
||||
|
||||
for shape in shapes:
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(shape, dtype=dtype, device=device)
|
||||
|
||||
try:
|
||||
# Benchmark both without CUDA graphs
|
||||
time_tilelang, y_ref, s_ref = benchmark_kernel(
|
||||
act_quant,
|
||||
x,
|
||||
block_size,
|
||||
scale_fmt,
|
||||
warmup=5,
|
||||
repeat=50,
|
||||
use_cuda_graph=False,
|
||||
)
|
||||
time_triton, y_triton, s_triton = benchmark_kernel(
|
||||
act_quant_triton,
|
||||
x,
|
||||
block_size,
|
||||
scale_fmt,
|
||||
warmup=5,
|
||||
repeat=50,
|
||||
use_cuda_graph=False,
|
||||
)
|
||||
|
||||
# Check accuracy
|
||||
passed, _ = check_accuracy(y_ref, s_ref, y_triton, s_triton)
|
||||
|
||||
speedup = time_tilelang / time_triton if time_triton > 0 else 0
|
||||
status = "✓ PASS" if passed else "✗ FAIL"
|
||||
|
||||
print(
|
||||
f"{str(shape):<20} {time_tilelang:<15.4f} {time_triton:<15.4f} "
|
||||
f"{speedup:<10.2f} {status}"
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"{str(shape):<20} ERROR: {str(e)}")
|
||||
|
||||
print("=" * 100)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run comprehensive benchmark
|
||||
if torch.cuda.is_available():
|
||||
print("\n" + "=" * 80)
|
||||
print("Running Comprehensive Benchmark with scale_fmt=None")
|
||||
print("=" * 80)
|
||||
test_act_quant_comprehensive_benchmark(scale_fmt=None)
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("Running Comprehensive Benchmark with scale_fmt!=None")
|
||||
print("=" * 80)
|
||||
test_act_quant_comprehensive_benchmark(scale_fmt="any")
|
||||
else:
|
||||
print("CUDA not available. Skipping tests.")
|
||||
@@ -0,0 +1,191 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MOE_NVFP4,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE,
|
||||
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestMoERunner(CustomTestCase):
|
||||
BASE_URL = DEFAULT_URL_FOR_TEST
|
||||
TIMEOUT = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
|
||||
DEFAULT_EVAL_KWARGS = {
|
||||
"eval_name": "mmlu",
|
||||
"num_examples": 5,
|
||||
"num_threads": 1,
|
||||
}
|
||||
|
||||
CONFIGS = {
|
||||
"moe_runner_auto": {
|
||||
"model": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--moe-runner-backend",
|
||||
"triton",
|
||||
"--attention-backend",
|
||||
"torch_native",
|
||||
"--sampling-backend",
|
||||
"pytorch",
|
||||
],
|
||||
},
|
||||
"moe_runner_triton": {
|
||||
"model": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--moe-runner-backend",
|
||||
"triton",
|
||||
"--attention-backend",
|
||||
"torch_native",
|
||||
"--sampling-backend",
|
||||
"pytorch",
|
||||
],
|
||||
},
|
||||
"moe_runner_triton_kernel": {
|
||||
"model": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--moe-runner-backend",
|
||||
"triton_kernel",
|
||||
"--attention-backend",
|
||||
"torch_native",
|
||||
"--sampling-backend",
|
||||
"pytorch",
|
||||
],
|
||||
},
|
||||
"moe_runner_flashinfer_cutlass": {
|
||||
"model": DEFAULT_MODEL_NAME_FOR_TEST_MOE_NVFP4, # requires model with modelopt_fp4 quantization
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--moe-runner-backend",
|
||||
"flashinfer_cutlass",
|
||||
"--attention-backend",
|
||||
"torch_native",
|
||||
"--sampling-backend",
|
||||
"pytorch",
|
||||
],
|
||||
},
|
||||
"moe_runner_deep_gemm": {
|
||||
"model": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--moe-runner-backend",
|
||||
"deep_gemm",
|
||||
"--attention-backend",
|
||||
"torch_native",
|
||||
"--sampling-backend",
|
||||
"pytorch",
|
||||
],
|
||||
},
|
||||
"moe_runner_flashinfer_trtllm": {
|
||||
"model": DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE, # modelopt_fp4 or fp8 quantization is required for Flashinfer trtllm MOE
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--moe-runner-backend",
|
||||
"flashinfer_trtllm",
|
||||
],
|
||||
},
|
||||
"moe_runner_flashinfer_mxfp4": {
|
||||
"model": DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE,
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--moe-runner-backend",
|
||||
"flashinfer_mxfp4",
|
||||
"--quantization",
|
||||
"mxfp4",
|
||||
"--attention-backend",
|
||||
"torch_native",
|
||||
"--sampling-backend",
|
||||
"pytorch",
|
||||
],
|
||||
},
|
||||
"moe_runner_flashinfer_cutedsl": {
|
||||
"model": DEFAULT_MODEL_NAME_FOR_TEST_MOE_NVFP4,
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--moe-runner-backend",
|
||||
"flashinfer_cutedsl",
|
||||
"--attention-backend",
|
||||
"torch_native",
|
||||
"--sampling-backend",
|
||||
"pytorch",
|
||||
],
|
||||
},
|
||||
"moe_runner_cutlass": {
|
||||
"model": DEFAULT_MODEL_NAME_FOR_TEST_MOE_NVFP4,
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--moe-runner-backend",
|
||||
"cutlass",
|
||||
"--attention-backend",
|
||||
"torch_native",
|
||||
"--sampling-backend",
|
||||
"pytorch",
|
||||
],
|
||||
},
|
||||
"moe_runner_speculative": {
|
||||
"model": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--moe-runner-backend",
|
||||
"triton",
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-draft-model-path",
|
||||
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
|
||||
"--speculative-moe-runner-backend",
|
||||
"triton",
|
||||
"--speculative-num-steps",
|
||||
"2",
|
||||
"--speculative-num-draft-tokens",
|
||||
"4",
|
||||
"--attention-backend",
|
||||
"torch_native",
|
||||
"--sampling-backend",
|
||||
"pytorch",
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
def _run_config(self, config: dict) -> None:
|
||||
model = config["model"]
|
||||
other_args = config.get("other_args", [])
|
||||
eval_kwargs = self.DEFAULT_EVAL_KWARGS
|
||||
|
||||
process = popen_launch_server(
|
||||
model,
|
||||
self.BASE_URL,
|
||||
timeout=self.TIMEOUT,
|
||||
other_args=other_args,
|
||||
)
|
||||
try:
|
||||
args = SimpleNamespace(
|
||||
base_url=self.BASE_URL,
|
||||
model=model,
|
||||
**eval_kwargs,
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreaterEqual(metrics["score"], 0.48)
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
|
||||
for _name, _cfg in TestMoERunner.CONFIGS.items():
|
||||
setattr(
|
||||
TestMoERunner,
|
||||
f"test_{_name}",
|
||||
(lambda self, cfg=_cfg: self._run_config(cfg)),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,761 @@
|
||||
import random
|
||||
import unittest
|
||||
from enum import Enum
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.lora.backend.chunked_backend import ChunkedSgmvLoRABackend
|
||||
from sglang.srt.lora.triton_ops import (
|
||||
chunked_sgmv_lora_expand_forward,
|
||||
chunked_sgmv_lora_shrink_forward,
|
||||
)
|
||||
from sglang.srt.lora.triton_ops.chunked_sgmv_expand import _chunked_lora_expand_kernel
|
||||
from sglang.srt.lora.triton_ops.chunked_sgmv_shrink import _chunked_lora_shrink_kernel
|
||||
from sglang.srt.lora.utils import LoRABatchInfo
|
||||
|
||||
CHUNK_SIZE = 16
|
||||
|
||||
|
||||
def reset_kernel_cache():
|
||||
_chunked_lora_shrink_kernel._clear_cache()
|
||||
_chunked_lora_expand_kernel._clear_cache()
|
||||
|
||||
|
||||
def safe_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
|
||||
"""Matrix multiplication with mixed precision handling for float16"""
|
||||
result = torch.matmul(a.float(), b.float())
|
||||
return result.to(a.dtype)
|
||||
|
||||
|
||||
class BatchComposition(Enum):
|
||||
UNIFORM = "uniform"
|
||||
MIXED = "mixed"
|
||||
SKEWED = "skewed"
|
||||
NONE = "_NO_LORA_"
|
||||
|
||||
|
||||
class BatchMode(Enum):
|
||||
PREFILL = "prefill"
|
||||
DECODE = "decode"
|
||||
|
||||
|
||||
def reference_sgmv_shrink(
|
||||
x: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
seq_lengths: List[int],
|
||||
lora_assignments: List[str],
|
||||
num_slices: int = 1,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Simple sequence-level reference implementation of SGMV shrink operation.
|
||||
|
||||
Args:
|
||||
x: (total_seq_len, input_dim) - Input activations
|
||||
weights: (num_loras, num_slices * max_rank, input_dim) - LoRA A weights
|
||||
batch_info: Batch information (only used for lora_ranks)
|
||||
seq_lengths: Length of each sequence
|
||||
lora_assignments: LoRA name for each sequence
|
||||
num_slices: Number of slices (3 for QKV, 2 for gate_up, 1 for others)
|
||||
|
||||
Returns:
|
||||
output: (total_seq_len, num_slices * max_rank) - Intermediate activations
|
||||
"""
|
||||
if weights.numel() == 0:
|
||||
total_seq_len = x.shape[0]
|
||||
return torch.zeros(total_seq_len, 0, dtype=x.dtype, device=x.device)
|
||||
|
||||
total_seq_len, input_dim = x.shape
|
||||
num_loras, weight_out_dim, _ = weights.shape
|
||||
max_rank = weight_out_dim // num_slices
|
||||
|
||||
output = torch.zeros(
|
||||
total_seq_len, num_slices * max_rank, dtype=x.dtype, device=x.device
|
||||
)
|
||||
|
||||
unique_loras = sorted(set(lora_assignments))
|
||||
lora_name_to_idx = {name: idx for idx, name in enumerate(unique_loras)}
|
||||
lora_ranks = batch_info.lora_ranks.cpu().numpy()
|
||||
|
||||
token_offset = 0
|
||||
for seq_len, lora_name in zip(seq_lengths, lora_assignments):
|
||||
if seq_len == 0:
|
||||
continue
|
||||
|
||||
lora_idx = lora_name_to_idx[lora_name]
|
||||
rank = lora_ranks[lora_idx]
|
||||
|
||||
if rank > 0:
|
||||
x_seq = x[token_offset : token_offset + seq_len, :]
|
||||
w_seq = weights[lora_idx, : num_slices * rank, :]
|
||||
|
||||
result = safe_matmul(x_seq, w_seq.t())
|
||||
output[token_offset : token_offset + seq_len, : num_slices * rank] = result
|
||||
|
||||
token_offset += seq_len
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def reference_sgmv_expand(
|
||||
x: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
batch_info: LoRABatchInfo,
|
||||
seq_lengths: List[int],
|
||||
lora_assignments: List[str],
|
||||
slice_offsets: torch.Tensor,
|
||||
max_slice_size: int,
|
||||
base_output: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Simple sequence-level reference implementation of SGMV expand operation.
|
||||
|
||||
Args:
|
||||
x: (total_seq_len, num_slices * max_rank) - Intermediate activations
|
||||
weights: (num_loras, output_dim, max_rank) - LoRA B weights
|
||||
batch_info: Batch information (only used for lora_ranks)
|
||||
seq_lengths: Length of each sequence
|
||||
lora_assignments: LoRA name for each sequence
|
||||
slice_offsets: Tensor defining slice boundaries
|
||||
max_slice_size: Maximum slice size for chunking
|
||||
base_output: Optional base output to accumulate into
|
||||
|
||||
Returns:
|
||||
output: (total_seq_len, total_output_dim) - Final output
|
||||
"""
|
||||
if weights.numel() == 0:
|
||||
total_seq_len = x.shape[0]
|
||||
total_output_dim = slice_offsets[-1].item() if len(slice_offsets) > 0 else 0
|
||||
return torch.zeros(
|
||||
total_seq_len, total_output_dim, dtype=x.dtype, device=x.device
|
||||
)
|
||||
|
||||
total_seq_len, _ = x.shape
|
||||
|
||||
num_slices = len(slice_offsets) - 1
|
||||
|
||||
if base_output is not None:
|
||||
output = base_output.clone()
|
||||
else:
|
||||
total_output_dim = slice_offsets[-1].item()
|
||||
output = torch.zeros(
|
||||
total_seq_len, total_output_dim, dtype=x.dtype, device=x.device
|
||||
)
|
||||
|
||||
unique_loras = sorted(set(lora_assignments))
|
||||
lora_name_to_idx = {name: idx for idx, name in enumerate(unique_loras)}
|
||||
lora_ranks = batch_info.lora_ranks.cpu().numpy()
|
||||
|
||||
token_offset = 0
|
||||
for seq_len, lora_name in zip(seq_lengths, lora_assignments):
|
||||
if seq_len == 0:
|
||||
continue
|
||||
|
||||
lora_idx = lora_name_to_idx[lora_name]
|
||||
lora_rank = lora_ranks[lora_idx]
|
||||
|
||||
if lora_rank > 0:
|
||||
# Extract sequence intermediate activations
|
||||
x_seq = x[
|
||||
token_offset : token_offset + seq_len, : num_slices * lora_rank
|
||||
] # (seq_len, num_slices * rank)
|
||||
|
||||
for slice_idx in range(num_slices):
|
||||
slice_start_input = slice_idx * lora_rank
|
||||
slice_end_input = (slice_idx + 1) * lora_rank
|
||||
|
||||
slice_start_output = slice_offsets[slice_idx].item()
|
||||
slice_end_output = slice_offsets[slice_idx + 1].item()
|
||||
|
||||
x_slice = x_seq[:, slice_start_input:slice_end_input] # (seq_len, rank)
|
||||
w_slice = weights[
|
||||
lora_idx, slice_start_output:slice_end_output, :lora_rank
|
||||
] # (slice_dim, rank)
|
||||
|
||||
result = safe_matmul(x_slice, w_slice.t()) # (seq_len, slice_dim)
|
||||
output[
|
||||
token_offset : token_offset + seq_len,
|
||||
slice_start_output:slice_end_output,
|
||||
] += result
|
||||
|
||||
token_offset += seq_len
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class TestChunkedSGMV(unittest.TestCase):
|
||||
|
||||
# Test configuration constants
|
||||
RTOL = 1e-3
|
||||
ATOL = 1e-3
|
||||
DEFAULT_BATCH_SIZE = 8
|
||||
|
||||
def _compare_shrink_outputs(
|
||||
self,
|
||||
chunked_output: torch.Tensor,
|
||||
reference_output: torch.Tensor,
|
||||
seq_lengths: List[int],
|
||||
lora_assignments: List[str],
|
||||
batch_info: LoRABatchInfo,
|
||||
num_slices: int,
|
||||
test_name: str,
|
||||
):
|
||||
"""
|
||||
Compare only the valid portions of shrink outputs.
|
||||
|
||||
The chunked SGMV shrink kernel only guarantees correctness for
|
||||
output[seq_start:seq_end, :rank * num_slices] for each sequence.
|
||||
"""
|
||||
# Create mapping from LoRA names to indices and ranks
|
||||
unique_loras = sorted(set(lora_assignments))
|
||||
lora_name_to_idx = {name: idx for idx, name in enumerate(unique_loras)}
|
||||
lora_ranks = batch_info.lora_ranks.cpu().numpy()
|
||||
|
||||
token_offset = 0
|
||||
for seq_idx, (seq_len, lora_name) in enumerate(
|
||||
zip(seq_lengths, lora_assignments)
|
||||
):
|
||||
if seq_len == 0:
|
||||
continue
|
||||
|
||||
lora_idx = lora_name_to_idx[lora_name]
|
||||
rank = lora_ranks[lora_idx]
|
||||
|
||||
if rank > 0:
|
||||
# Only compare the valid columns for this sequence
|
||||
valid_cols = num_slices * rank
|
||||
|
||||
chunked_seq = chunked_output[
|
||||
token_offset : token_offset + seq_len, :valid_cols
|
||||
]
|
||||
reference_seq = reference_output[
|
||||
token_offset : token_offset + seq_len, :valid_cols
|
||||
]
|
||||
|
||||
torch.testing.assert_close(
|
||||
chunked_seq,
|
||||
reference_seq,
|
||||
rtol=self.RTOL,
|
||||
atol=self.ATOL,
|
||||
msg=f"Shrink operation failed for {test_name}, sequence {seq_idx} ({lora_name})",
|
||||
)
|
||||
|
||||
token_offset += seq_len
|
||||
|
||||
def setUp(self):
|
||||
"""Set up common test parameters"""
|
||||
torch.manual_seed(42)
|
||||
random.seed(42)
|
||||
|
||||
self.device = torch.device("cuda")
|
||||
self.dtype = torch.float16
|
||||
self.input_dim = 2560 # Hidden dimension
|
||||
self.max_seq_len = 1024
|
||||
|
||||
# LoRA configurations: name -> (rank, output_q, output_k, output_v)
|
||||
self.lora_configs = {
|
||||
"lora_A": (8, 4096, 1024, 1024),
|
||||
"lora_B": (16, 4096, 1024, 1024),
|
||||
"lora_C": (32, 4096, 1024, 1024),
|
||||
"_NO_LORA_": (0, 4096, 1024, 1024),
|
||||
}
|
||||
|
||||
# QKV slice offsets: 4096 (Q) + 1024 (K) + 1024 (V) = 6144 total
|
||||
self.slice_offsets = torch.tensor(
|
||||
[0, 4096, 5120, 6144], dtype=torch.int32, device=self.device
|
||||
)
|
||||
self.max_slice_size = 4096
|
||||
|
||||
def generate_sequence_lengths(
|
||||
self,
|
||||
batch_size: int,
|
||||
batch_mode: BatchMode = BatchMode.PREFILL,
|
||||
min_len: int = 1,
|
||||
max_len: int = None,
|
||||
) -> List[int]:
|
||||
"""Generate sequence lengths for a batch based on mode"""
|
||||
if batch_mode == BatchMode.DECODE:
|
||||
return [1] * batch_size
|
||||
else:
|
||||
if max_len is None:
|
||||
max_len = self.max_seq_len
|
||||
return [random.randint(min_len, max_len) for _ in range(batch_size)]
|
||||
|
||||
def create_lora_weights(
|
||||
self, lora_name: str, include_missing_k: bool = False
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Create LoRA A and B weights for given configuration"""
|
||||
rank, out_q, out_k, out_v = self.lora_configs[lora_name]
|
||||
|
||||
if rank == 0:
|
||||
lora_a = torch.empty(
|
||||
0, self.input_dim, dtype=self.dtype, device=self.device
|
||||
)
|
||||
lora_b = torch.empty(
|
||||
out_q + out_k + out_v, 0, dtype=self.dtype, device=self.device
|
||||
)
|
||||
return lora_a, lora_b
|
||||
|
||||
# Create LoRA A weights (3 slices for QKV)
|
||||
lora_a = torch.randn(
|
||||
3 * rank, self.input_dim, dtype=self.dtype, device=self.device
|
||||
)
|
||||
|
||||
if include_missing_k:
|
||||
lora_a[rank : 2 * rank, :] = 0.0
|
||||
|
||||
# Create LoRA B weights (stacked Q, K, V)
|
||||
total_output_dim = out_q + out_k + out_v
|
||||
lora_b = torch.randn(
|
||||
total_output_dim, rank, dtype=self.dtype, device=self.device
|
||||
)
|
||||
|
||||
if include_missing_k:
|
||||
lora_b[out_q : out_q + out_k, :] = 0.0
|
||||
|
||||
return lora_a, lora_b
|
||||
|
||||
def create_batch_info(
|
||||
self,
|
||||
seq_lengths: List[int],
|
||||
lora_assignments: List[Optional[str]],
|
||||
batch_mode: BatchMode = BatchMode.PREFILL,
|
||||
) -> LoRABatchInfo:
|
||||
"""Create LoRABatchInfo using the same logic as chunked backend"""
|
||||
unique_loras = sorted(set(lora_assignments))
|
||||
lora_name_to_idx = {name: idx for idx, name in enumerate(unique_loras)}
|
||||
|
||||
seq_weight_indices = [lora_name_to_idx[name] for name in lora_assignments]
|
||||
|
||||
lora_ranks = [self.lora_configs[name][0] for name in unique_loras]
|
||||
|
||||
def create_mock_batch():
|
||||
# Create a minimal mock ForwardBatch for the test
|
||||
class MockForwardBatch:
|
||||
def __init__(self, batch_size, seq_lengths):
|
||||
self.batch_size = batch_size
|
||||
self.extend_seq_lens_cpu = seq_lengths
|
||||
self.forward_mode = MockForwardMode()
|
||||
|
||||
class MockForwardMode:
|
||||
def is_extend(self):
|
||||
return batch_mode == BatchMode.PREFILL
|
||||
|
||||
return MockForwardBatch(len(seq_lengths), seq_lengths)
|
||||
|
||||
mock_batch = create_mock_batch()
|
||||
|
||||
# Use the same functions as chunked backend
|
||||
permutation, weights_reordered = ChunkedSgmvLoRABackend._get_permutation(
|
||||
seq_weight_indices, mock_batch
|
||||
)
|
||||
|
||||
# Create a minimal backend instance to access _get_segments_info
|
||||
mock_server_args = type(
|
||||
"ServerArgs", (object,), {"max_lora_chunk_size": "MOCK_NEVER_USED"}
|
||||
)
|
||||
mock_backend = ChunkedSgmvLoRABackend(
|
||||
max_loras_per_batch=8, device=self.device, server_args=mock_server_args
|
||||
)
|
||||
weight_indices_list, seg_indptr = mock_backend._get_segments_info(
|
||||
weights_reordered,
|
||||
chunk_size=CHUNK_SIZE,
|
||||
)
|
||||
|
||||
scalings = [1.0] * len(unique_loras)
|
||||
seg_indptr_tensor = seg_indptr.to(self.device)
|
||||
weight_indices_tensor = weight_indices_list.to(self.device)
|
||||
lora_ranks_tensor = (
|
||||
torch.tensor(lora_ranks, dtype=torch.int32, device=self.device)
|
||||
if lora_ranks
|
||||
else torch.empty(0, dtype=torch.int32, device=self.device)
|
||||
)
|
||||
scalings_tensor = (
|
||||
torch.tensor(scalings, dtype=torch.float32, device=self.device)
|
||||
if scalings
|
||||
else torch.empty(0, dtype=torch.float32, device=self.device)
|
||||
)
|
||||
permutation_tensor = permutation.to(
|
||||
self.device, dtype=torch.int32
|
||||
) # Convert to int32 for LoRABatchInfo
|
||||
seq_lens_tensor = torch.tensor(
|
||||
seq_lengths, dtype=torch.int32, device=self.device
|
||||
)
|
||||
|
||||
return LoRABatchInfo(
|
||||
use_cuda_graph=False,
|
||||
bs=len(seq_lengths),
|
||||
num_segments=len(weight_indices_list), # Number of segments, not sequences!
|
||||
seg_indptr=seg_indptr_tensor,
|
||||
weight_indices=weight_indices_tensor,
|
||||
lora_ranks=lora_ranks_tensor,
|
||||
scalings=scalings_tensor,
|
||||
seg_lens=seq_lens_tensor, # Original sequence lengths for reference
|
||||
max_len=CHUNK_SIZE,
|
||||
permutation=permutation_tensor, # Token reordering permutation
|
||||
)
|
||||
|
||||
def stack_lora_weights(
|
||||
self, weight_list: List[torch.Tensor], is_lora_a: bool
|
||||
) -> torch.Tensor:
|
||||
"""Stack LoRA weights from different adapters into a single tensor"""
|
||||
if not weight_list:
|
||||
return torch.empty(0, 0, 0, dtype=self.dtype, device=self.device)
|
||||
|
||||
first_non_empty = next((w for w in weight_list if w.numel() > 0), None)
|
||||
if first_non_empty is None:
|
||||
return torch.empty(
|
||||
len(weight_list), 0, 0, dtype=self.dtype, device=self.device
|
||||
)
|
||||
if is_lora_a:
|
||||
# LoRA A: (slice_num * rank, input_dim) -> (num_loras, slice_num * max_rank, input_dim)
|
||||
max_rank = max(w.shape[0] // 3 if w.numel() > 0 else 0 for w in weight_list)
|
||||
final_shape = (len(weight_list), 3 * max_rank, self.input_dim)
|
||||
else:
|
||||
# LoRA B: (output_dim, rank) -> (num_loras, output_dim, max_rank)
|
||||
max_rank = max(w.shape[1] if w.numel() > 0 else 0 for w in weight_list)
|
||||
output_dim = first_non_empty.shape[0]
|
||||
final_shape = (len(weight_list), output_dim, max_rank)
|
||||
|
||||
stacked = torch.zeros(final_shape, dtype=self.dtype, device=self.device)
|
||||
|
||||
for i, weight in enumerate(weight_list):
|
||||
if weight.numel() > 0:
|
||||
if is_lora_a:
|
||||
stacked[i, : weight.shape[0], :] = weight
|
||||
else:
|
||||
stacked[i, :, : weight.shape[1]] = weight
|
||||
|
||||
return stacked
|
||||
|
||||
def create_test_batch(
|
||||
self,
|
||||
batch_composition: BatchComposition,
|
||||
batch_size: int,
|
||||
batch_mode: BatchMode = BatchMode.PREFILL,
|
||||
include_missing_k: bool = False,
|
||||
) -> Tuple[
|
||||
torch.Tensor,
|
||||
Dict[str, Tuple[torch.Tensor, torch.Tensor]],
|
||||
LoRABatchInfo,
|
||||
List[int],
|
||||
List[str],
|
||||
]:
|
||||
"""Create test batch with specified composition and mode"""
|
||||
|
||||
# Reset kernel cache to avoid cross-test contamination
|
||||
reset_kernel_cache()
|
||||
|
||||
seq_lengths = self.generate_sequence_lengths(
|
||||
batch_size, batch_mode, 1, self.max_seq_len
|
||||
)
|
||||
if batch_composition == BatchComposition.UNIFORM:
|
||||
lora_assignments = ["lora_A"] * batch_size
|
||||
elif batch_composition == BatchComposition.MIXED:
|
||||
lora_names = ["lora_A", "lora_B", "lora_C", None]
|
||||
lora_assignments = [
|
||||
lora_names[i % len(lora_names)] for i in range(batch_size)
|
||||
]
|
||||
elif batch_composition == BatchComposition.SKEWED:
|
||||
num_minority = max(1, batch_size // 8)
|
||||
lora_assignments = ["lora_A"] * num_minority + ["lora_B"] * (
|
||||
batch_size - num_minority
|
||||
)
|
||||
random.shuffle(lora_assignments)
|
||||
elif batch_composition == BatchComposition.NONE:
|
||||
lora_assignments = [None] * batch_size
|
||||
else:
|
||||
raise ValueError(f"Unknown batch composition: {batch_composition}")
|
||||
|
||||
total_seq_len = sum(seq_lengths)
|
||||
x = torch.randn(
|
||||
total_seq_len, self.input_dim, dtype=self.dtype, device=self.device
|
||||
)
|
||||
|
||||
normalized_assignments = [
|
||||
name if name is not None else "_NO_LORA_" for name in lora_assignments
|
||||
]
|
||||
unique_loras = set(normalized_assignments)
|
||||
weights = {}
|
||||
for lora_name in unique_loras:
|
||||
weights[lora_name] = self.create_lora_weights(lora_name, include_missing_k)
|
||||
|
||||
batch_info = self.create_batch_info(
|
||||
seq_lengths, normalized_assignments, batch_mode
|
||||
)
|
||||
|
||||
return x, weights, batch_info, seq_lengths, normalized_assignments
|
||||
|
||||
def run_test_comparison(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
weights: Dict[str, Tuple[torch.Tensor, torch.Tensor]],
|
||||
batch_info: LoRABatchInfo,
|
||||
seq_lengths: List[int],
|
||||
lora_assignments: List[str],
|
||||
test_name: str,
|
||||
):
|
||||
"""Run comparison between chunked and reference implementations"""
|
||||
if not weights: # Handle case with no LoRA weights
|
||||
return
|
||||
|
||||
# Stack LoRA A weights
|
||||
lora_a_weights = [weights[name][0] for name in sorted(weights.keys())]
|
||||
stacked_lora_a = self.stack_lora_weights(lora_a_weights, is_lora_a=True)
|
||||
|
||||
# Stack LoRA B weights
|
||||
lora_b_weights = [weights[name][1] for name in sorted(weights.keys())]
|
||||
stacked_lora_b = self.stack_lora_weights(lora_b_weights, is_lora_a=False)
|
||||
|
||||
# Test shrink operation
|
||||
chunked_shrink = chunked_sgmv_lora_shrink_forward(
|
||||
x, stacked_lora_a, batch_info, num_slices=3
|
||||
)
|
||||
reference_shrink = reference_sgmv_shrink(
|
||||
x, stacked_lora_a, batch_info, seq_lengths, lora_assignments, num_slices=3
|
||||
)
|
||||
|
||||
# Only compare valid portions of shrink output (first rank * num_slices columns per sequence)
|
||||
self._compare_shrink_outputs(
|
||||
chunked_shrink,
|
||||
reference_shrink,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
batch_info,
|
||||
num_slices=3,
|
||||
test_name=test_name,
|
||||
)
|
||||
|
||||
# Test expand operation
|
||||
chunked_expand = chunked_sgmv_lora_expand_forward(
|
||||
reference_shrink,
|
||||
stacked_lora_b,
|
||||
batch_info,
|
||||
self.slice_offsets,
|
||||
self.max_slice_size,
|
||||
base_output=None,
|
||||
)
|
||||
reference_expand = reference_sgmv_expand(
|
||||
reference_shrink,
|
||||
stacked_lora_b,
|
||||
batch_info,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
self.slice_offsets,
|
||||
self.max_slice_size,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(
|
||||
chunked_expand,
|
||||
reference_expand,
|
||||
rtol=self.RTOL,
|
||||
atol=self.ATOL,
|
||||
msg=f"Expand operation failed for {test_name}",
|
||||
)
|
||||
|
||||
# === Basic Operations Tests ===
|
||||
|
||||
def test_shrink_basic(self):
|
||||
"""Test basic shrink operation against PyTorch reference"""
|
||||
for batch_size in [1, 2, 16, 64]:
|
||||
with self.subTest(batch_size=batch_size):
|
||||
x, weights, batch_info, seq_lengths, lora_assignments = (
|
||||
self.create_test_batch(BatchComposition.UNIFORM, batch_size)
|
||||
)
|
||||
|
||||
lora_a_weights = [weights[name][0] for name in sorted(weights.keys())]
|
||||
stacked_lora_a = self.stack_lora_weights(lora_a_weights, is_lora_a=True)
|
||||
|
||||
chunked_shrink = chunked_sgmv_lora_shrink_forward(
|
||||
x, stacked_lora_a, batch_info, num_slices=3
|
||||
)
|
||||
reference_shrink = reference_sgmv_shrink(
|
||||
x,
|
||||
stacked_lora_a,
|
||||
batch_info,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
num_slices=3,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(
|
||||
chunked_shrink, reference_shrink, rtol=self.RTOL, atol=self.ATOL
|
||||
)
|
||||
|
||||
def test_expand_basic(self):
|
||||
"""Test basic expand operation against PyTorch reference"""
|
||||
for batch_size in [1, 2, 16, 64]:
|
||||
with self.subTest(batch_size=batch_size):
|
||||
x, weights, batch_info, seq_lengths, lora_assignments = (
|
||||
self.create_test_batch(BatchComposition.UNIFORM, batch_size)
|
||||
)
|
||||
|
||||
lora_a_weights = [weights[name][0] for name in sorted(weights.keys())]
|
||||
stacked_lora_a = self.stack_lora_weights(lora_a_weights, is_lora_a=True)
|
||||
|
||||
intermediate = reference_sgmv_shrink(
|
||||
x,
|
||||
stacked_lora_a,
|
||||
batch_info,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
num_slices=3,
|
||||
)
|
||||
|
||||
lora_b_weights = [weights[name][1] for name in sorted(weights.keys())]
|
||||
stacked_lora_b = self.stack_lora_weights(
|
||||
lora_b_weights, is_lora_a=False
|
||||
)
|
||||
|
||||
chunked_expand = chunked_sgmv_lora_expand_forward(
|
||||
intermediate,
|
||||
stacked_lora_b,
|
||||
batch_info,
|
||||
self.slice_offsets,
|
||||
self.max_slice_size,
|
||||
base_output=None,
|
||||
)
|
||||
reference_expand = reference_sgmv_expand(
|
||||
intermediate,
|
||||
stacked_lora_b,
|
||||
batch_info,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
self.slice_offsets,
|
||||
self.max_slice_size,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(
|
||||
chunked_expand, reference_expand, rtol=self.RTOL, atol=self.ATOL
|
||||
)
|
||||
|
||||
# === QKV Operations Test ===
|
||||
|
||||
def test_qkv_missing_projections(self):
|
||||
"""Test QKV operations with missing k_proj (Qwen3 scenario)"""
|
||||
for batch_size in [1, 2, 16, 64]:
|
||||
with self.subTest(batch_size=batch_size):
|
||||
x, weights, batch_info, seq_lengths, lora_assignments = (
|
||||
self.create_test_batch(
|
||||
BatchComposition.MIXED, batch_size, include_missing_k=True
|
||||
)
|
||||
)
|
||||
self.run_test_comparison(
|
||||
x,
|
||||
weights,
|
||||
batch_info,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
f"QKV missing k_proj batch_size={batch_size}",
|
||||
)
|
||||
|
||||
# === Batch Composition Tests ===
|
||||
|
||||
def test_uniform_lora_batch(self):
|
||||
"""All sequences use same LoRA, random sequence lengths"""
|
||||
for batch_size in [1, 2, 16, 64]:
|
||||
with self.subTest(batch_size=batch_size):
|
||||
x, weights, batch_info, seq_lengths, lora_assignments = (
|
||||
self.create_test_batch(BatchComposition.UNIFORM, batch_size)
|
||||
)
|
||||
self.run_test_comparison(
|
||||
x,
|
||||
weights,
|
||||
batch_info,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
f"uniform batch_size={batch_size}",
|
||||
)
|
||||
|
||||
def test_evenly_mixed_lora_batch(self):
|
||||
"""Sequences evenly distributed across LoRAs, random lengths"""
|
||||
for batch_size in [1, 2, 16, 64]:
|
||||
with self.subTest(batch_size=batch_size):
|
||||
x, weights, batch_info, seq_lengths, lora_assignments = (
|
||||
self.create_test_batch(BatchComposition.MIXED, batch_size)
|
||||
)
|
||||
self.run_test_comparison(
|
||||
x,
|
||||
weights,
|
||||
batch_info,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
f"mixed batch_size={batch_size}",
|
||||
)
|
||||
|
||||
def test_highly_skewed_lora_batch(self):
|
||||
"""Highly uneven LoRA distribution, random lengths"""
|
||||
for batch_size in [1, 2, 16, 64]:
|
||||
with self.subTest(batch_size=batch_size):
|
||||
x, weights, batch_info, seq_lengths, lora_assignments = (
|
||||
self.create_test_batch(BatchComposition.SKEWED, batch_size)
|
||||
)
|
||||
self.run_test_comparison(
|
||||
x,
|
||||
weights,
|
||||
batch_info,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
f"skewed batch_size={batch_size}",
|
||||
)
|
||||
|
||||
# === Decode Mode Tests ===
|
||||
|
||||
def test_decode_uniform_lora_batch(self):
|
||||
"""Decode mode: All sequences use same LoRA, all length 1"""
|
||||
for batch_size in [1, 2, 16, 64]:
|
||||
with self.subTest(batch_size=batch_size):
|
||||
x, weights, batch_info, seq_lengths, lora_assignments = (
|
||||
self.create_test_batch(
|
||||
BatchComposition.UNIFORM, batch_size, BatchMode.DECODE
|
||||
)
|
||||
)
|
||||
self.run_test_comparison(
|
||||
x,
|
||||
weights,
|
||||
batch_info,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
f"decode uniform batch_size={batch_size}",
|
||||
)
|
||||
|
||||
def test_decode_mixed_lora_batch(self):
|
||||
"""Decode mode: Sequences distributed across LoRAs, all length 1"""
|
||||
for batch_size in [1, 2, 16, 64]:
|
||||
with self.subTest(batch_size=batch_size):
|
||||
x, weights, batch_info, seq_lengths, lora_assignments = (
|
||||
self.create_test_batch(
|
||||
BatchComposition.MIXED, batch_size, BatchMode.DECODE
|
||||
)
|
||||
)
|
||||
self.run_test_comparison(
|
||||
x,
|
||||
weights,
|
||||
batch_info,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
f"decode mixed batch_size={batch_size}",
|
||||
)
|
||||
|
||||
def test_decode_skewed_lora_batch(self):
|
||||
"""Decode mode: Highly uneven LoRA distribution, all length 1"""
|
||||
for batch_size in [1, 2, 16, 64]:
|
||||
with self.subTest(batch_size=batch_size):
|
||||
x, weights, batch_info, seq_lengths, lora_assignments = (
|
||||
self.create_test_batch(
|
||||
BatchComposition.SKEWED, batch_size, BatchMode.DECODE
|
||||
)
|
||||
)
|
||||
self.run_test_comparison(
|
||||
x,
|
||||
weights,
|
||||
batch_info,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
f"decode skewed batch_size={batch_size}",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,108 @@
|
||||
# Copyright 2023-2024 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import unittest
|
||||
from typing import List
|
||||
|
||||
from utils import (
|
||||
ALL_OTHER_LORA_MODELS,
|
||||
CI_LORA_MODELS,
|
||||
DEFAULT_PROMPTS,
|
||||
TORCH_DTYPES,
|
||||
LoRAModelCase,
|
||||
run_lora_test_by_batch,
|
||||
run_lora_test_one_by_one,
|
||||
)
|
||||
|
||||
from sglang.test.test_utils import CustomTestCase, is_in_ci
|
||||
|
||||
TEST_CUDA_GRAPH_PADDING_PROMPTS = [
|
||||
"AI is a field of computer science focused on",
|
||||
"""
|
||||
### Instruction:
|
||||
Tell me about llamas and alpacas
|
||||
### Response:
|
||||
Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids (camels, dromedaries). Llamas live in the Andean mountains of South America where they graze on grasses and shrubs. Alpaca is another name for domesticated llama. The word "alpaca" comes from an Incan language meaning "golden fleece." Alpacas look very similar to llamas but are smaller than their wild relatives. Both species were used by ancient people as pack animals and for meat. Today both llamas and alpacas are raised primarily for their fiber which can be spun into yarn or knitted into clothing.
|
||||
### Question 2:
|
||||
What do you know about llamas?
|
||||
### Answer:
|
||||
""",
|
||||
"Computer science is the study of",
|
||||
]
|
||||
|
||||
|
||||
class TestLoRACudaGraph(CustomTestCase):
|
||||
|
||||
def _run_without_cuda_graph_on_model_cases(self, model_cases: List[LoRAModelCase]):
|
||||
# Since we have already enabled CUDA graph by default in other lora tests,
|
||||
# we only need to run lora tests without CUDA graph here.
|
||||
for model_case in model_cases:
|
||||
# If skip_long_prompt is True, filter out prompts longer than 1000 characters
|
||||
prompts = (
|
||||
DEFAULT_PROMPTS
|
||||
if not model_case.skip_long_prompt
|
||||
else [p for p in DEFAULT_PROMPTS if len(p) < 1000]
|
||||
)
|
||||
for torch_dtype in TORCH_DTYPES:
|
||||
run_lora_test_one_by_one(
|
||||
prompts,
|
||||
model_case,
|
||||
torch_dtype,
|
||||
max_new_tokens=32,
|
||||
disable_cuda_graph=True,
|
||||
test_tag="without_cuda_graph",
|
||||
)
|
||||
|
||||
def _run_cuda_graph_padding_on_model_cases(self, model_cases: List[LoRAModelCase]):
|
||||
for model_case in model_cases:
|
||||
# Run a batch size of 3, which will not be captured by CUDA graph and need padding
|
||||
prompts = TEST_CUDA_GRAPH_PADDING_PROMPTS
|
||||
for torch_dtype in TORCH_DTYPES:
|
||||
run_lora_test_by_batch(
|
||||
prompts,
|
||||
model_case,
|
||||
torch_dtype,
|
||||
max_new_tokens=32,
|
||||
disable_cuda_graph=False,
|
||||
test_tag="cuda_graph_padding",
|
||||
)
|
||||
|
||||
def test_ci_lora_models(self):
|
||||
self._run_without_cuda_graph_on_model_cases(CI_LORA_MODELS)
|
||||
self._run_cuda_graph_padding_on_model_cases(CI_LORA_MODELS)
|
||||
|
||||
def test_all_lora_models(self):
|
||||
if is_in_ci():
|
||||
return
|
||||
|
||||
# Retain ONLY_RUN check here
|
||||
filtered_models = []
|
||||
for model_case in ALL_OTHER_LORA_MODELS:
|
||||
if "ONLY_RUN" in os.environ and os.environ["ONLY_RUN"] != model_case.base:
|
||||
continue
|
||||
filtered_models.append(model_case)
|
||||
|
||||
self._run_without_cuda_graph_on_model_cases(filtered_models)
|
||||
self._run_cuda_graph_padding_on_model_cases(filtered_models)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
mp.set_start_method("spawn")
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
unittest.main(warnings="ignore")
|
||||
@@ -0,0 +1,63 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
MODELS = [
|
||||
SimpleNamespace(
|
||||
model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
|
||||
tp_size=8,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
|
||||
class TestLlama4LoRA(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
def test_bringup(self):
|
||||
for model in MODELS:
|
||||
try:
|
||||
process = popen_launch_server(
|
||||
model.model,
|
||||
self.base_url,
|
||||
timeout=3 * DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--enable-lora",
|
||||
"--max-lora-rank",
|
||||
"64",
|
||||
"--lora-target-modules",
|
||||
"all",
|
||||
"--tp-size",
|
||||
str(model.tp_size),
|
||||
"--context-length",
|
||||
"262144",
|
||||
"--attention-backend",
|
||||
"fa3",
|
||||
],
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error testing {model.model}: {e}")
|
||||
self.fail(f"Test failed for {model.model}: {e}")
|
||||
|
||||
finally:
|
||||
# Ensure process cleanup happens regardless of success/failure
|
||||
if process is not None and process.poll() is None:
|
||||
print(f"Cleaning up process {process.pid}")
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception as e:
|
||||
print(f"Error killing process: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,229 @@
|
||||
import random
|
||||
import unittest
|
||||
from typing import Sequence
|
||||
|
||||
from utils import TORCH_DTYPES, LoRAAdaptor, LoRAModelCase, ensure_reproducibility
|
||||
|
||||
from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration
|
||||
from sglang.srt.models.qwen3_vl_moe import Qwen3VLMoeForConditionalGeneration
|
||||
from sglang.test.runners import HFRunner, SRTRunner
|
||||
from sglang.test.test_utils import CustomTestCase, calculate_rouge_l
|
||||
|
||||
|
||||
class TestLoRAQwen3VLGating(CustomTestCase):
|
||||
"""Unit tests for should_apply_lora gating on Qwen3‑VL dense and MoE variants."""
|
||||
|
||||
def _assert_pattern(
|
||||
self, pattern, positives: Sequence[str], negatives: Sequence[str]
|
||||
):
|
||||
for name in positives:
|
||||
self.assertTrue(bool(pattern.match(name)), f"Expected to match: {name}")
|
||||
for name in negatives:
|
||||
self.assertFalse(bool(pattern.match(name)), f"Should not match: {name}")
|
||||
|
||||
def test_qwen3_vl_should_apply_lora_regex(self):
|
||||
positives = (
|
||||
"model.layers.0.self_attn.qkv_proj",
|
||||
"model.layers.1.self_attn.o_proj",
|
||||
"model.layers.2.mlp.gate_up_proj",
|
||||
"model.layers.3.mlp.down_proj",
|
||||
)
|
||||
negatives = (
|
||||
"visual.blocks.0.attn.qkv_proj",
|
||||
"model.layers.x.self_attn.qkv_proj",
|
||||
"model.layers.0.attn.qkv_proj",
|
||||
"model.layers.0.mlp.not_proj",
|
||||
"model.layers.0.self_attn.q_proj",
|
||||
)
|
||||
self._assert_pattern(
|
||||
Qwen3VLForConditionalGeneration._lora_pattern, positives, negatives
|
||||
)
|
||||
|
||||
def test_qwen3_vl_moe_should_apply_lora_regex(self):
|
||||
positives = (
|
||||
"model.layers.0.self_attn.qkv_proj",
|
||||
"model.layers.5.self_attn.o_proj",
|
||||
)
|
||||
negatives = (
|
||||
"model.layers.0.mlp.gate_up_proj",
|
||||
"model.layers.0.mlp.down_proj",
|
||||
"visual.blocks.0.attn.qkv_proj",
|
||||
"model.layers.x.self_attn.qkv_proj",
|
||||
"model.layers.0.attn.qkv_proj",
|
||||
)
|
||||
self._assert_pattern(
|
||||
Qwen3VLMoeForConditionalGeneration._lora_pattern_moe, positives, negatives
|
||||
)
|
||||
|
||||
|
||||
TEST_MULTIPLE_BATCH_PROMPTS = [
|
||||
"""
|
||||
### Instruction:
|
||||
Tell me about llamas and alpacas
|
||||
### Response:
|
||||
Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids (camels, dromedaries). Llamas live in the Andean mountains of South America where they graze on grasses and shrubs. Alpaca is another name for domesticated llama. The word "alpaca" comes from an Incan language meaning "golden fleece." Alpacas look very similar to llamas but are smaller than their wild relatives. Both species were used by ancient people as pack animals and for meat. Today both llamas and alpacas are raised primarily for their fiber which can be spun into yarn or knitted into clothing.
|
||||
### Question 2:
|
||||
What do you know about llamas?
|
||||
### Answer:
|
||||
""",
|
||||
"""
|
||||
### Instruction:
|
||||
Write a poem about the transformers Python library.
|
||||
Mention the word "large language models" in that poem.
|
||||
### Response:
|
||||
The Transformers are large language models,
|
||||
They're used to make predictions on text.
|
||||
""",
|
||||
"AI is a field of computer science focused on",
|
||||
"Computer science is the study of",
|
||||
"Write a short story.",
|
||||
"What are the main components of a computer?",
|
||||
]
|
||||
|
||||
|
||||
LORA_MODEL_VARIANTS = [
|
||||
(
|
||||
"Qwen3-VL",
|
||||
LoRAModelCase(
|
||||
base="Qwen/Qwen3-VL-4B-Instruct",
|
||||
adaptors=[
|
||||
LoRAAdaptor(
|
||||
name="mryufei/Qwen3-VL-4B-Instruct-trl-sft",
|
||||
prefill_tolerance=3e-1,
|
||||
),
|
||||
],
|
||||
max_loras_per_batch=1,
|
||||
),
|
||||
),
|
||||
# TODO: Move 30B MoE to 2 GPU runner
|
||||
# (
|
||||
# "Qwen3-VL-MoE",
|
||||
# LoRAModelCase(
|
||||
# base="Qwen/Qwen3-VL-30B-A3B-Instruct",
|
||||
# adaptors=[
|
||||
# LoRAAdaptor(
|
||||
# name="sosoai/qwen3_vl_30b_lora",
|
||||
# prefill_tolerance=3e-1,
|
||||
# ),
|
||||
# ],
|
||||
# max_loras_per_batch=1,
|
||||
# ),
|
||||
# ),
|
||||
]
|
||||
|
||||
LORA_MAX_NEW_TOKENS = 32
|
||||
|
||||
|
||||
def _run_lora_multiple_batch_on_model_cases(
|
||||
model_cases: Sequence[LoRAModelCase], *, max_new_tokens: int, variant_label: str
|
||||
):
|
||||
for model_case in model_cases:
|
||||
for torch_dtype in TORCH_DTYPES:
|
||||
backend = "csgmv"
|
||||
base_path = model_case.base
|
||||
lora_adapter_paths = [adaptor.name for adaptor in model_case.adaptors]
|
||||
|
||||
batches = [
|
||||
(
|
||||
[
|
||||
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
|
||||
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
|
||||
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
|
||||
],
|
||||
[None, lora_adapter_paths[0], None],
|
||||
),
|
||||
(
|
||||
[
|
||||
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
|
||||
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
|
||||
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
|
||||
],
|
||||
[lora_adapter_paths[0], None, None],
|
||||
),
|
||||
(
|
||||
[
|
||||
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
|
||||
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
|
||||
random.choice(TEST_MULTIPLE_BATCH_PROMPTS),
|
||||
],
|
||||
[None, None, None],
|
||||
),
|
||||
]
|
||||
|
||||
print(
|
||||
f"\n=== {variant_label} LoRA parity on '{base_path}', backend={backend}, dtype={torch_dtype} ==="
|
||||
)
|
||||
|
||||
ensure_reproducibility()
|
||||
srt_runner = SRTRunner(
|
||||
base_path,
|
||||
torch_dtype=torch_dtype,
|
||||
model_type="generation",
|
||||
lora_paths=lora_adapter_paths,
|
||||
max_loras_per_batch=model_case.max_loras_per_batch,
|
||||
lora_backend=backend,
|
||||
sleep_on_idle=True,
|
||||
attention_backend="torch_native",
|
||||
disable_radix_cache=True,
|
||||
)
|
||||
|
||||
ensure_reproducibility()
|
||||
hf_runner = HFRunner(
|
||||
base_path,
|
||||
torch_dtype=torch_dtype,
|
||||
model_type="generation",
|
||||
patch_model_do_sample_false=True,
|
||||
)
|
||||
|
||||
with srt_runner, hf_runner:
|
||||
for i, (prompts, lora_paths) in enumerate(batches):
|
||||
print(
|
||||
f"\n--- Running Batch {i + 1} --- prompts: {prompts}, lora_paths: {lora_paths}"
|
||||
)
|
||||
|
||||
srt_outputs = srt_runner.batch_forward(
|
||||
prompts,
|
||||
max_new_tokens=max_new_tokens,
|
||||
lora_paths=lora_paths,
|
||||
)
|
||||
|
||||
hf_outputs = hf_runner.forward(
|
||||
prompts,
|
||||
max_new_tokens=max_new_tokens,
|
||||
lora_paths=lora_paths,
|
||||
)
|
||||
|
||||
print("SRT outputs:", [s for s in srt_outputs.output_strs])
|
||||
print("HF outputs:", [s for s in hf_outputs.output_strs])
|
||||
|
||||
for srt_out, hf_out in zip(
|
||||
srt_outputs.output_strs, hf_outputs.output_strs
|
||||
):
|
||||
srt_str = srt_out.strip()
|
||||
hf_str = hf_out.strip()
|
||||
rouge_tol = model_case.rouge_l_tolerance
|
||||
rouge_score = calculate_rouge_l([srt_str], [hf_str])[0]
|
||||
if rouge_score < rouge_tol:
|
||||
raise AssertionError(
|
||||
f"ROUGE-L score {rouge_score} below tolerance {rouge_tol} "
|
||||
f"for base '{base_path}', adaptor '{lora_paths}', backend '{backend}', prompt: '{prompts}...'"
|
||||
)
|
||||
|
||||
print(f"--- Batch {i + 1} Comparison Passed --- ")
|
||||
|
||||
|
||||
class TestLoRAQwen3VLIntegration(CustomTestCase):
|
||||
"""Parity integration tests for Qwen3‑VL dense and MoE LoRA adapters."""
|
||||
|
||||
def test_ci_lora_models(self):
|
||||
for label, model_case in LORA_MODEL_VARIANTS:
|
||||
with self.subTest(variant=label):
|
||||
_run_lora_multiple_batch_on_model_cases(
|
||||
[model_case],
|
||||
max_new_tokens=LORA_MAX_NEW_TOKENS,
|
||||
variant_label=label,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,78 @@
|
||||
# Copyright 2023-2024 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
|
||||
import multiprocessing as mp
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.test.runners import HFRunner, SRTRunner
|
||||
from sglang.test.test_utils import get_similarities
|
||||
|
||||
TEXTS = "two Subway Series sandwiches with meats, cheese, lettuce, tomatoes, and onions on a black background, accompanied by the Subway Series logo, highlighting a new sandwich series."
|
||||
IMAGES = "https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild/resolve/main/images/023.jpg"
|
||||
MODELS = [
|
||||
("openai/clip-vit-large-patch14-336", 1e-5),
|
||||
]
|
||||
TORCH_DTYPES = [torch.float16]
|
||||
|
||||
|
||||
class TestClipModels(unittest.TestCase):
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
mp.set_start_method("spawn", force=True)
|
||||
|
||||
def assert_close_embeddings(self, model, prefill_tolerance, torch_dtype):
|
||||
|
||||
with HFRunner(
|
||||
model,
|
||||
torch_dtype=torch_dtype,
|
||||
model_type="embedding",
|
||||
) as hf_runner:
|
||||
hf_text_embeds = hf_runner.forward(prompts=TEXTS)
|
||||
hf_image_embeds = hf_runner.forward(image_data=IMAGES)
|
||||
|
||||
with SRTRunner(
|
||||
model,
|
||||
tp_size=1,
|
||||
torch_dtype=torch_dtype,
|
||||
model_type="embedding",
|
||||
) as srt_runner:
|
||||
text_embeds = srt_runner.forward(prompts=TEXTS)
|
||||
image_embeds = srt_runner.forward(prompts="padding", image_data=IMAGES)
|
||||
|
||||
text_similarity = get_similarities(
|
||||
text_embeds.embed_logits[0], hf_text_embeds.embed_logits[0]
|
||||
)
|
||||
image_similarity = get_similarities(
|
||||
image_embeds.embed_logits[0], hf_image_embeds.embed_logits[0]
|
||||
)
|
||||
print("text similarity diff", abs(text_similarity - 1))
|
||||
print("image similarity diff", abs(image_similarity - 1))
|
||||
assert torch.all(
|
||||
abs(text_similarity - 1) < prefill_tolerance
|
||||
), "embeddings are not all close"
|
||||
assert torch.all(
|
||||
abs(image_similarity - 1) < prefill_tolerance
|
||||
), "embeddings are not all close"
|
||||
|
||||
def test_accuracy(self):
|
||||
for model, prefill_tolerance in MODELS:
|
||||
for torch_dtype in TORCH_DTYPES:
|
||||
self.assert_close_embeddings(model, prefill_tolerance, torch_dtype)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,34 @@
|
||||
import unittest
|
||||
|
||||
from sglang.test.test_utils import CustomTestCase, is_in_ci, run_bench_one_batch
|
||||
|
||||
|
||||
class TestDummyGrok1(CustomTestCase):
|
||||
|
||||
def test_dummy_grok_1(self):
|
||||
_, output_throughput, _ = run_bench_one_batch(
|
||||
None,
|
||||
[
|
||||
"--model",
|
||||
"/dummy-grok",
|
||||
"--tokenizer-path",
|
||||
"Xenova/grok-1-tokenizer",
|
||||
"--batch-size",
|
||||
"2",
|
||||
"--tp",
|
||||
"2",
|
||||
"--quantization",
|
||||
"fp8",
|
||||
"--load-format",
|
||||
"dummy",
|
||||
"--json-model-override-args",
|
||||
'{"num_hidden_layers": 2}',
|
||||
],
|
||||
)
|
||||
|
||||
if is_in_ci():
|
||||
self.assertGreater(output_throughput, 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,146 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestFalconH1(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "tiiuae/Falcon-H1-0.5B-Instruct"
|
||||
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=[
|
||||
"--tensor-parallel-size",
|
||||
"1",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.74)
|
||||
|
||||
|
||||
class TestFalconH1TP4(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "tiiuae/Falcon-H1-0.5B-Instruct"
|
||||
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=[
|
||||
"--tensor-parallel-size",
|
||||
"4",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.74)
|
||||
|
||||
|
||||
class TestFalconH1NoGatedRMS(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "tiiuae/Falcon-H1-1.5B-Instruct"
|
||||
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=[
|
||||
"--tensor-parallel-size",
|
||||
"1",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.74)
|
||||
|
||||
|
||||
class TestFalconH1NoGatedTP4(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "tiiuae/Falcon-H1-1.5B-Instruct"
|
||||
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=[
|
||||
"--tensor-parallel-size",
|
||||
"4",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.74)
|
||||
@@ -0,0 +1,85 @@
|
||||
# Copyright 2023-2024 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
|
||||
|
||||
import multiprocessing as mp
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.test.runners import HFRunner, SRTRunner
|
||||
from sglang.test.test_utils import CustomTestCase, get_similarities
|
||||
|
||||
TEXTS = "two Subway Series sandwiches with meats, cheese, lettuce, tomatoes, and onions on a black background, accompanied by the Subway Series logo, highlighting a new sandwich series."
|
||||
IMAGES = "https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild/resolve/main/images/023.jpg"
|
||||
|
||||
|
||||
MODELS = [
|
||||
("Alibaba-NLP/gme-Qwen2-VL-2B-Instruct", 1e-3),
|
||||
]
|
||||
TORCH_DTYPES = [torch.float16]
|
||||
|
||||
|
||||
class TestQmeQwenModels(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
mp.set_start_method("spawn", force=True)
|
||||
|
||||
def assert_close_embeddings(self, model, prefill_tolerance, torch_dtype):
|
||||
|
||||
prompts_no_image = f"<|im_start|>system\nYou are a helpful assistant<|im_end|>\n<|im_start|>user\n{TEXTS}<|im_end|>\n<|im_start|>assistant\n<|endoftext|>"
|
||||
prompts_with_image = f"<|im_start|>system\nYou are a helpful assistant<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>assistant\n<|endoftext|>"
|
||||
with HFRunner(
|
||||
model,
|
||||
torch_dtype=torch_dtype,
|
||||
model_type="embedding",
|
||||
) as hf_runner:
|
||||
hf_text_embeddings = hf_runner.forward(prompts=[prompts_no_image])
|
||||
hf_image_embeddings = hf_runner.forward(
|
||||
prompts=[prompts_with_image], image_data=[IMAGES]
|
||||
)
|
||||
with SRTRunner(
|
||||
model,
|
||||
tp_size=1,
|
||||
torch_dtype=torch_dtype,
|
||||
model_type="embedding",
|
||||
) as srt_runner:
|
||||
srt_text_embeddings = srt_runner.forward(prompts=prompts_no_image)
|
||||
srt_image_embeddings = srt_runner.forward(
|
||||
prompts=prompts_with_image, image_data=IMAGES
|
||||
)
|
||||
|
||||
similarity = get_similarities(
|
||||
hf_text_embeddings.embed_logits[0], srt_text_embeddings.embed_logits[0]
|
||||
)
|
||||
print("texts similarity diff", abs(similarity - 1))
|
||||
assert torch.all(
|
||||
abs(similarity - 1) < prefill_tolerance
|
||||
), "embeddings are not all close"
|
||||
similarity = get_similarities(
|
||||
hf_image_embeddings.embed_logits[0], srt_image_embeddings.embed_logits[0]
|
||||
)
|
||||
print("images similarity diff", abs(similarity - 1))
|
||||
assert torch.all(
|
||||
abs(similarity - 1) < prefill_tolerance
|
||||
), "embeddings are not all close"
|
||||
|
||||
def test_accuracy(self):
|
||||
for model, prefill_tolerance in MODELS:
|
||||
for torch_dtype in TORCH_DTYPES:
|
||||
self.assert_close_embeddings(model, prefill_tolerance, torch_dtype)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,53 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestGrok(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "lmzheng/grok-1"
|
||||
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=[
|
||||
"--load-format",
|
||||
"dummy",
|
||||
"--json-model-override-args",
|
||||
'{"num_hidden_layers": 2}',
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=64,
|
||||
max_new_tokens=256,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
# It is dummy weights so we only assert the output throughput instead of accuracy.
|
||||
self.assertGreater(metrics["output_throughput"], 1000)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,73 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
MODELS = [
|
||||
SimpleNamespace(
|
||||
model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
|
||||
accuracy=0.9,
|
||||
tp_size=4,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class TestLlama4(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
def test_gsm8k(self):
|
||||
|
||||
for model in MODELS:
|
||||
try:
|
||||
process = popen_launch_server(
|
||||
model.model,
|
||||
self.base_url,
|
||||
timeout=3 * DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--chat-template",
|
||||
"llama-4",
|
||||
"--tp-size",
|
||||
str(model.tp_size),
|
||||
"--mem-fraction-static",
|
||||
"0.8",
|
||||
"--context-length",
|
||||
"8192",
|
||||
],
|
||||
)
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreaterEqual(metrics["accuracy"], model.accuracy)
|
||||
except Exception as e:
|
||||
print(f"Error testing {model.model}: {e}")
|
||||
self.fail(f"Test failed for {model.model}: {e}")
|
||||
|
||||
finally:
|
||||
# Ensure process cleanup happens regardless of success/failure
|
||||
if process is not None and process.poll() is None:
|
||||
print(f"Cleaning up process {process.pid}")
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception as e:
|
||||
print(f"Error killing process: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,58 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestMiMoMTP(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "XiaomiMiMo/MiMo-7B-RL"
|
||||
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=[
|
||||
"--trust-remote-code",
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-num-steps",
|
||||
"1",
|
||||
"--speculative-eagle-topk",
|
||||
"1",
|
||||
"--speculative-num-draft-tokens",
|
||||
"2",
|
||||
"--mem-fraction-static",
|
||||
"0.5",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.7)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,213 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestUnslothPhi4(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "unsloth/phi-4"
|
||||
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=[],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.78)
|
||||
|
||||
|
||||
class TestUnslothPhi4Bnb4bit(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "unsloth/phi-4-bnb-4bit"
|
||||
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=[
|
||||
"--load-format",
|
||||
"bitsandbytes",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.75)
|
||||
|
||||
|
||||
class TestUnslothPhi4UnslothBnb4bit(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "unsloth/phi-4-unsloth-bnb-4bit"
|
||||
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=[
|
||||
"--load-format",
|
||||
"bitsandbytes",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.75)
|
||||
|
||||
|
||||
class TestUnslothPhi4MiniInstruct(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "unsloth/Phi-4-mini-instruct"
|
||||
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=[],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.65)
|
||||
|
||||
|
||||
class TestUnslothPhi4MiniBnb4bit(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "unsloth/Phi-4-mini-instruct-bnb-4bit"
|
||||
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=[
|
||||
"--load-format",
|
||||
"bitsandbytes",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.6)
|
||||
|
||||
|
||||
class TestUnslothPhi4MiniUnslothBnb4bit(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "unsloth/Phi-4-mini-instruct-unsloth-bnb-4bit"
|
||||
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=[
|
||||
"--load-format",
|
||||
"bitsandbytes",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.6)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,87 @@
|
||||
import unittest
|
||||
|
||||
from nightly_utils import NightlyBenchmarkRunner
|
||||
|
||||
from sglang.test.test_utils import DEFAULT_URL_FOR_TEST, _parse_int_list_env
|
||||
|
||||
DEEPSEEK_V31_MODEL_PATH = "deepseek-ai/DeepSeek-V3.1"
|
||||
PROFILE_DIR = "performance_profiles_deepseek_v31"
|
||||
|
||||
|
||||
class TestNightlyDeepseekV31Performance(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEEPSEEK_V31_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.batch_sizes = [1, 1, 8, 16, 64]
|
||||
cls.input_lens = tuple(_parse_int_list_env("NIGHTLY_INPUT_LENS", "4096"))
|
||||
cls.output_lens = tuple(_parse_int_list_env("NIGHTLY_OUTPUT_LENS", "512"))
|
||||
|
||||
# Define variant configurations
|
||||
cls.variants = [
|
||||
{
|
||||
"name": "basic",
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"8",
|
||||
"--model-loader-extra-config",
|
||||
'{"enable_multithread_load": true}',
|
||||
],
|
||||
},
|
||||
{
|
||||
"name": "mtp",
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"8",
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-num-steps",
|
||||
"3",
|
||||
"--speculative-eagle-topk",
|
||||
"1",
|
||||
"--speculative-num-draft-tokens",
|
||||
"4",
|
||||
"--mem-frac",
|
||||
"0.7",
|
||||
"--model-loader-extra-config",
|
||||
'{"enable_multithread_load": true}',
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
cls.runner = NightlyBenchmarkRunner(PROFILE_DIR, cls.__name__, cls.base_url)
|
||||
cls.runner.setup_profile_directory()
|
||||
|
||||
def test_bench_one_batch(self):
|
||||
failed_variants = []
|
||||
|
||||
try:
|
||||
for variant_config in self.variants:
|
||||
with self.subTest(variant=variant_config["name"]):
|
||||
results, success = self.runner.run_benchmark_for_model(
|
||||
model_path=self.model,
|
||||
batch_sizes=self.batch_sizes,
|
||||
input_lens=self.input_lens,
|
||||
output_lens=self.output_lens,
|
||||
other_args=variant_config["other_args"],
|
||||
variant=variant_config["name"],
|
||||
)
|
||||
|
||||
if not success:
|
||||
failed_variants.append(variant_config["name"])
|
||||
|
||||
self.runner.add_report(results)
|
||||
finally:
|
||||
self.runner.write_final_report()
|
||||
|
||||
if failed_variants:
|
||||
raise AssertionError(
|
||||
f"Benchmark failed for {self.model} with the following variants: "
|
||||
f"{', '.join(failed_variants)}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,103 @@
|
||||
import unittest
|
||||
|
||||
from nightly_utils import NightlyBenchmarkRunner
|
||||
|
||||
from sglang.test.test_utils import DEFAULT_URL_FOR_TEST, _parse_int_list_env
|
||||
|
||||
DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp"
|
||||
PROFILE_DIR = "performance_profiles_deepseek_v32"
|
||||
|
||||
|
||||
class TestNightlyDeepseekV32Performance(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEEPSEEK_V32_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.batch_sizes = [1, 1, 8, 16, 64]
|
||||
cls.input_lens = tuple(_parse_int_list_env("NIGHTLY_INPUT_LENS", "4096"))
|
||||
cls.output_lens = tuple(_parse_int_list_env("NIGHTLY_OUTPUT_LENS", "512"))
|
||||
|
||||
# Define variant configurations
|
||||
cls.variants = [
|
||||
{
|
||||
"name": "basic",
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"8",
|
||||
"--model-loader-extra-config",
|
||||
'{"enable_multithread_load": true}',
|
||||
],
|
||||
},
|
||||
{
|
||||
"name": "mtp",
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"8",
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-num-steps",
|
||||
"3",
|
||||
"--speculative-eagle-topk",
|
||||
"1",
|
||||
"--speculative-num-draft-tokens",
|
||||
"4",
|
||||
"--mem-frac",
|
||||
"0.7",
|
||||
"--model-loader-extra-config",
|
||||
'{"enable_multithread_load": true}',
|
||||
],
|
||||
},
|
||||
{
|
||||
"name": "nsa",
|
||||
"other_args": [
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"8",
|
||||
"--attention-backend",
|
||||
"nsa",
|
||||
"--nsa-prefill-backend",
|
||||
"flashmla_sparse",
|
||||
"--nsa-decode-backend",
|
||||
"flashmla_kv",
|
||||
"--model-loader-extra-config",
|
||||
'{"enable_multithread_load": true}',
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
cls.runner = NightlyBenchmarkRunner(PROFILE_DIR, cls.__name__, cls.base_url)
|
||||
cls.runner.setup_profile_directory()
|
||||
|
||||
def test_bench_one_batch(self):
|
||||
failed_variants = []
|
||||
|
||||
try:
|
||||
for variant_config in self.variants:
|
||||
with self.subTest(variant=variant_config["name"]):
|
||||
results, success = self.runner.run_benchmark_for_model(
|
||||
model_path=self.model,
|
||||
batch_sizes=self.batch_sizes,
|
||||
input_lens=self.input_lens,
|
||||
output_lens=self.output_lens,
|
||||
other_args=variant_config["other_args"],
|
||||
variant=variant_config["name"],
|
||||
)
|
||||
|
||||
if not success:
|
||||
failed_variants.append(variant_config["name"])
|
||||
|
||||
self.runner.add_report(results)
|
||||
finally:
|
||||
self.runner.write_final_report()
|
||||
|
||||
if failed_variants:
|
||||
raise AssertionError(
|
||||
f"Benchmark failed for {self.model} with the following variants: "
|
||||
f"{', '.join(failed_variants)}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,124 @@
|
||||
import json
|
||||
import unittest
|
||||
import warnings
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP1,
|
||||
DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP2,
|
||||
DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1,
|
||||
DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP2,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
ModelLaunchSettings,
|
||||
check_evaluation_test_results,
|
||||
parse_models,
|
||||
popen_launch_server,
|
||||
write_results_to_json,
|
||||
)
|
||||
|
||||
MODEL_SCORE_THRESHOLDS = {
|
||||
"meta-llama/Llama-3.1-8B-Instruct": 0.82,
|
||||
"mistralai/Mistral-7B-Instruct-v0.3": 0.58,
|
||||
"deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct": 0.85,
|
||||
"google/gemma-2-27b-it": 0.91,
|
||||
"meta-llama/Llama-3.1-70B-Instruct": 0.95,
|
||||
"mistralai/Mixtral-8x7B-Instruct-v0.1": 0.616,
|
||||
"Qwen/Qwen2-57B-A14B-Instruct": 0.86,
|
||||
"neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8": 0.83,
|
||||
"neuralmagic/Mistral-7B-Instruct-v0.3-FP8": 0.54,
|
||||
"neuralmagic/DeepSeek-Coder-V2-Lite-Instruct-FP8": 0.835,
|
||||
"zai-org/GLM-4.5-Air-FP8": 0.75,
|
||||
# The threshold of neuralmagic/gemma-2-2b-it-FP8 should be 0.6, but this model has some accuracy regression.
|
||||
# The fix is tracked at https://github.com/sgl-project/sglang/issues/4324, we set it to 0.50, for now, to make CI green.
|
||||
"neuralmagic/gemma-2-2b-it-FP8": 0.50,
|
||||
"neuralmagic/Meta-Llama-3.1-70B-Instruct-FP8": 0.94,
|
||||
"neuralmagic/Mixtral-8x7B-Instruct-v0.1-FP8": 0.65,
|
||||
"neuralmagic/Qwen2-72B-Instruct-FP8": 0.94,
|
||||
"neuralmagic/Qwen2-57B-A14B-Instruct-FP8": 0.82,
|
||||
}
|
||||
|
||||
|
||||
# Do not use `CustomTestCase` since `test_mgsm_en_all_models` does not want retry
|
||||
class TestNightlyGsm8KEval(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.models = []
|
||||
models_tp1 = parse_models(
|
||||
DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1
|
||||
) + parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP1)
|
||||
for model_path in models_tp1:
|
||||
cls.models.append(ModelLaunchSettings(model_path, tp_size=1))
|
||||
|
||||
models_tp2 = parse_models(
|
||||
DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP2
|
||||
) + parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP2)
|
||||
for model_path in models_tp2:
|
||||
cls.models.append(ModelLaunchSettings(model_path, tp_size=2))
|
||||
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
def test_mgsm_en_all_models(self):
|
||||
warnings.filterwarnings(
|
||||
"ignore", category=ResourceWarning, message="unclosed.*socket"
|
||||
)
|
||||
is_first = True
|
||||
all_results = []
|
||||
for model_setup in self.models:
|
||||
with self.subTest(model=model_setup.model_path):
|
||||
other_args = list(model_setup.extra_args)
|
||||
|
||||
if model_setup.model_path == "meta-llama/Llama-3.1-70B-Instruct":
|
||||
other_args.extend(["--mem-fraction-static", "0.9"])
|
||||
|
||||
process = popen_launch_server(
|
||||
model=model_setup.model_path,
|
||||
other_args=other_args,
|
||||
base_url=self.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
)
|
||||
|
||||
try:
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=model_setup.model_path,
|
||||
eval_name="mgsm_en",
|
||||
num_examples=None,
|
||||
num_threads=1024,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
print(
|
||||
f"{'=' * 42}\n{model_setup.model_path} - metrics={metrics} score={metrics['score']}\n{'=' * 42}\n"
|
||||
)
|
||||
|
||||
write_results_to_json(
|
||||
model_setup.model_path, metrics, "w" if is_first else "a"
|
||||
)
|
||||
is_first = False
|
||||
|
||||
# 0.0 for empty latency
|
||||
all_results.append((model_setup.model_path, metrics["score"], 0.0))
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
try:
|
||||
with open("results.json", "r") as f:
|
||||
print("\nFinal Results from results.json:")
|
||||
print(json.dumps(json.load(f), indent=2))
|
||||
except Exception as e:
|
||||
print(f"Error reading results.json: {e}")
|
||||
|
||||
# Check all scores after collecting all results
|
||||
check_evaluation_test_results(
|
||||
all_results,
|
||||
self.__class__.__name__,
|
||||
model_accuracy_thresholds=MODEL_SCORE_THRESHOLDS,
|
||||
model_count=len(self.models),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,60 @@
|
||||
import unittest
|
||||
|
||||
from nightly_utils import NightlyBenchmarkRunner
|
||||
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
ModelLaunchSettings,
|
||||
_parse_int_list_env,
|
||||
parse_models,
|
||||
)
|
||||
|
||||
PROFILE_DIR = "performance_profiles_text_models"
|
||||
|
||||
|
||||
class TestNightlyTextModelsPerformance(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.models = []
|
||||
# TODO: replace with DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1 or other model lists
|
||||
for model_path in parse_models("meta-llama/Llama-3.1-8B-Instruct"):
|
||||
cls.models.append(ModelLaunchSettings(model_path, tp_size=1))
|
||||
for model_path in parse_models("Qwen/Qwen2-57B-A14B-Instruct"):
|
||||
cls.models.append(ModelLaunchSettings(model_path, tp_size=2))
|
||||
# (parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1), False, False),
|
||||
# (parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP2), False, True),
|
||||
# (parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP1), True, False),
|
||||
# (parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP2), True, True),
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.batch_sizes = [1, 1, 8, 16, 64]
|
||||
cls.input_lens = tuple(_parse_int_list_env("NIGHTLY_INPUT_LENS", "4096"))
|
||||
cls.output_lens = tuple(_parse_int_list_env("NIGHTLY_OUTPUT_LENS", "512"))
|
||||
cls.runner = NightlyBenchmarkRunner(PROFILE_DIR, cls.__name__, cls.base_url)
|
||||
cls.runner.setup_profile_directory()
|
||||
|
||||
def test_bench_one_batch(self):
|
||||
all_model_succeed = True
|
||||
|
||||
for model_setup in self.models:
|
||||
with self.subTest(model=model_setup.model_path):
|
||||
results, success = self.runner.run_benchmark_for_model(
|
||||
model_path=model_setup.model_path,
|
||||
batch_sizes=self.batch_sizes,
|
||||
input_lens=self.input_lens,
|
||||
output_lens=self.output_lens,
|
||||
other_args=model_setup.extra_args,
|
||||
)
|
||||
|
||||
if not success:
|
||||
all_model_succeed = False
|
||||
|
||||
self.runner.add_report(results)
|
||||
|
||||
self.runner.write_final_report()
|
||||
|
||||
if not all_model_succeed:
|
||||
raise AssertionError("Some models failed the perf tests.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,127 @@
|
||||
import json
|
||||
import unittest
|
||||
import warnings
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
ModelEvalMetrics,
|
||||
ModelLaunchSettings,
|
||||
check_evaluation_test_results,
|
||||
popen_launch_server,
|
||||
write_results_to_json,
|
||||
)
|
||||
|
||||
MODEL_THRESHOLDS = {
|
||||
# Conservative thresholds on 100 MMMU samples, especially for latency thresholds
|
||||
ModelLaunchSettings("deepseek-ai/deepseek-vl2-small"): ModelEvalMetrics(
|
||||
0.330, 56.1
|
||||
),
|
||||
ModelLaunchSettings("deepseek-ai/Janus-Pro-7B"): ModelEvalMetrics(0.285, 40.3),
|
||||
ModelLaunchSettings("Efficient-Large-Model/NVILA-8B-hf"): ModelEvalMetrics(
|
||||
0.270, 56.7
|
||||
),
|
||||
ModelLaunchSettings("Efficient-Large-Model/NVILA-Lite-2B-hf"): ModelEvalMetrics(
|
||||
0.270, 23.8
|
||||
),
|
||||
ModelLaunchSettings("google/gemma-3-4b-it"): ModelEvalMetrics(0.360, 10.9),
|
||||
ModelLaunchSettings("google/gemma-3n-E4B-it"): ModelEvalMetrics(0.360, 17.7),
|
||||
ModelLaunchSettings("mistral-community/pixtral-12b"): ModelEvalMetrics(0.360, 16.6),
|
||||
ModelLaunchSettings("moonshotai/Kimi-VL-A3B-Instruct"): ModelEvalMetrics(
|
||||
0.330, 22.3
|
||||
),
|
||||
ModelLaunchSettings("openbmb/MiniCPM-o-2_6"): ModelEvalMetrics(0.330, 29.3),
|
||||
ModelLaunchSettings("openbmb/MiniCPM-v-2_6"): ModelEvalMetrics(0.259, 36.3),
|
||||
ModelLaunchSettings("OpenGVLab/InternVL2_5-2B"): ModelEvalMetrics(0.300, 17.0),
|
||||
ModelLaunchSettings("Qwen/Qwen2-VL-7B-Instruct"): ModelEvalMetrics(0.310, 83.3),
|
||||
ModelLaunchSettings("Qwen/Qwen2.5-VL-7B-Instruct"): ModelEvalMetrics(0.340, 31.9),
|
||||
ModelLaunchSettings(
|
||||
"Qwen/Qwen3-VL-30B-A3B-Instruct", extra_args=["--tp=2"]
|
||||
): ModelEvalMetrics(0.29, 37.0),
|
||||
ModelLaunchSettings(
|
||||
"unsloth/Mistral-Small-3.1-24B-Instruct-2503"
|
||||
): ModelEvalMetrics(0.310, 16.7),
|
||||
ModelLaunchSettings("XiaomiMiMo/MiMo-VL-7B-RL"): ModelEvalMetrics(0.28, 32.0),
|
||||
ModelLaunchSettings("zai-org/GLM-4.1V-9B-Thinking"): ModelEvalMetrics(0.280, 30.4),
|
||||
}
|
||||
|
||||
|
||||
class TestNightlyVLMMmmuEval(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.models = list(MODEL_THRESHOLDS.keys())
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
def test_mmmu_vlm_models(self):
|
||||
warnings.filterwarnings(
|
||||
"ignore", category=ResourceWarning, message="unclosed.*socket"
|
||||
)
|
||||
is_first = True
|
||||
all_results = []
|
||||
|
||||
for model in self.models:
|
||||
model_path = model.model_path
|
||||
with self.subTest(model=model_path):
|
||||
process = popen_launch_server(
|
||||
model=model_path,
|
||||
base_url=self.base_url,
|
||||
other_args=model.extra_args,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
)
|
||||
try:
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=model_path,
|
||||
eval_name="mmmu",
|
||||
num_examples=100,
|
||||
num_threads=64,
|
||||
max_tokens=30,
|
||||
)
|
||||
|
||||
args.return_latency = True
|
||||
|
||||
metrics, latency = run_eval(args)
|
||||
|
||||
metrics["score"] = round(metrics["score"], 4)
|
||||
metrics["latency"] = round(latency, 4)
|
||||
print(
|
||||
f"{'=' * 42}\n{model_path} - metrics={metrics} score={metrics['score']}\n{'=' * 42}\n"
|
||||
)
|
||||
|
||||
write_results_to_json(model_path, metrics, "w" if is_first else "a")
|
||||
is_first = False
|
||||
|
||||
all_results.append(
|
||||
(model_path, metrics["score"], metrics["latency"])
|
||||
)
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
try:
|
||||
with open("results.json", "r") as f:
|
||||
print("\nFinal Results from results.json:")
|
||||
print(json.dumps(json.load(f), indent=2))
|
||||
except Exception as e:
|
||||
print(f"Error reading results: {e}")
|
||||
|
||||
model_accuracy_thresholds = {
|
||||
model.model_path: threshold.accuracy
|
||||
for model, threshold in MODEL_THRESHOLDS.items()
|
||||
}
|
||||
model_latency_thresholds = {
|
||||
model.model_path: threshold.eval_time
|
||||
for model, threshold in MODEL_THRESHOLDS.items()
|
||||
}
|
||||
check_evaluation_test_results(
|
||||
all_results,
|
||||
self.__class__.__name__,
|
||||
model_accuracy_thresholds=model_accuracy_thresholds,
|
||||
model_latency_thresholds=model_latency_thresholds,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,88 @@
|
||||
import os
|
||||
import unittest
|
||||
import warnings
|
||||
|
||||
from nightly_utils import NightlyBenchmarkRunner
|
||||
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
ModelLaunchSettings,
|
||||
_parse_int_list_env,
|
||||
parse_models,
|
||||
)
|
||||
|
||||
PROFILE_DIR = "performance_profiles_vlms"
|
||||
|
||||
MODEL_DEFAULTS = [
|
||||
# Keep conservative defaults. Can be overridden by env NIGHTLY_VLM_MODELS
|
||||
ModelLaunchSettings(
|
||||
"Qwen/Qwen2.5-VL-7B-Instruct",
|
||||
extra_args=["--mem-fraction-static=0.7"],
|
||||
),
|
||||
ModelLaunchSettings(
|
||||
"google/gemma-3-27b-it",
|
||||
),
|
||||
ModelLaunchSettings("Qwen/Qwen3-VL-30B-A3B-Instruct", extra_args=["--tp=2"]),
|
||||
# "OpenGVLab/InternVL2_5-2B",
|
||||
# buggy in official transformers impl
|
||||
# "openbmb/MiniCPM-V-2_6",
|
||||
]
|
||||
|
||||
|
||||
class TestNightlyVLMModelsPerformance(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
warnings.filterwarnings(
|
||||
"ignore", category=ResourceWarning, message="unclosed.*socket"
|
||||
)
|
||||
|
||||
nightly_vlm_models_str = os.environ.get("NIGHTLY_VLM_MODELS")
|
||||
if nightly_vlm_models_str:
|
||||
cls.models = []
|
||||
model_paths = parse_models(nightly_vlm_models_str)
|
||||
for model_path in model_paths:
|
||||
cls.models.append(ModelLaunchSettings(model_path))
|
||||
else:
|
||||
cls.models = MODEL_DEFAULTS
|
||||
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
cls.batch_sizes = _parse_int_list_env("NIGHTLY_VLM_BATCH_SIZES", "1,1,2,8,16")
|
||||
cls.input_lens = tuple(_parse_int_list_env("NIGHTLY_VLM_INPUT_LENS", "4096"))
|
||||
cls.output_lens = tuple(_parse_int_list_env("NIGHTLY_VLM_OUTPUT_LENS", "512"))
|
||||
cls.runner = NightlyBenchmarkRunner(PROFILE_DIR, cls.__name__, cls.base_url)
|
||||
cls.runner.setup_profile_directory()
|
||||
|
||||
def test_bench_one_batch(self):
|
||||
all_model_succeed = True
|
||||
|
||||
for model_setup in self.models:
|
||||
with self.subTest(model=model_setup.model_path):
|
||||
# VLMs need additional benchmark args for dataset and trust-remote-code
|
||||
extra_bench_args = [
|
||||
"--trust-remote-code",
|
||||
"--dataset-name=mmmu",
|
||||
]
|
||||
|
||||
results, success = self.runner.run_benchmark_for_model(
|
||||
model_path=model_setup.model_path,
|
||||
batch_sizes=self.batch_sizes,
|
||||
input_lens=self.input_lens,
|
||||
output_lens=self.output_lens,
|
||||
other_args=model_setup.extra_args,
|
||||
extra_bench_args=extra_bench_args,
|
||||
)
|
||||
|
||||
if not success:
|
||||
all_model_succeed = False
|
||||
|
||||
self.runner.add_report(results)
|
||||
|
||||
self.runner.write_final_report()
|
||||
|
||||
if not all_model_succeed:
|
||||
raise AssertionError("Some models failed the perf tests.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,266 @@
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import subprocess
|
||||
import sys
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
MODELS = [
|
||||
SimpleNamespace(model="Qwen/Qwen2.5-VL-7B-Instruct", mmmu_accuracy=0.60),
|
||||
]
|
||||
|
||||
|
||||
# Set default mem_fraction_static to 0.8
|
||||
DEFAULT_MEM_FRACTION_STATIC = 0.8
|
||||
|
||||
|
||||
class TestVLMPiecewiseCudaGraph(CustomTestCase):
|
||||
parsed_args = None # Class variable to store args
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# Removed argument parsing from here
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
cls.time_out = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
|
||||
|
||||
if cls.parsed_args is None:
|
||||
cls.parsed_args = SimpleNamespace(
|
||||
mem_fraction_static=DEFAULT_MEM_FRACTION_STATIC
|
||||
)
|
||||
|
||||
# Set OpenAI API key and base URL environment variables. Needed for lmm-evals to work.
|
||||
os.environ["OPENAI_API_KEY"] = cls.api_key
|
||||
os.environ["OPENAI_API_BASE"] = f"{cls.base_url}/v1"
|
||||
|
||||
def run_mmmu_eval(
|
||||
self,
|
||||
model_version: str,
|
||||
output_path: str,
|
||||
*,
|
||||
env: dict | None = None,
|
||||
):
|
||||
"""
|
||||
Evaluate a VLM on the MMMU validation set with lmms‑eval.
|
||||
Only `model_version` (checkpoint) and `chat_template` vary;
|
||||
We are focusing only on the validation set due to resource constraints.
|
||||
"""
|
||||
# -------- fixed settings --------
|
||||
model = "openai_compatible"
|
||||
tp = 1
|
||||
tasks = "mmmu_val"
|
||||
batch_size = 32
|
||||
log_suffix = "openai_compatible"
|
||||
os.makedirs(output_path, exist_ok=True)
|
||||
|
||||
# -------- compose --model_args --------
|
||||
model_args = f'model_version="{model_version}",' f"tp={tp}"
|
||||
|
||||
# -------- build command list --------
|
||||
cmd = [
|
||||
"python3",
|
||||
"-m",
|
||||
"lmms_eval",
|
||||
"--model",
|
||||
model,
|
||||
"--model_args",
|
||||
model_args,
|
||||
"--tasks",
|
||||
tasks,
|
||||
"--batch_size",
|
||||
str(batch_size),
|
||||
"--output_path",
|
||||
str(output_path),
|
||||
]
|
||||
|
||||
subprocess.run(
|
||||
cmd,
|
||||
check=True,
|
||||
timeout=3600,
|
||||
)
|
||||
|
||||
def _run_vlm_mmmu_test(
|
||||
self,
|
||||
model,
|
||||
output_path,
|
||||
test_name="",
|
||||
custom_env=None,
|
||||
log_level="info",
|
||||
capture_output=False,
|
||||
):
|
||||
"""
|
||||
Common method to run VLM MMMU benchmark test.
|
||||
Args:
|
||||
model: Model to test
|
||||
output_path: Path for output logs
|
||||
test_name: Optional test name for logging
|
||||
custom_env: Optional custom environment variables
|
||||
log_level: Log level for server (default: "info")
|
||||
capture_output: Whether to capture server stdout/stderr
|
||||
"""
|
||||
print(f"\nTesting model: {model.model}{test_name}")
|
||||
|
||||
process = None
|
||||
mmmu_accuracy = 0 # Initialize to handle potential exceptions
|
||||
server_output = ""
|
||||
|
||||
try:
|
||||
# Prepare environment variables
|
||||
process_env = os.environ.copy()
|
||||
if custom_env:
|
||||
process_env.update(custom_env)
|
||||
# if test vlm with cuda_ipc feature, open this env_var
|
||||
process_env["SGLANG_USE_CUDA_IPC_TRANSPORT"] = "1"
|
||||
|
||||
# Prepare stdout/stderr redirection if needed
|
||||
stdout_file = None
|
||||
stderr_file = None
|
||||
if capture_output:
|
||||
stdout_file = open("/tmp/server_stdout.log", "w")
|
||||
stderr_file = open("/tmp/server_stderr.log", "w")
|
||||
|
||||
# Launch server for testing
|
||||
process = popen_launch_server(
|
||||
model.model,
|
||||
base_url=self.base_url,
|
||||
timeout=self.time_out,
|
||||
api_key=self.api_key,
|
||||
other_args=[
|
||||
"--trust-remote-code",
|
||||
"--piecewise-cuda-graph-max-tokens",
|
||||
"8192",
|
||||
"--enable-piecewise-cuda-graph",
|
||||
"--tp=8",
|
||||
"--piecewise-cuda-graph-compiler=eager",
|
||||
"--disable-radix-cache",
|
||||
"--log-level",
|
||||
log_level,
|
||||
],
|
||||
env=process_env,
|
||||
return_stdout_stderr=(
|
||||
(stdout_file, stderr_file) if capture_output else None
|
||||
),
|
||||
)
|
||||
|
||||
# Run evaluation
|
||||
self.run_mmmu_eval(model.model, output_path)
|
||||
|
||||
# Get the result file
|
||||
# Search recursively for JSON result files (lmms-eval v0.4.1+ creates subdirectories)
|
||||
result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
|
||||
if not result_files:
|
||||
result_files = glob.glob(f"{output_path}/*.json")
|
||||
|
||||
if not result_files:
|
||||
raise FileNotFoundError(f"No JSON result files found in {output_path}")
|
||||
|
||||
result_file_path = result_files[0]
|
||||
|
||||
with open(result_file_path, "r") as f:
|
||||
result = json.load(f)
|
||||
print(f"Result{test_name}\n: {result}")
|
||||
|
||||
# Process the result
|
||||
mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
|
||||
print(
|
||||
f"Model {model.model} achieved accuracy{test_name}: {mmmu_accuracy:.4f}"
|
||||
)
|
||||
|
||||
# Capture server output if requested
|
||||
if capture_output and process:
|
||||
server_output = self._read_output_from_files()
|
||||
|
||||
# Assert performance meets expected threshold
|
||||
self.assertGreaterEqual(
|
||||
mmmu_accuracy,
|
||||
model.mmmu_accuracy,
|
||||
f"Model {model.model} accuracy ({mmmu_accuracy:.4f}) below expected threshold ({model.mmmu_accuracy:.4f}){test_name}",
|
||||
)
|
||||
|
||||
return server_output
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error testing {model.model}{test_name}: {e}")
|
||||
self.fail(f"Test failed for {model.model}{test_name}: {e}")
|
||||
|
||||
finally:
|
||||
# Ensure process cleanup happens regardless of success/failure
|
||||
if process is not None and process.poll() is None:
|
||||
print(f"Cleaning up process {process.pid}")
|
||||
try:
|
||||
kill_process_tree(process.pid)
|
||||
except Exception as e:
|
||||
print(f"Error killing process: {e}")
|
||||
|
||||
# clean up temporary files
|
||||
if capture_output:
|
||||
if stdout_file:
|
||||
stdout_file.close()
|
||||
if stderr_file:
|
||||
stderr_file.close()
|
||||
for filename in ["/tmp/server_stdout.log", "/tmp/server_stderr.log"]:
|
||||
try:
|
||||
if os.path.exists(filename):
|
||||
os.remove(filename)
|
||||
except Exception as e:
|
||||
print(f"Error removing {filename}: {e}")
|
||||
|
||||
def _read_output_from_files(self):
|
||||
output_lines = []
|
||||
|
||||
log_files = [
|
||||
("/tmp/server_stdout.log", "[STDOUT]"),
|
||||
("/tmp/server_stderr.log", "[STDERR]"),
|
||||
]
|
||||
for filename, tag in log_files:
|
||||
try:
|
||||
if os.path.exists(filename):
|
||||
with open(filename, "r") as f:
|
||||
for line in f:
|
||||
output_lines.append(f"{tag} {line.rstrip()}")
|
||||
except Exception as e:
|
||||
print(f"Error reading {tag.lower()} file: {e}")
|
||||
|
||||
return "\n".join(output_lines)
|
||||
|
||||
def test_vlm_mmmu_benchmark(self):
|
||||
"""Test VLM models against MMMU benchmark."""
|
||||
models_to_test = MODELS
|
||||
|
||||
if is_in_ci():
|
||||
models_to_test = [random.choice(MODELS)]
|
||||
|
||||
for model in models_to_test:
|
||||
self._run_vlm_mmmu_test(model, "./logs")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Define and parse arguments here, before unittest.main
|
||||
parser = argparse.ArgumentParser(description="Test VLM models")
|
||||
parser.add_argument(
|
||||
"--mem-fraction-static",
|
||||
type=float,
|
||||
help="Static memory fraction for the model",
|
||||
default=DEFAULT_MEM_FRACTION_STATIC,
|
||||
)
|
||||
|
||||
# Parse args intended for unittest
|
||||
args = parser.parse_args()
|
||||
|
||||
# Store the parsed args object on the class
|
||||
TestVLMPiecewiseCudaGraph.parsed_args = args
|
||||
|
||||
# Pass args to unittest
|
||||
unittest.main(argv=[sys.argv[0]])
|
||||
@@ -0,0 +1,289 @@
|
||||
import unittest
|
||||
|
||||
import openai
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestCacheReport(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.min_cached = 5
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=300,
|
||||
other_args=[
|
||||
"--chunked-prefill-size=40",
|
||||
"--enable-cache-report",
|
||||
],
|
||||
)
|
||||
cls.client = openai.Client(api_key="EMPTY", base_url=f"{cls.base_url}/v1")
|
||||
cls.aclient = openai.AsyncClient(api_key="EMPTY", base_url=f"{cls.base_url}/v1")
|
||||
|
||||
usage = cls.run_openai(cls, "1").usage
|
||||
# we can assume that our request is of size 1, plus the total template size
|
||||
# ideally we would like to know the begin size / end size of the template to be more precise
|
||||
total_template_size = usage.prompt_tokens - 1
|
||||
print(f"template size: {total_template_size}")
|
||||
usage2 = cls.run_openai(cls, "2").usage
|
||||
assert usage2.prompt_tokens_details.cached_tokens <= total_template_size
|
||||
cls.min_cached = max(
|
||||
usage2.prompt_tokens_details.cached_tokens,
|
||||
total_template_size - usage2.prompt_tokens_details.cached_tokens,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def run_decode(self, return_logprob=False, top_logprobs_num=0, n=1):
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
# we use an uncommon start to minimise the chance that the cache is hit by chance
|
||||
json={
|
||||
"text": "_ The capital of France is",
|
||||
"sampling_params": {
|
||||
"temperature": 0 if n == 1 else 0.5,
|
||||
"max_new_tokens": 128,
|
||||
"n": n,
|
||||
"stop_token_ids": [119690],
|
||||
},
|
||||
"stream": False,
|
||||
"return_logprob": return_logprob,
|
||||
"top_logprobs_num": top_logprobs_num,
|
||||
"logprob_start_len": 0,
|
||||
},
|
||||
)
|
||||
return response
|
||||
|
||||
def run_openai(self, message):
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
# {"role": "system", "content": "You are a helpful AI assistant"},
|
||||
{"role": "user", "content": message},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=100,
|
||||
)
|
||||
return response
|
||||
|
||||
async def run_openai_async(self, message):
|
||||
response = await self.aclient.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "user", "content": message},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=100,
|
||||
)
|
||||
return response
|
||||
|
||||
def cache_report_openai(self, message):
|
||||
response = self.run_openai(message)
|
||||
print(
|
||||
f"openai first request cached_tokens: {int(response.usage.prompt_tokens_details.cached_tokens)}"
|
||||
)
|
||||
first_cached_tokens = int(response.usage.prompt_tokens_details.cached_tokens)
|
||||
# assert int(response.usage.cached_tokens) == 0
|
||||
assert first_cached_tokens <= self.min_cached
|
||||
response = self.run_openai(message)
|
||||
cached_tokens = int(response.usage.prompt_tokens_details.cached_tokens)
|
||||
print(f"openai second request cached_tokens: {cached_tokens}")
|
||||
assert cached_tokens > 0
|
||||
assert cached_tokens == int(response.usage.prompt_tokens) - 1
|
||||
return first_cached_tokens
|
||||
|
||||
async def cache_report_openai_async(self, message):
|
||||
response = await self.run_openai_async(message)
|
||||
cached_tokens = int(response.usage.prompt_tokens_details.cached_tokens)
|
||||
prompt_tokens = int(response.usage.prompt_tokens)
|
||||
return cached_tokens, prompt_tokens
|
||||
|
||||
def test_generate(self):
|
||||
print("=" * 100)
|
||||
response = self.run_decode()
|
||||
# print(response.json())
|
||||
cached_tokens = int(response.json()["meta_info"]["cached_tokens"])
|
||||
print(f"sglang first request cached_tokens: {cached_tokens}")
|
||||
print(
|
||||
f"sglang first request prompt_tokens: {int(response.json()['meta_info']['prompt_tokens'])}"
|
||||
)
|
||||
# can't assure to be 0: depends on the initialisation request / if a template is used with the model
|
||||
assert cached_tokens < self.min_cached
|
||||
response = self.run_decode()
|
||||
cached_tokens = int(response.json()["meta_info"]["cached_tokens"])
|
||||
print(f"sglang second request cached_tokens: {cached_tokens}")
|
||||
print(
|
||||
f"sglang second request prompt_tokens: {int(response.json()['meta_info']['prompt_tokens'])}"
|
||||
)
|
||||
assert cached_tokens == int(response.json()["meta_info"]["prompt_tokens"]) - 1
|
||||
|
||||
def test_cache_split_prefill_openai(self):
|
||||
print("=" * 100)
|
||||
self.cache_report_openai(
|
||||
"€ This is a very long and unique text that should not be already cached, the twist is"
|
||||
" that it should be longer than the chunked-prefill-size, so it should be split among"
|
||||
" several prefill requests. Still, it shouldn't be cached"
|
||||
)
|
||||
|
||||
def test_cache_report_openai(self):
|
||||
print("=" * 100)
|
||||
# warm up the cache, for the template
|
||||
self.run_openai("Introduce the capital of France.")
|
||||
|
||||
first_cached_tokens_1 = self.run_openai(
|
||||
"How many sparrow do you need to lift a coconut?"
|
||||
).usage.prompt_tokens_details.cached_tokens
|
||||
|
||||
usage_2 = self.run_openai("* sing something about cats").usage
|
||||
first_cached_tokens_2 = usage_2.prompt_tokens_details.cached_tokens
|
||||
# first request may not have 0 cached tokens, but if they only have the template in common they
|
||||
# should be the same once the cache is warmed up
|
||||
assert first_cached_tokens_1 == first_cached_tokens_2
|
||||
|
||||
resp = self.run_openai("* sing something about cats and dogs")
|
||||
print(resp.usage)
|
||||
|
||||
resp = self.run_openai("* sing something about cats, please")
|
||||
print(resp.usage)
|
||||
assert (
|
||||
resp.usage.prompt_tokens_details.cached_tokens
|
||||
>= usage_2.prompt_tokens - self.min_cached
|
||||
)
|
||||
|
||||
# TODO: flaky test
|
||||
# def test_cache_report_openai_async(self):
|
||||
# print("=" * 100)
|
||||
|
||||
# async def run_test():
|
||||
# task0 = asyncio.create_task(
|
||||
# self.cache_report_openai_async(
|
||||
# "first request, to start the inference and let the next two request be started in the same batch"
|
||||
# )
|
||||
# )
|
||||
# await asyncio.sleep(1) # to force the first request to be started first
|
||||
# task1 = asyncio.create_task(
|
||||
# self.cache_report_openai_async(
|
||||
# "> can the same batch parallel request use the cache?"
|
||||
# )
|
||||
# )
|
||||
# task2 = asyncio.create_task(
|
||||
# self.cache_report_openai_async(
|
||||
# "> can the same batch parallel request use the cache?"
|
||||
# )
|
||||
# )
|
||||
# result0, result1, result2 = await asyncio.gather(task0, task1, task2)
|
||||
|
||||
# cached_tokens0, prompt_tokens0 = result0
|
||||
# cached_tokens1, prompt_tokens1 = result1
|
||||
# cached_tokens2, prompt_tokens2 = result2
|
||||
|
||||
# print(
|
||||
# f"Async request 0 - Cached tokens: {cached_tokens0}, Prompt tokens: {prompt_tokens0}"
|
||||
# )
|
||||
# print(
|
||||
# f"Async request 1 - Cached tokens: {cached_tokens1}, Prompt tokens: {prompt_tokens1}"
|
||||
# )
|
||||
# print(
|
||||
# f"Async request 2 - Cached tokens: {cached_tokens2}, Prompt tokens: {prompt_tokens2}"
|
||||
# )
|
||||
|
||||
# # Assert that no requests used the cache (because first is alone, and the next two are in the same batch)
|
||||
# # If a new optimisation limiting starting request with same prefix at the same time was added
|
||||
# # to maximise the cache hit, this would not be true
|
||||
# assert cached_tokens1 == cached_tokens2 == cached_tokens0
|
||||
|
||||
# asyncio.run(run_test())
|
||||
|
||||
def test_cache_salt_effectiveness(self):
|
||||
print("=" * 100)
|
||||
print("Testing cache_salt effectiveness")
|
||||
|
||||
# Use a unique message to avoid interference with other tests
|
||||
test_message = "What is the capital of Japan?"
|
||||
|
||||
# First request with cache_salt "salt1"
|
||||
response1 = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": test_message}],
|
||||
temperature=0,
|
||||
max_tokens=10,
|
||||
extra_body={"cache_salt": "salt1"},
|
||||
)
|
||||
cached_tokens_1_first = int(response1.usage.prompt_tokens_details.cached_tokens)
|
||||
prompt_tokens_1 = int(response1.usage.prompt_tokens)
|
||||
print(
|
||||
f"First request with salt1 - cached_tokens: {cached_tokens_1_first}, prompt_tokens: {prompt_tokens_1}"
|
||||
)
|
||||
|
||||
# Second request with same cache_salt "salt1" - should get cache hit
|
||||
response2 = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": test_message}],
|
||||
temperature=0,
|
||||
max_tokens=10,
|
||||
extra_body={"cache_salt": "salt1"},
|
||||
)
|
||||
cached_tokens_1_second = int(
|
||||
response2.usage.prompt_tokens_details.cached_tokens
|
||||
)
|
||||
print(
|
||||
f"Second request with salt1 - cached_tokens: {cached_tokens_1_second}, prompt_tokens: {prompt_tokens_1}"
|
||||
)
|
||||
|
||||
# Verify cache hit for same salt
|
||||
assert (
|
||||
cached_tokens_1_second > cached_tokens_1_first
|
||||
), "Should have cache hit with same cache_salt"
|
||||
assert (
|
||||
cached_tokens_1_second == prompt_tokens_1 - 1
|
||||
), "Should cache all prompt tokens except the last one"
|
||||
|
||||
# Third request with different cache_salt "salt2" - should not get cache hit
|
||||
response3 = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": test_message}],
|
||||
temperature=0,
|
||||
max_tokens=10,
|
||||
extra_body={"cache_salt": "salt2"},
|
||||
)
|
||||
cached_tokens_2_first = int(response3.usage.prompt_tokens_details.cached_tokens)
|
||||
print(f"First request with salt2 - cached_tokens: {cached_tokens_2_first}")
|
||||
|
||||
# Verify no cache hit for different salt (should be similar to first request with salt1)
|
||||
assert (
|
||||
cached_tokens_2_first <= cached_tokens_1_first + self.min_cached
|
||||
), "Different cache_salt should not share cache"
|
||||
|
||||
# Fourth request with same cache_salt "salt2" - should now get cache hit
|
||||
response4 = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": test_message}],
|
||||
temperature=0,
|
||||
max_tokens=10,
|
||||
extra_body={"cache_salt": "salt2"},
|
||||
)
|
||||
cached_tokens_2_second = int(
|
||||
response4.usage.prompt_tokens_details.cached_tokens
|
||||
)
|
||||
print(f"Second request with salt2 - cached_tokens: {cached_tokens_2_second}")
|
||||
|
||||
# Verify cache hit for salt2
|
||||
assert (
|
||||
cached_tokens_2_second == cached_tokens_2_first
|
||||
), "Should have cache hit with same cache_salt for salt2"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,105 @@
|
||||
import asyncio
|
||||
import unittest
|
||||
|
||||
import openai
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestContinuousUsageStats(CustomTestCase):
|
||||
@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=300)
|
||||
cls.client = openai.Client(api_key="EMPTY", base_url=f"{cls.base_url}/v1")
|
||||
cls.aclient = openai.AsyncClient(api_key="EMPTY", base_url=f"{cls.base_url}/v1")
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_continuous_usage_stats_enabled(self):
|
||||
stream = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": "What is machine learning?"}],
|
||||
stream=True,
|
||||
max_tokens=30,
|
||||
temperature=0,
|
||||
stream_options={"include_usage": True, "continuous_usage_stats": True},
|
||||
)
|
||||
|
||||
chunks_with_usage = 0
|
||||
chunks_with_content = 0
|
||||
last_usage = None
|
||||
|
||||
for chunk in stream:
|
||||
has_content = len(chunk.choices) > 0 and chunk.choices[0].delta.content
|
||||
if chunk.usage:
|
||||
chunks_with_usage += 1
|
||||
last_usage = chunk.usage
|
||||
if has_content:
|
||||
chunks_with_content += 1
|
||||
|
||||
assert chunks_with_content > 0
|
||||
assert chunks_with_usage >= chunks_with_content
|
||||
assert last_usage.prompt_tokens > 0
|
||||
assert last_usage.completion_tokens > 0
|
||||
assert (
|
||||
last_usage.total_tokens
|
||||
== last_usage.prompt_tokens + last_usage.completion_tokens
|
||||
)
|
||||
|
||||
async def test_continuous_usage_stats_async(self):
|
||||
stream = await self.aclient.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": "What is deep learning?"}],
|
||||
stream=True,
|
||||
max_tokens=30,
|
||||
temperature=0,
|
||||
stream_options={"include_usage": True, "continuous_usage_stats": True},
|
||||
)
|
||||
|
||||
chunks_with_usage = 0
|
||||
chunks_with_content = 0
|
||||
|
||||
async for chunk in stream:
|
||||
has_content = len(chunk.choices) > 0 and chunk.choices[0].delta.content
|
||||
if chunk.usage:
|
||||
chunks_with_usage += 1
|
||||
if has_content:
|
||||
chunks_with_content += 1
|
||||
|
||||
assert chunks_with_content > 0
|
||||
assert chunks_with_usage >= chunks_with_content
|
||||
|
||||
def test_continuous_usage_stats_disabled(self):
|
||||
stream = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": "What is AI?"}],
|
||||
stream=True,
|
||||
max_tokens=30,
|
||||
temperature=0,
|
||||
stream_options={"include_usage": True, "continuous_usage_stats": False},
|
||||
)
|
||||
|
||||
usage_chunks = []
|
||||
for chunk in stream:
|
||||
if chunk.usage:
|
||||
usage_chunks.append(chunk)
|
||||
|
||||
assert len(usage_chunks) == 1
|
||||
assert len(usage_chunks[0].choices) == 0
|
||||
|
||||
def test_async_runner(self):
|
||||
asyncio.run(self.test_continuous_usage_stats_async())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,126 @@
|
||||
"""
|
||||
python3 -m unittest test.srt.openai_server.features.test_structural_tag
|
||||
"""
|
||||
|
||||
import json
|
||||
import unittest
|
||||
from typing import Any
|
||||
|
||||
import openai
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
def setup_class(cls, backend: str):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
other_args = [
|
||||
"--max-running-requests",
|
||||
"10",
|
||||
"--grammar-backend",
|
||||
backend,
|
||||
]
|
||||
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
)
|
||||
|
||||
|
||||
class TestStructuralTagXGrammarBackend(CustomTestCase):
|
||||
model: str
|
||||
base_url: str
|
||||
process: Any
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
setup_class(cls, backend="xgrammar")
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_stag_constant_str_openai(self):
|
||||
client = openai.Client(api_key="EMPTY", base_url=f"{self.base_url}/v1")
|
||||
|
||||
# even when the answer is ridiculous, the model should follow the instruction
|
||||
answer = "The capital of France is Berlin."
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful AI assistant"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Introduce the capital of France. Return in a JSON format.",
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=128,
|
||||
response_format={
|
||||
"type": "structural_tag",
|
||||
"format": {
|
||||
"type": "const_string",
|
||||
"value": answer,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
text = response.choices[0].message.content
|
||||
self.assertEqual(text, answer)
|
||||
|
||||
def test_stag_json_schema_openai(self):
|
||||
client = openai.Client(api_key="EMPTY", base_url=f"{self.base_url}/v1")
|
||||
json_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string", "pattern": "^[\\w]+$"},
|
||||
"population": {"type": "integer"},
|
||||
},
|
||||
"required": ["name", "population"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful AI assistant"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Introduce the capital of France. Return in a JSON format.",
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=128,
|
||||
response_format={
|
||||
"type": "structural_tag",
|
||||
"format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": json_schema,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
text = response.choices[0].message.content
|
||||
try:
|
||||
js_obj = json.loads(text)
|
||||
except (TypeError, json.decoder.JSONDecodeError):
|
||||
print("JSONDecodeError", text)
|
||||
raise
|
||||
|
||||
self.assertIsInstance(js_obj["name"], str)
|
||||
self.assertIsInstance(js_obj["population"], int)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,114 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestFp8KvcacheBase(CustomTestCase):
|
||||
model_config = None
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if cls.model_config is None:
|
||||
raise NotImplementedError("model_config must be specified in subclass")
|
||||
|
||||
cls.model = cls.model_config["model_name"]
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
dirpath = os.path.dirname(__file__)
|
||||
config_file = os.path.join(dirpath, cls.model_config["config_filename"])
|
||||
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--kv-cache-dtype",
|
||||
"fp8_e4m3",
|
||||
"--quantization-param-path",
|
||||
config_file,
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class TestFp8KvcacheLlama(TestFp8KvcacheBase):
|
||||
model_config = {
|
||||
"model_name": DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
"config_filename": "kv_cache_scales_llama3_8b.json",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_mgsm_en(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mgsm_en",
|
||||
num_examples=None,
|
||||
num_threads=1024,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreater(metrics["score"], 0.80)
|
||||
|
||||
def test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreaterEqual(metrics["score"], 0.65)
|
||||
|
||||
|
||||
class TestFp8KvcacheQwen(TestFp8KvcacheBase):
|
||||
model_config = {
|
||||
"model_name": DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
|
||||
"config_filename": "kv_cache_scales_qwen2_1_5b.json",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_mgsm_en(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mgsm_en",
|
||||
num_examples=None,
|
||||
num_threads=1024,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreater(metrics["score"], 0.01)
|
||||
|
||||
def test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreaterEqual(metrics["score"], 0.3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,276 @@
|
||||
import multiprocessing
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import random
|
||||
import traceback
|
||||
import unittest
|
||||
from multiprocessing import Process
|
||||
|
||||
import torch
|
||||
from torch.distributed.device_mesh import init_device_mesh
|
||||
from torch.distributed.fsdp import CPUOffload
|
||||
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
|
||||
from torch.distributed.fsdp import MixedPrecision
|
||||
from torch.distributed.fsdp.api import (
|
||||
ShardedStateDictConfig,
|
||||
ShardingStrategy,
|
||||
StateDictType,
|
||||
)
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
from sglang.srt.entrypoints.verl_engine import VerlEngine
|
||||
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
||||
from sglang.test.runners import (
|
||||
HFRunner,
|
||||
SRTRunner,
|
||||
check_close_model_outputs,
|
||||
get_dtype_str,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase, find_available_port, is_in_ci
|
||||
|
||||
_MAX_NEW_TOKENS = 8
|
||||
_PROMPTS = ["1+1=2, 1+2=3, 1+3=4, 1+4=5, 1+5=", "1*1=1, 1*2=2, 1*3=3, 1*4=4, 1*5="]
|
||||
_TORCH_DTYPE = torch.float16
|
||||
|
||||
# Set to false to temporarily debug issues unrelated to weight update
|
||||
_ENABLE_UPDATE_WEIGHTS = True
|
||||
# _ENABLE_UPDATE_WEIGHTS = False
|
||||
|
||||
# TODO maybe we should add more other models? should we keep it in sync with test_generation_models.py?
|
||||
ALL_MODELS = [
|
||||
dict(model_path="meta-llama/Llama-3.2-1B-Instruct"),
|
||||
dict(model_path="Qwen/Qwen2-1.5B"),
|
||||
dict(model_path="allenai/OLMo-1B-0724-hf"),
|
||||
dict(model_path="allenai/OLMo-2-1124-7B-Instruct"),
|
||||
dict(
|
||||
model_path="ibm-granite/granite-3.0-2b-instruct",
|
||||
prefill_tolerance=0.22,
|
||||
decode_tolerance=0.22,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class TestVerlEngine(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
multiprocessing.set_start_method("spawn")
|
||||
|
||||
def assert_fragment_e2e_execution(
|
||||
self,
|
||||
index: int,
|
||||
model_path: str,
|
||||
mem_fraction_static: float = 0.4,
|
||||
dp_size: int = 1,
|
||||
tp_size: int = 2,
|
||||
tight_memory: bool = False,
|
||||
prefill_tolerance: float = 0.1,
|
||||
decode_tolerance: float = 0.1,
|
||||
):
|
||||
master_port = find_available_port(23456)
|
||||
|
||||
print(f"assert_fragment_e2e_execution START {index=} {model_path=}")
|
||||
|
||||
processes = []
|
||||
output_reader, output_writer = mp.Pipe(duplex=False)
|
||||
world_size = dp_size * tp_size
|
||||
for rank in range(world_size):
|
||||
p = Process(
|
||||
target=_run_subprocess,
|
||||
kwargs=dict(
|
||||
rank=rank,
|
||||
dp_size=dp_size,
|
||||
tp_size=tp_size,
|
||||
master_port=master_port,
|
||||
output_writer=output_writer,
|
||||
model_path=model_path,
|
||||
mem_fraction_static=mem_fraction_static,
|
||||
tight_memory=tight_memory,
|
||||
prefill_tolerance=prefill_tolerance,
|
||||
decode_tolerance=decode_tolerance,
|
||||
),
|
||||
)
|
||||
p.start()
|
||||
processes.append(p)
|
||||
|
||||
for _ in range(tp_size):
|
||||
self.assertTrue(
|
||||
output_reader.recv(),
|
||||
f"Subprocess has error, please see logs above. ({index=} {model_path=})",
|
||||
)
|
||||
|
||||
for p in processes:
|
||||
p.join()
|
||||
|
||||
def test_ci_models(self):
|
||||
ci_models = [random.choice(ALL_MODELS)]
|
||||
for index, model_info in enumerate(ci_models):
|
||||
self.assert_fragment_e2e_execution(index=index, **model_info)
|
||||
|
||||
def test_others(self):
|
||||
if is_in_ci():
|
||||
return
|
||||
|
||||
for index, model_info in enumerate(ALL_MODELS):
|
||||
self.assert_fragment_e2e_execution(index=index, **model_info)
|
||||
|
||||
# def test_adhoc(self):
|
||||
# self.assert_fragment_e2e_execution(index=0, model_path="meta-llama/Llama-3.2-1B-Instruct")
|
||||
|
||||
|
||||
def _run_subprocess(
|
||||
rank: int,
|
||||
dp_size: int,
|
||||
tp_size: int,
|
||||
master_port: int,
|
||||
output_writer,
|
||||
model_path: str,
|
||||
mem_fraction_static: float,
|
||||
tight_memory: bool,
|
||||
prefill_tolerance: float,
|
||||
decode_tolerance: float,
|
||||
):
|
||||
try:
|
||||
print(f"subprocess[{rank=}] Start {os.environ.get('CUDA_VISIBLE_DEVICES')=}")
|
||||
|
||||
os.environ["MASTER_ADDR"] = "localhost"
|
||||
os.environ["MASTER_PORT"] = str(master_port)
|
||||
torch.distributed.init_process_group(rank=rank, world_size=dp_size * tp_size)
|
||||
torch.cuda.set_device(rank)
|
||||
|
||||
base_gpu_id = rank // tp_size * tp_size
|
||||
|
||||
mesh_kwargs = dict(
|
||||
mesh_shape=(dp_size, tp_size, 1), mesh_dim_names=["dp", "tp", "pp"]
|
||||
)
|
||||
inference_device_mesh_device = init_device_mesh("cuda", **mesh_kwargs)
|
||||
inference_device_mesh_cpu = init_device_mesh("cpu", **mesh_kwargs)
|
||||
print(
|
||||
f"subprocess[{rank=},{base_gpu_id=}] {inference_device_mesh_device=} {inference_device_mesh_cpu=}"
|
||||
)
|
||||
|
||||
# hf model is used for comparison
|
||||
hf_model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path, torch_dtype=_TORCH_DTYPE, trust_remote_code=True
|
||||
).cuda()
|
||||
hf_tokenizer = get_tokenizer(model_path, trust_remote_code=True)
|
||||
|
||||
hf_outputs = HFRunner.forward_generation_raw(
|
||||
base_model=hf_model,
|
||||
prompts=_PROMPTS,
|
||||
max_new_tokens=_MAX_NEW_TOKENS,
|
||||
tokenizer=hf_tokenizer,
|
||||
lora_paths=None,
|
||||
torch_dtype=_TORCH_DTYPE,
|
||||
output_str_only=False,
|
||||
)
|
||||
print(
|
||||
f"subprocess[{rank=}] call hf.forward {hf_outputs=}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
if _ENABLE_UPDATE_WEIGHTS:
|
||||
if tight_memory:
|
||||
hf_model.cpu()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# test update weights
|
||||
print(f"subprocess[{rank=}] get_fsdp_state_dict", flush=True)
|
||||
fsdp_state_dict = _get_fsdp_state_dict(
|
||||
hf_model=hf_model, world_size=dp_size * tp_size
|
||||
)
|
||||
|
||||
engine = VerlEngine(
|
||||
model_path=model_path,
|
||||
load_format="dummy" if _ENABLE_UPDATE_WEIGHTS else "auto",
|
||||
mem_fraction_static=mem_fraction_static,
|
||||
random_seed=42,
|
||||
base_gpu_id=base_gpu_id,
|
||||
trust_remote_code=True,
|
||||
dtype=get_dtype_str(_TORCH_DTYPE),
|
||||
device_mesh_cpu=inference_device_mesh_cpu["tp"],
|
||||
)
|
||||
print(f"subprocess[{rank=}] {engine=}", flush=True)
|
||||
|
||||
if _ENABLE_UPDATE_WEIGHTS:
|
||||
print(f"subprocess[{rank=}] call update_weights_from_tensor", flush=True)
|
||||
engine.update_weights_from_tensor(
|
||||
[(k, v) for k, v in fsdp_state_dict.items()]
|
||||
)
|
||||
|
||||
for enable_batch in [False, True]:
|
||||
if enable_batch:
|
||||
fn = SRTRunner.batch_forward_generation_raw
|
||||
else:
|
||||
fn = SRTRunner.forward_generation_raw
|
||||
|
||||
srt_outputs = fn(
|
||||
prompts=_PROMPTS,
|
||||
max_new_tokens=_MAX_NEW_TOKENS,
|
||||
lora_paths=None,
|
||||
engine=engine,
|
||||
)
|
||||
print(
|
||||
f"subprocess[{rank=}] call srt.forward {enable_batch=} {srt_outputs=}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
check_close_model_outputs(
|
||||
hf_outputs=hf_outputs,
|
||||
srt_outputs=srt_outputs,
|
||||
prefill_tolerance=prefill_tolerance,
|
||||
decode_tolerance=decode_tolerance,
|
||||
rouge_l_tolerance=1,
|
||||
check_logprobs=not enable_batch,
|
||||
debug_text=f"{enable_batch=} {rank=}",
|
||||
)
|
||||
|
||||
execution_ok = True
|
||||
|
||||
except Exception as e:
|
||||
print(f"subprocess[{rank=}] has error: {e}", flush=True)
|
||||
traceback.print_exc()
|
||||
execution_ok = False
|
||||
|
||||
output_writer.send(execution_ok)
|
||||
output_writer.close()
|
||||
|
||||
if "engine" in locals() and engine is not None:
|
||||
engine.shutdown()
|
||||
print(f"subprocess[{rank=}] end", flush=True)
|
||||
|
||||
|
||||
# Adapted from https://github.com/volcengine/verl/blob/main/tests/rollout/run_fsdp_vllm.py
|
||||
def _get_fsdp_state_dict(hf_model, world_size: int):
|
||||
device_mesh = init_device_mesh(
|
||||
"cuda", mesh_shape=(world_size,), mesh_dim_names=["fsdp"]
|
||||
)
|
||||
|
||||
mixed_precision = MixedPrecision(
|
||||
param_dtype=torch.bfloat16,
|
||||
reduce_dtype=torch.float32,
|
||||
buffer_dtype=torch.float32,
|
||||
)
|
||||
fsdp_model = FSDP(
|
||||
hf_model,
|
||||
use_orig_params=True,
|
||||
auto_wrap_policy=None,
|
||||
device_id=torch.cuda.current_device(),
|
||||
sharding_strategy=ShardingStrategy.FULL_SHARD,
|
||||
mixed_precision=mixed_precision,
|
||||
cpu_offload=CPUOffload(offload_params=False),
|
||||
sync_module_states=False,
|
||||
device_mesh=device_mesh,
|
||||
)
|
||||
print(f"{fsdp_model=}")
|
||||
|
||||
FSDP.set_state_dict_type(
|
||||
fsdp_model,
|
||||
state_dict_type=StateDictType.SHARDED_STATE_DICT,
|
||||
state_dict_config=ShardedStateDictConfig(),
|
||||
)
|
||||
|
||||
return fsdp_model.state_dict()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,290 @@
|
||||
import multiprocessing
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import random
|
||||
import traceback
|
||||
import unittest
|
||||
from multiprocessing import Process
|
||||
|
||||
import torch
|
||||
from torch.distributed.device_mesh import init_device_mesh
|
||||
from torch.distributed.fsdp import CPUOffload
|
||||
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
|
||||
from torch.distributed.fsdp import MixedPrecision
|
||||
from torch.distributed.fsdp.api import (
|
||||
ShardedStateDictConfig,
|
||||
ShardingStrategy,
|
||||
StateDictType,
|
||||
)
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
from sglang.srt.entrypoints.verl_engine import VerlEngine
|
||||
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
||||
from sglang.test.runners import (
|
||||
HFRunner,
|
||||
SRTRunner,
|
||||
check_close_model_outputs,
|
||||
get_dtype_str,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase, find_available_port, is_in_ci
|
||||
|
||||
_MAX_NEW_TOKENS = 8
|
||||
_PROMPTS = ["1+1=2, 1+2=3, 1+3=4, 1+4=5, 1+5=", "1*1=1, 1*2=2, 1*3=3, 1*4=4, 1*5="]
|
||||
_TORCH_DTYPE = torch.float16
|
||||
|
||||
# Set to false to temporarily debug issues unrelated to weight update
|
||||
_ENABLE_UPDATE_WEIGHTS = True
|
||||
# _ENABLE_UPDATE_WEIGHTS = False
|
||||
|
||||
# TODO maybe we should add more other models? should we keep it in sync with test_generation_models.py?
|
||||
ALL_MODELS = [
|
||||
dict(
|
||||
model_path="Qwen/Qwen2.5-0.5B",
|
||||
dp_size=2,
|
||||
tp_size=2, # default to 2
|
||||
),
|
||||
dict(
|
||||
model_path="Qwen/Qwen2.5-14B-Instruct",
|
||||
mem_fraction_static=0.7,
|
||||
dp_size=2,
|
||||
tp_size=2,
|
||||
tight_memory=True,
|
||||
decode_tolerance=1.3,
|
||||
), # test_generation_models.py same config (qwen + tp=8) gives 1.22 decode error
|
||||
dict(
|
||||
model_path="THUDM/glm-4-9b-chat",
|
||||
mem_fraction_static=0.5,
|
||||
dp_size=2,
|
||||
tp_size=2,
|
||||
tight_memory=True,
|
||||
),
|
||||
# Fail to run these models in test_generation_models.py, need to fix that first
|
||||
# dict(model_path="openai-community/gpt2"),
|
||||
# dict(model_path="microsoft/Phi-3-small-8k-instruct"),
|
||||
]
|
||||
|
||||
|
||||
class TestVerlEngine(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
multiprocessing.set_start_method("spawn")
|
||||
|
||||
def assert_fragment_e2e_execution(
|
||||
self,
|
||||
index: int,
|
||||
model_path: str,
|
||||
mem_fraction_static: float = 0.4,
|
||||
dp_size: int = 1,
|
||||
tp_size: int = 2,
|
||||
tight_memory: bool = False,
|
||||
prefill_tolerance: float = 0.1,
|
||||
decode_tolerance: float = 0.1,
|
||||
):
|
||||
master_port = find_available_port(23456)
|
||||
|
||||
print(f"assert_fragment_e2e_execution START {index=} {model_path=}")
|
||||
|
||||
processes = []
|
||||
output_reader, output_writer = mp.Pipe(duplex=False)
|
||||
world_size = dp_size * tp_size
|
||||
for rank in range(world_size):
|
||||
p = Process(
|
||||
target=_run_subprocess,
|
||||
kwargs=dict(
|
||||
rank=rank,
|
||||
dp_size=dp_size,
|
||||
tp_size=tp_size,
|
||||
master_port=master_port,
|
||||
output_writer=output_writer,
|
||||
model_path=model_path,
|
||||
mem_fraction_static=mem_fraction_static,
|
||||
tight_memory=tight_memory,
|
||||
prefill_tolerance=prefill_tolerance,
|
||||
decode_tolerance=decode_tolerance,
|
||||
),
|
||||
)
|
||||
p.start()
|
||||
processes.append(p)
|
||||
|
||||
for _ in range(tp_size):
|
||||
self.assertTrue(
|
||||
output_reader.recv(),
|
||||
f"Subprocess has error, please see logs above. ({index=} {model_path=})",
|
||||
)
|
||||
|
||||
for p in processes:
|
||||
p.join()
|
||||
|
||||
def test_ci_models(self):
|
||||
ci_models = [random.choice(ALL_MODELS)]
|
||||
for index, model_info in enumerate(ci_models):
|
||||
self.assert_fragment_e2e_execution(index=index, **model_info)
|
||||
|
||||
def test_others(self):
|
||||
if is_in_ci():
|
||||
return
|
||||
|
||||
for index, model_info in enumerate(ALL_MODELS):
|
||||
self.assert_fragment_e2e_execution(index=index, **model_info)
|
||||
|
||||
# def test_adhoc(self):
|
||||
# self.assert_fragment_e2e_execution(index=0, model_path="meta-llama/Llama-3.2-1B-Instruct")
|
||||
|
||||
|
||||
def _run_subprocess(
|
||||
rank: int,
|
||||
dp_size: int,
|
||||
tp_size: int,
|
||||
master_port: int,
|
||||
output_writer,
|
||||
model_path: str,
|
||||
mem_fraction_static: float,
|
||||
tight_memory: bool,
|
||||
prefill_tolerance: float,
|
||||
decode_tolerance: float,
|
||||
):
|
||||
try:
|
||||
print(f"subprocess[{rank=}] Start {os.environ.get('CUDA_VISIBLE_DEVICES')=}")
|
||||
|
||||
os.environ["MASTER_ADDR"] = "localhost"
|
||||
os.environ["MASTER_PORT"] = str(master_port)
|
||||
torch.distributed.init_process_group(rank=rank, world_size=dp_size * tp_size)
|
||||
torch.cuda.set_device(rank)
|
||||
|
||||
base_gpu_id = rank // tp_size * tp_size
|
||||
|
||||
mesh_kwargs = dict(
|
||||
mesh_shape=(dp_size, tp_size, 1), mesh_dim_names=["dp", "tp", "pp"]
|
||||
)
|
||||
inference_device_mesh_device = init_device_mesh("cuda", **mesh_kwargs)
|
||||
inference_device_mesh_cpu = init_device_mesh("cpu", **mesh_kwargs)
|
||||
print(
|
||||
f"subprocess[{rank=},{base_gpu_id=}] {inference_device_mesh_device=} {inference_device_mesh_cpu=}"
|
||||
)
|
||||
|
||||
# hf model is used for comparison
|
||||
hf_model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path, torch_dtype=_TORCH_DTYPE, trust_remote_code=True
|
||||
).cuda()
|
||||
hf_tokenizer = get_tokenizer(model_path, trust_remote_code=True)
|
||||
|
||||
hf_outputs = HFRunner.forward_generation_raw(
|
||||
base_model=hf_model,
|
||||
prompts=_PROMPTS,
|
||||
max_new_tokens=_MAX_NEW_TOKENS,
|
||||
tokenizer=hf_tokenizer,
|
||||
lora_paths=None,
|
||||
torch_dtype=_TORCH_DTYPE,
|
||||
output_str_only=False,
|
||||
)
|
||||
print(
|
||||
f"subprocess[{rank=}] call hf.forward {hf_outputs=}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
if _ENABLE_UPDATE_WEIGHTS:
|
||||
if tight_memory:
|
||||
hf_model.cpu()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# test update weights
|
||||
print(f"subprocess[{rank=}] get_fsdp_state_dict", flush=True)
|
||||
fsdp_state_dict = _get_fsdp_state_dict(
|
||||
hf_model=hf_model, world_size=dp_size * tp_size
|
||||
)
|
||||
|
||||
engine = VerlEngine(
|
||||
model_path=model_path,
|
||||
load_format="dummy" if _ENABLE_UPDATE_WEIGHTS else "auto",
|
||||
mem_fraction_static=mem_fraction_static,
|
||||
random_seed=42,
|
||||
base_gpu_id=base_gpu_id,
|
||||
trust_remote_code=True,
|
||||
dtype=get_dtype_str(_TORCH_DTYPE),
|
||||
device_mesh_cpu=inference_device_mesh_cpu["tp"],
|
||||
)
|
||||
print(f"subprocess[{rank=}] {engine=}", flush=True)
|
||||
|
||||
if _ENABLE_UPDATE_WEIGHTS:
|
||||
print(f"subprocess[{rank=}] call update_weights_from_tensor", flush=True)
|
||||
engine.update_weights_from_tensor(
|
||||
[(k, v) for k, v in fsdp_state_dict.items()]
|
||||
)
|
||||
|
||||
for enable_batch in [False, True]:
|
||||
if enable_batch:
|
||||
fn = SRTRunner.batch_forward_generation_raw
|
||||
else:
|
||||
fn = SRTRunner.forward_generation_raw
|
||||
|
||||
srt_outputs = fn(
|
||||
prompts=_PROMPTS,
|
||||
max_new_tokens=_MAX_NEW_TOKENS,
|
||||
lora_paths=None,
|
||||
engine=engine,
|
||||
)
|
||||
print(
|
||||
f"subprocess[{rank=}] call srt.forward {enable_batch=} {srt_outputs=}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
check_close_model_outputs(
|
||||
hf_outputs=hf_outputs,
|
||||
srt_outputs=srt_outputs,
|
||||
prefill_tolerance=prefill_tolerance,
|
||||
decode_tolerance=decode_tolerance,
|
||||
rouge_l_tolerance=1,
|
||||
check_logprobs=not enable_batch,
|
||||
debug_text=f"{enable_batch=} {rank=}",
|
||||
)
|
||||
|
||||
execution_ok = True
|
||||
|
||||
except Exception as e:
|
||||
print(f"subprocess[{rank=}] has error: {e}", flush=True)
|
||||
traceback.print_exc()
|
||||
execution_ok = False
|
||||
|
||||
output_writer.send(execution_ok)
|
||||
output_writer.close()
|
||||
|
||||
if "engine" in locals() and engine is not None:
|
||||
engine.shutdown()
|
||||
print(f"subprocess[{rank=}] end", flush=True)
|
||||
|
||||
|
||||
# Adapted from https://github.com/volcengine/verl/blob/main/tests/rollout/run_fsdp_vllm.py
|
||||
def _get_fsdp_state_dict(hf_model, world_size: int):
|
||||
device_mesh = init_device_mesh(
|
||||
"cuda", mesh_shape=(world_size,), mesh_dim_names=["fsdp"]
|
||||
)
|
||||
|
||||
mixed_precision = MixedPrecision(
|
||||
param_dtype=torch.bfloat16,
|
||||
reduce_dtype=torch.float32,
|
||||
buffer_dtype=torch.float32,
|
||||
)
|
||||
fsdp_model = FSDP(
|
||||
hf_model,
|
||||
use_orig_params=True,
|
||||
auto_wrap_policy=None,
|
||||
device_id=torch.cuda.current_device(),
|
||||
sharding_strategy=ShardingStrategy.FULL_SHARD,
|
||||
mixed_precision=mixed_precision,
|
||||
cpu_offload=CPUOffload(offload_params=False),
|
||||
sync_module_states=False,
|
||||
device_mesh=device_mesh,
|
||||
)
|
||||
print(f"{fsdp_model=}")
|
||||
|
||||
FSDP.set_state_dict_type(
|
||||
fsdp_model,
|
||||
state_dict_type=StateDictType.SHARDED_STATE_DICT,
|
||||
state_dict_config=ShardedStateDictConfig(),
|
||||
)
|
||||
|
||||
return fsdp_model.state_dict()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,295 @@
|
||||
"""
|
||||
Unit tests for AsyncDynamicbatchTokenizer.
|
||||
|
||||
Tests the async dynamic batching functionality for tokenization,
|
||||
including batch efficiency, timeout handling, and error cases.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from sglang.srt.managers.async_dynamic_batch_tokenizer import AsyncDynamicbatchTokenizer
|
||||
|
||||
|
||||
class TestAsyncDynamicbatchTokenizer:
|
||||
"""Test suite for AsyncDynamicbatchTokenizer."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tokenizer(self):
|
||||
"""Create a mock tokenizer that behaves like HuggingFace tokenizer."""
|
||||
|
||||
def mock_encode(texts, **kwargs):
|
||||
is_single = isinstance(texts, str)
|
||||
if is_single:
|
||||
texts = [texts]
|
||||
|
||||
# Simulate tokenization - convert text to mock token ids
|
||||
input_ids = []
|
||||
token_type_ids = []
|
||||
|
||||
for text in texts:
|
||||
# Simple mock: text length determines number of tokens
|
||||
tokens = [i for i in range(len(text.split()))]
|
||||
input_ids.append(tokens)
|
||||
|
||||
if kwargs.get("return_token_type_ids", False):
|
||||
token_type_ids.append([0] * len(tokens))
|
||||
|
||||
result = {"input_ids": input_ids}
|
||||
if kwargs.get("return_token_type_ids", False):
|
||||
result["token_type_ids"] = token_type_ids
|
||||
|
||||
# For single inputs, return individual result (not wrapped in a list)
|
||||
if is_single:
|
||||
result = {"input_ids": input_ids[0]}
|
||||
if kwargs.get("return_token_type_ids", False):
|
||||
result["token_type_ids"] = token_type_ids[0]
|
||||
|
||||
# Create a proper BatchEncoding-like object that supports dict operations
|
||||
class MockBatchEncoding(dict):
|
||||
def __init__(self, data):
|
||||
super().__init__(data)
|
||||
for key, value in data.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
return MockBatchEncoding(result)
|
||||
|
||||
# Return the function directly - the AsyncDynamicbatchTokenizer will call it
|
||||
return mock_encode
|
||||
|
||||
@pytest.fixture
|
||||
def async_tokenizer(self, mock_tokenizer):
|
||||
"""Create AsyncDynamicbatchTokenizer instance."""
|
||||
return AsyncDynamicbatchTokenizer(
|
||||
tokenizer=mock_tokenizer, max_batch_size=4, batch_wait_timeout_s=0.01
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_request(self, async_tokenizer):
|
||||
"""Test tokenizing a single request."""
|
||||
text = "hello world"
|
||||
result = await async_tokenizer.encode(text)
|
||||
|
||||
assert "input_ids" in result
|
||||
assert result["input_ids"] == [0, 1] # 2 words -> 2 tokens
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_request_with_token_type_ids(self, async_tokenizer):
|
||||
"""Test tokenizing with token type IDs."""
|
||||
text = "hello world"
|
||||
result = await async_tokenizer.encode(text, return_token_type_ids=True)
|
||||
|
||||
assert "input_ids" in result
|
||||
assert "token_type_ids" in result
|
||||
assert result["input_ids"] == [0, 1]
|
||||
assert result["token_type_ids"] == [0, 0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_requests_same_kwargs(self, async_tokenizer):
|
||||
"""Test that concurrent requests with same kwargs get batched."""
|
||||
texts = ["hello world", "how are you", "fine thanks", "good morning"]
|
||||
|
||||
# Start all requests concurrently
|
||||
tasks = [async_tokenizer.encode(text) for text in texts]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
# Verify all results
|
||||
assert len(results) == 4
|
||||
for i, result in enumerate(results):
|
||||
assert "input_ids" in result
|
||||
expected_tokens = list(range(len(texts[i].split())))
|
||||
assert result["input_ids"] == expected_tokens
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_requests_different_kwargs(self, async_tokenizer):
|
||||
"""Test that requests with different kwargs are processed individually."""
|
||||
text1 = "hello world"
|
||||
text2 = "how are you"
|
||||
|
||||
# One with token_type_ids, one without
|
||||
task1 = async_tokenizer.encode(text1, return_token_type_ids=True)
|
||||
task2 = async_tokenizer.encode(text2)
|
||||
|
||||
result1, result2 = await asyncio.gather(task1, task2)
|
||||
|
||||
# First result should have token_type_ids
|
||||
assert "input_ids" in result1
|
||||
assert "token_type_ids" in result1
|
||||
assert result1["input_ids"] == [0, 1]
|
||||
assert result1["token_type_ids"] == [0, 0]
|
||||
|
||||
# Second result should not have token_type_ids
|
||||
assert "input_ids" in result2
|
||||
assert "token_type_ids" not in result2
|
||||
assert result2["input_ids"] == [0, 1, 2]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_timeout(self, async_tokenizer):
|
||||
"""Test that batching respects timeout."""
|
||||
# Send first request
|
||||
task1 = asyncio.create_task(async_tokenizer.encode("hello world"))
|
||||
|
||||
# Wait longer than batch timeout
|
||||
await asyncio.sleep(0.02) # Longer than 0.01s timeout
|
||||
|
||||
# Send second request
|
||||
task2 = asyncio.create_task(async_tokenizer.encode("how are you"))
|
||||
|
||||
results = await asyncio.gather(task1, task2)
|
||||
|
||||
# Both should complete successfully
|
||||
assert len(results) == 2
|
||||
assert results[0]["input_ids"] == [0, 1]
|
||||
assert results[1]["input_ids"] == [0, 1, 2]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_batch_size_limit(self, async_tokenizer):
|
||||
"""Test that batching respects max_batch_size."""
|
||||
# Send more requests than max_batch_size (4)
|
||||
texts = [f"text {i}" for i in range(6)]
|
||||
tasks = [async_tokenizer.encode(text) for text in texts]
|
||||
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
# All should complete successfully
|
||||
assert len(results) == 6
|
||||
for i, result in enumerate(results):
|
||||
assert "input_ids" in result
|
||||
assert result["input_ids"] == [0, 1] # "text i" -> 2 tokens
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_callable_interface(self, async_tokenizer):
|
||||
"""Test that the tokenizer is callable."""
|
||||
text = "hello world"
|
||||
result = await async_tokenizer(text)
|
||||
|
||||
assert "input_ids" in result
|
||||
assert result["input_ids"] == [0, 1]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lazy_initialization(self, mock_tokenizer):
|
||||
"""Test that initialization happens lazily."""
|
||||
tokenizer = AsyncDynamicbatchTokenizer(mock_tokenizer)
|
||||
|
||||
# Should not be initialized yet
|
||||
assert not tokenizer._initialized
|
||||
|
||||
# First encode should initialize
|
||||
await tokenizer.encode("hello")
|
||||
|
||||
# Should now be initialized
|
||||
assert tokenizer._initialized
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_handling_in_tokenizer(self, mock_tokenizer):
|
||||
"""Test error handling when tokenizer fails."""
|
||||
|
||||
# Create a new async tokenizer with a failing tokenizer
|
||||
def failing_tokenizer(*args, **kwargs):
|
||||
raise ValueError("Tokenizer error")
|
||||
|
||||
async_tokenizer = AsyncDynamicbatchTokenizer(
|
||||
tokenizer=failing_tokenizer, max_batch_size=4, batch_wait_timeout_s=0.01
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Tokenizer error"):
|
||||
await async_tokenizer.encode("hello world")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_processing_logs(self, async_tokenizer, caplog):
|
||||
"""Test that batch processing logs are generated."""
|
||||
caplog.set_level(logging.DEBUG)
|
||||
|
||||
# Send multiple requests to trigger batching
|
||||
tasks = [
|
||||
async_tokenizer.encode("hello world"),
|
||||
async_tokenizer.encode("how are you"),
|
||||
]
|
||||
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# Should have batch processing log
|
||||
assert any(
|
||||
"Processing dynamic batch of size" in record.message
|
||||
for record in caplog.records
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_queue_immediate_processing(self, async_tokenizer):
|
||||
"""Test that single requests are processed immediately when queue is empty."""
|
||||
start_time = time.time()
|
||||
result = await async_tokenizer.encode("hello world")
|
||||
end_time = time.time()
|
||||
|
||||
# Should complete quickly (much less than batch timeout)
|
||||
assert end_time - start_time < 0.005 # 5ms should be plenty
|
||||
assert result["input_ids"] == [0, 1]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_tokenizer_integration(self):
|
||||
"""Test with a real HuggingFace tokenizer."""
|
||||
try:
|
||||
# Use a small, fast tokenizer for testing
|
||||
real_tokenizer = AutoTokenizer.from_pretrained("gpt2")
|
||||
async_tokenizer = AsyncDynamicbatchTokenizer(
|
||||
tokenizer=real_tokenizer, max_batch_size=2, batch_wait_timeout_s=0.01
|
||||
)
|
||||
|
||||
text = "Hello, world!"
|
||||
result = await async_tokenizer.encode(text)
|
||||
|
||||
# Should get actual token IDs
|
||||
assert "input_ids" in result
|
||||
assert isinstance(result["input_ids"], list)
|
||||
assert len(result["input_ids"]) > 0
|
||||
assert all(isinstance(token_id, int) for token_id in result["input_ids"])
|
||||
|
||||
except Exception as e:
|
||||
pytest.skip(f"Real tokenizer test skipped: {e}")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_mixed_requests(self, async_tokenizer):
|
||||
"""Test mixing single and batched requests."""
|
||||
# Start some requests
|
||||
task1 = asyncio.create_task(async_tokenizer.encode("hello"))
|
||||
task2 = asyncio.create_task(async_tokenizer.encode("world"))
|
||||
|
||||
# Wait a bit
|
||||
await asyncio.sleep(0.005)
|
||||
|
||||
# Start more requests
|
||||
task3 = asyncio.create_task(async_tokenizer.encode("how are"))
|
||||
task4 = asyncio.create_task(async_tokenizer.encode("you doing"))
|
||||
|
||||
results = await asyncio.gather(task1, task2, task3, task4)
|
||||
|
||||
# All should complete successfully
|
||||
assert len(results) == 4
|
||||
for result in results:
|
||||
assert "input_ids" in result
|
||||
assert isinstance(result["input_ids"], list)
|
||||
|
||||
def test_cleanup_on_destruction(self, mock_tokenizer):
|
||||
"""Test that resources are cleaned up properly."""
|
||||
tokenizer = AsyncDynamicbatchTokenizer(mock_tokenizer)
|
||||
|
||||
# Mock the executor and task
|
||||
tokenizer._executor = Mock()
|
||||
tokenizer._batcher_task = Mock()
|
||||
tokenizer._batcher_task.done.return_value = False
|
||||
|
||||
# Call destructor
|
||||
tokenizer.__del__()
|
||||
|
||||
# Should cancel task and shutdown executor
|
||||
tokenizer._batcher_task.cancel.assert_called_once()
|
||||
tokenizer._executor.shutdown.assert_called_once_with(wait=False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
@@ -0,0 +1,364 @@
|
||||
"""
|
||||
Unit tests for AsyncMMDataProcessor.
|
||||
|
||||
Covers:
|
||||
- Async and sync processing paths
|
||||
- Concurrency limiting via semaphore
|
||||
- Per-call timeout behavior (async and sync)
|
||||
- Argument passthrough (images, audios, text/ids, request_obj, kwargs)
|
||||
- Error propagation and shutdown behavior
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.managers.async_mm_data_processor import AsyncMMDataProcessor
|
||||
|
||||
|
||||
class TestAsyncMMDataProcessor:
|
||||
"""Test suite for AsyncMMDataProcessor."""
|
||||
|
||||
@pytest.fixture
|
||||
def async_processor(self):
|
||||
"""Create a processor exposing an async process_mm_data_async."""
|
||||
|
||||
class AsyncProc:
|
||||
async def process_mm_data_async(
|
||||
self,
|
||||
*,
|
||||
image_data=None,
|
||||
audio_data=None,
|
||||
input_text=None,
|
||||
request_obj=None,
|
||||
**kwargs,
|
||||
):
|
||||
# Allow tests to simulate latency via kwargs
|
||||
delay = kwargs.get("delay_s", 0.0)
|
||||
if delay:
|
||||
await asyncio.sleep(delay)
|
||||
return {
|
||||
"path": "async",
|
||||
"images": image_data,
|
||||
"audios": audio_data,
|
||||
"text": input_text,
|
||||
"request": request_obj,
|
||||
"kwargs": kwargs,
|
||||
}
|
||||
|
||||
return AsyncProc()
|
||||
|
||||
@pytest.fixture
|
||||
def sync_processor(self):
|
||||
"""Provide a processor exposing a sync process_mm_data."""
|
||||
|
||||
class SyncProc:
|
||||
def process_mm_data(
|
||||
self,
|
||||
*,
|
||||
image_data=None,
|
||||
audio_data=None,
|
||||
input_text=None,
|
||||
request_obj=None,
|
||||
**kwargs,
|
||||
):
|
||||
delay = kwargs.get("delay_s", 0.0)
|
||||
if delay:
|
||||
# Simulate CPU/blocking work
|
||||
time.sleep(delay)
|
||||
return {
|
||||
"path": "sync",
|
||||
"images": image_data,
|
||||
"audios": audio_data,
|
||||
"text": input_text,
|
||||
"request": request_obj,
|
||||
"kwargs": kwargs,
|
||||
}
|
||||
|
||||
return SyncProc()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_path_basic(self, async_processor):
|
||||
"""Async processor should be awaited directly."""
|
||||
proc = AsyncMMDataProcessor(async_processor)
|
||||
out = await proc.process(
|
||||
image_data=["img1.png"],
|
||||
audio_data=["a.wav"],
|
||||
input_text_or_ids="hello",
|
||||
request_obj={"rid": 1},
|
||||
mode="fast",
|
||||
)
|
||||
assert out["path"] == "async"
|
||||
assert out["images"] == ["img1.png"]
|
||||
assert out["audios"] == ["a.wav"]
|
||||
assert out["text"] == "hello"
|
||||
assert out["request"] == {"rid": 1}
|
||||
assert out["kwargs"]["mode"] == "fast"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_fallback_basic(self, sync_processor):
|
||||
"""Sync processor should run in fallback executor."""
|
||||
proc = AsyncMMDataProcessor(sync_processor)
|
||||
out = await proc.process(
|
||||
image_data=[b"\x00\x01"],
|
||||
audio_data=None,
|
||||
input_text_or_ids=[1, 2, 3],
|
||||
request_obj="req-obj",
|
||||
role="user",
|
||||
)
|
||||
assert out["path"] == "sync"
|
||||
assert out["images"] == [b"\x00\x01"]
|
||||
assert out["audios"] is None
|
||||
assert out["text"] == [1, 2, 3]
|
||||
assert out["request"] == "req-obj"
|
||||
assert out["kwargs"]["role"] == "user"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_async(self, async_processor):
|
||||
"""Timeout should raise asyncio.TimeoutError for async path."""
|
||||
proc = AsyncMMDataProcessor(async_processor, timeout_s=0.01)
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
await proc.process(
|
||||
input_text_or_ids="slow",
|
||||
request_obj=None,
|
||||
delay_s=0.05, # longer than timeout
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_sync(self, sync_processor):
|
||||
"""Timeout should raise asyncio.TimeoutError for sync fallback path."""
|
||||
proc = AsyncMMDataProcessor(sync_processor, timeout_s=0.01)
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
await proc.process(
|
||||
input_text_or_ids="slow",
|
||||
request_obj=None,
|
||||
delay_s=0.05, # longer than timeout
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semaphore_release_after_timeout(self, sync_processor):
|
||||
"""
|
||||
If a call times out, the semaphore should be released so a subsequent call can proceed.
|
||||
Use >=2 fallback workers so the timed-out thread doesn't block the next call.
|
||||
"""
|
||||
proc = AsyncMMDataProcessor(
|
||||
sync_processor,
|
||||
max_concurrent_calls=2,
|
||||
timeout_s=0.01,
|
||||
)
|
||||
|
||||
# First call will time out
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
await proc.process(
|
||||
input_text_or_ids="slow1", request_obj=None, delay_s=0.05
|
||||
)
|
||||
|
||||
# Second call should be able to acquire the semaphore and complete
|
||||
out = await proc.process(input_text_or_ids="ok", request_obj=None, delay_s=0.0)
|
||||
assert out["text"] == "ok"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrency_limit_async(self):
|
||||
"""Ensure max_concurrent_calls caps concurrency for async path."""
|
||||
current = 0
|
||||
max_seen = 0
|
||||
|
||||
class AsyncProc:
|
||||
async def process_mm_data_async(self, **kwargs):
|
||||
nonlocal current, max_seen
|
||||
current += 1
|
||||
max_seen = max(max_seen, current)
|
||||
try:
|
||||
await asyncio.sleep(0.02)
|
||||
return {"ok": True}
|
||||
finally:
|
||||
current -= 1
|
||||
|
||||
proc = AsyncMMDataProcessor(AsyncProc(), max_concurrent_calls=2)
|
||||
|
||||
tasks = [
|
||||
proc.process(input_text_or_ids=f"t{i}", request_obj=None) for i in range(6)
|
||||
]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
assert max_seen <= 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrency_limit_sync(self):
|
||||
"""Ensure max_concurrent_calls caps concurrency for sync fallback path."""
|
||||
current = 0
|
||||
max_seen = 0
|
||||
lock = threading.Lock()
|
||||
|
||||
class SyncProc:
|
||||
def process_mm_data(self, **kwargs):
|
||||
nonlocal current, max_seen
|
||||
with lock:
|
||||
current += 1
|
||||
max_seen = max(max_seen, current)
|
||||
try:
|
||||
time.sleep(0.02)
|
||||
return {"ok": True}
|
||||
finally:
|
||||
with lock:
|
||||
current -= 1
|
||||
|
||||
proc = AsyncMMDataProcessor(SyncProc(), max_concurrent_calls=3)
|
||||
|
||||
tasks = [
|
||||
proc.process(input_text_or_ids=f"s{i}", request_obj=None) for i in range(9)
|
||||
]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
assert max_seen <= 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_from_async_processor(self):
|
||||
"""Exceptions raised by the async processor should propagate."""
|
||||
|
||||
class BadAsync:
|
||||
async def process_mm_data_async(self, **_):
|
||||
await asyncio.sleep(0)
|
||||
raise ValueError("async boom")
|
||||
|
||||
proc = AsyncMMDataProcessor(BadAsync())
|
||||
with pytest.raises(ValueError, match="async boom"):
|
||||
await proc.process(input_text_or_ids="x", request_obj=None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_from_sync_processor(self):
|
||||
"""Exceptions raised by the sync processor should propagate."""
|
||||
|
||||
class BadSync:
|
||||
def process_mm_data(self, **_):
|
||||
raise RuntimeError("sync boom")
|
||||
|
||||
proc = AsyncMMDataProcessor(BadSync())
|
||||
with pytest.raises(RuntimeError, match="sync boom"):
|
||||
await proc.process(input_text_or_ids="x", request_obj=None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_both_methods_raises(self):
|
||||
"""Processor missing both methods should raise at call time."""
|
||||
|
||||
class Empty:
|
||||
pass
|
||||
|
||||
proc = AsyncMMDataProcessor(Empty())
|
||||
with pytest.raises(
|
||||
RuntimeError, match="neither 'process_mm_data_async' nor 'process_mm_data'"
|
||||
):
|
||||
await proc.process(input_text_or_ids="x", request_obj=None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_attribute_not_coroutine_uses_sync_fallback(self):
|
||||
"""
|
||||
If `process_mm_data_async` exists but isn't a coroutine function,
|
||||
wrapper should treat it as sync and use `process_mm_data`.
|
||||
"""
|
||||
|
||||
class WeirdProc:
|
||||
# Not a coroutine function:
|
||||
def process_mm_data_async(self, **_):
|
||||
return {"path": "would-be-async"}
|
||||
|
||||
def process_mm_data(self, **_):
|
||||
return {"path": "sync"}
|
||||
|
||||
proc = AsyncMMDataProcessor(WeirdProc())
|
||||
out = await proc.process(input_text_or_ids="x", request_obj=None)
|
||||
assert out["path"] == "sync"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_kwargs_and_request_passthrough_async(self, async_processor):
|
||||
"""Extra kwargs and request_obj should be forwarded on async path."""
|
||||
proc = AsyncMMDataProcessor(async_processor)
|
||||
out = await proc.process(
|
||||
image_data=["i1", "i2"],
|
||||
audio_data=["a1"],
|
||||
input_text_or_ids="hello world",
|
||||
request_obj={"uid": 42},
|
||||
return_meta=True,
|
||||
delay_s=0.0,
|
||||
)
|
||||
assert out["images"] == ["i1", "i2"]
|
||||
assert out["audios"] == ["a1"]
|
||||
assert out["text"] == "hello world"
|
||||
assert out["request"] == {"uid": 42}
|
||||
assert out["kwargs"]["return_meta"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_kwargs_and_request_passthrough_sync(self, sync_processor):
|
||||
"""Extra kwargs and request_obj should be forwarded on sync path."""
|
||||
proc = AsyncMMDataProcessor(sync_processor)
|
||||
out = await proc.process(
|
||||
image_data=None,
|
||||
audio_data=[],
|
||||
input_text_or_ids=[101, 102],
|
||||
request_obj=("r", 7),
|
||||
lang="en",
|
||||
)
|
||||
assert out["images"] is None
|
||||
assert out["audios"] == []
|
||||
assert out["text"] == [101, 102]
|
||||
assert out["request"] == ("r", 7)
|
||||
assert out["kwargs"]["lang"] == "en"
|
||||
|
||||
def test_shutdown_on_sync_executor(self, sync_processor):
|
||||
"""Explicit shutdown should close fallback executor for sync path."""
|
||||
proc = AsyncMMDataProcessor(sync_processor)
|
||||
# Swap real executor for a mock to assert shutdown behavior
|
||||
proc.fallback_exec = Mock()
|
||||
proc.shutdown()
|
||||
proc.fallback_exec.shutdown.assert_called_once_with(wait=False)
|
||||
|
||||
def test_del_calls_shutdown(self, sync_processor, caplog):
|
||||
"""__del__ should best-effort shutdown without raising."""
|
||||
caplog.set_level(logging.DEBUG)
|
||||
proc = AsyncMMDataProcessor(sync_processor)
|
||||
proc.fallback_exec = Mock()
|
||||
# Simulate object destruction
|
||||
proc.__del__()
|
||||
proc.fallback_exec.shutdown.assert_called_once_with(wait=False)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_mixed_requests(self, async_processor):
|
||||
"""Mix different payloads and ensure all complete with valid outputs."""
|
||||
proc = AsyncMMDataProcessor(async_processor, max_concurrent_calls=4)
|
||||
|
||||
tasks = [
|
||||
proc.process(input_text_or_ids="t1", request_obj=1),
|
||||
proc.process(image_data=["i.png"], input_text_or_ids=[9, 8], request_obj=2),
|
||||
proc.process(
|
||||
audio_data=["v.wav"], input_text_or_ids="speech", request_obj=3
|
||||
),
|
||||
proc.process(
|
||||
image_data=[], audio_data=[], input_text_or_ids=None, request_obj=4
|
||||
),
|
||||
]
|
||||
outs = await asyncio.gather(*tasks)
|
||||
assert len(outs) == 4
|
||||
for out in outs:
|
||||
assert "path" in out
|
||||
assert out["path"] == "async"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_many_requests_values_match_inputs(self, sync_processor):
|
||||
"""For sync path, ensure each response corresponds to its specific input."""
|
||||
proc = AsyncMMDataProcessor(sync_processor, max_concurrent_calls=8)
|
||||
texts = [f"msg-{i}" for i in range(10)]
|
||||
tasks = [
|
||||
proc.process(input_text_or_ids=t, request_obj=i)
|
||||
for i, t in enumerate(texts)
|
||||
]
|
||||
outs = await asyncio.gather(*tasks)
|
||||
got = [o["text"] for o in outs]
|
||||
assert got == texts
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
@@ -0,0 +1,158 @@
|
||||
"""
|
||||
Test script to verify SGLang config file integration.
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
from sglang.srt.server_args import prepare_server_args
|
||||
from sglang.srt.server_args_config_parser import ConfigArgumentMerger
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def merger():
|
||||
"""Fixture providing a ConfigArgumentMerger instance."""
|
||||
return ConfigArgumentMerger()
|
||||
|
||||
|
||||
def test_server_args_config_parser(merger):
|
||||
"""Test the config parser functionality."""
|
||||
# Create a temporary config file
|
||||
config_data = {
|
||||
"model-path": "microsoft/DialoGPT-medium",
|
||||
"host": "0.0.0.0",
|
||||
"port": 30000,
|
||||
"tensor-parallel-size": 2,
|
||||
"trust-remote-code": False,
|
||||
"enable-metrics": True,
|
||||
"stream-output": True,
|
||||
"skip-server-warmup": False,
|
||||
"log-requests": True,
|
||||
"show-time-cost": True,
|
||||
"is-embedding": False,
|
||||
}
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
|
||||
yaml.dump(config_data, f)
|
||||
config_file = f.name
|
||||
|
||||
try:
|
||||
# Test config parser directly
|
||||
config_args = merger._parse_yaml_config(config_file)
|
||||
|
||||
# Test merging with CLI args
|
||||
cli_args = ["--config", config_file, "--max-running-requests", "128"]
|
||||
merged_args = merger.merge_config_with_args(cli_args)
|
||||
|
||||
# Verify the merged args contain both config and CLI values
|
||||
assert "--model-path" in merged_args
|
||||
assert "microsoft/DialoGPT-medium" in merged_args
|
||||
assert "--host" in merged_args
|
||||
assert "0.0.0.0" in merged_args
|
||||
assert "--port" in merged_args
|
||||
assert "30000" in merged_args
|
||||
assert "--tensor-parallel-size" in merged_args
|
||||
assert "2" in merged_args
|
||||
assert "--max-running-requests" in merged_args
|
||||
assert "128" in merged_args
|
||||
|
||||
# Test boolean arguments
|
||||
assert "--enable-metrics" in merged_args # True boolean
|
||||
assert "--stream-output" in merged_args # True boolean
|
||||
assert "--log-requests" in merged_args # True boolean
|
||||
assert "--show-time-cost" in merged_args # True boolean
|
||||
# False booleans should not be present (only add flag if True)
|
||||
assert "--trust-remote-code" not in merged_args # False boolean
|
||||
assert "--skip-server-warmup" not in merged_args # False boolean
|
||||
assert "--is-embedding" not in merged_args # False boolean
|
||||
|
||||
finally:
|
||||
os.unlink(config_file)
|
||||
|
||||
|
||||
def test_server_args_integration():
|
||||
"""Test the integration with server args."""
|
||||
# Create a temporary config file
|
||||
config_data = {
|
||||
"model-path": "microsoft/DialoGPT-medium",
|
||||
"host": "0.0.0.0",
|
||||
"port": 30000,
|
||||
"tensor-parallel-size": 1,
|
||||
"max-running-requests": 256,
|
||||
}
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
|
||||
yaml.dump(config_data, f)
|
||||
config_file = f.name
|
||||
|
||||
try:
|
||||
# Test with config file
|
||||
argv = ["--config", config_file]
|
||||
server_args = prepare_server_args(argv)
|
||||
|
||||
# Verify that config values were loaded
|
||||
assert server_args.model_path == "microsoft/DialoGPT-medium"
|
||||
assert server_args.host == "0.0.0.0"
|
||||
assert server_args.port == 30000
|
||||
assert server_args.tp_size == 1
|
||||
assert server_args.max_running_requests == 256
|
||||
|
||||
finally:
|
||||
os.unlink(config_file)
|
||||
|
||||
|
||||
def test_cli_override():
|
||||
"""Test that CLI arguments override config file values."""
|
||||
# Create a temporary config file
|
||||
config_data = {
|
||||
"model-path": "microsoft/DialoGPT-medium",
|
||||
"port": 30000,
|
||||
"tensor-parallel-size": 1,
|
||||
}
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
|
||||
yaml.dump(config_data, f)
|
||||
config_file = f.name
|
||||
|
||||
try:
|
||||
# Test CLI override (CLI should take precedence)
|
||||
argv = [
|
||||
"--config",
|
||||
config_file,
|
||||
"--port",
|
||||
"40000",
|
||||
"--tensor-parallel-size",
|
||||
"2",
|
||||
]
|
||||
server_args = prepare_server_args(argv)
|
||||
|
||||
# Verify that CLI values override config values
|
||||
assert server_args.model_path == "microsoft/DialoGPT-medium" # From config
|
||||
assert server_args.port == 40000 # From CLI (overrides config)
|
||||
assert server_args.tp_size == 2 # From CLI (overrides config)
|
||||
|
||||
finally:
|
||||
os.unlink(config_file)
|
||||
|
||||
|
||||
def test_error_handling():
|
||||
"""Test error handling for invalid config files."""
|
||||
# Test non-existent config file
|
||||
with pytest.raises(ValueError, match="Config file not found"):
|
||||
argv = ["--config", "non-existent.yaml"]
|
||||
prepare_server_args(argv)
|
||||
|
||||
# Test invalid YAML file
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
|
||||
f.write("invalid: yaml: content: [")
|
||||
invalid_yaml_file = f.name
|
||||
|
||||
try:
|
||||
with pytest.raises(Exception):
|
||||
argv = ["--config", invalid_yaml_file]
|
||||
prepare_server_args(argv)
|
||||
finally:
|
||||
os.unlink(invalid_yaml_file)
|
||||
@@ -0,0 +1,182 @@
|
||||
import os
|
||||
import random
|
||||
import socket
|
||||
import unittest
|
||||
from typing import Any
|
||||
|
||||
import ray
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.srt.distributed import init_distributed_environment
|
||||
from sglang.srt.distributed.communication_op import ( # noqa
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from sglang.srt.distributed.parallel_state import (
|
||||
get_tensor_model_parallel_group,
|
||||
graph_capture,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
def get_open_port() -> int:
|
||||
# try ipv4
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
except OSError:
|
||||
# try ipv6
|
||||
with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def multi_process_parallel(
|
||||
world_size: int,
|
||||
cls: Any,
|
||||
test_target: Any,
|
||||
) -> None:
|
||||
|
||||
# Using ray helps debugging the error when it failed
|
||||
# as compared to multiprocessing.
|
||||
# NOTE: We need to set working_dir for distributed tests,
|
||||
# otherwise we may get import errors on ray workers
|
||||
|
||||
ray.init(log_to_driver=True)
|
||||
|
||||
distributed_init_port = get_open_port()
|
||||
refs = []
|
||||
for rank in range(world_size):
|
||||
refs.append(test_target.remote(cls, world_size, rank, distributed_init_port))
|
||||
ray.get(refs)
|
||||
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
class TestCustomAllReduce(CustomTestCase):
|
||||
TEST_SIZES = [
|
||||
512,
|
||||
4096,
|
||||
32768,
|
||||
262144,
|
||||
2097152,
|
||||
16777216,
|
||||
33554432,
|
||||
67108864,
|
||||
] # 512B...32MB
|
||||
WORLD_SIZES = [2, 4, 6, 8]
|
||||
TEST_LOOP = 10
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
random.seed(42) # keep the deterministic seed
|
||||
|
||||
def test_graph_allreduce(self):
|
||||
for world_size in self.WORLD_SIZES:
|
||||
if world_size > torch.cuda.device_count():
|
||||
continue
|
||||
multi_process_parallel(world_size, self, self.graph_allreduce)
|
||||
|
||||
def test_eager_allreduce(self):
|
||||
for world_size in self.WORLD_SIZES:
|
||||
if world_size > torch.cuda.device_count():
|
||||
continue
|
||||
multi_process_parallel(world_size, self, self.eager_allreduce)
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def graph_allreduce(self, world_size, rank, distributed_init_port):
|
||||
del os.environ["CUDA_VISIBLE_DEVICES"]
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
distributed_init_method = f"tcp://localhost:{distributed_init_port}"
|
||||
init_distributed_environment(
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
distributed_init_method=distributed_init_method,
|
||||
local_rank=rank,
|
||||
)
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
group = get_tensor_model_parallel_group().device_group
|
||||
|
||||
# Set global server args to avoid "Global server args is not set yet!" error
|
||||
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
|
||||
|
||||
# A small all_reduce for warmup.
|
||||
# this is needed because device communicators might be created lazily
|
||||
# (e.g. NCCL). This will ensure that the communicator is initialized
|
||||
# before any communication happens, so that this group can be used for
|
||||
# graph capture immediately.
|
||||
data = torch.zeros(1)
|
||||
data = data.to(device=device)
|
||||
torch.distributed.all_reduce(data, group=group)
|
||||
torch.cuda.synchronize()
|
||||
del data
|
||||
|
||||
for sz in self.TEST_SIZES:
|
||||
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
||||
for _ in range(self.TEST_LOOP):
|
||||
with graph_capture() as graph_capture_context:
|
||||
# use integers so result matches NCCL exactly
|
||||
inp1 = torch.randint(
|
||||
1,
|
||||
16,
|
||||
(sz,),
|
||||
dtype=dtype,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
inp2 = torch.randint(
|
||||
1,
|
||||
16,
|
||||
(sz,),
|
||||
dtype=dtype,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(
|
||||
graph, stream=graph_capture_context.stream
|
||||
):
|
||||
out1 = tensor_model_parallel_all_reduce(inp1)
|
||||
# the input buffer is immediately modified to test
|
||||
# synchronization
|
||||
dist.all_reduce(inp1, group=group)
|
||||
out2 = tensor_model_parallel_all_reduce(inp2)
|
||||
dist.all_reduce(inp2, group=group)
|
||||
graph.replay()
|
||||
torch.testing.assert_close(out1, inp1)
|
||||
torch.testing.assert_close(out2, inp2)
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def eager_allreduce(self, world_size, rank, distributed_init_port):
|
||||
del os.environ["CUDA_VISIBLE_DEVICES"]
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
distributed_init_method = f"tcp://localhost:{distributed_init_port}"
|
||||
init_distributed_environment(
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
distributed_init_method=distributed_init_method,
|
||||
local_rank=rank,
|
||||
)
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
group = get_tensor_model_parallel_group().device_group
|
||||
|
||||
# Set global server args to avoid "Global server args is not set yet!" error
|
||||
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
|
||||
|
||||
for sz in self.TEST_SIZES:
|
||||
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
||||
for _ in range(self.TEST_LOOP):
|
||||
inp1 = torch.randint(
|
||||
1, 16, (sz,), dtype=dtype, device=torch.cuda.current_device()
|
||||
)
|
||||
out1 = tensor_model_parallel_all_reduce(inp1)
|
||||
dist.all_reduce(inp1, group=group)
|
||||
torch.testing.assert_close(out1, inp1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,318 @@
|
||||
"""
|
||||
Unit tests for DeepSeek chat template tool call handling.
|
||||
|
||||
Tests verify that the DeepSeek chat templates (v3, v3.1, v3.2) correctly handle
|
||||
both dict and string types for tool['function']['arguments'] without double-escaping,
|
||||
addressing issue #11700.
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from jinja2 import Template
|
||||
|
||||
|
||||
class TestDeepSeekChatTemplateToolCalls(unittest.TestCase):
|
||||
"""Test DeepSeek chat templates handle tool calls correctly."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Load all DeepSeek chat templates."""
|
||||
base_path = os.path.join(
|
||||
os.path.dirname(__file__), "..", "..", "examples", "chat_template"
|
||||
)
|
||||
|
||||
cls.templates = {}
|
||||
template_files = {
|
||||
"v3": "tool_chat_template_deepseekv3.jinja",
|
||||
"v3.1": "tool_chat_template_deepseekv31.jinja",
|
||||
"v3.2": "tool_chat_template_deepseekv32.jinja",
|
||||
}
|
||||
|
||||
for version, filename in template_files.items():
|
||||
template_path = os.path.join(base_path, filename)
|
||||
with open(template_path, "r") as f:
|
||||
template_content = f.read()
|
||||
cls.templates[version] = Template(template_content)
|
||||
|
||||
def _render_template(
|
||||
self, version, messages, tools=None, add_generation_prompt=True
|
||||
):
|
||||
"""Helper method to render a template with given messages and tools."""
|
||||
template = self.templates[version]
|
||||
|
||||
# Common template variables
|
||||
context = {
|
||||
"messages": messages,
|
||||
"add_generation_prompt": add_generation_prompt,
|
||||
"bos_token": "<|begin▁of▁sentence|>",
|
||||
}
|
||||
|
||||
if tools is not None:
|
||||
context["tools"] = tools
|
||||
|
||||
return template.render(**context)
|
||||
|
||||
def test_tool_arguments_as_dict(self):
|
||||
"""Test that tool arguments as dict are properly JSON-encoded (normal case)."""
|
||||
# This tests the normal case where arguments come from OpenAI API as dict
|
||||
|
||||
for version in ["v3", "v3.1", "v3.2"]:
|
||||
with self.subTest(version=version):
|
||||
messages = [
|
||||
{"role": "user", "content": "What's the weather in NYC?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": {
|
||||
"city": "New York",
|
||||
"unit": "celsius",
|
||||
}, # Dict
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather information",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {"type": "string"},
|
||||
"unit": {"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
output = self._render_template(version, messages, tools)
|
||||
|
||||
# Should contain properly formatted JSON (not double-escaped)
|
||||
self.assertIn('"city"', output, f"{version}: Should contain city key")
|
||||
self.assertIn(
|
||||
'"New York"', output, f"{version}: Should contain city value"
|
||||
)
|
||||
|
||||
# Should NOT contain double-escaped quotes
|
||||
self.assertNotIn(
|
||||
'\\"city\\"', output, f"{version}: Should not double-escape"
|
||||
)
|
||||
self.assertNotIn(
|
||||
'\\\\"', output, f"{version}: Should not have escaped backslashes"
|
||||
)
|
||||
|
||||
def test_tool_arguments_as_string(self):
|
||||
"""Test that tool arguments as string are used as-is (multi-round case)."""
|
||||
# This tests the multi-round function calling case from issue #11700
|
||||
# where arguments might already be JSON strings from previous model output
|
||||
|
||||
for version in ["v3", "v3.1", "v3.2"]:
|
||||
with self.subTest(version=version):
|
||||
messages = [
|
||||
{"role": "user", "content": "What's the stock price of NVDA?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_stock_info",
|
||||
"arguments": '{"symbol": "NVDA"}', # Already a JSON string
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_stock_info",
|
||||
"description": "Get stock information",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"symbol": {"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
output = self._render_template(version, messages, tools)
|
||||
|
||||
# Should contain the JSON string as-is
|
||||
self.assertIn(
|
||||
'{"symbol": "NVDA"}',
|
||||
output,
|
||||
f"{version}: Should contain JSON as-is",
|
||||
)
|
||||
|
||||
# Should NOT double-escape (the bug from issue #11700)
|
||||
# Bad output would look like: "{\"symbol\": \"NVDA\"}" or "{\\"symbol\\": \\"NVDA\\"}"
|
||||
self.assertNotIn(
|
||||
'{\\"symbol\\"', output, f"{version}: Should not double-escape"
|
||||
)
|
||||
self.assertNotIn(
|
||||
'"{\\"symbol', output, f"{version}: Should not wrap and escape"
|
||||
)
|
||||
|
||||
# Verify it's not triple-quoted or escaped
|
||||
self.assertNotIn(
|
||||
'""{"', output, f"{version}: Should not have extra quotes"
|
||||
)
|
||||
|
||||
def test_multiple_tool_calls_mixed_types(self):
|
||||
"""Test multiple tool calls with mixed dict and string argument types."""
|
||||
# This tests a complex scenario with multiple tools, some with dict args, some with string
|
||||
|
||||
for version in ["v3", "v3.1", "v3.2"]:
|
||||
with self.subTest(version=version):
|
||||
messages = [
|
||||
{"role": "user", "content": "Get weather and stock info"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": {"city": "Boston"}, # Dict
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_stock_info",
|
||||
"arguments": '{"symbol": "TSLA"}', # String
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_stock_info",
|
||||
"description": "Get stock info",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"symbol": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
output = self._render_template(version, messages, tools)
|
||||
|
||||
# First tool (dict) should be properly JSON-encoded
|
||||
self.assertIn(
|
||||
'"city"', output, f"{version}: First tool should have city key"
|
||||
)
|
||||
self.assertIn(
|
||||
'"Boston"',
|
||||
output,
|
||||
f"{version}: First tool should have Boston value",
|
||||
)
|
||||
|
||||
# Second tool (string) should be used as-is
|
||||
self.assertIn(
|
||||
'{"symbol": "TSLA"}',
|
||||
output,
|
||||
f"{version}: Second tool should use string as-is",
|
||||
)
|
||||
|
||||
# Neither should be double-escaped
|
||||
self.assertNotIn(
|
||||
'\\"city\\"',
|
||||
output,
|
||||
f"{version}: First tool should not double-escape",
|
||||
)
|
||||
self.assertNotIn(
|
||||
'\\"symbol\\"',
|
||||
output,
|
||||
f"{version}: Second tool should not double-escape",
|
||||
)
|
||||
|
||||
def test_tool_call_with_content(self):
|
||||
"""Test tool calls that also include content text."""
|
||||
# Some models include explanatory text along with tool calls
|
||||
|
||||
for version in ["v3", "v3.1", "v3.2"]:
|
||||
with self.subTest(version=version):
|
||||
messages = [
|
||||
{"role": "user", "content": "What's the weather?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Let me check the weather for you.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": {"city": "Seattle"},
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
output = self._render_template(version, messages, tools)
|
||||
|
||||
# Should contain both the content and the tool call
|
||||
self.assertIn(
|
||||
"Let me check the weather",
|
||||
output,
|
||||
f"{version}: Should include content",
|
||||
)
|
||||
self.assertIn(
|
||||
'"city"', output, f"{version}: Should include tool arguments"
|
||||
)
|
||||
self.assertNotIn(
|
||||
'\\"city\\"', output, f"{version}: Should not double-escape"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,99 @@
|
||||
"""
|
||||
Usage:
|
||||
cd test/src
|
||||
python3 -m unittest test_deepseek_v32_cp_single_node.TestDeepseekV32CP.test_a_gsm8k
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
write_github_step_summary,
|
||||
)
|
||||
|
||||
FULL_DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp"
|
||||
|
||||
|
||||
class TestDeepseekV32CP(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = FULL_DEEPSEEK_V32_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"8",
|
||||
"--dp",
|
||||
"2",
|
||||
"--enable-dp-attention",
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-num-steps",
|
||||
"3",
|
||||
"--speculative-eagle-topk",
|
||||
"1",
|
||||
"--speculative-num-draft-tokens",
|
||||
"4",
|
||||
"--mem-frac",
|
||||
"0.7",
|
||||
"--cuda-graph-max-bs",
|
||||
"32",
|
||||
"--max-running-requests",
|
||||
"32",
|
||||
"--enable-nsa-prefill-context-parallel",
|
||||
]
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_a_gsm8k(
|
||||
self,
|
||||
): # Append an "a" to make this test run first (alphabetically) to warm up the server
|
||||
requests.get(self.base_url + "/flush_cache")
|
||||
|
||||
args = SimpleNamespace(
|
||||
num_shots=8,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
parallel=32,
|
||||
max_new_tokens=512,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
server_info = requests.get(self.base_url + "/get_server_info")
|
||||
avg_spec_accept_length = server_info.json()["internal_states"][0][
|
||||
"avg_spec_accept_length"
|
||||
]
|
||||
print(f"{avg_spec_accept_length=}")
|
||||
|
||||
if is_in_ci():
|
||||
write_github_step_summary(
|
||||
f"### test_gsm8k (deepseek-v32 nsa-cp)\n"
|
||||
f'{metrics["accuracy"]=:.3f}\n'
|
||||
f"{avg_spec_accept_length=:.2f}\n"
|
||||
)
|
||||
self.assertGreater(metrics["accuracy"], 0.935)
|
||||
self.assertGreater(avg_spec_accept_length, 2.7)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,87 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_DEEPSEEK_NVFP4_MODEL_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
try_cached_model,
|
||||
)
|
||||
|
||||
|
||||
class TestDeepseekR1Nvfp4CuteDSLDeepEP(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = try_cached_model(DEFAULT_DEEPSEEK_NVFP4_MODEL_FOR_TEST)
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--disable-radix-cache",
|
||||
"--mem-fraction-static",
|
||||
"0.89",
|
||||
"--max-prefill-tokens",
|
||||
"16384",
|
||||
"--max-running-requests",
|
||||
"256",
|
||||
"--chunked-prefill-size",
|
||||
"1024",
|
||||
"--tp",
|
||||
"4",
|
||||
"--dp",
|
||||
"4",
|
||||
"--ep",
|
||||
"4",
|
||||
"--moe-dense-tp-size",
|
||||
"1",
|
||||
"--enable-dp-attention",
|
||||
"--quantization",
|
||||
"modelopt_fp4",
|
||||
"--attention-backend",
|
||||
"trtllm_mla",
|
||||
"--moe-a2a-backend",
|
||||
"deepep",
|
||||
"--moe-runner-backend",
|
||||
"flashinfer_cutedsl",
|
||||
"--deepep-mode",
|
||||
"low_latency",
|
||||
]
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
env={
|
||||
**os.environ,
|
||||
"SGLANG_DEEPEP_BF16_DISPATCH": "1",
|
||||
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
|
||||
"SGLANG_CUTEDSL_MOE_NVFP4_DISPATCH": "0",
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=512,
|
||||
parallel=512,
|
||||
max_new_tokens=512,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Eval accuracy of GSM8K: {metrics=}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.92)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,65 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestDoubleSparsity(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
dirpath = os.path.dirname(__file__)
|
||||
config_file = os.path.join(
|
||||
dirpath, "double-sparsity-config-Llama-3.1-8B-Instruct.json"
|
||||
)
|
||||
# NOTE: Generate the config file by running https://github.com/andy-yang-1/DoubleSparse/blob/main/evaluation/group_channel_config.py
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--enable-double-sparsity",
|
||||
"--ds-channel-config-path",
|
||||
config_file,
|
||||
"--ds-heavy-channel-num",
|
||||
"32",
|
||||
"--ds-heavy-channel-type",
|
||||
"k",
|
||||
"--ds-heavy-token-num",
|
||||
"512",
|
||||
"--ds-sparse-decode-threshold",
|
||||
"0",
|
||||
"--max-total-tokens",
|
||||
"200000",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreaterEqual(metrics["score"], 0.65)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,97 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
write_github_step_summary,
|
||||
)
|
||||
|
||||
FULL_DEEPSEEK_V3_FP4_MODEL_PATH = "nvidia/DeepSeek-V3-0324-FP4"
|
||||
|
||||
|
||||
class TestEagleDPAttnServerBase(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
os.environ["SGLANG_ENABLE_SPEC_V2"] = "1"
|
||||
cls.model = FULL_DEEPSEEK_V3_FP4_MODEL_PATH
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = [
|
||||
"--tp-size",
|
||||
"4",
|
||||
"--dp-size",
|
||||
"4",
|
||||
"--enable-dp-attention",
|
||||
"--attention-backend",
|
||||
"trtllm_mla",
|
||||
"--moe-runner-backend",
|
||||
"flashinfer_trtllm",
|
||||
"--quantization",
|
||||
"modelopt_fp4",
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-num-steps",
|
||||
"3",
|
||||
"--speculative-eagle-topk",
|
||||
"1",
|
||||
"--speculative-num-draft-tokens",
|
||||
"4",
|
||||
"--kv-cache-dtype",
|
||||
"fp8_e4m3",
|
||||
]
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
if "SGLANG_ENABLE_SPEC_V2" in os.environ:
|
||||
del os.environ["SGLANG_ENABLE_SPEC_V2"]
|
||||
|
||||
def test_a_gsm8k(
|
||||
self,
|
||||
): # Append an "a" to make this test run first (alphabetically) to warm up the server
|
||||
requests.get(self.base_url + "/flush_cache")
|
||||
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
server_info = requests.get(self.base_url + "/get_server_info")
|
||||
avg_spec_accept_length = server_info.json()["internal_states"][0][
|
||||
"avg_spec_accept_length"
|
||||
]
|
||||
print(f"{avg_spec_accept_length=}")
|
||||
|
||||
if is_in_ci():
|
||||
write_github_step_summary(
|
||||
f"### test_gsm8k (deepseek-v3-fp4 mtp)\n"
|
||||
f'{metrics["accuracy"]=:.3f}\n'
|
||||
f"{avg_spec_accept_length=:.2f}\n"
|
||||
)
|
||||
self.assertGreater(metrics["accuracy"], 0.94)
|
||||
self.assertGreater(avg_spec_accept_length, 2.04)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Executable
+101
@@ -0,0 +1,101 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class TestExpertDistribution(CustomTestCase):
|
||||
def test_expert_distribution_record(self):
|
||||
# TODO: Add tests for DeepEP gatherer (currently our CI cannot run that)
|
||||
for info in [
|
||||
dict(model_path="deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct"),
|
||||
dict(model_path="Qwen/Qwen1.5-MoE-A2.7B"),
|
||||
dict(model_path="Qwen/Qwen1.5-MoE-A2.7B", tp_size=2),
|
||||
dict(model_path="Qwen/Qwen1.5-MoE-A2.7B", mode="per_pass"),
|
||||
dict(model_path="Qwen/Qwen1.5-MoE-A2.7B", mode="per_token"),
|
||||
]:
|
||||
with self.subTest(info=info):
|
||||
self._execute_core(**info)
|
||||
|
||||
def _execute_core(self, model_path: str, mode: str = "stat", tp_size: int = 1):
|
||||
"""Test expert distribution record endpoints"""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
envs.SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR.set(tmp_dir)
|
||||
|
||||
process = popen_launch_server(
|
||||
model_path,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--trust-remote-code",
|
||||
"--tp-size",
|
||||
str(tp_size),
|
||||
"--expert-distribution-recorder-mode",
|
||||
mode,
|
||||
"--disable-cuda-graph",
|
||||
"--disable-overlap-schedule",
|
||||
],
|
||||
)
|
||||
|
||||
try:
|
||||
# Start recording
|
||||
response = requests.post(
|
||||
f"{DEFAULT_URL_FOR_TEST}/start_expert_distribution_record"
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
# Make some requests to generate expert distribution data
|
||||
response = requests.post(
|
||||
f"{DEFAULT_URL_FOR_TEST}/generate",
|
||||
json={
|
||||
"text": "The capital of France is",
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 32,
|
||||
},
|
||||
},
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
# Stop recording
|
||||
response = requests.post(
|
||||
f"{DEFAULT_URL_FOR_TEST}/stop_expert_distribution_record"
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
# Dump the recorded data
|
||||
response = requests.post(
|
||||
f"{DEFAULT_URL_FOR_TEST}/dump_expert_distribution_record"
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
# Check data rows
|
||||
data = torch.load(
|
||||
list(Path(tmp_dir).glob("*.pt"))[0], weights_only=True
|
||||
)
|
||||
print(f"{data=}")
|
||||
|
||||
if mode in ["per_pass", "per_token"]:
|
||||
self.assertGreater(len(data), 0, "Should contain data rows")
|
||||
else:
|
||||
logical_count = data["logical_count"]
|
||||
print(f"{logical_count.sum()=} {logical_count=}")
|
||||
self.assertTrue(logical_count.sum() > 0)
|
||||
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,255 @@
|
||||
import os
|
||||
import traceback
|
||||
import unittest
|
||||
from dataclasses import dataclass
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
import torch.distributed
|
||||
import torch.multiprocessing as mp
|
||||
from torch.multiprocessing import Process
|
||||
|
||||
from sglang.srt.eplb import expert_location_updater
|
||||
from sglang.test.test_utils import CustomTestCase, find_available_port
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
|
||||
@dataclass
|
||||
class _TestInfo:
|
||||
nnodes: int
|
||||
num_logical_experts: int
|
||||
num_physical_experts: int
|
||||
num_repeat: int = 5000
|
||||
|
||||
|
||||
class TestExpertLocationUpdater(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
mp.set_start_method("spawn", force=True)
|
||||
|
||||
def test_cpu(self):
|
||||
self._test_common(device="cpu")
|
||||
self._test_core(
|
||||
num_gpus=32,
|
||||
device="cpu",
|
||||
infos=[
|
||||
_TestInfo(
|
||||
nnodes=4,
|
||||
num_logical_experts=256,
|
||||
num_physical_experts=288,
|
||||
num_repeat=10000,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
def test_cpu_slow(self):
|
||||
if is_in_ci():
|
||||
return
|
||||
self._test_core(
|
||||
num_gpus=144,
|
||||
device="cpu",
|
||||
infos=[
|
||||
_TestInfo(
|
||||
nnodes=18,
|
||||
num_logical_experts=256,
|
||||
num_physical_experts=288,
|
||||
num_repeat=10000,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
def test_gpu(self):
|
||||
if is_in_ci():
|
||||
return
|
||||
self._test_common(device="cuda")
|
||||
|
||||
def _test_common(self, device):
|
||||
infos = []
|
||||
|
||||
for nnodes in [1, 2, 4]:
|
||||
for num_logical_experts in [2, 5, 20, 256]:
|
||||
for num_physical_experts in [8, 16, 256, 288]:
|
||||
if num_logical_experts > num_physical_experts:
|
||||
continue
|
||||
infos.append(
|
||||
_TestInfo(
|
||||
nnodes=nnodes,
|
||||
num_logical_experts=num_logical_experts,
|
||||
num_physical_experts=num_physical_experts,
|
||||
)
|
||||
)
|
||||
|
||||
self._test_core(num_gpus=8, device=device, infos=infos)
|
||||
|
||||
def _test_core(
|
||||
self,
|
||||
num_gpus: int,
|
||||
device: str,
|
||||
infos: List[_TestInfo],
|
||||
):
|
||||
master_port = find_available_port(23456)
|
||||
|
||||
processes = []
|
||||
output_reader, output_writer = mp.Pipe(duplex=False)
|
||||
for rank in range(num_gpus):
|
||||
p = Process(
|
||||
target=_run_subprocess,
|
||||
kwargs=dict(
|
||||
rank=rank,
|
||||
num_gpus=num_gpus,
|
||||
output_writer=output_writer,
|
||||
master_port=master_port,
|
||||
device=device,
|
||||
infos=infos,
|
||||
),
|
||||
)
|
||||
p.start()
|
||||
processes.append(p)
|
||||
|
||||
for _ in range(num_gpus):
|
||||
self.assertTrue(
|
||||
output_reader.recv(), f"Subprocess has error, please see logs above."
|
||||
)
|
||||
|
||||
for p in processes:
|
||||
p.join()
|
||||
|
||||
|
||||
def _run_subprocess(
|
||||
rank: int,
|
||||
num_gpus: int,
|
||||
master_port: int,
|
||||
device: str,
|
||||
infos: List[_TestInfo],
|
||||
output_writer,
|
||||
):
|
||||
try:
|
||||
os.environ["MASTER_ADDR"] = "localhost"
|
||||
os.environ["MASTER_PORT"] = str(master_port)
|
||||
|
||||
torch.random.manual_seed(42)
|
||||
torch.distributed.init_process_group(
|
||||
rank=rank,
|
||||
world_size=num_gpus,
|
||||
backend={"cpu": "gloo", "cuda": None}[device],
|
||||
)
|
||||
if device == "cuda":
|
||||
torch.cuda.set_device(f"cuda:{rank}")
|
||||
|
||||
for info in infos:
|
||||
_execute_test(info, rank=rank, num_gpus=num_gpus, device=device)
|
||||
|
||||
execution_ok = True
|
||||
except Exception as e:
|
||||
print(f"subprocess[{rank=}] has error: {e}", flush=True)
|
||||
traceback.print_exc()
|
||||
execution_ok = False
|
||||
|
||||
output_writer.send(execution_ok)
|
||||
output_writer.close()
|
||||
|
||||
|
||||
def _execute_test(info: _TestInfo, rank: int, num_gpus: int, device: str):
|
||||
if rank == 0:
|
||||
print(f"Test: {num_gpus=} {info=}", flush=True)
|
||||
|
||||
assert info.num_physical_experts % num_gpus == 0
|
||||
num_local_physical_experts = info.num_physical_experts // num_gpus
|
||||
assert num_gpus % info.nnodes == 0
|
||||
num_gpu_per_node = num_gpus // info.nnodes
|
||||
|
||||
def _create_routed_experts_weights(physical_to_logical_map):
|
||||
local_logical_expert_ids = physical_to_logical_map[
|
||||
rank * num_local_physical_experts : (rank + 1) * num_local_physical_experts
|
||||
].cpu()
|
||||
return [
|
||||
local_logical_expert_ids.to(device).clone(),
|
||||
torch.tensor(
|
||||
[
|
||||
[local_logical_expert_id * 10, local_logical_expert_id * 100]
|
||||
for local_logical_expert_id in local_logical_expert_ids.tolist()
|
||||
],
|
||||
device=device,
|
||||
),
|
||||
]
|
||||
|
||||
def _create_physical_to_logical_map():
|
||||
if rank == 0:
|
||||
ans = torch.concat(
|
||||
[
|
||||
torch.arange(0, info.num_logical_experts),
|
||||
torch.randint(
|
||||
0,
|
||||
info.num_logical_experts,
|
||||
(info.num_physical_experts - info.num_logical_experts,),
|
||||
),
|
||||
]
|
||||
)
|
||||
ans = ans[torch.randperm(ans.shape[0])]
|
||||
else:
|
||||
ans = torch.empty((info.num_physical_experts,), dtype=torch.int64)
|
||||
|
||||
assert ans.dtype == torch.int64 and ans.shape == (info.num_physical_experts,)
|
||||
ans = ans.to(device)
|
||||
torch.distributed.broadcast(ans, src=0)
|
||||
|
||||
return ans.cpu()
|
||||
|
||||
physical_to_logical_map = _create_physical_to_logical_map()
|
||||
routed_experts_weights = _create_routed_experts_weights(physical_to_logical_map)
|
||||
|
||||
for i in range(info.num_repeat):
|
||||
if rank == 0 and ((i % 500 == 0) or (i == info.num_repeat - 1)):
|
||||
print(f"Step {i}/{info.num_repeat}", flush=True)
|
||||
|
||||
new_physical_to_logical_map = _create_physical_to_logical_map()
|
||||
expect_new_weights = _create_routed_experts_weights(new_physical_to_logical_map)
|
||||
|
||||
output_logs = expert_location_updater.update_expert_weights_single_layer(
|
||||
routed_experts_weights=routed_experts_weights,
|
||||
temp_buffers=expert_location_updater.create_temp_buffers(
|
||||
routed_experts_weights
|
||||
),
|
||||
old_physical_to_logical_map=physical_to_logical_map.tolist(),
|
||||
new_physical_to_logical_map=new_physical_to_logical_map.tolist(),
|
||||
num_local_physical_experts=num_local_physical_experts,
|
||||
num_gpu_per_node=num_gpu_per_node,
|
||||
rank=rank,
|
||||
debug=True,
|
||||
)
|
||||
|
||||
local_has_error = not all(
|
||||
torch.all(x == y)
|
||||
for x, y in zip(routed_experts_weights, expect_new_weights, strict=True)
|
||||
)
|
||||
global_has_error = torch.tensor(local_has_error, device=device)
|
||||
torch.distributed.all_reduce(
|
||||
global_has_error, op=torch.distributed.ReduceOp.MAX
|
||||
)
|
||||
|
||||
if global_has_error.cpu().item():
|
||||
output_logs_str = "\n".join(output_logs)
|
||||
local_message = (
|
||||
f"===================== rank {rank} ============================\n"
|
||||
f"{num_gpus=} {info=}\n"
|
||||
f"{routed_experts_weights[0].tolist()=}\n"
|
||||
f"{expect_new_weights[0].tolist()=}\n"
|
||||
f"{physical_to_logical_map.tolist()=}\n"
|
||||
f"{new_physical_to_logical_map.tolist()=}\n"
|
||||
f"===logs===\n"
|
||||
f"{output_logs_str}\n"
|
||||
f"==============================================================\n"
|
||||
)
|
||||
|
||||
global_messages = ([None] * num_gpus) if rank == 0 else None
|
||||
torch.distributed.gather_object(local_message, global_messages, dst=0)
|
||||
|
||||
if rank == 0:
|
||||
print("\n\n".join(global_messages), flush=True)
|
||||
raise AssertionError(f"Error happens, see logs above")
|
||||
|
||||
physical_to_logical_map = new_physical_to_logical_map
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,72 @@
|
||||
import unittest
|
||||
|
||||
import openai
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class TestFimCompletion(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "deepseek-ai/deepseek-coder-1.3b-base"
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
other_args = ["--completion-template", "deepseek_coder"]
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
other_args=other_args,
|
||||
)
|
||||
cls.base_url += "/v1"
|
||||
cls.tokenizer = get_tokenizer(cls.model)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def run_fim_completion(self, number_of_completion):
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
prompt = "function sum(a: number, b: number): number{\n"
|
||||
suffix = "}"
|
||||
|
||||
prompt_input = self.tokenizer.encode(prompt) + self.tokenizer.encode(suffix)
|
||||
num_prompt_tokens = len(prompt_input) + 2
|
||||
|
||||
response = client.completions.create(
|
||||
model=self.model,
|
||||
prompt=prompt,
|
||||
suffix=suffix,
|
||||
temperature=0.3,
|
||||
max_tokens=32,
|
||||
stream=False,
|
||||
n=number_of_completion,
|
||||
)
|
||||
|
||||
print(response)
|
||||
print(len(response.choices))
|
||||
assert len(response.choices) == number_of_completion
|
||||
assert response.id
|
||||
assert response.created
|
||||
assert response.object == "text_completion"
|
||||
assert (
|
||||
response.usage.prompt_tokens == num_prompt_tokens
|
||||
), f"{response.usage.prompt_tokens} vs {num_prompt_tokens}"
|
||||
assert response.usage.completion_tokens > 0
|
||||
assert response.usage.total_tokens > 0
|
||||
|
||||
def test_fim_completion(self):
|
||||
for number_of_completion in [1, 3]:
|
||||
self.run_fim_completion(number_of_completion)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,305 @@
|
||||
"""
|
||||
Test forward_split_prefill functionality.
|
||||
|
||||
Usage:
|
||||
python3 -m unittest test_forward_split_prefill.TestForwardSplitPrefill
|
||||
or
|
||||
python3 test_forward_split_prefill.py
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.model_executor.model_runner import ModelRunner
|
||||
from sglang.srt.sampling.sampling_params import SamplingParams
|
||||
from sglang.srt.server_args import PortArgs, ServerArgs
|
||||
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
||||
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
||||
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
|
||||
|
||||
|
||||
class TestForwardSplitPrefill(CustomTestCase):
|
||||
"""Test cases for forward_split_prefill functionality."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Set up the test environment once for all tests."""
|
||||
cls.model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.tp_size = 1
|
||||
cls.device = "cuda"
|
||||
|
||||
# Initialize server args
|
||||
cls.server_args = ServerArgs(
|
||||
model_path=cls.model_path,
|
||||
tokenizer_path=cls.model_path,
|
||||
host="127.0.0.1",
|
||||
disable_cuda_graph=True, # Disable CUDA graph for testing split prefill
|
||||
disable_hybrid_swa_memory=True,
|
||||
port=30000,
|
||||
tp_size=cls.tp_size,
|
||||
mem_fraction_static=0.8,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
|
||||
cls.port_args = PortArgs.init_new(cls.server_args)
|
||||
|
||||
# Load model and tokenizer
|
||||
cls.model_config = ModelConfig.from_server_args(cls.server_args)
|
||||
cls.model_runner = ModelRunner(
|
||||
model_config=cls.model_config,
|
||||
mem_fraction_static=cls.server_args.mem_fraction_static,
|
||||
gpu_id=0,
|
||||
tp_rank=0,
|
||||
tp_size=cls.tp_size,
|
||||
pp_rank=0,
|
||||
pp_size=1,
|
||||
nccl_port=cls.port_args.nccl_port,
|
||||
server_args=cls.server_args,
|
||||
)
|
||||
|
||||
cls.tokenizer = get_tokenizer(
|
||||
cls.server_args.tokenizer_path,
|
||||
tokenizer_mode=cls.server_args.tokenizer_mode,
|
||||
trust_remote_code=cls.server_args.trust_remote_code,
|
||||
)
|
||||
|
||||
print(
|
||||
f"Test with model: {cls.model_path}, num_hidden_layers: {cls.model_config.num_hidden_layers}"
|
||||
)
|
||||
|
||||
def prepare_test_batch(self, batch_size=2, input_len=128, is_split_prefill=True):
|
||||
"""Prepare a test batch for split prefill testing."""
|
||||
# Create synthetic input
|
||||
input_ids = np.random.randint(10, 1000, (batch_size, input_len), dtype=np.int32)
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.0,
|
||||
max_new_tokens=8,
|
||||
)
|
||||
|
||||
reqs = []
|
||||
for i in range(batch_size):
|
||||
req = Req(
|
||||
rid=i,
|
||||
origin_input_text="",
|
||||
origin_input_ids=list(input_ids[i]),
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
req.fill_ids = req.origin_input_ids
|
||||
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
|
||||
req.logprob_start_len = len(req.origin_input_ids) - 1
|
||||
reqs.append(req)
|
||||
|
||||
# Create dummy tree_cache for tests (no prefix caching, just allocation)
|
||||
dummy_tree_cache = SimpleNamespace(
|
||||
page_size=1,
|
||||
device=self.model_runner.device,
|
||||
token_to_kv_pool_allocator=self.model_runner.token_to_kv_pool_allocator,
|
||||
)
|
||||
|
||||
batch = ScheduleBatch.init_new(
|
||||
reqs=reqs,
|
||||
req_to_token_pool=self.model_runner.req_to_token_pool,
|
||||
token_to_kv_pool_allocator=self.model_runner.token_to_kv_pool_allocator,
|
||||
tree_cache=dummy_tree_cache,
|
||||
model_config=self.model_config,
|
||||
enable_overlap=False,
|
||||
spec_algorithm=SpeculativeAlgorithm.NONE,
|
||||
enable_custom_logit_processor=False,
|
||||
)
|
||||
if is_split_prefill:
|
||||
batch.prepare_for_split_prefill()
|
||||
else:
|
||||
batch.prepare_for_extend()
|
||||
|
||||
# Create forward batch
|
||||
model_worker_batch = batch.get_model_worker_batch()
|
||||
forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
|
||||
|
||||
return forward_batch
|
||||
|
||||
def test_split_prefill_functionality(self):
|
||||
"""Test that split prefill can complete successfully."""
|
||||
print("\n=== Testing split prefill functionality ===")
|
||||
|
||||
forward_batch = self.prepare_test_batch(batch_size=2, input_len=64)
|
||||
|
||||
# Reset split index
|
||||
forward_batch.split_index = 0
|
||||
|
||||
# Test split prefill in chunks
|
||||
num_layers = self.model_config.num_hidden_layers
|
||||
chunk_size = max(1, num_layers // 4) # Split into 4 chunks
|
||||
|
||||
results = []
|
||||
split_count = 0
|
||||
|
||||
while forward_batch.split_index < num_layers:
|
||||
print(
|
||||
f"Processing split {split_count}, split_index: {forward_batch.split_index}"
|
||||
)
|
||||
|
||||
result = self.model_runner.forward_split_prefill(
|
||||
forward_batch=forward_batch,
|
||||
reinit_attn_backend=(split_count == 0),
|
||||
forward_count=chunk_size,
|
||||
)
|
||||
|
||||
results.append(result)
|
||||
split_count += 1
|
||||
|
||||
# Verify split_index is updated correctly
|
||||
expected_next_index = min(split_count * chunk_size, num_layers)
|
||||
self.assertEqual(forward_batch.split_index, expected_next_index)
|
||||
|
||||
# The last result should contain logits
|
||||
self.assertIsNotNone(results[-1], "Final split should return logits")
|
||||
print(f"Split prefill completed in {split_count} splits")
|
||||
|
||||
def test_split_prefill_vs_normal_prefill(self):
|
||||
"""Test that split prefill produces the same results as normal prefill."""
|
||||
print("\n=== Testing split prefill vs normal prefill consistency ===")
|
||||
|
||||
forward_batch_normal = self.prepare_test_batch(
|
||||
batch_size=2, input_len=128, is_split_prefill=False
|
||||
)
|
||||
forward_batch_split = self.prepare_test_batch(
|
||||
batch_size=2, input_len=128, is_split_prefill=True
|
||||
)
|
||||
|
||||
# Ensure same input
|
||||
forward_batch_split.input_ids = forward_batch_normal.input_ids.clone()
|
||||
forward_batch_split.positions = forward_batch_normal.positions.clone()
|
||||
|
||||
# Method 1: Normal extend (prefill)
|
||||
print("Running normal extend (prefill)...")
|
||||
normal_result = self.model_runner.forward_extend(forward_batch_normal)
|
||||
|
||||
# Method 2: Split prefill
|
||||
print("Running split prefill...")
|
||||
num_layers = self.model_config.num_hidden_layers
|
||||
chunk_size = max(1, num_layers // 3) # Split into 3 chunks
|
||||
|
||||
split_result = None
|
||||
|
||||
while forward_batch_split.split_index < num_layers:
|
||||
result = self.model_runner.forward_split_prefill(
|
||||
forward_batch=forward_batch_split,
|
||||
forward_count=chunk_size,
|
||||
)
|
||||
if result is not None:
|
||||
split_result = result
|
||||
|
||||
# Compare results
|
||||
self.assertIsNotNone(normal_result, "Normal prefill should return result")
|
||||
self.assertIsNotNone(split_result, "Split prefill should return result")
|
||||
|
||||
# Compare logits shapes
|
||||
self.assertEqual(
|
||||
normal_result.next_token_logits.shape,
|
||||
split_result.next_token_logits.shape,
|
||||
"Logits shapes should match",
|
||||
)
|
||||
|
||||
# Compare logits values (should be very close due to same computation)
|
||||
# Use a larger tolerance for numerical differences in split computation
|
||||
torch.testing.assert_close(
|
||||
normal_result.next_token_logits,
|
||||
split_result.next_token_logits,
|
||||
rtol=1e-3,
|
||||
atol=1e-3,
|
||||
msg="Split prefill and normal prefill should produce similar logits",
|
||||
)
|
||||
|
||||
print("✓ Split prefill and normal prefill produce consistent results")
|
||||
|
||||
def test_split_prefill_different_chunk_sizes(self):
|
||||
"""Test split prefill with different chunk sizes."""
|
||||
print("\n=== Testing split prefill with different chunk sizes ===")
|
||||
|
||||
num_layers = self.model_config.num_hidden_layers
|
||||
chunk_sizes = [1, 2, max(1, num_layers // 2), num_layers]
|
||||
|
||||
# Prepare identical batches for each test
|
||||
base_batch = self.prepare_test_batch(batch_size=1, input_len=16)
|
||||
base_input_ids = base_batch.input_ids.clone()
|
||||
base_positions = base_batch.positions.clone()
|
||||
|
||||
results = []
|
||||
|
||||
for chunk_size in chunk_sizes:
|
||||
if chunk_size > num_layers:
|
||||
continue
|
||||
|
||||
print(f"Testing chunk size: {chunk_size}")
|
||||
|
||||
# Prepare fresh batch
|
||||
forward_batch = self.prepare_test_batch(batch_size=1, input_len=16)
|
||||
forward_batch.input_ids = base_input_ids.clone()
|
||||
forward_batch.positions = base_positions.clone()
|
||||
forward_batch.split_index = 0
|
||||
|
||||
# Run split prefill
|
||||
split_result = None
|
||||
|
||||
while forward_batch.split_index < num_layers:
|
||||
result = self.model_runner.forward_split_prefill(
|
||||
forward_batch=forward_batch,
|
||||
forward_count=chunk_size,
|
||||
)
|
||||
if result is not None:
|
||||
split_result = result
|
||||
|
||||
self.assertIsNotNone(
|
||||
split_result,
|
||||
f"Split prefill should succeed with chunk_size={chunk_size}",
|
||||
)
|
||||
results.append(split_result)
|
||||
|
||||
# Compare all results should be identical (same input, same computation)
|
||||
if len(results) > 1:
|
||||
for i, result in enumerate(results[1:], 1):
|
||||
torch.testing.assert_close(
|
||||
results[0].next_token_logits,
|
||||
result.next_token_logits,
|
||||
rtol=1e-3,
|
||||
atol=1e-3,
|
||||
msg=f"Results with different chunk sizes should be identical (chunk_size {chunk_sizes[i]})",
|
||||
)
|
||||
|
||||
print("✓ All chunk sizes produce consistent results")
|
||||
|
||||
def test_split_prefill_edge_cases(self):
|
||||
"""Test edge cases for split prefill."""
|
||||
print("\n=== Testing split prefill edge cases ===")
|
||||
|
||||
# Test with single layer chunks
|
||||
forward_batch = self.prepare_test_batch(batch_size=1, input_len=8)
|
||||
|
||||
# Process one layer at a time
|
||||
num_layers = self.model_config.num_hidden_layers
|
||||
for layer_idx in range(num_layers):
|
||||
result = self.model_runner.forward_split_prefill(
|
||||
forward_batch=forward_batch,
|
||||
reinit_attn_backend=(layer_idx == 0),
|
||||
forward_count=1, # One layer at a time
|
||||
)
|
||||
|
||||
if layer_idx == num_layers - 1:
|
||||
# Last layer should return result
|
||||
self.assertIsNotNone(result, "Last layer should return logits")
|
||||
else:
|
||||
# Intermediate layers should return None
|
||||
self.assertIsNone(result, f"Layer {layer_idx} should return None")
|
||||
|
||||
print("✓ Single layer processing works correctly")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,183 @@
|
||||
import gc
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
from sglang.utils import terminate_process
|
||||
|
||||
|
||||
def _process_return(ret):
|
||||
if isinstance(ret, list) and len(ret) == 2:
|
||||
print(f"running assert_allclose on data parallel")
|
||||
np.testing.assert_allclose(ret[0], ret[1])
|
||||
return np.array(ret[0])
|
||||
return np.array(ret)
|
||||
|
||||
|
||||
class TestGetWeightsByName(CustomTestCase):
|
||||
|
||||
def init_hf_model(self, model_name, tie_word_embeddings):
|
||||
self.hf_model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name, torch_dtype="bfloat16", tie_word_embeddings=tie_word_embeddings
|
||||
).to("cuda:0")
|
||||
|
||||
def init_backend(self, backend, dp, tp, model_name):
|
||||
self.backend = backend
|
||||
self.dp = dp
|
||||
self.tp = tp
|
||||
if backend == "Engine":
|
||||
self.engine = sgl.Engine(
|
||||
model_path=model_name,
|
||||
random_seed=42,
|
||||
tp_size=tp,
|
||||
dp_size=dp,
|
||||
)
|
||||
else:
|
||||
self.process = popen_launch_server(
|
||||
model_name,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=(
|
||||
"--tp-size",
|
||||
str(tp),
|
||||
"--dp-size",
|
||||
str(dp),
|
||||
),
|
||||
)
|
||||
|
||||
def clean_up(self):
|
||||
del self.hf_model
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
if self.backend == "Engine":
|
||||
self.engine.shutdown()
|
||||
else:
|
||||
terminate_process(self.process)
|
||||
|
||||
def assert_tie_word_embeddings(self, truncate_size):
|
||||
print("assert_tie_word_embeddings")
|
||||
if self.backend == "Engine":
|
||||
backend_ret = _process_return(
|
||||
self.engine.get_weights_by_name("lm_head.weight", truncate_size)
|
||||
)
|
||||
else:
|
||||
backend_ret = _process_return(
|
||||
requests.get(
|
||||
f"{DEFAULT_URL_FOR_TEST}/get_weights_by_name",
|
||||
json={"name": "lm_head.weight", "truncate_size": truncate_size},
|
||||
).json()
|
||||
)
|
||||
print("assert_tie_word_embeddings of hf and backend")
|
||||
assert np.allclose(
|
||||
self.hf_model.get_parameter("model.embed_tokens.weight")
|
||||
.cpu()
|
||||
.detach()
|
||||
.float()
|
||||
.numpy()[:truncate_size],
|
||||
backend_ret,
|
||||
)
|
||||
assert np.allclose(
|
||||
self.hf_model.get_parameter("lm_head.weight")
|
||||
.cpu()
|
||||
.detach()
|
||||
.float()
|
||||
.numpy()[:truncate_size],
|
||||
self.hf_model.get_parameter("model.embed_tokens.weight")
|
||||
.cpu()
|
||||
.detach()
|
||||
.float()
|
||||
.numpy()[:truncate_size],
|
||||
)
|
||||
|
||||
def assert_weights_all_close(self, param_name, truncate_size):
|
||||
print(
|
||||
f"param_name: {param_name}, backend: {self.backend}, dp: {self.dp}, tp: {self.tp}"
|
||||
)
|
||||
param = self.hf_model.get_parameter(param_name)[:truncate_size]
|
||||
param_np = param.cpu().detach().float().numpy()
|
||||
|
||||
if self.backend == "Engine":
|
||||
engine_ret = self.engine.get_weights_by_name(param_name, truncate_size)
|
||||
engine_ret = _process_return(engine_ret)
|
||||
np.testing.assert_allclose(engine_ret, param_np, rtol=1e-5, atol=1e-5)
|
||||
|
||||
if self.backend == "Runtime":
|
||||
runtime_ret = requests.get(
|
||||
f"{DEFAULT_URL_FOR_TEST}/get_weights_by_name",
|
||||
json={"name": param_name, "truncate_size": truncate_size},
|
||||
).json()
|
||||
runtime_ret = _process_return(runtime_ret)
|
||||
np.testing.assert_allclose(runtime_ret, param_np, rtol=1e-5, atol=1e-5)
|
||||
|
||||
def test_get_weights_by_name(self):
|
||||
if is_in_ci():
|
||||
test_suits = [
|
||||
("Engine", 1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST),
|
||||
]
|
||||
else:
|
||||
test_suits = [
|
||||
("Runtime", 1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST),
|
||||
("Engine", 1, 1, DEFAULT_MODEL_NAME_FOR_TEST),
|
||||
]
|
||||
if torch.cuda.device_count() >= 2:
|
||||
test_suits.append(("Engine", 1, 2, DEFAULT_SMALL_MODEL_NAME_FOR_TEST))
|
||||
test_suits.append(("Runtime", 2, 1, DEFAULT_MODEL_NAME_FOR_TEST))
|
||||
|
||||
if torch.cuda.device_count() >= 4:
|
||||
test_suits.extend(
|
||||
[
|
||||
("Engine", 2, 2, DEFAULT_SMALL_MODEL_NAME_FOR_TEST),
|
||||
("Runtime", 2, 2, DEFAULT_MODEL_NAME_FOR_TEST),
|
||||
]
|
||||
)
|
||||
|
||||
parameters = [
|
||||
"model.embed_tokens.weight",
|
||||
"model.layers.0.input_layernorm.weight",
|
||||
"model.layers.1.self_attn.q_proj.weight",
|
||||
"model.layers.2.self_attn.k_proj.weight",
|
||||
"model.layers.3.self_attn.v_proj.weight",
|
||||
"model.layers.4.self_attn.o_proj.weight",
|
||||
"model.layers.5.mlp.gate_proj.weight",
|
||||
"model.layers.6.mlp.up_proj.weight",
|
||||
"model.layers.7.mlp.down_proj.weight",
|
||||
"model.layers.8.post_attention_layernorm.weight",
|
||||
"model.norm.weight",
|
||||
"lm_head.weight",
|
||||
]
|
||||
|
||||
truncate_size = 100
|
||||
|
||||
for test_suit in test_suits:
|
||||
if test_suit[-1] == DEFAULT_MODEL_NAME_FOR_TEST:
|
||||
tie_word_embeddings = False
|
||||
else:
|
||||
tie_word_embeddings = True
|
||||
|
||||
self.init_hf_model(test_suit[-1], tie_word_embeddings)
|
||||
self.init_backend(*test_suit)
|
||||
|
||||
for param_name in parameters:
|
||||
self.assert_weights_all_close(param_name, truncate_size)
|
||||
|
||||
if tie_word_embeddings:
|
||||
self.assert_tie_word_embeddings(truncate_size)
|
||||
|
||||
self.clean_up()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,28 @@
|
||||
import unittest
|
||||
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestHealthCheck(CustomTestCase):
|
||||
def test_health_check(self):
|
||||
"""Test that metrics endpoint returns data when enabled"""
|
||||
with self.assertRaises(TimeoutError):
|
||||
popen_launch_server(
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
timeout=60,
|
||||
other_args=[
|
||||
"--disable-cuda-graph",
|
||||
"--json-model-override-args",
|
||||
'{"architectures": ["LlamaForCausalLMForHealthTest"]}',
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,374 @@
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
import zmq
|
||||
from msgspec.msgpack import Decoder
|
||||
|
||||
from sglang.srt.disaggregation.kv_events import (
|
||||
AllBlocksCleared,
|
||||
BlockRemoved,
|
||||
BlockStored,
|
||||
KVEventBatch,
|
||||
)
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class TestKvEvents(CustomTestCase):
|
||||
def test_kv_events_enabled(self):
|
||||
"""Test that kv events are sent and received by subscriber data when enabled"""
|
||||
|
||||
# Launch kv events subscriber
|
||||
decoder = Decoder(type=KVEventBatch)
|
||||
context = zmq.Context()
|
||||
sub = context.socket(zmq.SUB)
|
||||
sub.connect("tcp://localhost:5557")
|
||||
topic = "kv-events"
|
||||
sub.setsockopt_string(zmq.SUBSCRIBE, topic)
|
||||
|
||||
# Launch sglang server
|
||||
process = popen_launch_server(
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--kv-events-config",
|
||||
'{"publisher": "zmq", "topic": "kv-events"}',
|
||||
"--max-total-tokens",
|
||||
32,
|
||||
"--cuda-graph-max-bs",
|
||||
2,
|
||||
"--enable-dp-attention",
|
||||
"--dp-size",
|
||||
1,
|
||||
],
|
||||
)
|
||||
|
||||
try:
|
||||
# Make some requests to generate some metrics
|
||||
response = requests.get(f"{DEFAULT_URL_FOR_TEST}/health_generate")
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
response = requests.post(
|
||||
f"{DEFAULT_URL_FOR_TEST}/generate",
|
||||
json={
|
||||
"text": "The capital of France is",
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 32,
|
||||
},
|
||||
},
|
||||
)
|
||||
response = requests.post(
|
||||
f"{DEFAULT_URL_FOR_TEST}/generate",
|
||||
json={
|
||||
"text": "The capital of Spain is",
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 32,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
# Expected events. These may be dependent on model used (meta-llama/Llama-3.2-1B-Instruct)
|
||||
expected_events = [
|
||||
# <begin> The capital city of France is
|
||||
BlockStored(
|
||||
block_hashes=[-6650323075460941099],
|
||||
parent_block_hash=5740354900026072187,
|
||||
token_ids=[128000, 791, 6864, 3363, 315, 9822, 374],
|
||||
block_size=7,
|
||||
lora_id=None,
|
||||
),
|
||||
# Paris. The Eiffel Tower
|
||||
BlockStored(
|
||||
block_hashes=[-7584018293207282755],
|
||||
parent_block_hash=-6650323075460941099,
|
||||
token_ids=[12366, 13, 578, 469, 3168, 301, 22703],
|
||||
block_size=7,
|
||||
lora_id=None,
|
||||
),
|
||||
BlockStored(
|
||||
block_hashes=[-8753497827991233192],
|
||||
parent_block_hash=5740354900026072187,
|
||||
token_ids=[0],
|
||||
block_size=1,
|
||||
lora_id=None,
|
||||
),
|
||||
BlockRemoved(block_hashes=[-6650323075460941099]),
|
||||
# <begin> The capital
|
||||
BlockStored(
|
||||
block_hashes=[-2697055055087824455],
|
||||
parent_block_hash=5740354900026072187,
|
||||
token_ids=[128000, 791, 6864],
|
||||
block_size=3,
|
||||
lora_id=None,
|
||||
),
|
||||
# city of France is
|
||||
BlockStored(
|
||||
block_hashes=[-7505627135785778022],
|
||||
parent_block_hash=-2697055055087824455,
|
||||
token_ids=[3363, 315, 9822, 374],
|
||||
block_size=4,
|
||||
lora_id=None,
|
||||
),
|
||||
# of France is
|
||||
BlockStored(
|
||||
block_hashes=[-3861108700662737012],
|
||||
parent_block_hash=-2697055055087824455,
|
||||
token_ids=[315, 9822, 374],
|
||||
block_size=3,
|
||||
lora_id=None,
|
||||
),
|
||||
BlockRemoved(block_hashes=[-7584018293207282755]),
|
||||
BlockRemoved(block_hashes=[-8753497827991233192]),
|
||||
BlockRemoved(block_hashes=[-7505627135785778022]),
|
||||
# Paris. The Eiffel Tower is located in Paris. The Eiffel Tower is a famous landmark in Paris
|
||||
BlockStored(
|
||||
block_hashes=[-3064341286825792715],
|
||||
parent_block_hash=-3861108700662737012,
|
||||
token_ids=[
|
||||
12366,
|
||||
13,
|
||||
578,
|
||||
469,
|
||||
3168,
|
||||
301,
|
||||
22703,
|
||||
374,
|
||||
7559,
|
||||
304,
|
||||
12366,
|
||||
13,
|
||||
578,
|
||||
469,
|
||||
3168,
|
||||
301,
|
||||
22703,
|
||||
374,
|
||||
264,
|
||||
11495,
|
||||
38350,
|
||||
304,
|
||||
12366,
|
||||
],
|
||||
block_size=23,
|
||||
lora_id=None,
|
||||
),
|
||||
BlockRemoved(block_hashes=[-3861108700662737012]),
|
||||
# of
|
||||
BlockStored(
|
||||
block_hashes=[6115672085296369592],
|
||||
parent_block_hash=-2697055055087824455,
|
||||
token_ids=[315],
|
||||
block_size=1,
|
||||
lora_id=None,
|
||||
),
|
||||
# France is
|
||||
BlockStored(
|
||||
block_hashes=[4208810872343132234],
|
||||
parent_block_hash=6115672085296369592,
|
||||
token_ids=[9822, 374],
|
||||
block_size=2,
|
||||
lora_id=None,
|
||||
),
|
||||
# Spain is
|
||||
BlockStored(
|
||||
block_hashes=[1675819893649989955],
|
||||
parent_block_hash=6115672085296369592,
|
||||
token_ids=[18157, 374],
|
||||
block_size=2,
|
||||
lora_id=None,
|
||||
),
|
||||
BlockRemoved(block_hashes=[-3064341286825792715]),
|
||||
# Madrid. The capital of France is Paris. The capital of Italy is Rome. The capital of Spain is Madrid.
|
||||
BlockStored(
|
||||
block_hashes=[-8505834929190027295],
|
||||
parent_block_hash=1675819893649989955,
|
||||
token_ids=[
|
||||
25048,
|
||||
13,
|
||||
578,
|
||||
6864,
|
||||
315,
|
||||
9822,
|
||||
374,
|
||||
12366,
|
||||
13,
|
||||
578,
|
||||
6864,
|
||||
315,
|
||||
15704,
|
||||
374,
|
||||
22463,
|
||||
13,
|
||||
578,
|
||||
6864,
|
||||
315,
|
||||
18157,
|
||||
374,
|
||||
25048,
|
||||
13,
|
||||
],
|
||||
block_size=23,
|
||||
lora_id=None,
|
||||
),
|
||||
]
|
||||
|
||||
# Get events
|
||||
events = []
|
||||
start = time.time()
|
||||
max_wait_s = 5
|
||||
while (
|
||||
len(events) < len(expected_events)
|
||||
and (time.time() - start) < max_wait_s
|
||||
):
|
||||
_, seq_bytes, payload = sub.recv_multipart()
|
||||
event_batch = decoder.decode(payload)
|
||||
for event in event_batch.events:
|
||||
events.append(event)
|
||||
|
||||
for expected in expected_events:
|
||||
self.assertIn(expected, events)
|
||||
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
def test_kv_events_attn_dp(self):
|
||||
"""Test that kv events are properly tagged with DP rank in attention DP mode"""
|
||||
|
||||
# Launch multiple subscribers for different DP ranks
|
||||
decoder = Decoder(type=KVEventBatch)
|
||||
context = zmq.Context()
|
||||
|
||||
# Subscribe to both DP rank endpoints
|
||||
sub_dp0 = context.socket(zmq.SUB)
|
||||
sub_dp0.connect("tcp://localhost:5557") # DP rank 0
|
||||
topic = "kv-events"
|
||||
sub_dp0.setsockopt_string(zmq.SUBSCRIBE, topic)
|
||||
|
||||
sub_dp1 = context.socket(zmq.SUB)
|
||||
sub_dp1.connect("tcp://localhost:5558") # DP rank 1 (offset by rank)
|
||||
sub_dp1.setsockopt_string(zmq.SUBSCRIBE, topic)
|
||||
|
||||
# Launch sglang server with DP attention enabled
|
||||
process = popen_launch_server(
|
||||
"silence09/DeepSeek-R1-Small-2layers",
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--kv-events-config",
|
||||
'{"publisher": "zmq", "topic": "kv-events"}',
|
||||
"--max-total-tokens",
|
||||
64,
|
||||
"--cuda-graph-max-bs",
|
||||
4,
|
||||
"--enable-dp-attention",
|
||||
"--dp-size",
|
||||
2,
|
||||
"--tp-size",
|
||||
2,
|
||||
],
|
||||
)
|
||||
|
||||
try:
|
||||
# Make requests to generate events
|
||||
response = requests.get(f"{DEFAULT_URL_FOR_TEST}/health_generate")
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
# Send multiple requests to trigger events from both DP ranks
|
||||
for i in range(4):
|
||||
response = requests.post(
|
||||
f"{DEFAULT_URL_FOR_TEST}/generate",
|
||||
json={
|
||||
"text": f"Request {i}: The capital of country {i} is",
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 16,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
# Collect events from both DP ranks
|
||||
events_dp0 = []
|
||||
events_dp1 = []
|
||||
start = time.time()
|
||||
max_wait_s = 10
|
||||
min_events_per_rank = 3 # Expect at least a few events from each rank
|
||||
|
||||
while (time.time() - start) < max_wait_s and (
|
||||
len(events_dp0) < min_events_per_rank
|
||||
or len(events_dp1) < min_events_per_rank
|
||||
):
|
||||
# Check DP rank 0
|
||||
if sub_dp0.poll(timeout=100): # 100ms timeout
|
||||
_, seq_bytes, payload = sub_dp0.recv_multipart()
|
||||
event_batch = decoder.decode(payload)
|
||||
print(
|
||||
f"DP Rank 0 - EventBatch: ts={event_batch.ts}, attn_dp_rank={event_batch.attn_dp_rank}"
|
||||
)
|
||||
self.assertEqual(
|
||||
event_batch.attn_dp_rank,
|
||||
0,
|
||||
"DP rank 0 events should have attn_dp_rank=0",
|
||||
)
|
||||
for event in event_batch.events:
|
||||
print(f" DP0 - {event}")
|
||||
events_dp0.append(event)
|
||||
|
||||
# Check DP rank 1
|
||||
if sub_dp1.poll(timeout=100): # 100ms timeout
|
||||
_, seq_bytes, payload = sub_dp1.recv_multipart()
|
||||
event_batch = decoder.decode(payload)
|
||||
print(
|
||||
f"DP Rank 1 - EventBatch: ts={event_batch.ts}, attn_dp_rank={event_batch.attn_dp_rank}"
|
||||
)
|
||||
self.assertEqual(
|
||||
event_batch.attn_dp_rank,
|
||||
1,
|
||||
"DP rank 1 events should have attn_dp_rank=1",
|
||||
)
|
||||
for event in event_batch.events:
|
||||
print(f" DP1 - {event}")
|
||||
events_dp1.append(event)
|
||||
|
||||
# Verify we got events from both DP ranks
|
||||
print(f"Collected {len(events_dp0)} events from DP rank 0")
|
||||
print(f"Collected {len(events_dp1)} events from DP rank 1")
|
||||
|
||||
self.assertGreaterEqual(
|
||||
len(events_dp0),
|
||||
min_events_per_rank,
|
||||
f"Expected at least {min_events_per_rank} events from DP rank 0",
|
||||
)
|
||||
self.assertGreaterEqual(
|
||||
len(events_dp1),
|
||||
min_events_per_rank,
|
||||
f"Expected at least {min_events_per_rank} events from DP rank 1",
|
||||
)
|
||||
|
||||
# Verify event types are as expected
|
||||
for events in [events_dp0, events_dp1]:
|
||||
for event in events:
|
||||
self.assertIsInstance(
|
||||
event,
|
||||
(BlockStored, BlockRemoved, AllBlocksCleared),
|
||||
f"Event should be a KV cache event, got {type(event)}",
|
||||
)
|
||||
|
||||
finally:
|
||||
sub_dp0.close()
|
||||
sub_dp1.close()
|
||||
context.term()
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,526 @@
|
||||
"""
|
||||
Logprobs Accuracy Test for SGLang
|
||||
|
||||
======================
|
||||
With deterministic/batch invariant kernels, we can ensure that SGLang produces exactly the same
|
||||
logprobs results for identical inputs. However, logprobs are highly sensitive to GPU hardware,
|
||||
kernels, torch versions, and other factors, so we cannot maintain a unified logprobs baseline
|
||||
across different machines.
|
||||
|
||||
This test is designed to be run locally by contributors to verify logprobs accuracy
|
||||
before making changes to related code.
|
||||
When submitting changes that affect logprobs computation, please:
|
||||
1. Generate baseline
|
||||
2. Run test
|
||||
3. Submit results
|
||||
|
||||
We really appreciate your effort and contribution to SGLang!
|
||||
|
||||
======================
|
||||
What does this test do?
|
||||
This test fetches 1000 samples from the ShareGPT dataset, generates logprobs for each sample,
|
||||
and saves them as a baseline. Then, by running the test mode, it validates the accuracy of
|
||||
logprobs by comparing them against the baseline.
|
||||
|
||||
This test ensures that:
|
||||
- the boundary of log probs requests are correct, eg, the index for tokens that required log probs are strictly followed
|
||||
- logprobs remain invariant between test runs, and also before and after your code changes;
|
||||
|
||||
======================
|
||||
Usage
|
||||
|
||||
Step 1: Generate Baseline (Before Code Changes)
|
||||
```bash
|
||||
python test/srt/test_logprobs.py gen
|
||||
```
|
||||
|
||||
Step 2: Test Against Baseline (After Code Changes)
|
||||
```bash
|
||||
python test/srt/test_logprobs.py test
|
||||
```
|
||||
This tests your changes against the locally generated baseline from Step 1.
|
||||
The test passes if the maximum and mean differences are within the tolerance thresholds.
|
||||
======================
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import pickle
|
||||
import random
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
import torch
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
# Configuration
|
||||
DENSE_MODEL_NAME = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
SHAREGPT_URL = (
|
||||
"https://huggingface.co/datasets/anon8231489123/"
|
||||
"ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json"
|
||||
)
|
||||
|
||||
# Hardware-specific configuration
|
||||
if torch.version.cuda is not None:
|
||||
print("Running on NVIDIA CUDA GPU")
|
||||
DENSE_TOLERANCE_MAX_DIFF = 1e-5
|
||||
DENSE_TOLERANCE_MEAN_DIFF = 1e-5
|
||||
else:
|
||||
print("No GPU backend (CPU only)")
|
||||
raise ValueError("No GPU backend (CPU only)")
|
||||
|
||||
# Common configuration
|
||||
TOP_K = 20
|
||||
NUM_SAMPLES = 1000
|
||||
LOGPROB_SAMPLE_RATIO = 0.5
|
||||
TEMPERATURE = 1.0
|
||||
MAX_LEN = 20000
|
||||
|
||||
# Default output files
|
||||
DEFAULT_BASELINE_PKL = "sglang_baseline_local.pkl"
|
||||
DEFAULT_META_JSON = "baseline_meta_preview.json"
|
||||
|
||||
# Default engine configuration
|
||||
DEFAULT_ENGINE_CONFIG = {
|
||||
"model_path": DENSE_MODEL_NAME,
|
||||
"random_seed": 42,
|
||||
"skip_tokenizer_init": True,
|
||||
"mem_fraction_static": 0.8,
|
||||
"enable_deterministic_inference": True,
|
||||
"attention_backend": "flashinfer",
|
||||
}
|
||||
|
||||
|
||||
def generate_baseline(
|
||||
baseline_file=DEFAULT_BASELINE_PKL,
|
||||
meta_file=DEFAULT_META_JSON,
|
||||
num_samples=NUM_SAMPLES,
|
||||
):
|
||||
"""Generate a local baseline for logprobs testing.
|
||||
|
||||
Args:
|
||||
baseline_file: Path to save the baseline pickle file
|
||||
meta_file: Path to save the metadata preview JSON file
|
||||
num_samples: Number of samples to generate
|
||||
"""
|
||||
print(f"SGLang version: {sgl.__version__}")
|
||||
print("Downloading ShareGPT dataset...")
|
||||
|
||||
# Download ShareGPT dataset
|
||||
try:
|
||||
response = requests.get(SHAREGPT_URL, timeout=30)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
print(f"Dataset size: {len(data)}")
|
||||
except requests.exceptions.RequestException as e:
|
||||
raise Exception(f"Failed to download ShareGPT dataset: {e}") from e
|
||||
|
||||
# Filter and prepare texts
|
||||
texts = []
|
||||
for s in data:
|
||||
if "conversations" in s and len(s["conversations"]) > 0:
|
||||
try:
|
||||
text = s["conversations"][0]["value"]
|
||||
if isinstance(text, str) and len(text) <= MAX_LEN and len(text) >= 5500:
|
||||
texts.append(text)
|
||||
if len(texts) >= num_samples * 40: # Get more samples for filtering
|
||||
break
|
||||
except (KeyError, IndexError, TypeError) as e:
|
||||
print(f"Warning: Skipping invalid conversation data: {e}")
|
||||
continue
|
||||
|
||||
if not texts:
|
||||
raise ValueError("No valid texts found in the dataset")
|
||||
|
||||
print(f"Loading tokenizer for {DENSE_MODEL_NAME}...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(DENSE_MODEL_NAME, use_fast=True)
|
||||
|
||||
rng = np.random.default_rng(42)
|
||||
|
||||
print(f"Launching SGLang Engine with {DENSE_MODEL_NAME}...")
|
||||
engine = sgl.Engine(
|
||||
model_path=DENSE_MODEL_NAME,
|
||||
attention_backend="flashinfer",
|
||||
enable_deterministic_inference=True,
|
||||
random_seed=42,
|
||||
skip_tokenizer_init=True,
|
||||
mem_fraction_static=0.8,
|
||||
max_running_requests=1,
|
||||
)
|
||||
|
||||
records = []
|
||||
prompt_lengths = []
|
||||
|
||||
try:
|
||||
for i, text in enumerate(texts):
|
||||
if len(records) >= num_samples:
|
||||
break
|
||||
|
||||
try:
|
||||
ids = tokenizer.encode(text, add_special_tokens=False)
|
||||
if len(ids) < 5:
|
||||
continue
|
||||
|
||||
start_pos = int(rng.integers(0, max(1, len(ids) - 3)))
|
||||
|
||||
outputs = engine.generate(
|
||||
input_ids=[ids],
|
||||
sampling_params={
|
||||
"temperature": 1.0,
|
||||
"top_p": 1.0,
|
||||
"top_k": TOP_K,
|
||||
"max_new_tokens": 1,
|
||||
},
|
||||
return_logprob=True,
|
||||
logprob_start_len=start_pos,
|
||||
top_logprobs_num=TOP_K,
|
||||
)
|
||||
meta = outputs[0]["meta_info"]
|
||||
|
||||
records.append(
|
||||
dict(id=i, text=text, ids=ids, start_pos=start_pos, meta=meta)
|
||||
)
|
||||
prompt_lengths.append(len(ids))
|
||||
|
||||
if (i + 1) % 50 == 0:
|
||||
print(f"Processed {len(records)}/{num_samples} samples")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Warning: Failed to process sample {i}: {e}")
|
||||
continue
|
||||
|
||||
if not records:
|
||||
raise RuntimeError(
|
||||
"Failed to generate any baseline records. Please check the warnings above for errors."
|
||||
)
|
||||
|
||||
# Save baseline files
|
||||
with open(baseline_file, "wb") as f:
|
||||
pickle.dump(records, f)
|
||||
with open(meta_file, "w", encoding="utf-8") as f:
|
||||
json.dump(records[:2], f, ensure_ascii=False, indent=2)
|
||||
|
||||
print(f"✅ Saved {len(records)} samples to {baseline_file}")
|
||||
print(f"✅ Meta preview saved to {meta_file}")
|
||||
|
||||
if prompt_lengths:
|
||||
avg_prompt_length = sum(prompt_lengths) / len(prompt_lengths)
|
||||
print(f"📊 Average prompt length: {avg_prompt_length:.2f} tokens")
|
||||
|
||||
finally:
|
||||
engine.shutdown()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
class TestLogprobsDense(unittest.TestCase):
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Set up the test class - initialize the engine once for all tests."""
|
||||
print(f"Launching SGLang Engine with {DENSE_MODEL_NAME}...")
|
||||
cls.engine = sgl.Engine(**DEFAULT_ENGINE_CONFIG)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""Clean up after all tests - shutdown the engine."""
|
||||
cls.engine.shutdown()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
@classmethod
|
||||
def restart_engine_with_config(cls, **kwargs):
|
||||
"""Create engine with custom configuration"""
|
||||
# Safely shutdown existing engine
|
||||
cls.engine.shutdown()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Set chunk size
|
||||
chunk_size = kwargs.pop("chunk_size", None)
|
||||
if chunk_size is not None:
|
||||
print(f"Setting chunk size to {chunk_size}")
|
||||
os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "True"
|
||||
os.environ["SGLANG_LOGITS_PROCESSER_CHUNK_SIZE"] = str(chunk_size)
|
||||
else:
|
||||
os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "False"
|
||||
|
||||
# Create engine with merged configuration
|
||||
engine_config = {**DEFAULT_ENGINE_CONFIG, **kwargs}
|
||||
cls.engine = sgl.Engine(**engine_config)
|
||||
|
||||
def load_test_data(self, baseline_file=None):
|
||||
"""Load test data from local baseline file. In test mode, only local baseline is supported."""
|
||||
if not baseline_file:
|
||||
raise ValueError("baseline_file is required in test mode")
|
||||
|
||||
if not os.path.exists(baseline_file):
|
||||
raise FileNotFoundError(
|
||||
f"Baseline file not found: {baseline_file}. Please run 'gen' mode first to generate the baseline."
|
||||
)
|
||||
|
||||
print(f"Loading local baseline from {baseline_file}...")
|
||||
try:
|
||||
with open(baseline_file, "rb") as f:
|
||||
records = pickle.load(f)
|
||||
print(f"Successfully loaded {len(records)} records from local baseline")
|
||||
return records
|
||||
except (IOError, pickle.PickleError) as e:
|
||||
raise Exception(f"Failed to load local baseline: {e}") from e
|
||||
|
||||
def compare_meta(self, baseline_meta, sglang_meta):
|
||||
"""Compare metadata between two outputs and return max and mean differences."""
|
||||
diffs = []
|
||||
for key in ["input_top_logprobs", "output_top_logprobs"]:
|
||||
baseline_logprobs, sglang_logprobs = baseline_meta[key], sglang_meta[key]
|
||||
self.assertEqual(
|
||||
len(baseline_logprobs),
|
||||
len(sglang_logprobs),
|
||||
f"Length of {key} is not equal, sglang did not return the correct number of log probs(should be top 20)",
|
||||
)
|
||||
for baseline_entry, sglang_entry in zip(baseline_logprobs, sglang_logprobs):
|
||||
if not baseline_entry or not sglang_entry:
|
||||
continue
|
||||
baseline_token_map = {tid: lp for lp, tid, _ in baseline_entry}
|
||||
sglang_token_map = {tid: lp for lp, tid, _ in sglang_entry}
|
||||
common_tokens = baseline_token_map.keys() & sglang_token_map.keys()
|
||||
self.assertGreaterEqual(
|
||||
len(common_tokens),
|
||||
TOP_K,
|
||||
f"there are only {len(common_tokens)} common topk tokens that matches",
|
||||
)
|
||||
for token_id in common_tokens:
|
||||
diffs.append(
|
||||
abs(baseline_token_map[token_id] - sglang_token_map[token_id])
|
||||
)
|
||||
if not diffs:
|
||||
return 0.0, 0.0
|
||||
return max(diffs), float(np.mean(diffs))
|
||||
|
||||
def test_logprobs_comparison(self, baseline_file=None):
|
||||
"""Test the logprobs comparison functionality with different parameter combinations."""
|
||||
# Load test data with retry mechanism
|
||||
records = self.load_test_data(baseline_file)
|
||||
|
||||
# Fast configs for CI
|
||||
test_configs = [
|
||||
{"num_samples": NUM_SAMPLES},
|
||||
{"num_samples": 42, "chunk_size": 1, "max_running_requests": 16},
|
||||
{"num_samples": 42, "chunk_size": 2, "max_running_requests": 16},
|
||||
{"num_samples": 42, "chunk_size": 3, "max_running_requests": 16},
|
||||
{"num_samples": NUM_SAMPLES, "chunk_size": 16, "max_running_requests": 128},
|
||||
{"num_samples": NUM_SAMPLES, "chunk_size": 128, "max_running_requests": 16},
|
||||
{"num_samples": NUM_SAMPLES, "chunk_size": 128, "max_running_requests": 8},
|
||||
{"num_samples": NUM_SAMPLES, "chunk_size": 128, "max_running_requests": 32},
|
||||
{
|
||||
"num_samples": NUM_SAMPLES,
|
||||
"chunk_size": 128,
|
||||
"max_running_requests": 128,
|
||||
},
|
||||
{"num_samples": NUM_SAMPLES, "chunk_size": 256, "max_running_requests": 8},
|
||||
{"num_samples": NUM_SAMPLES, "chunk_size": 256, "max_running_requests": 32},
|
||||
{
|
||||
"num_samples": NUM_SAMPLES,
|
||||
"chunk_size": 256,
|
||||
"max_running_requests": 128,
|
||||
},
|
||||
]
|
||||
|
||||
# Run tests
|
||||
for config in test_configs:
|
||||
with self.subTest(config=config):
|
||||
print(f"Testing with config: {config}")
|
||||
|
||||
# Sample records for this config
|
||||
test_records = random.sample(records, k=min(NUM_SAMPLES, len(records)))
|
||||
random.shuffle(test_records)
|
||||
|
||||
# Calculate how many samples should return logprobs
|
||||
logprob_count = int(len(test_records) * LOGPROB_SAMPLE_RATIO)
|
||||
print(
|
||||
f"Testing with {len(test_records)} samples, temperature={TEMPERATURE}"
|
||||
)
|
||||
print(
|
||||
f"Will return logprobs for {logprob_count} samples (ratio: {LOGPROB_SAMPLE_RATIO})"
|
||||
)
|
||||
|
||||
all_max, all_mean = [], []
|
||||
logprob_returned_count = 0
|
||||
|
||||
# Process all records at once
|
||||
input_ids = [rec["ids"] for rec in test_records]
|
||||
logprob_start_lens = [rec["start_pos"] for rec in test_records]
|
||||
|
||||
# Determine which samples should return logprobs (randomly selected)
|
||||
logprob_indices = set(
|
||||
random.sample(range(len(test_records)), logprob_count)
|
||||
)
|
||||
return_logprob_array = [
|
||||
sample_idx in logprob_indices
|
||||
for sample_idx in range(len(test_records))
|
||||
]
|
||||
|
||||
# Sampling param per request
|
||||
sampling_params = [
|
||||
{
|
||||
"temperature": TEMPERATURE,
|
||||
"top_p": 1.0,
|
||||
"top_k": TOP_K,
|
||||
"max_new_tokens": 1,
|
||||
}
|
||||
for _ in test_records
|
||||
]
|
||||
|
||||
# Some configs must restart the engine to take effect
|
||||
chunk_size = config.get("chunk_size", None)
|
||||
max_running_requests = config.get("max_running_requests", None)
|
||||
if chunk_size is not None or max_running_requests is not None:
|
||||
self.restart_engine_with_config(
|
||||
chunk_size=chunk_size,
|
||||
max_running_requests=max_running_requests,
|
||||
)
|
||||
|
||||
outputs = self.engine.generate(
|
||||
input_ids=input_ids,
|
||||
sampling_params=sampling_params,
|
||||
return_logprob=return_logprob_array,
|
||||
logprob_start_len=logprob_start_lens,
|
||||
top_logprobs_num=TOP_K,
|
||||
)
|
||||
|
||||
for sample_idx, (rec, output) in enumerate(zip(test_records, outputs)):
|
||||
# Only compare logprobs for samples that should have them
|
||||
if sample_idx in logprob_indices:
|
||||
# Safe access to meta_info and input_top_logprobs
|
||||
meta_info = output.get("meta_info")
|
||||
input_top_logprobs = (
|
||||
meta_info.get("input_top_logprobs") if meta_info else None
|
||||
)
|
||||
|
||||
self.assertIsNotNone(
|
||||
input_top_logprobs,
|
||||
f"return_logprob enabled on this sample, but input_top_logprobs is None (length: {len(input_top_logprobs) if input_top_logprobs is not None else 'N/A'})",
|
||||
)
|
||||
baseline_meta = rec["meta"]
|
||||
sglang_meta = meta_info
|
||||
|
||||
max_diff, mean_diff = self.compare_meta(
|
||||
baseline_meta, sglang_meta
|
||||
)
|
||||
all_max.append(max_diff)
|
||||
all_mean.append(mean_diff)
|
||||
logprob_returned_count += 1
|
||||
else:
|
||||
# Verify that logprobs were not returned for this sample
|
||||
meta_info = output.get("meta_info")
|
||||
input_top_logprobs = (
|
||||
meta_info.get("input_top_logprobs") if meta_info else None
|
||||
)
|
||||
output_token_ids_logprobs = (
|
||||
meta_info.get("output_token_ids_logprobs")
|
||||
if meta_info
|
||||
else None
|
||||
)
|
||||
|
||||
self.assertFalse(
|
||||
input_top_logprobs,
|
||||
f"return_logprob is disabled on this sample, Sample {sample_idx} should not have logprobs, content: {output_token_ids_logprobs}",
|
||||
)
|
||||
|
||||
max_of_max = max(all_max) if all_max else 0.0
|
||||
mean_of_mean = np.mean(all_mean) if all_mean else 0.0
|
||||
|
||||
print(f"max Δ={max_of_max:.6g}")
|
||||
print(f"mean Δ={mean_of_mean:.6g}")
|
||||
print(
|
||||
f"logprobs returned for {logprob_returned_count} samples (expected: {logprob_count})"
|
||||
)
|
||||
|
||||
# Verify correct number of logprobs returned
|
||||
self.assertEqual(
|
||||
logprob_returned_count,
|
||||
logprob_count,
|
||||
f"Expected {logprob_count} samples with logprobs, got {logprob_returned_count}",
|
||||
)
|
||||
|
||||
# Basic validation
|
||||
self.assertIsInstance(all_max, list)
|
||||
self.assertIsInstance(all_mean, list)
|
||||
self.assertGreater(
|
||||
len(all_max),
|
||||
0,
|
||||
f"No test samples processed for config {{'num_samples': {NUM_SAMPLES}, 'logprob_sample_ratio': {LOGPROB_SAMPLE_RATIO}, 'temperature': {TEMPERATURE}}}",
|
||||
)
|
||||
|
||||
# Tolerance checks with clear error messages
|
||||
failed_samples = []
|
||||
for sample_idx, (max_diff, mean_diff) in enumerate(
|
||||
zip(all_max, all_mean)
|
||||
):
|
||||
if max_diff > DENSE_TOLERANCE_MAX_DIFF:
|
||||
failed_samples.append(
|
||||
f"Sample {sample_idx}: max_diff={max_diff:.6g} > {DENSE_TOLERANCE_MAX_DIFF}"
|
||||
)
|
||||
if mean_diff > DENSE_TOLERANCE_MEAN_DIFF:
|
||||
failed_samples.append(
|
||||
f"Sample {sample_idx}: mean_diff={mean_diff:.6g} > {DENSE_TOLERANCE_MEAN_DIFF}"
|
||||
)
|
||||
|
||||
if failed_samples:
|
||||
self.fail(
|
||||
f"Config {{'num_samples': {NUM_SAMPLES}, 'logprob_sample_ratio': {LOGPROB_SAMPLE_RATIO}, 'temperature': {TEMPERATURE}}} - Tolerance exceeded in {len(failed_samples)} samples:\n"
|
||||
+ "\n".join(failed_samples[:5])
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to handle command line arguments and run either generation or testing."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="SGLang Logprobs Test and Baseline Generation"
|
||||
)
|
||||
parser.add_argument(
|
||||
"mode",
|
||||
choices=["gen", "test"],
|
||||
help="Mode to run: 'gen' to generate baseline, 'test' to run tests",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.mode == "gen":
|
||||
print("🚀 Generating baseline...")
|
||||
generate_baseline()
|
||||
print(f"\n✅ Baseline generation complete!")
|
||||
print(f"📁 Baseline saved to: {DEFAULT_BASELINE_PKL}")
|
||||
print(f"📁 Metadata preview saved to: {DEFAULT_META_JSON}")
|
||||
print(f"\n💡 Next steps:")
|
||||
print(f" 1. Make your code changes")
|
||||
print(f" 2. Run: python {__file__} test")
|
||||
|
||||
elif args.mode == "test":
|
||||
print("🧪 Running logprobs test...")
|
||||
if not os.path.exists(DEFAULT_BASELINE_PKL):
|
||||
print(f"❌ Baseline file not found: {DEFAULT_BASELINE_PKL}")
|
||||
print(f"💡 Generate baseline first by running:")
|
||||
print(f" python {__file__} gen")
|
||||
print(f" This will download ShareGPT data and generate a local baseline.")
|
||||
return 1
|
||||
|
||||
# Set environment variable for testing
|
||||
os.environ["RETURN_ORIGINAL_LOGPROB"] = "True"
|
||||
|
||||
# Create test instance and run
|
||||
test_instance = TestLogprobsDense()
|
||||
test_instance.setUpClass()
|
||||
try:
|
||||
test_instance.test_logprobs_comparison(baseline_file=DEFAULT_BASELINE_PKL)
|
||||
print("\n✅ Test completed successfully!")
|
||||
finally:
|
||||
test_instance.tearDownClass()
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
exit(main())
|
||||
@@ -0,0 +1,66 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestDeepseekTP2(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "lmsys/sglang-ci-dsv3-test"
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = ["--trust-remote-code"]
|
||||
if torch.cuda.is_available() and torch.version.cuda:
|
||||
other_args.extend(
|
||||
["--tp", "2", "--enable-torch-compile", "--cuda-graph-max-bs", "2"]
|
||||
)
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
self.assertGreater(metrics["accuracy"], 0.62)
|
||||
|
||||
def test_gsm8k_bs1(self):
|
||||
# test torch compile accuracy for bs=1
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=10,
|
||||
max_new_tokens=512,
|
||||
parallel=1,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
self.assertGreater(metrics["accuracy"], 0.62)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,58 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_MODELOPT_QUANT_ACCURACY_TEST_FP8,
|
||||
DEFAULT_MODEL_NAME_FOR_MODELOPT_QUANT_ACCURACY_TEST_FP8_REVISION,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestEvalFP8ModelOptQuantAccuracy(CustomTestCase):
|
||||
|
||||
def _run_test(self, model, other_args, expected_score):
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = other_args or []
|
||||
|
||||
process = popen_launch_server(
|
||||
model,
|
||||
base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
)
|
||||
|
||||
try:
|
||||
args = SimpleNamespace(
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
temperature=0.1,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreaterEqual(metrics["score"], expected_score)
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
@unittest.skipIf(
|
||||
torch.version.hip is not None, "modelopt quantization unsupported on ROCm"
|
||||
)
|
||||
def test_mmlu_offline_only(self):
|
||||
"""Test with offline quantization only."""
|
||||
self._run_test(
|
||||
model=DEFAULT_MODEL_NAME_FOR_MODELOPT_QUANT_ACCURACY_TEST_FP8,
|
||||
other_args=[
|
||||
"--revision",
|
||||
DEFAULT_MODEL_NAME_FOR_MODELOPT_QUANT_ACCURACY_TEST_FP8_REVISION,
|
||||
],
|
||||
expected_score=0.64,
|
||||
)
|
||||
@@ -0,0 +1,30 @@
|
||||
import unittest
|
||||
|
||||
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
|
||||
|
||||
from sglang.srt.layers.quantization.modelopt_quant import (
|
||||
ModelOptFp8Config,
|
||||
ModelOptFp8KVCacheMethod,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestModelOptFp8KVCacheMethod(CustomTestCase):
|
||||
def test_kv_cache_method_initialization(self):
|
||||
"""Test that ModelOptFp8KVCacheMethod can be instantiated and
|
||||
inherits from BaseKVCacheMethod."""
|
||||
# Create a ModelOptFp8Config object
|
||||
quant_config = ModelOptFp8Config(is_checkpoint_fp8_serialized=True)
|
||||
|
||||
# Instantiate the KV cache method
|
||||
kv_cache_method = ModelOptFp8KVCacheMethod(quant_config)
|
||||
|
||||
# Check inheritance
|
||||
self.assertIsInstance(kv_cache_method, BaseKVCacheMethod)
|
||||
|
||||
# Check that the quant_config is stored
|
||||
self.assertEqual(kv_cache_method.quant_config, quant_config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,40 @@
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
from sglang.srt.utils import prepare_model_and_tokenizer
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestDownloadFromModelScope(CustomTestCase):
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "iic/nlp_lstmcrf_word-segmentation_chinese-news"
|
||||
stat, output = subprocess.getstatusoutput("pip install modelscope")
|
||||
|
||||
cls.with_modelscope_environ = {k: v for k, v in os.environ.items()}
|
||||
cls.with_modelscope_environ["SGLANG_USE_MODELSCOPE"] = "True"
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
pass
|
||||
|
||||
def test_prepare_model_and_tokenizer(self):
|
||||
from modelscope.utils.file_utils import get_model_cache_root
|
||||
|
||||
model_cache_root = get_model_cache_root()
|
||||
if os.path.exists(model_cache_root):
|
||||
shutil.rmtree(model_cache_root)
|
||||
with mock.patch.dict(os.environ, self.with_modelscope_environ, clear=True):
|
||||
model_path, tokenizer_path = prepare_model_and_tokenizer(
|
||||
self.model, self.model
|
||||
)
|
||||
assert os.path.exists(os.path.join(model_path, "pytorch_model.bin"))
|
||||
assert os.path.exists(os.path.join(tokenizer_path, "config.json"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,196 @@
|
||||
"""For Now, MSCCL is only supported on TP16 and TP8 case
|
||||
|
||||
if [[ $RANK -eq 0 ]]; then
|
||||
ray start --block --head --port=6379 &
|
||||
python3 test_mscclpp.py;
|
||||
else
|
||||
ray start --block --address=${MASTER_ADDR}:6379;
|
||||
fi
|
||||
"""
|
||||
|
||||
import os
|
||||
import random
|
||||
import socket
|
||||
import unittest
|
||||
from typing import Any
|
||||
|
||||
import ray
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.srt.distributed import init_distributed_environment
|
||||
from sglang.srt.distributed.communication_op import ( # noqa
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from sglang.srt.distributed.parallel_state import (
|
||||
get_tensor_model_parallel_group,
|
||||
graph_capture,
|
||||
initialize_model_parallel,
|
||||
set_custom_all_reduce,
|
||||
set_mscclpp_all_reduce,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
def get_open_port() -> int:
|
||||
# try ipv4
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
except OSError:
|
||||
# try ipv6
|
||||
with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def multi_process_parallel(
|
||||
world_size: int,
|
||||
master_addr: str,
|
||||
cls: Any,
|
||||
test_target: Any,
|
||||
) -> None:
|
||||
|
||||
# Using ray helps debugging the error when it failed
|
||||
# as compared to multiprocessing.
|
||||
# NOTE: We need to set working_dir for distributed tests,
|
||||
# otherwise we may get import errors on ray workers
|
||||
|
||||
ray.init(log_to_driver=True)
|
||||
|
||||
distributed_init_port = get_open_port()
|
||||
refs = []
|
||||
for rank in range(world_size):
|
||||
refs.append(
|
||||
test_target.remote(
|
||||
cls, world_size, master_addr, rank, distributed_init_port
|
||||
)
|
||||
)
|
||||
ray.get(refs)
|
||||
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
class TestMSCCLAllReduce(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
random.seed(42)
|
||||
# 1KB to 1MB
|
||||
cls.test_sizes = [512, 4096, 32768, 262144, 524288]
|
||||
cls.world_sizes = [8]
|
||||
TEST_TP16 = int(os.getenv("SGL_MSCCLPP_TEST_TP16", "0"))
|
||||
if TEST_TP16:
|
||||
cls.world_sizes = [16]
|
||||
cls.test_loop = 10
|
||||
|
||||
def test_graph_allreduce(self):
|
||||
TEST_MASTER_ADDR = os.getenv("SGL_MSCCLPP_TEST_MASTER_ADDR", "localhost")
|
||||
for world_size in self.world_sizes:
|
||||
if world_size not in [8, 16]:
|
||||
continue
|
||||
multi_process_parallel(
|
||||
world_size, TEST_MASTER_ADDR, self, self.graph_allreduce
|
||||
)
|
||||
|
||||
def test_eager_allreduce(self):
|
||||
TEST_MASTER_ADDR = os.getenv("SGL_MSCCLPP_TEST_MASTER_ADDR", "localhost")
|
||||
for world_size in self.world_sizes:
|
||||
if world_size not in [8, 16]:
|
||||
continue
|
||||
multi_process_parallel(
|
||||
world_size, TEST_MASTER_ADDR, self, self.eager_allreduce
|
||||
)
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def graph_allreduce(self, world_size, master_addr, rank, distributed_init_port):
|
||||
del os.environ["CUDA_VISIBLE_DEVICES"]
|
||||
device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
|
||||
torch.cuda.set_device(device)
|
||||
distributed_init_method = f"tcp://{master_addr}:{distributed_init_port}"
|
||||
set_mscclpp_all_reduce(True)
|
||||
set_custom_all_reduce(False)
|
||||
init_distributed_environment(
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
distributed_init_method=distributed_init_method,
|
||||
local_rank=rank % torch.cuda.device_count(),
|
||||
)
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
group = get_tensor_model_parallel_group().device_group
|
||||
|
||||
# A small all_reduce for warmup.
|
||||
# this is needed because device communicators might be created lazily
|
||||
# (e.g. NCCL). This will ensure that the communicator is initialized
|
||||
# before any communication happens, so that this group can be used for
|
||||
# graph capture immediately.
|
||||
data = torch.zeros(1)
|
||||
data = data.to(device=device)
|
||||
torch.distributed.all_reduce(data, group=group)
|
||||
torch.cuda.synchronize()
|
||||
del data
|
||||
|
||||
for sz in self.test_sizes:
|
||||
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
||||
for _ in range(self.test_loop):
|
||||
with graph_capture() as graph_capture_context:
|
||||
# use integers so result matches NCCL exactly
|
||||
inp1 = torch.randint(
|
||||
1,
|
||||
16,
|
||||
(sz,),
|
||||
dtype=dtype,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
inp2 = torch.randint(
|
||||
1,
|
||||
16,
|
||||
(sz,),
|
||||
dtype=dtype,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(
|
||||
graph, stream=graph_capture_context.stream
|
||||
):
|
||||
out1 = tensor_model_parallel_all_reduce(inp1)
|
||||
# the input buffer is immediately modified to test
|
||||
# synchronization
|
||||
dist.all_reduce(inp1, group=group)
|
||||
out2 = tensor_model_parallel_all_reduce(inp2)
|
||||
dist.all_reduce(inp2, group=group)
|
||||
graph.replay()
|
||||
torch.testing.assert_close(out1, inp1)
|
||||
torch.testing.assert_close(out2, inp2)
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def eager_allreduce(self, world_size, master_addr, rank, distributed_init_port):
|
||||
del os.environ["CUDA_VISIBLE_DEVICES"]
|
||||
device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
|
||||
torch.cuda.set_device(device)
|
||||
distributed_init_method = f"tcp://{master_addr}:{distributed_init_port}"
|
||||
set_mscclpp_all_reduce(True)
|
||||
set_custom_all_reduce(False)
|
||||
init_distributed_environment(
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
distributed_init_method=distributed_init_method,
|
||||
local_rank=rank,
|
||||
)
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
group = get_tensor_model_parallel_group().device_group
|
||||
|
||||
for sz in self.test_sizes:
|
||||
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
||||
for _ in range(self.test_loop):
|
||||
inp1 = torch.randint(
|
||||
1, 16, (sz,), dtype=dtype, device=torch.cuda.current_device()
|
||||
)
|
||||
out1 = tensor_model_parallel_all_reduce(inp1)
|
||||
dist.all_reduce(inp1, group=group)
|
||||
torch.testing.assert_close(out1, inp1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,303 @@
|
||||
import multiprocessing
|
||||
import os
|
||||
import random
|
||||
import socket
|
||||
import unittest
|
||||
from typing import Any
|
||||
|
||||
import ray
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.srt import _custom_ops as ops
|
||||
from sglang.srt.distributed import init_distributed_environment
|
||||
from sglang.srt.distributed.communication_op import ( # noqa
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from sglang.srt.distributed.device_communicators.quick_all_reduce import (
|
||||
qr_rocm_arch_available,
|
||||
)
|
||||
from sglang.srt.distributed.parallel_state import (
|
||||
get_tensor_model_parallel_group,
|
||||
graph_capture,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
torch.manual_seed(42)
|
||||
random.seed(44) # keep the deterministic seed
|
||||
|
||||
|
||||
def get_open_port() -> int:
|
||||
# try ipv4
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
except OSError:
|
||||
# try ipv6
|
||||
with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def multi_process_parallel(
|
||||
world_size: int, cls: Any, test_target: Any, quant_mode: str
|
||||
) -> None:
|
||||
|
||||
# Using ray helps debugging the error when it failed
|
||||
# as compared to multiprocessing.
|
||||
# NOTE: We need to set working_dir for distributed tests,
|
||||
# otherwise we may get import errors on ray workers
|
||||
|
||||
ray.init(log_to_driver=True)
|
||||
|
||||
distributed_init_port = get_open_port()
|
||||
refs = []
|
||||
for rank in range(world_size):
|
||||
refs.append(
|
||||
test_target.remote(cls, world_size, rank, distributed_init_port, quant_mode)
|
||||
)
|
||||
ray.get(refs)
|
||||
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
class TestQuickAllReduce(CustomTestCase):
|
||||
TEST_SIZES = [
|
||||
2 * 1024 * 1024,
|
||||
4 * 1024 * 1024,
|
||||
8 * 1024 * 1024,
|
||||
16 * 1024 * 1024,
|
||||
32 * 1024 * 1024,
|
||||
]
|
||||
TEST_LOOP = 5
|
||||
# Too many configurations can lead to a test grid that is too large
|
||||
# The tp takes too long to boot,let's just choose 4 out of 12 configurations
|
||||
# WORLD_SIZES = [2, 4, 8]
|
||||
# QUANT_MODE = ["FP", "INT8", "INT6", "INT4"]
|
||||
QUANT_MODE_WORLD_SIZE_PART = [["FP", 8], ["INT4", 4], ["INT8", 2], ["INT6", 2]]
|
||||
|
||||
@unittest.skipIf(
|
||||
not qr_rocm_arch_available(),
|
||||
"Only test Quick AllReduce on ROCm architectures >= gfx94*",
|
||||
)
|
||||
def test_graph_allreduce(self):
|
||||
for quant_mode_world_size_part in self.QUANT_MODE_WORLD_SIZE_PART:
|
||||
quant_mode = quant_mode_world_size_part[0]
|
||||
world_size = quant_mode_world_size_part[1]
|
||||
if world_size > torch.cuda.device_count():
|
||||
continue
|
||||
multi_process_parallel(world_size, self, self.graph_allreduce, quant_mode)
|
||||
|
||||
@unittest.skipIf(
|
||||
not qr_rocm_arch_available(),
|
||||
"Only test Quick AllReduce on ROCm architectures >= gfx94*",
|
||||
)
|
||||
def test_eager_allreduce(self):
|
||||
for quant_mode_world_size_part in self.QUANT_MODE_WORLD_SIZE_PART:
|
||||
quant_mode = quant_mode_world_size_part[0]
|
||||
world_size = quant_mode_world_size_part[1]
|
||||
if world_size > torch.cuda.device_count():
|
||||
continue
|
||||
multi_process_parallel(world_size, self, self.eager_allreduce, quant_mode)
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def graph_allreduce(self, world_size, rank, distributed_init_port, quant_mode):
|
||||
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
|
||||
os.environ["ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_mode
|
||||
os.environ["ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "0"
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
distributed_init_method = f"tcp://localhost:{distributed_init_port}"
|
||||
init_distributed_environment(
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
distributed_init_method=distributed_init_method,
|
||||
local_rank=rank,
|
||||
)
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
group = get_tensor_model_parallel_group().device_group
|
||||
|
||||
# A small all_reduce for warmup.
|
||||
# this is needed because device communicators might be created lazily
|
||||
# (e.g. NCCL). This will ensure that the communicator is initialized
|
||||
# before any communication happens, so that this group can be used for
|
||||
# graph capture immediately.
|
||||
data = torch.zeros(1)
|
||||
data = data.to(device=device)
|
||||
torch.distributed.all_reduce(data, group=group)
|
||||
torch.cuda.synchronize()
|
||||
del data
|
||||
|
||||
for sz in self.TEST_SIZES:
|
||||
for dtype in [torch.float16, torch.bfloat16]:
|
||||
for _ in range(self.TEST_LOOP):
|
||||
with graph_capture() as graph_capture_context:
|
||||
# use integers so result matches NCCL exactly
|
||||
inp1 = torch.randint(
|
||||
1,
|
||||
23,
|
||||
(sz,),
|
||||
dtype=dtype,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
inp2 = torch.randint(
|
||||
-23,
|
||||
1,
|
||||
(sz,),
|
||||
dtype=dtype,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(
|
||||
graph, stream=graph_capture_context.stream
|
||||
):
|
||||
out1 = tensor_model_parallel_all_reduce(inp1)
|
||||
# the input buffer is immediately modified to test
|
||||
# synchronization
|
||||
dist.all_reduce(inp1, group=group)
|
||||
out2 = tensor_model_parallel_all_reduce(inp2)
|
||||
dist.all_reduce(inp2, group=group)
|
||||
graph.replay()
|
||||
atol = 1.25 * world_size
|
||||
rtol = 0.5 * world_size
|
||||
for inp, out in [[inp1, out1], [inp2, out2]]:
|
||||
torch.testing.assert_close(out, inp, atol=atol, rtol=rtol)
|
||||
# try:
|
||||
# torch.testing.assert_close(out, inp, atol=atol, rtol=rtol)
|
||||
# except AssertionError as e:
|
||||
# print("Max abs diff:", (out - inp).abs().max())
|
||||
# print("Max rel diff:", ((out - inp).abs() / inp.abs().clamp(min=1e-5)).max())
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def eager_allreduce(self, world_size, rank, distributed_init_port, quant_mode):
|
||||
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
|
||||
os.environ["ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_mode
|
||||
os.environ["ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "0"
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
distributed_init_method = f"tcp://localhost:{distributed_init_port}"
|
||||
init_distributed_environment(
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
distributed_init_method=distributed_init_method,
|
||||
local_rank=rank,
|
||||
)
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
group = get_tensor_model_parallel_group().device_group
|
||||
|
||||
for sz in self.TEST_SIZES:
|
||||
for dtype in [torch.float16, torch.bfloat16]:
|
||||
for _ in range(self.TEST_LOOP):
|
||||
inp1 = torch.randint(
|
||||
1,
|
||||
23,
|
||||
(sz,),
|
||||
dtype=dtype,
|
||||
device=torch.cuda.current_device(),
|
||||
)
|
||||
out1 = tensor_model_parallel_all_reduce(inp1)
|
||||
dist.all_reduce(inp1, group=group)
|
||||
atol = 1.25 * world_size
|
||||
rtol = 0.5 * world_size
|
||||
torch.testing.assert_close(out1, inp1, atol=atol, rtol=rtol)
|
||||
# try:
|
||||
# torch.testing.assert_close(out1, inp1, atol=atol, rtol=rtol)
|
||||
# except AssertionError as e:
|
||||
# print("Max abs diff:", (out1 - inp1).abs().max())
|
||||
# print("Max rel diff:", ((out1 - inp1).abs() / inp1.abs().clamp(min=1e-5)).max())
|
||||
|
||||
|
||||
def qr_variable_input(rank, world_size):
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.cuda.set_device(device)
|
||||
qr_max_size = None # MB
|
||||
_ptr = ops.init_custom_qr(rank, world_size, qr_max_size)
|
||||
ranks = []
|
||||
for i in range(world_size):
|
||||
ranks.append(i)
|
||||
dist.init_process_group(
|
||||
backend="nccl",
|
||||
init_method="tcp://127.0.0.1:29500",
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
)
|
||||
cpu_group = torch.distributed.new_group(ranks, backend="nccl")
|
||||
|
||||
handle = ops.qr_get_handle(_ptr)
|
||||
world_size = dist.get_world_size(group=cpu_group)
|
||||
handles = [None] * world_size
|
||||
dist.all_gather_object(handles, handle, group=cpu_group)
|
||||
ops.qr_open_handles(_ptr, handles)
|
||||
|
||||
num = 1
|
||||
s1 = 1024
|
||||
while num < 50000: # 50000 is sufficient to identify issues.
|
||||
dtype = torch.float16
|
||||
if num % 2 == 0:
|
||||
s2 = 1024
|
||||
inp1 = torch.zeros(
|
||||
(s1, s2), dtype=dtype, device=torch.cuda.current_device()
|
||||
)
|
||||
else:
|
||||
s2 = 2048
|
||||
inp1 = torch.ones((s1, s2), dtype=dtype, device=torch.cuda.current_device())
|
||||
result = torch.empty_like(inp1)
|
||||
# FP = 0 INT8 = 1 INT6 = 2 INT4 = 3 NONE = 4
|
||||
ops.qr_all_reduce(_ptr, inp1, result, 3, cast_bf2half=True)
|
||||
try:
|
||||
if inp1[0, 0] == 0:
|
||||
assert torch.all(result == 0)
|
||||
else:
|
||||
assert torch.all(result == world_size)
|
||||
except AssertionError:
|
||||
print("Assertion failed! Allreduce results are incorrect.")
|
||||
raise
|
||||
num += 1
|
||||
|
||||
|
||||
class TestQuickreduceVariableInput(CustomTestCase):
|
||||
"""
|
||||
When the tensor parallelism is set to 4 or 8, frequent changes
|
||||
in the input shape can cause QuickReduce to hang (this issue
|
||||
has been observed with the gpt_oss model).
|
||||
"""
|
||||
|
||||
TP_SIZES = [4, 8]
|
||||
|
||||
@unittest.skipIf(
|
||||
not qr_rocm_arch_available(),
|
||||
"Only test Quick AllReduce on ROCm architectures >= gfx94*",
|
||||
)
|
||||
def test_custom_quick_allreduce_variable_input(self):
|
||||
for tp_size in self.TP_SIZES:
|
||||
world_size = tp_size
|
||||
if world_size > torch.cuda.device_count():
|
||||
return
|
||||
|
||||
multiprocessing.set_start_method("spawn", force=True)
|
||||
# 90s is enough
|
||||
timeout = 90
|
||||
processes = []
|
||||
for rank in range(tp_size):
|
||||
p = multiprocessing.Process(
|
||||
target=qr_variable_input, args=(rank, tp_size)
|
||||
)
|
||||
p.start()
|
||||
processes.append((rank, p))
|
||||
for rank, p in processes:
|
||||
p.join(timeout=timeout)
|
||||
if p.is_alive():
|
||||
for r, proc in processes:
|
||||
if proc.is_alive():
|
||||
proc.terminate()
|
||||
proc.join()
|
||||
raise RuntimeError(
|
||||
f"QuickReduce hang detected after {timeout} seconds!"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,183 @@
|
||||
"""
|
||||
python3 -m unittest test_sagemaker_server.TestSageMakerServer.test_chat_completion
|
||||
"""
|
||||
|
||||
import json
|
||||
import unittest
|
||||
|
||||
import 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_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestSageMakerServer(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-123456"
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
api_key=cls.api_key,
|
||||
)
|
||||
cls.tokenizer = get_tokenizer(DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def run_chat_completion(self, logprobs, parallel_sample_num):
|
||||
data = {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful AI assistant"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is the capital of France? Answer in a few words.",
|
||||
},
|
||||
],
|
||||
"temperature": 0,
|
||||
"logprobs": logprobs is not None and logprobs > 0,
|
||||
"top_logprobs": logprobs,
|
||||
"n": parallel_sample_num,
|
||||
}
|
||||
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
response = requests.post(
|
||||
f"{self.base_url}/invocations", json=data, headers=headers
|
||||
).json()
|
||||
|
||||
if logprobs:
|
||||
assert isinstance(
|
||||
response["choices"][0]["logprobs"]["content"][0]["top_logprobs"][0][
|
||||
"token"
|
||||
],
|
||||
str,
|
||||
)
|
||||
|
||||
ret_num_top_logprobs = len(
|
||||
response["choices"][0]["logprobs"]["content"][0]["top_logprobs"]
|
||||
)
|
||||
assert (
|
||||
ret_num_top_logprobs == logprobs
|
||||
), f"{ret_num_top_logprobs} vs {logprobs}"
|
||||
|
||||
assert len(response["choices"]) == parallel_sample_num
|
||||
assert response["choices"][0]["message"]["role"] == "assistant"
|
||||
assert isinstance(response["choices"][0]["message"]["content"], str)
|
||||
assert response["id"]
|
||||
assert response["created"]
|
||||
assert response["usage"]["prompt_tokens"] > 0
|
||||
assert response["usage"]["completion_tokens"] > 0
|
||||
assert response["usage"]["total_tokens"] > 0
|
||||
|
||||
def run_chat_completion_stream(self, logprobs, parallel_sample_num=1):
|
||||
data = {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful AI assistant"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is the capital of France? Answer in a few words.",
|
||||
},
|
||||
],
|
||||
"temperature": 0,
|
||||
"logprobs": logprobs is not None and logprobs > 0,
|
||||
"top_logprobs": logprobs,
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
"n": parallel_sample_num,
|
||||
}
|
||||
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
response = requests.post(
|
||||
f"{self.base_url}/invocations", json=data, stream=True, headers=headers
|
||||
)
|
||||
|
||||
is_firsts = {}
|
||||
for line in response.iter_lines():
|
||||
line = line.decode("utf-8").replace("data: ", "")
|
||||
if len(line) < 1 or line == "[DONE]":
|
||||
continue
|
||||
print(f"value: {line}")
|
||||
line = json.loads(line)
|
||||
usage = line.get("usage")
|
||||
if usage is not None:
|
||||
assert usage["prompt_tokens"] > 0
|
||||
assert usage["completion_tokens"] > 0
|
||||
assert usage["total_tokens"] > 0
|
||||
continue
|
||||
|
||||
index = line.get("choices")[0].get("index")
|
||||
data = line.get("choices")[0].get("delta")
|
||||
|
||||
if is_firsts.get(index, True):
|
||||
assert data["role"] == "assistant"
|
||||
is_firsts[index] = False
|
||||
continue
|
||||
|
||||
# Skip chunks that are just empty placeholders, usually at stream end/stop
|
||||
if data.get("content") is None:
|
||||
continue
|
||||
|
||||
if logprobs:
|
||||
assert line.get("choices")[0].get("logprobs")
|
||||
assert isinstance(
|
||||
line.get("choices")[0]
|
||||
.get("logprobs")
|
||||
.get("content")[0]
|
||||
.get("top_logprobs")[0]
|
||||
.get("token"),
|
||||
str,
|
||||
)
|
||||
assert isinstance(
|
||||
line.get("choices")[0]
|
||||
.get("logprobs")
|
||||
.get("content")[0]
|
||||
.get("top_logprobs"),
|
||||
list,
|
||||
)
|
||||
ret_num_top_logprobs = len(
|
||||
line.get("choices")[0]
|
||||
.get("logprobs")
|
||||
.get("content")[0]
|
||||
.get("top_logprobs")
|
||||
)
|
||||
assert (
|
||||
ret_num_top_logprobs == logprobs
|
||||
), f"{ret_num_top_logprobs} vs {logprobs}"
|
||||
|
||||
assert isinstance(data["content"], str)
|
||||
assert line["id"]
|
||||
assert line["created"]
|
||||
|
||||
for index in [i for i in range(parallel_sample_num)]:
|
||||
assert not is_firsts.get(
|
||||
index, True
|
||||
), f"index {index} is not found in the response"
|
||||
|
||||
def test_chat_completion(self):
|
||||
for logprobs in [None, 5]:
|
||||
for parallel_sample_num in [1, 2]:
|
||||
self.run_chat_completion(logprobs, parallel_sample_num)
|
||||
|
||||
def test_chat_completion_stream(self):
|
||||
for logprobs in [None, 5]:
|
||||
for parallel_sample_num in [1, 2]:
|
||||
self.run_chat_completion_stream(logprobs, parallel_sample_num)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,203 @@
|
||||
import unittest
|
||||
|
||||
from sglang.srt.managers.schedule_batch import Req
|
||||
from sglang.srt.managers.schedule_policy import (
|
||||
CacheAgnosticPolicy,
|
||||
CacheAwarePolicy,
|
||||
SchedulePolicy,
|
||||
)
|
||||
from sglang.srt.mem_cache.radix_cache import RadixCache
|
||||
from sglang.srt.sampling.sampling_params import SamplingParams
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestSchedulePolicy(CustomTestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.tree_cache = RadixCache(None, None, False)
|
||||
|
||||
def test_init_with_cache_aware_policy(self):
|
||||
policy = SchedulePolicy(
|
||||
policy="lpm",
|
||||
tree_cache=self.tree_cache,
|
||||
enable_hierarchical_cache=True,
|
||||
enable_priority_scheduling=False,
|
||||
schedule_low_priority_values_first=False,
|
||||
)
|
||||
self.assertEqual(policy.policy, CacheAwarePolicy.LPM)
|
||||
|
||||
def test_init_with_cache_agnostic_policy(self):
|
||||
policy = SchedulePolicy(
|
||||
policy="fcfs",
|
||||
tree_cache=self.tree_cache,
|
||||
enable_hierarchical_cache=True,
|
||||
enable_priority_scheduling=False,
|
||||
schedule_low_priority_values_first=False,
|
||||
)
|
||||
self.assertEqual(policy.policy, CacheAgnosticPolicy.FCFS)
|
||||
|
||||
def test_init_with_unknown_policy(self):
|
||||
with self.assertRaises(ValueError):
|
||||
SchedulePolicy(
|
||||
policy="invalid",
|
||||
tree_cache=self.tree_cache,
|
||||
enable_hierarchical_cache=True,
|
||||
enable_priority_scheduling=False,
|
||||
schedule_low_priority_values_first=False,
|
||||
)
|
||||
|
||||
def test_init_with_disabled_cache(self):
|
||||
disabled_tree_cache = RadixCache(None, None, disable=True, page_size=1)
|
||||
policy = SchedulePolicy(
|
||||
policy="lpm",
|
||||
tree_cache=disabled_tree_cache,
|
||||
enable_hierarchical_cache=True,
|
||||
enable_priority_scheduling=False,
|
||||
schedule_low_priority_values_first=False,
|
||||
)
|
||||
self.assertEqual(policy.policy, CacheAgnosticPolicy.FCFS)
|
||||
|
||||
def test_calc_priority_fcfs(self):
|
||||
tree_cache = RadixCache(None, None, False)
|
||||
waiting_queue = [
|
||||
Req(1, "a b", [1, 2], SamplingParams()),
|
||||
Req(3, "a b c", [1, 2, 3], SamplingParams()),
|
||||
Req(2, "a", [1], SamplingParams()),
|
||||
]
|
||||
|
||||
policy = SchedulePolicy(
|
||||
policy="fcfs",
|
||||
tree_cache=tree_cache,
|
||||
enable_hierarchical_cache=True,
|
||||
enable_priority_scheduling=False,
|
||||
schedule_low_priority_values_first=False,
|
||||
)
|
||||
policy.calc_priority(waiting_queue)
|
||||
# Check if FCFS keeps the original order
|
||||
self.assertEqual(waiting_queue[0].rid, 1)
|
||||
self.assertEqual(waiting_queue[1].rid, 3)
|
||||
self.assertEqual(waiting_queue[2].rid, 2)
|
||||
|
||||
def test_calc_priority_priority_enabled_fcfs_scheduling(self):
|
||||
tree_cache = RadixCache(None, None, False)
|
||||
|
||||
waiting_queue = [
|
||||
Req(1, "a b", [1, 2], SamplingParams()),
|
||||
Req(3, "a b c", [1, 2, 3], SamplingParams()),
|
||||
Req(2, "a", [1], SamplingParams()),
|
||||
]
|
||||
waiting_queue[0].priority, waiting_queue[0].queue_time_start = 1, 1
|
||||
waiting_queue[1].priority, waiting_queue[1].queue_time_start = 0, 1
|
||||
waiting_queue[2].priority, waiting_queue[2].queue_time_start = 0, 0
|
||||
|
||||
policy = SchedulePolicy(
|
||||
policy="fcfs",
|
||||
tree_cache=tree_cache,
|
||||
enable_hierarchical_cache=True,
|
||||
enable_priority_scheduling=True,
|
||||
schedule_low_priority_values_first=False,
|
||||
)
|
||||
policy.calc_priority(waiting_queue)
|
||||
# Check if priority enabled fcfs ordering is applied.
|
||||
self.assertEqual(waiting_queue[0].rid, 1)
|
||||
self.assertEqual(waiting_queue[1].rid, 2)
|
||||
self.assertEqual(waiting_queue[2].rid, 3)
|
||||
|
||||
def test_calc_priority_priority_enabled_fcfs_scheduling_with_low_priority_values_first(
|
||||
self,
|
||||
):
|
||||
tree_cache = RadixCache(None, None, False)
|
||||
|
||||
waiting_queue = [
|
||||
Req(1, "a b", [1, 2], SamplingParams()),
|
||||
Req(3, "a b c", [1, 2, 3], SamplingParams()),
|
||||
Req(2, "a", [1], SamplingParams()),
|
||||
]
|
||||
waiting_queue[0].priority, waiting_queue[0].queue_time_start = -1, 0
|
||||
waiting_queue[1].priority, waiting_queue[1].queue_time_start = 0, 1
|
||||
waiting_queue[2].priority, waiting_queue[2].queue_time_start = 0, 0
|
||||
|
||||
policy = SchedulePolicy(
|
||||
policy="fcfs",
|
||||
tree_cache=tree_cache,
|
||||
enable_hierarchical_cache=True,
|
||||
enable_priority_scheduling=True,
|
||||
schedule_low_priority_values_first=True,
|
||||
)
|
||||
policy.calc_priority(waiting_queue)
|
||||
# Check if priority enabled fcfs ordering is applied.
|
||||
self.assertEqual(waiting_queue[0].rid, 1)
|
||||
self.assertEqual(waiting_queue[1].rid, 2)
|
||||
self.assertEqual(waiting_queue[2].rid, 3)
|
||||
|
||||
def test_calc_priority_longest_output_first_scheduling(self):
|
||||
tree_cache = RadixCache(None, None, False)
|
||||
|
||||
waiting_queue = [
|
||||
Req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1000)),
|
||||
Req(3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10)),
|
||||
Req(2, "a", [1], SamplingParams(max_new_tokens=100)),
|
||||
]
|
||||
|
||||
policy = SchedulePolicy(
|
||||
policy="lof",
|
||||
tree_cache=tree_cache,
|
||||
enable_hierarchical_cache=True,
|
||||
enable_priority_scheduling=False,
|
||||
schedule_low_priority_values_first=False,
|
||||
)
|
||||
policy.calc_priority(waiting_queue)
|
||||
# Check if priority enabled fcfs ordering is applied.
|
||||
self.assertEqual(waiting_queue[0].rid, 1)
|
||||
self.assertEqual(waiting_queue[1].rid, 2)
|
||||
self.assertEqual(waiting_queue[2].rid, 3)
|
||||
|
||||
def test_calc_priority_priority_enabled_longest_output_first_scheduling(self):
|
||||
tree_cache = RadixCache(None, None, False)
|
||||
|
||||
waiting_queue = [
|
||||
Req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1), priority=1),
|
||||
Req(3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10), priority=0),
|
||||
Req(2, "a", [1], SamplingParams(max_new_tokens=100), priority=0),
|
||||
]
|
||||
|
||||
policy = SchedulePolicy(
|
||||
policy="lof",
|
||||
tree_cache=tree_cache,
|
||||
enable_hierarchical_cache=True,
|
||||
enable_priority_scheduling=True,
|
||||
schedule_low_priority_values_first=False,
|
||||
)
|
||||
policy.calc_priority(waiting_queue)
|
||||
# Check if priority enabled fcfs ordering is applied.
|
||||
self.assertEqual(waiting_queue[0].rid, 1)
|
||||
self.assertEqual(waiting_queue[1].rid, 2)
|
||||
self.assertEqual(waiting_queue[2].rid, 3)
|
||||
|
||||
def test_calc_priority_priority_enabled_longest_output_first_scheduling_with_low_priority_values_first(
|
||||
self,
|
||||
):
|
||||
tree_cache = RadixCache(None, None, False)
|
||||
|
||||
waiting_queue = [
|
||||
Req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1), priority=0),
|
||||
Req(3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10), priority=1),
|
||||
Req(2, "a", [1], SamplingParams(max_new_tokens=100), priority=1),
|
||||
]
|
||||
|
||||
policy = SchedulePolicy(
|
||||
policy="lof",
|
||||
tree_cache=tree_cache,
|
||||
enable_hierarchical_cache=True,
|
||||
enable_priority_scheduling=True,
|
||||
schedule_low_priority_values_first=True,
|
||||
)
|
||||
policy.calc_priority(waiting_queue)
|
||||
# Check if priority enabled fcfs ordering is applied.
|
||||
self.assertEqual(waiting_queue[0].rid, 1)
|
||||
self.assertEqual(waiting_queue[1].rid, 2)
|
||||
self.assertEqual(waiting_queue[2].rid, 3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,786 @@
|
||||
"""
|
||||
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.TestSessionControlVision.test_session_control
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
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.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
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 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()
|
||||
@@ -0,0 +1,60 @@
|
||||
import unittest
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
|
||||
|
||||
|
||||
class TestSRTEngineWithQuantArgs(CustomTestCase):
|
||||
|
||||
def test_1_quantization_args(self):
|
||||
|
||||
# we only test fp8 because other methods are currently dependent on vllm. We can add other methods back to test after vllm dependency is resolved.
|
||||
quantization_args_list = [
|
||||
# "awq",
|
||||
"fp8",
|
||||
# "gptq",
|
||||
# "marlin",
|
||||
# "gptq_marlin",
|
||||
# "awq_marlin",
|
||||
# "bitsandbytes",
|
||||
# "gguf",
|
||||
]
|
||||
|
||||
prompt = "Today is a sunny day and I like"
|
||||
model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
sampling_params = {"temperature": 0, "max_new_tokens": 8}
|
||||
|
||||
for quantization_args in quantization_args_list:
|
||||
engine = sgl.Engine(
|
||||
model_path=model_path, random_seed=42, quantization=quantization_args
|
||||
)
|
||||
engine.generate(prompt, sampling_params)
|
||||
engine.shutdown()
|
||||
|
||||
def test_2_torchao_args(self):
|
||||
|
||||
# we don't test int8dq because currently there is conflict between int8dq and capture cuda graph
|
||||
torchao_args_list = [
|
||||
# "int8dq",
|
||||
"int8wo",
|
||||
"fp8wo",
|
||||
"fp8dq-per_tensor",
|
||||
"fp8dq-per_row",
|
||||
] + [f"int4wo-{group_size}" for group_size in [32, 64, 128, 256]]
|
||||
|
||||
prompt = "Today is a sunny day and I like"
|
||||
model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
sampling_params = {"temperature": 0, "max_new_tokens": 8}
|
||||
|
||||
for torchao_config in torchao_args_list:
|
||||
engine = sgl.Engine(
|
||||
model_path=model_path, random_seed=42, torchao_config=torchao_config
|
||||
)
|
||||
engine.generate(prompt, sampling_params)
|
||||
engine.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,121 @@
|
||||
"""
|
||||
Unit tests for enable_tokenizer_batch_encode feature.
|
||||
|
||||
This tests the batch tokenization functionality which allows processing
|
||||
multiple text inputs in a single batch for improved performance.
|
||||
|
||||
Usage:
|
||||
python3 -m unittest test_tokenizer_batch_encode.TestTokenizerBatchEncode.test_batch_validation_constraints
|
||||
python3 -m unittest test_tokenizer_batch_encode.TestTokenizerBatchEncodeUnit.test_batch_tokenize_and_process_logic
|
||||
python3 -m unittest test_tokenizer_batch_encode.TestTokenizerBatchEncodeLogic.test_batch_processing_path
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from sglang.srt.managers.io_struct import GenerateReqInput
|
||||
from sglang.srt.managers.tokenizer_manager import TokenizerManager
|
||||
from sglang.srt.server_args import PortArgs, ServerArgs
|
||||
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
|
||||
class TestTokenizerBatchEncode(unittest.TestCase):
|
||||
"""Test cases for tokenizer batch encoding validation and setup."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test fixtures."""
|
||||
self.server_args = ServerArgs(
|
||||
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
enable_tokenizer_batch_encode=True,
|
||||
)
|
||||
self.port_args = PortArgs.init_new(self.server_args)
|
||||
|
||||
with patch("zmq.asyncio.Context"), patch(
|
||||
"sglang.srt.utils.get_zmq_socket"
|
||||
), patch(
|
||||
"sglang.srt.utils.hf_transformers_utils.get_tokenizer"
|
||||
) as mock_tokenizer:
|
||||
|
||||
mock_tokenizer.return_value = Mock(vocab_size=32000)
|
||||
self.tokenizer_manager = TokenizerManager(self.server_args, self.port_args)
|
||||
|
||||
def test_batch_encode_enabled(self):
|
||||
"""Test that batch encoding is enabled when configured."""
|
||||
self.assertTrue(self.server_args.enable_tokenizer_batch_encode)
|
||||
|
||||
def test_batch_encode_disabled(self):
|
||||
"""Test that batch encoding can be disabled."""
|
||||
server_args_disabled = ServerArgs(
|
||||
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
enable_tokenizer_batch_encode=False,
|
||||
)
|
||||
self.assertFalse(server_args_disabled.enable_tokenizer_batch_encode)
|
||||
|
||||
def test_multimodal_input_validation(self):
|
||||
"""Test that multimodal inputs are rejected in batch mode."""
|
||||
req = GenerateReqInput(text="test", image_data=["dummy"])
|
||||
req.contains_mm_input = Mock(return_value=True)
|
||||
|
||||
batch_obj = Mock()
|
||||
batch_obj.__getitem__ = lambda self, i: req
|
||||
|
||||
self.tokenizer_manager.is_generation = True
|
||||
|
||||
with self.assertRaises(ValueError) as cm:
|
||||
self.tokenizer_manager._validate_batch_tokenization_constraints(
|
||||
1, batch_obj
|
||||
)
|
||||
|
||||
self.assertIn("multimodal", str(cm.exception))
|
||||
|
||||
def test_pretokenized_input_validation(self):
|
||||
"""Test that pre-tokenized inputs are rejected in batch mode."""
|
||||
req = GenerateReqInput(input_ids=[1, 2, 3])
|
||||
|
||||
batch_obj = Mock()
|
||||
batch_obj.__getitem__ = lambda self, i: req
|
||||
|
||||
with self.assertRaises(ValueError) as cm:
|
||||
self.tokenizer_manager._validate_batch_tokenization_constraints(
|
||||
1, batch_obj
|
||||
)
|
||||
|
||||
self.assertIn("pre-tokenized", str(cm.exception))
|
||||
|
||||
def test_input_embeds_validation(self):
|
||||
"""Test that input embeds are rejected in batch mode."""
|
||||
req = GenerateReqInput(input_embeds=[0.1, 0.2])
|
||||
|
||||
batch_obj = Mock()
|
||||
batch_obj.__getitem__ = lambda self, i: req
|
||||
|
||||
with self.assertRaises(ValueError) as cm:
|
||||
self.tokenizer_manager._validate_batch_tokenization_constraints(
|
||||
1, batch_obj
|
||||
)
|
||||
|
||||
self.assertIn("input_embeds", str(cm.exception))
|
||||
|
||||
def test_valid_text_only_requests_pass_validation(self):
|
||||
"""Test that valid text-only requests pass validation."""
|
||||
# Create valid requests (text-only)
|
||||
requests = []
|
||||
for i in range(3):
|
||||
req = GenerateReqInput(text=f"test text {i}")
|
||||
req.contains_mm_input = Mock(return_value=False)
|
||||
requests.append(req)
|
||||
|
||||
batch_obj = Mock()
|
||||
batch_obj.__getitem__ = Mock(side_effect=lambda i: requests[i])
|
||||
|
||||
# Should not raise any exception
|
||||
try:
|
||||
self.tokenizer_manager._validate_batch_tokenization_constraints(
|
||||
3, batch_obj
|
||||
)
|
||||
except Exception as e:
|
||||
self.fail(f"Validation failed for valid text-only requests: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
@@ -0,0 +1,386 @@
|
||||
"""
|
||||
Unit tests for TokenizerManager helper methods.
|
||||
|
||||
This tests the refactored tokenization functionality including input format detection,
|
||||
tokenizer input preparation, and result extraction logic.
|
||||
|
||||
Usage:
|
||||
python3 -m unittest test_tokenizer_manager.TestInputFormatDetection
|
||||
python3 -m unittest test_tokenizer_manager.TestTokenizerInputPreparation
|
||||
python3 -m unittest test_tokenizer_manager.TestTokenizerResultExtraction
|
||||
python3 -m unittest test_tokenizer_manager.TestTokenizerManagerIntegration
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from sglang.srt.managers.tokenizer_manager import TokenizerManager
|
||||
from sglang.srt.server_args import PortArgs, ServerArgs
|
||||
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
|
||||
class TestInputFormatDetection(unittest.TestCase):
|
||||
"""Test cases for _detect_input_format method."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test fixtures."""
|
||||
with patch("sglang.srt.utils.get_device", return_value="cpu"):
|
||||
self.server_args = ServerArgs(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
|
||||
self.port_args = PortArgs.init_new(self.server_args)
|
||||
|
||||
with patch("zmq.asyncio.Context"), patch(
|
||||
"sglang.srt.utils.get_zmq_socket"
|
||||
), patch(
|
||||
"sglang.srt.utils.hf_transformers_utils.get_tokenizer"
|
||||
) as mock_tokenizer:
|
||||
mock_tokenizer.return_value = Mock(vocab_size=32000)
|
||||
self.tokenizer_manager = TokenizerManager(self.server_args, self.port_args)
|
||||
|
||||
def test_detect_single_string(self):
|
||||
"""Test detection of single string input."""
|
||||
text = "Hello world"
|
||||
result = self.tokenizer_manager._detect_input_format(
|
||||
text, is_cross_encoder=False
|
||||
)
|
||||
self.assertEqual(result, "single_string")
|
||||
|
||||
def test_detect_single_string_cross_encoder_disabled(self):
|
||||
"""Test single string with cross_encoder disabled still returns single_string."""
|
||||
text = "Hello world"
|
||||
result = self.tokenizer_manager._detect_input_format(
|
||||
text, is_cross_encoder=True
|
||||
)
|
||||
self.assertEqual(result, "single_string")
|
||||
|
||||
def test_detect_batch_strings(self):
|
||||
"""Test detection of batch string inputs."""
|
||||
texts = ["Hello", "World", "How are you?"]
|
||||
result = self.tokenizer_manager._detect_input_format(
|
||||
texts, is_cross_encoder=False
|
||||
)
|
||||
self.assertEqual(result, "batch_strings")
|
||||
|
||||
def test_detect_batch_strings_cross_encoder_disabled(self):
|
||||
"""Test batch strings with cross_encoder disabled."""
|
||||
texts = ["Hello", "World"]
|
||||
result = self.tokenizer_manager._detect_input_format(
|
||||
texts, is_cross_encoder=True
|
||||
)
|
||||
self.assertEqual(result, "batch_strings")
|
||||
|
||||
def test_detect_cross_encoder_single_pair(self):
|
||||
"""Test detection of cross-encoder single pair."""
|
||||
texts = [["query text", "document text"]]
|
||||
result = self.tokenizer_manager._detect_input_format(
|
||||
texts, is_cross_encoder=True
|
||||
)
|
||||
self.assertEqual(result, "cross_encoder_pairs")
|
||||
|
||||
def test_detect_cross_encoder_multiple_pairs(self):
|
||||
"""Test detection of cross-encoder multiple pairs."""
|
||||
texts = [["q1", "d1"], ["q2", "d2"], ["q3", "d3"]]
|
||||
result = self.tokenizer_manager._detect_input_format(
|
||||
texts, is_cross_encoder=True
|
||||
)
|
||||
self.assertEqual(result, "cross_encoder_pairs")
|
||||
|
||||
def test_detect_cross_encoder_disabled_with_pairs(self):
|
||||
"""Test pairs with cross_encoder disabled should return batch_strings."""
|
||||
texts = [["query", "document"]]
|
||||
result = self.tokenizer_manager._detect_input_format(
|
||||
texts, is_cross_encoder=False
|
||||
)
|
||||
self.assertEqual(result, "batch_strings")
|
||||
|
||||
def test_detect_empty_list(self):
|
||||
"""Test detection with empty list."""
|
||||
texts = []
|
||||
result = self.tokenizer_manager._detect_input_format(
|
||||
texts, is_cross_encoder=True
|
||||
)
|
||||
self.assertEqual(result, "batch_strings")
|
||||
|
||||
def test_detect_malformed_cross_encoder_pairs(self):
|
||||
"""Test malformed cross-encoder pairs (not length 2)."""
|
||||
texts = [["query only"]] # Single element, not a pair
|
||||
result = self.tokenizer_manager._detect_input_format(
|
||||
texts, is_cross_encoder=True
|
||||
)
|
||||
self.assertEqual(result, "batch_strings")
|
||||
|
||||
texts = [["query", "doc", "extra"]] # Three elements, not a pair
|
||||
result = self.tokenizer_manager._detect_input_format(
|
||||
texts, is_cross_encoder=True
|
||||
)
|
||||
self.assertEqual(result, "batch_strings")
|
||||
|
||||
|
||||
class TestTokenizerInputPreparation(unittest.TestCase):
|
||||
"""Test cases for _prepare_tokenizer_input method."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test fixtures."""
|
||||
with patch("sglang.srt.utils.get_device", return_value="cpu"):
|
||||
self.server_args = ServerArgs(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
|
||||
self.port_args = PortArgs.init_new(self.server_args)
|
||||
|
||||
with patch("zmq.asyncio.Context"), patch(
|
||||
"sglang.srt.utils.get_zmq_socket"
|
||||
), patch(
|
||||
"sglang.srt.utils.hf_transformers_utils.get_tokenizer"
|
||||
) as mock_tokenizer:
|
||||
mock_tokenizer.return_value = Mock(vocab_size=32000)
|
||||
self.tokenizer_manager = TokenizerManager(self.server_args, self.port_args)
|
||||
|
||||
def test_prepare_single_string_input(self):
|
||||
"""Test preparation of single string input."""
|
||||
text = "Hello world"
|
||||
result = self.tokenizer_manager._prepare_tokenizer_input(text, "single_string")
|
||||
self.assertEqual(result, ["Hello world"])
|
||||
|
||||
def test_prepare_batch_strings_input(self):
|
||||
"""Test preparation of batch strings input."""
|
||||
texts = ["Hello", "World", "Test"]
|
||||
result = self.tokenizer_manager._prepare_tokenizer_input(texts, "batch_strings")
|
||||
self.assertEqual(result, ["Hello", "World", "Test"])
|
||||
|
||||
def test_prepare_cross_encoder_pairs_input(self):
|
||||
"""Test preparation of cross-encoder pairs input."""
|
||||
texts = [["query1", "doc1"], ["query2", "doc2"]]
|
||||
result = self.tokenizer_manager._prepare_tokenizer_input(
|
||||
texts, "cross_encoder_pairs"
|
||||
)
|
||||
self.assertEqual(result, [["query1", "doc1"], ["query2", "doc2"]])
|
||||
|
||||
def test_prepare_cross_encoder_single_pair_input(self):
|
||||
"""Test preparation of single cross-encoder pair."""
|
||||
texts = [["query text", "document text"]]
|
||||
result = self.tokenizer_manager._prepare_tokenizer_input(
|
||||
texts, "cross_encoder_pairs"
|
||||
)
|
||||
self.assertEqual(result, [["query text", "document text"]])
|
||||
|
||||
def test_prepare_unknown_input_format(self):
|
||||
"""Test preparation with unknown input format falls back to returning as-is."""
|
||||
texts = ["test"]
|
||||
result = self.tokenizer_manager._prepare_tokenizer_input(
|
||||
texts, "unknown_format"
|
||||
)
|
||||
self.assertEqual(result, ["test"])
|
||||
|
||||
|
||||
class TestTokenizerResultExtraction(unittest.TestCase):
|
||||
"""Test cases for _extract_tokenizer_results method."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test fixtures."""
|
||||
with patch("sglang.srt.utils.get_device", return_value="cpu"):
|
||||
self.server_args = ServerArgs(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
|
||||
self.port_args = PortArgs.init_new(self.server_args)
|
||||
|
||||
with patch("zmq.asyncio.Context"), patch(
|
||||
"sglang.srt.utils.get_zmq_socket"
|
||||
), patch(
|
||||
"sglang.srt.utils.hf_transformers_utils.get_tokenizer"
|
||||
) as mock_tokenizer:
|
||||
mock_tokenizer.return_value = Mock(vocab_size=32000)
|
||||
self.tokenizer_manager = TokenizerManager(self.server_args, self.port_args)
|
||||
|
||||
def test_extract_single_string_results(self):
|
||||
"""Test extraction for single string input."""
|
||||
input_ids = [[101, 2129, 102]]
|
||||
token_type_ids = [[0, 0, 0]]
|
||||
|
||||
result_input_ids, result_token_type_ids = (
|
||||
self.tokenizer_manager._extract_tokenizer_results(
|
||||
input_ids, token_type_ids, "single_string", original_batch_size=1
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(result_input_ids, [101, 2129, 102])
|
||||
self.assertEqual(result_token_type_ids, [0, 0, 0])
|
||||
|
||||
def test_extract_single_cross_encoder_results(self):
|
||||
"""Test extraction for single cross-encoder pair."""
|
||||
input_ids = [[101, 2129, 102, 4068, 102]]
|
||||
token_type_ids = [[0, 0, 0, 1, 1]]
|
||||
|
||||
result_input_ids, result_token_type_ids = (
|
||||
self.tokenizer_manager._extract_tokenizer_results(
|
||||
input_ids, token_type_ids, "cross_encoder_pairs", original_batch_size=1
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(result_input_ids, [101, 2129, 102, 4068, 102])
|
||||
self.assertEqual(result_token_type_ids, [0, 0, 0, 1, 1])
|
||||
|
||||
def test_extract_batch_results(self):
|
||||
"""Test extraction for batch inputs."""
|
||||
input_ids = [[101, 2129, 102], [101, 4068, 102]]
|
||||
token_type_ids = [[0, 0, 0], [0, 0, 0]]
|
||||
|
||||
result_input_ids, result_token_type_ids = (
|
||||
self.tokenizer_manager._extract_tokenizer_results(
|
||||
input_ids, token_type_ids, "batch_strings", original_batch_size=2
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(result_input_ids, [[101, 2129, 102], [101, 4068, 102]])
|
||||
self.assertEqual(result_token_type_ids, [[0, 0, 0], [0, 0, 0]])
|
||||
|
||||
def test_extract_multiple_cross_encoder_results(self):
|
||||
"""Test extraction for multiple cross-encoder pairs."""
|
||||
input_ids = [[101, 2129, 102, 4068, 102], [101, 7592, 102, 2088, 102]]
|
||||
token_type_ids = [[0, 0, 0, 1, 1], [0, 0, 0, 1, 1]]
|
||||
|
||||
result_input_ids, result_token_type_ids = (
|
||||
self.tokenizer_manager._extract_tokenizer_results(
|
||||
input_ids, token_type_ids, "cross_encoder_pairs", original_batch_size=2
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
result_input_ids, [[101, 2129, 102, 4068, 102], [101, 7592, 102, 2088, 102]]
|
||||
)
|
||||
self.assertEqual(result_token_type_ids, [[0, 0, 0, 1, 1], [0, 0, 0, 1, 1]])
|
||||
|
||||
def test_extract_empty_results(self):
|
||||
"""Test extraction with empty results."""
|
||||
input_ids = []
|
||||
token_type_ids = None
|
||||
|
||||
result_input_ids, result_token_type_ids = (
|
||||
self.tokenizer_manager._extract_tokenizer_results(
|
||||
input_ids, token_type_ids, "single_string", original_batch_size=1
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(result_input_ids, [])
|
||||
self.assertIsNone(result_token_type_ids)
|
||||
|
||||
def test_extract_with_none_token_type_ids(self):
|
||||
"""Test extraction when token_type_ids is None."""
|
||||
input_ids = [[101, 2129, 102]]
|
||||
token_type_ids = None
|
||||
|
||||
result_input_ids, result_token_type_ids = (
|
||||
self.tokenizer_manager._extract_tokenizer_results(
|
||||
input_ids, token_type_ids, "single_string", original_batch_size=1
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(result_input_ids, [101, 2129, 102])
|
||||
self.assertIsNone(result_token_type_ids)
|
||||
|
||||
|
||||
class TestTokenizerManagerIntegration(unittest.TestCase):
|
||||
"""Integration tests combining multiple helper methods."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test fixtures."""
|
||||
with patch("sglang.srt.utils.get_device", return_value="cpu"):
|
||||
self.server_args = ServerArgs(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
|
||||
self.port_args = PortArgs.init_new(self.server_args)
|
||||
|
||||
with patch("zmq.asyncio.Context"), patch(
|
||||
"sglang.srt.utils.get_zmq_socket"
|
||||
), patch(
|
||||
"sglang.srt.utils.hf_transformers_utils.get_tokenizer"
|
||||
) as mock_tokenizer:
|
||||
mock_tokenizer.return_value = Mock(vocab_size=32000)
|
||||
self.tokenizer_manager = TokenizerManager(self.server_args, self.port_args)
|
||||
|
||||
def test_full_workflow_single_string(self):
|
||||
"""Test complete workflow for single string input."""
|
||||
text = "Hello world"
|
||||
|
||||
# Step 1: Detect format
|
||||
input_format = self.tokenizer_manager._detect_input_format(
|
||||
text, is_cross_encoder=False
|
||||
)
|
||||
self.assertEqual(input_format, "single_string")
|
||||
|
||||
# Step 2: Prepare input
|
||||
tokenizer_input = self.tokenizer_manager._prepare_tokenizer_input(
|
||||
text, input_format
|
||||
)
|
||||
self.assertEqual(tokenizer_input, ["Hello world"])
|
||||
|
||||
# Step 3: Extract results (simulated tokenizer output)
|
||||
mock_input_ids = [[101, 2129, 4248, 102]]
|
||||
mock_token_type_ids = None
|
||||
|
||||
result_input_ids, result_token_type_ids = (
|
||||
self.tokenizer_manager._extract_tokenizer_results(
|
||||
mock_input_ids, mock_token_type_ids, input_format, original_batch_size=1
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(result_input_ids, [101, 2129, 4248, 102])
|
||||
self.assertIsNone(result_token_type_ids)
|
||||
|
||||
def test_full_workflow_cross_encoder_pairs(self):
|
||||
"""Test complete workflow for cross-encoder pairs."""
|
||||
texts = [
|
||||
["How many people live in Berlin?", "Berlin is well known for its museums."]
|
||||
]
|
||||
|
||||
# Step 1: Detect format
|
||||
input_format = self.tokenizer_manager._detect_input_format(
|
||||
texts, is_cross_encoder=True
|
||||
)
|
||||
self.assertEqual(input_format, "cross_encoder_pairs")
|
||||
|
||||
# Step 2: Prepare input
|
||||
tokenizer_input = self.tokenizer_manager._prepare_tokenizer_input(
|
||||
texts, input_format
|
||||
)
|
||||
self.assertEqual(tokenizer_input, texts)
|
||||
|
||||
# Step 3: Extract results (simulated tokenizer output for cross-encoder)
|
||||
mock_input_ids = [[101, 2129, 2116, 102, 4068, 2003, 102]]
|
||||
mock_token_type_ids = [[0, 0, 0, 0, 1, 1, 1]]
|
||||
|
||||
result_input_ids, result_token_type_ids = (
|
||||
self.tokenizer_manager._extract_tokenizer_results(
|
||||
mock_input_ids, mock_token_type_ids, input_format, original_batch_size=1
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(result_input_ids, [101, 2129, 2116, 102, 4068, 2003, 102])
|
||||
self.assertEqual(result_token_type_ids, [0, 0, 0, 0, 1, 1, 1])
|
||||
|
||||
def test_full_workflow_batch_strings(self):
|
||||
"""Test complete workflow for batch strings."""
|
||||
texts = ["Hello", "World", "Test"]
|
||||
|
||||
# Step 1: Detect format
|
||||
input_format = self.tokenizer_manager._detect_input_format(
|
||||
texts, is_cross_encoder=False
|
||||
)
|
||||
self.assertEqual(input_format, "batch_strings")
|
||||
|
||||
# Step 2: Prepare input
|
||||
tokenizer_input = self.tokenizer_manager._prepare_tokenizer_input(
|
||||
texts, input_format
|
||||
)
|
||||
self.assertEqual(tokenizer_input, ["Hello", "World", "Test"])
|
||||
|
||||
# Step 3: Extract results (simulated tokenizer output)
|
||||
mock_input_ids = [[101, 7592, 102], [101, 2088, 102], [101, 2774, 102]]
|
||||
mock_token_type_ids = None
|
||||
|
||||
result_input_ids, result_token_type_ids = (
|
||||
self.tokenizer_manager._extract_tokenizer_results(
|
||||
mock_input_ids, mock_token_type_ids, input_format, original_batch_size=3
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
result_input_ids, [[101, 7592, 102], [101, 2088, 102], [101, 2774, 102]]
|
||||
)
|
||||
self.assertIsNone(result_token_type_ids)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
@@ -0,0 +1,49 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest test_torch_flex_attention_backend.TestTorchFlexAttnBackend.test_gsm8k
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestTorchFlexAttnBackend(CustomTestCase):
|
||||
def test_gsm8k(self):
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
process = popen_launch_server(
|
||||
model,
|
||||
base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=["--attention-backend", "flex_attention"],
|
||||
)
|
||||
|
||||
try:
|
||||
args = SimpleNamespace(
|
||||
num_shots=8,
|
||||
data_path=None,
|
||||
num_questions=100,
|
||||
parallel=10,
|
||||
max_new_tokens=512,
|
||||
host="http://127.0.0.1",
|
||||
port=int(base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["accuracy"], 0.62)
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,30 @@
|
||||
import unittest
|
||||
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
run_bench_offline_throughput,
|
||||
)
|
||||
|
||||
|
||||
class TestTorchTP(CustomTestCase):
|
||||
def test_torch_native_llama(self):
|
||||
output_throughput = run_bench_offline_throughput(
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
[
|
||||
"--tp",
|
||||
"2",
|
||||
# This cannot run anymore with the new torch version.
|
||||
# "--json-model-override-args",
|
||||
# '{"architectures": ["TorchNativeLlamaForCausalLM"]}',
|
||||
"--disable-cuda-graph",
|
||||
],
|
||||
)
|
||||
|
||||
if is_in_ci():
|
||||
self.assertGreater(output_throughput, 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,272 @@
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
import unittest
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import requests
|
||||
import zmq
|
||||
|
||||
from sglang import Engine
|
||||
from sglang.srt.tracing.trace import *
|
||||
from sglang.srt.utils import get_zmq_socket, 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,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Req:
|
||||
rid: int
|
||||
trace_context: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class TestTrace(CustomTestCase):
|
||||
def __launch_otel_jaeger(self):
|
||||
cmd = [
|
||||
"docker",
|
||||
"compose",
|
||||
"-f",
|
||||
"../../examples/monitoring/tracing_compose.yaml",
|
||||
"up",
|
||||
"-d",
|
||||
]
|
||||
proc = subprocess.run(cmd)
|
||||
|
||||
if proc.returncode != 0:
|
||||
print("launch opentelemetry collector and jaeger docker err")
|
||||
return False
|
||||
return True
|
||||
|
||||
def __stop_otel_jaeger(self):
|
||||
cmd = [
|
||||
"docker",
|
||||
"compose",
|
||||
"-f",
|
||||
"../../examples/monitoring/tracing_compose.yaml",
|
||||
"down",
|
||||
]
|
||||
proc = subprocess.run(cmd)
|
||||
|
||||
if proc.returncode != 0:
|
||||
print("stop opentelemetry collector and jaeger docker err")
|
||||
return False
|
||||
return True
|
||||
|
||||
def __clear_trace_file(self):
|
||||
try:
|
||||
os.remove("/tmp/otel_trace.json")
|
||||
except:
|
||||
pass
|
||||
|
||||
def test_trace_enable(self):
|
||||
self.__clear_trace_file()
|
||||
assert self.__launch_otel_jaeger()
|
||||
|
||||
process = popen_launch_server(
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=["--enable-trace", "--otlp-traces-endpoint", "0.0.0.0:4317"],
|
||||
)
|
||||
|
||||
try:
|
||||
# Make some requests to generate trace data
|
||||
response = requests.get(f"{DEFAULT_URL_FOR_TEST}/health_generate")
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
response = requests.post(
|
||||
f"{DEFAULT_URL_FOR_TEST}/generate",
|
||||
json={
|
||||
"text": "The capital of France is",
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 32,
|
||||
},
|
||||
"stream": True,
|
||||
},
|
||||
stream=True,
|
||||
)
|
||||
for _ in response.iter_lines(decode_unicode=False):
|
||||
pass
|
||||
|
||||
# sleep for a few seconds to wait for opentelemetry collector to asynchronously export data to file.
|
||||
time.sleep(10)
|
||||
|
||||
# check trace file
|
||||
assert os.path.isfile("/tmp/otel_trace.json"), "trace file not exist"
|
||||
assert os.path.getsize("/tmp/otel_trace.json") > 0, "trace file is empty"
|
||||
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
assert self.__stop_otel_jaeger()
|
||||
|
||||
def test_trace_engine_enable(self):
|
||||
self.__clear_trace_file()
|
||||
assert self.__launch_otel_jaeger()
|
||||
|
||||
prompt = "Today is a sunny day and I like"
|
||||
model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
sampling_params = {"temperature": 0, "max_new_tokens": 8}
|
||||
|
||||
engine = Engine(
|
||||
model_path=model_path,
|
||||
random_seed=42,
|
||||
enable_trace=True,
|
||||
otlp_traces_endpoint="localhost:4317",
|
||||
)
|
||||
|
||||
try:
|
||||
engine.generate(prompt, sampling_params)
|
||||
|
||||
# sleep for a few seconds to wait for opentelemetry collector to asynchronously export data to file.
|
||||
time.sleep(10)
|
||||
|
||||
# check trace file
|
||||
assert os.path.isfile("/tmp/otel_trace.json"), "trace file not exist"
|
||||
assert os.path.getsize("/tmp/otel_trace.json") > 0, "trace file is empty"
|
||||
finally:
|
||||
engine.shutdown()
|
||||
assert self.__stop_otel_jaeger()
|
||||
|
||||
def test_trace_engine_encode(self):
|
||||
self.__clear_trace_file()
|
||||
assert self.__launch_otel_jaeger()
|
||||
|
||||
prompt = "Today is a sunny day and I like"
|
||||
model_path = "Qwen/Qwen2-7B"
|
||||
|
||||
engine = Engine(
|
||||
model_path=model_path,
|
||||
random_seed=42,
|
||||
enable_trace=True,
|
||||
otlp_traces_endpoint="localhost:4317",
|
||||
is_embedding=True,
|
||||
)
|
||||
|
||||
try:
|
||||
engine.encode(prompt)
|
||||
|
||||
# sleep for a few seconds to wait for opentelemetry collector to asynchronously export data to file.
|
||||
time.sleep(10)
|
||||
|
||||
# check trace file
|
||||
assert os.path.isfile("/tmp/otel_trace.json"), "trace file not exist"
|
||||
assert os.path.getsize("/tmp/otel_trace.json") > 0, "trace file is empty"
|
||||
finally:
|
||||
engine.shutdown()
|
||||
assert self.__stop_otel_jaeger()
|
||||
|
||||
def test_slice_trace_simple(self):
|
||||
self.__clear_trace_file()
|
||||
assert self.__launch_otel_jaeger()
|
||||
try:
|
||||
process_tracing_init("0.0.0.0:4317", "test")
|
||||
trace_set_thread_info("Test")
|
||||
trace_req_start(0)
|
||||
trace_slice_start("test slice", 0)
|
||||
time.sleep(1)
|
||||
trace_slice_end("test slice", 0)
|
||||
trace_req_finish(0)
|
||||
|
||||
# sleep for a few seconds to wait for opentelemetry collector to asynchronously export data to file.
|
||||
time.sleep(10)
|
||||
# check trace file
|
||||
assert os.path.isfile("/tmp/otel_trace.json"), "trace file not exist"
|
||||
assert os.path.getsize("/tmp/otel_trace.json") > 0, "trace file is empty"
|
||||
finally:
|
||||
assert self.__stop_otel_jaeger()
|
||||
|
||||
def test_slice_trace_complex(self):
|
||||
self.__clear_trace_file()
|
||||
assert self.__launch_otel_jaeger()
|
||||
try:
|
||||
process_tracing_init("0.0.0.0:4317", "test")
|
||||
trace_set_thread_info("Test")
|
||||
trace_req_start(0)
|
||||
trace_slice_start("", 0, anonymous=True)
|
||||
time.sleep(1)
|
||||
trace_slice_end("slice A", 0, auto_next_anon=True)
|
||||
time.sleep(1)
|
||||
trace_slice_end("slice B", 0, auto_next_anon=True)
|
||||
time.sleep(1)
|
||||
trace_slice_end("slice C", 0, thread_finish_flag=True)
|
||||
trace_req_finish(0)
|
||||
|
||||
# sleep for a few seconds to wait for opentelemetry collector to asynchronously export data to file.
|
||||
time.sleep(10)
|
||||
# check trace file
|
||||
assert os.path.isfile("/tmp/otel_trace.json"), "trace file not exist"
|
||||
assert os.path.getsize("/tmp/otel_trace.json") > 0, "trace file is empty"
|
||||
finally:
|
||||
assert self.__stop_otel_jaeger()
|
||||
|
||||
def test_trace_context_propagete(self):
|
||||
def __process_work():
|
||||
process_tracing_init("0.0.0.0:4317", "test")
|
||||
trace_set_thread_info("Sub Process")
|
||||
|
||||
context = zmq.Context(2)
|
||||
recv_from_main = get_zmq_socket(
|
||||
context, zmq.PULL, "ipc:///tmp/zmq_test.ipc", True
|
||||
)
|
||||
|
||||
try:
|
||||
req = recv_from_main.recv_pyobj()
|
||||
trace_set_proc_propagate_context(req.rid, req.trace_context)
|
||||
trace_slice_start("work", req.rid)
|
||||
time.sleep(1)
|
||||
trace_slice_end("work", req.rid, thread_finish_flag=True)
|
||||
finally:
|
||||
recv_from_main.close()
|
||||
context.term()
|
||||
|
||||
self.__clear_trace_file()
|
||||
assert self.__launch_otel_jaeger()
|
||||
|
||||
context = zmq.Context(2)
|
||||
send_to_subproc = get_zmq_socket(
|
||||
context, zmq.PUSH, "ipc:///tmp/zmq_test.ipc", False
|
||||
)
|
||||
try:
|
||||
process_tracing_init("0.0.0.0:4317", "test")
|
||||
trace_set_thread_info("Main Process")
|
||||
|
||||
subproc = mp.Process(target=__process_work)
|
||||
subproc.start()
|
||||
|
||||
# sleep for a few second to ensure subprocess init
|
||||
time.sleep(1)
|
||||
|
||||
req = Req(rid=0)
|
||||
trace_req_start(req.rid)
|
||||
trace_slice_start("dispatch", req.rid)
|
||||
time.sleep(1)
|
||||
req.trace_context = trace_get_proc_propagate_context(req.rid)
|
||||
send_to_subproc.send_pyobj(req)
|
||||
trace_slice_end("dispatch", req.rid)
|
||||
|
||||
subproc.join()
|
||||
trace_req_finish(req.rid)
|
||||
|
||||
# sleep for a few seconds to wait for opentelemetry collector to asynchronously export data to file.
|
||||
time.sleep(10)
|
||||
# check trace file
|
||||
assert os.path.isfile("/tmp/otel_trace.json"), "trace file not exist"
|
||||
assert os.path.getsize("/tmp/otel_trace.json") > 0, "trace file is empty"
|
||||
|
||||
finally:
|
||||
send_to_subproc.close()
|
||||
context.term()
|
||||
assert self.__stop_otel_jaeger()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,259 @@
|
||||
import random
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.triton_ops.decode_attention import (
|
||||
decode_attention_fwd_grouped,
|
||||
)
|
||||
from sglang.srt.layers.attention.triton_ops.rocm_mla_decode_rope import (
|
||||
decode_attention_fwd_grouped_rope,
|
||||
)
|
||||
from sglang.srt.layers.rotary_embedding import DeepseekScalingRotaryEmbedding
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestTritonAttentionMLA(CustomTestCase):
|
||||
|
||||
def _set_all_seeds(self, seed):
|
||||
"""Set all random seeds for reproducibility."""
|
||||
random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
def setUp(self):
|
||||
# Set seeds before each test method
|
||||
self._set_all_seeds(42)
|
||||
|
||||
def preprocess_kv_cache(self, kv_cache, kv_lora_rank):
|
||||
latent_cache = kv_cache
|
||||
v_input = latent_cache[..., :kv_lora_rank]
|
||||
v_input = v_input.contiguous().unsqueeze(1)
|
||||
k_input = latent_cache.unsqueeze(1)
|
||||
k_input[..., :kv_lora_rank] = v_input
|
||||
|
||||
return k_input, v_input
|
||||
|
||||
def input_helper(
|
||||
self,
|
||||
B,
|
||||
H,
|
||||
S,
|
||||
kv_lora_rank,
|
||||
rotary_dim,
|
||||
qk_rope_head_dim,
|
||||
num_kv_splits,
|
||||
dtype,
|
||||
device,
|
||||
rope_base=10,
|
||||
rope_max_seq_len=16384,
|
||||
rope_scaling=1.0,
|
||||
is_neox_style=False,
|
||||
):
|
||||
q = torch.randn(
|
||||
B, H, kv_lora_rank + qk_rope_head_dim, device=device, dtype=dtype
|
||||
)
|
||||
kv_cache = torch.randn(
|
||||
B * S, kv_lora_rank + qk_rope_head_dim, dtype=dtype, device=device
|
||||
)
|
||||
kv_indptr = torch.arange(B + 1, device=device) * S
|
||||
kv_indices = torch.arange(B * S, device=device)
|
||||
attn_logits = torch.empty(
|
||||
B, H, num_kv_splits, kv_lora_rank + 1, dtype=dtype, device=device
|
||||
)
|
||||
rotary_emb = DeepseekScalingRotaryEmbedding(
|
||||
qk_rope_head_dim,
|
||||
rotary_dim,
|
||||
rope_max_seq_len,
|
||||
rope_base,
|
||||
is_neox_style,
|
||||
rope_scaling,
|
||||
q.dtype,
|
||||
device="cpu",
|
||||
).cuda()
|
||||
positions = torch.tensor([S], device=device).unsqueeze(0).repeat(B, 1)
|
||||
|
||||
return kv_indptr, kv_indices, q, kv_cache, attn_logits, rotary_emb, positions
|
||||
|
||||
def ref_compute_full_fwd(
|
||||
self,
|
||||
q,
|
||||
k_input,
|
||||
v_input,
|
||||
kv_lora_rank,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
num_kv_splits,
|
||||
sm_scale,
|
||||
logit_cap,
|
||||
rotary_emb,
|
||||
positions,
|
||||
use_rope,
|
||||
device="cuda",
|
||||
):
|
||||
|
||||
B, H = q.shape[0], q.shape[1]
|
||||
S = kv_indptr[1].item()
|
||||
qk_rope_head_dim = k_input.shape[-1] - kv_lora_rank
|
||||
|
||||
q_input = torch.empty(B, H, kv_lora_rank + qk_rope_head_dim, dtype=q.dtype).to(
|
||||
device
|
||||
)
|
||||
q_nope_out, q_pe = q.split([kv_lora_rank, qk_rope_head_dim], dim=-1)
|
||||
k_pe_t = k_input.view(B, 1, S, -1)[:, :, -1:, kv_lora_rank:]
|
||||
|
||||
if use_rope:
|
||||
q_pe, k_pe_t = rotary_emb(positions, q_pe.unsqueeze(2), k_pe_t)
|
||||
q_pe = q_pe.squeeze()
|
||||
|
||||
k_input.view(B, 1, S, -1)[:, :, -1:, kv_lora_rank:] = k_pe_t
|
||||
|
||||
q_input[..., :kv_lora_rank] = q_nope_out
|
||||
q_input[..., kv_lora_rank:] = q_pe
|
||||
|
||||
B, H = q_input.shape[0], q_input.shape[1]
|
||||
kv_lora_rank = v_input.shape[-1]
|
||||
device = q_input.device
|
||||
|
||||
attn_logits = torch.empty(
|
||||
B, H, num_kv_splits, kv_lora_rank + 1, dtype=q_input.dtype, device=device
|
||||
)
|
||||
o = torch.empty(B, H, kv_lora_rank, dtype=q_input.dtype, device=device)
|
||||
|
||||
decode_attention_fwd_grouped(
|
||||
q_input,
|
||||
k_input,
|
||||
v_input,
|
||||
o,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
attn_logits,
|
||||
num_kv_splits,
|
||||
sm_scale,
|
||||
logit_cap,
|
||||
)
|
||||
|
||||
return attn_logits, o, k_pe_t.squeeze()
|
||||
|
||||
def _test_rocm_fused_mla_kernel(
|
||||
self,
|
||||
B,
|
||||
H,
|
||||
S,
|
||||
kv_lora_rank,
|
||||
qk_rope_head_dim,
|
||||
rotary_dim,
|
||||
dtype,
|
||||
use_rope,
|
||||
is_neox_style,
|
||||
num_kv_splits=2,
|
||||
sm_scale=1.0,
|
||||
logit_cap=0.0,
|
||||
device="cuda",
|
||||
):
|
||||
kv_indptr, kv_indices, q, kv_cache, attn_logits, rotary_emb, positions = (
|
||||
self.input_helper(
|
||||
B,
|
||||
H,
|
||||
S,
|
||||
kv_lora_rank,
|
||||
rotary_dim,
|
||||
qk_rope_head_dim,
|
||||
num_kv_splits,
|
||||
dtype,
|
||||
device=device,
|
||||
is_neox_style=is_neox_style,
|
||||
)
|
||||
)
|
||||
|
||||
k_input, v_input = self.preprocess_kv_cache(kv_cache, kv_lora_rank)
|
||||
k_pe_tokens = torch.empty(
|
||||
B, qk_rope_head_dim, dtype=kv_cache.dtype, device=device
|
||||
)
|
||||
tri_o = torch.empty(B, H, kv_lora_rank, dtype=kv_cache.dtype, device=device)
|
||||
|
||||
decode_attention_fwd_grouped_rope(
|
||||
q,
|
||||
k_input,
|
||||
v_input,
|
||||
tri_o,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
k_pe_tokens if use_rope else None,
|
||||
kv_lora_rank,
|
||||
rotary_dim if use_rope else None,
|
||||
rotary_emb.cos_sin_cache if use_rope else None,
|
||||
positions if use_rope else None,
|
||||
attn_logits,
|
||||
num_kv_splits,
|
||||
sm_scale,
|
||||
logit_cap,
|
||||
use_rope,
|
||||
is_neox_style,
|
||||
)
|
||||
|
||||
tri_logits = attn_logits
|
||||
|
||||
# reference
|
||||
ref_logits, ref_o, ref_k_pe_tokens = self.ref_compute_full_fwd(
|
||||
q,
|
||||
k_input,
|
||||
v_input,
|
||||
kv_lora_rank,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
num_kv_splits,
|
||||
sm_scale,
|
||||
logit_cap,
|
||||
rotary_emb,
|
||||
positions,
|
||||
use_rope,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
if use_rope:
|
||||
torch.testing.assert_close(
|
||||
ref_k_pe_tokens, k_pe_tokens.squeeze(), atol=1e-2, rtol=1e-2
|
||||
)
|
||||
torch.testing.assert_close(ref_logits, tri_logits, atol=1e-2, rtol=1e-2)
|
||||
torch.testing.assert_close(ref_o, tri_o, atol=1e-2, rtol=1e-2)
|
||||
|
||||
def test_grouped_rocm_fused_mla(self):
|
||||
configs = [
|
||||
(1, 128, 2048, 512, 64, 64),
|
||||
(1, 128, 2048, 512, 128, 64),
|
||||
(1, 128, 2048, 512, 127, 64),
|
||||
(1, 128, 2050, 512, 127, 64),
|
||||
(1, 128, 2050, 512, 128, 64),
|
||||
(8, 128, 2048, 512, 64, 64),
|
||||
(8, 128, 2048, 512, 128, 64),
|
||||
(8, 128, 2048, 512, 127, 64),
|
||||
(8, 128, 2050, 512, 127, 64),
|
||||
(8, 128, 2050, 512, 128, 64),
|
||||
]
|
||||
dtypes = [torch.bfloat16, torch.float32]
|
||||
use_rope_list = [True, False]
|
||||
is_neox_style_list = [True, False]
|
||||
|
||||
for B, H, S, kv_lora_rank, qk_rope_head_dim, rotary_dim in configs:
|
||||
for dtype in dtypes:
|
||||
for use_rope in use_rope_list:
|
||||
for is_neox_style in is_neox_style_list:
|
||||
self._test_rocm_fused_mla_kernel(
|
||||
B,
|
||||
H,
|
||||
S,
|
||||
kv_lora_rank,
|
||||
qk_rope_head_dim,
|
||||
rotary_dim,
|
||||
dtype,
|
||||
use_rope,
|
||||
is_neox_style,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,192 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from sglang.srt.layers.activation import SiluAndMul
|
||||
from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.triton_kernels import TritonKernelsQuantInfo
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
|
||||
from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat
|
||||
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestFusedMOE(CustomTestCase):
|
||||
NUM_EXPERTS = [8, 64]
|
||||
TOP_KS = [2, 4]
|
||||
|
||||
@staticmethod
|
||||
def create_random_cuda_tensor(shape, dtype, mean=0, std=0.01):
|
||||
"""Create a random CUDA tensor
|
||||
|
||||
Args:
|
||||
shape: Tensor shape
|
||||
dtype: Data type
|
||||
mean: Mean value
|
||||
std: Standard deviation
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Randomly initialized CUDA tensor
|
||||
"""
|
||||
return torch.empty(shape, dtype=dtype, device="cuda").normal_(mean, std)
|
||||
|
||||
def get_tolerance(self, dtype):
|
||||
"""Get tolerance values for different data types
|
||||
|
||||
Args:
|
||||
dtype: Data type
|
||||
|
||||
Returns:
|
||||
tuple: (relative tolerance, absolute tolerance)
|
||||
"""
|
||||
if dtype == torch.float32:
|
||||
return 1e-5, 1e-5
|
||||
elif dtype in [torch.float16, torch.bfloat16]:
|
||||
return 1e-5, 1e-5
|
||||
else:
|
||||
return 1e-2, 1e-2 # Default values for other types
|
||||
|
||||
def torch_naive_moe(
|
||||
self,
|
||||
a,
|
||||
w1,
|
||||
w2,
|
||||
score,
|
||||
topk,
|
||||
return_per_expert: bool = False,
|
||||
):
|
||||
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
|
||||
|
||||
B, D = a.shape
|
||||
a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
|
||||
out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
|
||||
score = torch.softmax(score, dim=-1, dtype=torch.float32)
|
||||
topk_weight, topk_ids = torch.topk(score, topk)
|
||||
topk_weight = topk_weight.view(-1)
|
||||
topk_ids = topk_ids.view(-1)
|
||||
|
||||
if w1.dtype == torch.float8_e4m3fn:
|
||||
w1_compute = w1.to(a.dtype)
|
||||
w2_compute = w2.to(a.dtype)
|
||||
else:
|
||||
w1_compute = w1
|
||||
w2_compute = w2
|
||||
|
||||
for i in range(w1_compute.shape[0]):
|
||||
mask = topk_ids == i
|
||||
if mask.sum():
|
||||
out[mask] = SiluAndMul()(
|
||||
a[mask] @ w1_compute[i].transpose(0, 1)
|
||||
) @ w2_compute[i].transpose(0, 1)
|
||||
|
||||
weighted = out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(
|
||||
out.dtype
|
||||
)
|
||||
|
||||
if return_per_expert:
|
||||
return weighted
|
||||
|
||||
return weighted.sum(dim=1)
|
||||
|
||||
def _test_case(self, m, n, k, e, topk, dtype):
|
||||
rtol, atol = self.get_tolerance(dtype)
|
||||
|
||||
a = self.create_random_cuda_tensor((m, k), dtype)
|
||||
w1 = self.create_random_cuda_tensor((e, 2 * n, k), dtype)
|
||||
w2 = self.create_random_cuda_tensor((e, k, n), dtype)
|
||||
w1_tri = w1.clone()
|
||||
w2_tri = w2.clone()
|
||||
w1_tri = w1_tri.transpose(-2, -1).contiguous()
|
||||
w2_tri = w2_tri.transpose(-2, -1).contiguous()
|
||||
score = self.create_random_cuda_tensor((m, e), dtype)
|
||||
|
||||
topk_op = TopK(
|
||||
top_k=topk,
|
||||
renormalize=False,
|
||||
use_grouped_topk=False,
|
||||
)
|
||||
topk_op.topk_config.output_format = TopKOutputFormat.TRITON_KERNEL
|
||||
triton_topk_output = topk_op.forward_cuda(
|
||||
hidden_states=a,
|
||||
router_logits=score,
|
||||
)
|
||||
|
||||
quant_info = TritonKernelsQuantInfo(w13_weight=w1_tri, w2_weight=w2_tri)
|
||||
|
||||
dispatch_output = StandardDispatchOutput(
|
||||
hidden_states=a, topk_output=triton_topk_output
|
||||
)
|
||||
|
||||
torch_per_expert = self.torch_naive_moe(
|
||||
a, w1, w2, score, topk, return_per_expert=True
|
||||
)
|
||||
torch_combined = torch_per_expert.sum(dim=1)
|
||||
|
||||
def run_runner(config):
|
||||
runner = MoeRunner(MoeRunnerBackend.TRITON_KERNELS, config)
|
||||
result = runner.run(dispatch_output, quant_info)
|
||||
return result.hidden_states
|
||||
|
||||
# Combined output (no_combine=False)
|
||||
non_fused_config = MoeRunnerConfig(inplace=False)
|
||||
non_fused_output = run_runner(non_fused_config)
|
||||
torch.testing.assert_close(
|
||||
non_fused_output, torch_combined, rtol=rtol, atol=atol
|
||||
)
|
||||
|
||||
# Per-expert output (no_combine=True)
|
||||
non_fused_no_combine_config = MoeRunnerConfig(
|
||||
inplace=False, no_combine=True, top_k=topk
|
||||
)
|
||||
non_fused_no_combine_output = run_runner(non_fused_no_combine_config)
|
||||
torch.testing.assert_close(
|
||||
non_fused_no_combine_output, torch_per_expert, rtol=rtol, atol=atol
|
||||
)
|
||||
|
||||
def test_various_configurations(self):
|
||||
m_values = [1, 32, 64, 256]
|
||||
n_values = [128, 1024]
|
||||
k_values = [128, 512, 1024]
|
||||
dtypes = [torch.bfloat16]
|
||||
|
||||
# Calculate total number of tests
|
||||
total_tests = (
|
||||
len(m_values)
|
||||
* len(n_values)
|
||||
* len(k_values)
|
||||
* len(self.NUM_EXPERTS)
|
||||
* len(self.TOP_KS)
|
||||
* len(dtypes)
|
||||
)
|
||||
|
||||
# Create progress bar
|
||||
with tqdm(total=total_tests, desc="Running MoE tests") as pbar:
|
||||
for m in m_values:
|
||||
for n in n_values:
|
||||
for k in k_values:
|
||||
for e in self.NUM_EXPERTS:
|
||||
for topk in self.TOP_KS:
|
||||
for dtype in dtypes:
|
||||
with self.subTest(
|
||||
m=m,
|
||||
n=n,
|
||||
k=k,
|
||||
e=e,
|
||||
topk=topk,
|
||||
dtype=dtype,
|
||||
):
|
||||
self._test_case(
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
e,
|
||||
topk,
|
||||
dtype,
|
||||
)
|
||||
torch.cuda.empty_cache()
|
||||
pbar.update(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,246 @@
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.activation import SiluAndMul
|
||||
from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_moe
|
||||
from sglang.srt.layers.moe.topk import TopKConfig, select_experts
|
||||
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
||||
|
||||
NUM_EXPERTS = [8, 64]
|
||||
TOP_KS = [2, 6]
|
||||
|
||||
|
||||
def quantize_weights(
|
||||
w: torch.Tensor,
|
||||
quant_type: str,
|
||||
group_size: Optional[int],
|
||||
zero_points: bool = False,
|
||||
ref_zero_points_after_scales: bool = False,
|
||||
):
|
||||
assert quant_type in ["w4a16", "w4a16b8", "w8a16", "w8a16b128"]
|
||||
assert not zero_points or group_size is not None, (
|
||||
"to have group zero points, group_size must be provided "
|
||||
"(-1 group_size is channelwise)"
|
||||
)
|
||||
|
||||
orig_device = w.device
|
||||
orig_type = w.dtype
|
||||
size_k, size_n = w.shape
|
||||
|
||||
assert w.is_floating_point(), "w must be float"
|
||||
|
||||
if group_size == -1:
|
||||
group_size = size_k
|
||||
|
||||
# Reshape to [groupsize, -1]
|
||||
if group_size is not None and group_size < size_k:
|
||||
w = w.reshape((-1, group_size, size_n))
|
||||
w = w.permute(1, 0, 2)
|
||||
w = w.reshape((group_size, -1))
|
||||
|
||||
# Compute scale for each group
|
||||
max_val = torch.max(w, 0, keepdim=True).values
|
||||
min_val = torch.min(w, 0, keepdim=True).values
|
||||
|
||||
if quant_type == "w4a16":
|
||||
max_q_val = 15
|
||||
min_q_val = 0
|
||||
elif quant_type == "w4a16b8":
|
||||
max_q_val = 7
|
||||
min_q_val = -1
|
||||
elif quant_type == "w8a16":
|
||||
max_q_val = 255
|
||||
min_q_val = 0
|
||||
elif quant_type == "w8a16b128":
|
||||
max_q_val = 127
|
||||
min_q_val = -128
|
||||
|
||||
w_s = torch.Tensor([1.0]).to(w.device) # unscaled case
|
||||
maybe_w_zp = None
|
||||
if group_size is not None:
|
||||
if zero_points:
|
||||
w_s = (max_val - min_val).clamp(min=1e-5) / max_q_val
|
||||
maybe_w_zp = (
|
||||
torch.round(torch.abs(min_val / w_s)).clamp(min_q_val, max_q_val).int()
|
||||
)
|
||||
else:
|
||||
# If the bias is such that there are no possible negative/positive
|
||||
# values, set the max value to inf to avoid divide by 0
|
||||
w_s = torch.max(
|
||||
abs(max_val / (max_q_val if max_q_val != 0 else torch.inf)),
|
||||
abs(min_val / (min_q_val if min_q_val != 0 else torch.inf)),
|
||||
)
|
||||
|
||||
# Quantize
|
||||
w_q = torch.round(w / w_s).int() + (maybe_w_zp if zero_points else 0)
|
||||
w_q = torch.clamp(w_q, min_q_val, max_q_val)
|
||||
|
||||
# Compute ref (dequantized)
|
||||
# For some kernels (namely Machete) the zero-points are applied after the
|
||||
# scales are applied, for this case computing the reference in similar way
|
||||
# allows us to use tighter error tolerances in our unit tests.
|
||||
if ref_zero_points_after_scales and maybe_w_zp is not None:
|
||||
w_ref = w_q.to(orig_type) * w_s - maybe_w_zp.to(orig_type) * w_s
|
||||
else:
|
||||
w_ref = (w_q - (maybe_w_zp if zero_points else 0)).to(orig_type) * w_s
|
||||
|
||||
if quant_type == "w4a16b8":
|
||||
w_q += 8
|
||||
elif quant_type == "w8a16b128":
|
||||
w_q += 128
|
||||
|
||||
# Restore original shapes
|
||||
if group_size is not None and group_size < size_k:
|
||||
|
||||
def reshape_w(w):
|
||||
w = w.reshape((group_size, -1, size_n))
|
||||
w = w.permute(1, 0, 2)
|
||||
w = w.reshape((size_k, size_n)).contiguous()
|
||||
return w
|
||||
|
||||
w_q = reshape_w(w_q)
|
||||
w_ref = reshape_w(w_ref)
|
||||
w_s = w_s.reshape((-1, size_n)).contiguous()
|
||||
|
||||
if maybe_w_zp is not None:
|
||||
maybe_w_zp = maybe_w_zp.reshape((-1, size_n)).contiguous()
|
||||
maybe_w_zp = maybe_w_zp.to(device=orig_device)
|
||||
|
||||
return (
|
||||
w_ref.to(device=orig_device),
|
||||
w_q.to(device=orig_device),
|
||||
w_s if group_size is not None else None,
|
||||
maybe_w_zp,
|
||||
)
|
||||
|
||||
|
||||
def torch_moe(a, w1, w2, score, topk):
|
||||
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
|
||||
|
||||
B, D = a.shape
|
||||
a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
|
||||
out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
|
||||
score = torch.softmax(score, dim=-1, dtype=torch.float32)
|
||||
topk_weight, topk_ids = torch.topk(score, topk)
|
||||
topk_weight = topk_weight.view(-1)
|
||||
topk_ids = topk_ids.view(-1)
|
||||
for i in range(w1.shape[0]):
|
||||
mask = topk_ids == i
|
||||
if mask.sum():
|
||||
out[mask] = SiluAndMul()(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(
|
||||
0, 1
|
||||
)
|
||||
return (
|
||||
out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
|
||||
).sum(dim=1)
|
||||
|
||||
|
||||
# fork from https://github.com/vllm-project/vllm/blob/main/tests/kernels/test_moe.py
|
||||
@pytest.mark.parametrize("m", [1, 32, 222])
|
||||
@pytest.mark.parametrize("n", [128, 1024, 2048])
|
||||
@pytest.mark.parametrize("k", [128, 1024])
|
||||
@pytest.mark.parametrize("e", NUM_EXPERTS)
|
||||
@pytest.mark.parametrize("topk", TOP_KS)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("group_size", [64, 128])
|
||||
@pytest.mark.parametrize("has_zp", [True, False])
|
||||
@pytest.mark.parametrize("weight_bits", [8]) # [4, 8])
|
||||
def test_fused_moe_wn16(
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
e: int,
|
||||
topk: int,
|
||||
dtype: torch.dtype,
|
||||
group_size: int,
|
||||
has_zp: bool,
|
||||
weight_bits: int,
|
||||
):
|
||||
print(m, n, k, e, topk, dtype, group_size, has_zp, weight_bits)
|
||||
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
|
||||
w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
|
||||
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
|
||||
score = torch.randn((m, e), device="cuda", dtype=dtype)
|
||||
|
||||
if weight_bits == 4:
|
||||
pack_factor = 2
|
||||
quant_type = "w4a16" if has_zp else "w4a16b8"
|
||||
elif weight_bits == 8:
|
||||
pack_factor = 1
|
||||
quant_type = "w8a16" if has_zp else "w8a16b128"
|
||||
|
||||
w1_ref = w1.clone()
|
||||
w2_ref = w2.clone()
|
||||
w1_qweight = torch.empty(
|
||||
(e, 2 * n, k // pack_factor), device="cuda", dtype=torch.uint8
|
||||
)
|
||||
w2_qweight = torch.empty((e, k, n // pack_factor), device="cuda", dtype=torch.uint8)
|
||||
w1_scales = torch.empty((e, 2 * n, k // group_size), device="cuda", dtype=dtype)
|
||||
w2_scales = torch.empty((e, k, n // group_size), device="cuda", dtype=dtype)
|
||||
w1_qzeros = torch.empty(
|
||||
(e, 2 * n // pack_factor, k // group_size), device="cuda", dtype=torch.uint8
|
||||
)
|
||||
w2_qzeros = torch.empty(
|
||||
(e, k // pack_factor, n // group_size), device="cuda", dtype=torch.uint8
|
||||
)
|
||||
|
||||
for i in range(e * 2):
|
||||
expert_id = i % e
|
||||
if i // e == 0:
|
||||
w, w_ref, w_qweight, w_scales, w_qzeros = (
|
||||
w1,
|
||||
w1_ref,
|
||||
w1_qweight,
|
||||
w1_scales,
|
||||
w1_qzeros,
|
||||
)
|
||||
else:
|
||||
w, w_ref, w_qweight, w_scales, w_qzeros = (
|
||||
w2,
|
||||
w2_ref,
|
||||
w2_qweight,
|
||||
w2_scales,
|
||||
w2_qzeros,
|
||||
)
|
||||
weight, qweight, scales, qzeros = quantize_weights(
|
||||
w[expert_id].T, quant_type, group_size, has_zp, False
|
||||
)
|
||||
weight = weight.T
|
||||
qweight = qweight.T.contiguous().to(torch.uint8)
|
||||
scales = scales.T
|
||||
if has_zp:
|
||||
qzeros = qzeros.T.contiguous().to(torch.uint8)
|
||||
if weight_bits == 4:
|
||||
qweight = qweight[:, 1::2] * 16 + qweight[:, ::2]
|
||||
if has_zp:
|
||||
qzeros = qzeros[1::2, :] * 16 + qzeros[::2, :]
|
||||
|
||||
w_ref[expert_id] = weight
|
||||
w_qweight[expert_id] = qweight
|
||||
w_scales[expert_id] = scales
|
||||
if has_zp:
|
||||
w_qzeros[expert_id] = qzeros
|
||||
|
||||
topk_output = select_experts(
|
||||
hidden_states=a,
|
||||
router_logits=score,
|
||||
topk_config=TopKConfig(top_k=topk),
|
||||
)
|
||||
|
||||
triton_output = fused_moe(
|
||||
a,
|
||||
w1_qweight,
|
||||
w2_qweight,
|
||||
topk_output,
|
||||
use_int4_w4a16=weight_bits == 4,
|
||||
use_int8_w8a16=weight_bits == 8,
|
||||
w1_scale=w1_scales,
|
||||
w2_scale=w2_scales,
|
||||
w1_zp=w1_qzeros if has_zp else None,
|
||||
w2_zp=w2_qzeros if has_zp else None,
|
||||
block_shape=[0, group_size],
|
||||
)
|
||||
torch_output = torch_moe(a, w1_ref, w2_ref, score, topk)
|
||||
torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
|
||||
@@ -0,0 +1,152 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardMode
|
||||
from sglang.srt.two_batch_overlap import (
|
||||
compute_split_seq_index,
|
||||
compute_split_token_index,
|
||||
)
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_ENABLE_THINKING_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestTwoBatchOverlap(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
with envs.SGLANG_ENABLE_JIT_DEEPGEMM.override(False):
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"2",
|
||||
"--dp",
|
||||
"2",
|
||||
"--enable-dp-attention",
|
||||
"--moe-a2a-backend",
|
||||
"deepep",
|
||||
"--deepep-mode",
|
||||
"normal",
|
||||
"--disable-cuda-graph", # DeepEP normal does not support CUDA Graph
|
||||
"--enable-two-batch-overlap",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_generate_single_prompt(self):
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
# we use an uncommon start to minimise the chance that the cache is hit by chance
|
||||
json={
|
||||
"text": "_ 1+1=2, 1+2=3, 1+3=4, 1+4=",
|
||||
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
|
||||
},
|
||||
)
|
||||
print(f"{response.json()=}")
|
||||
self.assertEqual(response.json()["text"], "5, 1+5=6")
|
||||
|
||||
def test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreater(metrics["score"], 0.5)
|
||||
|
||||
|
||||
class TestTwoBatchOverlapUnitTest(unittest.TestCase):
|
||||
def test_compute_split_seq_and_token_index(self):
|
||||
for num_tokens, expect in [
|
||||
(0, 0),
|
||||
(100, 50),
|
||||
(99, 49),
|
||||
]:
|
||||
actual = compute_split_seq_index(
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
num_tokens=num_tokens,
|
||||
extend_lens=None,
|
||||
token_num_per_seq=1,
|
||||
)
|
||||
self.assertEqual(actual, expect)
|
||||
|
||||
for extend_lens, expect in [
|
||||
([], (0, 0)),
|
||||
([42], (0, 21)),
|
||||
([42, 999], (1, 520)),
|
||||
([999, 42], (0, 520)),
|
||||
([498, 502], (1, 498)),
|
||||
([4096, 4096, 4096, 4096], (2, 8192)),
|
||||
([4095, 4096, 4096, 4096, 1], (2, 8191)),
|
||||
([1, 4095, 4096, 4096, 4096], (3, 8192)),
|
||||
([4097, 4096, 4096, 4095, 1], (2, 8193)),
|
||||
([1, 1, 1, 1, 99999], (4, 50001)),
|
||||
([99999, 1, 1, 1, 1], (0, 50001)),
|
||||
]:
|
||||
actual_seq_idx = compute_split_seq_index(
|
||||
forward_mode=ForwardMode.EXTEND,
|
||||
num_tokens=None,
|
||||
extend_lens=extend_lens,
|
||||
token_num_per_seq=None,
|
||||
)
|
||||
actual_token_idx = compute_split_token_index(
|
||||
split_seq_index=actual_seq_idx,
|
||||
forward_mode=ForwardMode.EXTEND,
|
||||
extend_seq_lens=extend_lens,
|
||||
token_num_per_seq=None,
|
||||
)
|
||||
actual = (actual_seq_idx, actual_token_idx)
|
||||
print(f"{extend_lens=} {expect=} {actual=}")
|
||||
self.assertEqual(actual, expect)
|
||||
|
||||
|
||||
class TestQwen3TwoBatchOverlap(TestTwoBatchOverlap):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_ENABLE_THINKING_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.api_key = "sk-1234"
|
||||
with envs.SGLANG_ENABLE_JIT_DEEPGEMM.override(False):
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"2",
|
||||
"--dp",
|
||||
"2",
|
||||
"--enable-dp-attention",
|
||||
"--moe-a2a-backend",
|
||||
"deepep",
|
||||
"--deepep-mode",
|
||||
"normal",
|
||||
"--disable-cuda-graph", # DeepEP normal does not support CUDA Graph
|
||||
"--enable-two-batch-overlap",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
python3 -m unittest test_vertex_endpoint.TestVertexEndpoint.test_vertex_generate
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from http import HTTPStatus
|
||||
|
||||
import requests
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class TestVertexEndpoint(CustomTestCase):
|
||||
@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=["--cuda-graph-max-bs", 2],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def run_generate(self, parameters):
|
||||
data = {
|
||||
"instances": [
|
||||
{"text": "The capital of France is"},
|
||||
{"text": "The capital of China is"},
|
||||
],
|
||||
"parameters": parameters,
|
||||
}
|
||||
response = requests.post(self.base_url + "/vertex_generate", json=data)
|
||||
response_json = response.json()
|
||||
assert len(response_json["predictions"]) == len(data["instances"])
|
||||
return response_json
|
||||
|
||||
def test_vertex_generate(self):
|
||||
for parameters in [None, {"sampling_params": {"max_new_tokens": 4}}]:
|
||||
self.run_generate(parameters)
|
||||
|
||||
def test_vertex_generate_fail(self):
|
||||
data = {
|
||||
"instances": [
|
||||
{"prompt": "The capital of France is"},
|
||||
],
|
||||
}
|
||||
response = requests.post(self.base_url + "/vertex_generate", json=data)
|
||||
assert response.status_code == HTTPStatus.BAD_REQUEST
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,321 @@
|
||||
"""
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from transformers import AutoModel, AutoProcessor, AutoTokenizer
|
||||
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
|
||||
from sglang.srt.managers.mm_utils import embed_mm_inputs, init_mm_embedding_cache
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
Modality,
|
||||
MultimodalDataItem,
|
||||
MultimodalInputs,
|
||||
)
|
||||
from sglang.srt.model_executor.model_runner import ModelRunner
|
||||
from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
|
||||
from sglang.srt.parser.conversation import generate_chat_conv
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.test_utils import download_image_with_retry
|
||||
|
||||
|
||||
# Test the logits output between HF and SGLang
|
||||
class VisionLLMLogitsBase(unittest.IsolatedAsyncioTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.image_url = "https://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true"
|
||||
cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
cls.model_path = ""
|
||||
cls.chat_template = ""
|
||||
cls.processor = ""
|
||||
cls.main_image = download_image_with_retry(cls.image_url)
|
||||
|
||||
def compare_outputs(self, sglang_output: torch.Tensor, hf_output: torch.Tensor):
|
||||
# Convert to float32 for numerical stability if needed
|
||||
hf = hf_output.float()
|
||||
sg = sglang_output.float()
|
||||
|
||||
# Basic shape and dtype comparison
|
||||
print("\n=== Basic Properties ===")
|
||||
print(f"Shapes match: {hf.shape == sg.shape}")
|
||||
print(f"HF shape: {hf.shape}, SGLang shape: {sg.shape}")
|
||||
print(f"HF dtype: {hf.dtype}, SGLang dtype: {sg.dtype}")
|
||||
|
||||
# Move tensors to CPU for numpy operations
|
||||
hf_np = hf.cpu().numpy()
|
||||
sg_np = sg.cpu().numpy()
|
||||
|
||||
# Statistical metrics
|
||||
print("\n=== Statistical Metrics ===")
|
||||
print(f"Mean absolute difference: {torch.mean(torch.abs(hf - sg)).item():.6f}")
|
||||
print(f"Max absolute difference: {torch.max(torch.abs(hf - sg)).item():.6f}")
|
||||
print(f"Mean squared error: {torch.mean((hf - sg) ** 2).item():.6f}")
|
||||
print(
|
||||
f"Root mean squared error: {torch.sqrt(torch.mean((hf - sg) ** 2)).item():.6f}"
|
||||
)
|
||||
|
||||
# Cosine similarity (across feature dimension)
|
||||
cos_sim = F.cosine_similarity(hf, sg)
|
||||
print(f"Mean cosine similarity: {torch.mean(cos_sim).item():.6f}")
|
||||
print(f"Min cosine similarity: {torch.min(cos_sim).item():.6f}")
|
||||
|
||||
# Find largest absolute differences
|
||||
print("\n=== Largest Absolute Differences ===")
|
||||
diffs = torch.abs(hf - sg)
|
||||
flat_diffs = diffs.flatten()
|
||||
|
||||
# Get indices of top 10 differences
|
||||
top_k = 10
|
||||
top_values, top_flat_indices = torch.topk(flat_diffs, top_k)
|
||||
|
||||
# Convert flat indices to multidimensional indices
|
||||
top_indices = np.unravel_index(top_flat_indices.cpu().numpy(), diffs.shape)
|
||||
|
||||
print(f"\nTop {top_k} largest absolute differences:")
|
||||
print(
|
||||
"Index".ljust(30)
|
||||
+ "Difference".ljust(15)
|
||||
+ "HF Value".ljust(15)
|
||||
+ "SGLang Value"
|
||||
)
|
||||
print("-" * 75)
|
||||
|
||||
for i in range(top_k):
|
||||
# Get the index tuple for this difference
|
||||
idx = tuple(dim[i] for dim in top_indices)
|
||||
diff_val = top_values[i].item()
|
||||
hf_val = hf[idx].item()
|
||||
sg_val = sg[idx].item()
|
||||
|
||||
# Format the index tuple and values
|
||||
idx_str = str(idx)
|
||||
print(f"{idx_str:<30}{diff_val:<15.6f}{hf_val:<15.6f}{sg_val:.6f}")
|
||||
|
||||
np.testing.assert_allclose(hf_np, sg_np)
|
||||
|
||||
def get_completion_request(self) -> ChatCompletionRequest:
|
||||
json_str = f"""
|
||||
{{
|
||||
"model": "{self.model_path}",
|
||||
"messages": [
|
||||
{{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{{
|
||||
"type": "image_url",
|
||||
"image_url": {{
|
||||
"url": "{self.image_url}"
|
||||
}}
|
||||
}},
|
||||
{{
|
||||
"type": "text",
|
||||
"text": "What's in this picture?"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
]
|
||||
}}
|
||||
"""
|
||||
|
||||
return ChatCompletionRequest.model_validate_json(json_str)
|
||||
|
||||
def get_processor_output(self, req: Optional[ChatCompletionRequest] = None):
|
||||
if req is None:
|
||||
req = self.get_completion_request()
|
||||
conv = generate_chat_conv(req, template_name=self.chat_template)
|
||||
text = conv.get_prompt()
|
||||
|
||||
# Process inputs using processor
|
||||
# FIXME: the formal arguments may differ
|
||||
inputs = self.processor(
|
||||
text=[text],
|
||||
images=[self.main_image],
|
||||
return_tensors="pt",
|
||||
).to(self.device)
|
||||
|
||||
return inputs
|
||||
|
||||
def get_sglang_model(self):
|
||||
self.model_runner = ModelRunner(
|
||||
model_config=ModelConfig(self.model_path, model_override_args="{}"),
|
||||
mem_fraction_static=0.8,
|
||||
gpu_id=0,
|
||||
tp_rank=0,
|
||||
tp_size=1,
|
||||
pp_rank=0,
|
||||
pp_size=1,
|
||||
nccl_port=12435,
|
||||
server_args=ServerArgs(
|
||||
model_path=self.model_path,
|
||||
disable_cuda_graph=True,
|
||||
),
|
||||
)
|
||||
return self.model_runner.model
|
||||
|
||||
|
||||
class TestMiniCPMV2_6Logits(VisionLLMLogitsBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model_path = "openbmb/MiniCPM-V-2_6"
|
||||
cls.tokenizer = AutoTokenizer.from_pretrained(
|
||||
cls.model_path, trust_remote_code=True
|
||||
)
|
||||
cls.processor = AutoProcessor.from_pretrained(
|
||||
cls.model_path, trust_remote_code=True
|
||||
)
|
||||
cls.chat_template = "minicpmv"
|
||||
|
||||
cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
cls.hf_model = (
|
||||
AutoModel.from_pretrained(
|
||||
cls.model_path, torch_dtype=torch.bfloat16, trust_remote_code=True
|
||||
)
|
||||
.eval()
|
||||
.to(cls.device)
|
||||
)
|
||||
init_mm_embedding_cache()
|
||||
|
||||
async def test_vlm_embedding_output(self):
|
||||
"""
|
||||
Compares the embedding output of vlm
|
||||
"""
|
||||
inputs = self.get_processor_output()
|
||||
|
||||
with torch.no_grad():
|
||||
# hf
|
||||
model_inputs = {
|
||||
"input_ids": inputs.input_ids,
|
||||
"image_bound": inputs.image_bound,
|
||||
"pixel_values": inputs.pixel_values,
|
||||
"tgt_sizes": inputs.tgt_sizes,
|
||||
}
|
||||
(hf_output, _) = self.hf_model.get_vllm_embedding(
|
||||
model_inputs,
|
||||
)
|
||||
hf_output = hf_output.squeeze(0)
|
||||
|
||||
# sglang
|
||||
model = self.get_sglang_model()
|
||||
input_ids = inputs["input_ids"].to(self.device).flatten()
|
||||
|
||||
pixel_values = inputs["pixel_values"]
|
||||
tgt_sizes = inputs["tgt_sizes"]
|
||||
pixel_values_flat: List[torch.Tensor] = []
|
||||
tgt_sizes_flat: List[torch.Tensor] = []
|
||||
for pixel_b, tgt_b in zip(pixel_values, tgt_sizes):
|
||||
# per image
|
||||
if len(pixel_b) != len(tgt_b):
|
||||
raise ValueError(
|
||||
"Inconsistent N lengths, found: "
|
||||
f"{len(pixel_b)} vs {len(tgt_b)}"
|
||||
)
|
||||
for pixel_n, tgt_n in zip(pixel_b, tgt_b):
|
||||
pixel_values_flat += [pixel_n]
|
||||
tgt_sizes_flat += [tgt_n]
|
||||
|
||||
im_start_id, im_end_id = (
|
||||
self.tokenizer.im_start_id,
|
||||
self.tokenizer.im_end_id,
|
||||
)
|
||||
slice_start_id, slice_end_id = (
|
||||
self.tokenizer.slice_start_id,
|
||||
self.tokenizer.slice_end_id,
|
||||
)
|
||||
|
||||
image_offsets = BaseMultimodalProcessor.get_mm_items_offset_by_pair(
|
||||
input_ids=input_ids, mm_start_id=im_start_id, mm_end_id=im_end_id
|
||||
)
|
||||
slice_offsets = BaseMultimodalProcessor.get_mm_items_offset_by_pair(
|
||||
input_ids=input_ids, mm_start_id=slice_start_id, mm_end_id=slice_end_id
|
||||
)
|
||||
image_offsets.extend(slice_offsets)
|
||||
image_offsets = sorted(image_offsets)
|
||||
|
||||
sglang_output = embed_mm_inputs(
|
||||
mm_inputs_list=[
|
||||
MultimodalInputs(
|
||||
mm_items=[
|
||||
MultimodalDataItem(
|
||||
feature=pixel_values_flat,
|
||||
offsets=image_offsets,
|
||||
tgt_size=tgt_sizes_flat,
|
||||
modality=Modality.IMAGE,
|
||||
pad_value=self.processor.tokenizer.unk_token_id,
|
||||
)
|
||||
]
|
||||
),
|
||||
],
|
||||
extend_prefix_lens=[0],
|
||||
extend_seq_lens=[input_ids.shape[0]],
|
||||
input_ids=input_ids,
|
||||
input_embedding=model.get_input_embeddings(),
|
||||
multimodal_model=model,
|
||||
placeholder_tokens={
|
||||
Modality.IMAGE: self.processor.tokenizer.unk_token_id,
|
||||
},
|
||||
)
|
||||
|
||||
self.compare_outputs(sglang_output, hf_output)
|
||||
|
||||
|
||||
class TestMiniCPMV4Logits(VisionLLMLogitsBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model_path = "openbmb/MiniCPM-V-4"
|
||||
cls.tokenizer = AutoTokenizer.from_pretrained(
|
||||
cls.model_path, trust_remote_code=True
|
||||
)
|
||||
cls.processor = AutoProcessor.from_pretrained(
|
||||
cls.model_path, trust_remote_code=True
|
||||
)
|
||||
cls.chat_template = "minicpmv"
|
||||
|
||||
cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
cls.hf_model = (
|
||||
AutoModel.from_pretrained(
|
||||
cls.model_path, torch_dtype=torch.bfloat16, trust_remote_code=True
|
||||
)
|
||||
.eval()
|
||||
.to(cls.device)
|
||||
)
|
||||
init_mm_embedding_cache()
|
||||
|
||||
async def test_vlm_embedding_output(self):
|
||||
"""
|
||||
Compares the embedding output of vlm
|
||||
"""
|
||||
inputs = self.get_processor_output()
|
||||
|
||||
with torch.no_grad():
|
||||
# hf
|
||||
model_inputs = {
|
||||
"input_ids": inputs.input_ids,
|
||||
"image_bound": inputs.image_bound,
|
||||
"pixel_values": inputs.pixel_values,
|
||||
"tgt_sizes": inputs.tgt_sizes,
|
||||
}
|
||||
hf_output = self.hf_model.get_input_embeddings()(inputs.input_ids)
|
||||
|
||||
# sglang
|
||||
model = self.get_model()
|
||||
sglang_output = self.vlm_func(
|
||||
model,
|
||||
input_ids=inputs.input_ids.to(self.device),
|
||||
pixel_values=inputs.pixel_values,
|
||||
image_bound=inputs.image_bound.to(self.device),
|
||||
tgt_sizes=inputs.tgt_sizes.to(self.device),
|
||||
input_embedding=model.get_input_embeddings(),
|
||||
multimodal_model=model,
|
||||
placeholder_tokens={
|
||||
Modality.IMAGE: self.processor.tokenizer.unk_token_id,
|
||||
},
|
||||
)
|
||||
|
||||
self.compare_outputs(sglang_output, hf_output)
|
||||
@@ -0,0 +1,61 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest test_wave_attention_backend.TestWaveAttnBackend.test_mmlu
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
run_bench_one_batch,
|
||||
)
|
||||
|
||||
|
||||
class TestWaveAttnBackend(unittest.TestCase):
|
||||
def test_latency(self):
|
||||
_, output_throughput, _ = run_bench_one_batch(
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
[
|
||||
"--attention-backend",
|
||||
"wave",
|
||||
"--enable-torch-compile",
|
||||
],
|
||||
)
|
||||
|
||||
if is_in_ci():
|
||||
self.assertGreater(output_throughput, 153)
|
||||
|
||||
def _test_mmlu(self):
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
process = popen_launch_server(
|
||||
model,
|
||||
base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=["--attention-backend", "wave"],
|
||||
)
|
||||
|
||||
try:
|
||||
args = SimpleNamespace(
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreaterEqual(metrics["score"], 0.65)
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,227 @@
|
||||
"""
|
||||
Test weight version functionality.
|
||||
|
||||
This test suite verifies the weight_version feature implementation including:
|
||||
1. Default weight_version setting
|
||||
2. /get_weight_version endpoint
|
||||
3. /update_weight_version endpoint
|
||||
4. /generate request meta_info contains weight_version
|
||||
5. OpenAI API response metadata contains weight_version
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestWeightVersion(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Start server once for all tests with custom weight version."""
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = "http://127.0.0.1:30000"
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
base_url=cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--weight-version",
|
||||
"test_version_1.0",
|
||||
"--attention-backend",
|
||||
"flashinfer",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""Terminate server after all tests complete."""
|
||||
if cls.process:
|
||||
cls.process.terminate()
|
||||
|
||||
def test_weight_version_comprehensive(self):
|
||||
"""Comprehensive test for all weight_version functionality."""
|
||||
|
||||
response = requests.get(f"{self.base_url}/get_model_info")
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertIn("weight_version", data)
|
||||
self.assertEqual(data["weight_version"], "test_version_1.0")
|
||||
|
||||
response = requests.get(f"{self.base_url}/get_weight_version")
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertIn("weight_version", data)
|
||||
self.assertEqual(data["weight_version"], "test_version_1.0")
|
||||
|
||||
request_data = {
|
||||
"text": "Hello, how are you?",
|
||||
"sampling_params": {
|
||||
"temperature": 0.0,
|
||||
"max_new_tokens": 5,
|
||||
},
|
||||
}
|
||||
response = requests.post(f"{self.base_url}/generate", json=request_data)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertIn("meta_info", data)
|
||||
self.assertIn("weight_version", data["meta_info"])
|
||||
self.assertEqual(data["meta_info"]["weight_version"], "test_version_1.0")
|
||||
|
||||
request_data = {
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
response = requests.post(
|
||||
f"{self.base_url}/v1/chat/completions", json=request_data
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertIn("metadata", data)
|
||||
self.assertIn("weight_version", data["metadata"])
|
||||
self.assertEqual(data["metadata"]["weight_version"], "test_version_1.0")
|
||||
|
||||
request_data = {
|
||||
"model": self.model,
|
||||
"prompt": "Hello",
|
||||
"max_tokens": 5,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
response = requests.post(f"{self.base_url}/v1/completions", json=request_data)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertIn("metadata", data)
|
||||
self.assertIn("weight_version", data["metadata"])
|
||||
self.assertEqual(data["metadata"]["weight_version"], "test_version_1.0")
|
||||
|
||||
update_data = {
|
||||
"new_version": "updated_version_2.0",
|
||||
"abort_all_requests": False,
|
||||
}
|
||||
response = requests.post(
|
||||
f"{self.base_url}/update_weight_version", json=update_data
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual(data["new_version"], "updated_version_2.0")
|
||||
|
||||
response = requests.get(f"{self.base_url}/get_weight_version")
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data["weight_version"], "updated_version_2.0")
|
||||
|
||||
gen_data = {
|
||||
"text": "Test persistence",
|
||||
"sampling_params": {"temperature": 0.0, "max_new_tokens": 3},
|
||||
}
|
||||
response = requests.post(f"{self.base_url}/generate", json=gen_data)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data["meta_info"]["weight_version"], "updated_version_2.0")
|
||||
|
||||
chat_data = {
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": "Test"}],
|
||||
"max_tokens": 3,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
response = requests.post(f"{self.base_url}/v1/chat/completions", json=chat_data)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data["metadata"]["weight_version"], "updated_version_2.0")
|
||||
|
||||
update_data = {"new_version": "final_version_3.0", "abort_all_requests": True}
|
||||
response = requests.post(
|
||||
f"{self.base_url}/update_weight_version", json=update_data
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual(data["new_version"], "final_version_3.0")
|
||||
|
||||
# Check /get_weight_version
|
||||
response = requests.get(f"{self.base_url}/get_weight_version")
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(response.json()["weight_version"], "final_version_3.0")
|
||||
|
||||
# Check /get_model_info
|
||||
response = requests.get(f"{self.base_url}/get_model_info")
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(response.json()["weight_version"], "final_version_3.0")
|
||||
|
||||
# Check /generate meta_info
|
||||
response = requests.post(
|
||||
f"{self.base_url}/generate",
|
||||
json={
|
||||
"text": "Final test",
|
||||
"sampling_params": {"temperature": 0.0, "max_new_tokens": 2},
|
||||
},
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(
|
||||
response.json()["meta_info"]["weight_version"], "final_version_3.0"
|
||||
)
|
||||
|
||||
# Check OpenAI chat metadata
|
||||
response = requests.post(
|
||||
f"{self.base_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": "Final"}],
|
||||
"max_tokens": 2,
|
||||
"temperature": 0.0,
|
||||
},
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(
|
||||
response.json()["metadata"]["weight_version"], "final_version_3.0"
|
||||
)
|
||||
|
||||
print("All weight_version functionality tests passed!")
|
||||
|
||||
def test_update_weight_version_with_weight_updates(self):
|
||||
"""Test that weight_version can be updated along with weight updates using real model data."""
|
||||
print("Testing weight_version update with real weight operations...")
|
||||
|
||||
# Get current model info for reference
|
||||
model_info_response = requests.get(f"{self.base_url}/get_model_info")
|
||||
self.assertEqual(model_info_response.status_code, 200)
|
||||
current_model_path = model_info_response.json()["model_path"]
|
||||
|
||||
update_data = {
|
||||
"model_path": current_model_path,
|
||||
"load_format": "auto",
|
||||
"abort_all_requests": False,
|
||||
"weight_version": "disk_update_v2.0.0",
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
f"{self.base_url}/update_weights_from_disk", json=update_data
|
||||
)
|
||||
self.assertEqual(
|
||||
response.status_code,
|
||||
200,
|
||||
f"update_weights_from_disk failed with status {response.status_code}",
|
||||
)
|
||||
|
||||
# Verify version was updated
|
||||
version_response = requests.get(f"{self.base_url}/get_weight_version")
|
||||
self.assertEqual(version_response.status_code, 200)
|
||||
self.assertEqual(
|
||||
version_response.json()["weight_version"], "disk_update_v2.0.0"
|
||||
)
|
||||
|
||||
print("Weight update with weight_version test completed!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user