test_chunked_prefill.py 1.96 KB
Newer Older
1
2
3
4
5
import unittest
from types import SimpleNamespace

from sglang.srt.utils import kill_child_process
from sglang.test.run_eval import run_eval
6
7
from sglang.test.test_utils import (
    DEFAULT_MODEL_NAME_FOR_TEST,
8
9
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
    DEFAULT_URL_FOR_TEST,
10
11
    popen_launch_server,
)
12
13


14
class TestChunkedPrefill(unittest.TestCase):
15
16
17
18
    def run_mmlu(
        self, disable_radix_cache, enable_mixed_chunk, chunked_prefill_size=32
    ):
        other_args = ["--chunked-prefill-size", str(chunked_prefill_size)]
19
20
21
        if disable_radix_cache:
            other_args += ["--disable-radix-cache"]

22
23
24
        if enable_mixed_chunk:
            other_args += ["--enable-mixed-chunk"]

25
        model = DEFAULT_MODEL_NAME_FOR_TEST
26
        base_url = DEFAULT_URL_FOR_TEST
27
28
29
        process = popen_launch_server(
            model,
            base_url,
30
            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
31
            other_args=other_args,
32
33
34
        )

        args = SimpleNamespace(
35
36
            base_url=base_url,
            model=model,
37
            eval_name="mmlu",
38
            num_examples=64,
39
            num_threads=32,
40
41
        )

42
43
        try:
            metrics = run_eval(args)
44
            assert metrics["score"] >= 0.65
45
46
47
48
        finally:
            kill_child_process(process.pid)

    def test_chunked_prefill(self):
49
50
51
52
        self.run_mmlu(disable_radix_cache=False, enable_mixed_chunk=False)

    def test_mixed_chunked_prefill(self):
        self.run_mmlu(disable_radix_cache=False, enable_mixed_chunk=True)
53
54

    def test_chunked_prefill_without_radix_cache(self):
55
56
57
58
        self.run_mmlu(disable_radix_cache=True, enable_mixed_chunk=False)

    def test_mixed_chunked_prefill_without_radix_cache(self):
        self.run_mmlu(disable_radix_cache=True, enable_mixed_chunk=True)
59

60
61
62
63
64
    def test_no_chunked_prefill(self):
        self.run_mmlu(
            disable_radix_cache=False, enable_mixed_chunk=False, chunked_prefill_size=-1
        )

65
66

if __name__ == "__main__":
Lianmin Zheng's avatar
Lianmin Zheng committed
67
    unittest.main()