[None][feat] Add multi-turn support for trtllm-bench

tensorrt_llm/bench/benchmark/throughput.py

@@ -434,6 +434,12 @@ def throughput_command(
 
         post_proc_params = None  # No detokenization
 
+        has_multi_turn = any(r.is_multi_turn for r in requests)
+        multi_turn_tokenizer = tokenizer if has_multi_turn else None
+        if has_multi_turn:
+            logger.info("Multi-turn requests detected. Turns will be processed "
+                        "sequentially within each request.")
+
         # Perform warmup if requested.
         if options.warmup > 0:
             logger.info("Setting up for warmup...")
@@ -446,7 +452,8 @@ def throughput_command(
                                 warmup_dataset,
                                 False,
                                 options.concurrency,
-                                modality=options.modality))
+                                modality=options.modality,
+                                tokenizer=multi_turn_tokenizer))
             # WAR: IterationResult is a singleton tied to the executor.
             # Since the benchmark calls asyncio.run() multiple times (e.g., during warmup),
             # we must reset it to ensure it attaches to the correct event loop.
@@ -463,7 +470,8 @@ def throughput_command(
                                 options.streaming,
                                 options.concurrency,
                                 iteration_writer.full_address,
-                                modality=options.modality))
+                                modality=options.modality,
+                                tokenizer=multi_turn_tokenizer))
 
         logger.info("Benchmark done. Reporting results...")
         if options.modality is not None:

tensorrt_llm/bench/benchmark/utils/asynchronous.py

@@ -7,6 +7,7 @@
 from typing import Dict, List, Optional, Set, Tuple
 
 import tqdm
+from transformers import PreTrainedTokenizer
 from zmq import PUSH
 from zmq.asyncio import Context
 
@@ -28,7 +29,8 @@ def __init__(self,
                  outbox: asyncio.Queue[PerfItemTuple],
                  streaming: bool,
                  concurrency: int = -1,
-                 modality: Optional[str] = None) -> None:
+                 modality: Optional[str] = None,
+                 tokenizer: Optional[PreTrainedTokenizer] = None) -> None:
         self.llm = llm
         self._inbox: asyncio.Queue[Tuple[InferenceRequest,
                                          SamplingParams]] = asyncio.Queue()
@@ -45,18 +47,27 @@ def __init__(self,
         self.streaming = streaming
         self.request_seen = asyncio.Event()
         self.modality = modality
+        self.tokenizer = tokenizer
 
     async def process_request(self, request: InferenceRequest,
                               sampling_params: SamplingParams,
                               post_proc_params: PostprocParams):
-        # Set up sampling params with inference request
+        if request.is_multi_turn and self.tokenizer is not None:
+            await self._process_multi_turn_request(request, sampling_params,
+                                                   post_proc_params)
+        else:
+            await self._process_single_request(request, sampling_params,
+                                               post_proc_params)
+
+    async def _process_single_request(self, request: InferenceRequest,
+                                      sampling_params: SamplingParams,
+                                      post_proc_params: PostprocParams):
         self.request_seen.set()
         sampling_params.max_tokens = request.output_tokens
 
         async with semaphore_guard(self._concurrency_semaphore):
             request_start_timestamp = time.perf_counter_ns()
             time_on_first_token = None
-            # Schedule the request in the LLM API (asynchronously)
             logger.debug(f"request.lora_request: {request.lora_request}")
             output: RequestOutput = self.llm.generate_async(
                 request.input_ids if self.modality is None else request.prompt,
@@ -70,12 +81,10 @@ async def process_request(self, request: InferenceRequest,
                         time_on_first_token = time.perf_counter_ns()
                         response = stream_output
             else:
-                # Wait for the response to return to us.
                 response: RequestOutput = await output.aresult()
 
         response_end_timestamp = time.perf_counter_ns()
 
-        # Mark that the response returned. Construct a record to send to statistics.
         tokens = list(chain(*(beam.token_ids for beam in response.outputs)))
         request_perf_item = PerfItemTuple(
             start_timestamp=request_start_timestamp,
@@ -89,7 +98,75 @@ async def process_request(self, request: InferenceRequest,
             time_on_first_token=time_on_first_token,
         )
 
-        # Register the new request perf items in the outbound queue for statistics keeping
+        await self._outbox.put(request_perf_item)
+
+    async def _process_multi_turn_request(self, request: InferenceRequest,
+                                          sampling_params: SamplingParams,
+                                          post_proc_params: PostprocParams):
+        """Process a multi-turn request by iterating through turns sequentially.
+
+        Each turn builds on the previous conversation context: the model's
+        response from turn N is appended to the messages before encoding
+        turn N+1.  All turns within a single request share one concurrency
+        slot so that the conversation history stays consistent.
+        """
+        self.request_seen.set()
+        sampling_params.max_tokens = request.output_tokens
+        tokenizer = self.tokenizer
+
+        messages: List[dict] = []
+        total_input_tokens = 0
+        all_output_tokens: List[int] = []
+
+        async with semaphore_guard(self._concurrency_semaphore):
+            request_start_timestamp = time.perf_counter_ns()
+            time_on_first_token = None
+            last_response = None
+
+            for turn_id, question in enumerate(request.turns):
+                messages.append({"role": "user", "content": question})
+
+                input_ids = tokenizer.apply_chat_template(
+                    messages, add_generation_prompt=True)
+
+                output: RequestOutput = self.llm.generate_async(
+                    input_ids,
+                    sampling_params=sampling_params,
+                    _postproc_params=post_proc_params,
+                    streaming=False)
+                response: RequestOutput = await output.aresult()
+
+                if turn_id == 0 and time_on_first_token is None:
+                    time_on_first_token = time.perf_counter_ns()
+
+                turn_tokens = list(
+                    chain(*(beam.token_ids for beam in response.outputs)))
+                all_output_tokens.extend(turn_tokens)
+                total_input_tokens += len(input_ids)
+
+                assistant_text = tokenizer.decode(turn_tokens,
+                                                  skip_special_tokens=True)
+                messages.append({
+                    "role": "assistant",
+                    "content": assistant_text
+                })
+
+                last_response = response
+
+        response_end_timestamp = time.perf_counter_ns()
+
+        request_perf_item = PerfItemTuple(
+            start_timestamp=request_start_timestamp,
+            end_timestamp=response_end_timestamp,
+            request_id=last_response.id,
+            num_input_tokens=total_input_tokens,
+            response_is_final=last_response.finished,
+            error=False,
+            tokens=all_output_tokens,
+            decoding_iteration=last_response.decoding_iter,
+            time_on_first_token=time_on_first_token,
+        )
+
         await self._outbox.put(request_perf_item)
 
     def _raise_for_failed_tasks(self):
@@ -256,6 +333,7 @@ async def async_benchmark(
     concurrency: int = -1,
     iteration_log_addr: str = None,
     modality: Optional[str] = None,
+    tokenizer: Optional[PreTrainedTokenizer] = None,
 ) -> StatsKeeper:
     outbox = asyncio.Queue()
     statistics = StatsKeeper()
@@ -266,7 +344,8 @@ async def async_benchmark(
                          outbox,
                          streaming,
                          concurrency=concurrency,
-                         modality=modality)
+                         modality=modality,
+                         tokenizer=tokenizer)
     enqueue_task: Optional[asyncio.Task] = None
     try:
         backend.run(iteration_addr=iteration_log_addr)

tensorrt_llm/bench/dataclasses/general.py

@@ -24,14 +24,22 @@ class InferenceRequest(BaseModel):
     input_ids: Optional[List[int]] = Field(
         alias=AliasChoices("input_ids", "logits"))
     lora_request: Optional[LoRARequest] = None
+    turns: Optional[List[str]] = None
+    category: Optional[str] = None
+    question_id: Optional[int] = None
 
     @model_validator(mode="after")
     def verify_prompt_and_logits(self) -> InferenceRequest:
-        if self.prompt is None and self.input_ids is None:
+        if self.prompt is None and self.input_ids is None and self.turns is None:
             raise ValueError(
-                f"Both prompt and input_ids for {self.task_id} are both None.")
+                f"prompt, input_ids, and turns for task {self.task_id} are all None."
+            )
         return self
 
+    @property
+    def is_multi_turn(self) -> bool:
+        return self.turns is not None and len(self.turns) > 1
+
 
 class DatasetMetadata(BaseModel):
     isl_stats: PercentileStats

tensorrt_llm/bench/utils/data.py

@@ -51,7 +51,7 @@ def create_dataset_from_stream(
         tokenizer (PreTrainedTokenizer): HuggingFace tokenizer.
         stream (TextIO): Stream of input requests.
         max_input_length (int, optional): Maximum input length to cap prompts to. Defaults to 0.
-        max_output_length (int, optional): Maximum output length to cap prompts to.. Defaults to 0.
+        max_output_length (int, optional): Maximum output length to cap prompts to. Defaults to 0.
         num_requests (int, optional): Number of requests to limit to. Defaults to 0.
 
     Returns:
@@ -86,30 +86,55 @@ def create_dataset_from_stream(
     all_logits = []
     task_ids = []
     lora_requests = []
+    all_turns = []
+    all_categories = []
+    all_question_ids = []
     while (line := stream.readline()) and len(task_ids) < max_requests:
-        # We expect the data to come in as a JSON string.
-        # For example:
+        # We support two JSONL formats:
+        #
+        # 1. Standard single-turn format:
         # {"task_id": 1, "prompt": "Generate an infinite response to the following:
         # There once was a man who.", "output_tokens": 1000}
         #
+        # 2. Multi-turn format (e.g. MT-Bench question.jsonl):
+        # {"question_id": 81, "category": "writing", "turns": ["Write a blog post...", "Rewrite..."]}
+        # When "turns" is present, the first turn is used as the prompt for
+        # tokenization/metadata, and all turns are stored for sequential
+        # multi-turn benchmarking.  "output_tokens" is required, same as
+        # single-turn format.
+        #
         # For multimodal data, the data should be of the form:
-        # {"task_id": 1, "prompt": "Generate an infinite response to the following:
-        # There once was a man who.", "output_tokens": 1000,
+        # {"task_id": 1, "prompt": "...", "output_tokens": 1000,
         # "media_paths": ["/path/to/image1.jpg", "/path/to/image2.jpg"]}
         #
         # For LoRA data, the data should be of the form:
-        # {"task_id": 1, "prompt": "Generate an infinite response to the following:
-        # There once was a man who.", "output_tokens": 1000,
+        # {"task_id": 1, "prompt": "...", "output_tokens": 1000,
         # "lora_request": {"lora_name": "my_lora", "lora_int_id": 1, "lora_path": "/path/to/lora"}}
         #
         # Each line should be a complete JSON dictionary with no indentation
-        # or newline characters. The task_id field is required.
+        # or newline characters.
         data = json.loads(line)
-        prompts.append(data.get("prompt"))
-        media_paths.append(data.get("media_paths", None))
-        all_logits.append(data.get("input_ids", data.get("logits", None)))
-        all_osl.append(data.get("output_tokens"))
-        task_ids.append(data.get("task_id"))
+
+        turns = data.get("turns")
+        if turns is not None and isinstance(turns, list):
+            prompts.append(data.get("prompt") or turns[0])
+            media_paths.append(data.get("media_paths", None))
+            all_logits.append(data.get("input_ids", data.get("logits", None)))
+            all_turns.append(turns)
+            all_categories.append(data.get("category"))
+            all_question_ids.append(data.get("question_id"))
+            all_osl.append(data.get("output_tokens"))
+            task_ids.append(
+                data.get("task_id", data.get("question_id", len(task_ids))))
+        else:
+            prompts.append(data.get("prompt"))
+            media_paths.append(data.get("media_paths", None))
+            all_logits.append(data.get("input_ids", data.get("logits", None)))
+            all_turns.append(None)
+            all_categories.append(None)
+            all_question_ids.append(None)
+            all_osl.append(data.get("output_tokens"))
+            task_ids.append(data.get("task_id"))
 
         # Parse LoRA request if present
         lora_data = data.get("lora_request", None)
@@ -138,8 +163,9 @@ def create_dataset_from_stream(
 
     all_isl = []
     all_seq_len = []
-    for prompt, logits, osl, task_id, lora_request in zip(
-            prompts, all_logits, all_osl, task_ids, lora_requests):
+    for prompt, logits, osl, task_id, lora_request, turns, category, question_id in zip(
+            prompts, all_logits, all_osl, task_ids, lora_requests, all_turns,
+            all_categories, all_question_ids):
         if modality is not None:
             # NOTE: we cannot tokenize multi-modal data, handled by preprocessor
             #       so the actual sequence length is unknown until the model is run
@@ -151,14 +177,18 @@ def create_dataset_from_stream(
             logits = tokenize(prompt)["input_ids"] if logits is None else logits
             cur_isl = len(logits)
         all_isl.append(cur_isl)
-        all_seq_len.append(cur_isl + osl)
+        num_turns = len(turns) if turns is not None else 1
+        all_seq_len.append(cur_isl + num_turns * osl)
 
         request = InferenceRequest(
             task_id=task_id,
             prompt=prompt,
             output_tokens=output_limiter(osl),
             input_ids=logits,
             lora_request=lora_request,
+            turns=turns,
+            category=category,
+            question_id=question_id,
         )
         dataset.append(request)