feat: per-dataset max_new_tokens override

[roborluo]

What does this PR do?

When running a combined performance + accuracy benchmark in a single --mode both invocation, the two phases want opposite generation caps, but today the harness only exposes one global model_params.max_new_tokens:

• Performance phase needs a small cap. max_new_tokens is sent to the server as the per-request max_tokens, and a disaggregated decode scheduler reserves/plans decode-KV for that declared upper bound — even though generation actually stops at EOS far sooner. A large cap (e.g. 32768) over-reserves decode KV (~3.2× vs 10240), starves admittable decode slots at high concurrency, and triggers KV-transfer-timeout storms on the context→gen path. A realistic small cap avoids this.
• Accuracy phase needs a large cap, otherwise long reasoning outputs get truncated and scores are artificially deflated. This matches the MLPerf Inference gpt-oss-120b reference, where the performance and accuracy workloads use different token settings — see language/gpt-oss-120b → Model and Dataset download. Without a per-dataset override, you cannot satisfy both in one --mode both run.

Type of change

• Bug fix
• New feature
• Documentation update
• Refactor/cleanup

Related issues

N/A

Testing

• Tests added/updated
• All tests pass locally
• Manual testing completed

Checklist

• Code follows project style
• Pre-commit hooks pass
• Documentation updated (if needed)

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[gemini-code-assist]

Code Review

This pull request introduces a per-dataset max_new_tokens override capability to allow performance and accuracy datasets to use different token limits, falling back to the global model_params when unset. The feedback suggests encapsulating the override logic into a helper method get_model_params on the Dataset class to eliminate code duplication across the accuracy and performance dataset loading paths, and adding corresponding unit tests for this helper.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Use the new get_model_params helper method on the Dataset configuration model to simplify the override logic and eliminate duplication.

Suggested change

	# Per-dataset max_new_tokens override (falls back to global model_params).
	acc_model_params = (
	config.model_params
	if acc_cfg.max_new_tokens is None
	else config.model_params.model_copy(
	update={"max_new_tokens": acc_cfg.max_new_tokens}
	)
	)
	acc_model_params = acc_cfg.get_model_params(config.model_params)

[gemini-code-assist]

Use the new get_model_params helper method on the Dataset configuration model to simplify the override logic and eliminate duplication.

Suggested change

	# Per-dataset max_new_tokens override (falls back to global model_params).
	perf_model_params = (
	config.model_params
	if perf_cfg.max_new_tokens is None
	else config.model_params.model_copy(
	update={"max_new_tokens": perf_cfg.max_new_tokens}
	)
	)
	perf_model_params = perf_cfg.get_model_params(config.model_params)

[arekay-nv]

@roborluo Can you look at #344 which addresses the same issue. We can consolidate the two here and merge this one. I think the other one has also the templates correctly populated which you are failing in CI.