feat: Enable per dataset max-osl (#344)

* Enable per dataset max-osl

---------

Signed-off-by: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com>

Author: arekay-nv

examples/04_GPTOSS120B_Example/Readme.md

@@ -57,6 +57,39 @@ The config uses `api_type: openai_completions`, which routes to `/v1/completions
 token IDs (`prompt: [id, id, ...]`). This applies the Harmony format client-side and bypasses vLLM's
 chat template, producing the same token sequence as the SGLang path and matching SGLang accuracy scores.
 
+#### MLPerf per-dataset output budgets
+
+The MLPerf Inference gpt-oss-120b rules use **different max output lengths** for the performance and
+accuracy phases (see the
+[MLCommons reference](https://github.com/mlcommons/inference/tree/master/language/gpt-oss-120b)):
+
+| Phase       | `max_output_len` | `reasoning_effort` |
+| ----------- | ---------------- | ------------------ |
+| performance | 10240            | low                |
+| accuracy    | 32768            | high               |
+
+[`vllm_gptoss_120b_per_dataset_osl_example.yaml`](vllm_gptoss_120b_per_dataset_osl_example.yaml)
+expresses this in a single run using per-dataset `generation_config_override`: the global
+`model_params.max_new_tokens` sets the perf budget (10240) and each accuracy dataset overrides it to
+32768 so the harmony CoT is not truncated before the `assistantfinal` answer. This config sets only the
+output budget: `reasoning_effort` is not a YAML field. With `api_type: openai_completions` the Harmony
+format is applied client-side, so reasoning effort is encoded into the prompt tokens — pre-tokenized in
+the perf dataset, and fixed at `high` by the adapter's `Harmonize()` for the accuracy datasets.
+
+> **Reasoning effort caveat.** This config controls only the output budget, not reasoning effort. The
+> accuracy datasets carry text prompts, so the client-side `Harmonize()` encodes them at its default
+> `reasoning_effort=high` — matching the MLPerf accuracy setting. The perf dataset is consumed
+> pre-tokenized (`parser.input_tokens`), so `Harmonize()` is skipped and its reasoning effort is fixed at
+> parquet-build time (the perf parquet is supplied by the LLM task-force — see "Getting the Dataset").
+> MLPerf's perf phase uses `reasoning_effort=low`, so a perf-compliant run depends on that parquet being
+> built with `low`; it is not something this client YAML (or the repo's default `Harmonize()`, which
+> encodes `high`) can set. Treat this example as a demonstration of the per-dataset OSL split.
+
+```bash
+uv run inference-endpoint benchmark from-config \
+  -c examples/04_GPTOSS120B_Example/vllm_gptoss_120b_per_dataset_osl_example.yaml
+```
+
 ### vllm bench serve (Reference Comparison)
 
 `vllm bench serve` supports custom datasets only in `jsonl` format. To convert the parquet file:

examples/04_GPTOSS120B_Example/vllm_gptoss_120b_per_dataset_osl_example.yaml

@@ -0,0 +1,89 @@
+# gpt-oss-120b — MLPerf-style per-dataset output budgets in a single run.
+#
+# The MLPerf Inference gpt-oss-120b rules use DIFFERENT max output lengths for the
+# performance and accuracy phases:
+#   - performance: max_output_len = 10240  (reasoning_effort low)
+#   - accuracy:    max_output_len = 32768  (reasoning_effort high)
+# See https://github.com/mlcommons/inference/tree/master/language/gpt-oss-120b
+#
+# This config expresses that with the per-dataset `generation_config_override`
+# mechanism: model_params.max_new_tokens sets the global perf budget (10240) and the
+# accuracy datasets override it to 32768 so their harmony CoT is not truncated before
+# the `assistantfinal` answer. Only the OUTPUT budget is set here; reasoning_effort is
+# applied client-side by the `openai_completions` adapter's Harmony encoding — baked
+# into the pre-tokenized perf dataset, and fixed at "high" (the Harmonize default) for
+# the text-prompt accuracy datasets.
+name: "gpt-oss-120b-per-dataset-osl"
+version: "1.0"
+type: "online"
+timeout: 9000
+
+model_params:
+  name: "openai/gpt-oss-120b"
+  temperature: 1.0
+  top_p: 1.0
+  top_k: 0
+  skip_special_tokens: false
+  max_new_tokens: 10240 # global default = MLPerf perf budget; accuracy datasets override to 32768
+  streaming: "on"
+
+datasets:
+  - name: "perf-test" # PERF phase: capped at the MLPerf perf budget (10240)
+    type: "performance"
+    path: "examples/04_GPTOSS120B_Example/data/perf_eval_ref.parquet"
+    parser:
+      input_tokens: "input_tokens"
+    generation_config_override:
+      max_new_tokens: 10240 # same as the global default; stated per-dataset so both phases use the mechanism
+  - name: "aime25::gptoss" # ACCURACY: raised to the MLPerf accuracy budget (32768)
+    type: "accuracy"
+    generation_config_override:
+      max_new_tokens: 32768
+    accuracy_config:
+      eval_method: "pass_at_1"
+      ground_truth: "answer"
+      extractor: "boxed_math_extractor"
+      num_repeats: 8
+  - name: "gpqa::gptoss"
+    type: "accuracy"
+    generation_config_override:
+      max_new_tokens: 32768
+    accuracy_config:
+      eval_method: "pass_at_1"
+      extractor: "abcd_extractor"
+      ground_truth: "ground_truth"
+      num_repeats: 5
+  - name: "livecodebench::gptoss"
+    type: "accuracy"
+    generation_config_override:
+      max_new_tokens: 32768
+    accuracy_config:
+      eval_method: "code_bench_scorer"
+      extractor: "python_code_extractor"
+      num_repeats: 3
+
+settings:
+  runtime:
+    min_duration_ms: 30000
+    max_duration_ms: 14400000 # 4h whole-session deadline (perf + all accuracy phases)
+    scheduler_random_seed: 42
+    dataloader_random_seed: 42
+    n_samples_to_issue: 2000 # PERF phase only; accuracy phases issue their own sample counts
+
+  load_pattern:
+    type: "concurrency"
+    target_concurrency: 1024
+
+  client:
+    num_workers: 16
+    worker_initialization_timeout: 300.0
+    log_level: "WARN"
+    worker_gc_mode: "disabled"
+
+endpoint_config:
+  endpoints:
+    - "http://localhost:8000"
+  api_key: null
+  api_type: "openai_completions"
+
+report_dir: "results/vllm_gptoss_120b_per_dataset_osl/"

src/inference_endpoint/commands/benchmark/execute.py

@@ -333,9 +333,11 @@ def _load_datasets(
                 acc_cfg.accuracy_config.extras or {},
             )
         )
-        ds.load(
-            api_type=config.endpoint_config.api_type, model_params=config.model_params
-        )
+        # Value/api-type validity of the override is already enforced at config
+        # construction (BenchmarkConfig validates each dataset's effective params),
+        # so this cannot raise for a validated config.
+        ds_model_params = acc_cfg.effective_generation_config(config.model_params)
+        ds.load(api_type=config.endpoint_config.api_type, model_params=ds_model_params)
         logger.info(f"Loaded {ds} - {ds.num_samples()} samples")
 
     if not accuracy_cfgs:
@@ -349,16 +351,18 @@ def _load_datasets(
         if len(performance_cfgs) > 1:
             raise InputValidationError("Multiple performance datasets not supported")
         perf_cfg = performance_cfgs[0]
+        # Override validity is enforced at config construction (see accuracy loop).
+        perf_model_params = perf_cfg.effective_generation_config(config.model_params)
         try:
             dataloader = DataLoaderFactory.create_loader(perf_cfg)
             dataloader.load(
                 api_type=config.endpoint_config.api_type,
-                model_params=config.model_params,
+                model_params=perf_model_params,
             )
             logger.info(f"Loaded {dataloader.num_samples()} samples")
         except FileNotFoundError as e:
             raise InputValidationError(
-                f"Dataset file not found: {performance_cfgs[0].path}"
+                f"Dataset file not found: {perf_cfg.path}"
             ) from e
         except Exception as e:
             raise SetupError(f"Failed to load dataset: {e}") from e

src/inference_endpoint/config/schema.py

@@ -59,6 +59,45 @@ class SystemDefaults(BaseModel):
     DEFAULT_METRIC: ClassVar[metrics.Metric] = metrics.Throughput(0.0)
 
 
+def _deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
+    """Recursively merge ``override`` into ``base`` and return the result.
+
+    For overlapping keys whose values are both dicts, recurse; otherwise the
+    override value wins. Mutates a *copy* — callers can safely pass model_dump()
+    output. Used by ``Dataset.effective_generation_config`` so a sparse nested
+    override (e.g. ``{osl_distribution: {max: 512}}``) preserves siblings.
+    """
+    out = dict(base)
+    for k, v in override.items():
+        if isinstance(v, dict) and isinstance(out.get(k), dict):
+            out[k] = _deep_merge(out[k], v)
+        else:
+            out[k] = v
+    return out
+
+
+# ModelParams fields that drive the single global tokenizer / MetricsAggregator
+# (launched once from top-level model_params), so a per-dataset override would
+# desync ISL/OSL/TTFT/TPOT accounting without changing what is measured. Rejected
+# as generation_config_override keys — they are per-run/identity, not per-dataset.
+_METRICS_DECOUPLED_OVERRIDE_KEYS = frozenset({"name", "streaming", "tokenizer_name"})
+
+
+def _non_default_completion_controls(mp: ModelParams) -> list[str]:
+    """Completion-only ModelParams controls set to a non-default value.
+
+    ``min_new_tokens``/``skip_special_tokens`` are only honored by the
+    ``openai_completions`` adapter; ``BenchmarkConfig`` rejects them for other
+    ``api_type``s. Shared by the top-level and per-dataset-override checks so
+    both config surfaces validate identically.
+    """
+    checks = {
+        "min_new_tokens": mp.min_new_tokens != 1,
+        "skip_special_tokens": not mp.skip_special_tokens,
+    }
+    return [name for name, non_default in checks.items() if non_default]
+
+
 class LoadPatternType(str, Enum):
     """Load pattern types."""
 
@@ -430,6 +469,36 @@ class Dataset(BaseModel):
     agentic_inference: AgenticInferenceConfig | None = Field(
         None, description="Agentic inference conversation configuration"
     )
+    # Per-dataset generation config is a first-class capability: different
+    # accuracy datasets legitimately want different generation settings (e.g.
+    # per-dataset max OSL or top_p, as seen in DS-V4), and dataset-scoping also
+    # enables per-dataset dynamic OSL distributions. Only generation knobs are
+    # overridable — per-run/identity fields (`_METRICS_DECOUPLED_OVERRIDE_KEYS`:
+    # name / streaming / tokenizer_name) drive the single global tokenizer and
+    # MetricsAggregator, so overriding them per-dataset would desync ISL/OSL/
+    # TTFT/TPOT accounting; they are rejected at validation.
+    #
+    # TODO(post-mortem): split ModelParams into a per-run ModelIdentity and a
+    # GenerationConfig, so the override surface is exactly the generation fields
+    # and identity fields cannot be named here at all. Field/method names use
+    # "generation_config" to keep that migration mechanical.
+    #
+    # Nested dicts (`osl_distribution`, `chat_template_kwargs`) are deep-merged
+    # so sparse overrides preserve sibling defaults.
+    generation_config_override: dict[str, Any] | None = Field(
+        None,
+        description=(
+            "Per-dataset overrides for the top-level model_params (sparse — "
+            "only the fields you want to override). Merged on top of "
+            "BenchmarkConfig.model_params at dataset-load time. Useful for "
+            "MLPerf-style runs where accuracy and performance use different "
+            "output budgets in the same fleet, e.g. "
+            "generation_config_override: {max_new_tokens: 32768, "
+            "temperature: 0.0}. NOTE: per-run/identity keys (`name`, "
+            "`streaming`, `tokenizer_name`) are rejected here — set them on "
+            "top-level model_params."
+        ),
+    )
 
     @model_validator(mode="after")
     def _auto_derive_name(self) -> Self:
@@ -438,6 +507,53 @@ def _auto_derive_name(self) -> Self:
             object.__setattr__(self, "name", Path(self.path).stem)
         return self
 
+    @model_validator(mode="after")
+    def _validate_generation_config_override(self) -> Self:
+        """Fail fast on unknown keys and on per-run/identity keys the single
+        global tokenizer / MetricsAggregator would ignore. Override *values*
+        are validated at merge time (see ``effective_generation_config``)
+        because cross-field validation needs the base ``ModelParams`` from
+        ``BenchmarkConfig``.
+        """
+        if self.generation_config_override:
+            keys = set(self.generation_config_override)
+            valid = set(ModelParams.model_fields)
+            bad = sorted(keys - valid)
+            if bad:
+                raise ValueError(
+                    f"Dataset '{self.name}': unknown keys in "
+                    f"generation_config_override: {bad}. "
+                    f"Valid keys: {sorted(valid)}"
+                )
+            decoupled = sorted(keys & _METRICS_DECOUPLED_OVERRIDE_KEYS)
+            if decoupled:
+                raise ValueError(
+                    f"Dataset '{self.name}': generation_config_override keys "
+                    f"{decoupled} are not honored per-dataset — the single "
+                    "global tokenizer / metrics aggregator is launched from "
+                    "top-level model_params, so a per-dataset value would "
+                    "desync ISL/OSL/TTFT/TPOT accounting. Set them on "
+                    "top-level model_params instead."
+                )
+        return self
+
+    def effective_generation_config(self, base: ModelParams) -> ModelParams:
+        """Return base merged with this dataset's generation-config overrides.
+
+        Nested dicts are deep-merged so a sparse nested override preserves
+        sibling defaults (e.g. ``{osl_distribution: {max: 512}}`` keeps the
+        base ``type/mean/std/min``). The merged dict is re-validated through
+        ``ModelParams.model_validate`` so type-invalid scalar overrides (e.g.
+        ``temperature: 'hot'``) are rejected. Note that this only catches
+        scalar invalidity — a sparse nested override whose merged result
+        passes default-validation will not raise (callers that need stricter
+        nested validation should set ``base`` to an explicit instance).
+        """
+        if not self.generation_config_override:
+            return base
+        merged = _deep_merge(base.model_dump(), self.generation_config_override)
+        return ModelParams.model_validate(merged)
+
 
 class AccuracyConfig(BaseModel):
     """Accuracy configuration.
@@ -988,30 +1104,47 @@ def _resolve_and_validate(self) -> Self:
 
         # TODO(vir): Move API-type-specific validation out of this generic
         # cross-model validator and into the selected adapter. Requires a larger refactor.
-        non_default_completion_controls = {
-            "model_params.min_new_tokens": self.model_params.min_new_tokens != 1,
-            "model_params.skip_special_tokens": not self.model_params.skip_special_tokens,
+        #
+        # Completion-only controls must be gated by api_type for BOTH the
+        # top-level model_params AND every per-dataset generation_config_override,
+        # so the two config surfaces validate identically. Merge each dataset's
+        # effective params once here (parse time) — this also surfaces
+        # value-invalid overrides before setup produces side effects — and reuse
+        # the result for the agentic-inference check below.
+        effective_by_dataset: dict[int, ModelParams] = {
+            id(dataset): dataset.effective_generation_config(self.model_params)
+            for dataset in self.datasets
+            if dataset.generation_config_override
         }
-        configured_controls = [
-            name
-            for name, is_non_default in non_default_completion_controls.items()
-            if is_non_default
+        completion_control_surfaces: list[tuple[str, ModelParams]] = [
+            ("model_params", self.model_params)
         ]
-        if (
-            configured_controls
-            and self.endpoint_config.api_type != APIType.OPENAI_COMPLETIONS
-        ):
-            controls = " and ".join(configured_controls)
-            verb = "requires" if len(configured_controls) == 1 else "require"
-            required_api_type = "endpoint_config.api_type=openai_completions"
-            raise ValueError(f"{controls} {verb} {required_api_type}")
-        if configured_controls and any(
-            dataset.agentic_inference is not None for dataset in self.datasets
-        ):
-            raise ValueError(
-                "OpenAI text-completion generation controls are not supported "
-                "for agentic inference datasets"
-            )
+        for dataset in self.datasets:
+            effective = effective_by_dataset.get(id(dataset))
+            if effective is not None:
+                completion_control_surfaces.append(
+                    (
+                        f"datasets['{dataset.name}'].generation_config_override",
+                        effective,
+                    )
+                )
+        for prefix, mp in completion_control_surfaces:
+            controls = _non_default_completion_controls(mp)
+            if controls and self.endpoint_config.api_type != APIType.OPENAI_COMPLETIONS:
+                names = " and ".join(f"{prefix}.{name}" for name in controls)
+                verb = "requires" if len(controls) == 1 else "require"
+                raise ValueError(
+                    f"{names} {verb} endpoint_config.api_type=openai_completions"
+                )
+        for dataset in self.datasets:
+            if dataset.agentic_inference is None:
+                continue
+            effective = effective_by_dataset.get(id(dataset), self.model_params)
+            if _non_default_completion_controls(effective):
+                raise ValueError(
+                    "OpenAI text-completion generation controls are not supported "
+                    "for agentic inference datasets"
+                )
 
         # --- Validate (cross-model checks only; sub-models self-validate) ---
         if self.type == TestType.SUBMISSION and not self.benchmark_mode:

src/inference_endpoint/config/templates/concurrency_template_full.yaml

@@ -31,6 +31,7 @@ datasets:  # Dataset configs
   generate_params: null  # Dataset-specific parameters passed to the generate() method
   accuracy_config: null  # Accuracy evaluation settings
   agentic_inference: null  # Agentic inference conversation configuration
+  generation_config_override: null  # Per-dataset overrides for the top-level model_params (sparse — only the fields you want to override). Merged on top of BenchmarkConfig.model_params at dataset-load time. Useful for MLPerf-style runs where accuracy and performance use different output budgets in the same fleet, e.g. generation_config_override: {max_new_tokens: 32768, temperature: 0.0}. NOTE: per-run/identity keys (`name`, `streaming`, `tokenizer_name`) are rejected here — set them on top-level model_params.
 - name: accuracy
   type: accuracy  # Dataset purpose: performance or accuracy | options: performance, accuracy
   path: '<DATASET_PATH eg: tests/assets/datasets/ds_samples.jsonl>'  # Dataset file path
@@ -47,6 +48,7 @@ datasets:  # Dataset configs
     extractor: boxed_math_extractor  # Answer extractor (abcd_extractor, boxed_math_extractor, identity_extractor, python_code_extractor)
     num_repeats: 1  # Repeat dataset N times for evaluation
   agentic_inference: null  # Agentic inference conversation configuration
+  generation_config_override: null  # Per-dataset overrides for the top-level model_params (sparse — only the fields you want to override). Merged on top of BenchmarkConfig.model_params at dataset-load time. Useful for MLPerf-style runs where accuracy and performance use different output budgets in the same fleet, e.g. generation_config_override: {max_new_tokens: 32768, temperature: 0.0}. NOTE: per-run/identity keys (`name`, `streaming`, `tokenizer_name`) are rejected here — set them on top-level model_params.
 settings:
   runtime:
     min_duration_ms: 600000  # Min duration (ms, or with suffix: 600s, 10m)