Test reorganization: Move tests to manual/ (#13610)

This commit is contained in:
alisonshao
2025-11-20 13:41:58 -08:00
committed by GitHub
parent ada8ce1fd0
commit 6b262ac839
74 changed files with 0 additions and 74 deletions
@@ -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()
+116
View File
@@ -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.")
+191
View File
@@ -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()
+108
View File
@@ -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")
+63
View File
@@ -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()
+229
View File
@@ -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()
+78
View File
@@ -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()
+146
View File
@@ -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()
+53
View File
@@ -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()
+73
View File
@@ -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()
+58
View File
@@ -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()
+213
View File
@@ -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()
+127
View File
@@ -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()
+88
View File
@@ -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()
+114
View File
@@ -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()
+276
View File
@@ -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()
+290
View File
@@ -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__])
+364
View 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__])
+158
View 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)
+182
View 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()
+318
View File
@@ -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()
+65
View File
@@ -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()
+101
View File
@@ -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()
+255
View File
@@ -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()
+72
View File
@@ -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()
+305
View File
@@ -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()
+183
View File
@@ -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()
+28
View File
@@ -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()
+374
View File
@@ -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()
+526
View File
@@ -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())
+66
View File
@@ -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()
+58
View File
@@ -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,
)
+30
View File
@@ -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()
+196
View File
@@ -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()
+303
View File
@@ -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()
+183
View File
@@ -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()
+203
View File
@@ -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()
+786
View File
@@ -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()
+121
View File
@@ -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)
+386
View File
@@ -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()
+30
View File
@@ -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()
+272
View File
@@ -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()
+192
View File
@@ -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()
+246
View File
@@ -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)
+152
View File
@@ -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()
+64
View File
@@ -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()
+321
View File
@@ -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()
+227
View File
@@ -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()