From f1eacc22537134111ff86a28e030d88c2c3fb29d Mon Sep 17 00:00:00 2001 From: Matthew Romano Date: Fri, 11 Sep 2026 11:02:02 -0400 Subject: [PATCH] fix(llama): preserve full eos_token_id list through to the runtime HF configs may declare eos_token_id as a list of stop tokens (Llama 3.1+, MiniCPM5). The llama family kept only the first id, which breaks stopping for Llama 3.1-Instruct since its real per-turn stop token is the *last* id in the list, not the first. Also exclude the new minicpm5-2b benchmark manifest from the release performance suite (functional/e2e qualification is present, but no matching release-performance workload or receipt exists yet). Signed-off-by: Matthew Romano --- apps/benchmark/performance/release.yaml | 5 ++ families/llama/config.py | 20 +++++- families/llama/model.py | 7 ++- families/llama/runtime/plugin.cpp | 9 ++- .../llama/tests/manifests/minicpm5-2b.json | 21 +++++++ families/llama/tests/test_config.py | 37 +++++++++++ families/llama/tests/test_runtime_config.py | 63 +++++++++++++++++++ 7 files changed, 155 insertions(+), 7 deletions(-) create mode 100644 families/llama/tests/manifests/minicpm5-2b.json create mode 100644 families/llama/tests/test_config.py create mode 100644 families/llama/tests/test_runtime_config.py diff --git a/apps/benchmark/performance/release.yaml b/apps/benchmark/performance/release.yaml index da6b5118d9..ab07275790 100644 --- a/apps/benchmark/performance/release.yaml +++ b/apps/benchmark/performance/release.yaml @@ -105,6 +105,11 @@ excluded_profiles: nano width, which is the only one whose class-attention norm covers a single token, but the release-performance workload and receipt were collected only for xcit_tiny_12_p16_224. + - model: minicpm5-2b + reason: >- + Functional and HF reference-parity e2e qualification is present, but + this change does not add a matching release-performance workload or + receipt. entries: - id: detr.detect diff --git a/families/llama/config.py b/families/llama/config.py index 7af3570d36..cf570fb218 100644 --- a/families/llama/config.py +++ b/families/llama/config.py @@ -10,6 +10,20 @@ from pathlib import Path +def _as_id_list(value: object) -> tuple[int, ...]: + """Normalize an HF eos_token_id-shaped value (missing, a single id, or + a list of ids — e.g. Llama 3.1+, MiniCPM5) into a tuple of ints. + + Checking for None (rather than falsiness) matters here: a legitimate + id of 0 must survive, unlike a missing value. + """ + if value is None: + return () + if isinstance(value, (list, tuple)): + return tuple(int(v) for v in value) + return (int(value),) + + @dataclass class ModelConfig: """Parsed model architecture from HF config.json.""" @@ -26,6 +40,7 @@ class ModelConfig: rope_theta: float = 10000.0 bos_token_id: int = -1 eos_token_id: int = -1 + eos_token_ids: tuple[int, ...] = () pad_token_id: int = -1 tie_word_embeddings: bool = False max_position_embeddings: int = 8192 @@ -173,6 +188,8 @@ def from_json(text: str) -> ModelConfig: if not architectures and architecture: architectures = [architecture] + eos_token_ids = _as_id_list(d.get("eos_token_id")) + return ModelConfig( model_type=d.get("model_type", "") or architecture, architectures=architectures, @@ -185,7 +202,8 @@ def from_json(text: str) -> ModelConfig: rms_norm_eps=eps, rope_theta=rope_theta, bos_token_id=d.get("bos_token_id", -1) or -1, - eos_token_id=d.get("eos_token_id", -1) or -1, + eos_token_id=eos_token_ids[0] if eos_token_ids else -1, + eos_token_ids=eos_token_ids, pad_token_id=d.get("pad_token_id", -1) or -1, tie_word_embeddings=d.get("tie_word_embeddings", False), max_position_embeddings=d.get("max_position_embeddings", diff --git a/families/llama/model.py b/families/llama/model.py index 706039da11..d7dbe186c5 100644 --- a/families/llama/model.py +++ b/families/llama/model.py @@ -11,7 +11,7 @@ from .build_routing import native_kv_architecture_capability, native_kv_build_capability from .checkpoint_mapper import WeightDict, load_standard_weights -from .config import ModelConfig +from .config import ModelConfig, _as_id_list from .dual_profile_decoder_builder import build_dual_profile_decoder_engine from .native_kv_contract import validate_native_kv_weights from .standard_decoder_builder import build_standard_decoder_engine @@ -98,7 +98,7 @@ def _runtime_config(model_dir: Path, config: ModelConfig, **updates) -> dict: "num_key_value_heads": config.num_key_value_heads, "head_dim": config.head_dim, "bos_token_id": config.bos_token_id, - "eos_token_id": config.eos_token_id, + "eos_token_id": list(config.eos_token_ids) or [-1], "pad_token_id": config.pad_token_id, } runtime.update(config.raw.get("_native_kv_cache_metadata", {})) @@ -108,7 +108,8 @@ def _runtime_config(model_dir: Path, config: ModelConfig, **updates) -> dict: if not isinstance(generation, dict): raise ValueError("generation_config.json must contain one JSON object") if "eos_token_id" in generation: - runtime["eos_token_id"] = generation["eos_token_id"] + eos_ids = _as_id_list(generation["eos_token_id"]) + runtime["eos_token_id"] = list(eos_ids) if eos_ids else [-1] runtime.update(updates) return runtime diff --git a/families/llama/runtime/plugin.cpp b/families/llama/runtime/plugin.cpp index 9a4d3f4408..ae25aa73d7 100644 --- a/families/llama/runtime/plugin.cpp +++ b/families/llama/runtime/plugin.cpp @@ -32,7 +32,7 @@ struct RuntimeConfig { std::int32_t head_dim; std::int32_t vocab_size; std::int32_t bos_token_id; - std::int32_t eos_token_id; + std::vector eos_token_ids; std::int32_t pad_token_id; std::int32_t max_cache_length; std::string precision; @@ -84,13 +84,16 @@ RuntimeConfig parse_runtime_config(const BundleReader& bundle) { require_value(json, "head_dim"), require_value(json, "vocab_size"), require_value(json, "bos_token_id"), - require_value(json, "eos_token_id"), + require_value>(json, "eos_token_id"), require_value(json, "pad_token_id"), require_value(json, "max_cache_length"), require_value(json, "precision"), require_value(json, "decoder_engine_layout"), dynamic_kv_cache, }; + if (config.eos_token_ids.empty()) { + throw std::runtime_error("llama runtime.json has an empty 'eos_token_id' list"); + } if (config.hidden_size <= 0 || config.num_layers <= 0 || config.num_heads <= 0 || config.num_key_value_heads <= 0 || config.head_dim <= 0 || config.vocab_size <= 0 || config.max_cache_length <= 0 || @@ -225,7 +228,7 @@ ITask* create(const FamilyContext& context) { LlamaTextGenConfig text_config; text_config.vocab_size = config.vocab_size; text_config.id_bos = config.bos_token_id; - text_config.id_eos = config.eos_token_id; + text_config.id_eos_ids = config.eos_token_ids; text_config.chat_template_format = llama_detect_chat_template_format(chat_template(context.reader)); text_config.prefill_max_length = prefill_token_limit(*modules.prefill); diff --git a/families/llama/tests/manifests/minicpm5-2b.json b/families/llama/tests/manifests/minicpm5-2b.json new file mode 100644 index 0000000000..1701ab50bd --- /dev/null +++ b/families/llama/tests/manifests/minicpm5-2b.json @@ -0,0 +1,21 @@ +{ + "name": "minicpm5-2b", + "hf_id": "openbmb/MiniCPM5-2B", + "bundle": "minicpm5-2b.bundle", + "family": "llama", + "task": "text_generation", + "precision": "bf16", + "trust_remote_code": false, + "testcases": [ + { + "name": "minicpm5-2b", + "reference_precision": "fp32", + "prompt": "What is the capital of France? Answer in one word.", + "max_new_tokens": 20, + "use_chat_template": true, + "enable_thinking": false + } + ], + "max_sequence_length": 512, + "tensor_parallel_size": 1 +} diff --git a/families/llama/tests/test_config.py b/families/llama/tests/test_config.py new file mode 100644 index 0000000000..8e8df4ce57 --- /dev/null +++ b/families/llama/tests/test_config.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from families.llama.config import ModelConfig + + +def test_scalar_eos_token_id() -> None: + config = ModelConfig.create_tiny(model_type="llama", eos_token_id=1) + + assert config.eos_token_id == 1 + assert config.eos_token_ids == (1,) + + +def test_list_eos_token_id_preserves_every_id() -> None: + # HF configs allow eos_token_id to be a list of stop tokens (e.g. Llama + # 3.1+, MiniCPM5). Every id must reach the runtime bundle: the model may + # naturally emit any one of them as its real per-turn stop token (for + # Llama 3.1-Instruct, that's <|eot_id|>, the *last* id in the list, not + # the first), so truncating to one id would silently break stopping. + config = ModelConfig.create_tiny(model_type="llama", eos_token_id=[1, 130073]) + + assert config.eos_token_id == 1 + assert config.eos_token_ids == (1, 130073) + + +def test_empty_eos_token_id_list_falls_back() -> None: + config = ModelConfig.create_tiny(model_type="llama", eos_token_id=[]) + + assert config.eos_token_id == -1 + assert config.eos_token_ids == () + + +def test_missing_eos_token_id_falls_back() -> None: + config = ModelConfig.create_tiny(model_type="llama") + + assert config.eos_token_id == -1 + assert config.eos_token_ids == () diff --git a/families/llama/tests/test_runtime_config.py b/families/llama/tests/test_runtime_config.py new file mode 100644 index 0000000000..7a2102136d --- /dev/null +++ b/families/llama/tests/test_runtime_config.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""_runtime_config()'s generation_config.json override path. + +Importing families.llama.model drags in the TensorRT import chain (via +standard_decoder_builder.py), so this lives separately from the +TensorRT-free families.llama.tests.test_config. +""" + +from __future__ import annotations + +import json + +import pytest + +trt = pytest.importorskip("tensorrt") + +from ..config import ModelConfig # noqa: E402 +from ..model import _runtime_config # noqa: E402 + + +pytestmark = [pytest.mark.gpu, pytest.mark.trt] + + +def test_generation_config_overrides_with_full_eos_list(tmp_path) -> None: + # The checkpoint's own config.json disagrees with generation_config.json + # (as real checkpoints do) so this only passes if the override path is + # actually exercised, not the config.json path. + config = ModelConfig.create_tiny(model_type="llama", eos_token_id=1) + (tmp_path / "generation_config.json").write_text( + json.dumps({"eos_token_id": [128001, 128008, 128009]}), + encoding="utf-8", + ) + + runtime = _runtime_config(tmp_path, config, precision="bf16", max_cache_length=128, + decoder_engine_layout="split") + + assert runtime["eos_token_id"] == [128001, 128008, 128009] + + +def test_generation_config_empty_eos_list_falls_back(tmp_path) -> None: + # An empty override list must not reach the bundle as `null` or `[]` — + # the C++ loader rejects an empty eos_token_id list. + config = ModelConfig.create_tiny(model_type="llama", eos_token_id=1) + (tmp_path / "generation_config.json").write_text( + json.dumps({"eos_token_id": []}), + encoding="utf-8", + ) + + runtime = _runtime_config(tmp_path, config, precision="bf16", max_cache_length=128, + decoder_engine_layout="split") + + assert runtime["eos_token_id"] == [-1] + + +def test_no_generation_config_keeps_config_json_eos_list(tmp_path) -> None: + config = ModelConfig.create_tiny(model_type="llama", eos_token_id=[1, 130073]) + + runtime = _runtime_config(tmp_path, config, precision="bf16", max_cache_length=128, + decoder_engine_layout="split") + + assert runtime["eos_token_id"] == [1, 130073]