264 lines
8.1 KiB
Python
264 lines
8.1 KiB
Python
import os
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
|
|
import requests
|
|
|
|
from sglang.srt.utils import kill_process_tree
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
|
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,
|
|
)
|
|
|
|
register_cuda_ci(est_time=900, suite="nightly-8-gpu-h200", nightly=True)
|
|
|
|
DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp"
|
|
|
|
# Global list to collect results
|
|
TEST_RESULTS = []
|
|
|
|
|
|
class TestDeepseekV32_TP(CustomTestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.model = DEEPSEEK_V32_MODEL_PATH
|
|
cls.base_url = DEFAULT_URL_FOR_TEST
|
|
# Pure TP configuration without --dp and --enable-dp-attention
|
|
other_args = [
|
|
"--trust-remote-code",
|
|
"--attention-backend",
|
|
"nsa",
|
|
"--nsa-prefill-backend",
|
|
"flashmla_sparse",
|
|
"--nsa-decode-backend",
|
|
"flashmla_kv",
|
|
"--tp",
|
|
"8",
|
|
]
|
|
cls.process = popen_launch_server(
|
|
cls.model,
|
|
cls.base_url,
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=other_args,
|
|
)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
kill_process_tree(cls.process.pid)
|
|
|
|
def test_a_gsm8k(
|
|
self,
|
|
): # Append an "a" to make this test run first (alphabetically) to warm up the server
|
|
args = SimpleNamespace(
|
|
num_shots=20,
|
|
data_path=None,
|
|
num_questions=1400,
|
|
parallel=1400,
|
|
max_new_tokens=512,
|
|
host="http://127.0.0.1",
|
|
port=int(self.base_url.split(":")[-1]),
|
|
)
|
|
metrics = run_eval_few_shot_gsm8k(args)
|
|
print(f"{metrics=}")
|
|
|
|
if is_in_ci():
|
|
TEST_RESULTS.append(
|
|
{
|
|
"variant": "pure_tp",
|
|
"prefill_backend": "flashmla_sparse",
|
|
"decode_backend": "flashmla_kv",
|
|
"kv_cache": "fp16",
|
|
"accuracy": metrics["accuracy"],
|
|
}
|
|
)
|
|
self.assertGreater(metrics["accuracy"], 0.935)
|
|
|
|
|
|
class TestDeepseekV32_Partial_TP(CustomTestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.model = DEEPSEEK_V32_MODEL_PATH
|
|
cls.base_url = DEFAULT_URL_FOR_TEST
|
|
# Partial TP configuration with dp=4 and dp-attention enabled
|
|
other_args = [
|
|
"--trust-remote-code",
|
|
"--attention-backend",
|
|
"nsa",
|
|
"--nsa-prefill-backend",
|
|
"flashmla_sparse",
|
|
"--nsa-decode-backend",
|
|
"flashmla_kv",
|
|
"--tp",
|
|
"8",
|
|
"--dp",
|
|
"4",
|
|
"--enable-dp-attention",
|
|
]
|
|
cls.process = popen_launch_server(
|
|
cls.model,
|
|
cls.base_url,
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=other_args,
|
|
)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
kill_process_tree(cls.process.pid)
|
|
|
|
def test_a_gsm8k(
|
|
self,
|
|
): # Append an "a" to make this test run first (alphabetically) to warm up the server
|
|
args = SimpleNamespace(
|
|
num_shots=20,
|
|
data_path=None,
|
|
num_questions=1400,
|
|
parallel=1400,
|
|
max_new_tokens=512,
|
|
host="http://127.0.0.1",
|
|
port=int(self.base_url.split(":")[-1]),
|
|
)
|
|
metrics = run_eval_few_shot_gsm8k(args)
|
|
print(f"{metrics=}")
|
|
|
|
if is_in_ci():
|
|
TEST_RESULTS.append(
|
|
{
|
|
"variant": "partial_tp",
|
|
"prefill_backend": "flashmla_sparse",
|
|
"decode_backend": "flashmla_kv",
|
|
"kv_cache": "fp16",
|
|
"accuracy": metrics["accuracy"],
|
|
}
|
|
)
|
|
self.assertGreater(metrics["accuracy"], 0.935)
|
|
|
|
|
|
class TestDeepseekV32_TP_MTP(CustomTestCase):
|
|
"""Test DeepSeek V3.2 with pure TP + MTP (EAGLE speculative decoding)."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.model = DEEPSEEK_V32_MODEL_PATH
|
|
cls.base_url = DEFAULT_URL_FOR_TEST
|
|
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",
|
|
]
|
|
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(
|
|
num_shots=20,
|
|
data_path=None,
|
|
num_questions=1400,
|
|
parallel=1400,
|
|
max_new_tokens=512,
|
|
host="http://127.0.0.1",
|
|
port=int(self.base_url.split(":")[-1]),
|
|
)
|
|
metrics = run_eval_few_shot_gsm8k(args)
|
|
print(f"{metrics=}")
|
|
|
|
server_info = requests.get(self.base_url + "/get_server_info")
|
|
avg_spec_accept_length = server_info.json()["internal_states"][0][
|
|
"avg_spec_accept_length"
|
|
]
|
|
print(f"{avg_spec_accept_length=}")
|
|
|
|
if is_in_ci():
|
|
TEST_RESULTS.append(
|
|
{
|
|
"variant": "tp_mtp",
|
|
"prefill_backend": "flashmla_sparse",
|
|
"decode_backend": "flashmla_kv",
|
|
"kv_cache": "fp16",
|
|
"accuracy": metrics["accuracy"],
|
|
"avg_spec_accept_length": avg_spec_accept_length,
|
|
}
|
|
)
|
|
self.assertGreater(metrics["accuracy"], 0.935)
|
|
self.assertGreater(avg_spec_accept_length, 2.5)
|
|
|
|
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"{acc_length=:.2f} {speed=:.2f}")
|
|
|
|
if is_in_ci():
|
|
# Update last result with speed data
|
|
if TEST_RESULTS and TEST_RESULTS[-1]["variant"] == "tp_mtp":
|
|
TEST_RESULTS[-1]["speed"] = speed
|
|
|
|
# Write the summary table after all tests complete
|
|
_write_summary_table()
|
|
|
|
self.assertGreater(acc_length, 2.5)
|
|
self.assertGreater(speed, 110)
|
|
|
|
|
|
def _format_optional_metric(value, fmt=".2f", suffix=""):
|
|
"""Format an optional metric value, returning '-' if not available."""
|
|
if value is None:
|
|
return "-"
|
|
return f"{value:{fmt}}{suffix}"
|
|
|
|
|
|
def _write_summary_table():
|
|
"""Write a markdown table with all test results."""
|
|
if not TEST_RESULTS:
|
|
return
|
|
|
|
gpu_config = os.getenv("GPU_CONFIG", "8-gpu-h200")
|
|
|
|
# Build table header - keep original columns + add MTP-specific ones
|
|
summary = (
|
|
f"### {DEEPSEEK_V32_MODEL_PATH} GSM8K Accuracy (TP Tests) [{gpu_config}]\n\n"
|
|
)
|
|
summary += "| Variant | Prefill Backend | Decode Backend | KV Cache | Accuracy | Spec Acc Len | Speed |\n"
|
|
summary += "|---------|-----------------|----------------|----------|----------|--------------|-------|\n"
|
|
|
|
# Add each result as a row
|
|
for result in TEST_RESULTS:
|
|
summary += (
|
|
f"| {result['variant']} | {result['prefill_backend']} | "
|
|
f"{result['decode_backend']} | {result['kv_cache']} | "
|
|
f"{result['accuracy']:.3f} | "
|
|
f"{_format_optional_metric(result.get('avg_spec_accept_length'))} | "
|
|
f"{_format_optional_metric(result.get('speed'), '.1f', ' tok/s')} |\n"
|
|
)
|
|
|
|
write_github_step_summary(summary)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|