[bench] Add agentic-trace multi-turn dataset to bench_serving (#29215)

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

docs_new/docs/developer_guide/bench_serving.mdx

@@ -66,6 +66,7 @@ Select with `--dataset-name`:
 - `generated-shared-prefix`: synthetic dataset with shared long system prompts and short questions
 - `mmmu`: samples from MMMU (Math split) and includes images
 - `speed-bench`: [SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) (**SPEculative Evaluation Dataset**) — a unified benchmark for evaluating [Speculative Decoding (SD)](https://arxiv.org/abs/2604.09557) algorithms. Uses the Throughput split, which provides fixed-length input sequences (1K–32K tokens) grouped into three output-entropy categories (`low_entropy`, `mixed`, `high_entropy`). Requires a pre-downloaded JSONL file passed via `--dataset-path`.
+- `agentic-trace`: replays pre-built multi-turn agentic traces (e.g. OpenHands / SWE-smith). Each conversation is replayed round by round, feeding the server's real assistant reply back into the next round's history. Requires a chat backend (`--backend sglang-oai-chat`) and a trace JSON passed via `--dataset-path`.
 
 Common dataset flags:
 
@@ -93,6 +94,13 @@ Image dataset flags (for `image`):
 - `--image-format`: Image format (jpeg or png)
 - `--image-content`: Image content type (random or blank)
 
+Agentic trace flags (for `agentic-trace`):
+
+- `--dataset-path`: path to the pre-built trace JSON
+- `--sharegpt-output-len`: per-turn output length (default: 220)
+- `--dataset-offset`: rotate the conversation list by this many entries before sampling, so successive sweep steps start on fresh conversations
+- `--agentic-max-turns`: cap each conversation to at most this many turns (useful for small, fast profiling runs)
+
 SPEED-Bench flags (for `speed-bench`):
 
 - `--dataset-path PATH`: path to the pre-downloaded SPEED-Bench Throughput JSONL (e.g., `throughput_1k.jsonl`). Use the [SPEED-Bench measurement framework](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/specdec_bench) to generate it.

python/sglang/benchmark/datasets/__init__.py

@@ -1,5 +1,6 @@
 from typing import Dict, Type
 
+from sglang.benchmark.datasets.agentic_trace import AgenticTraceDataset
 from sglang.benchmark.datasets.autobench import AutoBenchmarkDataset
 from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
 from sglang.benchmark.datasets.custom import CustomDataset
@@ -16,6 +17,7 @@ from sglang.benchmark.datasets.sharegpt import ShareGPTDataset
 from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset
 
 DATASET_MAPPING: Dict[str, Type[BaseDataset]] = {
+    "agentic-trace": AgenticTraceDataset,
     "autobench": AutoBenchmarkDataset,
     "sharegpt": ShareGPTDataset,
     "custom": CustomDataset,

python/sglang/benchmark/datasets/agentic_trace.py

@@ -0,0 +1,114 @@
+import json
+import os
+from argparse import Namespace
+from dataclasses import dataclass
+from typing import List, Optional
+
+import numpy as np
+from transformers import PreTrainedTokenizerBase
+
+from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
+
+# Per-turn output length when --sharegpt-output-len is not given; matches the
+# ~220-token average assistant reply of OpenHands-style agentic traces.
+DEFAULT_AGENTIC_OUTPUT_LEN = 220
+
+
+@dataclass
+class AgenticTraceDataset(BaseDataset):
+    """Multi-turn agentic trace loader (e.g. OpenHands / SWE-smith traces).
+
+    Expects a trace JSON of the shape::
+
+        {
+          "metadata": {...},
+          "conversations": [
+            [   # one conversation == a list of turns
+              {"messages": [{"role": "system", ...}, {"role": "user", ...}],
+               "prompt_tokens": 73821},
+              {"messages": [{"role": "user", ...}], "prompt_tokens": 74894},
+              ...
+            ],
+            ...
+          ]
+        }
+
+    Each turn's ``messages`` holds only the new non-assistant messages for that
+    turn. One conversation becomes one :class:`DatasetRow` whose ``prompt`` is
+    the list of per-turn message deltas; ``bench_serving`` detects this shape as
+    multi-turn and replays each conversation round by round, feeding the
+    server's real assistant reply back into the next round's history.
+
+    Use with a chat backend (``--backend sglang-oai-chat``).
+    """
+
+    dataset_path: str
+    num_requests: int
+    fixed_output_len: Optional[int]
+    offset: int
+    max_turns: Optional[int]
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "AgenticTraceDataset":
+        return cls(
+            dataset_path=args.dataset_path,
+            num_requests=args.num_prompts,
+            fixed_output_len=args.sharegpt_output_len,
+            offset=args.dataset_offset,
+            max_turns=args.agentic_max_turns,
+        )
+
+    def load(
+        self, tokenizer: PreTrainedTokenizerBase, model_id=None
+    ) -> List[DatasetRow]:
+        if not os.path.isfile(self.dataset_path):
+            raise FileNotFoundError(f"Dataset not found at {self.dataset_path}")
+
+        with open(self.dataset_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        conversations = data.get("conversations", [])
+        if not conversations:
+            raise ValueError(f"No 'conversations' found in {self.dataset_path}.")
+
+        offset = self.offset % len(conversations)
+        if offset:
+            conversations = conversations[offset:] + conversations[:offset]
+
+        output_len = self.fixed_output_len or DEFAULT_AGENTIC_OUTPUT_LEN
+
+        filtered_dataset: List[DatasetRow] = []
+        for conversation in conversations:
+            if self.num_requests > 0 and len(filtered_dataset) >= self.num_requests:
+                break
+
+            prompt = [turn["messages"] for turn in conversation if turn.get("messages")]
+            if self.max_turns:
+                prompt = prompt[: self.max_turns]
+            if not prompt:
+                continue
+
+            # Informational only: multi-turn replay ignores per-row prompt_len.
+            prompt_len = int(conversation[0].get("prompt_tokens", 0))
+
+            filtered_dataset.append(
+                DatasetRow(
+                    prompt=prompt,
+                    prompt_len=prompt_len,
+                    output_len=output_len,
+                )
+            )
+
+        if not filtered_dataset:
+            raise ValueError(
+                f"No usable conversations loaded from {self.dataset_path}."
+            )
+
+        num_turns = [len(row.prompt) for row in filtered_dataset]
+        print(
+            f"#Conversations: {len(filtered_dataset)} "
+            f"(offset={offset}, turns/conv min={min(num_turns)} "
+            f"max={max(num_turns)} avg={np.mean(num_turns):.1f})"
+        )
+        print(f"#Output tokens per turn: {output_len}")
+        return filtered_dataset

python/sglang/benchmark/serving.py

@@ -2158,6 +2158,7 @@ def cli_main():
         type=str,
         default="sharegpt",
         choices=[
+            "agentic-trace",
             "autobench",
             "sharegpt",
             "custom",
@@ -2176,6 +2177,21 @@ def cli_main():
     parser.add_argument(
         "--dataset-path", type=str, default="", help="Path to the dataset."
     )
+    parser.add_argument(
+        "--dataset-offset",
+        type=int,
+        default=0,
+        help="Rotate the conversation list by this many entries before sampling "
+        "(agentic-trace dataset), so successive sweep steps start on fresh "
+        "conversations.",
+    )
+    parser.add_argument(
+        "--agentic-max-turns",
+        type=int,
+        default=None,
+        help="Cap each conversation to at most this many turns (agentic-trace "
+        "dataset). Default: use all turns in the trace.",
+    )
     parser.add_argument(
         "--speed-bench-category",
         type=str,

test/registered/bench_fn/test_benchmark_datasets_api.py

@@ -19,6 +19,10 @@ from tokenizers.pre_tokenizers import Whitespace
 from transformers import PreTrainedTokenizerFast
 
 from sglang.benchmark.datasets import DATASET_MAPPING, get_dataset
+from sglang.benchmark.datasets.agentic_trace import (
+    DEFAULT_AGENTIC_OUTPUT_LEN,
+    AgenticTraceDataset,
+)
 from sglang.benchmark.datasets.common import DatasetRow, gen_mm_prompt
 from sglang.benchmark.datasets.custom import sample_custom_requests
 from sglang.benchmark.datasets.generated_shared_prefix import (
@@ -148,6 +152,8 @@ def make_args(**overrides):
         "mooncake_workload": "conversation",
         "speed_bench_category": None,
         "speed_bench_output_len": 512,
+        "dataset_offset": 0,
+        "agentic_max_turns": None,
     }
     args.update(overrides)
     return SimpleNamespace(**args)
@@ -284,6 +290,37 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase):
                 f.write(json.dumps(row) + "\n")
         return str(path)
 
+    def _write_agentic_trace_json(self):
+        trace = {
+            "metadata": {"source": "test"},
+            "conversations": [
+                [
+                    {
+                        "messages": [
+                            {"role": "system", "content": "You are an agent."},
+                            {"role": "user", "content": "Fix the bug."},
+                        ],
+                        "prompt_tokens": 100,
+                    },
+                    {
+                        "messages": [{"role": "user", "content": "Tool output: ok."}],
+                        "prompt_tokens": 200,
+                    },
+                    {"messages": []},
+                ],
+                [
+                    {
+                        "messages": [{"role": "user", "content": "Run the tests."}],
+                        "prompt_tokens": 50,
+                    },
+                ],
+            ],
+        }
+        path = self.tmpdir_path / "agentic_trace.json"
+        with open(path, "w") as f:
+            json.dump(trace, f)
+        return str(path)
+
     async def _collect_mooncake_rows(self, records):
         out = []
         async for row in get_mooncake_request_over_time(
@@ -517,8 +554,69 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase):
         with self.assertRaises(ValueError):
             SpeedBenchDataset.from_args(args)
 
+    def test_agentic_trace_sampler(self):
+        dataset_path = self._write_agentic_trace_json()
+        args = make_args(
+            dataset_name="agentic-trace",
+            dataset_path=dataset_path,
+            num_prompts=10,
+        )
+        dataset = AgenticTraceDataset.from_args(args)
+        rows = dataset.load(self.tokenizer)
+        self.assertEqual(len(rows), 2)
+        self.assertTrue(all(isinstance(row, DatasetRow) for row in rows))
+        self.assertTrue(
+            all(row.output_len == DEFAULT_AGENTIC_OUTPUT_LEN for row in rows)
+        )
+        # Multi-turn shape: prompt is a list of per-turn message lists, with
+        # the empty third turn of the first conversation dropped.
+        self.assertEqual(len(rows[0].prompt), 2)
+        self.assertEqual(len(rows[1].prompt), 1)
+        self.assertEqual(rows[0].prompt[0][0]["role"], "system")
+        self.assertEqual(rows[0].prompt_len, 100)
+        self.assertEqual(rows[1].prompt_len, 50)
+
+    def test_agentic_trace_offset_and_max_turns(self):
+        dataset_path = self._write_agentic_trace_json()
+        args = make_args(
+            dataset_name="agentic-trace",
+            dataset_path=dataset_path,
+            num_prompts=10,
+            sharegpt_output_len=64,
+            dataset_offset=1,
+            agentic_max_turns=1,
+        )
+        dataset = AgenticTraceDataset.from_args(args)
+        rows = dataset.load(self.tokenizer)
+        self.assertEqual(len(rows), 2)
+        # offset=1 rotates the second (single-turn) conversation to the front.
+        self.assertEqual(rows[0].prompt_len, 50)
+        self.assertTrue(all(len(row.prompt) == 1 for row in rows))
+        self.assertTrue(all(row.output_len == 64 for row in rows))
+
+    def test_agentic_trace_invalid_input_raises(self):
+        args = make_args(
+            dataset_name="agentic-trace",
+            dataset_path=str(self.tmpdir_path / "missing.json"),
+            num_prompts=1,
+        )
+        with self.assertRaises(FileNotFoundError):
+            AgenticTraceDataset.from_args(args).load(self.tokenizer)
+
+        empty_path = self.tmpdir_path / "empty_trace.json"
+        with open(empty_path, "w") as f:
+            json.dump({"metadata": {}, "conversations": []}, f)
+        args = make_args(
+            dataset_name="agentic-trace",
+            dataset_path=str(empty_path),
+            num_prompts=1,
+        )
+        with self.assertRaises(ValueError):
+            AgenticTraceDataset.from_args(args).load(self.tokenizer)
+
     def test_dataset_mapping_and_dispatch(self):
         expected = {
+            "agentic-trace",
             "sharegpt",
             "custom",
             "openai",
@@ -603,6 +701,15 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase):
         self.assertEqual(len(speed_bench_rows), 2)
         self.assertTrue(all(isinstance(row, DatasetRow) for row in speed_bench_rows))
 
+        agentic_args = make_args(
+            dataset_name="agentic-trace",
+            dataset_path=self._write_agentic_trace_json(),
+            num_prompts=2,
+        )
+        agentic_rows = get_dataset(agentic_args, self.tokenizer, model_id="dummy-model")
+        self.assertEqual(len(agentic_rows), 2)
+        self.assertTrue(all(isinstance(row, DatasetRow) for row in agentic_rows))
+
     def test_get_dataset_unknown_dataset(self):
         args = make_args(dataset_name="not-a-dataset")
         with self.assertRaises(ValueError):