chore: vendor sglang v0.5.10 snapshot

This commit is contained in:
2026-04-24 12:29:36 +00:00
parent 78f0d15221
commit bded08301f
4308 changed files with 1200894 additions and 2 deletions

View File

@@ -0,0 +1,94 @@
import unittest
from types import SimpleNamespace
from urllib.parse import urlparse
from sglang.srt.environ import envs
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_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
TEST_MODEL_MATRIX = {
"/root/.cache/modelscope/hub/models/vllm-ascend/DeepSeek-R1-0528-W8A8": {
"accuracy": 0.95,
"latency": 1000,
"output_throughput": 6,
},
}
class TestAscendDeepSeekMTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.models = TEST_MODEL_MATRIX.keys()
cls.base_url = DEFAULT_URL_FOR_TEST
cls.url = urlparse(DEFAULT_URL_FOR_TEST)
cls.common_args = [
"--trust-remote-code",
"--attention-backend",
"ascend",
"--mem-fraction-static",
0.8,
"--disable-radix-cache",
"--chunked-prefill-size",
32768,
"--tp-size",
16,
"--dp-size",
2,
"--enable-dp-attention",
"--speculative-algorithm",
"NEXTN",
"--speculative-num-steps",
1,
"--speculative-eagle-topk",
1,
"--speculative-num-draft-tokens",
2,
]
envs.SGLANG_NPU_USE_MLAPO.set(True)
envs.SGLANG_ENABLE_SPEC_V2.set(True)
envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.set(True)
def test_a_gsm8k(self):
for model in self.models:
with self.subTest(model=model):
print(f"##=== Testing accuracy: {model} ===##")
process = popen_launch_server(
model,
self.base_url,
timeout=2400,
other_args=[
*self.common_args,
],
)
try:
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=1319,
max_new_tokens=512,
parallel=128,
host=f"http://{self.url.hostname}",
port=int(self.url.port),
)
metrics = run_eval_few_shot_gsm8k(args)
self.assertGreaterEqual(
metrics["accuracy"],
TEST_MODEL_MATRIX[model]["accuracy"],
)
finally:
kill_process_tree(process.pid)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,102 @@
"""
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",
],
)
@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()

View File

@@ -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()

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()

View File

@@ -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()

File diff suppressed because one or more lines are too long

View File

@@ -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")

View File

@@ -0,0 +1,445 @@
# Copy from deepseek-ai/DeepEP/tests/test_internode.py
import os
import time
# noinspection PyUnresolvedReferences
import deep_ep
# Test compatibility with low latency functions
import test_deepep_low_latency
import torch
import torch.distributed as dist
from sglang.test.test_deepep_utils import (
bench,
calc_diff,
create_grouped_scores,
init_dist,
inplace_unique,
per_token_cast_back,
per_token_cast_to_fp8,
)
def test_main(
num_sms: int,
local_rank: int,
num_local_ranks: int,
num_ranks: int,
num_nodes: int,
rank: int,
buffer: deep_ep.Buffer,
group: dist.ProcessGroup,
):
# Settings
num_tokens, hidden, num_topk_groups, num_topk, num_experts = (
4096,
7168,
min(num_nodes, 4),
8,
(256 // num_ranks) * num_ranks,
)
assert num_experts % num_ranks == 0 and num_local_ranks == 8
if local_rank == 0:
print(
f"[config] num_tokens={num_tokens}, hidden={hidden}, num_topk_groups={num_topk_groups}, num_topk={num_topk}",
flush=True,
)
# Random data
x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * rank
x_pure_rand = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
x_e4m3 = per_token_cast_to_fp8(x)
scores = (
torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs()
+ 1
)
group_scores = scores.view(num_tokens, num_nodes, -1).amax(dim=-1)
group_idx = torch.topk(
group_scores, k=num_topk_groups, dim=-1, sorted=False
).indices
masked_scores = create_grouped_scores(scores, group_idx, num_nodes)
topk_idx = torch.topk(masked_scores, num_topk, dim=-1, largest=True, sorted=False)[
1
]
topk_weights = (
torch.ones((num_tokens, num_topk), dtype=torch.float32, device="cuda") * rank
)
topk_weights_pure_rand = torch.randn(
(num_tokens, num_topk), dtype=torch.float32, device="cuda"
)
rank_idx = topk_idx // (num_experts // num_ranks)
rank_idx.masked_fill_(topk_idx == -1, -1)
inplace_unique(rank_idx, num_ranks)
rdma_rank_idx = rank_idx // num_local_ranks
rdma_rank_idx.masked_fill_(rank_idx == -1, -1)
inplace_unique(rdma_rank_idx, num_nodes)
# RDMA dispatch counts
rdma_idx = topk_idx // (num_experts // num_nodes)
rdma_idx.masked_fill_(topk_idx == -1, -1)
inplace_unique(rdma_idx, num_nodes)
num_rdma_token_sent = rdma_idx.ne(-1).sum().item()
# Expert meta
num_tokens_per_expert = torch.zeros((num_experts,), dtype=torch.int, device="cuda")
for i in range(num_experts):
num_tokens_per_expert[i] = (topk_idx == i).sum()
gbl_num_tokens_per_expert = num_tokens_per_expert.clone()
dist.all_reduce(gbl_num_tokens_per_expert, group=group)
# Rank layout meta
num_tokens_per_rank = torch.empty((num_ranks,), dtype=torch.int, device="cuda")
num_tokens_per_rdma_rank = torch.empty((num_nodes,), dtype=torch.int, device="cuda")
token_idx_in_rank = torch.full(
(num_ranks, num_tokens), -1, dtype=torch.long, device="cuda"
)
for i in range(num_ranks):
num_tokens_per_rank[i] = (rank_idx == i).sum()
token_sel = (rank_idx == i).max(dim=-1)[0]
count = token_sel.sum().item()
tokens = torch.sort(token_sel.to(torch.int), descending=True)[1]
tokens[:count] = torch.sort(tokens[:count])[0]
token_idx_in_rank[i][tokens[:count]] = torch.arange(
count, dtype=torch.long, device="cuda"
)
for i in range(num_nodes):
num_tokens_per_rdma_rank[i] = (rdma_rank_idx == i).sum()
token_idx_in_rank = token_idx_in_rank.T.contiguous().to(torch.int)
is_token_in_rank = token_idx_in_rank >= 0
gbl_num_tokens_per_rank = num_tokens_per_rank.clone()
dist.all_reduce(gbl_num_tokens_per_rank, group=group)
(
ref_num_tokens_per_rank,
ref_num_tokens_per_rdma_rank,
ref_num_tokens_per_expert,
ref_is_token_in_rank,
_,
) = buffer.get_dispatch_layout(topk_idx, num_experts)
assert torch.allclose(ref_num_tokens_per_rank, num_tokens_per_rank)
assert torch.allclose(ref_num_tokens_per_rdma_rank, num_tokens_per_rdma_rank)
assert torch.allclose(ref_num_tokens_per_expert, num_tokens_per_expert)
assert torch.allclose(ref_is_token_in_rank, is_token_in_rank)
t = bench(lambda: buffer.get_dispatch_layout(topk_idx, num_experts))[0]
if local_rank == 0:
print(f"[layout] Kernel performance: {t * 1000:.3f} ms", flush=True)
print("", flush=True)
group.barrier()
time.sleep(1)
# Config
rdma_buffer_size, nvl_buffer_size = 128, (720 if num_ranks in (144, 160) else 512)
config = deep_ep.Config(num_sms, 8, nvl_buffer_size, 16, rdma_buffer_size)
# Test dispatch
# noinspection PyShadowingNames
def check_data(check_x, recv_gbl_rank_prefix_sum):
assert torch.allclose(check_x.amin(dim=1), check_x.amax(dim=1))
check_start = 0
for i in range(num_ranks):
check_end = recv_gbl_rank_prefix_sum[i].item()
assert (check_x[check_start:check_end, :].int() - i).sum().item() == 0
check_start = check_end
for previous_mode in (False, True):
for async_mode in (False, True):
for current_x in (x_pure_rand, x, x_e4m3):
for with_topk in (False, True):
if local_rank == 0:
print(
f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
flush=True,
end="",
)
dispatch_args = {
"x": current_x,
"num_tokens_per_rank": num_tokens_per_rank,
"num_tokens_per_rdma_rank": num_tokens_per_rdma_rank,
"is_token_in_rank": is_token_in_rank,
"num_tokens_per_expert": num_tokens_per_expert,
"config": config,
"async_finish": async_mode,
}
if with_topk:
dispatch_args.update(
{
"topk_idx": topk_idx,
"topk_weights": (
topk_weights_pure_rand
if current_x is x_pure_rand
else topk_weights
),
}
)
if previous_mode:
dispatch_args.update({"previous_event": buffer.capture()})
(
recv_x,
recv_topk_idx,
recv_topk_weights,
recv_num_tokens_per_expert_list,
handle,
event,
) = buffer.dispatch(**dispatch_args)
event.current_stream_wait() if async_mode else ()
recv_x = (
per_token_cast_back(*recv_x)
if isinstance(recv_x, tuple)
else recv_x
)
# Checks
recv_gbl_rank_prefix_sum = handle[-4]
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
0
), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
assert (
gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
== recv_num_tokens_per_expert_list
)
if current_x is not x_pure_rand:
check_data(recv_x, recv_gbl_rank_prefix_sum)
if with_topk:
# Check `topk_idx`
assert (
recv_topk_idx.eq(-1)
| (
(recv_topk_idx >= 0)
& (recv_topk_idx < (num_experts // num_ranks))
)
).sum().item() == recv_topk_idx.numel()
for i, count in enumerate(recv_num_tokens_per_expert_list):
assert recv_topk_idx.eq(i).sum().item() == count
# Check `topk_weights`
if current_x is not x_pure_rand:
recv_topk_weights[recv_topk_idx.eq(-1)] = (
recv_topk_weights.amax(dim=1, keepdim=True).expand_as(
recv_topk_weights
)[recv_topk_idx.eq(-1)]
)
check_data(recv_topk_weights, recv_gbl_rank_prefix_sum)
# Test cached dispatch (must without top-k staffs)
if not with_topk:
dispatch_args = {
"x": current_x,
"handle": handle,
"config": config,
"async_finish": async_mode,
}
if previous_mode:
dispatch_args.update({"previous_event": buffer.capture()})
recv_x, _, _, _, _, event = buffer.dispatch(**dispatch_args)
event.current_stream_wait() if async_mode else ()
recv_x = (
per_token_cast_back(*recv_x)
if isinstance(recv_x, tuple)
else recv_x
)
if current_x is not x_pure_rand:
check_data(recv_x, recv_gbl_rank_prefix_sum)
# Test combine
combine_args = {
"x": recv_x,
"handle": handle,
"config": config,
"async_finish": async_mode,
}
if with_topk:
combine_args.update({"topk_weights": recv_topk_weights})
if previous_mode:
dispatch_args.update({"previous_event": buffer.capture()})
combined_x, combined_topk_weights, event = buffer.combine(
**combine_args
)
event.current_stream_wait() if async_mode else ()
check_x = combined_x.float() / is_token_in_rank.sum(
dim=1
).unsqueeze(1)
ref_x = x_pure_rand if current_x is x_pure_rand else x
assert calc_diff(check_x, ref_x) < 5e-6
if with_topk:
check_topk_weights = (
combined_topk_weights
if (current_x is x_pure_rand)
else (
combined_topk_weights
/ is_token_in_rank.sum(dim=1).unsqueeze(1)
)
)
ref_topk_weights = (
topk_weights_pure_rand
if current_x is x_pure_rand
else topk_weights
)
assert calc_diff(check_topk_weights, ref_topk_weights) < 1e-9
# For later tuning
dispatch_bf16_rdma_send_bytes = num_rdma_token_sent * hidden * 2
dispatch_bf16_nvl_recv_bytes = recv_x.numel() * 2
combine_bf16_nvl_send_bytes = dispatch_bf16_nvl_recv_bytes
combine_bf16_rdma_recv_bytes = dispatch_bf16_rdma_send_bytes
if local_rank == 0:
print(" passed", flush=True)
if local_rank == 0:
print("", flush=True)
# Tune dispatch performance
best_dispatch_results = None
fp8_factor = (1 + 4 / 128) / 2
for current_x in (x_e4m3, x):
best_time, best_results = 1e10, None
rdma_send_bytes = (
(dispatch_bf16_rdma_send_bytes * fp8_factor)
if isinstance(current_x, tuple)
else dispatch_bf16_rdma_send_bytes
)
nvl_recv_bytes = (
(dispatch_bf16_nvl_recv_bytes * fp8_factor)
if isinstance(current_x, tuple)
else dispatch_bf16_nvl_recv_bytes
)
for nvl_chunk_size in range(4, 33, 4):
for rdma_chunk_size in range(4, 33, 4):
config = deep_ep.Config(
num_sms,
nvl_chunk_size,
nvl_buffer_size,
rdma_chunk_size,
rdma_buffer_size,
)
tune_args = {"x": current_x, "handle": handle, "config": config}
t = bench(lambda: buffer.dispatch(**tune_args))[0]
if t < best_time:
best_time, best_results = t, (
num_sms,
nvl_chunk_size,
rdma_chunk_size,
)
if local_rank == 0:
print(
f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}, RDMA chunk {rdma_chunk_size}: {rdma_send_bytes / 1e9 / t:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / t:.2f} GB/s (NVL) ",
flush=True,
)
if local_rank == 0:
print(
f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
flush=True,
)
print("", flush=True)
if isinstance(current_x, tuple):
# Gather FP8 the best config from rank 0
best_dispatch_results = torch.tensor(
[best_results[0], best_results[1], best_results[2]],
dtype=torch.int32,
device="cuda",
)
all_best_fp8_results_list = [
torch.zeros_like(best_dispatch_results)
for _ in range(torch.distributed.get_world_size())
]
dist.all_gather(
all_best_fp8_results_list, best_dispatch_results, group=group
)
best_dispatch_results = all_best_fp8_results_list[0].tolist()
dispatch_config = deep_ep.Config(
best_dispatch_results[0],
best_dispatch_results[1],
nvl_buffer_size,
best_dispatch_results[2],
rdma_buffer_size,
)
dispatch_args = {
"x": x,
"num_tokens_per_rank": num_tokens_per_rank,
"num_tokens_per_rdma_rank": num_tokens_per_rdma_rank,
"is_token_in_rank": is_token_in_rank,
"num_tokens_per_expert": num_tokens_per_expert,
"config": dispatch_config if dispatch_config is not None else config,
}
recv_x, _, _, _, handle, _ = buffer.dispatch(**dispatch_args)
# Tune combine performance
best_time, best_results = 1e10, None
for nvl_chunk_size in range(1, 5, 1):
for rdma_chunk_size in range(8, 33, 4):
config = deep_ep.Config(
num_sms,
nvl_chunk_size,
nvl_buffer_size,
rdma_chunk_size,
rdma_buffer_size,
)
tune_args = {"x": recv_x, "handle": handle, "config": config}
t = bench(lambda: buffer.combine(**tune_args))[0]
if local_rank == 0:
print(
f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}, RDMA chunk {rdma_chunk_size}: {combine_bf16_rdma_recv_bytes / 1e9 / t:.2f} GB/s (RDMA), {combine_bf16_nvl_send_bytes / 1e9 / t:.2f} GB/s (NVL) ",
flush=True,
)
if t < best_time:
best_time, best_results = t, (
num_sms,
nvl_chunk_size,
rdma_chunk_size,
)
if local_rank == 0:
print(
f"[tuning] Best combine: SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {combine_bf16_rdma_recv_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {combine_bf16_nvl_send_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
flush=True,
)
print("", flush=True)
# noinspection PyUnboundLocalVariable
def test_loop(local_rank: int, num_local_ranks: int):
num_nodes = int(os.getenv("WORLD_SIZE", 1))
rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
test_ll_compatibility = False
if test_ll_compatibility:
ll_num_tokens, ll_hidden, ll_num_experts, ll_num_topk = 16, 5120, 256, 9
buffer = deep_ep.Buffer(
group,
int(1e9),
int(1e9),
low_latency_mode=test_ll_compatibility,
num_qps_per_rank=(ll_num_experts // num_ranks if test_ll_compatibility else 1),
)
assert num_local_ranks == 8 and num_ranks > 8
torch.manual_seed(rank)
for i in (24,):
test_main(
i, local_rank, num_local_ranks, num_ranks, num_nodes, rank, buffer, group
)
if local_rank == 0:
print("", flush=True)
# Test compatibility with low latency functions
if test_ll_compatibility:
buffer.clean_low_latency_buffer(ll_num_tokens, ll_hidden, ll_num_experts)
test_deepep_low_latency.test_main(
ll_num_tokens,
ll_hidden,
ll_num_experts,
ll_num_topk,
rank,
num_ranks,
group,
buffer,
seed=1,
)
if __name__ == "__main__":
num_processes = 8
torch.multiprocessing.spawn(test_loop, args=(num_processes,), nprocs=num_processes)

View File

@@ -0,0 +1,378 @@
# Copy from deepseek-ai/DeepEP/tests/test_intranode.py
import time
# noinspection PyUnresolvedReferences
import deep_ep
# Test compatibility with low latency functions
import test_deepep_low_latency
import torch
import torch.distributed as dist
from sglang.test.test_deepep_utils import (
bench,
calc_diff,
init_dist,
inplace_unique,
per_token_cast_back,
per_token_cast_to_fp8,
)
def test_main(
num_sms: int,
local_rank: int,
num_ranks: int,
rank: int,
buffer: deep_ep.Buffer,
group: dist.ProcessGroup,
):
# Settings
num_tokens, hidden, num_topk, num_experts = (
4096,
7168,
8,
(256 // num_ranks) * num_ranks,
)
assert num_experts % num_ranks == 0
if local_rank == 0:
print(
f"[config] num_tokens={num_tokens}, hidden={hidden}, num_topk={num_topk}",
flush=True,
)
# Random data
x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * rank
x_pure_rand = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
x_e4m3 = per_token_cast_to_fp8(x)
scores = (
torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs()
+ 1
)
topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)[1]
topk_weights = (
torch.ones((num_tokens, num_topk), dtype=torch.float32, device="cuda") * rank
)
topk_weights_pure_rand = torch.randn(
(num_tokens, num_topk), dtype=torch.float32, device="cuda"
)
rank_idx = topk_idx // (num_experts // num_ranks)
rank_idx.masked_fill_(topk_idx == -1, -1)
inplace_unique(rank_idx, num_ranks)
# Expert meta
num_tokens_per_expert = torch.zeros((num_experts,), dtype=torch.int, device="cuda")
for i in range(num_experts):
num_tokens_per_expert[i] = (topk_idx == i).sum()
gbl_num_tokens_per_expert = num_tokens_per_expert.clone()
dist.all_reduce(gbl_num_tokens_per_expert, group=group)
# Rank layout meta
num_tokens_per_rank = torch.empty((num_ranks,), dtype=torch.int, device="cuda")
token_idx_in_rank = torch.full(
(num_ranks, num_tokens), -1, dtype=torch.long, device="cuda"
)
for i in range(num_ranks):
num_tokens_per_rank[i] = (rank_idx == i).sum()
token_sel = (rank_idx == i).max(dim=-1)[0]
count = token_sel.sum().item()
tokens = torch.sort(token_sel.to(torch.int), descending=True)[1]
tokens[:count] = torch.sort(tokens[:count])[0]
token_idx_in_rank[i][tokens[:count]] = torch.arange(
count, dtype=torch.long, device="cuda"
)
token_idx_in_rank = token_idx_in_rank.T.contiguous().to(torch.int)
is_token_in_rank = token_idx_in_rank >= 0
gbl_num_tokens_per_rank = num_tokens_per_rank.clone()
dist.all_reduce(gbl_num_tokens_per_rank, group=group)
ref_num_tokens_per_rank, _, ref_num_tokens_per_expert, ref_is_token_in_rank, _ = (
buffer.get_dispatch_layout(topk_idx, num_experts)
)
assert torch.allclose(ref_num_tokens_per_rank, num_tokens_per_rank)
assert torch.allclose(ref_num_tokens_per_expert, num_tokens_per_expert)
assert torch.allclose(ref_is_token_in_rank, is_token_in_rank)
t = bench(lambda: buffer.get_dispatch_layout(topk_idx, num_experts))[0]
if local_rank == 0:
print(f"[layout] Kernel performance: {t * 1000:.3f} ms", flush=True)
print("", flush=True)
group.barrier()
time.sleep(1)
# Config
nvl_buffer_size = 256
config = deep_ep.Config(num_sms, 8, nvl_buffer_size)
# Test dispatch
# noinspection PyShadowingNames
def check_data(check_x, rank_prefix_matrix):
assert torch.allclose(check_x.amin(dim=1), check_x.amax(dim=1))
check_start = 0
for i in range(num_ranks):
check_end = rank_prefix_matrix[i][rank].item()
assert (check_x[check_start:check_end, :].int() - i).sum().item() == 0
check_start = check_end
for previous_mode in (False, True):
for async_mode in (False, True):
for current_x in (x_pure_rand, x, x_e4m3):
for with_topk in (False, True):
if local_rank == 0:
print(
f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
flush=True,
end="",
)
dispatch_args = {
"x": current_x,
"num_tokens_per_rank": num_tokens_per_rank,
"is_token_in_rank": is_token_in_rank,
"num_tokens_per_expert": num_tokens_per_expert,
"config": config,
"async_finish": async_mode,
}
if with_topk:
dispatch_args.update(
{
"topk_idx": topk_idx,
"topk_weights": (
topk_weights_pure_rand
if current_x is x_pure_rand
else topk_weights
),
}
)
if previous_mode:
dispatch_args.update({"previous_event": buffer.capture()})
(
recv_x,
recv_topk_idx,
recv_topk_weights,
recv_num_tokens_per_expert_list,
handle,
event,
) = buffer.dispatch(**dispatch_args)
event.current_stream_wait() if async_mode else ()
recv_x = (
per_token_cast_back(*recv_x)
if isinstance(recv_x, tuple)
else recv_x
)
# Checks
rank_prefix_matrix = handle[0]
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
0
), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
assert (
gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
== recv_num_tokens_per_expert_list
)
if current_x is not x_pure_rand:
check_data(recv_x, rank_prefix_matrix)
if with_topk:
# Check `topk_idx`
assert (
recv_topk_idx.eq(-1)
| (
(recv_topk_idx >= 0)
& (recv_topk_idx < (num_experts // num_ranks))
)
).sum().item() == recv_topk_idx.numel()
for i, count in enumerate(recv_num_tokens_per_expert_list):
assert recv_topk_idx.eq(i).sum().item() == count
# Check `topk_weights`
if current_x is not x_pure_rand:
recv_topk_weights[recv_topk_idx.eq(-1)] = (
recv_topk_weights.amax(dim=1, keepdim=True).expand_as(
recv_topk_weights
)[recv_topk_idx.eq(-1)]
)
check_data(recv_topk_weights, rank_prefix_matrix)
# Test cached dispatch (must without top-k staffs)
if not with_topk:
dispatch_args = {
"x": current_x,
"handle": handle,
"config": config,
"async_finish": async_mode,
}
if previous_mode:
dispatch_args.update({"previous_event": buffer.capture()})
recv_x, _, _, _, _, event = buffer.dispatch(**dispatch_args)
event.current_stream_wait() if async_mode else ()
recv_x = (
per_token_cast_back(*recv_x)
if isinstance(recv_x, tuple)
else recv_x
)
if current_x is not x_pure_rand:
check_data(recv_x, rank_prefix_matrix)
# Test combine
combine_args = {
"x": recv_x,
"handle": handle,
"config": config,
"async_finish": async_mode,
}
if with_topk:
combine_args.update({"topk_weights": recv_topk_weights})
if previous_mode:
dispatch_args.update({"previous_event": buffer.capture()})
combined_x, combined_topk_weights, event = buffer.combine(
**combine_args
)
event.current_stream_wait() if async_mode else ()
check_x = combined_x.float() / is_token_in_rank.sum(
dim=1
).unsqueeze(1)
ref_x = x_pure_rand if current_x is x_pure_rand else x
assert calc_diff(check_x, ref_x) < 5e-6
if with_topk:
check_topk_weights = (
combined_topk_weights
if (current_x is x_pure_rand)
else (
combined_topk_weights
/ is_token_in_rank.sum(dim=1).unsqueeze(1)
)
)
ref_topk_weights = (
topk_weights_pure_rand
if current_x is x_pure_rand
else topk_weights
)
assert calc_diff(check_topk_weights, ref_topk_weights) < 1e-9
# For later tuning
dispatch_bf16_nvl_recv_bytes = recv_x.numel() * 2
combine_bf16_nvl_send_bytes = dispatch_bf16_nvl_recv_bytes
if local_rank == 0:
print(" passed", flush=True)
if local_rank == 0:
print("", flush=True)
# Tune dispatch performance
best_dispatch_results = None
fp8_factor = (1 + 4 / 128) / 2
for current_x in (x_e4m3, x):
best_time, best_results = 1e10, None
nvl_recv_bytes = (
(dispatch_bf16_nvl_recv_bytes * fp8_factor)
if isinstance(current_x, tuple)
else dispatch_bf16_nvl_recv_bytes
)
for nvl_chunk_size in range(4, 33, 4):
config = deep_ep.Config(num_sms, nvl_chunk_size, nvl_buffer_size)
tune_args = {"x": current_x, "handle": handle, "config": config}
t = bench(lambda: buffer.dispatch(**tune_args))[0]
if t < best_time:
best_time, best_results = t, (num_sms, nvl_chunk_size)
if local_rank == 0:
print(
f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}: {nvl_recv_bytes / 1e9 / t:.2f} GB/s (NVL) ",
flush=True,
)
if local_rank == 0:
print(
f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
flush=True,
)
print("", flush=True)
if isinstance(current_x, tuple):
# Gather FP8 the best config from rank 0
best_dispatch_results = torch.tensor(
[best_results[0], best_results[1]], dtype=torch.int32, device="cuda"
)
all_best_fp8_results_list = [
torch.zeros_like(best_dispatch_results)
for _ in range(torch.distributed.get_world_size())
]
dist.all_gather(
all_best_fp8_results_list, best_dispatch_results, group=group
)
best_dispatch_results = all_best_fp8_results_list[0].tolist()
dispatch_config = deep_ep.Config(
best_dispatch_results[0], best_dispatch_results[1], nvl_buffer_size
)
dispatch_args = {
"x": x,
"num_tokens_per_rank": num_tokens_per_rank,
"is_token_in_rank": is_token_in_rank,
"num_tokens_per_expert": num_tokens_per_expert,
"config": dispatch_config if dispatch_config is not None else config,
}
recv_x, _, _, _, handle, _ = buffer.dispatch(**dispatch_args)
# Tune combine performance
best_time, best_results = 1e10, None
for nvl_chunk_size in range(1, 7, 1):
config = deep_ep.Config(num_sms, nvl_chunk_size, nvl_buffer_size)
tune_args = {"x": recv_x, "handle": handle, "config": config}
t = bench(lambda: buffer.combine(**tune_args))[0]
if local_rank == 0:
print(
f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}: {combine_bf16_nvl_send_bytes / 1e9 / t:.2f} GB/s (NVL) ",
flush=True,
)
if t < best_time:
best_time, best_results = t, (num_sms, nvl_chunk_size)
if local_rank == 0:
print(
f"[tuning] Best combine: SMs {best_results[0]}, NVL chunk {best_results[1]}: {combine_bf16_nvl_send_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
flush=True,
)
print("", flush=True)
# noinspection PyUnboundLocalVariable
def test_loop(local_rank: int, num_local_ranks: int):
rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
test_ll_compatibility, num_rdma_bytes = False, 0
if test_ll_compatibility:
ll_num_tokens, ll_hidden, ll_num_experts, ll_num_topk = 16, 5120, 256, 9
num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
ll_num_tokens, ll_hidden, num_ranks, ll_num_experts
)
buffer = deep_ep.Buffer(
group,
int(1e9),
num_rdma_bytes,
low_latency_mode=test_ll_compatibility,
num_qps_per_rank=(ll_num_experts // num_ranks if test_ll_compatibility else 1),
)
torch.manual_seed(rank)
for i in (24,):
test_main(i, local_rank, num_ranks, rank, buffer, group)
if local_rank == 0:
print("", flush=True)
# Test compatibility with low latency functions
if test_ll_compatibility:
buffer.clean_low_latency_buffer(ll_num_tokens, ll_hidden, ll_num_experts)
test_deepep_low_latency.test_main(
ll_num_tokens,
ll_hidden,
ll_num_experts,
ll_num_topk,
rank,
num_ranks,
group,
buffer,
seed=1,
)
if __name__ == "__main__":
num_processes = 8
torch.multiprocessing.spawn(test_loop, args=(num_processes,), nprocs=num_processes)

View File

@@ -0,0 +1,325 @@
# Copy from deepseek-ai/DeepEP/tests/test_low_latency.py
import random
from functools import partial
import deep_ep
import torch
import torch.distributed as dist
from sglang.test.test_deepep_utils import (
bench,
bench_kineto,
calc_diff,
hash_tensor,
init_dist,
per_token_cast_back,
)
def test_main(
num_tokens: int,
hidden: int,
num_experts: int,
num_topk: int,
rank: int,
num_ranks: int,
group: dist.ProcessGroup,
buffer: deep_ep.Buffer,
seed: int = 0,
):
torch.manual_seed(seed + rank)
random.seed(seed + rank)
assert num_experts % num_ranks == 0
num_local_experts = num_experts // num_ranks
# NOTES: the integers greater than 256 exceeds the BF16 precision limit
rank_offset = 128
assert (
num_ranks - rank_offset < 257
), "Too many ranks (exceeding test precision limit)"
x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * (
rank - rank_offset
)
x[:, -128:] = torch.arange(num_tokens, device="cuda").to(torch.bfloat16).view(-1, 1)
scores = (
torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs()
+ 1
)
topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=True)[1]
topk_weights = torch.randn(
(num_tokens, num_topk), dtype=torch.float32, device="cuda"
).abs()
# Randomly mask some positions
for i in range(10):
topk_idx[random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)] = (
-1
)
# Check dispatch correctness
do_check = True
hash_value, num_times = 0, 0
for return_recv_hook in (False, True):
for dispatch_use_fp8 in (False, True):
num_times += 1
for i in range((num_times % 2) + 1):
packed_recv_x, packed_recv_count, handle, event, hook = (
buffer.low_latency_dispatch(
x,
topk_idx,
num_tokens,
num_experts,
use_fp8=dispatch_use_fp8,
async_finish=not return_recv_hook,
return_recv_hook=return_recv_hook,
)
)
hook() if return_recv_hook else event.current_stream_wait()
packed_recv_x = (
(packed_recv_x[0], packed_recv_x[1].contiguous())
if dispatch_use_fp8
else packed_recv_x
)
simulated_gemm_x = (
per_token_cast_back(
packed_recv_x[0].view(-1, hidden),
packed_recv_x[1].view(-1, hidden // 128),
).view(packed_recv_x[0].shape)
if dispatch_use_fp8
else packed_recv_x.clone()
)
all_topk_idx = torch.empty(
(num_ranks, num_tokens, num_topk), dtype=topk_idx.dtype, device="cuda"
)
dist.all_gather_into_tensor(all_topk_idx, topk_idx, group=group)
for i in range(num_local_experts if do_check else 0):
expert_id = rank * num_local_experts + i
recv_x = (
per_token_cast_back(packed_recv_x[0][i], packed_recv_x[1][i])
if dispatch_use_fp8
else packed_recv_x[i]
)
recv_count, recv_src_info, recv_layout_range = (
packed_recv_count[i],
handle[0][i],
handle[1][i],
)
# Check expert indices
int_mask = (2**32) - 1
num_valid_tokens = recv_count.item()
assert (
num_valid_tokens == (recv_layout_range & int_mask).sum().item()
), f"{num_valid_tokens} != {recv_layout_range & int_mask}.sum().item()"
assert (
num_valid_tokens == (all_topk_idx == expert_id).sum().item()
), f"{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}"
# Check received data
recv_x = recv_x[:num_valid_tokens]
recv_x_amin = recv_x[:, :-128].amin(dim=-1)
recv_src_info = recv_src_info[:num_valid_tokens]
assert torch.equal(recv_x_amin, recv_x[:, :-128].amax(dim=-1))
assert (
recv_x[:, -128:] - recv_src_info.view(-1, 1) % num_tokens
).sum().item() == 0
for j in range(num_ranks):
begin_idx, count = (recv_layout_range[j] >> 32).item(), (
recv_layout_range[j] & int_mask
).item()
assert (recv_x_amin == j - rank_offset).sum().item() == (
all_topk_idx[j] == expert_id
).sum().item()
assert (
recv_x[begin_idx : begin_idx + count][:-128] - j
).sum().item() == 0
if dispatch_use_fp8:
hash_value ^= hash_tensor(packed_recv_x[0][i, :num_valid_tokens])
hash_value ^= hash_tensor(packed_recv_x[1][i, :num_valid_tokens])
else:
hash_value ^= hash_tensor(packed_recv_x[i, :num_valid_tokens])
# Check combine correctness
for zero_copy in (False, True):
if zero_copy:
buffer.get_next_low_latency_combine_buffer(handle)[
:, :, :
] = simulated_gemm_x
out = torch.empty(
(num_tokens, hidden), dtype=torch.bfloat16, device="cuda"
)
combined_x, event, hook = buffer.low_latency_combine(
simulated_gemm_x,
topk_idx,
topk_weights,
handle,
async_finish=not return_recv_hook,
zero_copy=zero_copy,
return_recv_hook=return_recv_hook,
out=out,
)
hook() if return_recv_hook else event.current_stream_wait()
if do_check:
diff = calc_diff(
x
* topk_weights.masked_fill(topk_idx == -1, 0)
.sum(dim=1)
.view(-1, 1),
combined_x,
)
assert torch.isnan(combined_x).sum().item() == 0
assert diff < 1e-5, f"Error: {diff=}, {zero_copy=}"
hash_value ^= hash_tensor(combined_x)
def create_test_cast_with_outliers(num_outliers):
tmp = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
tmp /= tmp.abs().amax(dim=1).view(-1, 1)
assert tmp.abs().amax().item() <= 1
# Create some amax outliers
for i in range(num_outliers):
tmp[random.randint(0, num_tokens - 1)] *= 1e3
return tmp
# noinspection PyShadowingNames
def large_gemm_with_hook(hook):
mat_0 = torch.randn((8192, 8192), dtype=torch.float)
mat_1 = torch.randn((8192, 8192), dtype=torch.float)
mat_0 @ mat_1
hook()
# noinspection PyShadowingNames
def test_func(zero_copy: bool, return_recv_hook: bool):
recv_x, recv_count, handle, event, hook = buffer.low_latency_dispatch(
x,
topk_idx,
num_tokens,
num_experts,
async_finish=False,
return_recv_hook=return_recv_hook,
)
large_gemm_with_hook(hook) if return_recv_hook else None
if zero_copy:
buffer.get_next_low_latency_combine_buffer(handle)[
:, :, :
] = simulated_gemm_x
combined_x, event, hook = buffer.low_latency_combine(
simulated_gemm_x,
topk_idx,
topk_weights,
handle,
zero_copy=zero_copy,
return_recv_hook=return_recv_hook,
)
large_gemm_with_hook(hook) if return_recv_hook else None
# Calculate bandwidth
num_fp8_bytes, num_bf16_bytes = (hidden + hidden / 128 * 4 + 16), hidden * 2
num_dispatch_comm_bytes, num_combine_comm_bytes = 0, 0
for i in range(num_tokens):
num_selections = (topk_idx[i] != -1).sum().item()
num_dispatch_comm_bytes += num_fp8_bytes * num_selections
num_combine_comm_bytes += num_bf16_bytes * num_selections
# Dispatch + combine testing
avg_t, min_t, max_t = bench(
partial(test_func, zero_copy=False, return_recv_hook=False)
)
print(
f"[rank {rank}] Dispatch + combine bandwidth: {(num_dispatch_comm_bytes + num_combine_comm_bytes) / 1e9 / avg_t:.2f} GB/s, "
f"avg_t={avg_t * 1e6:.2f} us, min_t={min_t * 1e6:.2f} us, max_t={max_t * 1e6:.2f} us",
flush=True,
)
# Separate profiling
for return_recv_hook in (False, True):
group.barrier()
dispatch_t, combine_t = bench_kineto(
partial(test_func, zero_copy=True, return_recv_hook=return_recv_hook),
kernel_names=("dispatch", "combine"),
barrier_comm_profiling=True,
suppress_kineto_output=True,
)
if not return_recv_hook:
print(
f"[rank {rank}] Dispatch bandwidth: {num_dispatch_comm_bytes / 1e9 / dispatch_t:.2f} GB/s, avg_t={dispatch_t * 1e6:.2f} us | "
f"Combine bandwidth: {num_combine_comm_bytes / 1e9 / combine_t:.2f} GB/s, avg_t={combine_t * 1e6:.2f} us",
flush=True,
)
else:
print(
f"[rank {rank}] Dispatch send/recv time: {dispatch_t * 2 * 1e6:.2f} us | "
f"Combine send/recv time: {combine_t * 2 * 1e6:.2f} us",
flush=True,
)
return hash_value
# noinspection PyUnboundLocalVariable
def test_loop(local_rank: int, num_local_ranks: int):
rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
num_tokens, hidden, num_topk, num_experts = 128, 7168, 8, 288
num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
num_tokens, hidden, num_ranks, num_experts
)
if local_rank == 0:
print(f"Allocating buffer size: {num_rdma_bytes / 1e6} MB ...", flush=True)
buffer = deep_ep.Buffer(
group,
num_rdma_bytes=num_rdma_bytes,
low_latency_mode=True,
num_qps_per_rank=num_experts // num_ranks,
)
test_main(
num_tokens,
hidden,
num_experts,
num_topk,
rank,
num_ranks,
group,
buffer,
seed=1,
)
do_pressure_test = False
for seed in range(int(1e9) if do_pressure_test else 0):
if local_rank == 0:
print(f"Testing with seed {seed} ...", flush=True)
ref_hash = test_main(
num_tokens,
hidden,
num_experts,
num_topk,
rank,
num_ranks,
group,
buffer,
seed=seed,
)
for i in range(20):
assert (
test_main(
num_tokens,
hidden,
num_experts,
num_topk,
rank,
num_ranks,
group,
buffer,
seed=seed,
)
== ref_hash
), f"Error: seed={seed}"
if __name__ == "__main__":
# TODO: you may modify NUMA binding for less CPU overhead
num_processes = 8
torch.multiprocessing.spawn(test_loop, args=(num_processes,), nprocs=num_processes)

154
third_party/sglang/test/manual/ep/test_eplb.py vendored Executable file
View File

@@ -0,0 +1,154 @@
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
import sglang as sgl
from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
class _BaseTestDynamicEPLB(CustomTestCase):
extra_args = []
@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),
envs.SGLANG_EXPERT_LOCATION_UPDATER_CANARY.override(True),
):
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",
"--enable-eplb",
"--ep-num-redundant-experts",
"4",
"--eplb-rebalance-num-iterations",
"50",
"--expert-distribution-recorder-buffer-size",
"50",
# TODO pr-chain: enable later
# "--enable-expert-distribution-metrics",
# TODO auto determine these flags
"--expert-distribution-recorder-mode",
"stat",
"--ep-dispatch-algorithm",
"static",
*cls.extra_args,
],
)
@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.assertGreater(metrics["score"], 0.5)
class TestDynamicEPLBSimple(_BaseTestDynamicEPLB):
pass
class TestDynamicEPLBMultiChunk(_BaseTestDynamicEPLB):
extra_args = ["--eplb-rebalance-layers-per-chunk", "1"]
class TestStaticEPLB(CustomTestCase):
def test_save_expert_distribution_and_init_expert_location(self):
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
with tempfile.TemporaryDirectory() as tmp_dir:
engine_kwargs = dict(
model_path=DEFAULT_MLA_MODEL_NAME_FOR_TEST,
trust_remote_code=True,
ep_num_redundant_experts=4,
enable_dp_attention=True,
moe_a2a_backend="deepep",
disable_cuda_graph=True,
expert_distribution_recorder_mode="stat",
tp_size=2,
dp_size=2,
log_level="info",
# TODO pr-chain: enable later
# enable_expert_distribution_metrics=True,
)
print(f"Action: start engine")
envs.SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR.set(tmp_dir)
engine = sgl.Engine(
**engine_kwargs,
disable_overlap_schedule=True,
)
engine.start_expert_distribution_record()
self._assert_engine_generate_correct(engine)
print(f"Action: dump_expert_distribution_record")
engine.dump_expert_distribution_record()
snapshot_path = list(Path(tmp_dir).glob("*.pt"))[0]
assert snapshot_path is not None
print(f"{snapshot_path=}")
print(f"Action: shutdown engine")
engine.shutdown()
del engine
print(f"Action: start engine with init_expert_location")
engine = sgl.Engine(
**engine_kwargs,
init_expert_location=str(snapshot_path),
port=21000,
# TODO auto determine these flags
ep_dispatch_algorithm="static",
)
self._assert_engine_generate_correct(engine)
print(f"Action: shutdown engine")
engine.shutdown()
del engine
def _assert_engine_generate_correct(self, engine: sgl.Engine):
output = engine.generate(
prompt=["1+1=2, 2+2=4", "One plus one is two, two plus two is four"],
sampling_params=dict(max_new_tokens=8, temperature=0.0),
)
print(f"engine.generate {output=}")
self.assertEqual(
[x["text"] for x in output],
[", 4+4=8,", ", four plus four is eight, eight"],
)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,116 @@
import json
import unittest
from types import SimpleNamespace
from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
class TestPureTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MLA_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=[
"--trust-remote-code",
"--tp",
"2",
"--moe-a2a-backend",
"deepep",
"--disable-cuda-graph",
],
)
@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.assertGreater(metrics["score"], 0.5)
class TestDPAttn(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",
# Test custom config
"--deepep-config",
json.dumps(
{
"normal_dispatch": {
"num_sms": 20,
"num_max_nvl_chunked_send_tokens": 16,
"num_max_nvl_chunked_recv_tokens": 256,
"num_max_rdma_chunked_send_tokens": 6,
"num_max_rdma_chunked_recv_tokens": 128,
},
"normal_combine": {
"num_sms": 20,
"num_max_nvl_chunked_send_tokens": 6,
"num_max_nvl_chunked_recv_tokens": 256,
"num_max_rdma_chunked_send_tokens": 6,
"num_max_rdma_chunked_recv_tokens": 128,
},
}
),
],
)
@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.assertGreater(metrics["score"], 0.5)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,75 @@
"""
Usage:
python -m unittest test_moe_deepep_eval_accuracy_large.TestMoEDeepEPEvalAccuracyLarge.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_DEEPEP_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
class TestMoEDeepEPEvalAccuracyLarge(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_DEEPEP_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=[
"--trust-remote-code",
"--tp",
"8",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"128",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=64,
num_shots=8,
)
metrics = run_eval(args)
print(f"Eval accuracy of GSM8K: {metrics=}")
self.assertGreater(metrics["score"], 0.93)
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)
print(f"Eval accuracy of MMLU: {metrics=}")
self.assertGreater(metrics["score"], 0.87)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,143 @@
import time
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.run_eval import run_eval
from sglang.test.server_fixtures.disaggregation_fixture import get_rdma_devices_args
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
CustomTestCase,
popen_launch_pd_server,
)
ib_devices = get_rdma_devices_args()
class TestBackup(CustomTestCase):
extra_args = []
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_port = 20000
cls.base_url = f"http://127.0.0.1:{cls.base_port}"
cls.num_processes = 2
# TODO (stage 100): in the future, implement a specified multiprocess launcher
cls.processes = [
popen_launch_pd_server(
cls.model,
f"http://127.0.0.1:{cls.base_port + i}",
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--enable-dp-attention",
"--dp",
"4",
"--elastic-ep-backend",
"mooncake",
"--mooncake-ib-device",
ib_devices,
"--moe-a2a-backend",
"mooncake",
"--deepep-mode",
"low_latency",
"--moe-dense-tp-size",
"1",
"--enable-dp-lm-head",
"--enable-two-batch-overlap",
"--disable-custom-all-reduce",
"--enable-elastic-expert-backup",
"--enable-eplb",
"--eplb-rebalance-num-iterations",
"50",
"--chunked-prefill-size",
"512",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"512",
"--mem-fraction-static",
"0.5",
"--dist-init-addr",
"127.0.0.1:5000",
"--nnodes",
f"{cls.num_processes}",
"--node-rank",
f"{i}",
"--base-gpu-id",
f"{i * 2}",
],
)
for i in range(cls.num_processes)
]
server_ready = [False] * cls.num_processes
start_time = time.perf_counter()
with requests.Session() as session:
while (
time.perf_counter() - start_time < DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
and not all(server_ready)
):
for i, process in enumerate(cls.processes):
return_code = process.poll()
if return_code is not None:
# Server failed to start (non-zero exit code) or crashed
raise Exception(
f"Server process exited with code {return_code}. "
"Check server logs for errors."
)
try:
headers = {
"Content-Type": "application/json; charset=utf-8",
}
response = session.get(
f"http://127.0.0.1:{cls.base_port + i}/health_generate",
headers=headers,
)
if response.status_code == 200:
server_ready[i] = True
except requests.RequestException:
pass
return_code = process.poll()
if return_code is not None:
raise Exception(
f"Server unexpectedly exits ({return_code=}). Usually there will be error logs describing the cause far above this line."
)
time.sleep(10)
if not all(server_ready):
for process in cls.processes:
kill_process_tree(process.pid)
raise TimeoutError("Server failed to start within the timeout period.")
@classmethod
def tearDownClass(cls):
for process in cls.processes:
kill_process_tree(process.pid)
def test_gsm8k(self):
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(metrics)
self.assertGreater(metrics["score"], 0.60)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,115 @@
import os
import time
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.server_fixtures.disaggregation_fixture import get_rdma_devices_args
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
TEST_MODEL = os.environ.get("NIXL_EP_TEST_MODEL", DEFAULT_MODEL_NAME_FOR_TEST_MLA)
os.environ.setdefault("SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK", "1024")
ib_devices = get_rdma_devices_args()
NIXL_COMMON = [
"--trust-remote-code",
"--moe-a2a-backend",
"nixl",
"--deepep-mode",
"low_latency",
"--tp",
"8",
"--mem-fraction-static",
"0.78",
]
DP_ATTN = ["--dp", "8", "--enable-dp-attention"]
ELASTIC_NIXL = [
"--elastic-ep-backend",
"nixl",
"--enable-eplb",
"--ep-num-redundant-experts",
"24",
]
ELASTIC_MOONCAKE = [
"--elastic-ep-backend",
"mooncake",
"--mooncake-ib-device",
ib_devices,
"--enable-eplb",
"--ep-num-redundant-experts",
"24",
]
class _EPTestBase(CustomTestCase):
server_args: list[str] = []
@classmethod
def setUpClass(cls):
cls.model = TEST_MODEL
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=cls.server_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
cls.process.wait(timeout=15)
time.sleep(2)
def _run_gsm8k(self):
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(metrics)
return metrics
def test_gsm8k(self):
metrics = self._run_gsm8k()
self.assertGreater(metrics["score"], 0.60)
class TestNixlEPTP(_EPTestBase):
server_args = [*NIXL_COMMON]
class TestNixlEPDPAttn(_EPTestBase):
server_args = [*NIXL_COMMON, *DP_ATTN]
class TestNixlEPElasticEP(_EPTestBase):
server_args = [*NIXL_COMMON, *DP_ATTN, *ELASTIC_NIXL]
class TestNixlMoeMooncakeElasticEP(_EPTestBase):
server_args = [*NIXL_COMMON, *DP_ATTN, *ELASTIC_MOONCAKE]
pkill_process_1 = "sglang::scheduler_DP1_TP8_EP8"
def test_gsm8k_fault_1(self):
os.system(f"pkill -f {self.pkill_process_1}")
metrics = self._run_gsm8k()
self.assertGreater(metrics["score"], 0.60)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,257 @@
import os
import random
import tempfile
import unittest
from typing import Dict
import requests
from sglang.benchmark.utils import get_tokenizer
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
)
class DisaggregationHiCacheBase(PDDisaggregationServerBase):
"""Base class for disaggregation with HiCache tests"""
@classmethod
def setUpClass(cls):
super(DisaggregationHiCacheBase, cls).setUpClass()
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
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", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
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.
res = requests.post(
f"{self.prefill_url}/flush_cache",
params={"timeout": 30},
timeout=40,
)
res.raise_for_status()
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()

View File

@@ -0,0 +1,214 @@
"""
Usage:
python3 -m unittest test_pp_with_hicache.TestPPWithHiCache.test_eval_accuracy
"""
import os
import subprocess
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.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
find_available_port,
popen_launch_server,
)
class TestPPWithHiCache(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.base_url = f"http://127.0.0.1:{find_available_port(23337)}"
parsed_url = urlparse(cls.base_url)
cls.base_host = parsed_url.hostname
cls.base_port = str(parsed_url.port)
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
cls._start_mooncake_services()
server_args_dict = {
"--enable-hierarchical-cache": True,
"--mem-fraction-static": 0.6,
"--hicache-ratio": 1.2,
"--page-size": 64,
"--enable-cache-report": True,
"--hicache-storage-prefetch-policy": "wait_complete",
"--hicache-storage-backend": "mooncake",
"--tp-size": 2,
"--pp-size": 2,
"--chunked-prefill-size": 256,
"--hicache-mem-layout": "page_first",
}
final_server_args = []
for key, value in server_args_dict.items():
final_server_args.append(str(key))
if value is not True:
final_server_args.append(str(value))
env_vars = {**os.environ, **cls._mooncake_env()}
try:
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=final_server_args,
env=env_vars,
)
except Exception:
cls._stop_mooncake_services()
raise
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process"):
kill_process_tree(cls.process.pid)
cls._stop_mooncake_services()
@classmethod
def _start_mooncake_services(cls):
try:
import mooncake.http_metadata_server # type: ignore # noqa: F401
except Exception as exc: # pragma: no cover - environment dependent
raise unittest.SkipTest(
f"Mooncake metadata server module unavailable: {exc}"
) from exc
cls._mooncake_master_port = find_available_port(50051)
cls._mooncake_metadata_port = find_available_port(8080)
try:
cls._mooncake_metadata_process = subprocess.Popen(
[
"python3",
"-m",
"mooncake.http_metadata_server",
"--port",
str(cls._mooncake_metadata_port),
],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
preexec_fn=os.setsid,
)
except (FileNotFoundError, subprocess.SubprocessError) as exc:
cls._stop_mooncake_services()
raise unittest.SkipTest(
f"Could not start Mooncake metadata service: {exc}"
) from exc
try:
cls._mooncake_master_process = subprocess.Popen(
["mooncake_master", "--port", str(cls._mooncake_master_port)],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
preexec_fn=os.setsid,
)
except (FileNotFoundError, subprocess.SubprocessError) as exc:
cls._stop_mooncake_services()
raise unittest.SkipTest(f"Could not start mooncake_master: {exc}") from exc
if not cls._wait_for_mooncake_ready():
cls._stop_mooncake_services()
raise unittest.SkipTest("Mooncake services did not become ready in time")
@classmethod
def _stop_mooncake_services(cls):
for attr in ("_mooncake_metadata_process", "_mooncake_master_process"):
proc = getattr(cls, attr, None)
if proc:
try:
os.killpg(os.getpgid(proc.pid), 9)
proc.wait(timeout=5)
except Exception:
pass
cls._mooncake_metadata_process = None
cls._mooncake_master_process = None
@classmethod
def _mooncake_env(cls):
return {
"MOONCAKE_MASTER": f"127.0.0.1:{cls._mooncake_master_port}",
"MOONCAKE_PROTOCOL": "tcp",
"MC_MS_AUTO_DISC": "0",
"MOONCAKE_DEVICE": "",
"MOONCAKE_TE_META_DATA_SERVER": f"http://127.0.0.1:{cls._mooncake_metadata_port}/metadata",
"MOONCAKE_GLOBAL_SEGMENT_SIZE": "4294967296",
"SGLANG_ENABLE_DETERMINISTIC_INFERENCE": "1",
}
@classmethod
def _wait_for_mooncake_ready(cls, timeout: int = 30) -> bool:
start_time = time.time()
while time.time() - start_time < timeout:
metadata_ready = False
master_ready = False
if (
getattr(cls, "_mooncake_metadata_process", None)
and cls._mooncake_metadata_process.poll() is None
):
try:
resp = requests.get(
f"http://127.0.0.1:{cls._mooncake_metadata_port}/metadata",
timeout=2,
)
print(resp)
metadata_ready = True
except requests.RequestException:
metadata_ready = False
if (
getattr(cls, "_mooncake_master_process", None)
and cls._mooncake_master_process.poll() is None
):
if time.time() - start_time > 3:
master_ready = True
if metadata_ready and master_ready:
return True
time.sleep(1.5)
return False
def flush_cache(self):
res = requests.post(
f"{self.base_url}/flush_cache",
params={"timeout": 30},
timeout=40,
)
res.raise_for_status()
def test_eval_accuracy(self):
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=40,
num_threads=24,
)
metrics_initial = run_eval(args)
self.assertGreater(metrics_initial["score"], 0.6)
self.flush_cache()
metrics_cached = run_eval(args)
self.assertGreater(metrics_cached["score"], 0.6)
accuracy_diff = abs(metrics_initial["score"] - metrics_cached["score"])
self.assertLess(accuracy_diff, 0.05)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,64 @@
#!/usr/bin/env python3
import argparse
try:
import mooncake
BENCH_TOOL_PATH = f"{mooncake.__path__[0]}/transfer_engine_bench"
print(f"Mooncake is installed. Bench tool path:\n{BENCH_TOOL_PATH}")
except ImportError:
BENCH_TOOL_PATH = None
print("Mooncake is not installed.")
exit(0)
def run_cmd(args):
cmd = [BENCH_TOOL_PATH]
if args.initiator:
cmd += ["--mode=initiator"]
elif args.target:
cmd += ["--mode=target"]
if args.metadata_server:
cmd += [f"--metadata_server={args.metadata_server}"]
if args.mc_segment_id:
cmd += [f"--segment_id={args.mc_segment_id}"]
if args.device:
cmd += [f"--device_name={args.device}"]
if args.bench_h2h:
cmd += ["--use_vram=false"]
cmd += ["--auto_discovery"]
print(f"Executing command: {' '.join(cmd)}")
import subprocess
try:
subprocess.run(cmd, check=True)
except subprocess.CalledProcessError as e:
print(f"Command failed with error: {e}")
exit(1)
def main():
parser = argparse.ArgumentParser()
group = parser.add_mutually_exclusive_group()
group.add_argument("--initiator", action="store_true", help="Run as initiator")
group.add_argument("--target", action="store_true", help="Run as target")
parser.add_argument("--metadata-server", type=str, default="P2PHANDSHAKE")
parser.add_argument("--mc-segment-id", type=str, default=None)
parser.add_argument("--bench-h2h", action="store_true")
parser.add_argument("--device", type=str, default="mlx5_0")
args = parser.parse_args()
print("Running Mooncake transfer engine benchmark...")
if not args.initiator and not args.target:
parser.error("Please specify --initiator or --target")
if args.initiator and args.mc_segment_id is None:
parser.error("Please specify --mc-segment-id for initiator")
run_cmd(args)
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,430 @@
#!/usr/bin/env python3
"""
Test script for validating Mooncake transfer-engine gating and initialization.
Tests the Mooncake-related branches in the current model-runner flow.
This test verifies:
1. MooncakeTransferEngine initialization conditions
2. Different server argument combinations that trigger mooncake TE
3. Mooncake transfer engine initialization with hostname, gpu_id, and ib_device
Usage:
# Run from project root on 2 GPUs
CUDA_VISIBLE_DEVICES=0,1 python test/manual/kv_transfer/test_mooncake_transfer_engine_init.py
"""
import argparse
import multiprocessing
import os
import sys
import time
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Optional
from unittest.mock import patch
@dataclass
class ServerArgs:
"""Mock ServerArgs for testing."""
disaggregation_mode: str = "null"
disaggregation_transfer_backend: str = "mooncake"
enable_hierarchical_cache: bool = False
hicache_storage_backend: str = "mooncake"
encoder_only: bool = False
language_only: bool = False
encoder_transfer_backend: str = "mooncake"
enable_elastic_expert_backup: bool = False
elastic_ep_backend: Optional[str] = None
disaggregation_ib_device: Optional[str] = None
mooncake_ib_device: Optional[str] = None
def test_mooncake_te_condition(server_args: ServerArgs) -> bool:
"""
Test the condition logic for using MooncakeTransferEngine.
"""
from sglang.srt.model_executor.model_runner import ModelRunner
dummy_runner = SimpleNamespace(server_args=server_args, gpu_id=0)
init_called = False
def _fake_init_mooncake_transfer_engine(*, hostname, gpu_id, ib_device):
nonlocal init_called
init_called = True
return SimpleNamespace(
hostname=hostname,
gpu_id=gpu_id,
ib_device=ib_device,
)
with patch(
"sglang.srt.distributed.device_communicators.mooncake_transfer_engine.init_mooncake_transfer_engine",
side_effect=_fake_init_mooncake_transfer_engine,
), patch(
"sglang.srt.model_executor.model_runner.get_local_ip_auto",
return_value="127.0.0.1",
):
ModelRunner.init_shared_mooncake_transfer_engine(dummy_runner)
return init_called
def run_mooncake_init(
rank: int,
world_size: int,
master_port: int,
args: argparse.Namespace,
server_args: ServerArgs,
):
"""Worker function for testing mooncake transfer engine initialization."""
os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = str(master_port)
os.environ["RANK"] = str(rank)
os.environ["WORLD_SIZE"] = str(world_size)
os.environ["LOCAL_RANK"] = str(rank)
# Import before try block to avoid NameError in finally
import torch
import torch.distributed as dist
dist_initialized = False
try:
# Initialize distributed environment
print(f"[Rank {rank}] Initializing distributed environment...")
dist.init_process_group(
backend="nccl",
world_size=world_size,
rank=rank,
init_method=f"tcp://127.0.0.1:{master_port}",
device_id=rank,
)
dist_initialized = True
# Set device
torch.cuda.set_device(rank)
# Sync to ensure all ranks are ready
dist.barrier()
print(f"[Rank {rank}] Distributed initialization complete.")
# Test the condition logic
use_mooncake_te = test_mooncake_te_condition(server_args)
print(f"[Rank {rank}] use_mooncake_te = {use_mooncake_te}")
if use_mooncake_te:
print(f"[Rank {rank}] Attempting to initialize MooncakeTransferEngine...")
from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
init_mooncake_transfer_engine,
)
from sglang.srt.utils import get_local_ip_auto
ib_device = (
server_args.disaggregation_ib_device or server_args.mooncake_ib_device
)
print(f"[Rank {rank}] IB device: {ib_device}")
# Always actually initialize mooncake
engine = init_mooncake_transfer_engine(
hostname=get_local_ip_auto(),
gpu_id=rank,
ib_device=ib_device,
)
print(f"[Rank {rank}] Session ID: {engine.get_session_id()}")
print(f"[Rank {rank}] MooncakeTransferEngine initialized successfully!")
dist.barrier()
print(f"[Rank {rank}] Test completed successfully!")
sys.exit(0)
except ImportError as e:
print(f"[Rank {rank}] Mooncake not available (ImportError): {e}")
sys.exit(1)
except Exception as e:
print(f"[Rank {rank}] Test failed with error: {e}")
import traceback
traceback.print_exc()
sys.exit(1)
finally:
# Cleanup
if dist_initialized and dist.is_initialized():
dist.destroy_process_group()
print(f"[Rank {rank}] Process group destroyed.")
def run_test(args: argparse.Namespace, server_args: ServerArgs) -> bool:
"""Run the mooncake transfer engine test."""
# Set CUDA visible devices
cuda_devices = args.cuda_visible_devices.split(",")
world_size = len(cuda_devices)
if world_size < 2:
print("ERROR: This test requires at least 2 GPUs.")
print(
"Usage: CUDA_VISIBLE_DEVICES=0,1 python test/manual/kv_transfer/test_mooncake_transfer_engine_init.py"
)
sys.exit(1)
# Check GPU availability
import torch
if not torch.cuda.is_available():
print("ERROR: CUDA is not available")
sys.exit(1)
available_gpus = torch.cuda.device_count()
if world_size > available_gpus:
print(f"ERROR: Requested {world_size} GPUs but only {available_gpus} available")
sys.exit(1)
print(f"Testing with {world_size} GPUs: {cuda_devices}")
print()
# Print server args configuration
print("ServerArgs configuration:")
for key, value in vars(server_args).items():
print(f" {key}: {value}")
print()
# Check if mooncake should be used
use_mooncake_te = test_mooncake_te_condition(server_args)
print(f"use_mooncake_te = {use_mooncake_te}")
print()
# Find a free port
import socket
with socket.socket() as s:
s.bind(("", 0))
master_port = s.getsockname()[1]
print(f"Using master port: {master_port}")
# Spawn worker processes
ctx = multiprocessing.get_context("spawn")
processes = []
for rank in range(world_size):
p = ctx.Process(
target=run_mooncake_init,
args=(rank, world_size, master_port, args, server_args),
)
p.start()
processes.append(p)
# Wait for all processes to complete
success = True
for i, p in enumerate(processes):
p.join(timeout=60)
if p.exitcode != 0:
print(f"Process {i} failed with exit code: {p.exitcode}")
success = False
# Cleanup any remaining processes
for p in processes:
if p.is_alive():
print(f"Process {p.pid} is still alive, terminating...")
p.terminate()
p.join(timeout=5)
return success
def test_condition_logic():
"""Test the condition logic for different server argument combinations."""
print("=" * 60)
print("Testing condition logic for use_mooncake_te")
print("=" * 60)
print()
original_hicache_reuse = os.environ.get("SGLANG_HICACHE_MOONCAKE_REUSE_TE")
passed = 0
failed = 0
try:
test_cases = [
# (name, env_value, server_args, expected_result)
(
"PD disaggregation with mooncake",
None,
ServerArgs(
disaggregation_mode="prefill",
disaggregation_transfer_backend="mooncake",
),
True,
),
(
"PD disaggregation without mooncake",
None,
ServerArgs(
disaggregation_mode="prefill",
disaggregation_transfer_backend="other",
),
False,
),
(
"No disaggregation",
None,
ServerArgs(),
False,
),
(
"HiCache with mooncake (env=False)",
"0",
ServerArgs(
enable_hierarchical_cache=True,
hicache_storage_backend="mooncake",
),
False,
),
(
"HiCache with mooncake (env=True)",
"1",
ServerArgs(
enable_hierarchical_cache=True,
hicache_storage_backend="mooncake",
),
True,
),
(
"Encoder only with mooncake",
None,
ServerArgs(encoder_only=True, encoder_transfer_backend="mooncake"),
True,
),
(
"Language only with mooncake",
None,
ServerArgs(language_only=True, encoder_transfer_backend="mooncake"),
True,
),
(
"Elastic expert backup with backend",
None,
ServerArgs(
enable_elastic_expert_backup=True,
elastic_ep_backend="mooncake",
),
True,
),
(
"Elastic expert backup without backend",
None,
ServerArgs(enable_elastic_expert_backup=True, elastic_ep_backend=None),
False,
),
]
for name, env_value, server_args, expected in test_cases:
if env_value is None:
os.environ.pop("SGLANG_HICACHE_MOONCAKE_REUSE_TE", None)
else:
os.environ["SGLANG_HICACHE_MOONCAKE_REUSE_TE"] = env_value
result = test_mooncake_te_condition(server_args)
status = "PASS" if result == expected else "FAIL"
if result == expected:
passed += 1
else:
failed += 1
print(f"{status}: {name}")
print(f" Expected: {expected}, Got: {result}")
print()
finally:
if original_hicache_reuse is None:
os.environ.pop("SGLANG_HICACHE_MOONCAKE_REUSE_TE", None)
else:
os.environ["SGLANG_HICACHE_MOONCAKE_REUSE_TE"] = original_hicache_reuse
print(f"Condition logic tests: {passed} passed, {failed} failed")
print()
return failed == 0
def main():
parser = argparse.ArgumentParser(
description="Validate Mooncake transfer-engine gating and initialization"
)
parser.add_argument(
"--cuda-visible-devices",
type=str,
default="0,1",
help="CUDA visible devices (default: 0,1)",
)
parser.add_argument(
"--test-case",
type=str,
choices=[
"pd_disaggregation",
"hicache",
"encoder_only",
"language_only",
"elastic_ep",
],
default="pd_disaggregation",
help="Test case to run",
)
args = parser.parse_args()
print("=" * 60)
print("Mooncake Transfer Engine Init Test")
print("=" * 60)
print()
# First run condition logic tests
condition_passed = test_condition_logic()
if not condition_passed:
print("Condition logic tests failed, skipping distributed test.")
sys.exit(1)
# Configure server args based on test case
server_args = ServerArgs()
if args.test_case == "pd_disaggregation":
server_args.disaggregation_mode = "prefill"
server_args.disaggregation_transfer_backend = "mooncake"
elif args.test_case == "hicache":
server_args.enable_hierarchical_cache = True
server_args.hicache_storage_backend = "mooncake"
os.environ["SGLANG_HICACHE_MOONCAKE_REUSE_TE"] = "1"
elif args.test_case == "encoder_only":
server_args.encoder_only = True
server_args.encoder_transfer_backend = "mooncake"
elif args.test_case == "language_only":
server_args.language_only = True
server_args.encoder_transfer_backend = "mooncake"
elif args.test_case == "elastic_ep":
server_args.enable_elastic_expert_backup = True
server_args.elastic_ep_backend = "mooncake"
start_time = time.time()
success = run_test(args, server_args)
elapsed_time = time.time() - start_time
print()
print("=" * 60)
if success:
print(f"TEST PASSED (elapsed: {elapsed_time:.2f}s)")
else:
print(f"TEST FAILED (elapsed: {elapsed_time:.2f}s)")
print("=" * 60)
sys.exit(0 if success else 1)
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,51 @@
import unittest
import sglang as sgl
from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST, CustomTestCase
class TestBind(CustomTestCase):
backend = None
@classmethod
def setUpClass(cls):
cls.backend = sgl.Runtime(model_path=DEFAULT_MODEL_NAME_FOR_TEST)
sgl.set_default_backend(cls.backend)
@classmethod
def tearDownClass(cls):
cls.backend.shutdown()
def test_bind(self):
@sgl.function
def few_shot_qa(s, prompt, question):
s += prompt
s += "Q: What is the capital of France?\n"
s += "A: Paris\n"
s += "Q: " + question + "\n"
s += "A:" + sgl.gen("answer", stop="\n")
few_shot_qa_2 = few_shot_qa.bind(
prompt="The following are questions with answers.\n\n"
)
tracer = few_shot_qa_2.trace()
print(tracer.last_node.print_graph_dfs() + "\n")
def test_cache(self):
@sgl.function
def few_shot_qa(s, prompt, question):
s += prompt
s += "Q: What is the capital of France?\n"
s += "A: Paris\n"
s += "Q: " + question + "\n"
s += "A:" + sgl.gen("answer", stop="\n")
few_shot_qa_2 = few_shot_qa.bind(
prompt="Answer the following questions as if you were a 5-year-old kid.\n\n"
)
few_shot_qa_2.cache(self.backend)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,91 @@
import unittest
import numpy as np
from sglang.lang.choices import (
greedy_token_selection,
token_length_normalized,
unconditional_likelihood_normalized,
)
from sglang.test.test_utils import CustomTestCase
MOCK_CHOICES_INPUT_DATA = {
"choices": [
"organ", # ["organ"]
"organism", # ["organ", "ism"]
"antidisestablishmentarianism", # ["ant", "id", "is", "est", "ablish", "ment", "arian", "ism"]
],
"normalized_prompt_logprobs": [-0.1, -0.2, -0.05],
"input_token_logprobs": [
[[-0.1, 1, None]],
[[-0.1, 1, None], [-0.3, 2, None]],
[
[-0.4, 3, None],
[-0.25, 4, None],
[-0.1, 5, None],
[-0.01, 6, None],
[-0.01, 7, None],
[-0.01, 8, None],
[-0.01, 9, None],
[-0.01, 2, None],
],
],
"output_token_logprobs": [
[[-0.1, 10, None]],
[[-0.1, 10, None]],
[[-0.1, 10, None]],
],
"unconditional_token_logprobs": [
[[None, 1, None]],
[[None, 1, None], [-1.4, 2, None]],
[
[None, 3, None],
[-0.25, 4, None],
[-0.1, 5, None],
[-0.01, 6, None],
[-0.01, 7, None],
[-0.01, 8, None],
[-0.01, 9, None],
[-0.01, 2, None],
],
],
}
class TestChoices(CustomTestCase):
def test_token_length_normalized(self):
"""Confirm 'antidisestablishmentarianism' is selected due to high confidences for
its later tokens resulting in highest token length normalized prompt logprob."""
decision = token_length_normalized(**MOCK_CHOICES_INPUT_DATA)
assert decision.decision == "antidisestablishmentarianism"
def test_greedy_token_selection(self):
"""Confirm 'organ' is selected due it having the joint highest initial token
logprob, and a higher average logprob than organism's second token."""
decision = greedy_token_selection(**MOCK_CHOICES_INPUT_DATA)
assert decision.decision == "organ"
assert np.allclose(
decision.meta_info["greedy_logprob_matrix"],
[
[-0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1],
[-0.1, -0.3, -0.2, -0.2, -0.2, -0.2, -0.2, -0.2],
[-0.4, -0.25, -0.1, -0.01, -0.01, -0.01, -0.01, -0.01],
],
atol=0.01,
)
def test_unconditional_likelihood_normalized(self):
"""Confirm 'organism' is selected due to it having the highest average token logprob
once normalized by the unconditional token logprobs."""
decision = unconditional_likelihood_normalized(**MOCK_CHOICES_INPUT_DATA)
assert decision.decision == "organism"
assert np.allclose(
decision.meta_info["normalized_unconditional_prompt_logprobs"],
[-0.1, 0.5, -0.05],
atol=0.01,
)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,138 @@
import argparse
from enum import Enum
from pydantic import BaseModel, constr
import sglang as sgl
from sglang.srt.constrained.outlines_backend import build_regex_from_object
from sglang.test.test_utils import (
add_common_sglang_args_and_parse,
select_sglang_backend,
)
IP_REGEX = r"((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)"
ip_jump_forward = (
r"The google's DNS sever address is "
+ IP_REGEX
+ r" and "
+ IP_REGEX
+ r". "
+ r"The google's website domain name is "
+ r"www\.(\w)+\.(\w)+"
+ r"."
)
# fmt: off
@sgl.function
def regex_gen(s):
s += "Q: What is the IP address of the Google DNS servers?\n"
s += "A: " + sgl.gen(
"answer",
max_tokens=128,
temperature=0,
regex=ip_jump_forward,
)
# fmt: on
json_jump_forward = (
r"""The information about Hogwarts is in the following JSON format\.\n"""
+ r"""\n\{\n"""
+ r""" "name": "[\w\d\s]*",\n"""
+ r""" "country": "[\w\d\s]*",\n"""
+ r""" "latitude": [-+]?[0-9]*\.?[0-9]+,\n"""
+ r""" "population": [-+]?[0-9]+,\n"""
+ r""" "top 3 landmarks": \["[\w\d\s]*", "[\w\d\s]*", "[\w\d\s]*"\],\n"""
+ r"""\}\n"""
)
# fmt: off
@sgl.function
def json_gen(s):
s += sgl.gen(
"json",
max_tokens=128,
temperature=0,
regex=json_jump_forward,
)
# fmt: on
class Weapon(str, Enum):
sword = "sword"
axe = "axe"
mace = "mace"
spear = "spear"
bow = "bow"
crossbow = "crossbow"
class Armor(str, Enum):
leather = "leather"
chainmail = "chainmail"
plate = "plate"
class Character(BaseModel):
name: constr(max_length=10)
age: int
armor: Armor
weapon: Weapon
strength: int
@sgl.function
def character_gen(s):
s += "Give me a character description who is a wizard.\n"
s += sgl.gen(
"character",
max_tokens=128,
temperature=0,
regex=build_regex_from_object(Character),
)
def main(args):
# Select backend
backend = select_sglang_backend(args)
sgl.set_default_backend(backend)
state = regex_gen.run(temperature=0)
print("=" * 20, "IP TEST", "=" * 20)
print(state.text())
state = json_gen.run(temperature=0)
print("=" * 20, "JSON TEST", "=" * 20)
print(state.text())
state = character_gen.run(temperature=0)
print("=" * 20, "CHARACTER TEST", "=" * 20)
print(state.text())
if __name__ == "__main__":
parser = argparse.ArgumentParser()
args = add_common_sglang_args_and_parse(parser)
main(args)
# ==================== IP TEST ====================
# Q: What is the IP address of the Google DNS servers?
# A: The google's DNS sever address is 8.8.8.8 and 8.8.4.4. The google's website domain name is www.google.com.
# ==================== JSON TEST ====================
# The information about Hogwarts is in the following JSON format.
# {
# "name": "Hogwarts School of Witchcraft and Wizardry",
# "country": "Scotland",
# "latitude": 55.566667,
# "population": 1000,
# "top 3 landmarks": ["Hogwarts Castle", "The Great Hall", "The Forbidden Forest"],
# }
# ==================== CHARACTER TEST ====================
# Give me a character description who is a wizard.
# { "name" : "Merlin", "age" : 500, "armor" : "chainmail" , "weapon" : "sword" , "strength" : 10 }

View File

@@ -0,0 +1,92 @@
import unittest
from sglang import OpenAI, set_default_backend
from sglang.test.test_programs import (
test_chat_completion_speculative,
test_completion_speculative,
test_decode_int,
test_decode_json,
test_expert_answer,
test_few_shot_qa,
test_image_qa,
test_mt_bench,
test_parallel_decoding,
test_parallel_encoding,
test_react,
test_select,
test_stream,
test_tool_use,
)
from sglang.test.test_utils import CustomTestCase
class TestOpenAIBackend(CustomTestCase):
instruct_backend = None
chat_backend = None
chat_vision_backend = None
@classmethod
def setUpClass(cls):
cls.instruct_backend = OpenAI("gpt-3.5-turbo-instruct")
cls.chat_backend = OpenAI("gpt-3.5-turbo")
cls.chat_vision_backend = OpenAI("gpt-4-turbo")
def test_few_shot_qa(self):
set_default_backend(self.instruct_backend)
test_few_shot_qa()
def test_mt_bench(self):
set_default_backend(self.chat_backend)
test_mt_bench()
def test_select(self):
set_default_backend(self.instruct_backend)
test_select(check_answer=True)
def test_decode_int(self):
set_default_backend(self.instruct_backend)
test_decode_int()
def test_decode_json(self):
set_default_backend(self.instruct_backend)
test_decode_json()
def test_expert_answer(self):
set_default_backend(self.instruct_backend)
test_expert_answer()
def test_tool_use(self):
set_default_backend(self.instruct_backend)
test_tool_use()
def test_react(self):
set_default_backend(self.instruct_backend)
test_react()
def test_parallel_decoding(self):
set_default_backend(self.instruct_backend)
test_parallel_decoding()
def test_parallel_encoding(self):
set_default_backend(self.instruct_backend)
test_parallel_encoding()
def test_image_qa(self):
set_default_backend(self.chat_vision_backend)
test_image_qa()
def test_stream(self):
set_default_backend(self.instruct_backend)
test_stream()
def test_completion_speculative(self):
set_default_backend(self.instruct_backend)
test_completion_speculative()
def test_chat_completion_speculative(self):
set_default_backend(self.chat_backend)
test_chat_completion_speculative()
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,68 @@
"""
Tests for the separate_reasoning functionality in sglang.
Usage:
python3 -m unittest test/lang/test_separate_reasoning.py
"""
import unittest
from sglang import gen, separate_reasoning
from sglang.lang.ir import SglExprList, SglSeparateReasoning
from sglang.test.test_utils import CustomTestCase
class TestSeparateReasoning(CustomTestCase):
def test_separate_reasoning_creation(self):
"""Test that SglSeparateReasoning objects are created correctly."""
# Test with valid model type and gen expression
test_gen = gen("test")
expr = separate_reasoning(test_gen, model_type="deepseek-r1")
self.assertIsInstance(expr, SglExprList)
self.assertEqual(len(expr.expr_list), 2)
self.assertEqual(expr.expr_list[0], test_gen)
reasoning_expr = expr.expr_list[1]
self.assertIsInstance(reasoning_expr, SglSeparateReasoning)
self.assertEqual(reasoning_expr.model_type, "deepseek-r1")
self.assertEqual(reasoning_expr.name, "test_reasoning_content")
# Test with another valid model type
expr = separate_reasoning(test_gen, model_type="qwen3")
self.assertIsInstance(expr, SglExprList)
self.assertEqual(expr.expr_list[1].model_type, "qwen3")
def test_separate_reasoning_name_processing(self):
"""Test that separate_reasoning correctly processes names."""
test_gen = gen("test_var")
expr = separate_reasoning(test_gen, model_type="deepseek-r1")
reasoning_expr = expr.expr_list[1]
self.assertEqual(reasoning_expr.name, "test_var_reasoning_content")
# Test the process_name_for_reasoning method
self.assertEqual(
reasoning_expr.process_name_for_reasoning("another_var"),
"another_var_reasoning_content",
)
def test_separate_reasoning_repr(self):
"""Test the string representation of SglSeparateReasoning."""
test_gen = gen("test_var")
expr = separate_reasoning(test_gen, model_type="deepseek-r1")
reasoning_expr = expr.expr_list[1]
self.assertEqual(
repr(reasoning_expr),
"SeparateReasoning(model_type=deepseek-r1, name=test_var_reasoning_content)",
)
def test_separate_reasoning_with_invalid_model_type(self):
"""Test that separate_reasoning accepts any model type during creation."""
# Create with invalid model type
test_gen = gen("test")
expr = separate_reasoning(test_gen, model_type="invalid-model")
self.assertIsInstance(expr, SglExprList)
self.assertEqual(expr.expr_list[1].model_type, "invalid-model")
# The actual validation happens in the ReasoningParser constructor
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,193 @@
"""
Tests for the execution of separate_reasoning functionality in sglang.
Usage:
python3 -m unittest test/lang/test_separate_reasoning_execution.py
"""
import threading
import unittest
from unittest.mock import MagicMock, patch
from sglang.lang.interpreter import StreamExecutor
from sglang.lang.ir import SglGen, SglSeparateReasoning
from sglang.test.test_utils import CustomTestCase
# Helper function to create events that won't block program exit
def create_daemon_event():
event = threading.Event()
return event
class MockReasoningParser:
def __init__(self, model_type):
self.model_type = model_type
self.parse_non_stream_called = False
self.parse_stream_chunk_called = False
def parse_non_stream(self, full_text):
self.parse_non_stream_called = True
# Simulate parsing by adding a prefix to indicate reasoning
reasoning = f"[REASONING from {self.model_type}]: {full_text}"
normal_text = f"[NORMAL from {self.model_type}]: {full_text}"
return reasoning, normal_text
def parse_stream_chunk(self, chunk_text):
self.parse_stream_chunk_called = True
# Simulate parsing by adding a prefix to indicate reasoning
reasoning = f"[REASONING from {self.model_type}]: {chunk_text}"
normal_text = f"[NORMAL from {self.model_type}]: {chunk_text}"
return reasoning, normal_text
class TestSeparateReasoningExecution(CustomTestCase):
def setUp(self):
"""Set up for the test."""
super().setUp()
# Store any events created during the test
self.events = []
def tearDown(self):
"""Clean up any threads that might have been created during the test."""
super().tearDown()
# Set all events to ensure any waiting threads are released
for event in self.events:
event.set()
def tearDown(self):
super().tearDown()
# wake up all threads
for ev in self.events:
ev.set()
@patch("sglang.srt.parser.reasoning_parser.ReasoningParser")
def test_execute_separate_reasoning(self, mock_parser_class):
"""Test that _execute_separate_reasoning correctly calls the ReasoningParser."""
# Setup mock parser
mock_parser = MockReasoningParser("deepseek-r1")
mock_parser_class.return_value = mock_parser
# Create a mock backend to avoid AttributeError in __del__
mock_backend = MagicMock()
# Create a StreamExecutor with necessary setup
executor = StreamExecutor(
backend=mock_backend,
arguments={},
default_sampling_para={},
chat_template={
"role_map": {"user": "user", "assistant": "assistant"}
}, # Simple chat template
stream=False,
use_thread=False,
)
# Set up the executor with a variable and its value
var_name = "test_var"
reasoning_name = f"{var_name}_reasoning_content"
var_value = "Test content"
executor.variables = {var_name: var_value}
# Create events and track them for cleanup
var_event = create_daemon_event()
reasoning_event = create_daemon_event()
self.events.extend([var_event, reasoning_event])
executor.variable_event = {var_name: var_event, reasoning_name: reasoning_event}
executor.variable_event[var_name].set() # Mark as ready
# Set up the current role
executor.cur_role = "assistant"
executor.cur_role_begin_pos = 0
executor.text_ = var_value
# Create a gen expression and a separate_reasoning expression
gen_expr = SglGen(var_name)
expr = SglSeparateReasoning("deepseek-r1", expr=gen_expr)
# Execute separate_reasoning
executor._execute_separate_reasoning(expr)
# Verify that the parser was created with the correct model type
mock_parser_class.assert_called_once_with("deepseek-r1")
# Verify that parse_non_stream was called
self.assertTrue(mock_parser.parse_non_stream_called)
# Verify that the variables were updated correctly
reasoning_name = f"{var_name}_reasoning_content"
self.assertIn(reasoning_name, executor.variables)
self.assertEqual(
executor.variables[reasoning_name],
f"[REASONING from deepseek-r1]: {var_value}",
)
self.assertEqual(
executor.variables[var_name], f"[NORMAL from deepseek-r1]: {var_value}"
)
# Verify that the variable event was set
self.assertIn(reasoning_name, executor.variable_event)
self.assertTrue(executor.variable_event[reasoning_name].is_set())
# Verify that the text was updated
self.assertEqual(executor.text_, f"[NORMAL from deepseek-r1]: {var_value}")
@patch("sglang.srt.parser.reasoning_parser.ReasoningParser")
def test_reasoning_parser_integration(self, mock_parser_class):
"""Test the integration between separate_reasoning and ReasoningParser."""
# Setup mock parsers for different model types
deepseek_parser = MockReasoningParser("deepseek-r1")
qwen_parser = MockReasoningParser("qwen3")
# Configure the mock to return different parsers based on model type
def get_parser(model_type):
if model_type == "deepseek-r1":
return deepseek_parser
elif model_type == "qwen3":
return qwen_parser
else:
raise ValueError(f"Unsupported model type: {model_type}")
mock_parser_class.side_effect = get_parser
# Test with DeepSeek-R1 model
test_text = "This is a test"
reasoning, normal_text = deepseek_parser.parse_non_stream(test_text)
self.assertEqual(reasoning, f"[REASONING from deepseek-r1]: {test_text}")
self.assertEqual(normal_text, f"[NORMAL from deepseek-r1]: {test_text}")
# Test with Qwen3 model
reasoning, normal_text = qwen_parser.parse_non_stream(test_text)
self.assertEqual(reasoning, f"[REASONING from qwen3]: {test_text}")
self.assertEqual(normal_text, f"[NORMAL from qwen3]: {test_text}")
@patch("sglang.srt.parser.reasoning_parser.ReasoningParser")
def test_reasoning_parser_invalid_model(self, mock_parser_class):
"""Test that ReasoningParser raises an error for invalid model types."""
# Configure the mock to raise an error for invalid model types
def get_parser(model_type):
if model_type in ["deepseek-r1", "qwen3"]:
return MockReasoningParser(model_type)
elif model_type is None:
raise ValueError("Model type must be specified")
else:
raise ValueError(f"Unsupported model type: {model_type}")
mock_parser_class.side_effect = get_parser
with self.assertRaises(ValueError) as context:
mock_parser_class("invalid-model")
self.assertIn("Unsupported model type", str(context.exception))
with self.assertRaises(ValueError) as context:
mock_parser_class(None)
self.assertIn("Model type must be specified", str(context.exception))
if __name__ == "__main__":
unittest.main()

View File

@@ -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.")

View File

@@ -0,0 +1,191 @@
import torch
from sglang.srt.layers.attention.nsa.index_buf_accessor import (
_get_k_and_s_triton_kernel,
)
def golden_torch_gen(
seq_len_tensor: torch.Tensor,
buffer_indexer: torch.Tensor,
buffer: torch.Tensor,
index_head_dim,
page_size,
):
dim_split = page_size * index_head_dim
torch_k_out = buffer[:, 0:dim_split]
torch_s_out = buffer[:, dim_split:]
torch_k_out = torch_k_out.reshape(-1, 128)
torch_s_out = torch_s_out.reshape(-1, 4)
batch = seq_len_tensor.shape[0]
index_list = []
for i in range(batch):
seq_len = seq_len_tensor[i].item()
buffer_index_ = buffer_indexer[i]
align_seq_len = ((seq_len + page_size - 1) / page_size) * page_size
needed_block_num = int((seq_len + page_size - 1) / page_size)
for j in range(needed_block_num):
block_idx = buffer_index_[j].item()
start_idx = block_idx * page_size
end_idx = 0
if j == (needed_block_num - 1):
end_idx = block_idx * page_size + (
seq_len - (needed_block_num - 1) * page_size
)
else:
end_idx = (block_idx + 1) * page_size
index_tensor = (
torch.arange(start=start_idx, end=end_idx, step=1)
.type(torch.int32)
.cuda()
)
index_list.append(index_tensor)
index_list_ = torch.cat(index_list, dim=0)
torch_k_out = torch.index_select(torch_k_out, dim=0, index=index_list_)
torch_s_out = torch.index_select(torch_s_out, dim=0, index=index_list_)
return torch_k_out, torch_s_out
def get_k_and_s_triton():
index_head_dim = 128
page_size = 64
num_page = 128
s_offset_in_page = page_size * index_head_dim
seq_len_tensor = torch.tensor(
[256, 267, 215, 32, 129], dtype=torch.int64, device="cuda"
) # 4 + 5 + 3 + 1 + 3 block
buffer_indexer = torch.tensor(
[
[1, 2, 3, 4, 0],
[7, 6, 5, 8, 9],
[10, 11, 12, 0, 0],
[13, 0, 0, 0, 0],
[14, 15, 16, 0, 0],
],
dtype=torch.int32,
device="cuda",
)
seq_len_sum = seq_len_tensor.sum()
batch = seq_len_tensor.shape[0]
triton_k_out = torch.empty(
(seq_len_sum, index_head_dim), dtype=torch.uint8, device="cuda"
)
triton_s_out = torch.empty((seq_len_sum, 4), dtype=torch.uint8, device="cuda")
buffer = torch.randint(
0,
num_page,
(num_page, page_size * index_head_dim + page_size * 4),
device="cuda",
).type(torch.uint8)
_, buf_numel_per_page = buffer.shape
_, page_indice_batch_offset = buffer_indexer.shape
max_seq_len = seq_len_tensor.max().item()
BLOCK_SIZE = 256
BLOCK_SIZE_K = 128
num_token_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
num_k_threads = (index_head_dim + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
grid = (batch, num_token_blocks, num_k_threads)
seq_num_pow2 = 1
while seq_num_pow2 < batch:
seq_num_pow2 *= 2
# acc test =====================
_get_k_and_s_triton_kernel[grid](
buf_ptr=buffer,
page_indices_ptr=buffer_indexer,
k_out_ptr=triton_k_out,
s_out_ptr=triton_s_out,
seq_len_ptr=seq_len_tensor,
seq_len_num_pow=seq_num_pow2,
page_size=page_size,
buf_numel_per_page=buf_numel_per_page,
index_head_dim=index_head_dim,
s_offset_in_page=s_offset_in_page,
page_indice_batch_offset=page_indice_batch_offset,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_SIZE_K=BLOCK_SIZE_K,
)
torch_k_out, torch_s_out = golden_torch_gen(
seq_len_tensor=seq_len_tensor,
buffer_indexer=buffer_indexer,
buffer=buffer,
index_head_dim=index_head_dim,
page_size=page_size,
)
torch.testing.assert_close(
triton_k_out, torch_k_out, rtol=0, atol=0, msg="k outputs differ!"
)
torch.testing.assert_close(
triton_s_out, torch_s_out, rtol=0, atol=0, msg="s outputs differ!"
)
print("_get_k_and_s_triton_kernel test pass")
# perf test =====================
import time
torch.cuda.synchronize()
for _ in range(10):
_get_k_and_s_triton_kernel[grid](
buf_ptr=buffer,
page_indices_ptr=buffer_indexer,
k_out_ptr=triton_k_out,
s_out_ptr=triton_s_out,
seq_len_ptr=seq_len_tensor,
seq_len_num_pow=seq_num_pow2,
page_size=page_size,
buf_numel_per_page=buf_numel_per_page,
index_head_dim=index_head_dim,
s_offset_in_page=s_offset_in_page,
page_indice_batch_offset=page_indice_batch_offset,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_SIZE_K=BLOCK_SIZE_K,
)
torch.cuda.synchronize()
start_time = time.perf_counter()
_get_k_and_s_triton_kernel[grid](
buf_ptr=buffer,
page_indices_ptr=buffer_indexer,
k_out_ptr=triton_k_out,
s_out_ptr=triton_s_out,
seq_len_ptr=seq_len_tensor,
seq_len_num_pow=seq_num_pow2,
page_size=page_size,
buf_numel_per_page=buf_numel_per_page,
index_head_dim=index_head_dim,
s_offset_in_page=s_offset_in_page,
page_indice_batch_offset=page_indice_batch_offset,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_SIZE_K=BLOCK_SIZE_K,
)
end_time = time.perf_counter()
print(
f"_get_k_and_s_triton_kernel triton kernel infer time is {((end_time-start_time)*1000):.4f} ms\n"
)
if __name__ == "__main__":
if not torch.cuda.is_available():
print("CUDA not available. Skipping tests.")
exit(0)
print("Start test cases...\n")
get_k_and_s_triton()
print("End test cases...\n")

View File

@@ -0,0 +1,591 @@
"""
Correctness tests for NSA Indexer K/S Buffer Access with Fused Triton Kernels.
This test verifies that the optimized Triton implementations (GetK, GetS, GetKAndS)
produce identical results to the torch_fast baseline implementations.
Test coverage:
- GetK.triton() vs GetK.torch_fast()
- GetS.triton() vs GetS.torch_fast()
- GetKAndS.triton() vs separate GetK.torch_fast() + GetS.torch_fast()
"""
import pytest
import torch
from sglang.srt.layers.attention.nsa.index_buf_accessor import GetK, GetKAndS, GetS
class MockNSATokenToKVPool:
"""Mock pool object that mimics NSATokenToKVPool for testing."""
def __init__(
self,
page_size: int = 64,
index_head_dim: int = 128,
quant_block_size: int = 128,
device: str = "cuda",
):
self.page_size = page_size
self.index_head_dim = index_head_dim
self.quant_block_size = quant_block_size
self.device = device
def create_test_buffer(
num_pages: int,
page_size: int = 64,
index_head_dim: int = 128,
device: str = "cuda",
) -> torch.Tensor:
"""
Create a test buffer mimicking the K/S buffer structure.
Buffer layout per page:
- First page_size * index_head_dim bytes: K data (fp8, stored as uint8)
- Next page_size * 4 bytes: S data (fp32 scales, stored as uint8)
Args:
num_pages: Number of pages to allocate
page_size: Tokens per page (typically 64)
index_head_dim: Dimension of K vectors (typically 128)
device: Device to allocate on
Returns:
Buffer of shape (num_pages, page_size * index_head_dim + page_size * 4)
"""
buf_numel_per_page = page_size * index_head_dim + page_size * 4
buf = torch.randint(
0, 256, (num_pages, buf_numel_per_page), dtype=torch.uint8, device=device
)
return buf
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
class TestGetK:
"""Test cases for GetK.triton() correctness."""
@pytest.mark.parametrize("num_pages", [1, 2, 4, 8, 16])
@pytest.mark.parametrize("seq_len", [64, 128, 256, 512, 1024])
@pytest.mark.parametrize("page_size", [64])
@pytest.mark.parametrize("index_head_dim", [128])
def test_getk_correctness(self, num_pages, seq_len, page_size, index_head_dim):
"""Test GetK.triton() produces same output as GetK.torch_fast()."""
device = torch.device("cuda")
# Ensure seq_len doesn't exceed available pages
max_seq_len = num_pages * page_size
seq_len = min(seq_len, max_seq_len)
# Create mock pool
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
# Create test buffer
buf = create_test_buffer(
num_pages=num_pages,
page_size=page_size,
index_head_dim=index_head_dim,
device=device,
)
# Create page indices
num_pages_needed = (seq_len + page_size - 1) // page_size
page_indices = torch.randint(
0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device
)
# Run both implementations
output_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
output_triton = GetK.triton(pool, buf, seq_len, page_indices)
# Verify shapes
assert output_torch.shape == (seq_len, index_head_dim)
assert output_triton.shape == (seq_len, index_head_dim)
assert output_torch.dtype == torch.uint8
assert output_triton.dtype == torch.uint8
# Compare results (should be exact match)
torch.testing.assert_close(
output_triton, output_torch, rtol=0, atol=0, msg="GetK outputs differ"
)
def test_getk_sequential_pages(self):
"""Test GetK with sequential page indices."""
device = torch.device("cuda")
page_size = 64
index_head_dim = 128
num_pages = 10
seq_len = 320 # 5 pages
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
# Sequential page indices [0, 1, 2, 3, 4]
page_indices = torch.arange(5, dtype=torch.int32, device=device)
output_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
output_triton = GetK.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(output_triton, output_torch, rtol=0, atol=0)
def test_getk_repeated_pages(self):
"""Test GetK with repeated page indices."""
device = torch.device("cuda")
page_size = 64
index_head_dim = 128
num_pages = 5
seq_len = 192 # 3 pages
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
# Repeated page indices [2, 2, 2]
page_indices = torch.full((3,), 2, dtype=torch.int32, device=device)
output_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
output_triton = GetK.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(output_triton, output_torch, rtol=0, atol=0)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
class TestGetS:
"""Test cases for GetS.triton() correctness."""
@pytest.mark.parametrize("num_pages", [1, 2, 4, 8, 16])
@pytest.mark.parametrize("seq_len", [64, 128, 256, 512, 1024])
@pytest.mark.parametrize("page_size", [64])
@pytest.mark.parametrize("index_head_dim", [128])
def test_gets_correctness(self, num_pages, seq_len, page_size, index_head_dim):
"""Test GetS.triton() produces same output as GetS.torch_fast()."""
device = torch.device("cuda")
# Ensure seq_len doesn't exceed available pages
max_seq_len = num_pages * page_size
seq_len = min(seq_len, max_seq_len)
# Create mock pool
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
# Create test buffer
buf = create_test_buffer(
num_pages=num_pages,
page_size=page_size,
index_head_dim=index_head_dim,
device=device,
)
# Create page indices
num_pages_needed = (seq_len + page_size - 1) // page_size
page_indices = torch.randint(
0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device
)
# Run both implementations
output_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
output_triton = GetS.triton(pool, buf, seq_len, page_indices)
# Verify shapes
assert output_torch.shape == (seq_len, 4)
assert output_triton.shape == (seq_len, 4)
assert output_torch.dtype == torch.uint8
assert output_triton.dtype == torch.uint8
# Compare results (should be exact match)
torch.testing.assert_close(
output_triton, output_torch, rtol=0, atol=0, msg="GetS outputs differ"
)
def test_gets_sequential_pages(self):
"""Test GetS with sequential page indices."""
device = torch.device("cuda")
page_size = 64
index_head_dim = 128
num_pages = 10
seq_len = 320 # 5 pages
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
# Sequential page indices [0, 1, 2, 3, 4]
page_indices = torch.arange(5, dtype=torch.int32, device=device)
output_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
output_triton = GetS.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(output_triton, output_torch, rtol=0, atol=0)
def test_gets_repeated_pages(self):
"""Test GetS with repeated page indices."""
device = torch.device("cuda")
page_size = 64
index_head_dim = 128
num_pages = 5
seq_len = 192 # 3 pages
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
# Repeated page indices [2, 2, 2]
page_indices = torch.full((3,), 2, dtype=torch.int32, device=device)
output_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
output_triton = GetS.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(output_triton, output_torch, rtol=0, atol=0)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
class TestGetKAndS:
"""Test cases for GetKAndS.triton() correctness."""
@pytest.mark.parametrize("num_pages", [1, 2, 4, 8, 16])
@pytest.mark.parametrize("seq_len", [64, 128, 256, 512, 1024])
@pytest.mark.parametrize("page_size", [64])
@pytest.mark.parametrize("index_head_dim", [128])
def test_get_k_and_s_correctness(
self, num_pages, seq_len, page_size, index_head_dim
):
"""Test GetKAndS.triton() produces same output as separate torch_fast calls."""
device = torch.device("cuda")
# Ensure seq_len doesn't exceed available pages
max_seq_len = num_pages * page_size
seq_len = min(seq_len, max_seq_len)
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
# Create mock pool
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
# Create test buffer
buf = create_test_buffer(
num_pages=num_pages,
page_size=page_size,
index_head_dim=index_head_dim,
device=device,
)
# Create page indices
num_pages_needed = (seq_len + page_size - 1) // page_size
page_indices = torch.randint(
0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device
)
page_indices_ = page_indices.unsqueeze(0)
# Run baseline: separate torch_fast calls
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
# Run fused Triton implementation
k_triton, s_triton = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
# Verify shapes
assert k_torch.shape == (seq_len, index_head_dim)
assert s_torch.shape == (seq_len, 4)
assert k_triton.shape == (seq_len, index_head_dim)
assert s_triton.shape == (seq_len, 4)
# Verify dtypes
assert k_torch.dtype == torch.uint8
assert s_torch.dtype == torch.uint8
assert k_triton.dtype == torch.uint8
assert s_triton.dtype == torch.uint8
# Compare K results
torch.testing.assert_close(
k_triton, k_torch, rtol=0, atol=0, msg="GetKAndS K outputs differ"
)
# Compare S results
torch.testing.assert_close(
s_triton, s_torch, rtol=0, atol=0, msg="GetKAndS S outputs differ"
)
def test_get_k_and_s_sequential_pages(self):
"""Test GetKAndS with sequential page indices."""
device = torch.device("cuda")
page_size = 64
index_head_dim = 128
num_pages = 10
seq_len = 320 # 5 pages
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
# Sequential page indices [0, 1, 2, 3, 4]
page_indices = torch.arange(5, dtype=torch.int32, device=device)
page_indices_ = page_indices.unsqueeze(0)
# Baseline
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
# Fused
k_triton, s_triton = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
def test_get_k_and_s_repeated_pages(self):
"""Test GetKAndS with repeated page indices."""
device = torch.device("cuda")
page_size = 64
index_head_dim = 128
num_pages = 5
seq_len = 192 # 3 pages
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
# Repeated page indices [2, 2, 2]
page_indices = torch.full((3,), 2, dtype=torch.int32, device=device)
page_indices_ = page_indices.unsqueeze(0)
# Baseline
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
# Fused
k_triton, s_triton = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
def test_get_k_and_s_partial_page(self):
"""Test GetKAndS when seq_len is not a multiple of page_size."""
device = torch.device("cuda")
page_size = 64
index_head_dim = 128
num_pages = 5
seq_len = 100 # Not a multiple of 64
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
num_pages_needed = (seq_len + page_size - 1) // page_size
page_indices = torch.arange(num_pages_needed, dtype=torch.int32, device=device)
page_indices_ = page_indices.unsqueeze(0)
# Baseline
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
# Fused
k_triton, s_triton = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
# Should handle partial pages correctly
torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
class TestEdgeCases:
"""Test edge cases and boundary conditions."""
def test_single_token(self):
"""Test with seq_len=1 (single token)."""
device = torch.device("cuda")
page_size = 64
index_head_dim = 128
num_pages = 2
seq_len = 1
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
page_indices = torch.tensor([0], dtype=torch.int32, device=device)
page_indices_ = page_indices.unsqueeze(0)
# Test GetK
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
k_triton = GetK.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
# Test GetS
s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
s_triton = GetS.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
# Test GetKAndS
k_triton2, s_triton2 = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0)
def test_exact_page_boundary(self):
"""Test when seq_len exactly matches page boundaries."""
device = torch.device("cuda")
page_size = 64
index_head_dim = 128
num_pages = 5
seq_len = 192 # Exactly 3 pages
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
page_indices = torch.arange(3, dtype=torch.int32, device=device)
page_indices_ = page_indices.unsqueeze(0)
# Test GetK
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
k_triton = GetK.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
# Test GetS
s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
s_triton = GetS.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
# Test GetKAndS
k_triton2, s_triton2 = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0)
def test_large_seq_len(self):
"""Test with large sequence length."""
device = torch.device("cuda")
page_size = 64
index_head_dim = 128
num_pages = 100
seq_len = 4096 # 64 pages
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
num_pages_needed = (seq_len + page_size - 1) // page_size
page_indices = torch.randint(
0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device
)
page_indices_ = page_indices.unsqueeze(0)
# Test GetK
k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
k_triton = GetK.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
# Test GetS
s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
s_triton = GetS.triton(pool, buf, seq_len, page_indices)
torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
# Test GetKAndS
k_triton2, s_triton2 = GetKAndS.triton(
pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
)
torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0)
torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0)
def print_test_summary():
"""Print a summary message about the test suite."""
print("\n" + "=" * 80)
print("NSA Indexer K/S Buffer Accessor Correctness Tests")
print("=" * 80)
print("Testing Triton implementations against torch_fast baseline:")
print(" - GetK.triton() vs GetK.torch_fast()")
print(" - GetS.triton() vs GetS.torch_fast()")
print(" - GetKAndS.triton() vs separate GetK/GetS torch_fast() calls")
print("=" * 80)
print()
if __name__ == "__main__":
# Run tests manually
if not torch.cuda.is_available():
print("CUDA not available. Skipping tests.")
exit(0)
print_test_summary()
# Run a few sample tests
print("Running sample correctness tests...\n")
# Test GetK
print("Testing GetK...")
test_getk = TestGetK()
test_getk.test_getk_correctness(
num_pages=4, seq_len=256, page_size=64, index_head_dim=128
)
test_getk.test_getk_sequential_pages()
print("✓ GetK tests passed\n")
# Test GetS
print("Testing GetS...")
test_gets = TestGetS()
test_gets.test_gets_correctness(
num_pages=4, seq_len=256, page_size=64, index_head_dim=128
)
test_gets.test_gets_sequential_pages()
print("✓ GetS tests passed\n")
# Test GetKAndS
print("Testing GetKAndS SeqLen=256...")
test_get_k_and_s = TestGetKAndS()
test_get_k_and_s.test_get_k_and_s_correctness(
num_pages=4, seq_len=256, page_size=64, index_head_dim=128
)
test_get_k_and_s.test_get_k_and_s_sequential_pages()
test_get_k_and_s.test_get_k_and_s_partial_page()
print("✓ GetKAndS SeqLen=256 tests passed\n")
print("Testing GetKAndS SeqLen=128K...")
test_get_k_and_s = TestGetKAndS()
test_get_k_and_s.test_get_k_and_s_correctness(
num_pages=2048, seq_len=131072, page_size=64, index_head_dim=128
)
test_get_k_and_s.test_get_k_and_s_sequential_pages()
test_get_k_and_s.test_get_k_and_s_partial_page()
print("✓ GetKAndS SeqLen=128K tests passed\n")
# Test edge cases
print("Testing edge cases...")
test_edge = TestEdgeCases()
test_edge.test_single_token()
test_edge.test_exact_page_boundary()
test_edge.test_large_seq_len()
print("✓ Edge case tests passed\n")
print("=" * 80)
print("All correctness tests passed successfully!")
print("=" * 80)

View File

@@ -0,0 +1,211 @@
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_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_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
class TestMoERunner(CustomTestCase):
BASE_URL = DEFAULT_URL_FOR_TEST
TIMEOUT = 6000
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_cutlass_fp8": {
"model": DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE,
"timeout": 3600,
"other_args": [
"--trust-remote-code",
"--moe-runner-backend",
"cutlass",
"--attention-backend",
"triton",
"--sampling-backend",
"pytorch",
"--disable-cuda-graph",
],
},
"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
env = dict(os.environ)
env["SGLANG_ENABLE_JIT_DEEPGEMM"] = "1"
env["SGLANG_JIT_DEEPGEMM_PRECOMPILE"] = "0"
env.update(config.get("env_overrides", {}))
timeout = config.get("timeout", self.TIMEOUT)
process = popen_launch_server(
model,
self.BASE_URL,
timeout=timeout,
other_args=other_args,
env=env,
)
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()

View File

@@ -0,0 +1,116 @@
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_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
class TestMoERunner4GPU(CustomTestCase):
BASE_URL = DEFAULT_URL_FOR_TEST
TIMEOUT = 6000
DEFAULT_EVAL_KWARGS = {
"eval_name": "mmlu",
"num_examples": 5,
"num_threads": 1,
}
CONFIGS = {
"moe_runner_cutlass_w4a8": {
"model": "tencent/DeepSeek-V3.1-Terminus-W4AFP8", # FP8 W8A8 MoE model
"other_args": [
"--trust-remote-code",
"--moe-runner-backend",
"cutlass",
"--attention-backend",
"triton",
"--sampling-backend",
"pytorch",
"--tp-size",
"4",
],
},
"moe_runner_cutlass_w4a8_deepep_normal": {
"model": "tencent/DeepSeek-V3.1-Terminus-W4AFP8", # FP8 W8A8 MoE model
"other_args": [
"--trust-remote-code",
"--moe-runner-backend",
"cutlass",
"--moe-a2a-backend",
"deepep",
"--deepep-mode",
"normal",
"--attention-backend",
"triton",
"--sampling-backend",
"pytorch",
"--tp-size",
"4",
],
},
"moe_runner_cutlass_w4a8_deepep_ll": {
"model": "tencent/DeepSeek-V3.1-Terminus-W4AFP8", # FP8 W8A8 MoE model
"env_overrides": {"SGLANG_DEEPEP_BF16_DISPATCH": "1"},
"other_args": [
"--trust-remote-code",
"--moe-runner-backend",
"cutlass",
"--moe-a2a-backend",
"deepep",
"--deepep-mode",
"low_latency",
"--attention-backend",
"triton",
"--sampling-backend",
"pytorch",
"--tp-size",
"4",
],
},
}
def _run_config(self, config: dict) -> None:
model = config["model"]
other_args = config.get("other_args", [])
eval_kwargs = self.DEFAULT_EVAL_KWARGS
env = dict(os.environ)
env["SGLANG_ENABLE_JIT_DEEPGEMM"] = "1"
env["SGLANG_JIT_DEEPGEMM_PRECOMPILE"] = "0"
env.update(config.get("env_overrides", {}))
timeout = config.get("timeout", self.TIMEOUT)
process = popen_launch_server(
model,
self.BASE_URL,
timeout=timeout,
other_args=other_args,
env=env,
)
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 TestMoERunner4GPU.CONFIGS.items():
setattr(
TestMoERunner4GPU,
f"test_{_name}",
(lambda self, cfg=_cfg: self._run_config(cfg)),
)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,107 @@
# 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 sglang.test.lora_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")

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()

View File

@@ -0,0 +1,218 @@
import random
import unittest
import torch
from sglang.srt.lora.torch_ops.lora_ops import sgemm_lora_a_fwd, sgemm_lora_b_fwd
from sglang.test.lora_utils import reference_sgmv_expand, reference_sgmv_shrink
from sglang.test.test_utils import CustomTestCase
class TestLoraOps(CustomTestCase):
def test_sgemm_lora_a_fwd(self):
batch_size = 2
input_dim = 1024
num_loras = 3
dtype = torch.float32
possible_lora_ranks = [8, 16, 32, 64, 128, 256]
lora_ranks = random.sample(
possible_lora_ranks,
counts=[num_loras] * len(possible_lora_ranks),
k=num_loras,
)
max_lora_rank = max(lora_ranks)
possible_lora_scaling = [0.25, 0.5, 1.0, 2.0, 4.0]
lora_scaling = random.sample(
possible_lora_scaling,
counts=[num_loras] * len(possible_lora_scaling),
k=num_loras,
)
inputs = torch.randn(batch_size, input_dim, dtype=dtype)
lora_a_weights = torch.randn(num_loras, max_lora_rank, input_dim, dtype=dtype)
lora_indices_tensor = torch.randint(
num_loras, (batch_size,), dtype=torch.int32, device="cpu"
)
seq_len_tensor = torch.ones(batch_size, dtype=torch.int32, device="cpu")
lora_ranks_tensor = torch.tensor(lora_ranks, dtype=torch.int32, device="cpu")
lora_scaling_tensor = torch.tensor(
lora_scaling, dtype=torch.float16, device="cpu"
)
expect_output = reference_sgmv_shrink(
inputs,
lora_a_weights,
lora_indices_tensor,
seq_len_tensor,
lora_ranks_tensor,
lora_scaling_tensor,
)
actual_output = sgemm_lora_a_fwd(
inputs,
lora_a_weights,
lora_indices_tensor,
seq_len_tensor,
lora_ranks_tensor,
lora_scaling_tensor,
)
self.assertTrue(torch.allclose(actual_output, expect_output))
def test_sgemm_lora_b_fwd(self):
batch_size = 2
output_dim = 1024
num_loras = 3
dtype = torch.float32
possible_lora_ranks = [8, 16, 32, 64, 128, 256]
lora_ranks = random.sample(
possible_lora_ranks,
counts=[num_loras] * len(possible_lora_ranks),
k=num_loras,
)
max_lora_rank = max(lora_ranks)
inputs = torch.randn(batch_size, max_lora_rank, dtype=dtype)
lora_b_weights = torch.randn(num_loras, output_dim, max_lora_rank, dtype=dtype)
lora_ranks_tensor = torch.tensor(lora_ranks, dtype=torch.int32, device="cpu")
seq_len_tensor = torch.ones(batch_size, dtype=torch.int32, device="cpu")
lora_indices_tensor = torch.randint(
num_loras, (batch_size,), dtype=torch.int32, device="cpu"
)
slice_offsets = torch.tensor([0, output_dim], dtype=torch.int32, device="cpu")
expect_output = reference_sgmv_expand(
inputs,
lora_b_weights,
lora_indices_tensor,
seq_len_tensor,
lora_ranks_tensor,
slice_offsets,
)
actual_output = sgemm_lora_b_fwd(
inputs,
lora_b_weights,
lora_indices_tensor,
seq_len_tensor,
lora_ranks_tensor,
slice_offsets,
)
self.assertTrue(torch.allclose(actual_output, expect_output))
def test_sgemm_lora_a_fwd_expand(self):
batch_size = 2
input_dim = 1024
num_loras = 3
dtype = torch.float32
possible_lora_ranks = [8, 16, 32, 64, 128, 256]
lora_ranks = random.sample(
possible_lora_ranks,
counts=[num_loras] * len(possible_lora_ranks),
k=num_loras,
)
max_lora_rank = max(lora_ranks)
possible_lora_scaling = [0.25, 0.5, 1.0, 2.0, 4.0]
lora_scaling = random.sample(
possible_lora_scaling,
counts=[num_loras] * len(possible_lora_scaling),
k=num_loras,
)
seq_len_tensor = torch.randint(
num_loras, (batch_size,), dtype=torch.int32, device="cpu"
)
seq_len = sum(seq_len_tensor)
inputs = torch.randn(seq_len, input_dim, dtype=dtype)
lora_a_weights = torch.randn(num_loras, max_lora_rank, input_dim, dtype=dtype)
lora_indices_tensor = torch.randint(
num_loras, (batch_size,), dtype=torch.int32, device="cpu"
)
lora_ranks_tensor = torch.tensor(lora_ranks, dtype=torch.int32, device="cpu")
lora_scaling_tensor = torch.tensor(
lora_scaling, dtype=torch.float16, device="cpu"
)
expect_output = reference_sgmv_shrink(
inputs,
lora_a_weights,
lora_indices_tensor,
seq_len_tensor,
lora_ranks_tensor,
lora_scaling_tensor,
)
actual_output = sgemm_lora_a_fwd(
inputs,
lora_a_weights,
lora_indices_tensor,
seq_len_tensor,
lora_ranks_tensor,
lora_scaling_tensor,
)
self.assertTrue(torch.allclose(actual_output, expect_output))
def test_sgemm_lora_b_fwd_expand(self):
batch_size = 2
output_dim = 1024
num_loras = 3
dtype = torch.float32
possible_lora_ranks = [8, 16, 32, 64, 128, 256]
lora_ranks = random.sample(
possible_lora_ranks,
counts=[num_loras] * len(possible_lora_ranks),
k=num_loras,
)
max_lora_rank = max(lora_ranks)
seq_len_tensor = torch.randint(
num_loras, (batch_size,), dtype=torch.int32, device="cpu"
)
seq_len = sum(seq_len_tensor)
inputs = torch.randn(seq_len, max_lora_rank, dtype=dtype)
lora_b_weights = torch.randn(num_loras, output_dim, max_lora_rank, dtype=dtype)
lora_ranks_tensor = torch.tensor(lora_ranks, dtype=torch.int32, device="cpu")
lora_indices_tensor = torch.randint(
num_loras, (batch_size,), dtype=torch.int32, device="cpu"
)
slice_offsets = torch.tensor([0, output_dim], dtype=torch.int32, device="cpu")
expect_output = reference_sgmv_expand(
inputs,
lora_b_weights,
lora_indices_tensor,
seq_len_tensor,
lora_ranks_tensor,
slice_offsets,
)
actual_output = sgemm_lora_b_fwd(
inputs,
lora_b_weights,
lora_indices_tensor,
seq_len_tensor,
lora_ranks_tensor,
slice_offsets,
)
self.assertTrue(torch.allclose(actual_output, expect_output))
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,52 @@
# Copyright 2023-2025 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
from sglang.test.lora_utils import (
CI_MULTI_LORA_MODELS,
LORA_MODELS_QWEN3,
run_lora_multiple_batch_on_model_cases,
)
from sglang.test.test_utils import CustomTestCase
class TestLoRASpecDecoding(CustomTestCase):
def test_qwen(self):
run_lora_multiple_batch_on_model_cases(
LORA_MODELS_QWEN3,
attention_backend="triton",
use_spec_decoding=True,
disable_cuda_graph=True,
enable_deterministic_inference=True,
)
def test_llama(self):
run_lora_multiple_batch_on_model_cases(
CI_MULTI_LORA_MODELS,
attention_backend="triton",
use_spec_decoding=True,
disable_cuda_graph=True,
enable_deterministic_inference=True,
)
if __name__ == "__main__":
try:
mp.set_start_method("spawn")
except RuntimeError:
pass
unittest.main(warnings="ignore")

View File

@@ -0,0 +1,246 @@
import unittest
import torch
from sglang.srt.lora.backend.torch_backend import TorchNativeLoRABackend
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.test.lora_utils import reference_sgmv_expand, reference_sgmv_shrink
from sglang.test.test_utils import CustomTestCase
class TestTorchNativeLoRABackend(CustomTestCase):
device = "cpu"
# set duplicate weights to test merging during prepare_lora_batch
weight_indices = [0, 0, 1]
lora_ranks = [1, 1]
scalings = [1.0, 0.5]
seq_lens = [1, 1, 1]
use_cuda_graph = False
forward_batch = ForwardBatch(
forward_mode=ForwardMode.EXTEND,
batch_size=3,
input_ids=torch.tensor([[1], [2], [3]], dtype=torch.int32),
req_pool_indices=None,
seq_lens=None,
out_cache_loc=None,
seq_lens_sum=3,
extend_seq_lens=torch.tensor(seq_lens, dtype=torch.int32),
extend_seq_lens_cpu=seq_lens,
)
@classmethod
def setUpClass(cls):
cls.backend = TorchNativeLoRABackend(max_loras_per_batch=2, device=cls.device)
cls.backend.prepare_lora_batch(
forward_batch=cls.forward_batch,
weight_indices=cls.weight_indices,
lora_ranks=cls.lora_ranks,
scalings=cls.scalings,
use_cuda_graph=cls.use_cuda_graph,
)
def test_run_lora_a_sgemm(self):
batch_size = 3
input_dim = 4
output_dim = 6
num_loras = 3
dtype = torch.float32
x = torch.randn(batch_size, input_dim, dtype=dtype)
weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
weight_indices_tensor = torch.tensor(
self.weight_indices, dtype=torch.int32, device=self.device
)
seg_len_tensor = torch.tensor(
self.seq_lens, dtype=torch.int32, device=self.device
)
lora_ranks_tensor = torch.tensor(
self.lora_ranks, dtype=torch.int32, device=self.device
)
scalings_tensor = torch.tensor(
self.scalings, dtype=torch.float, device=self.device
)
expect_output = reference_sgmv_shrink(
x,
weights,
weight_indices_tensor,
seg_len_tensor,
lora_ranks_tensor,
scalings_tensor,
)
actual_output = self.backend.run_lora_a_sgemm(x, weights)
self.assertTrue(torch.allclose(actual_output, expect_output))
def test_run_lora_b_sgemm(self):
batch_size = 3
input_dim = 6
output_dim = 4
num_loras = 3
dtype = torch.float32
x = torch.randn(batch_size, input_dim, dtype=dtype)
weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
_, weight_out_dim, _ = weights.shape
weight_indices_tensor = torch.tensor(
self.weight_indices, dtype=torch.int32, device=self.device
)
seg_len_tensor = torch.tensor(
self.seq_lens, dtype=torch.int32, device=self.device
)
lora_ranks_tensor = torch.tensor(
self.lora_ranks, dtype=torch.int32, device=self.device
)
expect_output = reference_sgmv_expand(
x,
weights,
weight_indices_tensor,
seg_len_tensor,
lora_ranks_tensor,
slice_offsets=torch.tensor(
[0, weight_out_dim], dtype=torch.int32, device="cpu"
),
)
actual_output = self.backend.run_lora_b_sgemm(x, weights)
self.assertTrue(torch.allclose(actual_output, expect_output))
def test_run_qkv_lora(self):
batch_size = 3
num_loras = 3
input_dim = 6
output_offset = [0, 3, 6, 9]
output_dim = output_offset[-1]
num_slices = len(output_offset) - 1 # 3 slices for Q, K, V
max_lora_rank = max(self.lora_ranks)
dtype = torch.float32
x = torch.randn(batch_size, input_dim, dtype=dtype)
output_offset_cpu = torch.tensor(output_offset, dtype=torch.int32)
qkv_lora_a = torch.randn(
num_loras, max_lora_rank * num_slices, input_dim, dtype=dtype
)
qkv_lora_b = torch.randn(
num_loras, output_dim, max_lora_rank * num_slices, dtype=dtype
)
weight_indices_tensor = torch.tensor(
self.weight_indices, dtype=torch.int32, device=self.device
)
seg_len_tensor = torch.tensor(
self.seq_lens, dtype=torch.int32, device=self.device
)
lora_ranks_tensor = torch.tensor(
self.lora_ranks, dtype=torch.int32, device=self.device
)
scalings_tensor = torch.tensor(
self.scalings, dtype=torch.float, device=self.device
)
expect_lora_a_output = reference_sgmv_shrink(
x,
qkv_lora_a,
weight_indices_tensor,
seg_len_tensor,
lora_ranks_tensor,
scalings_tensor,
num_slices,
)
expect_output = reference_sgmv_expand(
expect_lora_a_output,
qkv_lora_b,
weight_indices_tensor,
seg_len_tensor,
lora_ranks_tensor,
output_offset_cpu,
)
actual_output = self.backend.run_qkv_lora(
x, qkv_lora_a, qkv_lora_b, None, output_offset_cpu, 0
)
self.assertTrue(torch.allclose(actual_output, expect_output))
def test_run_gate_up_lora(self):
batch_size = 3
input_dim = 6
output_dim = 4
num_loras = 3
dtype = torch.float32
max_lora_rank = max(self.lora_ranks)
num_slices = 2
x = torch.randn(batch_size, input_dim, dtype=dtype)
gate_up_lora_a = torch.randn(
num_loras, max_lora_rank * num_slices, input_dim, dtype=dtype
)
gate_up_lora_b = torch.randn(
num_loras, output_dim, max_lora_rank * num_slices, dtype=dtype
)
_, weight_out_dim, _ = gate_up_lora_b.shape
slice_size = weight_out_dim // num_slices
output_offset = torch.tensor(
[0, slice_size, weight_out_dim], dtype=torch.int32, device="cpu"
)
weight_indices_tensor = torch.tensor(
self.weight_indices, dtype=torch.int32, device=self.device
)
seg_len_tensor = torch.tensor(
self.seq_lens, dtype=torch.int32, device=self.device
)
lora_ranks_tensor = torch.tensor(
self.lora_ranks, dtype=torch.int32, device=self.device
)
scalings_tensor = torch.tensor(
self.scalings, dtype=torch.float, device=self.device
)
expect_lora_a_output = reference_sgmv_shrink(
x,
gate_up_lora_a,
weight_indices_tensor,
seg_len_tensor,
lora_ranks_tensor,
scalings_tensor,
num_slices,
)
expect_output = reference_sgmv_expand(
expect_lora_a_output,
gate_up_lora_b,
weight_indices_tensor,
seg_len_tensor,
lora_ranks_tensor,
slice_offsets=output_offset,
)
actual_output = self.backend.run_gate_up_lora(x, gate_up_lora_a, gate_up_lora_b)
self.assertTrue(torch.allclose(actual_output, expect_output))
if __name__ == "__main__":
unittest.main()

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()

View File

@@ -0,0 +1,146 @@
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,
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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 0.74)

View File

@@ -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()

View File

@@ -0,0 +1,53 @@
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_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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=64,
num_threads=128,
)
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()

View File

@@ -0,0 +1,70 @@
import unittest
from types import SimpleNamespace
import requests
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,
CustomTestCase,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
class TestKimiK2Thinking(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = "moonshotai/Kimi-K2-Thinking"
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"8",
"--trust-remote-code",
"--tool-call-parser",
"kimi_k2",
"--reasoning-parser",
"kimi_k2",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
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,
):
requests.get(self.base_url + "/flush_cache")
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (Kimi-K2-Thinking)\n" f'{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], 0.95)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,72 @@
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_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(
base_url=self.base_url,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreaterEqual(metrics["score"], 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()

View File

@@ -0,0 +1,88 @@
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.send_one import BenchArgs, send_one_prompt
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,
)
MISTRAL_LARGE3_MODEL_PATH = "mistralai/Mistral-Large-3-675B-Instruct-2512"
class TestMistralLarge3Basic(CustomTestCase):
@classmethod
def setUpClass(cls):
# Set environment variable to disable JIT DeepGemm
os.environ["SGLANG_ENABLE_JIT_DEEPGEMM"] = "0"
cls.model = MISTRAL_LARGE3_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--tp",
"8",
"--attention-backend",
"trtllm_mla",
"--model-loader-extra-config",
'{"enable_multithread_load": true}',
"--chat-template",
"mistral",
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
# Clean up environment variable
if "SGLANG_ENABLE_JIT_DEEPGEMM" in os.environ:
del os.environ["SGLANG_ENABLE_JIT_DEEPGEMM"]
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=1400,
num_threads=1400,
num_shots=8,
)
metrics = run_eval(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (mistral-large-3)\n" f'{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], 0.90)
def test_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
acc_length, speed = send_one_prompt(args)
print(f"{speed=:.2f}")
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (mistral-large-3)\n" f"{speed=:.2f} token/s\n"
)
self.assertGreater(speed, 50)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,58 @@
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_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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 0.7)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,213 @@
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_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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 0.6)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,86 @@
import unittest
from sglang.test.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, variant=variant_config["name"])
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()

View File

@@ -0,0 +1,127 @@
import unittest
from sglang.test.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"
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",
"--dp",
"8",
"--enable-dp-attention",
"--model-loader-extra-config",
'{"enable_multithread_load": true}',
],
},
{
"name": "mtp",
"other_args": [
"--trust-remote-code",
"--tp",
"8",
"--dp",
"8",
"--enable-dp-attention",
"--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",
"--dp",
"8",
"--enable-dp-attention",
"--attention-backend",
"nsa",
"--nsa-prefill-backend",
"flashmla_sparse",
"--nsa-decode-backend",
"flashmla_kv",
"--model-loader-extra-config",
'{"enable_multithread_load": true}',
],
},
{
"name": "pure_tp",
"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, variant=variant_config["name"])
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()

View File

@@ -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()

View File

@@ -0,0 +1,59 @@
import unittest
from sglang.test.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()

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()

View File

@@ -0,0 +1,87 @@
import os
import unittest
import warnings
from sglang.test.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()

View File

@@ -0,0 +1,262 @@
import argparse
import glob
import json
import os
import random
import sys
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
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 lmmseval.
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),
]
_run_lmms_eval_with_retry(cmd, 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",
"--enforce-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]])

View File

@@ -0,0 +1,268 @@
import argparse
import glob
import json
import os
import random
import sys
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
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),
SimpleNamespace(model="Qwen/Qwen3-VL-8B-Instruct", mmmu_accuracy=0.60),
]
# Set default mem_fraction_static to 0.8
DEFAULT_MEM_FRACTION_STATIC = 0.8
class TestVLMViTCudaGraph(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
cls.enable_vit_cuda_graph = "1"
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"
os.environ["SGLANG_VIT_ENABLE_CUDA_GRAPH"] = cls.enable_vit_cuda_graph
def run_mmmu_eval(
self,
model_version: str,
output_path: str,
*,
env: dict | None = None,
):
"""
Evaluate a VLM on the MMMU validation set with lmmseval.
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),
]
_run_lmms_eval_with_retry(cmd, 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"
process_env["SGLANG_VIT_ENABLE_CUDA_GRAPH"] = "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=[
"--mm-attention-backend",
"fa3",
"--enforce-piecewise-cuda-graph",
"--piecewise-cuda-graph-max-tokens",
"8192",
"--chunked-prefill-size",
"8192",
"--disable-radix-cache",
"--disable-overlap-schedule",
"--piecewise-cuda-graph-compiler",
"eager",
],
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
TestVLMViTCudaGraph.parsed_args = args
# Pass args to unittest
unittest.main(argv=[sys.argv[0]])

View File

@@ -0,0 +1,258 @@
import argparse
import glob
import json
import os
import random
import sys
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
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/Qwen3-VL-30B-A3B-Instruct", mmmu_accuracy=0.51),
]
# Set default mem_fraction_static to 0.8
DEFAULT_MEM_FRACTION_STATIC = 0.8
class TestVLMViTFlashinferCudnn(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 lmmseval.
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),
]
_run_lmms_eval_with_retry(cmd, 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=[
"--mm-attention-backend",
"flashinfer_cudnn",
"--chunked-prefill-size",
"8192",
"--disable-radix-cache",
],
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
TestVLMViTFlashinferCudnn.parsed_args = args
# Pass args to unittest
unittest.main(argv=[sys.argv[0]])

View File

@@ -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()

View File

@@ -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()

View File

@@ -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()

View File

@@ -0,0 +1,88 @@
import unittest
from types import SimpleNamespace
from sglang.test.run_eval import run_eval
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
)
class TestDisaggregationPiecewiseCudaGraph(PDDisaggregationServerBase):
"""Test piecewise CUDA graph support in disaggregation prefill server"""
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
# Start servers
cls.start_prefill()
cls.start_decode()
# Wait for both to be ready
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
cls.launch_lb()
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"1",
"--enforce-piecewise-cuda-graph",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"1",
"--base-gpu-id",
"1",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_gsm8k_accuracy(self):
"""Verify that piecewise cuda graph works correctly in prefill server"""
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"GSM8K accuracy with piecewise cuda graph: {metrics['score']:.3f}")
self.assertGreater(metrics["score"], 0.62)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,42 @@
{
"model_type": "llama",
"kv_cache": {
"dtype": "float8_e4m3fn",
"scaling_factor": {
"0": {
"0": 1,
"1": 1,
"2": 1,
"3": 1,
"4": 1,
"5": 1,
"6": 1,
"7": 1,
"8": 1,
"9": 1,
"10": 1,
"11": 1,
"12": 1,
"13": 1,
"14": 1,
"15": 1,
"16": 1,
"17": 1,
"18": 1,
"19": 1,
"20": 1,
"21": 1,
"22": 1,
"23": 1,
"24": 1,
"25": 1,
"26": 1,
"27": 1,
"28": 1,
"29": 1,
"30": 1,
"31": 1
}
}
}
}

View File

@@ -0,0 +1,42 @@
{
"model_type": "llama",
"kv_cache": {
"dtype": "float8_e4m3fn",
"scaling_factor": {
"0": {
"0": 0.0408,
"1": 0.0503,
"2": 0.0667,
"3": 0.0909,
"4": 0.1135,
"5": 0.127,
"6": 0.1768,
"7": 0.1488,
"8": 0.1135,
"9": 0.1203,
"10": 0.1013,
"11": 0.0842,
"12": 0.1231,
"13": 0.1096,
"14": 0.1221,
"15": 0.1013,
"16": 0.1067,
"17": 0.0952,
"18": 0.0899,
"19": 0.097,
"20": 0.087,
"21": 0.0994,
"22": 0.0904,
"23": 0.1013,
"24": 0.1019,
"25": 0.1053,
"26": 0.1,
"27": 0.0894,
"28": 0.1013,
"29": 0.1488,
"30": 0.0766,
"31": 0.0821
}
}
}
}

View File

@@ -0,0 +1,38 @@
{
"model_type": "qwen",
"kv_cache": {
"dtype": "float8_e4m3fn",
"scaling_factor": {
"0": {
"0": 0.9846,
"1": 0.0645,
"2": 0.0731,
"3": 0.0800,
"4": 0.0748,
"5": 0.0780,
"6": 0.0702,
"7": 0.0894,
"8": 0.0410,
"9": 0.0758,
"10": 0.0556,
"11": 0.0731,
"12": 0.0899,
"13": 0.0780,
"14": 0.1441,
"15": 0.0914,
"16": 0.5614,
"17": 0.1067,
"18": 0.0537,
"19": 0.0658,
"20": 0.0523,
"21": 0.0533,
"22": 0.0699,
"23": 0.0635,
"24": 0.0588,
"25": 0.0884,
"26": 0.0947,
"27": 0.1032
}
}
}
}

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()

View File

@@ -0,0 +1,296 @@
"""
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 sys
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__":
sys.exit(pytest.main([__file__]))

View File

@@ -0,0 +1,166 @@
"""
Test script to verify SGLang config file integration.
"""
import argparse
import os
import sys
import tempfile
import pytest
import yaml
from sglang.srt.server_args import ServerArgs, prepare_server_args
from sglang.srt.server_args_config_parser import ConfigArgumentMerger
@pytest.fixture
def merger():
"""Fixture providing a ConfigArgumentMerger instance."""
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
return ConfigArgumentMerger(parser)
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,
"incremental-streaming-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 "--incremental-streaming-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)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

View File

@@ -0,0 +1,204 @@
#!/usr/bin/env python3
"""
Test cross-node scheduler_infos synchronization for remote weight loading.
Simulates multi-node setups on a single machine using different GPU subsets.
Validates that scheduler_infos are correctly synced across nodes via Gloo.
IMPORTANT: For multi-node tests, start both nodes within a few seconds of each
other to avoid port binding conflicts (they share the same network namespace).
Test cases:
- tp4_nodes2: TP=4 across 2 nodes, validates basic cross-node sync
- dp2_single_node: DP=2 with dp_attention on single node
- dp2_tp2_nodes2: DP=2, TP=4 across 2 nodes with dp_attention
Usage (multi-node):
Terminal 1: python test_cross_node_scheduler_info_sync.py --test-case tp4_nodes2 --node-rank 0
Terminal 2: python test_cross_node_scheduler_info_sync.py --test-case tp4_nodes2 --node-rank 1
Terminal 3: python test_cross_node_scheduler_info_sync.py --test-case tp4_nodes2 --test-only
Usage (single-node):
Terminal 1: python test_cross_node_scheduler_info_sync.py --test-case dp2_single_node --node-rank 0
Terminal 2: python test_cross_node_scheduler_info_sync.py --test-case dp2_single_node --test-only
Requirements: 4 GPUs on single machine
"""
import argparse
import socket
import subprocess
import sys
import time
from dataclasses import dataclass
from typing import List
import requests
from sglang.test.test_utils import (
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
)
@dataclass
class TestCase:
name: str
tp_size: int
dp_size: int
nnodes: int
gpus_per_node: int
expected_ranks: int
extra_args: List[str]
TEST_CASES = {
"tp4_nodes2": TestCase(
name="tp4_nodes2",
tp_size=4,
dp_size=1,
nnodes=2,
gpus_per_node=2,
expected_ranks=4,
extra_args=[],
),
"dp2_single_node": TestCase(
name="dp2_single_node",
tp_size=2,
dp_size=2,
nnodes=1,
gpus_per_node=2,
expected_ranks=2,
extra_args=["--enable-dp-attention", "--dp", "2", "--attention-backend", "fa3"],
),
"dp2_tp2_nodes2": TestCase(
name="dp2_tp2_nodes2",
tp_size=4,
dp_size=2,
nnodes=2,
gpus_per_node=2,
expected_ranks=4,
extra_args=["--enable-dp-attention", "--dp", "2", "--attention-backend", "fa3"],
),
}
TEST_CASE_MODELS = {
"tp4_nodes2": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
"dp2_single_node": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
"dp2_tp2_nodes2": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
}
def get_local_ip() -> str:
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
try:
s.connect(("8.8.8.8", 80))
return s.getsockname()[0]
except Exception:
return "127.0.0.1"
finally:
s.close()
def launch_node(
test_case: TestCase, node_rank: int, model_path: str, dist_init_addr: str
):
cmd = [
sys.executable,
"-m",
"sglang.launch_server",
"--model-path",
model_path,
"--tp",
str(test_case.tp_size),
"--port",
str(30000 + node_rank * 100),
"--host",
"0.0.0.0",
"--remote-instance-weight-loader-start-seed-via-transfer-engine",
]
if test_case.nnodes > 1:
cmd.extend(
[
"--nnodes",
str(test_case.nnodes),
"--node-rank",
str(node_rank),
"--dist-init-addr",
dist_init_addr,
"--base-gpu-id",
str(node_rank * test_case.gpus_per_node),
]
)
cmd.extend(test_case.extra_args)
print(f"[Node {node_rank}] {' '.join(cmd)}")
subprocess.run(cmd)
def test_api(test_case: TestCase) -> bool:
base_url = "http://127.0.0.1:30000"
print(f"Testing {test_case.name}: expecting {test_case.expected_ranks} ranks")
for _ in range(60):
try:
if requests.get(f"{base_url}/health", timeout=2).status_code == 200:
break
except Exception:
pass
time.sleep(2)
else:
print("ERROR: Server not ready")
return False
all_passed = True
for rank in range(test_case.expected_ranks):
try:
resp = requests.get(
f"{base_url}/get_remote_instance_transfer_engine_info",
params={"rank": rank},
timeout=5,
)
status = "" if resp.status_code == 200 else ""
print(f"{status} Rank {rank}: {resp.status_code}")
if resp.status_code != 200:
all_passed = False
except Exception as e:
print(f"✗ Rank {rank}: {e}")
all_passed = False
print("PASSED" if all_passed else "FAILED")
return all_passed
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--test-case", type=str, choices=list(TEST_CASES.keys()), required=True
)
parser.add_argument("--node-rank", type=int, choices=[0, 1])
parser.add_argument("--model-path", type=str, default=None)
parser.add_argument("--dist-init-addr", type=str, default=None)
parser.add_argument("--test-only", action="store_true")
args = parser.parse_args()
test_case = TEST_CASES[args.test_case]
model_path = args.model_path or TEST_CASE_MODELS.get(
args.test_case, DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT
)
if args.test_only:
sys.exit(0 if test_api(test_case) else 1)
if test_case.nnodes == 1:
launch_node(test_case, 0, model_path, "")
return
if args.node_rank is None:
print(f"Usage: --node-rank 0 or 1, then --test-only in another terminal")
sys.exit(0)
dist_init_addr = args.dist_init_addr or f"{get_local_ip()}:20000"
launch_node(test_case, args.node_rank, model_path, dist_init_addr)
if __name__ == "__main__":
main()

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()

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()

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()

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()

View File

@@ -0,0 +1,258 @@
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.srt.utils import get_device
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=get_device())
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}")
if device == "xpu":
torch.xpu.set_device(f"xpu:{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()

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()

View File

@@ -0,0 +1,307 @@
"""
Test forward_split_prefill functionality.
Usage:
python3 -m unittest test_forward_split_prefill.TestForwardSplitPrefill
or
python3 test_forward_split_prefill.py
"""
import unittest
import numpy as np
import torch
from sglang.bench_one_batch import TreeCacheNamespace
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 import get_device
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 = get_device()
# 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,
moe_ep_rank=0,
moe_ep_size=1,
)
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.logprob_start_len = -1
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
reqs.append(req)
# Create dummy tree_cache for tests (no prefix caching, just allocation)
dummy_tree_cache = TreeCacheNamespace(
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,
)
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()

View File

@@ -0,0 +1,185 @@
import gc
import unittest
import numpy as np
import requests
from transformers import AutoModelForCausalLM
import sglang as sgl
from sglang.srt.utils import get_device
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,
empty_gpu_cache,
get_gpu_count,
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(get_device())
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()
empty_gpu_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 get_gpu_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 get_gpu_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()

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()

View File

@@ -0,0 +1,292 @@
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_MLA_MODEL_NAME_FOR_TEST,
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,
},
},
)
# Get events
events = []
start = time.time()
max_wait_s = 5
min_events_expected = 5 # Expect at least some events
while (
len(events) < min_events_expected and (time.time() - start) < max_wait_s
):
if sub.poll(timeout=100): # 100ms timeout
_, seq_bytes, payload = sub.recv_multipart()
event_batch = decoder.decode(payload)
for event in event_batch.events:
events.append(event)
# Verify we received events
self.assertGreater(
len(events), 0, "Should have received at least one KV cache event"
)
# Track which blocks were stored and removed
stored_blocks = {} # hash -> BlockStored event
removed_hashes = set()
# Validate event structure and relationships
for event in events:
self.assertIsInstance(
event,
(BlockStored, BlockRemoved, AllBlocksCleared),
f"Event should be a KV cache event, got {type(event)}",
)
if isinstance(event, BlockStored):
# Validate BlockStored structure
self.assertIsInstance(event.block_hashes, list)
self.assertEqual(
len(event.block_hashes), 1, "Should have one hash per block"
)
self.assertIsInstance(event.token_ids, list)
self.assertEqual(
event.block_size,
len(event.token_ids),
"block_size should match token_ids length",
)
self.assertIsNone(
event.lora_id, "lora_id should be None for basic test"
)
# Store this block for later validation
block_hash = event.block_hashes[0]
stored_blocks[block_hash] = event
# If parent_block_hash is set, verify it was stored earlier
if event.parent_block_hash is not None:
# Parent should either be in stored_blocks or could be from a previous request
pass # Don't strictly enforce this as root blocks may have synthetic parents
elif isinstance(event, BlockRemoved):
# Validate BlockRemoved structure
self.assertIsInstance(event.block_hashes, list)
self.assertEqual(
len(event.block_hashes), 1, "Should have one hash per block"
)
removed_hashes.add(event.block_hashes[0])
# Verify we got both BlockStored and BlockRemoved events
self.assertGreater(
len(stored_blocks), 0, "Should have at least one BlockStored event"
)
# BlockRemoved events may not always occur in this short test, so just check if they do occur
# that they reference previously stored blocks
for removed_hash in removed_hashes:
# It's OK if the removed block wasn't in our stored_blocks
# (it could have been stored before we started listening)
pass
finally:
sub.close()
context.term()
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(
DEFAULT_MLA_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",
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()

View File

@@ -0,0 +1,527 @@
"""
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/manual/test_logprobs.py gen
```
Step 2: Test Against Baseline (After Code Changes)
```bash
python test/manual/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
num_samples = config.get("num_samples", NUM_SAMPLES)
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())

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.run_eval import run_eval
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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
self.assertGreater(metrics["score"], 0.62)
def test_gsm8k_bs1(self):
# test torch compile accuracy for bs=1
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=10,
num_threads=1,
)
metrics = run_eval(args)
self.assertGreater(metrics["score"], 0.62)
if __name__ == "__main__":
unittest.main()

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,
)

View File

@@ -0,0 +1,29 @@
import unittest
from sglang.srt.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()

View File

@@ -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()

View File

@@ -0,0 +1,293 @@
import os
import subprocess
import unittest
import requests
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
)
class TestMoriTransferEngineE2E(PDDisaggregationServerBase):
"""
Run:
SGLANG_MORI_MANUAL_E2E=1 python3 test/manual/test_mori_transfer_engine_e2e.py
Optional:
- SGLANG_MORI_E2E_TEST_MODEL: override model (defaults to a small test model)
- SGLANG_TEST_PD_DISAGG_DEVICES: RDMA devices string, e.g. "mlx5_roce0,mlx5_roce4"
"""
@classmethod
def setUpClass(cls):
if os.environ.get("SGLANG_MORI_MANUAL_E2E", "") not in ("1", "true", "True"):
raise unittest.SkipTest(
"Set SGLANG_MORI_MANUAL_E2E=1 to run this manual MORI E2E test."
)
try:
import torch
if not torch.cuda.is_available():
raise unittest.SkipTest("torch.cuda is not available.")
except Exception as e:
raise unittest.SkipTest(f"torch is not available/usable: {e}")
# Force the disaggregation fixture to use MORI backend in local/manual runs.
os.environ["SGLANG_TEST_PD_DISAGG_BACKEND"] = "mori"
super().setUpClass()
cls.model = os.environ.get(
"SGLANG_MORI_E2E_TEST_MODEL", DEFAULT_SMALL_MODEL_NAME_FOR_TEST
)
cls.start_prefill()
cls.start_decode()
cls.wait_server_ready(
cls.prefill_url + "/health",
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
process=cls.process_prefill,
)
cls.wait_server_ready(
cls.decode_url + "/health",
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
process=cls.process_decode,
)
cls.launch_lb()
@classmethod
def tearDownClass(cls):
os.environ.pop("SGLANG_TEST_PD_DISAGG_BACKEND", None)
super().tearDownClass()
@classmethod
def launch_lb(cls):
lb_command = [
"python3",
"-m",
"sglang_router.launch_router",
"--pd-disaggregation",
"--mini-lb",
"--prefill",
cls.prefill_url,
"--decode",
cls.decode_url,
"--host",
cls.base_host,
"--port",
cls.lb_port,
]
print("Starting load balancer:", " ".join(lb_command))
cls.process_lb = subprocess.Popen(lb_command, stdout=None, stderr=None)
cls.wait_server_ready(
cls.lb_url + "/health",
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
process=cls.process_lb,
)
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"1",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"1",
"--base-gpu-id",
"1",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_generate_smoke(self):
resp = requests.post(
self.lb_url + "/generate",
json={
"text": "Hello",
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
},
timeout=120,
)
self.assertEqual(resp.status_code, 200, resp.text)
out = resp.json()
self.assertIn("text", out)
self.assertIsInstance(out["text"], str)
self.assertGreater(len(out["text"]), 0)
class TestMoriTransferEngineTPMismatchE2E(PDDisaggregationServerBase):
"""Manual MORI PD-disaggregation E2E with TP mismatch.
Scenario:
- prefill: tp=2 (GPU 0-1)
- decode: tp=4 (GPU 2-5)
Manual-only and requires >= 6 visible GPUs.
"""
_PORT_DELTA = 10
@classmethod
def setUpClass(cls):
if os.environ.get("SGLANG_MORI_MANUAL_E2E", "") not in ("1", "true", "True"):
raise unittest.SkipTest(
"Set SGLANG_MORI_MANUAL_E2E=1 to run this manual MORI E2E test."
)
try:
import torch
if not torch.cuda.is_available():
raise unittest.SkipTest("torch.cuda is not available.")
if torch.cuda.device_count() < 6:
raise unittest.SkipTest(
"TP-mismatch test requires >= 6 visible GPUs (prefill tp=2 + decode tp=4)."
)
except Exception as e:
raise unittest.SkipTest(f"torch is not available/usable: {e}")
os.environ["SGLANG_TEST_PD_DISAGG_BACKEND"] = "mori"
super().setUpClass()
# Shift ports to avoid clashing with TestMoriTransferEngineE2E.
cls.lb_port = str(int(cls.lb_port) + cls._PORT_DELTA)
cls.prefill_port = str(int(cls.prefill_port) + cls._PORT_DELTA)
cls.decode_port = str(int(cls.decode_port) + cls._PORT_DELTA)
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}"
cls.model = os.environ.get(
"SGLANG_MORI_E2E_TEST_MODEL", DEFAULT_SMALL_MODEL_NAME_FOR_TEST
)
cls.start_prefill()
cls.start_decode()
cls.wait_server_ready(
cls.prefill_url + "/health",
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
process=cls.process_prefill,
)
cls.wait_server_ready(
cls.decode_url + "/health",
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
process=cls.process_decode,
)
cls.launch_lb()
@classmethod
def tearDownClass(cls):
os.environ.pop("SGLANG_TEST_PD_DISAGG_BACKEND", None)
super().tearDownClass()
@classmethod
def launch_lb(cls):
lb_command = [
"python3",
"-m",
"sglang_router.launch_router",
"--pd-disaggregation",
"--mini-lb",
"--prefill",
cls.prefill_url,
"--decode",
cls.decode_url,
"--host",
cls.base_host,
"--port",
cls.lb_port,
]
print("Starting load balancer:", " ".join(lb_command))
cls.process_lb = subprocess.Popen(lb_command, stdout=None, stderr=None)
cls.wait_server_ready(
cls.lb_url + "/health",
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
process=cls.process_lb,
)
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"2",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"4",
"--base-gpu-id",
"2",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_generate_smoke_tp_mismatch(self):
resp = requests.post(
self.lb_url + "/generate",
json={
"text": "Hello",
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
},
timeout=120,
)
self.assertEqual(resp.status_code, 200, resp.text)
out = resp.json()
self.assertIn("text", out)
self.assertIsInstance(out["text"], str)
self.assertGreater(len(out["text"]), 0)
if __name__ == "__main__":
unittest.main()

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()

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
import sglang.srt.distributed.device_communicators.custom_all_reduce_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()

View File

@@ -0,0 +1,461 @@
"""Integration tests for RayEngine and Ray HTTP server (requires GPU + Ray).
Tests the Ray actor scheduler backend:
- Offline inference via Engine(use_ray=True) inside a Ray actor on a placement group
- Error paths in RayEngine._launch_scheduler_processes()
- HTTP server launched via --use-ray flag
Usage:
# 1-GPU tests
python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflineTP1 -v -s
python -m pytest test/manual/test_ray_engine.py::TestRayEngineErrors -v -s
python -m pytest test/manual/test_ray_engine.py::TestRayHTTPServerTP1 -v -s
# 2-GPU tests
python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflineTP2 -v -s
python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflinePP2 -v -s
"""
from __future__ import annotations
import os
import time
import unittest
import torch
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
# Allow overriding the model via env var for environments without gated access
_MODEL = os.environ.get("SGLANG_TEST_MODEL", DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
try:
import ray
from ray.runtime_env import RuntimeEnv
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
# Prevent Ray from overriding CUDA_VISIBLE_DEVICES so that all GPUs
# remain visible inside actors regardless of num_gpus allocation.
_env_vars = {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}
if os.environ.get("HF_TOKEN"):
_env_vars["HF_TOKEN"] = os.environ["HF_TOKEN"]
_RAY_RUNTIME_ENV = RuntimeEnv(env_vars=_env_vars)
_has_ray = True
except ImportError:
_has_ray = False
_RAY_RUNTIME_ENV = None
_NUM_GPUS = torch.cuda.device_count()
_SAMPLING_PARAMS = {"max_new_tokens": 32, "temperature": 0.0}
_PROMPTS = [
"The capital of France is",
"Explain quantum computing in simple terms:",
"Write a haiku about programming:",
"What is 2 + 2?",
]
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _create_engine_on_pg(tp_size, pp_size=1, model=_MODEL, extra_kwargs=None):
"""Create an EngineActor on a placement group and wait for it to be ready.
Returns (engine_actor, placement_group).
"""
@ray.remote
class EngineActor:
def __init__(self, **kwargs):
from sglang.srt.ray.engine import RayEngine
self.engine = RayEngine(**kwargs)
def is_ready(self):
return True
def generate(self, prompt, sampling_params):
return self.engine.generate(prompt=prompt, sampling_params=sampling_params)
def shutdown(self):
if self.engine:
self.engine.shutdown()
self.engine = None
total_gpus = tp_size * pp_size
pg = placement_group(
[{"CPU": 1, "GPU": total_gpus}],
strategy="STRICT_PACK",
)
ray.get(pg.ready())
kwargs = dict(
model_path=model,
tp_size=tp_size,
pp_size=pp_size,
)
if extra_kwargs:
kwargs.update(extra_kwargs)
actor = EngineActor.options(
num_cpus=1,
num_gpus=0,
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=pg,
placement_group_bundle_index=0,
),
).remote(**kwargs)
ray.get(actor.is_ready.remote(), timeout=600)
return actor, pg
def _cleanup(actor, pg):
"""Shutdown engine actor and remove placement group."""
try:
ray.get(actor.shutdown.remote(), timeout=60)
except Exception:
pass
try:
ray.util.remove_placement_group(pg)
except Exception:
pass
# ---------------------------------------------------------------------------
# Tests: Offline TP=1
# ---------------------------------------------------------------------------
@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU")
class TestRayEngineOfflineTP1(unittest.TestCase):
@classmethod
def setUpClass(cls):
if not ray.is_initialized():
ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
cls.actor, cls.pg = _create_engine_on_pg(tp_size=1)
@classmethod
def tearDownClass(cls):
_cleanup(cls.actor, cls.pg)
ray.shutdown()
def test_offline_generate(self):
result = ray.get(
self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS)
)
self.assertIn("text", result)
self.assertGreater(len(result["text"]), 0)
print(f"Generated: {result['text'][:200]}")
def test_batch_generate(self):
for prompt in _PROMPTS:
result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
self.assertIn("text", result)
self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}")
def test_deterministic(self):
prompt = "The meaning of life is"
r1 = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
r2 = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
self.assertEqual(r1["text"], r2["text"])
# ---------------------------------------------------------------------------
# Tests: Offline TP=2
# ---------------------------------------------------------------------------
@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs")
class TestRayEngineOfflineTP2(unittest.TestCase):
@classmethod
def setUpClass(cls):
if not ray.is_initialized():
ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
cls.actor, cls.pg = _create_engine_on_pg(tp_size=2)
@classmethod
def tearDownClass(cls):
_cleanup(cls.actor, cls.pg)
ray.shutdown()
def test_offline_generate_tp2(self):
result = ray.get(
self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS)
)
self.assertIn("text", result)
self.assertGreater(len(result["text"]), 0)
print(f"Generated (TP=2): {result['text'][:200]}")
def test_batch_generate_tp2(self):
for prompt in _PROMPTS:
result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
self.assertIn("text", result)
self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}")
# ---------------------------------------------------------------------------
# Tests: Offline PP=2
# ---------------------------------------------------------------------------
@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs")
class TestRayEngineOfflinePP2(unittest.TestCase):
@classmethod
def setUpClass(cls):
if not ray.is_initialized():
ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
cls.actor, cls.pg = _create_engine_on_pg(tp_size=1, pp_size=2)
@classmethod
def tearDownClass(cls):
_cleanup(cls.actor, cls.pg)
ray.shutdown()
def test_offline_generate_pp2(self):
result = ray.get(
self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS)
)
self.assertIn("text", result)
self.assertGreater(len(result["text"]), 0)
print(f"Generated (PP=2): {result['text'][:200]}")
def test_batch_generate_pp2(self):
for prompt in _PROMPTS:
result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
self.assertIn("text", result)
self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}")
# ---------------------------------------------------------------------------
# Tests: Error paths
# ---------------------------------------------------------------------------
@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU")
class TestRayEngineErrors(unittest.TestCase):
@classmethod
def setUpClass(cls):
if not ray.is_initialized():
ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
@classmethod
def tearDownClass(cls):
ray.shutdown()
def test_dp_greater_than_1_raises(self):
"""RayEngine with dp_size > 1 should raise NotImplementedError."""
@ray.remote
class _BadActor:
def try_create(self):
from sglang.srt.ray.engine import RayEngine
try:
RayEngine(
model_path=_MODEL,
tp_size=1,
dp_size=2,
use_ray=True,
)
return None
except (NotImplementedError, RuntimeError) as e:
return str(e)
pg = placement_group([{"CPU": 1, "GPU": 1}], strategy="STRICT_PACK")
ray.get(pg.ready())
actor = _BadActor.options(
num_cpus=1,
num_gpus=0,
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=pg,
placement_group_bundle_index=0,
),
).remote()
try:
error_msg = ray.get(actor.try_create.remote(), timeout=120)
self.assertIsNotNone(error_msg, "Expected error but RayEngine created OK")
self.assertIn("dp_size", error_msg.lower())
finally:
ray.util.remove_placement_group(pg)
def test_missing_placement_group_raises(self):
"""RayEngine without a placement group should raise RuntimeError."""
@ray.remote(num_gpus=1)
def _try_create_without_pg():
from sglang.srt.ray.engine import RayEngine
try:
RayEngine(
model_path=_MODEL,
tp_size=1,
use_ray=True,
)
return None
except RuntimeError as e:
return str(e)
error_msg = ray.get(_try_create_without_pg.remote(), timeout=120)
self.assertIsNotNone(
error_msg, "Expected RuntimeError but RayEngine created OK"
)
self.assertIn("placement group", error_msg.lower())
# ---------------------------------------------------------------------------
# Tests: HTTP server
# ---------------------------------------------------------------------------
@unittest.skipUnless(_has_ray, "ray is not installed")
@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU")
class TestRayHTTPServerTP1(unittest.TestCase):
"""Test the Ray HTTP server path (launch_server.py --use-ray).
Launches the server inside a Ray task on a placement group (mirrors
examples/anyscale/driver_online.py) and sends HTTP requests to it.
"""
@classmethod
def setUpClass(cls):
import requests as req_lib
if not ray.is_initialized():
ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
cls.port = 30100
cls.pg = placement_group(
[{"CPU": 1, "GPU": 1}],
strategy="STRICT_PACK",
)
ray.get(cls.pg.ready())
pg_strategy = PlacementGroupSchedulingStrategy(
placement_group=cls.pg,
placement_group_bundle_index=0,
)
# Resolve the node IP where the server will run
@ray.remote(num_cpus=0, num_gpus=0)
def _get_ip():
return ray.util.get_node_ip_address()
cls.node_ip = ray.get(_get_ip.options(scheduling_strategy=pg_strategy).remote())
cls.base_url = f"http://{cls.node_ip}:{cls.port}"
# Launch server as a Ray task (blocks until server exits)
@ray.remote
def _launch(**kwargs):
from sglang.srt.ray.http_server import launch_server
from sglang.srt.server_args import ServerArgs
launch_server(ServerArgs(**kwargs))
cls.server_ref = _launch.options(
num_cpus=1,
num_gpus=0,
scheduling_strategy=pg_strategy,
).remote(
model_path=_MODEL,
tp_size=1,
port=cls.port,
host="0.0.0.0",
use_ray=True,
)
# Wait for health check
t0 = time.time()
timeout = 600
healthy = False
while time.time() - t0 < timeout:
ready, _ = ray.wait([cls.server_ref], timeout=0)
if ready:
try:
ray.get(cls.server_ref)
except Exception as e:
raise RuntimeError(f"Server task crashed: {e}") from e
raise RuntimeError("Server task exited before becoming healthy")
try:
if req_lib.get(f"{cls.base_url}/health", timeout=5).status_code == 200:
healthy = True
break
except req_lib.exceptions.RequestException:
pass
time.sleep(3)
if not healthy:
ray.cancel(cls.server_ref, force=True)
raise RuntimeError(f"Server did not become healthy within {timeout}s")
@classmethod
def tearDownClass(cls):
try:
ray.cancel(cls.server_ref, force=True)
except Exception:
pass
try:
ray.util.remove_placement_group(cls.pg)
except Exception:
pass
ray.shutdown()
def test_health_endpoint(self):
import requests
resp = requests.get(f"{self.base_url}/health", timeout=10)
self.assertEqual(resp.status_code, 200)
def test_generate_endpoint(self):
import requests
resp = requests.post(
f"{self.base_url}/generate",
json={
"text": "The capital of France is",
"sampling_params": _SAMPLING_PARAMS,
},
timeout=60,
)
resp.raise_for_status()
data = resp.json()
self.assertIn("text", data)
self.assertGreater(len(data["text"]), 0)
print(f"HTTP response: {data['text'][:200]}")
def test_generate_multiple(self):
import requests
for prompt in _PROMPTS:
resp = requests.post(
f"{self.base_url}/generate",
json={
"text": prompt,
"sampling_params": _SAMPLING_PARAMS,
},
timeout=60,
)
resp.raise_for_status()
data = resp.json()
self.assertIn("text", data)
self.assertGreater(len(data["text"]), 0, f"Empty output for: {prompt}")
if __name__ == "__main__":
unittest.main()

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()

View File

@@ -0,0 +1,314 @@
import unittest
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
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.create_simulated()
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):
tree_cache = RadixCache.create_simulated(disable=True)
policy = SchedulePolicy(
policy="lpm",
tree_cache=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.create_simulated()
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.create_simulated()
r1 = Req(1, "a b", [1, 2], SamplingParams())
r2 = Req(3, "a b c", [1, 2, 3], SamplingParams())
r3 = Req(2, "a", [1], SamplingParams())
r1.priority, r1.time_stats.wait_queue_entry_time = 1, 1
r2.priority, r2.time_stats.wait_queue_entry_time = 0, 1
r3.priority, r3.time_stats.wait_queue_entry_time = 0, 0
waiting_queue = [r1, r2, r3]
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.create_simulated()
r1 = Req(1, "a b", [1, 2], SamplingParams())
r2 = Req(3, "a b c", [1, 2, 3], SamplingParams())
r3 = Req(2, "a", [1], SamplingParams())
r1.priority, r1.time_stats.wait_queue_entry_time = -1, 1
r2.priority, r2.time_stats.wait_queue_entry_time = 0, 1
r3.priority, r3.time_stats.wait_queue_entry_time = 0, 0
waiting_queue = [r1, r2, r3]
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.create_simulated()
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.create_simulated()
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.create_simulated()
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)
def test_calc_priority_routing_key_scheduling(self):
"""Test routing-key policy: prioritize by routing key frequency in running batch."""
tree_cache = RadixCache.create_simulated()
running_reqs = [
Req("r1", "a", [1], SamplingParams(), routing_key="key_a"),
Req("r2", "b", [2], SamplingParams(), routing_key="key_a"),
Req("r3", "c", [3], SamplingParams(), routing_key="key_b"),
]
running_batch = ScheduleBatch(reqs=running_reqs)
waiting_queue = [
Req("w1", "d", [4], SamplingParams(), routing_key="key_b"),
Req("w2", "e", [5], SamplingParams(), routing_key="key_a"),
Req("w3", "f", [6], SamplingParams(), routing_key="key_c"),
]
policy = SchedulePolicy(
policy="routing-key",
tree_cache=tree_cache,
enable_hierarchical_cache=False,
enable_priority_scheduling=False,
schedule_low_priority_values_first=False,
)
policy.calc_priority(waiting_queue, running_batch)
self.assertEqual(waiting_queue[0].rid, "w2")
self.assertEqual(waiting_queue[1].rid, "w1")
self.assertEqual(waiting_queue[2].rid, "w3")
def test_calc_priority_routing_key_tie_break_by_lexicographic_order(self):
"""Test routing-key policy: tie-break by lexicographic order."""
tree_cache = RadixCache.create_simulated()
running_reqs = [
Req("r1", "a", [1], SamplingParams(), routing_key="key_b"),
Req("r2", "b", [2], SamplingParams(), routing_key="key_a"),
]
running_batch = ScheduleBatch(reqs=running_reqs)
waiting_queue = [
Req("w1", "d", [4], SamplingParams(), routing_key="key_b"),
Req("w2", "e", [5], SamplingParams(), routing_key="key_a"),
]
policy = SchedulePolicy(
policy="routing-key",
tree_cache=tree_cache,
enable_hierarchical_cache=False,
enable_priority_scheduling=False,
schedule_low_priority_values_first=False,
)
policy.calc_priority(waiting_queue, running_batch)
self.assertEqual(waiting_queue[0].rid, "w2")
self.assertEqual(waiting_queue[1].rid, "w1")
def test_calc_priority_routing_key_no_match_deprioritized(self):
"""Test routing-key policy: requests without matching routing keys are deprioritized."""
tree_cache = RadixCache.create_simulated()
running_reqs = [
Req("r1", "a", [1], SamplingParams(), routing_key="key_a"),
Req("r2", "b", [2], SamplingParams(), routing_key="key_b"),
Req("r3", "c", [3], SamplingParams(), routing_key="key_c"),
]
running_batch = ScheduleBatch(reqs=running_reqs)
waiting_queue = [
Req("w1", "d", [4], SamplingParams(), routing_key="key_d"),
Req("w2", "e", [5], SamplingParams(), routing_key="key_e"),
Req("w3", "f", [6], SamplingParams(), routing_key="key_c"),
]
policy = SchedulePolicy(
policy="routing-key",
tree_cache=tree_cache,
enable_hierarchical_cache=False,
enable_priority_scheduling=False,
schedule_low_priority_values_first=False,
)
policy.calc_priority(waiting_queue, running_batch)
self.assertEqual(waiting_queue[0].rid, "w3")
self.assertEqual(waiting_queue[1].rid, "w1")
self.assertEqual(waiting_queue[2].rid, "w2")
def test_calc_priority_routing_key_empty_running_batch(self):
"""Test routing-key policy: empty running batch keeps original order."""
tree_cache = RadixCache.create_simulated()
running_batch = ScheduleBatch(reqs=[])
waiting_queue = [
Req("w1", "d", [4], SamplingParams(), routing_key="key_a"),
Req("w2", "e", [5], SamplingParams(), routing_key="key_b"),
Req("w3", "f", [6], SamplingParams(), routing_key="key_c"),
]
policy = SchedulePolicy(
policy="routing-key",
tree_cache=tree_cache,
enable_hierarchical_cache=False,
enable_priority_scheduling=False,
schedule_low_priority_values_first=False,
)
policy.calc_priority(waiting_queue, running_batch)
self.assertEqual(waiting_queue[0].rid, "w1")
self.assertEqual(waiting_queue[1].rid, "w2")
self.assertEqual(waiting_queue[2].rid, "w3")
if __name__ == "__main__":
unittest.main()

View File

@@ -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()

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)

View File

@@ -0,0 +1,408 @@
"""
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 InputFormat, 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, InputFormat.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, InputFormat.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, InputFormat.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, InputFormat.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, InputFormat.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, InputFormat.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, InputFormat.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, InputFormat.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, InputFormat.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, InputFormat.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, InputFormat.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, InputFormat.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, InputFormat.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, InputFormat.CROSS_ENCODER_PAIRS
)
self.assertEqual(result, [["query text", "document text"]])
def test_prepare_batch_strings_input_format_passthrough(self):
"""Batch strings should pass through unchanged."""
texts = ["test"]
result = self.tokenizer_manager._prepare_tokenizer_input(
texts, InputFormat.BATCH_STRINGS
)
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,
InputFormat.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,
InputFormat.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,
InputFormat.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,
InputFormat.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,
InputFormat.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,
InputFormat.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, InputFormat.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, InputFormat.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, InputFormat.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)

View File

@@ -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.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 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(
base_url=base_url,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=100,
num_threads=10,
num_shots=8,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 0.62)
finally:
kill_process_tree(process.pid)
if __name__ == "__main__":
unittest.main()

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()

View File

@@ -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()

View File

@@ -0,0 +1,252 @@
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
from sglang.srt.utils import get_device
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):
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,
):
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
print(m, n, k, e, topk, dtype, group_size, has_zp, weight_bits)
a = torch.randn((m, k), device=get_device(), dtype=dtype) / 10
w1 = torch.randn((e, 2 * n, k), device=get_device(), dtype=dtype) / 10
w2 = torch.randn((e, k, n), device=get_device(), dtype=dtype) / 10
score = torch.randn((m, e), device=get_device(), 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=get_device(), dtype=torch.uint8
)
w2_qweight = torch.empty(
(e, k, n // pack_factor), device=get_device(), dtype=torch.uint8
)
w1_scales = torch.empty(
(e, 2 * n, k // group_size), device=get_device(), dtype=dtype
)
w2_scales = torch.empty((e, k, n // group_size), device=get_device(), dtype=dtype)
w1_qzeros = torch.empty(
(e, 2 * n // pack_factor, k // group_size),
device=get_device(),
dtype=torch.uint8,
)
w2_qzeros = torch.empty(
(e, k // pack_factor, n // group_size), device=get_device(), 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)

View File

@@ -0,0 +1,481 @@
"""
Unit tests for TRTLLM FP8 KV cache fusion kernel.
"""
import unittest
import torch
from sglang.srt.layers.attention.triton_ops.trtllm_fp8_kv_kernel import (
fused_fp8_set_kv_buffer,
)
from sglang.test.test_utils import CustomTestCase
class TestTRTLLMFP8KVKernel(CustomTestCase):
"""Test fused FP8 KV cache write kernel correctness."""
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA not available")
if torch.cuda.get_device_capability()[0] < 9:
raise unittest.SkipTest("FP8 requires compute capability >= 9.0")
def _test_kernel_correctness(
self,
num_tokens,
num_kv_heads,
head_dim,
page_size,
use_scale,
input_ndim,
cache_ndim,
):
"""Compare Triton kernel output against naive implementation."""
device = torch.device("cuda")
dtype = torch.bfloat16
# Create input tensors
if input_ndim == 3:
k = torch.randn(
num_tokens, num_kv_heads, head_dim, device=device, dtype=dtype
)
v = torch.randn(
num_tokens, num_kv_heads, head_dim, device=device, dtype=dtype
)
else:
k = torch.randn(
num_tokens, num_kv_heads * head_dim, device=device, dtype=dtype
)
v = torch.randn(
num_tokens, num_kv_heads * head_dim, device=device, dtype=dtype
)
# Create cache tensors (use FP8 to match real runtime behavior)
num_pages = 128
total_slots = num_pages * page_size
cache_dtype = torch.float8_e4m3fn
if cache_ndim == 3:
k_cache_triton = torch.zeros(
total_slots, num_kv_heads, head_dim, device=device, dtype=cache_dtype
)
v_cache_triton = torch.zeros(
total_slots, num_kv_heads, head_dim, device=device, dtype=cache_dtype
)
k_cache_naive = torch.zeros(
total_slots, num_kv_heads, head_dim, device=device, dtype=cache_dtype
)
v_cache_naive = torch.zeros(
total_slots, num_kv_heads, head_dim, device=device, dtype=cache_dtype
)
else:
k_cache_triton = torch.zeros(
num_pages,
page_size,
num_kv_heads,
head_dim,
device=device,
dtype=cache_dtype,
)
v_cache_triton = torch.zeros(
num_pages,
page_size,
num_kv_heads,
head_dim,
device=device,
dtype=cache_dtype,
)
k_cache_naive = torch.zeros(
num_pages,
page_size,
num_kv_heads,
head_dim,
device=device,
dtype=cache_dtype,
)
v_cache_naive = torch.zeros(
num_pages,
page_size,
num_kv_heads,
head_dim,
device=device,
dtype=cache_dtype,
)
# Create cache locations (ensure unique indices to avoid race conditions)
cache_loc = torch.randperm(total_slots, device=device, dtype=torch.int32)[
:num_tokens
]
# Optional scales
k_scale = 0.5 if use_scale else None
v_scale = 0.75 if use_scale else None
# Run Triton kernel
fused_fp8_set_kv_buffer(
k.clone(),
v.clone(),
k_cache_triton,
v_cache_triton,
cache_loc,
k_scale,
v_scale,
page_size,
use_triton=True,
)
# Run naive fallback
fused_fp8_set_kv_buffer(
k.clone(),
v.clone(),
k_cache_naive,
v_cache_naive,
cache_loc,
k_scale,
v_scale,
page_size,
use_triton=False,
)
# Compare results (bit-exact match expected)
self.assertTrue(
torch.equal(k_cache_triton, k_cache_naive),
"K cache mismatch between Triton and naive",
)
self.assertTrue(
torch.equal(v_cache_triton, v_cache_naive),
"V cache mismatch between Triton and naive",
)
def test_basic_3d_input_3d_cache(self):
"""Test basic case: 3D input, 3D cache, no scale."""
self._test_kernel_correctness(
num_tokens=16,
num_kv_heads=8,
head_dim=128,
page_size=16,
use_scale=False,
input_ndim=3,
cache_ndim=3,
)
def test_basic_3d_input_4d_cache(self):
"""Test basic case: 3D input, 4D cache, no scale."""
self._test_kernel_correctness(
num_tokens=16,
num_kv_heads=8,
head_dim=128,
page_size=16,
use_scale=False,
input_ndim=3,
cache_ndim=4,
)
def test_with_scale_3d_cache(self):
"""Test with scale: 3D input, 3D cache."""
self._test_kernel_correctness(
num_tokens=16,
num_kv_heads=8,
head_dim=128,
page_size=16,
use_scale=True,
input_ndim=3,
cache_ndim=3,
)
def test_with_scale_4d_cache(self):
"""Test with scale: 3D input, 4D cache."""
self._test_kernel_correctness(
num_tokens=16,
num_kv_heads=8,
head_dim=128,
page_size=16,
use_scale=True,
input_ndim=3,
cache_ndim=4,
)
def test_2d_input_3d_cache(self):
"""Test 2D input (flattened): 2D input, 3D cache."""
self._test_kernel_correctness(
num_tokens=16,
num_kv_heads=8,
head_dim=128,
page_size=16,
use_scale=False,
input_ndim=2,
cache_ndim=3,
)
def test_2d_input_4d_cache(self):
"""Test 2D input (flattened): 2D input, 4D cache."""
self._test_kernel_correctness(
num_tokens=16,
num_kv_heads=8,
head_dim=128,
page_size=16,
use_scale=False,
input_ndim=2,
cache_ndim=4,
)
def test_single_token(self):
"""Test edge case: single token."""
self._test_kernel_correctness(
num_tokens=1,
num_kv_heads=8,
head_dim=128,
page_size=16,
use_scale=True,
input_ndim=3,
cache_ndim=3,
)
def test_large_batch(self):
"""Test larger batch size."""
self._test_kernel_correctness(
num_tokens=128,
num_kv_heads=16,
head_dim=64,
page_size=16,
use_scale=True,
input_ndim=3,
cache_ndim=4,
)
def test_different_head_dims(self):
"""Test different head dimensions."""
for head_dim in [64, 128]:
self._test_kernel_correctness(
num_tokens=16,
num_kv_heads=8,
head_dim=head_dim,
page_size=16,
use_scale=False,
input_ndim=3,
cache_ndim=3,
)
def test_empty_input(self):
"""Test edge case: empty input (0 tokens)."""
device = torch.device("cuda")
dtype = torch.bfloat16
num_kv_heads = 8
head_dim = 128
page_size = 16
num_tokens = 0
# Empty inputs
k = torch.randn(num_tokens, num_kv_heads, head_dim, device=device, dtype=dtype)
v = torch.randn(num_tokens, num_kv_heads, head_dim, device=device, dtype=dtype)
# Cache (use FP8 to match real runtime behavior)
total_slots = 128
k_cache = torch.zeros(
total_slots,
num_kv_heads,
head_dim,
device=device,
dtype=torch.float8_e4m3fn,
)
v_cache = torch.zeros(
total_slots,
num_kv_heads,
head_dim,
device=device,
dtype=torch.float8_e4m3fn,
)
# Empty cache locations
cache_loc = torch.empty(num_tokens, device=device, dtype=torch.int32)
# Should not crash
fused_fp8_set_kv_buffer(
k,
v,
k_cache,
v_cache,
cache_loc,
k_scale=None,
v_scale=None,
page_size=page_size,
)
def test_fp8_kv_kernel_accepts_tensor_scales(self):
"""
Regression test for B200 Triton compilation issue.
This test ensures that fused_fp8_set_kv_buffer correctly handles
k_scale/v_scale when they are 0-dimensional tensors (torch.nn.Parameter).
Previously, Triton would treat 0-D tensor arguments as pointers,
causing a type error when performing "1.0 / k_scale" inside the kernel.
The fix converts tensor scales to Python floats in the wrapper.
"""
device = torch.device("cuda")
num_tokens = 4
num_kv_heads = 2
head_dim = 64
page_size = 16
total_slots = page_size
k = torch.randn(
num_tokens, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16
)
v = torch.randn_like(k)
k_cache = torch.empty(
total_slots,
num_kv_heads,
head_dim,
device=device,
dtype=torch.float8_e4m3fn,
)
v_cache = torch.empty_like(k_cache)
cache_loc = torch.arange(num_tokens, device=device, dtype=torch.int32)
# Use 0D tensor form of scale to reproduce the original bug scenario
k_scale = torch.tensor(1.0, device=device, dtype=torch.float32)
v_scale = torch.tensor(1.0, device=device, dtype=torch.float32)
# Old code would trigger Triton's IncompatibleTypeError here
# New code should handle this gracefully by converting to float
fused_fp8_set_kv_buffer(
k,
v,
k_cache,
v_cache,
cache_loc,
k_scale=k_scale,
v_scale=v_scale,
page_size=page_size,
use_triton=True,
)
# If we get here without exception, the regression is fixed
def test_fp8_kv_kernel_cuda_graph_compatible(self):
"""
Regression test for CUDA graph capture compatibility.
This test ensures that fused_fp8_set_kv_buffer works correctly within
CUDA graph capture, which is used in production for performance.
Previously, float(k_scale) caused GPU→CPU synchronization, triggering
cudaErrorStreamCaptureUnsupported during graph capture. The fix computes
inverse scales purely on GPU using tensor operations.
"""
device = torch.device("cuda")
num_tokens = 4
num_kv_heads = 2
head_dim = 64
page_size = 16
total_slots = page_size
k = torch.randn(
num_tokens, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16
)
v = torch.randn_like(k)
k_cache = torch.empty(
total_slots,
num_kv_heads,
head_dim,
device=device,
dtype=torch.float8_e4m3fn,
)
v_cache = torch.empty_like(k_cache)
cache_loc = torch.arange(num_tokens, device=device, dtype=torch.int32)
# Use 0D tensor scales (like nn.Parameter) to reproduce production scenario
k_scale = torch.tensor(1.0, device=device, dtype=torch.float32)
v_scale = torch.tensor(1.0, device=device, dtype=torch.float32)
# Test that kernel works under CUDA graph capture
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
# Old code would fail here with cudaErrorStreamCaptureUnsupported
# New code should succeed because all operations stay on GPU
fused_fp8_set_kv_buffer(
k,
v,
k_cache,
v_cache,
cache_loc,
k_scale=k_scale,
v_scale=v_scale,
page_size=page_size,
use_triton=True,
)
# Replay the graph to verify it works
graph.replay()
# If we get here without exception, CUDA graph compatibility is confirmed
def test_fp8_kv_kernel_cuda_graph_compatible_no_scale(self):
"""
Regression test for CUDA graph capture compatibility without scales.
This test ensures that fused_fp8_set_kv_buffer works correctly within
CUDA graph capture when k_scale/v_scale are None (use_provided_scale=False).
Previously, the code created new GPU tensors (torch.tensor(1.0, device=...))
during graph capture, triggering cudaErrorStreamCaptureUnsupported.
The fix passes dummy pointers when use_provided_scale=False, as the kernel
uses constant 1.0 and Triton optimizes away the pointer loads.
"""
device = torch.device("cuda")
num_tokens = 4
num_kv_heads = 2
head_dim = 64
page_size = 16
total_slots = page_size
k = torch.randn(
num_tokens, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16
)
v = torch.randn_like(k)
k_cache = torch.empty(
total_slots,
num_kv_heads,
head_dim,
device=device,
dtype=torch.float8_e4m3fn,
)
v_cache = torch.empty_like(k_cache)
cache_loc = torch.arange(num_tokens, device=device, dtype=torch.int32)
# Test that kernel works under CUDA graph capture WITHOUT scales
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
# No k_scale/v_scale provided - use_provided_scale=False branch
# Old code would fail here with cudaErrorStreamCaptureUnsupported
# New code should succeed by using dummy pointers
fused_fp8_set_kv_buffer(
k,
v,
k_cache,
v_cache,
cache_loc,
page_size=page_size,
use_triton=True,
)
# Replay the graph to verify it works
graph.replay()
# If we get here without exception, no-scale CUDA graph compatibility is confirmed
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,152 @@
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.batch_overlap.two_batch_overlap import (
compute_split_seq_index,
compute_split_token_index,
)
from sglang.srt.environ import envs
from sglang.srt.model_executor.forward_batch_info import ForwardMode
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()

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()

View File

@@ -0,0 +1,320 @@
""" """
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)

View File

@@ -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()

View File

@@ -0,0 +1,186 @@
"""
Unit tests for weight validation and cache cleanup logic.
Tests the fix for issue #14754 - ensuring that missing shards do not trigger
entire cache deletion, which can cause race conditions in multi-process scenarios.
"""
import json
import os
import struct
import tempfile
import unittest
from sglang.srt.model_loader.ci_weight_validation import (
_check_index_files_exist,
_validate_sharded_model,
)
class TestWeightValidation(unittest.TestCase):
"""Tests for weight validation functions."""
def test_validate_sharded_model_missing_shard(self):
"""
Test that missing shards are detected correctly.
This is the core test for issue #14754 fix: when a shard is missing,
the validation should return is_valid=False with an error message
containing "Missing", but corrupted_files should be empty (indicating
this is a missing shard issue, not a corruption issue).
This distinction is critical because:
- Missing shards: should NOT delete cache (other processes may be using it)
- Corrupted files: should delete only the corrupted files selectively
"""
with tempfile.TemporaryDirectory() as tmpdir:
# Create partial shards (missing shard 3)
for i in [1, 2]: # Missing shard 3
open(
os.path.join(tmpdir, f"model-0000{i}-of-00003.safetensors"), "w"
).close()
# Create index file
index_data = {
"weight_map": {
"layer1": "model-00001-of-00003.safetensors",
"layer2": "model-00002-of-00003.safetensors",
"layer3": "model-00003-of-00003.safetensors",
}
}
with open(os.path.join(tmpdir, "model.safetensors.index.json"), "w") as f:
json.dump(index_data, f)
weight_files = [
os.path.join(tmpdir, f"model-0000{i}-of-00003.safetensors")
for i in [1, 2]
]
is_valid, error_msg, corrupted_files = _validate_sharded_model(
tmpdir, weight_files
)
self.assertFalse(is_valid)
self.assertIn("Missing", error_msg)
# CRITICAL: corrupted_files should be empty for missing shards
# This is what prevents entire cache deletion
self.assertEqual(corrupted_files, [])
def test_validate_sharded_model_all_present(self):
"""Test that complete shards pass validation."""
with tempfile.TemporaryDirectory() as tmpdir:
# Create all shards with valid safetensors header
for i in [1, 2, 3]:
filepath = os.path.join(tmpdir, f"model-0000{i}-of-00003.safetensors")
# Create a minimal valid safetensors file
# Header: 8 bytes for header size + JSON header
header = b'{"__metadata__":{}}'
header_size = len(header)
with open(filepath, "wb") as f:
f.write(struct.pack("<Q", header_size))
f.write(header)
# Create index file
index_data = {
"weight_map": {
"layer1": "model-00001-of-00003.safetensors",
"layer2": "model-00002-of-00003.safetensors",
"layer3": "model-00003-of-00003.safetensors",
}
}
with open(os.path.join(tmpdir, "model.safetensors.index.json"), "w") as f:
json.dump(index_data, f)
weight_files = [
os.path.join(tmpdir, f"model-0000{i}-of-00003.safetensors")
for i in [1, 2, 3]
]
is_valid, error_msg, corrupted_files = _validate_sharded_model(
tmpdir, weight_files
)
self.assertTrue(is_valid)
self.assertIsNone(error_msg)
self.assertEqual(corrupted_files, [])
def test_validate_sharded_model_corrupted_shard(self):
"""
Test that corrupted shards are detected and returned in corrupted_files.
This tests the other branch: when a file exists but is corrupted
(invalid safetensors format), it should be added to corrupted_files
so that selective cleanup can remove just that file.
"""
with tempfile.TemporaryDirectory() as tmpdir:
# Create shard 1 as valid
filepath1 = os.path.join(tmpdir, "model-00001-of-00003.safetensors")
header = b'{"__metadata__":{}}'
with open(filepath1, "wb") as f:
f.write(struct.pack("<Q", len(header)))
f.write(header)
# Create shard 2 as corrupted (invalid header)
filepath2 = os.path.join(tmpdir, "model-00002-of-00003.safetensors")
with open(filepath2, "wb") as f:
f.write(b"invalid data that is not a valid safetensors file")
# Create shard 3 as valid
filepath3 = os.path.join(tmpdir, "model-00003-of-00003.safetensors")
with open(filepath3, "wb") as f:
f.write(struct.pack("<Q", len(header)))
f.write(header)
# Create index file
index_data = {
"weight_map": {
"layer1": "model-00001-of-00003.safetensors",
"layer2": "model-00002-of-00003.safetensors",
"layer3": "model-00003-of-00003.safetensors",
}
}
with open(os.path.join(tmpdir, "model.safetensors.index.json"), "w") as f:
json.dump(index_data, f)
weight_files = [filepath1, filepath2, filepath3]
is_valid, error_msg, corrupted_files = _validate_sharded_model(
tmpdir, weight_files
)
self.assertFalse(is_valid)
self.assertIn("Corrupt", error_msg)
# The corrupted file should be identified
self.assertEqual(len(corrupted_files), 1)
self.assertIn("model-00002-of-00003.safetensors", corrupted_files[0])
def test_broken_index_symlink_detected(self):
"""
Test that broken index symlinks are detected and cause validation to fail.
When an index file is a symlink pointing to a non-existent blob,
validation should fail (to trigger re-download) rather than silently
continuing and causing timeout during actual loading.
"""
with tempfile.TemporaryDirectory() as tmpdir:
# Create a broken symlink for the index file
index_path = os.path.join(tmpdir, "model.safetensors.index.json")
non_existent_blob = os.path.join(tmpdir, "blobs", "nonexistent_hash")
os.symlink(non_existent_blob, index_path)
# Verify it's a broken symlink
self.assertTrue(os.path.islink(index_path))
self.assertFalse(os.path.exists(index_path))
# Check should fail for broken symlink
is_valid, error_msg = _check_index_files_exist(tmpdir)
self.assertFalse(is_valid)
self.assertIn("Broken", error_msg)
# The broken symlink should have been cleaned up
self.assertFalse(os.path.exists(index_path))
self.assertFalse(os.path.islink(index_path))
if __name__ == "__main__":
unittest.main()

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()

View File

@@ -0,0 +1,161 @@
"""
Test Whisper model with CUDA graph support.
This test verifies that:
1. Whisper model works correctly with CUDA graph enabled (default)
2. Cross-attention KV cache is properly managed through RadixAttention
3. Output is consistent between CUDA graph and non-CUDA-graph modes
Usage:
python test_whisper_cuda_graph.py
Requires:
- A GPU with sufficient memory
- openai-whisper model (e.g., openai/whisper-large-v3)
- An audio file or URL for testing
"""
import io
import unittest
import requests
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,
)
WHISPER_MODEL = "openai/whisper-large-v3"
TEST_AUDIO_URL = "https://huggingface.co/datasets/Narsil/asr_dummy/resolve/main/1.flac"
TEST_AUDIO_LOCAL = "/tmp/test_whisper_audio.flac"
def get_audio_bytes():
"""Get audio bytes, downloading if necessary."""
import os
if os.path.exists(TEST_AUDIO_LOCAL):
with open(TEST_AUDIO_LOCAL, "rb") as f:
return f.read()
resp = requests.get(TEST_AUDIO_URL, timeout=30)
resp.raise_for_status()
with open(TEST_AUDIO_LOCAL, "wb") as f:
f.write(resp.content)
return resp.content
class TestWhisperCudaGraph(CustomTestCase):
"""Test Whisper with CUDA graph enabled (default behavior)."""
@classmethod
def setUpClass(cls):
cls.model = WHISPER_MODEL
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=[
"--served-model-name",
"whisper",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def _transcribe(self, language="en"):
"""Send a transcription request via OpenAI-compatible audio endpoint."""
audio_bytes = get_audio_bytes()
response = requests.post(
self.base_url + "/v1/audio/transcriptions",
files={"file": ("audio.ogg", io.BytesIO(audio_bytes), "audio/ogg")},
data={
"model": "whisper",
"language": language,
},
)
self.assertEqual(response.status_code, 200, response.text)
return response.json()
def test_basic_transcription(self):
"""Test that basic transcription works with CUDA graph."""
result = self._transcribe()
self.assertIn("text", result)
text = result["text"]
self.assertTrue(len(text) > 0, "Transcription should not be empty")
print(f"Transcription: {text}")
def test_multiple_sequential_requests(self):
"""Test multiple sequential requests to verify CUDA graph replay consistency."""
results = []
for i in range(3):
result = self._transcribe()
self.assertIn("text", result)
results.append(result["text"])
print(f"Request {i+1}: {result['text'][:80]}...")
# All transcriptions of the same audio should be identical
for i in range(1, len(results)):
self.assertEqual(
results[0],
results[i],
f"Transcription {i+1} differs from first transcription",
)
def test_transcription_quality(self):
"""Test that transcription quality is reasonable (contains expected words)."""
result = self._transcribe()
text = result["text"].lower()
# The test audio is a LibriSpeech sample about stew for dinner
self.assertIn("stew", text, f"Expected 'stew' in transcription: {text}")
self.assertIn("dinner", text, f"Expected 'dinner' in transcription: {text}")
print(f"Quality check passed: {result['text'][:80]}...")
class TestWhisperNoCudaGraph(CustomTestCase):
"""Test Whisper with CUDA graph explicitly disabled for comparison."""
@classmethod
def setUpClass(cls):
cls.model = WHISPER_MODEL
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=[
"--served-model-name",
"whisper",
"--disable-cuda-graph",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_basic_transcription_no_cuda_graph(self):
"""Test that transcription works without CUDA graph (baseline)."""
audio_bytes = get_audio_bytes()
response = requests.post(
self.base_url + "/v1/audio/transcriptions",
files={"file": ("audio.ogg", io.BytesIO(audio_bytes), "audio/ogg")},
data={
"model": "whisper",
"language": "en",
},
)
self.assertEqual(response.status_code, 200, response.text)
result = response.json()
self.assertIn("text", result)
self.assertTrue(len(result["text"]) > 0)
print(f"No CUDA graph transcription: {result['text'][:80]}...")
if __name__ == "__main__":
unittest.main(verbosity=3)

Some files were not shown because too many files have changed in this diff Show More