test: fix Task 1 CPU unit coverage

This commit is contained in:
2026-05-08 00:29:33 +08:00
parent 95bacb8862
commit 36ee803aa7
2 changed files with 14 additions and 3 deletions

View File

@@ -1,7 +1,12 @@
import sys
import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock
for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"):
if _mod not in sys.modules:
sys.modules[_mod] = MagicMock()
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder
from sglang.srt.mem_cache.base_prefix_cache import (
@@ -453,7 +458,6 @@ class TestPrefillAdder(CustomTestCase):
running_batch,
page_size=64,
rem_input_tokens=4096,
rem_total_tokens=4096,
)
req = self.create_mock_req("req", priority=0, max_new_tokens=16)
req.extend_input_len = 256
@@ -474,7 +478,7 @@ class TestPrefillAdder(CustomTestCase):
self.assertNotEqual(result, AddReqResult.NO_TOKEN)
params = self.mock_tree_cache.init_load_back.call_args.args[0]
self.assertEqual(params.mem_quota, 256)
self.assertEqual(params.mem_quota, 320)
if __name__ == "__main__":

View File

@@ -17,6 +17,13 @@ Usage:
python -m pytest test_radix_cache_unit.py::TestRadixCache::test_insert_basic
"""
import sys
import unittest.mock
for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"):
if _mod not in sys.modules:
sys.modules[_mod] = unittest.mock.MagicMock()
from sglang.srt.mem_cache.common import available_and_evictable_str
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
@@ -27,7 +34,6 @@ register_amd_ci(est_time=5, suite="stage-b-test-1-gpu-small-amd")
import random
import time
import unittest
import unittest.mock
import torch
@@ -743,6 +749,7 @@ class TestRadixCache(unittest.TestCase):
# Should have 1 page (split at page_size=2)
self.assertEqual(len(node.hash_value), 1)
@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
def test_memory_allocated(self):
keys, values = [], []