[None][test] Add InferenceMAX-style GSM8K accuracy eval mode

tensorrt_llm/commands/eval.py

@@ -21,9 +21,9 @@
 from .. import LLM as PyTorchLLM
 from ..evaluate import (AALCR, AIME2025, AIME2026, GSM8K, HLE, MMLU, MMMU,
                         ArenaHard, CnnDailymail, CoVoST2, GPQADiamond,
-                        GPQAExtended, GPQAMain, GPQANemoSkills, IFBench,
-                        JsonModeEval, LongBenchV1, LongBenchV2, MMMUPro,
-                        SciCode)
+                        GPQAExtended, GPQAMain, GPQANemoSkills,
+                        GSM8KInferenceMax, IFBench, JsonModeEval, LongBenchV1,
+                        LongBenchV2, MMMUPro, SciCode)
 from ..llmapi import KvCacheConfig
 from ..llmapi.llm_args import TorchLlmArgs
 from ..llmapi.llm_utils import update_llm_args_with_extra_options
@@ -216,6 +216,7 @@ def main(ctx, model: str, tokenizer: Optional[str],
 main.add_command(CnnDailymail.command)
 main.add_command(MMLU.command)
 main.add_command(GSM8K.command)
+main.add_command(GSM8KInferenceMax.command)
 main.add_command(GPQADiamond.command)
 main.add_command(GPQAMain.command)
 main.add_command(GPQAExtended.command)

tensorrt_llm/evaluate/__init__.py

@@ -18,7 +18,8 @@
 from .covost2 import CoVoST2
 from .json_mode_eval import JsonModeEval
 from .lm_eval import (AIME2025, AIME2026, GSM8K, MMMU, GPQADiamond,
-                      GPQAExtended, GPQAMain, LongBenchV1, MMMUPro)
+                      GPQAExtended, GPQAMain, GSM8KInferenceMax, LongBenchV1,
+                      MMMUPro)
 from .longbench_v2 import LongBenchV2
 from .mmlu import MMLU
 from .nemo_skills_eval import (AALCR, HLE, ArenaHard, GPQANemoSkills, IFBench,
@@ -28,6 +29,7 @@
     "CnnDailymail",
     "MMLU",
     "GSM8K",
+    "GSM8KInferenceMax",
     "GPQADiamond",
     "GPQAMain",
     "GPQAExtended",

tensorrt_llm/evaluate/lm_eval.py

@@ -475,7 +475,9 @@ def __init__(self,
                  output_path: Optional[str] = None,
                  output_dir: Optional[str] = None,
                  post_process_fn: Optional[Callable[[str], str]] = None,
-                 preserve_caller_max_tokens: bool = False):
+                 preserve_caller_max_tokens: bool = False,
+                 shuffle_dataset: bool = True,
+                 fewshot_random_seed: Optional[int] = None):
         try:
             import lm_eval
         except ImportError as e:
@@ -509,6 +511,12 @@ def __init__(self,
         # larger than the lm-eval task's max_gen_toks. Used by thinking
         # models (e.g. Kimi K2.5) whose CoT output exceeds the task default.
         self.preserve_caller_max_tokens = preserve_caller_max_tokens
+        # shuffle_dataset=False keeps every split in its published order, and
+        # fewshot_random_seed overrides random_seed for the few-shot sampler.
+        # Together they reproduce the lm-eval CLI's exemplar selection
+        # (unshuffled data, sampler seed 1234) for external comparability.
+        self.shuffle_dataset = shuffle_dataset
+        self.fewshot_random_seed = fewshot_random_seed
 
         task_manager = TaskManager(
             include_path=f"{os.path.dirname(__file__)}/lm_eval_tasks")
@@ -529,13 +537,17 @@ def _adjust_config(task_dict, random_seed):
                     }
                 else:
                     # NOTE: Few-shot random seed
-                    task_obj.set_fewshot_seed(seed=random_seed)
+                    fewshot_seed = (self.fewshot_random_seed
+                                    if self.fewshot_random_seed is not None else
+                                    random_seed)
+                    task_obj.set_fewshot_seed(seed=fewshot_seed)
                     adjusted_task_dict[task_name] = task_obj
 
                     # NOTE: Shuffle dataset
-                    data = adjusted_task_dict[task_name].dataset
-                    for split in data.keys():
-                        data[split] = data[split].shuffle(random_seed)
+                    if self.shuffle_dataset:
+                        data = adjusted_task_dict[task_name].dataset
+                        for split in data.keys():
+                            data[split] = data[split].shuffle(random_seed)
 
             return adjusted_task_dict
 
@@ -809,6 +821,124 @@ def command(ctx, **kwargs) -> None:
         GSM8K.command_harness(ctx, **kwargs)
 
 
+class GSM8KInferenceMax(LmEvalEvaluator):
+    """GSM8K under the InferenceMAX (SemiAnalysisAI/InferenceX) protocol.
+
+    Mirrors the public InferenceMAX accuracy eval (utils/evals/gsm8k.yaml +
+    ``run_lm_eval`` there): chat template applied with 5-shot multiturn
+    context, an explicit "#### [number]" answer-format instruction in
+    ``doc_to_text``, a generation budget large enough for thinking models
+    (``max_gen_toks: 12288``), and strict ``#### N`` extraction. Scores are
+    directly comparable to inferencex.semianalysis.com/evaluation, unlike the
+    completion-format ``gsm8k`` task whose scores run several points lower on
+    chat/thinking-tuned models.
+
+    Task yaml lives under ``tensorrt_llm/evaluate/lm_eval_tasks/
+    gsm8k_inferencemax/`` (upstream lm-eval ships no such variant).
+
+    Exemplar fidelity: the lm-eval CLI (and therefore InferenceMAX) samples
+    the 5 few-shot exemplars from the unshuffled train split with sampler
+    seed 1234. Defaults below reproduce that; the usual dataset shuffle
+    would silently swap in different exemplars and shift scores.
+    """
+
+    def __init__(self, **kwargs):
+        kwargs.setdefault("shuffle_dataset", False)
+        kwargs.setdefault("fewshot_random_seed", 1234)
+        super().__init__("gsm8k_inferencemax", **kwargs)
+
+    @click.command("gsm8k_inferencemax")
+    @click.option("--dataset_path",
+                  type=str,
+                  default=None,
+                  help="The path to GSM8K 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",
+                  type=click.BOOL,
+                  default=True,
+                  show_default=True,
+                  help="Whether to apply chat template. Default True — the "
+                  "InferenceMAX protocol always evaluates through the chat "
+                  "template.")
+    @click.option(
+        "--chat_template_kwargs",
+        type=str,
+        default=None,
+        callback=lambda ctx, param, value: json.loads(value) if value else None,
+        help='Chat template kwargs as JSON string. Default None keeps the '
+        'template defaults (e.g. adaptive thinking), matching InferenceMAX.')
+    @click.option("--fewshot_as_multiturn",
+                  type=click.BOOL,
+                  default=True,
+                  show_default=True,
+                  help="Apply fewshot as multiturn. Default True — lm-eval "
+                  "auto-enables multiturn fewshots under a chat template, and "
+                  "InferenceMAX inherits that default.")
+    @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=12288,
+        show_default=True,
+        help="Maximum generation length. Mirrors the task yaml "
+        "(``max_gen_toks: 12288``, the InferenceMAX budget of 16384 context "
+        "minus 4096 prompt); thinking models need the room.")
+    @click.option("--temperature",
+                  type=float,
+                  default=None,
+                  help="Sampling temperature. Overrides task yaml gen_kwargs.")
+    @click.option(
+        "--top_p",
+        type=float,
+        default=None,
+        help="Nucleus sampling top_p. Overrides task yaml gen_kwargs.")
+    @click.option("--top_k",
+                  type=int,
+                  default=None,
+                  help="Top-k sampling. Overrides task yaml gen_kwargs.")
+    @click.option("--sampling_seed",
+                  type=int,
+                  default=None,
+                  help="Random seed for generation sampling.")
+    @click.option("--log_samples",
+                  is_flag=True,
+                  default=False,
+                  help="Log sample outputs for debugging.")
+    @click.option("--output_path",
+                  type=str,
+                  default=None,
+                  help="Path to save evaluation results.")
+    @click.option("--output_dir",
+                  type=str,
+                  default=None,
+                  help="Directory to save the task infos.")
+    @click.pass_context
+    @staticmethod
+    def command(ctx, **kwargs) -> None:
+        if kwargs.get("fewshot_as_multiturn", False):
+            assert kwargs.get(
+                "apply_chat_template", False
+            ), "apply_chat_template must be True when fewshot_as_multiturn is True"
+        GSM8KInferenceMax.command_harness(ctx, **kwargs)
+
+
 class GPQADiamond(LmEvalEvaluator):
 
     def __init__(self, **kwargs):

tensorrt_llm/evaluate/lm_eval_tasks/gsm8k_inferencemax/gsm8k_inferencemax.yaml

@@ -0,0 +1,54 @@
+# GSM8K under the InferenceMAX protocol. Mirrors
+# https://github.com/SemiAnalysisAI/InferenceX utils/evals/gsm8k.yaml (itself
+# derived from upstream lm-evaluation-harness gsm8k.yaml with doc_to_text
+# changed so chat models answer in an extractable format; see
+# lm-evaluation-harness#3411). Differences from that file: the task is renamed
+# to avoid clashing with the upstream registry, and max_gen_toks pins the
+# InferenceMAX generation budget (16384 context minus 4096 prompt) so
+# thinking-model CoT is not truncated at lm-eval's 256-token default.
+tag:
+  - math_word_problems
+task: gsm8k_inferencemax
+dataset_path: openai/gsm8k
+dataset_name: main
+output_type: generate_until
+training_split: train
+fewshot_split: train
+test_split: test
+doc_to_text: "Question: {{question}}\nEnd your response with the answer on the last line, formatted as: #### [number]\nAnswer:"
+doc_to_target: "{{answer}}"
+metric_list:
+  - metric: exact_match
+    aggregation: mean
+    higher_is_better: true
+    ignore_case: true
+    ignore_punctuation: false
+    regexes_to_ignore:
+      - ","
+      - "\\$"
+      - "(?s).*#### "
+      - "\\.$"
+generation_kwargs:
+  until:
+    - "</s>"
+    - "<|im_end|>"
+  do_sample: false
+  temperature: 0.0
+  max_gen_toks: 12288
+repeats: 1
+num_fewshot: 5
+filter_list:
+  - name: "strict-match"
+    filter:
+      - function: "regex"
+        group_select: -1
+        regex_pattern: "#### (\\-?[0-9\\.\\,]+)"
+      - function: "take_first"
+  - name: "flexible-extract"
+    filter:
+      - function: "regex"
+        group_select: -1
+        regex_pattern: "(-?[$0-9.,]{2,})|(-?[0-9]+)"
+      - function: "take_first"
+metadata:
+  version: 3.0

tests/integration/defs/accuracy/accuracy_core.py

@@ -78,6 +78,13 @@ class HypothesisTestingParams:
     beta: float = 0.2
     sigma: float = 50.0
     higher_is_better: bool = True
+    # Free-text provenance for the reference row (e.g. an externally
+    # published score to compare against); printed in the report only.
+    reference_note: Optional[str] = None
+    # Explicit pass/fail cutoff. When set, it replaces the hypothesis-test
+    # threshold computed from sigma/alpha — use for benchmark-style floors
+    # where a fixed, human-readable gate is preferred.
+    threshold_override: Optional[float] = None
     theta: float = field(init=False)
     threshold: float = field(init=False)
 
@@ -86,16 +93,34 @@ def __post_init__(self) -> None:
                                    sigma=self.sigma,
                                    alpha=self.alpha,
                                    beta=self.beta)
-        self.threshold = compute_threshold(
-            self.num_samples,
-            self.ref_accuracy,
-            sigma=self.sigma,
-            alpha=self.alpha,
-            higher_is_better=self.higher_is_better)
+        if self.threshold_override is not None:
+            self.threshold = self.threshold_override
+        else:
+            self.threshold = compute_threshold(
+                self.num_samples,
+                self.ref_accuracy,
+                sigma=self.sigma,
+                alpha=self.alpha,
+                higher_is_better=self.higher_is_better)
 
     def report(self, accuracy: Optional[float] = None) -> str:
         metric_name = self.metric_name.upper()
-        report = f"""===========================================================
+        reference_line = f"Reference {self.metric_name}: {self.ref_accuracy:.3f}"
+        if self.reference_note is not None:
+            reference_line += f" ({self.reference_note})"
+        if self.threshold_override is not None:
+            # The hypothesis-test statistics do not apply to an explicit
+            # floor; printing them would suggest a different threshold.
+            report = f"""===========================================================
+= {metric_name} CHECK
+===========================================================
+#Samples: {self.num_samples}
+Higher is better: {self.higher_is_better}
+{reference_line}
+Threshold (explicit floor): {self.threshold:.3f}
+==========================================================="""
+        else:
+            report = f"""===========================================================
 = {metric_name} HYPOTHESIS TESTING
 ===========================================================
 Alpha (Type I:  False Positive): {self.alpha:.3f}
@@ -104,7 +129,7 @@ def report(self, accuracy: Optional[float] = None) -> str:
 #Samples: {self.num_samples}
 Higher is better: {self.higher_is_better}
 Theta (Minimum detectable effect): {self.theta:.3f}
-Reference {self.metric_name}: {self.ref_accuracy:.3f}
+{reference_line}
 Threshold: {self.threshold:.3f}
 ==========================================================="""
         if accuracy is not None:
@@ -190,7 +215,9 @@ def get_hypothesis_testing_params(self,
             sigma=entry.get("sigma", self.SIGMA),
             num_samples=entry.get("num_samples", self.NUM_SAMPLES),
             higher_is_better=entry.get("higher_is_better",
-                                       self.HIGHER_IS_BETTER))
+                                       self.HIGHER_IS_BETTER),
+            reference_note=entry.get("reference_note"),
+            threshold_override=entry.get("threshold"))
 
     def evaluate(self,
                  llm: Union[PyTorchLLM, AutoDeployLLM],
@@ -414,6 +441,36 @@ class GSM8K(AccuracyTask):
     EVALUATE_KWARGS = dict(scores_filter=None)
 
 
+class GSM8KInferenceMax(AccuracyTask):
+    # GSM8K under the InferenceMAX (SemiAnalysisAI/InferenceX) protocol: chat
+    # template + 5-shot multiturn + "#### [number]" format instruction + a
+    # 12288-token generation budget for thinking models, gated on strict
+    # "#### N" extraction. Scores are directly comparable to
+    # inferencex.semianalysis.com/evaluation; the completion-format GSM8K
+    # above measures the same model several points lower on chat/thinking
+    # models. Substantially slower than GSM8K: thinking output dominates
+    # (minutes with spec decoding + CUDA graphs, tens of minutes eager).
+    DATASET = "gsm8k_inferencemax"
+    DATASET_DIR = f"{llm_models_root()}/datasets/openai/gsm8k"
+
+    ALPHA = 0.05
+    BETA = 0.2
+    SIGMA = 50
+    NUM_SAMPLES = 1319  # Full sample
+
+    MAX_INPUT_LEN = 4096
+    MAX_OUTPUT_LEN = 12288
+
+    EVALUATOR_CLS = tensorrt_llm.evaluate.GSM8KInferenceMax
+    EVALUATOR_KWARGS = dict(dataset_path=DATASET_DIR,
+                            random_seed=0,
+                            apply_chat_template=True,
+                            fewshot_as_multiturn=True)
+
+    # InferenceMAX publishes em_strict as the primary score.
+    EVALUATE_KWARGS = dict(scores_filter="exact_match,strict-match")
+
+
 class GPQADiamond(AccuracyTask):
     DATASET = "gpqa_diamond"
     DATASET_DIR = f"{llm_models_root()}/datasets/gpqa"