Add mixed-prefix gsm8k eval and its CPU unit test (#27502)

Author: fzyzcjy

python/sglang/test/run_eval.py

@@ -178,6 +178,17 @@ def run_eval(args):
             num_shots=getattr(args, "num_shots", 5),
             data_path=getattr(args, "gsm8k_data_path", None),
         )
+    elif args.eval_name == "mixed_prefix_gsm8k":
+        from sglang.test.simple_eval_mixed_prefix_gsm8k import MixedPrefixGSM8KEval
+
+        eval_obj = MixedPrefixGSM8KEval(
+            num_examples=args.num_examples,
+            num_threads=args.num_threads,
+            num_shots=args.num_shots,
+            secondary_pool_size=args.mixed_prefix_gsm8k_secondary_pool_size,
+            data_path=args.gsm8k_data_path,
+            seed=args.mixed_prefix_gsm8k_seed,
+        )
     else:
         raise ValueError(f"Invalid eval name: {args.eval_name}")
 
@@ -367,6 +378,18 @@ def run_eval(args):
         default=None,
         help="Path to GSM8K data file (e.g., test.jsonl)",
     )
+    parser.add_argument(
+        "--mixed-prefix-gsm8k-secondary-pool-size",
+        type=int,
+        default=15,
+        help="Size of secondary example pool for eval_name=mixed_prefix_gsm8k (default: 15)",
+    )
+    parser.add_argument(
+        "--mixed-prefix-gsm8k-seed",
+        type=int,
+        default=42,
+        help="Seed for per-question random sampling in mixed_prefix_gsm8k (default: 42)",
+    )
 
     args = parser.parse_args()
 

python/sglang/test/simple_eval_gsm8k.py

@@ -56,20 +56,29 @@ def __init__(
         else:
             filename = download_and_cache_file(GSM8K_URL)
 
-        self._lines = list(read_jsonl(filename))
-        self._few_shot_prompt = get_few_shot_examples(self._lines, num_shots)
-
+        all_lines = list(read_jsonl(filename))
+        pool_size = self._setup_prefix_pool(all_lines, num_shots)
         # The evaluation data should not include the few-shot examples to prevent data leakage.
-        self._lines = self._lines[num_shots:]
+        self._lines = all_lines[pool_size:]
         if num_examples is not None:
+            # Slice caps silently when num_examples exceeds the available lines,
+            # matching upstream: callers like test_basic_sanity_eagle3 pass a
+            # num_examples larger than the dataset on purpose.
             self._lines = self._lines[:num_examples]
 
+    def _setup_prefix_pool(self, all_lines: list, num_shots: int) -> int:
+        self._few_shot_prompt = get_few_shot_examples(all_lines, num_shots)
+        return num_shots
+
+    def _build_prefix(self, idx: int) -> str:
+        return self._few_shot_prompt
+
     def __call__(self, sampler: SamplerBase) -> EvalResult:
         def fn(idx: int) -> SingleEvalResult:
             question = get_one_example(self._lines, idx, include_answer=False)
             correct_answer = get_answer_value(self._lines[idx]["answer"])
 
-            prompt_content = self._few_shot_prompt + question
+            prompt_content = self._build_prefix(idx) + question
             prompt_messages = [
                 sampler._pack_message(content=prompt_content, role="user")
             ]

python/sglang/test/simple_eval_mixed_prefix_gsm8k.py

@@ -0,0 +1,51 @@
+import random
+from typing import Optional
+
+from sglang.test.simple_eval_gsm8k import GSM8KEval, get_one_example
+
+
+class MixedPrefixGSM8KEval(GSM8KEval):
+
+    def __init__(
+        self,
+        num_examples: Optional[int],
+        num_threads: int,
+        num_shots: int,
+        secondary_pool_size: int,
+        data_path: Optional[str],
+        seed: int,
+    ):
+        self._secondary_pool_size = secondary_pool_size
+        self._seed = seed
+        super().__init__(
+            num_examples=num_examples,
+            num_threads=num_threads,
+            num_shots=num_shots,
+            data_path=data_path,
+        )
+
+    def _setup_prefix_pool(self, all_lines: list, num_shots: int) -> int:
+        overall_pool_size = num_shots + self._secondary_pool_size
+        if len(all_lines) < overall_pool_size + 1:
+            raise ValueError(
+                f"GSM8K dataset has {len(all_lines)} examples but mixed-prefix "
+                f"eval needs at least {overall_pool_size + 1} "
+                f"(num_shots {num_shots} + secondary "
+                f"{self._secondary_pool_size} + 1 test)."
+            )
+        self._primary_shots = all_lines[:num_shots]
+        self._secondary_pool = all_lines[num_shots:overall_pool_size]
+        return overall_pool_size
+
+    def _build_prefix(self, idx: int) -> str:
+        rng = random.Random(self._seed + idx)
+        num_primary = rng.randint(0, self._num_shots)
+        secondary_size = rng.randint(0, self._secondary_pool_size)
+        secondary_indices = rng.sample(range(len(self._secondary_pool)), secondary_size)
+        primary = self._primary_shots[:num_primary]
+        secondary = [self._secondary_pool[i] for i in secondary_indices]
+        combined = primary + secondary
+        return "".join(
+            get_one_example(combined, i, include_answer=True) + "\n\n"
+            for i in range(len(combined))
+        )

test/registered/unit/bench/test_mixed_prefix_gsm8k.py

@@ -0,0 +1,177 @@
+import json
+import os
+import tempfile
+import unittest
+from typing import List, Tuple
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.simple_eval_gsm8k import get_one_example
+from sglang.test.simple_eval_mixed_prefix_gsm8k import MixedPrefixGSM8KEval
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-b-test-cpu")
+
+
+def _write_synthetic_dataset(path: str, n: int) -> None:
+    with open(path, "w") as f:
+        for i in range(n):
+            f.write(
+                json.dumps(
+                    {
+                        "question": f"Synthetic question {i}: what is {i} + {i}?",
+                        "answer": f"The answer is {2 * i}. #### {2 * i}",
+                    }
+                )
+                + "\n"
+            )
+
+
+class TestMixedPrefixGSM8KEval(CustomTestCase):
+    NUM_SHOTS = 4
+    SECONDARY_POOL_SIZE = 12
+    NUM_EXAMPLES = 40
+
+    @classmethod
+    def setUpClass(cls):
+        cls._tmpdir = tempfile.TemporaryDirectory()
+        cls._data_path = os.path.join(cls._tmpdir.name, "synthetic.jsonl")
+        _write_synthetic_dataset(cls._data_path, 100)
+
+    @classmethod
+    def tearDownClass(cls):
+        cls._tmpdir.cleanup()
+
+    def _make_eval(self, seed: int = 42, num_examples=None):
+        return MixedPrefixGSM8KEval(
+            num_examples=(
+                num_examples if num_examples is not None else self.NUM_EXAMPLES
+            ),
+            num_threads=1,
+            num_shots=self.NUM_SHOTS,
+            secondary_pool_size=self.SECONDARY_POOL_SIZE,
+            data_path=self._data_path,
+            seed=seed,
+        )
+
+    def _primary_lines(self, evaluator) -> List[str]:
+        return [
+            get_one_example(evaluator._primary_shots, j, include_answer=True) + "\n\n"
+            for j in range(self.NUM_SHOTS)
+        ]
+
+    def _decompose(self, evaluator, prefix: str) -> Tuple[int, List[str]]:
+        k = 0
+        for line in self._primary_lines(evaluator):
+            if prefix.startswith(line):
+                prefix = prefix[len(line) :]
+                k += 1
+            else:
+                break
+        remainder_questions: List[str] = []
+        if prefix:
+            chunks = prefix.split("\n\n")
+            for chunk in chunks:
+                if chunk.startswith("Question: "):
+                    q_text = chunk[len("Question: ") :].split("\nAnswer:")[0]
+                    remainder_questions.append(q_text)
+        return k, remainder_questions
+
+    def test_primary_segment_is_strict_prefix_of_primary_shots(self):
+        e = self._make_eval()
+        for i in range(self.NUM_EXAMPLES):
+            k, _ = self._decompose(e, e._build_prefix(i))
+            self.assertGreaterEqual(k, 0)
+            self.assertLessEqual(k, self.NUM_SHOTS)
+
+    def test_remainder_questions_come_from_secondary_pool(self):
+        e = self._make_eval()
+        secondary_qs = {item["question"] for item in e._secondary_pool}
+        for i in range(self.NUM_EXAMPLES):
+            _, remainder = self._decompose(e, e._build_prefix(i))
+            for q in remainder:
+                self.assertIn(q, secondary_qs)
+
+    def test_remainder_no_duplicates_within_one_query(self):
+        e = self._make_eval()
+        for i in range(self.NUM_EXAMPLES):
+            _, remainder = self._decompose(e, e._build_prefix(i))
+            self.assertEqual(
+                len(remainder),
+                len(set(remainder)),
+                f"query {i} has duplicate secondary samples",
+            )
+
+    def test_remainder_size_within_secondary_pool_bound(self):
+        e = self._make_eval()
+        for i in range(self.NUM_EXAMPLES):
+            _, remainder = self._decompose(e, e._build_prefix(i))
+            self.assertGreaterEqual(len(remainder), 0)
+            self.assertLessEqual(len(remainder), self.SECONDARY_POOL_SIZE)
+
+    def test_primary_depth_takes_multiple_values(self):
+        e = self._make_eval()
+        ks = {
+            self._decompose(e, e._build_prefix(i))[0] for i in range(self.NUM_EXAMPLES)
+        }
+        self.assertGreater(len(ks), 2, f"k values seen: {ks}")
+
+    def test_secondary_size_takes_multiple_values(self):
+        e = self._make_eval()
+        sizes = {
+            len(self._decompose(e, e._build_prefix(i))[1])
+            for i in range(self.NUM_EXAMPLES)
+        }
+        self.assertGreater(len(sizes), 2, f"sizes seen: {sizes}")
+
+    def test_two_queries_share_min_primary_prefix(self):
+        e = self._make_eval()
+        lines = self._primary_lines(e)
+        prefixes = [e._build_prefix(i) for i in range(self.NUM_EXAMPLES)]
+        ks = [self._decompose(e, p)[0] for p in prefixes]
+        for i in range(self.NUM_EXAMPLES):
+            for j in range(i + 1, self.NUM_EXAMPLES):
+                shared = "".join(lines[: min(ks[i], ks[j])])
+                self.assertTrue(prefixes[i].startswith(shared))
+                self.assertTrue(prefixes[j].startswith(shared))
+
+    def test_build_prefix_is_deterministic(self):
+        a = self._make_eval(seed=42)
+        b = self._make_eval(seed=42)
+        for i in range(self.NUM_EXAMPLES):
+            self.assertEqual(a._build_prefix(i), b._build_prefix(i))
+
+    def test_seed_actually_matters(self):
+        a = self._make_eval(seed=42)
+        b = self._make_eval(seed=43)
+        differences = sum(
+            1
+            for i in range(self.NUM_EXAMPLES)
+            if a._build_prefix(i) != b._build_prefix(i)
+        )
+        self.assertGreater(differences, self.NUM_EXAMPLES // 2)
+
+    def test_pools_and_test_lines_pairwise_disjoint(self):
+        e = self._make_eval(num_examples=None)
+        primary_qs = {item["question"] for item in e._primary_shots}
+        secondary_qs = {item["question"] for item in e._secondary_pool}
+        test_qs = {item["question"] for item in e._lines}
+        self.assertEqual(primary_qs & secondary_qs, set())
+        self.assertEqual(primary_qs & test_qs, set())
+        self.assertEqual(secondary_qs & test_qs, set())
+
+    def test_insufficient_dataset_raises(self):
+        tiny = os.path.join(self._tmpdir.name, "tiny.jsonl")
+        _write_synthetic_dataset(tiny, n=5)
+        with self.assertRaises(ValueError):
+            MixedPrefixGSM8KEval(
+                num_examples=1,
+                num_threads=1,
+                num_shots=self.NUM_SHOTS,
+                secondary_pool_size=self.SECONDARY_POOL_SIZE,
+                data_path=tiny,
+                seed=42,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()