[None][chore] Guided decoding benchmarking

requirements-dev.txt

@@ -28,6 +28,6 @@ jieba==0.42.1
 rouge==1.0.1
 pytest-rerunfailures
 ruff==0.9.4
-lm_eval[api]==0.4.8
+lm_eval[api]==0.4.9
 docstring_parser
 genai-perf==0.0.13

tensorrt_llm/_torch/pyexecutor/grammar_matcher.py

@@ -194,7 +194,11 @@ def create(
                 grammar = llguidance.LLMatcher.grammar_from_json_schema(
                     '{"type": "object"}')
             case GuidedDecodingParams.GuideType.JSON_SCHEMA:
-                grammar = llguidance.LLMatcher.grammar_from_json_schema(guide)
+                kwargs = {}
+                if os.environ.get("TRTLLM_XGUIDANCE_LENIENT") == "1":
+                    kwargs["overrides"] = {"lenient": True}
+                grammar = llguidance.LLMatcher.grammar_from_json_schema(
+                    guide, **kwargs)
             case GuidedDecodingParams.GuideType.REGEX:
                 grammar = llguidance.LLMatcher.grammar_from_regex(guide)
             case GuidedDecodingParams.GuideType.EBNF_GRAMMAR:

tensorrt_llm/_torch/pyexecutor/guided_decoder.py

@@ -1,4 +1,5 @@
 import math
+import os
 from dataclasses import dataclass
 from queue import Queue
 from typing import Iterable, List, Optional, Tuple
@@ -15,6 +16,9 @@
 from .llm_request import LlmRequest
 from .scheduler import ScheduledRequests
 
+TRTLLM_DISABLE_DRAFT_GUIDED_DECODING = os.environ.get(
+    "TRTLLM_DISABLE_DRAFT_GUIDED_DECODING", "0") == "1"
+
 
 @dataclass(slots=True)
 class GuidedRequest:
@@ -492,6 +496,9 @@ def add_draft_batch(self,
                         new_tokens: torch.Tensor,
                         num_accepted_tokens: torch.Tensor,
                         draft_step: int = 0) -> None:
+        if TRTLLM_DISABLE_DRAFT_GUIDED_DECODING:
+            if draft_step > 0:
+                return
         batch_size = len(self.requests)
         assert new_tokens.size(0) == batch_size
         self.new_tokens[0, :batch_size].copy_(new_tokens, non_blocking=True)
@@ -531,6 +538,16 @@ def execute_draft_batch(self,
                             logits: torch.Tensor,
                             d2t: Optional[torch.Tensor] = None,
                             draft_step: int = 0) -> None:
+        if TRTLLM_DISABLE_DRAFT_GUIDED_DECODING:
+            if draft_step == 0:
+                with torch.cuda.stream(self.stream):
+                    torch.cuda.current_stream().wait_event(self.token_event)
+                    self.fetch_draft_batch(draft_step=draft_step)
+                    self.rollback_rejected_tokens()
+                    self.bitmask_event.record()
+                torch.cuda.current_stream().wait_event(self.bitmask_event)
+            return
+
         with torch.cuda.stream(self.stream):
             torch.cuda.current_stream().wait_event(self.token_event)
             self.fetch_draft_batch(draft_step=draft_step)

tensorrt_llm/_torch/speculative/model_drafter.py

@@ -9,7 +9,8 @@
 from tensorrt_llm.logger import logger
 
 from ..attention_backend.trtllm import TrtllmAttention
-from ..pyexecutor.guided_decoder import GuidedDecoder
+from ..pyexecutor.guided_decoder import (TRTLLM_DISABLE_DRAFT_GUIDED_DECODING,
+                                         GuidedDecoder)
 from ..pyexecutor.handle_logits import HandleLogits
 from ..pyexecutor.llm_request import LlmRequest, LlmRequestState
 from ..pyexecutor.resource_manager import (BaseResourceManager, ResourceManager,
@@ -28,15 +29,20 @@
 
 # Place the tool function here to avoid circular import
 def get_draft_model_prompt(spec_dec_mode: SpeculativeDecodingMode,
-                           input_tokens: torch.Tensor) -> torch.Tensor:
+                           request: LlmRequest) -> List[int]:
     """
     Can be used to modify prompts for speculative algorithms that need to update tokens
     before drafting.
     """
+    draft_input_tokens = request.get_tokens(0)
     if spec_dec_mode.is_eagle3() or spec_dec_mode.is_mtp_eagle():
         # EAGLE3 always throws away the first token when processing draft inputs
-        return input_tokens[1:]
-    return input_tokens
+        draft_input_tokens = draft_input_tokens[1:]
+    if request.is_context_init_state:
+        # A draft request's prompt is its target request's prompt adding the first golden token.
+        # Add a fake golden token here since the real one has not been generated.
+        draft_input_tokens.append(0)
+    return draft_input_tokens
 
 
 class ModelDrafter(Drafter):
@@ -70,7 +76,10 @@ def __init__(
         self.max_draft_tokens = max_draft_tokens
         # Sampling
         self.sampler = sampler
-        self.guided_decoder = guided_decoder
+        if TRTLLM_DISABLE_DRAFT_GUIDED_DECODING:
+            self.guided_decoder = None
+        else:
+            self.guided_decoder = guided_decoder
 
         self.use_static_draft_loop = draft_model_engine.model_is_wrapped
         if self.use_static_draft_loop:
@@ -163,7 +172,7 @@ def _create_draft_request_for_request(
         num_draft_tokens, num_accepted_tokens = self._initialize_draft_tokens(
             request)
         input_tokens = get_draft_model_prompt(self.spec_config.spec_dec_mode,
-                                              request.get_tokens(0))
+                                              request)
 
         # First time seeing this request - context request
         if request.max_beam_num_tokens - 1 == request.py_prompt_len:
@@ -235,9 +244,8 @@ def _prepare_draft_batch(
                 # We hit this path if we're doing chunked prefill. The target model processed
                 # a prefill chunk on the last iteration. Now, we need to fill in the KV cache
                 # for the draft model too.
-                all_tokens = request.get_tokens(0)
                 input_tokens = get_draft_model_prompt(
-                    self.spec_config.spec_dec_mode, all_tokens)
+                    self.spec_config.spec_dec_mode, request)
 
                 new_request = self._create_context_request(
                     request, input_tokens)

tensorrt_llm/bench/benchmark/utils/asynchronous.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import copy
 import time
 from contextlib import asynccontextmanager
 from itertools import chain
@@ -50,7 +51,9 @@ async def process_request(self, request: InferenceRequest,
                               post_proc_params: PostprocParams):
         # Set up sampling params with inference request
         self.request_seen.set()
+        sampling_params = copy.deepcopy(sampling_params)
         sampling_params.max_tokens = request.output_tokens
+        sampling_params.guided_decoding = request.guided_decoding_params
 
         async with semaphore_guard(self._concurrency_semaphore):
             request_start_timestamp = time.perf_counter_ns()

tensorrt_llm/bench/dataclasses/general.py

@@ -8,6 +8,7 @@
 
 from tensorrt_llm.bench.dataclasses.statistics import PercentileStats
 from tensorrt_llm.executor.request import LoRARequest
+from tensorrt_llm.sampling_params import GuidedDecodingParams
 
 
 class BenchmarkEnvironment(BaseModel):
@@ -23,6 +24,7 @@ class InferenceRequest(BaseModel):
     input_ids: Optional[List[int]] = Field(
         alias=AliasChoices("input_ids", "logits"))
     lora_request: Optional[LoRARequest] = None
+    guided_decoding_params: Optional[GuidedDecodingParams] = None
 
     @model_validator(mode="after")
     def verify_prompt_and_logits(self) -> InferenceRequest:

tensorrt_llm/bench/utils/data.py

@@ -9,6 +9,7 @@
 from tensorrt_llm.bench.dataclasses.statistics import PercentileStats
 from tensorrt_llm.executor.request import LoRARequest
 from tensorrt_llm.inputs import default_multimodal_input_loader
+from tensorrt_llm.sampling_params import GuidedDecodingParams
 
 
 def initialize_tokenizer(model_name: str) -> PreTrainedTokenizer:
@@ -86,6 +87,7 @@ def create_dataset_from_stream(
     all_logits = []
     task_ids = []
     lora_requests = []
+    all_guided_decoding_params = []
     while (line := stream.readline()) and len(task_ids) < max_requests:
         # We expect the data to come in as a JSON string.
         # For example:
@@ -121,6 +123,14 @@ def create_dataset_from_stream(
         else:
             lora_requests.append(None)
 
+        # Parse guided decoding params if present
+        if (guided_decoding_params :=
+                data.get("guided_decoding_params")) is not None:
+            all_guided_decoding_params.append(
+                GuidedDecodingParams(**guided_decoding_params))
+        else:
+            all_guided_decoding_params.append(None)
+
     if modality is not None:
         # Multimodal data need extra preprocessing
         assert modality in [
@@ -138,8 +148,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, guided_decoding_params in zip(
+            prompts, all_logits, all_osl, task_ids, lora_requests,
+            all_guided_decoding_params):
         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
@@ -159,6 +170,7 @@ def create_dataset_from_stream(
             output_tokens=output_limiter(osl),
             input_ids=logits,
             lora_request=lora_request,
+            guided_decoding_params=guided_decoding_params,
         )
         dataset.append(request)
 

tensorrt_llm/commands/eval.py

@@ -21,7 +21,8 @@
 from .. import LLM as PyTorchLLM
 from .._tensorrt_engine import LLM
 from ..evaluate import (GSM8K, MMLU, MMMU, CnnDailymail, GPQADiamond,
-                        GPQAExtended, GPQAMain, JsonModeEval)
+                        GPQAExtended, GPQAMain, HumanEval, JsonModeEval,
+                        JsonSchemaBenchEasy)
 from ..llmapi import BuildConfig, KvCacheConfig
 from ..llmapi.llm_utils import update_llm_args_with_extra_options
 from ..logger import logger, severity_map
@@ -159,6 +160,8 @@ def main(ctx, model: str, tokenizer: Optional[str], log_level: str,
 main.add_command(GPQAExtended.command)
 main.add_command(JsonModeEval.command)
 main.add_command(MMMU.command)
+main.add_command(HumanEval.command)
+main.add_command(JsonSchemaBenchEasy.command)
 
 if __name__ == "__main__":
     main()

tensorrt_llm/evaluate/__init__.py

@@ -15,10 +15,11 @@
 
 from .cnn_dailymail import CnnDailymail
 from .json_mode_eval import JsonModeEval
-from .lm_eval import GSM8K, MMMU, GPQADiamond, GPQAExtended, GPQAMain
+from .lm_eval import (GSM8K, MMMU, GPQADiamond, GPQAExtended, GPQAMain,
+                      HumanEval, JsonSchemaBenchEasy)
 from .mmlu import MMLU
 
 __all__ = [
     "CnnDailymail", "MMLU", "GSM8K", "GPQADiamond", "GPQAMain", "GPQAExtended",
-    "JsonModeEval", "MMMU"
+    "JsonModeEval", "MMMU", "HumanEval", "JsonSchemaBenchEasy"
 ]

tensorrt_llm/evaluate/interface.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import copy
+import json
 import random
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, List, Optional, Union
@@ -107,6 +108,24 @@ def evaluate(self,
         profiler.reset("trtllm exec")
 
         score = self.compute_score(results, references, *zip(*auxiliaries))
+
+        from dataclasses import asdict
+        num_repeats = 100
+        with open("bench_requests.jsonl", "w") as f:
+            for i, output in enumerate(outputs * num_repeats):
+                bench_request = {
+                    "task_id": i,
+                    "prompt": output.prompt,
+                    "input_ids": output.prompt_token_ids,
+                    "output_tokens": output.sampling_params.max_tokens
+                }
+                guided_decoding_params = output.sampling_params.guided_decoding
+                if guided_decoding_params is not None:
+                    bench_request["guided_decoding_params"] = asdict(
+                        guided_decoding_params)
+                f.write(json.dumps(bench_request))
+                f.write("\n")
+
         return score
 
     @staticmethod

tensorrt_llm/evaluate/lm_eval.py

@@ -395,7 +395,8 @@ def evaluate(self,
             limit=self.num_samples,
             apply_chat_template=self.apply_chat_template,
             fewshot_as_multiturn=self.fewshot_as_multiturn,
-            system_instruction=self.system_prompt)
+            system_instruction=self.system_prompt,
+            confirm_run_unsafe_code=True)
         # Normalize scores to range 0~100
         scores = results["results"][self.task_name]
         for metric in scores.keys():
@@ -617,6 +618,100 @@ def command(ctx, **kwargs) -> None:
         GPQAExtended.command_harness(ctx, **kwargs)
 
 
+class HumanEval(LmEvalEvaluator):
+
+    def __init__(self, **kwargs):
+        super().__init__("humaneval", **kwargs)
+
+    @click.command("humaneval")
+    @click.option("--dataset_path",
+                  type=str,
+                  default=None,
+                  help="The path to HumanEval dataset. "
+                  "If unspecified, the dataset is downloaded from HF hub.")
+    @click.option(
+        "--num_samples",
+        type=int,
+        default=None,
+        help="Number of samples to run the evaluation; None means full dataset."
+    )
+    @click.option("--random_seed",
+                  type=int,
+                  default=0,
+                  help="Random seed for dataset processing.")
+    @click.option("--apply_chat_template",
+                  is_flag=True,
+                  default=False,
+                  help="Whether to apply chat template.")
+    @click.option("--fewshot_as_multiturn",
+                  is_flag=True,
+                  default=False,
+                  help="Apply fewshot as multiturn.")
+    @click.option("--system_prompt",
+                  type=str,
+                  default=None,
+                  help="System prompt.")
+    @click.option("--max_input_length",
+                  type=int,
+                  default=4096,
+                  help="Maximum prompt length.")
+    @click.option("--max_output_length",
+                  type=int,
+                  default=2048,
+                  help="Maximum generation length.")
+    @click.pass_context
+    @staticmethod
+    def command(ctx, **kwargs) -> None:
+        HumanEval.command_harness(ctx, **kwargs)
+
+
+class JsonSchemaBenchEasy(LmEvalEvaluator):
+
+    def __init__(self, **kwargs):
+        super().__init__("jsonschema_bench_easy", **kwargs)
+
+    @click.command("jsonschema_bench_easy")
+    @click.option("--dataset_path",
+                  type=str,
+                  default=None,
+                  help="The path to JsonSchemaBench dataset. "
+                  "If unspecified, the dataset is downloaded from HF hub.")
+    @click.option(
+        "--num_samples",
+        type=int,
+        default=None,
+        help="Number of samples to run the evaluation; None means full dataset."
+    )
+    @click.option("--random_seed",
+                  type=int,
+                  default=0,
+                  help="Random seed for dataset processing.")
+    @click.option("--apply_chat_template",
+                  is_flag=True,
+                  default=False,
+                  help="Whether to apply chat template.")
+    @click.option("--fewshot_as_multiturn",
+                  is_flag=True,
+                  default=False,
+                  help="Apply fewshot as multiturn.")
+    @click.option("--system_prompt",
+                  type=str,
+                  default=None,
+                  help="System prompt.")
+    @click.option("--max_input_length",
+                  type=int,
+                  default=4096,
+                  help="Maximum prompt length.")
+    @click.option("--max_output_length",
+                  type=int,
+                  default=2048,
+                  help="Maximum generation length.")
+    @click.pass_context
+    @staticmethod
+    def command(ctx, **kwargs) -> None:
+        JsonSchemaBenchEasy.command_harness(ctx, **kwargs)
+
+
 class MMMU(LmEvalEvaluator):
 
     def __init__(self, **kwargs):