minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CI] Trim the base-c 4-gpu-h100 stage from 5 shards to 4 (#35407)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-19 00:48:07 -07:00
committed by GitHub
parent e614121866
commit ccbe380028
12 changed files with 502 additions and 1099 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/test/kits/pd_parity_kit.py
+70
View File
@@ -0,0 +1,70 @@
import time
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.server_fixtures.disaggregation_fixture import assert_process_healthy
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_server,
)
class PDLogprobParityMixin:
# Mix in before the PD server fixture, which owns the P/D launches.
reference_parallel_args = []
baseline_args = []
@staticmethod
def generate(base_url):
response = requests.post(
base_url + "/generate",
json={
"input_ids": [1] + [100 + i % 1000 for i in range(256)],
"sampling_params": {
"temperature": 0,
"max_new_tokens": 4,
"ignore_eos": True,
},
"return_logprob": True,
"top_logprobs_num": 5,
},
timeout=120,
)
response.raise_for_status()
return response.json()["meta_info"]
def test_logprob_parity(self):
baseline = popen_launch_server(
self.model,
self.lb_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=self.reference_parallel_args
+ ["--trust-remote-code"]
+ self.baseline_args,
env=self.extra_prefill_env,
)
try:
reference = self.generate(self.lb_url)
finally:
kill_process_tree(baseline.pid, wait_timeout=60)
time.sleep(5)
self.launch_all()
disaggregated = self.generate(self.lb_url)
reference_logprobs = reference["output_token_logprobs"]
disaggregated_logprobs = disaggregated["output_token_logprobs"]
self.assertEqual(
[item[1] for item in reference_logprobs],
[item[1] for item in disaggregated_logprobs],
)
self.assertEqual(len(reference_logprobs), 4)
for reference_item, disaggregated_item in zip(
reference_logprobs, disaggregated_logprobs
):
self.assertAlmostEqual(reference_item[0], disaggregated_item[0], delta=0.05)
assert_process_healthy(self, "load balancer", self.process_lb, self.lb_url)
assert_process_healthy(self, "prefill", self.process_prefill, self.prefill_url)
assert_process_healthy(self, "decode", self.process_decode, self.decode_url)
test/registered/disaggregation/test_disaggregation_kimi_linear.py
+13 -75
View File
@@ -1,24 +1,17 @@
import time
import unittest
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.pd_parity_kit import PDLogprobParityMixin
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
assert_process_healthy,
)
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_server,
)
register_cuda_ci(est_time=480, stage="base-c", runner_config="4-gpu-h100")
KIMI_LINEAR_MODEL = "yujiepan/kimi-linear-tiny-random"
SERVER_ENV = {"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM": "0"}
SERVER_ARGS = [
DETERMINISTIC_ARGS = [
"--skip-tokenizer-init",
"--random-seed",
"1",
@@ -34,74 +27,19 @@ SERVER_ARGS = [
]
class TestKimiLinearHeterogeneousTPDisaggregation(PDDisaggregationServerBase):
class TestKimiLinearHeterogeneousTPDisaggregation(
PDLogprobParityMixin, PDDisaggregationServerBase
):
model = KIMI_LINEAR_MODEL
extra_prefill_env = SERVER_ENV
extra_decode_env = SERVER_ENV
prefill_tp_size = 2
decode_tp_size = 1
decode_base_gpu_id = 2
reference_parallel_args = ["--tp-size", "2"]
extra_prefill_args = SERVER_ARGS
extra_decode_args = SERVER_ARGS
extra_prefill_env = SERVER_ENV
extra_decode_env = SERVER_ENV
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = KIMI_LINEAR_MODEL
@staticmethod
def generate(base_url):
response = requests.post(
base_url + "/generate",
json={
"input_ids": [1] + [100 + i % 1000 for i in range(256)],
"sampling_params": {
"temperature": 0,
"max_new_tokens": 4,
"ignore_eos": True,
},
"return_logprob": True,
"top_logprobs_num": 5,
},
timeout=120,
)
response.raise_for_status()
return response.json()["meta_info"]
def test_logprob_parity(self):
baseline = popen_launch_server(
self.model,
self.lb_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=self.reference_parallel_args
+ ["--trust-remote-code"]
+ SERVER_ARGS,
env=SERVER_ENV,
)
try:
reference = self.generate(self.lb_url)
finally:
kill_process_tree(baseline.pid, wait_timeout=60)
time.sleep(5)
self.launch_all()
disaggregated = self.generate(self.lb_url)
reference_logprobs = reference["output_token_logprobs"]
disaggregated_logprobs = disaggregated["output_token_logprobs"]
self.assertEqual(
[item[1] for item in reference_logprobs],
[item[1] for item in disaggregated_logprobs],
)
self.assertEqual(len(reference_logprobs), 4)
for reference_item, disaggregated_item in zip(
reference_logprobs, disaggregated_logprobs
):
self.assertAlmostEqual(reference_item[0], disaggregated_item[0], delta=0.05)
assert_process_healthy(self, "load balancer", self.process_lb, self.lb_url)
assert_process_healthy(self, "prefill", self.process_prefill, self.prefill_url)
assert_process_healthy(self, "decode", self.process_decode, self.decode_url)
baseline_args = DETERMINISTIC_ARGS
extra_prefill_args = DETERMINISTIC_ARGS
extra_decode_args = DETERMINISTIC_ARGS
class TestKimiLinearPipelineDisaggregation(TestKimiLinearHeterogeneousTPDisaggregation):
@@ -109,7 +47,7 @@ class TestKimiLinearPipelineDisaggregation(TestKimiLinearHeterogeneousTPDisaggre
decode_tp_size = 1
decode_base_gpu_id = 2
reference_parallel_args = ["--tp-size", "1", "--pp-size", "2"]
extra_prefill_args = SERVER_ARGS + ["--pp-size", "2"]
extra_prefill_args = DETERMINISTIC_ARGS + ["--pp-size", "2"]
if __name__ == "__main__":
test/registered/disaggregation/test_disaggregation_unified_memory.py
+16 -101
View File
@@ -1,53 +1,24 @@
"""PD disaggregation with --enable-unified-memory (MLA hybrid-Mamba).
Guards the unified-memory PD transfer scheme end to end: whole page-envelope
KV registration (`UnifiedMLATokenToKVPool.get_contiguous_buf_infos`), whole
slot-envelope KDA/mamba state transfer, virtual->physical index translation at
the prefill send / decode prealloc sites, and the compaction move gate. A
regression in any of them shifts the decode-side KV/state bytes and breaks
logprob parity with the non-PD unified-memory reference.
`--attention-backend` is deliberately NOT pinned, matching
`models_e2e/test_kimi_linear_unified_memory.py`, which documents that pinning
hides defects reachable only under the resolved default. The transferred bytes
are backend-independent, so the default (fa3 on this suite's H100 runner) covers
this file's subject either way. The linear-attn/Mamba backends stay pinned to
triton -- the page-major layout requires them.
`--enable-deterministic-inference` is deliberately NOT set. It would only guard
against batch-shape-dependent kernel variation, and the reference and P+D paths
run the same shapes: measured, two fresh servers on separate GPUs produce
bit-identical logits without it. Setting it would narrow the test to the
batch-invariant op set and a non-default sampling backend -- a less
representative config -- and couple a PD-transfer test to the deterministic code
path, so a defect there would fail this file for an unrelated reason.
"""
import time
import unittest
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.pd_parity_kit import PDLogprobParityMixin
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
assert_process_healthy,
)
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_server,
)
register_cuda_ci(est_time=900, stage="base-c", runner_config="4-gpu-h100")
register_cuda_ci(est_time=900, stage="base-b", runner_config="2-gpu-large")
KIMI_LINEAR_MODEL = "yujiepan/kimi-linear-tiny-random"
SERVER_ENV = {"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM": "0"}
SERVER_ARGS = [
# --attention-backend and --enable-deterministic-inference are deliberately
# absent: bytes are backend-independent and both paths run identical shapes.
UNIFIED_MEMORY_ARGS = [
"--skip-tokenizer-init",
"--random-seed",
"1",
"--enable-unified-memory",
# Page-major layout requires triton for both.
"--linear-attn-backend",
"triton",
"--mamba-backend",
@@ -63,75 +34,19 @@ SERVER_ARGS = [
]
class TestUnifiedMemoryDisaggregation(PDDisaggregationServerBase):
class TestUnifiedMemoryDisaggregation(PDLogprobParityMixin, PDDisaggregationServerBase):
"""1 prefill + 1 decode, both with --enable-unified-memory, vs a non-PD
unified-memory reference server."""
model = KIMI_LINEAR_MODEL
extra_prefill_env = SERVER_ENV
extra_decode_env = SERVER_ENV
prefill_tp_size = 1
decode_tp_size = 1
decode_base_gpu_id = 1
extra_prefill_args = SERVER_ARGS
extra_decode_args = SERVER_ARGS
extra_prefill_env = SERVER_ENV
extra_decode_env = SERVER_ENV
baseline_args = SERVER_ARGS
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.model = KIMI_LINEAR_MODEL
@staticmethod
def generate(base_url):
response = requests.post(
base_url + "/generate",
json={
"input_ids": [1] + [100 + i % 1000 for i in range(256)],
"sampling_params": {
"temperature": 0,
"max_new_tokens": 4,
"ignore_eos": True,
},
"return_logprob": True,
"top_logprobs_num": 5,
},
timeout=120,
)
response.raise_for_status()
return response.json()["meta_info"]
def test_logprob_parity(self):
baseline = popen_launch_server(
self.model,
self.lb_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--trust-remote-code"] + self.baseline_args,
env=SERVER_ENV,
)
try:
reference = self.generate(self.lb_url)
finally:
kill_process_tree(baseline.pid, wait_timeout=60)
time.sleep(5)
self.launch_all()
disaggregated = self.generate(self.lb_url)
reference_logprobs = reference["output_token_logprobs"]
disaggregated_logprobs = disaggregated["output_token_logprobs"]
self.assertEqual(
[item[1] for item in reference_logprobs],
[item[1] for item in disaggregated_logprobs],
)
self.assertEqual(len(reference_logprobs), 4)
for reference_item, disaggregated_item in zip(
reference_logprobs, disaggregated_logprobs
):
self.assertAlmostEqual(reference_item[0], disaggregated_item[0], delta=0.05)
assert_process_healthy(self, "load balancer", self.process_lb, self.lb_url)
assert_process_healthy(self, "prefill", self.process_prefill, self.prefill_url)
assert_process_healthy(self, "decode", self.process_decode, self.decode_url)
baseline_args = UNIFIED_MEMORY_ARGS
extra_prefill_args = UNIFIED_MEMORY_ARGS
extra_decode_args = UNIFIED_MEMORY_ARGS
class TestUnifiedMemoryDisaggregationChunkedPrefill(TestUnifiedMemoryDisaggregation):
@@ -142,10 +57,10 @@ class TestUnifiedMemoryDisaggregationChunkedPrefill(TestUnifiedMemoryDisaggregat
chunk size so any parity break isolates to the PD transfer.
"""
_chunked_args = SERVER_ARGS + ["--chunked-prefill-size", "64"]
_chunked_args = UNIFIED_MEMORY_ARGS + ["--chunked-prefill-size", "64"]
baseline_args = _chunked_args
extra_prefill_args = _chunked_args
extra_decode_args = _chunked_args
baseline_args = _chunked_args
if __name__ == "__main__":
test/registered/ep/test_deepep_small.py
+1 -265
View File
@@ -16,7 +16,7 @@ from sglang.test.test_utils import (
popen_launch_server,
)
register_cuda_ci(est_time=478, stage="base-c", runner_config="4-gpu-h100")
register_cuda_ci(est_time=270, stage="base-c", runner_config="4-gpu-h100")
class TestPureDP(CustomTestCase):
@@ -66,51 +66,6 @@ class TestPureDP(CustomTestCase):
self.assertGreater(metrics["score"], 0.60)
class TestHybridDPTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
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",
"4",
"--enable-dp-attention",
"--dp",
"2",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs-decode",
"128",
"--max-running-requests",
"256",
],
)
@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(metrics)
self.assertGreater(metrics["score"], 0.60)
class TestTP(CustomTestCase):
@classmethod
def setUpClass(cls):
@@ -153,55 +108,6 @@ class TestTP(CustomTestCase):
self.assertGreater(metrics["score"], 0.60)
@unittest.skip("covered in test_deepep_large.py")
class TestNoGatherdBuffer(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
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",
"4",
"--enable-dp-attention",
"--dp",
"4",
"--moe-dense-tp-size",
"1",
"--enable-dp-lm-head",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs-decode",
"32",
"--max-running-requests",
"512",
],
)
@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(metrics)
self.assertGreater(metrics["score"], 0.60)
class TestTBO(CustomTestCase):
@classmethod
def setUpClass(cls):
@@ -254,176 +160,6 @@ class TestTBO(CustomTestCase):
self.assertGreater(metrics["score"], 0.60)
class TestTBOWithTPAttn(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
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",
"4",
"--moe-a2a-backend",
"deepep",
"--enable-two-batch-overlap",
"--cuda-graph-max-bs-decode",
"128",
"--max-running-requests",
"512",
"--mem-fraction-static", # temp fix as DeepEP buffer is too large.
"0.7",
],
env={
**os.environ,
"SGLANG_TBO_DEBUG": "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(metrics)
self.assertGreater(metrics["score"], 0.60)
# There exists bug when using MTP + TBO + attn_tp_size > 1, currently skip that case.
# @unittest.skip("covered in TestMTPWithTPAttnAndTBO")
class TestTBOWithTPAttnAndDenseDP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
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",
"4",
"--moe-dense-tp-size",
"1",
"--moe-a2a-backend",
"deepep",
"--enable-two-batch-overlap",
"--cuda-graph-max-bs-decode",
"128",
"--max-running-requests",
"512",
"--mem-fraction-static", # temp fix as DeepEP buffer is too large.
"0.7",
],
env={
**os.environ,
"SGLANG_TBO_DEBUG": "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(metrics)
self.assertGreater(metrics["score"], 0.60)
@unittest.skip("covered in TestMTPWithTBO")
class TestMTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
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",
"4",
"--enable-dp-attention",
"--dp",
"2",
"--enable-dp-lm-head",
"--moe-a2a-backend",
"deepep",
"--speculative-algo",
"EAGLE",
"--speculative-draft-model-path",
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
"--speculative-num-steps",
"2",
"--speculative-eagle-topk",
"3",
"--speculative-num-draft-tokens",
"3",
"--cuda-graph-max-bs-decode",
"32",
"--max-running-requests",
"64",
],
)
@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(metrics)
self.assertGreater(metrics["score"], 0.60)
server_info = requests.get(self.base_url + "/server_info")
avg_spec_accept_length = server_info.json()["internal_states"][0][
"avg_spec_accept_length"
]
print(
f"###test_gsm8k (deepseek-v3 mtp + dp + tbo):\n"
f"accuracy={metrics['score']=:.3f}\n"
f"{avg_spec_accept_length=:.3f}\n"
)
self.assertGreater(avg_spec_accept_length, 2.1)
class TestMTPWithTBO(CustomTestCase):
@classmethod
def setUpClass(cls):
test/registered/ep/test_deepep_small_extra.py
+167
View File
@@ -0,0 +1,167 @@
import os
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.run_eval import run_eval
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,
)
register_cuda_ci(est_time=210, stage="extra-b", runner_config="4-gpu-h100")
class TestHybridDPTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
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",
"4",
"--enable-dp-attention",
"--dp",
"2",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs-decode",
"128",
"--max-running-requests",
"256",
],
)
@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(metrics)
self.assertGreater(metrics["score"], 0.60)
class TestTBOWithTPAttn(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
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",
"4",
"--moe-a2a-backend",
"deepep",
"--enable-two-batch-overlap",
"--cuda-graph-max-bs-decode",
"128",
"--max-running-requests",
"512",
"--mem-fraction-static", # temp fix as DeepEP buffer is too large.
"0.7",
],
env={
**os.environ,
"SGLANG_TBO_DEBUG": "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(metrics)
self.assertGreater(metrics["score"], 0.60)
# There exists bug when using MTP + TBO + attn_tp_size > 1, currently skip that case.
# @unittest.skip("covered in TestMTPWithTPAttnAndTBO")
class TestTBOWithTPAttnAndDenseDP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
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",
"4",
"--moe-dense-tp-size",
"1",
"--moe-a2a-backend",
"deepep",
"--enable-two-batch-overlap",
"--cuda-graph-max-bs-decode",
"128",
"--max-running-requests",
"512",
"--mem-fraction-static", # temp fix as DeepEP buffer is too large.
"0.7",
],
env={
**os.environ,
"SGLANG_TBO_DEBUG": "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(metrics)
self.assertGreater(metrics["score"], 0.60)
if __name__ == "__main__":
unittest.main()
test/registered/hicache/test_hicache_storage_3fs_backend.py
+1 -1
View File
@@ -13,7 +13,7 @@ from test_hicache_storage_file_backend import HiCacheStorageBaseMixin
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=300, stage="base-c", runner_config="4-gpu-h100")
register_cuda_ci(est_time=300, stage="base-b", runner_config="2-gpu-large")
register_amd_ci(est_time=300, suite="base-b-test-2-gpu-large")
test/registered/models_e2e/test_qwen3_next_models.py
+1 -42
View File
@@ -1,4 +1,3 @@
import os
import unittest
from sglang.test.ci.ci_register import register_cuda_ci
@@ -7,7 +6,7 @@ from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin
from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
register_cuda_ci(est_time=500, stage="base-c", runner_config="4-gpu-h100")
register_cuda_ci(est_time=260, stage="base-c", runner_config="4-gpu-h100")
QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
@@ -54,45 +53,5 @@ class TestQwen3NextLazyExtraBufferLargePage(
other_args = _make_args(page_size=2, track_interval=2)
class TestQwen3NextLazyExtraBufferAllocFail(KLDivergenceMixin, DefaultServerBase):
model = QWEN3_NEXT_MODEL
cache_chunk_size = 64
kl_div_thres = 0.002
other_args = _make_args(page_size=1, track_interval=2)
@classmethod
def setUpClass(cls):
os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1"
os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1"
super().setUpClass()
@classmethod
def tearDownClass(cls):
super().tearDownClass()
os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None)
os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None)
class TestQwen3NextLazyExtraBufferLargePageAllocFail(
KLDivergenceMixin, DefaultServerBase
):
model = QWEN3_NEXT_MODEL
cache_chunk_size = 64
kl_div_thres = 0.002
other_args = _make_args(page_size=2, track_interval=2)
@classmethod
def setUpClass(cls):
os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1"
os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1"
super().setUpClass()
@classmethod
def tearDownClass(cls):
super().tearDownClass()
os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None)
os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None)
if __name__ == "__main__":
unittest.main()
test/registered/models_e2e/test_qwen3_next_models_extra.py
+76
View File
@@ -0,0 +1,76 @@
import os
import unittest
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
register_cuda_ci(est_time=250, stage="extra-b", runner_config="4-gpu-h100")
QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
_COMMON_ARGS = [
"--trust-remote-code",
"--tp-size",
"4",
"--chunked-prefill-size",
"2048",
"--mamba-scheduler-strategy",
"extra_buffer_lazy",
"--attention-backend",
"triton",
]
def _make_args(*, page_size=1, track_interval=2):
return [
*_COMMON_ARGS,
"--mamba-track-interval",
str(track_interval),
"--page-size",
str(page_size),
]
class TestQwen3NextLazyExtraBufferAllocFail(KLDivergenceMixin, DefaultServerBase):
model = QWEN3_NEXT_MODEL
cache_chunk_size = 64
kl_div_thres = 0.002
other_args = _make_args(page_size=1, track_interval=2)
@classmethod
def setUpClass(cls):
os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1"
os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1"
super().setUpClass()
@classmethod
def tearDownClass(cls):
super().tearDownClass()
os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None)
os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None)
class TestQwen3NextLazyExtraBufferLargePageAllocFail(
KLDivergenceMixin, DefaultServerBase
):
model = QWEN3_NEXT_MODEL
cache_chunk_size = 64
kl_div_thres = 0.002
other_args = _make_args(page_size=2, track_interval=2)
@classmethod
def setUpClass(cls):
os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1"
os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1"
super().setUpClass()
@classmethod
def tearDownClass(cls):
super().tearDownClass()
os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None)
os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None)
if __name__ == "__main__":
unittest.main()
test/registered/pp/test_pp_gemma4.py
+153
View File
@@ -0,0 +1,153 @@
import time
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP,
DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
is_in_amd_ci,
is_in_ci,
popen_launch_server,
)
# tp=1 pp=2 -- two GPUs.
register_cuda_ci(est_time=220, stage="base-b", runner_config="2-gpu-large")
@unittest.skipIf(
is_in_amd_ci(),
"Gemma4 PP not yet validated on AMD",
)
class TestGemma4PPAccuracy(unittest.TestCase):
"""End-to-end PP=2 accuracy gate for Gemma4 multimodal.
Gemma4 has full-attention layers with head_dim=512 (FA's max is 256), so
sglang auto-selects the triton attention backend; no manual flag needed.
The 26B BF16 model splits to ~26 GB per stage under PP=2, well within an
H100's 80 GB.
"""
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP
cls.base_url = "http://127.0.0.1:23333"
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
1,
"--pp-size",
2,
"--trust-remote-code",
"--enable-multimodal",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
# Gemma4 is instruction-tuned and doesn't follow few-shot completion
# prompts well — use the chat API (default in run_eval), which scores
# ~0.98 on this model vs ~0.44 with api="completion".
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
num_examples=200,
num_threads=32,
)
metrics = run_eval(args)
print(f"{metrics=}")
# Chat-API baseline ~0.98; gate well below to absorb sample-noise
# without missing a real PP-routing regression (pre-PP-fix the model
# produced garbage outputs scoring ≈ 0).
self.assertGreaterEqual(metrics["score"], 0.90)
# Wait a little bit so that the memory check happens.
time.sleep(4)
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
def test_mmmu(self):
# Multimodal accuracy gate covering the vision_tower → embed_vision
# (first rank) → PP-proxy handoff → LM tail (last rank) chain.
# Measured 0.71 on 200 examples; full eval (~900 questions) takes
# ~5-7 min on H100 so this is manual-only.
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="mmmu",
num_examples=None,
num_threads=32,
)
metrics = run_eval(args)
print(f"{metrics=}")
# Measured 0.72 on this setup; published Gemma-4-26B MMMU lies in
# 0.69-0.73. Gate 0.65 leaves ~5 SE of headroom (SE on 900 binary
# samples ≈ 0.015) while still catching mid-grade vision/PP
# regressions, not just complete breakage.
self.assertGreater(metrics["score"], 0.65)
@unittest.skipIf(
is_in_amd_ci(),
"Gemma4 PP not yet validated on AMD",
)
class TestGemma4PLEPPAccuracy(unittest.TestCase):
"""PP=2 coverage for Gemma4 PLE variants (per_layer_inputs proxy path).
26B-A4B has ``hidden_size_per_layer_input=0`` so the default Gemma4 PP
test never crosses the PLE branch. Cuda graph + PLE corrupts outputs
(the runner's hardcoded ``{hidden_states, residual}`` PP-proxy schema
drops ``per_layer_inputs``), so this test pins the eager configuration.
"""
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP
cls.base_url = "http://127.0.0.1:23339"
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
1,
"--pp-size",
2,
"--trust-remote-code",
"--enable-multimodal",
# Required for PLE under PP — see Gemma4TextModel guard.
"--disable-cuda-graph",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
# Eager-path baseline ~0.92; gate 0.80 catches PLE breakage
# (corruption collapses score to ~0).
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
num_examples=100,
num_threads=32,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreaterEqual(metrics["score"], 0.80)
time.sleep(4)
if __name__ == "__main__":
unittest.main()
test/registered/pp/test_pp_single_node.py
+1 -135
View File
@@ -23,18 +23,15 @@ from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP,
DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_amd_ci,
is_in_ci,
popen_launch_server,
run_bench_one_batch_server,
)
register_cuda_ci(est_time=500, stage="base-c", runner_config="4-gpu-h100")
register_cuda_ci(est_time=280, stage="base-c", runner_config="4-gpu-h100")
register_amd_ci(est_time=500, suite="stage-c-test-4-gpu-amd")
@@ -145,137 +142,6 @@ class TestDPAttentionDP2PP2(CustomTestCase):
self.assertGreater(metrics["score"], 0.8)
@unittest.skipIf(
is_in_amd_ci(),
"Gemma4 PP not yet validated on AMD",
)
class TestGemma4PPAccuracy(unittest.TestCase):
"""End-to-end PP=2 accuracy gate for Gemma4 multimodal.
Gemma4 has full-attention layers with head_dim=512 (FA's max is 256), so
sglang auto-selects the triton attention backend; no manual flag needed.
The 26B BF16 model splits to ~26 GB per stage under PP=2, well within an
H100's 80 GB.
"""
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP
cls.base_url = "http://127.0.0.1:23333"
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
1,
"--pp-size",
2,
"--trust-remote-code",
"--enable-multimodal",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
# Gemma4 is instruction-tuned and doesn't follow few-shot completion
# prompts well — use the chat API (default in run_eval), which scores
# ~0.98 on this model vs ~0.44 with api="completion".
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
num_examples=200,
num_threads=32,
)
metrics = run_eval(args)
print(f"{metrics=}")
# Chat-API baseline ~0.98; gate well below to absorb sample-noise
# without missing a real PP-routing regression (pre-PP-fix the model
# produced garbage outputs scoring ≈ 0).
self.assertGreaterEqual(metrics["score"], 0.90)
# Wait a little bit so that the memory check happens.
time.sleep(4)
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
def test_mmmu(self):
# Multimodal accuracy gate covering the vision_tower → embed_vision
# (first rank) → PP-proxy handoff → LM tail (last rank) chain.
# Measured 0.71 on 200 examples; full eval (~900 questions) takes
# ~5-7 min on H100 so this is manual-only.
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="mmmu",
num_examples=None,
num_threads=32,
)
metrics = run_eval(args)
print(f"{metrics=}")
# Measured 0.72 on this setup; published Gemma-4-26B MMMU lies in
# 0.69-0.73. Gate 0.65 leaves ~5 SE of headroom (SE on 900 binary
# samples ≈ 0.015) while still catching mid-grade vision/PP
# regressions, not just complete breakage.
self.assertGreater(metrics["score"], 0.65)
@unittest.skipIf(
is_in_amd_ci(),
"Gemma4 PP not yet validated on AMD",
)
class TestGemma4PLEPPAccuracy(unittest.TestCase):
"""PP=2 coverage for Gemma4 PLE variants (per_layer_inputs proxy path).
26B-A4B has ``hidden_size_per_layer_input=0`` so the default Gemma4 PP
test never crosses the PLE branch. Cuda graph + PLE corrupts outputs
(the runner's hardcoded ``{hidden_states, residual}`` PP-proxy schema
drops ``per_layer_inputs``), so this test pins the eager configuration.
"""
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP
cls.base_url = "http://127.0.0.1:23339"
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
1,
"--pp-size",
2,
"--trust-remote-code",
"--enable-multimodal",
# Required for PLE under PP — see Gemma4TextModel guard.
"--disable-cuda-graph",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
# Eager-path baseline ~0.92; gate 0.80 catches PLE breakage
# (corruption collapses score to ~0).
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
num_examples=100,
num_threads=32,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreaterEqual(metrics["score"], 0.80)
time.sleep(4)
class TestPPMixedChunk(CustomTestCase):
@classmethod
def setUpClass(cls):
test/registered/rl/test_release_memory_occupation.py
-477
View File
@@ -1,477 +0,0 @@
"""Test memory release and resume operations for SGLang engine in hybrid RL training.
This test suite evaluates the SGLang engine's memory management capabilities, focusing
on releasing and resuming memory occupation for KV cache and model weights. It simulates
an RL workflow where the SGLang engine acts as a rollout engine for experience collection.
The process involves initializing the engine, sending a small number of requests to simulate
rollout, releasing memory to mimic offloading during RL training, resuming memory occupation,
updating weights with a trained HuggingFace model, and verifying the updated weights.
Detailed in our proposal (https://github.com/sgl-project/sglang/pull/7099), two test cases
are included:
1. Basic Release and Resume: Uses a lower mem_fraction_static (0.6) to control memory allocation
and avoid OOM errors carefully. This test simulates a scenario without multi-stage memory management,
ensuring the engine can release and resume memory occupation while maintaining functionality after
weight updates.
2. Multi-Stage Release and Resume: Employs a higher mem_fraction_static (0.85) to simulate higher
memory pressure, leveraging multi-stage memory management. It sequentially releases and resumes
KV cache and model weights, verifying memory deallocation and reallocation at each stage, and
ensuring correct weight updates and text generation.
3. Tensor Parallel Tests: Tests memory release and resume operations with different tensor parallel
configurations (tp=1, tp=2) to ensure proper memory management in distributed settings. For different
data parallel size, we test it in verl.
NOTE: This test is temporarily disabled.
"""
import os
import time
import unittest
from transformers import AutoModelForCausalLM
import sglang as sgl
from sglang.srt.constants import (
GPU_MEMORY_TYPE_CUDA_GRAPH,
GPU_MEMORY_TYPE_KV_CACHE,
GPU_MEMORY_TYPE_WEIGHTS,
)
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE,
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
CustomTestCase,
empty_gpu_cache,
get_gpu_count,
get_gpu_memory_gb,
)
register_cuda_ci(
est_time=200,
stage="base-c",
runner_config="4-gpu-h100",
disabled="Temporarily disabled - needs investigation",
)
# (temporarily) set to true to observe memory usage in nvidia-smi more clearly
_DEBUG_EXTRA = False
class TestReleaseMemoryOccupation(CustomTestCase):
def _setup_engine(
self,
model_name,
mem_fraction_static=0.8,
tp_size=1,
ep_size=1,
enable_weights_cpu_backup=False,
):
"""Common setup for engine and HF model."""
os.environ["SGLANG_MEMORY_SAVER_CUDA_GRAPH"] = "1"
engine = sgl.Engine(
model_path=model_name,
random_seed=42,
enable_memory_saver=True,
mem_fraction_static=mem_fraction_static,
tp_size=tp_size,
ep_size=ep_size,
enable_weights_cpu_backup=enable_weights_cpu_backup,
# disable_cuda_graph=True, # for debugging only
)
return engine
def _common_test_params(self):
"""Common test parameters."""
return {
"prompt": "Today is a sunny day and I like",
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
"expect_output_before_update_weights": " to spend it outdoors. I decided to",
"expect_output_after_update_weights": " to go for a walk. I like",
"prompt_moe": "The weather is nice today, and I want to",
"sampling_params_moe": {"temperature": 0, "max_new_tokens": 16},
"expect_output_before_update_weights_moe": " go to the park. I have a picnic basket, a book, and a",
"expect_output_after_update_weights_moe": " go to the park. I have a lot of things to do, but I",
"prompt_hybrid_mamba": "The weather is nice today, and I want to",
"sampling_params_hybrid_mamba": {"temperature": 0, "max_new_tokens": 16},
"expect_output_before_update_weights_hybrid_mamba": " go out for a walk. But I don't know what to wear. Can",
"expect_output_after_update_weights_hybrid_mamba": " go out for a walk. But I don't know what to wear. Can",
}
def _test_initial_generation(
self, engine, prompt, sampling_params, expect_output_before_update_weights
):
"""Test initial generation and memory allocation."""
print("generate (#1)")
outputs = engine.generate(prompt, sampling_params)["text"]
self.assertEqual(outputs, expect_output_before_update_weights)
if _DEBUG_EXTRA:
time.sleep(3)
def test_release_and_resume_occupation(self):
# Without multi-stage release and resume, we need to carefully control the memory fraction to avoid OOM
model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
assert get_gpu_count() >= 2, "Need at least 2 GPUs for tensor parallel tests"
for tp_size in [1, 2]:
print(f"Testing tp_size={tp_size} for test_release_and_resume_occupation")
engine = self._setup_engine(
model_name=model_name, mem_fraction_static=0.6, tp_size=tp_size
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt"],
params["sampling_params"],
params["expect_output_before_update_weights"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation()
gpu_memory_usage_after_release = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release,
gpu_memory_usage_before_release,
)
print(
f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
engine.resume_memory_occupation()
print(
f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
)
hf_model_new = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map=get_device(),
)
engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
# destroy the hf model
del hf_model_new
empty_gpu_cache()
print("generate (#2)")
outputs = engine.generate(params["prompt"], params["sampling_params"])[
"text"
]
self.assertEqual(outputs, params["expect_output_after_update_weights"])
engine.shutdown()
def test_release_and_resume_occupation_with_weights_cpu_backup(self):
# Test release and resume occupation with weights CPU backup
model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
print("Testing test_release_and_resume_occupation_with_weights_cpu_backup")
engine = self._setup_engine(
model_name=model_name,
mem_fraction_static=0.6,
enable_weights_cpu_backup=True,
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt"],
params["sampling_params"],
params["expect_output_before_update_weights"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation()
gpu_memory_usage_after_release = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release,
gpu_memory_usage_before_release,
)
print(
f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
engine.resume_memory_occupation()
print(
f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
)
print("generate post resume")
outputs = engine.generate(params["prompt"], params["sampling_params"])["text"]
self.assertEqual(outputs, params["expect_output_before_update_weights"])
engine.shutdown()
def test_multi_stage_release_and_resume(self):
# With multi-stage release and resume, we can set the memory fraction to 0.85 without concern of OOM
model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
for tp_size in [1, 2]:
if tp_size == 2 and get_gpu_count() < 2:
continue
print(f"Testing tp_size={tp_size} for test_multi_stage_release_and_resume")
os.environ["SGLANG_MEMORY_SAVER_CUDA_GRAPH"] = "1"
engine = sgl.Engine(
model_path=model_name,
random_seed=42,
enable_memory_saver=True,
mem_fraction_static=0.85, # Higher memory pressure
tp_size=tp_size,
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt"],
params["sampling_params"],
params["expect_output_before_update_weights"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_KV_CACHE])
gpu_memory_usage_after_release_kv_cache = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release_kv_cache,
gpu_memory_usage_before_release,
)
engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS])
gpu_memory_usage_after_release_weights = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release_weights,
gpu_memory_usage_after_release_kv_cache,
)
engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH])
gpu_memory_usage_after_release_cuda_graph = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release_cuda_graph,
gpu_memory_usage_after_release_weights,
)
print(f"Release took {time.perf_counter() - t:.2f}s")
print(
f"Memory: {gpu_memory_usage_before_release:.1f} → {gpu_memory_usage_after_release_kv_cache:.1f} → {gpu_memory_usage_after_release_weights:.1f} → {gpu_memory_usage_after_release_cuda_graph:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
gpu_memory_usage_before_resume = get_gpu_memory_gb()
# gpu_memory_usage_after_release_weights and gpu_memory_usage_before_resume should be close
self.assertAlmostEqual(
gpu_memory_usage_after_release_weights,
gpu_memory_usage_before_resume,
delta=3.0,
)
print(f"Resume weights took {time.perf_counter() - t:.2f}s")
engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH])
gpu_memory_usage_after_resume_cuda_graph = get_gpu_memory_gb()
self.assertGreater(
gpu_memory_usage_after_resume_cuda_graph,
gpu_memory_usage_before_resume,
)
engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS])
gpu_memory_usage_after_resume_weights = get_gpu_memory_gb()
self.assertGreater(
gpu_memory_usage_after_resume_weights,
gpu_memory_usage_after_resume_cuda_graph,
)
# Update weights from a trained model to serving engine, and then destroy the trained model
hf_model_new = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map=get_device(),
)
gpu_memory_usage_after_loaded_hf_model = get_gpu_memory_gb()
engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
# destroy the hf model
del hf_model_new
empty_gpu_cache()
engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_KV_CACHE])
gpu_memory_usage_after_resume_kv_cache = get_gpu_memory_gb()
self.assertGreater(
gpu_memory_usage_after_resume_kv_cache,
gpu_memory_usage_after_resume_weights,
)
print(f"Resume + update took {time.perf_counter() - t:.2f}s")
print(
f"Memory: {gpu_memory_usage_before_resume:.1f} → {gpu_memory_usage_after_resume_cuda_graph:.1f} → {gpu_memory_usage_after_resume_weights:.1f} → {gpu_memory_usage_after_loaded_hf_model:.1f} → {gpu_memory_usage_after_resume_kv_cache:.1f} GB"
)
print("generate (#2)")
outputs = engine.generate(params["prompt"], params["sampling_params"])[
"text"
]
self.assertEqual(outputs, params["expect_output_after_update_weights"])
engine.shutdown()
def test_moe_model_release_and_resume(self):
# Test with MoE model
model_name = DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT
tp_size = ep_size = 2
print(
f"Testing tp_size={tp_size} and ep_size={ep_size} for test_moe_model_release_and_resume"
)
engine = sgl.Engine(
model_path=model_name,
random_seed=42,
enable_memory_saver=True,
mem_fraction_static=0.5,
tp_size=tp_size,
ep_size=ep_size,
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt_moe"],
params["sampling_params_moe"],
params["expect_output_before_update_weights_moe"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation()
gpu_memory_usage_after_release = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release,
gpu_memory_usage_before_release,
)
print(
f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
engine.resume_memory_occupation()
print(
f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
)
hf_model_new = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map=get_device(),
)
engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
# destroy the hf model
del hf_model_new
empty_gpu_cache()
print("generate (#2)")
outputs = engine.generate(params["prompt_moe"], params["sampling_params_moe"])[
"text"
]
self.assertEqual(outputs, params["expect_output_after_update_weights_moe"])
engine.shutdown()
def test_hybrid_mamba_model_release_and_resume(self):
# Test with Hybrid Mamba model
model_name = DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST
tp_size = 4
print(
f"Testing tp_size={tp_size} for test_hybrid_mamba_model_release_and_resume"
)
engine = sgl.Engine(
model_path=model_name,
random_seed=42,
enable_memory_saver=True,
tp_size=tp_size,
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt_hybrid_mamba"],
params["sampling_params_hybrid_mamba"],
params["expect_output_before_update_weights_hybrid_mamba"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation()
gpu_memory_usage_after_release = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release,
gpu_memory_usage_before_release,
)
print(
f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
engine.resume_memory_occupation()
print(
f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
)
engine.update_weights_from_disk(model_name)
# destroy the hf model
empty_gpu_cache()
print("generate (#2)")
outputs = engine.generate(
params["prompt_hybrid_mamba"], params["sampling_params_hybrid_mamba"]
)["text"]
self.assertEqual(
outputs, params["expect_output_after_update_weights_hybrid_mamba"]
)
engine.shutdown()
if __name__ == "__main__":
unittest.main()
test/registered/unit/layers/test_flashinfer_comm_fusion.py
+3 -3
View File
@@ -10,9 +10,9 @@ from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-h100")
register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-b200")
register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-gb300")
# Collectives are mocked and world_size is a plain int, so the world_size=4
# cases need one real CUDA device.
register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
class _FakeWorkspace: