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

[Mamba] extra buffer lazy support (#27118)

Browse Source 
Co-authored-by: YAMY <74099316+YAMY1234@users.noreply.github.com>
This commit is contained in:
Hanming Lu
2026-06-03 12:42:11 -07:00
committed by GitHub
co-authored by YAMY
parent ac99794e64
commit e0b692600f
14 changed files with 429 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/4-gpu-models/test_qwen3_next_models.py
+33 -2
View File
@@ -19,11 +19,42 @@ class TestQwen3Next(
"--tp-size",
"4",
"--chunked-prefill-size",
"2048",
"1024",
"--mamba-scheduler-strategy",
"extra_buffer",
"--mamba-track-interval",
"128",
"2",
"--page-size",
"1",
"--attention-backend",
"triton",
"--moe-runner-backend",
"triton",
]
class TestQwen3NextLazyExtraBuffer(
GSM8KMixin, KLDivergenceMixin, PrefixCacheBranchingMixin, DefaultServerBase
):
model = QWEN3_NEXT_MODEL
cache_chunk_size = 64
gsm8k_accuracy_thres = 0.93
kl_div_thres = 0.0025
other_args = [
"--tp-size",
"4",
"--chunked-prefill-size",
"1024",
"--mamba-scheduler-strategy",
"extra_buffer_lazy",
"--mamba-track-interval",
"2",
"--page-size",
"1",
"--attention-backend",
"triton",
"--moe-runner-backend",
"triton",
]
test/registered/models_e2e/test_qwen3_next_models.py
+104
View File
@@ -0,0 +1,104 @@
import os
import unittest
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
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=250, stage="base-c", 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 TestQwen3NextLazyExtraBuffer(
GSM8KMixin, KLDivergenceMixin, PrefixCacheBranchingMixin, DefaultServerBase
):
model = QWEN3_NEXT_MODEL
cache_chunk_size = 64
gsm8k_accuracy_thres = 0.93
kl_div_thres = 0.001
other_args = _make_args(page_size=1, track_interval=2)
class TestQwen3NextLazyExtraBufferLargePage(
GSM8KMixin, KLDivergenceMixin, PrefixCacheBranchingMixin, DefaultServerBase
):
model = QWEN3_NEXT_MODEL
cache_chunk_size = 64
gsm8k_accuracy_thres = 0.93
kl_div_thres = 0.001
other_args = _make_args(page_size=2, track_interval=2)
@unittest.skip("Manual-only: forces all lazy preallocs to fail")
class TestQwen3NextLazyExtraBufferAllocFail(
GSM8KMixin, KLDivergenceMixin, PrefixCacheBranchingMixin, DefaultServerBase
):
model = QWEN3_NEXT_MODEL
cache_chunk_size = 64
gsm8k_accuracy_thres = 0.93
kl_div_thres = 0.001
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)
@unittest.skip("Manual-only: forces all lazy preallocs to fail")
class TestQwen3NextLazyExtraBufferLargePageAllocFail(
GSM8KMixin, KLDivergenceMixin, PrefixCacheBranchingMixin, DefaultServerBase
):
model = QWEN3_NEXT_MODEL
cache_chunk_size = 64
gsm8k_accuracy_thres = 0.93
kl_div_thres = 0.001
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()