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

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 @@ if __name__ == "__main__":
         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 @@ class GSM8KEval(Eval):
         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))
+        )