diff --git a/python/tensorrt_model_connect/families/locateanything/plugin.py b/python/tensorrt_model_connect/families/locateanything/plugin.py index 90c2419346..76b1425c6b 100644 --- a/python/tensorrt_model_connect/families/locateanything/plugin.py +++ b/python/tensorrt_model_connect/families/locateanything/plugin.py @@ -170,6 +170,14 @@ def get_bundle_config_overrides(self, config: ModelConfig) -> dict | None: "model_type": "locateanything", "runtime_strategy": "locateanything_vision_language", "embed_input": True, + "vocab_size": config.vocab_size, + "hidden_size": config.hidden_size, + "num_hidden_layers": config.num_hidden_layers, + "num_attention_heads": config.num_attention_heads, + "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, } diff --git a/src/runtime/models/locateanything/MODEL.toml b/src/runtime/models/locateanything/MODEL.toml index 64ddf15b3c..8e15770e9f 100644 --- a/src/runtime/models/locateanything/MODEL.toml +++ b/src/runtime/models/locateanything/MODEL.toml @@ -6,5 +6,6 @@ runtime_library = "libtrtmc_model_locateanything.so" runtime_plugins = ["plugin.cpp|register_locateanything_plugin"] runtime_strategies = ["locateanything_vision_language"] runtime_tests = [ - "test_locateanything_vl_pipeline|test_locateanything_vl_pipeline.cpp|trtmc_model_locateanything,trtmc_backend_trt|_|REQUIRES_TRT", + "test_locateanything_runtime_config_contract|test_locateanything_runtime_config_contract.cpp|_|_|_", + "test_locateanything_vl_pipeline|test_locateanything_vl_pipeline.cpp|trtmc_model_locateanything,trtmc_backend_trt|_|REQUIRES_TRT,REQUIRES_GPU", ] diff --git a/tests/cpp/models/locateanything/test_locateanything_runtime_config_contract.cpp b/tests/cpp/models/locateanything/test_locateanything_runtime_config_contract.cpp new file mode 100644 index 0000000000..5703c9ae0e --- /dev/null +++ b/tests/cpp/models/locateanything/test_locateanything_runtime_config_contract.cpp @@ -0,0 +1,81 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// CPU-only consumer contract for LocateAnything's serialized runtime config. + +#include "trtmc/runtime/pipeline_plugin.h" + +#include +#include +#include +#include + +int main() { + // The nested objects mirror nvidia/LocateAnything-3B's Hugging Face config. + // The duplicated top-level decoder fields are the final bundle contract + // consumed by the strict production C++ parser. + const std::string config = R"({ + "vocab_size": 152681, + "hidden_size": 2048, + "num_hidden_layers": 36, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "head_dim": 128, + "bos_token_id": 151643, + "eos_token_id": 151645, + "model_type": "locateanything", + "text_config": { + "model_type": "qwen2", + "vocab_size": 152681, + "hidden_size": 2048, + "intermediate_size": 11008, + "num_hidden_layers": 36, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "bos_token_id": 151643, + "eos_token_id": 151645, + "max_position_embeddings": 32768, + "rope_theta": 1000000.0 + }, + "vision_config": { + "model_type": "moonvit", + "hidden_size": 1152, + "intermediate_size": 4304, + "num_hidden_layers": 27, + "num_attention_heads": 16, + "patch_size": 14, + "merge_kernel_size": [2, 2] + }, + "runtime_strategy": "locateanything_vision_language", + "precision": "fp16", + "tokenizer_add_special_tokens": 0 + })"; + + const auto parsed = trtmc::parse_base_config(config, 384); + bool ok = true; + const auto check = [&](bool condition, const char* name) { + if (!condition) { + std::cerr << "FAIL: " << name << '\n'; + ok = false; + } + }; + + check(parsed.runtime_strategy == "locateanything_vision_language", "runtime strategy"); + check(parsed.precision == "fp16", "precision"); + check(parsed.vocab_size == 152681, "vocabulary size"); + check(parsed.hidden_size == 2048, "hidden size"); + check(parsed.num_layers == 36, "layer count"); + check(parsed.num_heads == 16, "attention head count"); + check(parsed.num_kv_heads == 2, "KV head count"); + check(parsed.head_dim == 128, "head dimension"); + check(parsed.attention_size == 2048, "attention width"); + check(parsed.max_cache_length == 384, "cache length override"); + check(parsed.id_bos == 151643, "BOS token"); + check(parsed.id_eos == 151645, "EOS token"); + check(parsed.id_eos_ids == std::vector({151645}), "EOS token list"); + check(parsed.tokenizer_add_special_tokens_present, "tokenizer flag presence"); + check(!parsed.tokenizer_add_special_tokens, "tokenizer flag value"); + return ok ? 0 : 1; +} diff --git a/tests/e2e/models/locateanything/test_locateanything_family_plugin_weights.py b/tests/e2e/models/locateanything/test_locateanything_family_plugin_weights.py index c9a80a586e..eca5a155bd 100644 --- a/tests/e2e/models/locateanything/test_locateanything_family_plugin_weights.py +++ b/tests/e2e/models/locateanything/test_locateanything_family_plugin_weights.py @@ -10,6 +10,8 @@ from __future__ import annotations import importlib +import json +from unittest.mock import patch from tests.builder.family_plugin_test_support import ( ModelConfig, @@ -82,6 +84,8 @@ def _write_locateanything_config(self, tmp_path): "num_hidden_layers": self.LAYERS, "num_attention_heads": self.HEADS, "num_key_value_heads": self.KV_HEADS, + "bos_token_id": 151643, + "eos_token_id": 151645, "rms_norm_eps": 1e-6, "rope_theta": 1000000.0, "max_position_embeddings": 32768, @@ -162,3 +166,81 @@ def test_get_vl_config(self, tmp_path): assert vl_cfg["box_start_token_id"] == 151668 assert "{image_pads}" in vl_cfg["vl_prompt_template"] assert "{prompt}" in vl_cfg["vl_prompt_template"] + + def test_mock_bundle_serializes_decoder_geometry_at_top_level(self, tmp_path): + """Mock engines while exercising the real LocateAnything bundle boundary.""" + from tensorrt_model_connect.engine_builder import build_bundle + from tensorrt_model_connect.families.locateanything.plugin import ( + plugin as production_plugin, + ) + + self._write_locateanything_config(tmp_path) + source_config = json.loads( + (tmp_path / "config.json").read_text(encoding="utf-8") + ) + + class MockLocateAnythingPlugin: + name = "locateanything" + runtime_strategy = "locateanything_vision_language" + embed_input = True + requires_tokenizer = False + + @staticmethod + def load_weights(_model_dir, _config): + return {} + + @staticmethod + def build_engine(_config, _weights, _max_cache_length, **_kwargs): + return b"MOCK_DECODER_PLAN" + + @staticmethod + def build_vision_engine(_model_dir, _config, _weights, **_kwargs): + return b"MOCK_VISION_PLAN" + + @staticmethod + def get_vl_config(config): + return production_plugin.get_vl_config(config) + + @staticmethod + def get_bundle_config_overrides(config): + return production_plugin.get_bundle_config_overrides(config) + + with ( + patch( + "tensorrt_model_connect.engine_builder.find_plugin", + return_value=MockLocateAnythingPlugin(), + ), + patch( + "tensorrt_model_connect.engine_builder._get_trt_version", + return_value="11.1.0", + ), + patch( + "tensorrt_model_connect.engine_builder._get_gpu_name", + return_value="CPU unit mock", + ), + patch("tensorrt_model_connect.engine_builder.write_bundle") as write_bundle, + ): + build_bundle( + str(tmp_path), + str(tmp_path / "locateanything.bundle"), + max_cache_length=384, + ) + + sections = { + section.name: section.data for section in write_bundle.call_args.args[2] + } + runtime_config = json.loads(sections["config.json"]) + decoder_config = source_config["text_config"] + assert runtime_config["text_config"] == decoder_config + decoder_contract = { + "vocab_size": self.VOCAB, + "hidden_size": self.HIDDEN, + "num_hidden_layers": self.LAYERS, + "num_attention_heads": self.HEADS, + "num_key_value_heads": self.KV_HEADS, + "head_dim": self.HIDDEN // self.HEADS, + "bos_token_id": 151643, + "eos_token_id": 151645, + } + assert all(key not in source_config for key in decoder_contract) + assert {key: runtime_config[key] for key in decoder_contract} == decoder_contract