EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Enable torch.compile for triton backend (#1422)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-09-14 15:38:37 -07:00
committed by GitHub
parent e3fc4658f4
commit 9463bc1385
9 changed files with 134 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
test/srt/test_bench_latency.py
View File

@@ -1,4 +1,3 @@
import os
import subprocess
import unittest
@@ -6,77 +5,25 @@ from sglang.srt.utils import kill_child_process
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_MOE_MODEL_NAME_FOR_TEST,
is_in_ci,
run_bench_latency,
)
class TestBenchLatency(unittest.TestCase):
def test_default(self):
command = [
"python3",
"-m",
"sglang.bench_latency",
"--model-path",
DEFAULT_MODEL_NAME_FOR_TEST,
"--batch-size",
"1",
"--input",
"128",
"--output",
"8",
]
process = subprocess.Popen(
command, stdout=subprocess.PIPE, stderr=subprocess.PIPE
)
output_throughput = run_bench_latency(DEFAULT_MODEL_NAME_FOR_TEST, [])
try:
stdout, stderr = process.communicate()
output = stdout.decode()
error = stderr.decode()
print(f"Output: {output}")
print(f"Error: {error}")
lastline = output.split("\n")[-3]
value = float(lastline.split(" ")[-2])
if os.getenv("SGLANG_IS_IN_CI", "false") == "true":
assert value > 130
finally:
kill_child_process(process.pid)
if is_in_ci():
assert output_throughput > 130, f"{output_throughput=}"
def test_moe_default(self):
command = [
"python3",
"-m",
"sglang.bench_latency",
"--model",
DEFAULT_MOE_MODEL_NAME_FOR_TEST,
"--batch-size",
"1",
"--input",
"128",
"--output",
"8",
"--tp",
"2",
]
process = subprocess.Popen(
command, stdout=subprocess.PIPE, stderr=subprocess.PIPE
output_throughput = run_bench_latency(
DEFAULT_MOE_MODEL_NAME_FOR_TEST, ["--tp", "2"]
)
try:
stdout, stderr = process.communicate()
output = stdout.decode()
error = stderr.decode()
print(f"Output: {output}")
print(f"Error: {error}")
lastline = output.split("\n")[-3]
value = float(lastline.split(" ")[-2])
if os.getenv("SGLANG_IS_IN_CI", "false") == "true":
assert value > 125
finally:
kill_child_process(process.pid)
if is_in_ci():
assert output_throughput > 125, f"{output_throughput=}"
if __name__ == "__main__":

16
test/srt/test_bench_serving.py
View File

@@ -1,9 +1,9 @@
import os
import unittest
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_MOE_MODEL_NAME_FOR_TEST,
is_in_ci,
run_bench_serving,
)
@@ -18,7 +18,7 @@ class TestBenchServing(unittest.TestCase):
other_server_args=[],
)
if os.getenv("SGLANG_IS_IN_CI", "false") == "true":
if is_in_ci():
assert res["output_throughput"] > 2600
def test_offline_throughput_without_radix_cache(self):
@@ -29,7 +29,7 @@ class TestBenchServing(unittest.TestCase):
other_server_args=["--disable-radix-cache"],
)
if os.getenv("SGLANG_IS_IN_CI", "false") == "true":
if is_in_ci():
assert res["output_throughput"] > 2800
def test_offline_throughput_without_chunked_prefill(self):
@@ -40,7 +40,7 @@ class TestBenchServing(unittest.TestCase):
other_server_args=["--chunked-prefill-size", "-1"],
)
if os.getenv("SGLANG_IS_IN_CI", "false") == "true":
if is_in_ci():
assert res["output_throughput"] > 2600
def test_offline_throughput_with_triton_attention_backend(self):
@@ -56,7 +56,7 @@ class TestBenchServing(unittest.TestCase):
],
)
if os.getenv("SGLANG_IS_IN_CI", "false") == "true":
if is_in_ci():
assert res["output_throughput"] > 2600
def test_online_latency_default(self):
@@ -67,7 +67,7 @@ class TestBenchServing(unittest.TestCase):
other_server_args=[],
)
if os.getenv("SGLANG_IS_IN_CI", "false") == "true":
if is_in_ci():
assert res["median_e2e_latency_ms"] < 12000
assert res["median_ttft_ms"] < 80
assert res["median_itl_ms"] < 12
@@ -80,7 +80,7 @@ class TestBenchServing(unittest.TestCase):
other_server_args=["--tp", "2"],
)
if os.getenv("SGLANG_IS_IN_CI", "false") == "true":
if is_in_ci():
assert res["output_throughput"] > 1850
def test_moe_offline_throughput_without_radix_cache(self):
@@ -91,7 +91,7 @@ class TestBenchServing(unittest.TestCase):
other_server_args=["--tp", "2", "--disable-radix-cache"],
)
if os.getenv("SGLANG_IS_IN_CI", "false") == "true":
if is_in_ci():
assert res["output_throughput"] > 1950

6
test/srt/test_moe_eval_accuracy_large.py
View File

@@ -42,7 +42,7 @@ class TestEvalAccuracyLarge(unittest.TestCase):
)
metrics = run_eval(args)
assert metrics["score"] >= 0.625, f"{metrics}"
assert metrics["score"] >= 0.62, f"{metrics}"
def test_human_eval(self):
args = SimpleNamespace(
@@ -54,7 +54,7 @@ class TestEvalAccuracyLarge(unittest.TestCase):
)
metrics = run_eval(args)
assert metrics["score"] >= 0.425, f"{metrics}"
assert metrics["score"] >= 0.42, f"{metrics}"
def test_mgsm_en(self):
args = SimpleNamespace(
@@ -66,7 +66,7 @@ class TestEvalAccuracyLarge(unittest.TestCase):
)
metrics = run_eval(args)
assert metrics["score"] >= 0.625, f"{metrics}"
assert metrics["score"] >= 0.62, f"{metrics}"
if __name__ == "__main__":

55
test/srt/test_triton_attn_backend.py
View File

@@ -1,3 +1,4 @@
import subprocess
import unittest
from types import SimpleNamespace
@@ -7,37 +8,49 @@ 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_latency,
)
class TestTritonAttnBackend(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
def test_latency(self):
output_throughput = run_bench_latency(
DEFAULT_MODEL_NAME_FOR_TEST,
[
"--attention-backend",
"triton",
"--enable-torch-compile",
],
)
if is_in_ci():
assert output_throughput > 155, f"{output_throughput=}"
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", "triton"],
)
@classmethod
def tearDownClass(cls):
kill_child_process(cls.process.pid)
try:
args = SimpleNamespace(
base_url=base_url,
model=model,
eval_name="mmlu",
num_examples=64,
num_threads=32,
)
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)
assert metrics["score"] >= 0.65
metrics = run_eval(args)
assert metrics["score"] >= 0.65
finally:
kill_child_process(process.pid)
if __name__ == "__main__":