diff --git a/test/registered/attention/test_mamba_unittest.py b/test/registered/attention/test_mamba_unittest.py index 8c6869e37..0a0f45fb8 100644 --- a/test/registered/attention/test_mamba_unittest.py +++ b/test/registered/attention/test_mamba_unittest.py @@ -11,6 +11,7 @@ from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=9, suite="stage-b-test-small-1-gpu") @@ -130,6 +131,9 @@ class TestMamba(unittest.TestCase): assert req_to_token_pool.mamba_pool.available_size() == mamba_cache_size - 1 def test_mamba_radix_cache_1(self): + set_global_server_args_for_scheduler( + ServerArgs(model_path="dummy", page_size=1) + ) # kv cache size = 128 dtype = torch.bfloat16