Add long context dataset benchmark support (#227)

Author: Lumosis

benchmarks/benchmark_serving.py

@@ -279,6 +279,21 @@ def load_openorca_dataset_pkl(
   return [(prompt, output) for prompt, output in zip(prompts, outputs)]
 
 
+def load_longcontext_dataset_pkl(
+    dataset_path: str,
+) -> list[tuple[Any, Any]]:
+  assert os.path.isfile(dataset_path)
+
+  # read pickle file
+  data = pandas.read_pickle(dataset_path)
+
+  samples = []
+  for _, row in data.iterrows():
+    samples.append((row["input"], row["ref_output"]))
+
+  return samples
+
+
 def load_mmlu_dataset_csv(dataset_path: str) -> tuple[Any, dict[str, str]]:
   assert dataset_path != ""
   dataset = []
@@ -837,7 +852,14 @@ def parse_args() -> argparse.Namespace:
       "--dataset",
       type=str,
       default="test",
-      choices=["test", "sharegpt", "openorca", "mmlu", "math500"],
+      choices=[
+          "test",
+          "sharegpt",
+          "openorca",
+          "mmlu",
+          "math500",
+          "longcontext",
+      ],
       help="The dataset name.",
   )
   parser.add_argument("--dataset-path", type=str, help="Path to the dataset.")
@@ -1057,6 +1079,10 @@ def main(args: argparse.Namespace):
       dataset = load_math500_dataset(
           args.dataset_path,
       )
+    elif args.dataset == "longcontext":
+      dataset = load_longcontext_dataset_pkl(
+          args.dataset_path,
+      )
     else:
       raise ValueError(
           f"Fatal Error: Uncognized input parameters: {args.dataset}"

jetstream/core/orchestrator.py

@@ -628,7 +628,7 @@ def _prefill_thread(self, idx: int):
                   is_bos,
                   prefill_engine.max_prefill_length,
                   prefill_engine.use_chunked_prefill,
-                  prefill_engine.chunk_size,
+                  prefill_engine.prefill_chunk_size,
               )
           )
           prefill_result = None
@@ -649,7 +649,7 @@ def _prefill_thread(self, idx: int):
             t_l_array = jnp.expand_dims(
                 jnp.arange(
                     0,
-                    chunk_num * prefill_engine.chunk_size
+                    chunk_num * prefill_engine.prefill_chunk_size
                     + true_lengths_of_chunks[chunk_num],
                 ),
                 1,

jetstream/engine/engine_api.py

@@ -445,11 +445,6 @@ def colocated_cpus(self) -> Union[list[CpuDevices], None]:
   def use_chunked_prefill(self) -> bool:
     return self._downstream_engine.use_chunked_prefill
 
-  @property
-  def chunk_size(self) -> bool:
-    """Maximum prefill length."""
-    return self._downstream_engine.prefill_chunk_size
-
   @property
   def prefill_chunk_size(self) -> int:
     """Maximum prefill length."""

jetstream/engine/mock_engine.py

@@ -500,11 +500,6 @@ def use_chunked_prefill(self) -> bool:
     """Maximum prefill length."""
     return self._use_chunked_prefill
 
-  @property
-  def chunk_size(self) -> bool:
-    """Maximum prefill length."""
-    return 2
-
   @property
   def prefill_chunk_size(self) -> int:
     """Maximum prefill length."""

jetstream/tests/core/test_orchestrator.py

@@ -108,9 +108,9 @@ async def test_orchestrator_chunked_prefill(self, interleaved_mode: bool):
         max_tokens=3,
     )
     iterator = client.Decode(request)
-    # chr of [266, 332, 415].
-    expected_text = ["B", "R", "g", ""]
-    expected_token_ids = [66, 82, 103, None]
+    # chr of [135, 168, 210].
+    expected_text = ["\x87", "¨", "Ò", ""]
+    expected_token_ids = [135, 168, 210, None]
     counter = 0
     async for resp in iterator:
       output_text = resp.stream_content.samples[0].text