diff --git a/apps/aiperf_qual/config/tasks.yaml b/apps/aiperf_qual/config/tasks.yaml index 69234ed6a1..b6e2276f7f 100644 --- a/apps/aiperf_qual/config/tasks.yaml +++ b/apps/aiperf_qual/config/tasks.yaml @@ -135,9 +135,10 @@ benchmarks: min_new_tokens: 64, gate: {margin: 2.0}} tasks: - text_generation: + text_generation: &generate_text absolute: [mmlu] performance: {output_grader: parity_token_exact, output_grader_params: {min_prefix: 8, accept_equal_text: true}} + conditional_text_generation: *generate_text vision_language_generation: # OCRBench alone (the owner's choice, October 2026): it caught both VLM conversion defects MMStar caught diff --git a/families/bart/model.py b/families/bart/model.py index e0bc8f4e0a..3a18ef3e13 100644 --- a/families/bart/model.py +++ b/families/bart/model.py @@ -944,8 +944,8 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: if request.context_parallel_size != 1: raise ValueError("this family does not support context parallelism") - if request.task != "text_generation": - raise ValueError("bart supports only task=text_generation") + if request.task not in ("conditional_text_generation", "text_generation"): + raise ValueError("bart supports only task=conditional_text_generation or text_generation") model_dir = Path(request.model_dir) config = ModelConfig.from_dir(model_dir) diff --git a/families/bart/runtime/CMakeLists.txt b/families/bart/runtime/CMakeLists.txt index 8512aa9e84..1471e58399 100644 --- a/families/bart/runtime/CMakeLists.txt +++ b/families/bart/runtime/CMakeLists.txt @@ -6,6 +6,7 @@ add_library(trtmc_model_bart SHARED decode_runtime.cpp distributed_runtime.cpp kv_cache.cpp + pipeline.cpp plugin.cpp plugin_helpers.cpp ) @@ -74,4 +75,43 @@ if(TRTMC_BUILD_TESTS) -Wall -Wextra -Wpedantic ) add_test(NAME bart_cross_kv_alloc COMMAND test_bart_cross_kv_alloc) + + add_executable(test_bart_task_contract + ${PROJECT_SOURCE_DIR}/families/bart/tests/cpp/test_task_contract.cpp + pipeline.cpp + decode_runtime.cpp + ) + target_include_directories(test_bart_task_contract PRIVATE + ${PROJECT_SOURCE_DIR} + ${PROJECT_SOURCE_DIR}/core/runtime/include + ) + target_include_directories(test_bart_task_contract SYSTEM PRIVATE + ${TRTMC_CUDA_INCLUDE_DIR} + ) + target_link_libraries(test_bart_task_contract PRIVATE + trtmc_core + nlohmann_json::nlohmann_json + ) + target_compile_options(test_bart_task_contract PRIVATE + -Wall -Wextra -Wpedantic -Werror + ) + add_test(NAME bart_task_contract COMMAND test_bart_task_contract) + set_tests_properties(bart_task_contract PROPERTIES LABELS "cpu") + + foreach(_language IN ITEMS c cpp) + add_executable(test_bart_sdk_${_language} + ${PROJECT_SOURCE_DIR}/families/bart/tests/sdk_consumer.${_language} + ) + target_link_libraries(test_bart_sdk_${_language} PRIVATE trtmc_c) + target_compile_options(test_bart_sdk_${_language} PRIVATE + -Wall -Wextra -Wpedantic -Werror + ) + set_target_properties(test_bart_sdk_${_language} PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}" + ) + endforeach() + + add_dependencies(test_bart_task_contract + test_bart_sdk_c test_bart_sdk_cpp + ) endif() diff --git a/families/bart/runtime/pipeline.cpp b/families/bart/runtime/pipeline.cpp new file mode 100644 index 0000000000..1cf0c2cd7d --- /dev/null +++ b/families/bart/runtime/pipeline.cpp @@ -0,0 +1,254 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "families/bart/runtime/pipeline.h" + +#include "families/bart/runtime/decode_runtime.h" + +#include +#include +#include +#include + +namespace trtmc { + +BartPipeline::BartPipeline(std::unique_ptr encoder, std::unique_ptr decoder, + std::unique_ptr state, int32_t hidden_size, + int32_t num_decoder_layers, int32_t max_source_length, + int32_t decoder_start_token_id, int32_t eos_token_id, + int32_t bos_token_id, int32_t pad_token_id, int32_t vocab_size, + cudaStream_t stream, std::shared_ptr tokenizer, + std::string task, std::string model_id) + : encoder_(std::move(encoder)), decoder_(std::move(decoder)), state_(std::move(state)), + hidden_size_(hidden_size), num_decoder_layers_(num_decoder_layers), + max_source_length_(max_source_length), decoder_start_token_id_(decoder_start_token_id), + eos_token_id_(eos_token_id), bos_token_id_(bos_token_id), pad_token_id_(pad_token_id), + vocab_size_(vocab_size), stream_(stream), tokenizer_(std::move(tokenizer)), + task_(std::move(task)), model_id_(std::move(model_id)) { + if (!encoder_ || !encoder_->ok()) + throw std::runtime_error("BartPipeline: invalid encoder"); + if (!decoder_ || !decoder_->ok()) + throw std::runtime_error("BartPipeline: invalid decoder"); + if (!state_ || !state_->ok()) + throw std::runtime_error("BartPipeline: invalid state"); + if (!tokenizer_) + throw std::runtime_error("BartPipeline: no tokenizer configured"); + if (vocab_size_ <= 0) + throw std::runtime_error("BartPipeline: vocabulary size must be positive"); + + cross_kv_bytes_ = static_cast(max_source_length_) * + static_cast(hidden_size_) * sizeof(float); + cross_k_ptrs_.resize(static_cast(num_decoder_layers_)); + cross_v_ptrs_.resize(static_cast(num_decoder_layers_)); + for (int32_t i = 0; i < num_decoder_layers_; ++i) { + const std::size_t layer = static_cast(i); + if (cross_k_ptrs_[layer].allocate(cross_kv_bytes_) != cudaSuccess) + throw std::runtime_error("BartPipeline: unable to allocate cross-attention key buffer"); + if (cross_v_ptrs_[layer].allocate(cross_kv_bytes_) != cudaSuccess) + throw std::runtime_error( + "BartPipeline: unable to allocate cross-attention value buffer"); + } +} + +std::vector BartPipeline::task_bindings() { + return {internal::bind(*this, bart_config_fields)}; +} + +internal::TextResult BartPipeline::run(const internal::ConditionalTextGenerationRequest& request, + internal::ConfigView config) { + internal::validate_config(bart_config_fields, config); + int32_t max_tokens = 128; + if (const auto val = + internal::config_get(config, bart_config_fields, "max_new_tokens")) { + if (*val <= 0 || *val > std::numeric_limits::max()) + throw internal::ConfigError( + "BartPipeline: max_new_tokens must be positive and fit in int32"); + max_tokens = static_cast(*val); + } + if (const auto val = internal::config_get(config, bart_config_fields, "temperature")) { + if (*val < 0.0) + throw internal::ConfigError("BartPipeline: temperature must be non-negative"); + } + if (const auto val = internal::config_get(config, bart_config_fields, "top_k")) { + if (*val < 0) + throw internal::ConfigError("BartPipeline: top_k must be non-negative"); + } + + auto [padded, copy_len] = prepare_encoder_input_semantic(request.source); + + run_encoder(padded, copy_len); + setup_cross_attention(); + + auto output_ids = run_decoder(max_tokens); + + internal::TextResult out; + out.token_ids = std::move(output_ids); + if (tokenizer_ && !out.token_ids.empty()) + out.text = tokenizer_->decode(out.token_ids); + return out; +} + +TextResult BartPipeline::generate(const std::string& prompt, const TextGenerationConfig& cfg) { + auto [padded, copy_len] = prepare_encoder_input(prompt); + if (copy_len == 0) + return TextResult{"[empty input]", {}}; + + run_encoder(padded, copy_len); + setup_cross_attention(); + + int32_t max_tokens = cfg.max_new_tokens > 0 ? cfg.max_new_tokens : 128; + auto output_ids = run_decoder(max_tokens); + + TextResult out; + out.token_ids = std::move(output_ids); + if (tokenizer_ && !out.token_ids.empty()) + out.text = tokenizer_->decode(out.token_ids); + return out; +} + +std::pair, int32_t> +BartPipeline::prepare_encoder_input_semantic(const internal::TextSource& source) { + std::vector ids; + if (const auto* view = std::get_if(&source)) { + if (view->empty()) + throw std::invalid_argument("BartPipeline: empty input"); + if (!tokenizer_) + throw std::runtime_error("BartPipeline: no tokenizer configured"); + ids = tokenizer_->encode(std::string(*view)); + if (ids.empty()) + throw std::invalid_argument("BartPipeline: empty input"); + } else { + const auto span = std::get>(source); + if (span.empty()) + throw std::invalid_argument("BartPipeline: empty input"); + for (const auto id : span) { + if (id < 0 || id >= vocab_size_) + throw std::invalid_argument("BartPipeline: token id outside model vocabulary"); + } + ids.assign(span.begin(), span.end()); + } + + // Add BOS/EOS if not present + if (bos_token_id_ >= 0 && (ids.empty() || ids.front() != bos_token_id_)) + ids.insert(ids.begin(), bos_token_id_); + if (eos_token_id_ >= 0 && (ids.empty() || ids.back() != eos_token_id_)) + ids.push_back(eos_token_id_); + + if (static_cast(ids.size()) > max_source_length_) + throw std::invalid_argument("BartPipeline: input exceeds maximum sequence length"); + + int32_t copy_len = static_cast(ids.size()); + std::vector padded(static_cast(max_source_length_), pad_token_id_); + std::copy_n(ids.begin(), ids.size(), padded.begin()); + return {std::move(padded), copy_len}; +} + +std::pair, int32_t> +BartPipeline::prepare_encoder_input(const std::string& prompt) { + std::vector ids; + if (tokenizer_) + ids = tokenizer_->encode(prompt); + if (ids.empty()) + return {{}, 0}; + + // Add BOS/EOS if the native tokenizer didn't + if (bos_token_id_ >= 0 && (ids.empty() || ids.front() != bos_token_id_)) + ids.insert(ids.begin(), bos_token_id_); + if (eos_token_id_ >= 0 && (ids.empty() || ids.back() != eos_token_id_)) + ids.push_back(eos_token_id_); + + int32_t copy_len = std::min(static_cast(ids.size()), max_source_length_); + std::vector padded(static_cast(max_source_length_), pad_token_id_); + std::copy_n(ids.begin(), copy_len, padded.begin()); + return {std::move(padded), copy_len}; +} + +void BartPipeline::run_encoder(const std::vector& padded_ids, int32_t actual_len) { + actual_enc_len_ = actual_len; + + // Build attention mask: 0.0 for valid positions, -1e9 for padding + encoder_attention_mask_.assign(static_cast(max_source_length_), -1e9f); + for (int32_t i = 0; i < actual_len; ++i) + encoder_attention_mask_[static_cast(i)] = 0.0f; + + TensorMap inputs; + Tensor ids_tensor; + ids_tensor.data = const_cast(padded_ids.data()); + ids_tensor.shape = {max_source_length_}; + ids_tensor.dtype = DType::kInt32; + inputs["input_ids"] = ids_tensor; + + // Provide attention mask if the encoder expects it + Tensor mask_tensor; + if (encoder_->has_input("attention_mask")) { + mask_tensor.data = encoder_attention_mask_.data(); + mask_tensor.shape = {max_source_length_}; + mask_tensor.dtype = DType::kFloat32; + inputs["attention_mask"] = mask_tensor; + } + + encoder_->forward_async(inputs); + encoder_->sync(); +} + +void BartPipeline::setup_cross_attention() { + void* enc_out = encoder_->device_ptr("encoder_output"); + if (!enc_out) + throw std::runtime_error("BartPipeline: no encoder_output"); + for (int32_t i = 0; i < num_decoder_layers_; ++i) { + auto idx = static_cast(i); + cudaMemcpy(cross_k_ptrs_[idx].get(), enc_out, cross_kv_bytes_, cudaMemcpyDeviceToDevice); + cudaMemcpy(cross_v_ptrs_[idx].get(), enc_out, cross_kv_bytes_, cudaMemcpyDeviceToDevice); + } + for (int32_t i = 0; i < num_decoder_layers_; ++i) { + std::string s = "_" + std::to_string(i); + decoder_->bind_external("cross_k" + s, cross_k_ptrs_[static_cast(i)].get()); + decoder_->bind_external("cross_v" + s, cross_v_ptrs_[static_cast(i)].get()); + } +} + +std::vector BartPipeline::run_decoder(int32_t max_new_tokens) { + state_->reset(); + state_->bind_to(*decoder_); + std::vector output_ids; + std::vector logits; + int32_t current_token = decoder_start_token_id_; + for (int32_t step = 0; step < max_new_tokens; ++step) { + run_decoder_step(current_token, logits); + int32_t next = bart_select_argmax_token(logits); + if (next == eos_token_id_) + break; + output_ids.push_back(next); + current_token = next; + } + return output_ids; +} + +void BartPipeline::run_decoder_step(int32_t token_id, std::vector& logits) { + Tensor token_tensor; + token_tensor.data = &token_id; + token_tensor.shape = {1}; + token_tensor.dtype = DType::kInt32; + TensorMap inputs; + inputs["token_id"] = token_tensor; + Tensor cross_mask_tensor; + if (decoder_->has_input("cross_attention_mask")) { + cross_mask_tensor.data = encoder_attention_mask_.data(); + cross_mask_tensor.shape = {max_source_length_}; + cross_mask_tensor.dtype = DType::kFloat32; + inputs["cross_attention_mask"] = cross_mask_tensor; + } + state_->prepare_step(inputs); + TensorMap outputs = decoder_->forward(inputs); + auto it = outputs.find("logits"); + if (it == outputs.end()) + throw std::runtime_error("BartPipeline: no logits output"); + auto num = it->second.numel(); + logits.resize(static_cast(num)); + std::memcpy(logits.data(), it->second.data, num * sizeof(float)); + state_->advance(); +} + +} // namespace trtmc diff --git a/families/bart/runtime/pipeline.h b/families/bart/runtime/pipeline.h new file mode 100644 index 0000000000..037302697b --- /dev/null +++ b/families/bart/runtime/pipeline.h @@ -0,0 +1,88 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "families/bart/runtime/device_buffer.h" +#include "families/bart/runtime/inference_state.h" +#include "families/bart/runtime/tokenizer.h" +#include "trtmc/internal/config.h" +#include "trtmc/internal/model.h" +#include "trtmc/internal/text.h" +#include "trtmc/runtime/trt_module.h" +#include "trtmc/task.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace trtmc { + +inline constexpr internal::ConfigField bart_config_fields[] = { + {"max_new_tokens", internal::ConfigKind::I64, std::int64_t{128}, "Generation token cap"}, + {"temperature", internal::ConfigKind::F64, 1.0, "Sampling temperature"}, + {"top_k", internal::ConfigKind::I64, std::int64_t{1}, "Top-k sampling cutoff"}, +}; + +class BartPipeline final : public internal::IModel, + public internal::IConditionalTextGeneration, + public ITextGeneration { + public: + BartPipeline(std::unique_ptr encoder, std::unique_ptr decoder, + std::unique_ptr state, int32_t hidden_size, + int32_t num_decoder_layers, int32_t max_source_length, + int32_t decoder_start_token_id, int32_t eos_token_id, int32_t bos_token_id, + int32_t pad_token_id, int32_t vocab_size, cudaStream_t stream, + std::shared_ptr tokenizer, + std::string task = std::string(internal::IConditionalTextGeneration::kTask), + std::string model_id = ""); + + const char* task() const noexcept override { return task_.c_str(); } + std::vector task_bindings() override; + + // Semantic Task SDK + internal::TextResult run(const internal::ConditionalTextGenerationRequest& request, + internal::ConfigView config) override; + + // Legacy ITextGeneration + TextResult generate(const std::string& prompt, const TextGenerationConfig& cfg) override; + int32_t default_max_new_tokens() const override { return 128; } + + private: + std::pair, int32_t> prepare_encoder_input(const std::string& prompt); + std::pair, int32_t> + prepare_encoder_input_semantic(const internal::TextSource& source); + void run_encoder(const std::vector& padded_ids, int32_t actual_len); + void setup_cross_attention(); + std::vector run_decoder(int32_t max_new_tokens); + void run_decoder_step(int32_t token_id, std::vector& logits); + + std::unique_ptr encoder_; + std::unique_ptr decoder_; + std::unique_ptr state_; + int32_t hidden_size_; + int32_t num_decoder_layers_; + int32_t max_source_length_; + int32_t decoder_start_token_id_; + int32_t eos_token_id_; + int32_t bos_token_id_; + int32_t pad_token_id_; + int32_t vocab_size_; + cudaStream_t stream_; + std::shared_ptr tokenizer_; + std::string task_; + std::string model_id_; + std::vector cross_k_ptrs_; + std::vector cross_v_ptrs_; + std::vector encoder_attention_mask_; + std::size_t cross_kv_bytes_{0}; + int32_t actual_enc_len_{0}; +}; + +} // namespace trtmc diff --git a/families/bart/runtime/plugin.cpp b/families/bart/runtime/plugin.cpp index c0639033c0..96ecc614ad 100644 --- a/families/bart/runtime/plugin.cpp +++ b/families/bart/runtime/plugin.cpp @@ -10,6 +10,7 @@ #include "families/bart/runtime/device_buffer.h" #include "families/bart/runtime/distributed_runtime.h" #include "families/bart/runtime/kv_cache.h" +#include "families/bart/runtime/pipeline.h" #include "families/bart/runtime/plugin_helpers.h" #include "families/bart/runtime/runtime_config.h" #include "trtmc/runtime/family_factory.h" @@ -22,6 +23,7 @@ #include #include #include +#include #include namespace trtmc { @@ -136,191 +138,6 @@ void validate_runtime_config_json(std::string_view json) { } // namespace bart -class BartPipeline final : public ITextGeneration { - public: - BartPipeline(std::unique_ptr encoder, std::unique_ptr decoder, - std::unique_ptr state, int32_t hidden_size, - int32_t num_decoder_layers, int32_t max_source_length, - int32_t decoder_start_token_id, int32_t eos_token_id, int32_t bos_token_id, - int32_t pad_token_id, cudaStream_t stream, std::shared_ptr tokenizer, - std::string model_id_str) - : encoder_(std::move(encoder)), decoder_(std::move(decoder)), state_(std::move(state)), - hidden_size_(hidden_size), num_decoder_layers_(num_decoder_layers), - max_source_length_(max_source_length), decoder_start_token_id_(decoder_start_token_id), - eos_token_id_(eos_token_id), bos_token_id_(bos_token_id), pad_token_id_(pad_token_id), - stream_(stream), tokenizer_(std::move(tokenizer)), model_id_(std::move(model_id_str)) { - if (!encoder_ || !encoder_->ok()) - throw std::runtime_error("BartPipeline: invalid encoder"); - if (!decoder_ || !decoder_->ok()) - throw std::runtime_error("BartPipeline: invalid decoder"); - if (!state_ || !state_->ok()) - throw std::runtime_error("BartPipeline: invalid state"); - - cross_kv_bytes_ = static_cast(max_source_length_) * - static_cast(hidden_size_) * sizeof(float); - cross_k_ptrs_.resize(static_cast(num_decoder_layers_)); - cross_v_ptrs_.resize(static_cast(num_decoder_layers_)); - for (int32_t i = 0; i < num_decoder_layers_; ++i) { - const std::size_t layer = static_cast(i); - if (cross_k_ptrs_[layer].allocate(cross_kv_bytes_) != cudaSuccess) - throw std::runtime_error( - "BartPipeline: unable to allocate cross-attention key buffer"); - if (cross_v_ptrs_[layer].allocate(cross_kv_bytes_) != cudaSuccess) - throw std::runtime_error( - "BartPipeline: unable to allocate cross-attention value buffer"); - } - } - - TextResult generate(const std::string& prompt, const TextGenerationConfig& cfg) override { - auto [padded, copy_len] = prepare_encoder_input(prompt); - if (copy_len == 0) - return TextResult{"[empty input]", {}}; - - run_encoder(padded, copy_len); - setup_cross_attention(); - - int32_t max_tokens = cfg.max_new_tokens > 0 ? cfg.max_new_tokens : 128; - auto output_ids = run_decoder(max_tokens); - - TextResult out; - out.token_ids = std::move(output_ids); - if (tokenizer_ && !out.token_ids.empty()) - out.text = tokenizer_->decode(out.token_ids); - return out; - } - - int32_t default_max_new_tokens() const override { return 128; } - - private: - std::pair, int32_t> prepare_encoder_input(const std::string& prompt) { - std::vector ids; - if (tokenizer_) - ids = tokenizer_->encode(prompt); - if (ids.empty()) - return {{}, 0}; - - // Add BOS/EOS if the native tokenizer didn't - if (bos_token_id_ >= 0 && (ids.empty() || ids.front() != bos_token_id_)) - ids.insert(ids.begin(), bos_token_id_); - if (eos_token_id_ >= 0 && (ids.empty() || ids.back() != eos_token_id_)) - ids.push_back(eos_token_id_); - - int32_t copy_len = std::min(static_cast(ids.size()), max_source_length_); - std::vector padded(static_cast(max_source_length_), pad_token_id_); - std::copy_n(ids.begin(), copy_len, padded.begin()); - return {std::move(padded), copy_len}; - } - - void run_encoder(const std::vector& padded_ids, int32_t actual_len) { - actual_enc_len_ = actual_len; - - // Build attention mask: 0.0 for valid positions, -1e9 for padding - encoder_attention_mask_.assign(static_cast(max_source_length_), -1e9f); - for (int32_t i = 0; i < actual_len; ++i) - encoder_attention_mask_[static_cast(i)] = 0.0f; - - TensorMap inputs; - Tensor ids_tensor; - ids_tensor.data = const_cast(padded_ids.data()); - ids_tensor.shape = {max_source_length_}; - ids_tensor.dtype = DType::kInt32; - inputs["input_ids"] = ids_tensor; - - // Provide attention mask if the encoder expects it - Tensor mask_tensor; - if (encoder_->has_input("attention_mask")) { - mask_tensor.data = encoder_attention_mask_.data(); - mask_tensor.shape = {max_source_length_}; - mask_tensor.dtype = DType::kFloat32; - inputs["attention_mask"] = mask_tensor; - } - - encoder_->forward_async(inputs); - encoder_->sync(); - } - - void setup_cross_attention() { - void* enc_out = encoder_->device_ptr("encoder_output"); - if (!enc_out) - throw std::runtime_error("BartPipeline: no encoder_output"); - for (int32_t i = 0; i < num_decoder_layers_; ++i) { - auto idx = static_cast(i); - cudaMemcpy(cross_k_ptrs_[idx].get(), enc_out, cross_kv_bytes_, - cudaMemcpyDeviceToDevice); - cudaMemcpy(cross_v_ptrs_[idx].get(), enc_out, cross_kv_bytes_, - cudaMemcpyDeviceToDevice); - } - for (int32_t i = 0; i < num_decoder_layers_; ++i) { - std::string s = "_" + std::to_string(i); - decoder_->bind_external("cross_k" + s, - cross_k_ptrs_[static_cast(i)].get()); - decoder_->bind_external("cross_v" + s, - cross_v_ptrs_[static_cast(i)].get()); - } - } - - std::vector run_decoder(int32_t max_new_tokens) { - state_->reset(); - state_->bind_to(*decoder_); - std::vector output_ids; - std::vector logits; - int32_t current_token = decoder_start_token_id_; - for (int32_t step = 0; step < max_new_tokens; ++step) { - run_decoder_step(current_token, logits); - int32_t next = bart_select_argmax_token(logits); - if (next == eos_token_id_) - break; - output_ids.push_back(next); - current_token = next; - } - return output_ids; - } - - void run_decoder_step(int32_t token_id, std::vector& logits) { - Tensor token_tensor; - token_tensor.data = &token_id; - token_tensor.shape = {1}; - token_tensor.dtype = DType::kInt32; - TensorMap inputs; - inputs["token_id"] = token_tensor; - Tensor cross_mask_tensor; - if (decoder_->has_input("cross_attention_mask")) { - cross_mask_tensor.data = encoder_attention_mask_.data(); - cross_mask_tensor.shape = {max_source_length_}; - cross_mask_tensor.dtype = DType::kFloat32; - inputs["cross_attention_mask"] = cross_mask_tensor; - } - state_->prepare_step(inputs); - TensorMap outputs = decoder_->forward(inputs); - auto it = outputs.find("logits"); - if (it == outputs.end()) - throw std::runtime_error("BartPipeline: no logits output"); - auto num = it->second.numel(); - logits.resize(static_cast(num)); - std::memcpy(logits.data(), it->second.data, num * sizeof(float)); - state_->advance(); - } - - std::unique_ptr encoder_; - std::unique_ptr decoder_; - std::unique_ptr state_; - int32_t hidden_size_; - int32_t num_decoder_layers_; - int32_t max_source_length_; - int32_t decoder_start_token_id_; - int32_t eos_token_id_; - int32_t bos_token_id_; - int32_t pad_token_id_; - int32_t actual_enc_len_{0}; - cudaStream_t stream_; - std::shared_ptr tokenizer_; - std::string model_id_; - std::vector cross_k_ptrs_; - std::vector cross_v_ptrs_; - std::vector encoder_attention_mask_; - std::size_t cross_kv_bytes_{0}; -}; - ITask* create_bart(const FamilyContext& context) { const BartRuntimeConfig config = parse_runtime_config(bart::require_text_section(context.reader, "runtime.json")); @@ -349,11 +166,19 @@ ITask* create_bart(const FamilyContext& context) { if (!state->ok()) throw std::runtime_error("bart failed to create KV cache"); + std::string task = std::string(internal::IConditionalTextGeneration::kTask); + if (!context.reader.info().task.empty()) { + task = context.reader.info().task; + if (task != internal::IConditionalTextGeneration::kTask && task != "text_generation") { + throw std::runtime_error("bart does not implement task: " + task); + } + } + return new BartPipeline(std::move(encoder), std::move(decoder), std::move(state), config.hidden_size, config.decoder_layers, config.max_source_length, config.decoder_start_token_id, config.eos_token_id, config.bos_token_id, - config.pad_token_id, stream, bart::create_tokenizer(context.reader), - std::string{}); + config.pad_token_id, config.vocab_size, stream, + bart::create_tokenizer(context.reader), std::move(task), std::string{}); } } // namespace trtmc diff --git a/families/bart/support.py b/families/bart/support.py index 77a093dfce..e7805dd58d 100644 --- a/families/bart/support.py +++ b/families/bart/support.py @@ -8,6 +8,6 @@ describe = family_support( model_types=("bart", "mbart"), - tasks=("text_generation",), - default_task="text_generation", + tasks=("conditional_text_generation", "text_generation"), + default_task="conditional_text_generation", ) diff --git a/families/bart/tests/cpp/test_task_contract.cpp b/families/bart/tests/cpp/test_task_contract.cpp new file mode 100644 index 0000000000..923e2ca8cd --- /dev/null +++ b/families/bart/tests/cpp/test_task_contract.cpp @@ -0,0 +1,577 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "families/bart/runtime/pipeline.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +// CPU CUDA stubs so tests run without a GPU. +extern "C" { + +cudaError_t cudaMalloc(void** devPtr, size_t size) { + if (size == 0) { + *devPtr = nullptr; + return cudaSuccess; + } + *devPtr = std::malloc(size); + return *devPtr ? cudaSuccess : cudaErrorMemoryAllocation; +} + +cudaError_t cudaFree(void* devPtr) { + std::free(devPtr); + return cudaSuccess; +} + +cudaError_t cudaMemcpy(void* dst, const void* src, size_t count, cudaMemcpyKind kind) { + (void)kind; + if (count > 0 && dst && src) { + std::memcpy(dst, src, count); + } + return cudaSuccess; +} + +} // extern "C" + +namespace { + +constexpr int32_t kHiddenSize = 32; +constexpr int32_t kDecoderLayers = 2; +constexpr int32_t kMaxSourceLength = 64; +constexpr int32_t kDecoderStartTokenId = 2; +constexpr int32_t kEosTokenId = 2; +constexpr int32_t kBosTokenId = 0; +constexpr int32_t kPadTokenId = 1; +constexpr int32_t kVocabSize = 100; + +class MockTokenizer final : public trtmc::ITokenizer { + public: + std::vector encode(const std::string& text) const override { + ++encode_calls; + std::vector ids; + ids.reserve(text.size()); + for (const unsigned char byte : text) { + ids.push_back(static_cast(byte % 90 + 3)); + } + return ids; + } + + std::string decode(const std::vector& ids) const override { + ++decode_calls; + std::string text; + text.reserve(ids.size()); + for (const auto id : ids) { + text.push_back(static_cast('a' + (id % 26))); + } + return text; + } + + std::int32_t id_for_token(std::string_view) const override { return 0; } + std::string token_for_id(std::int32_t) const override { return {}; } + + mutable int encode_calls{0}; + mutable int decode_calls{0}; +}; + +class MockEncoderModule final : public trtmc::ITrtModule { + public: + MockEncoderModule() { + enc_out_buffer_.assign(static_cast(kMaxSourceLength) * kHiddenSize, 0.5f); + } + + trtmc::TensorMap forward(const trtmc::TensorMap&) override { return {}; } + trtmc::DeviceTensorMap forward_device(const trtmc::DeviceTensorMap&) override { return {}; } + void forward_device_async(const trtmc::DeviceTensorMap&) override {} + + void forward_async(const trtmc::TensorMap& inputs) override { + ++calls; + if (inputs.find("input_ids") != inputs.end()) { + const auto& tensor = inputs.at("input_ids"); + last_input_numel = static_cast(tensor.numel()); + } + } + + void sync() override {} + cudaStream_t stream() const override { return nullptr; } + void enable_cuda_graph() override {} + bool cuda_graph_active() const override { return false; } + bool cuda_graph_captured() const override { return false; } + int32_t profile_idx() const override { return 0; } + std::vector input_info() const override { return {}; } + std::vector output_info() const override { return {}; } + bool has_input(const std::string& name) const override { return name == "attention_mask"; } + bool has_output(const std::string& name) const override { return name == "encoder_output"; } + trtmc::DType tensor_dtype(const std::string&) const override { return trtmc::DType::kFloat32; } + std::vector tensor_shape(const std::string&) const override { + return {kMaxSourceLength, kHiddenSize}; + } + std::vector input_profile_shape(const std::string&, int32_t, + trtmc::ProfileShapeSelector) const override { + return {}; + } + int32_t optimization_profile_count() const override { return 1; } + + void* device_ptr(const std::string& name) override { + if (name == "encoder_output") + return enc_out_buffer_.data(); + return nullptr; + } + + void bind_external(const std::string&, void*) override {} + void bind_external(const std::string&, void*, const std::vector&) override {} + int32_t input_rank(const std::string&) const override { return 1; } + bool input_is_dynamic(const std::string&) const override { return false; } + void reset_execution_context() override {} + void set_timing_label(std::string) override {} + bool ok() const override { return is_ok; } + void keep_alive(std::shared_ptr) override {} + + bool is_ok{true}; + int calls{0}; + int32_t last_input_numel{0}; + std::vector enc_out_buffer_; +}; + +class MockDecoderModule final : public trtmc::ITrtModule { + public: + MockDecoderModule() { logits_buffer_.assign(kVocabSize, 0.0f); } + + trtmc::TensorMap forward(const trtmc::TensorMap&) override { + ++calls; + std::fill(logits_buffer_.begin(), logits_buffer_.end(), 0.0f); + if (calls <= stop_at_step) { + logits_buffer_[static_cast(emit_token)] = 10.0f; + } else { + logits_buffer_[static_cast(kEosTokenId)] = 10.0f; + } + + trtmc::Tensor logits_tensor; + logits_tensor.data = logits_buffer_.data(); + logits_tensor.shape = {kVocabSize}; + logits_tensor.dtype = trtmc::DType::kFloat32; + return {{"logits", logits_tensor}}; + } + + trtmc::DeviceTensorMap forward_device(const trtmc::DeviceTensorMap&) override { return {}; } + void forward_device_async(const trtmc::DeviceTensorMap&) override {} + void forward_async(const trtmc::TensorMap&) override {} + void sync() override {} + cudaStream_t stream() const override { return nullptr; } + void enable_cuda_graph() override {} + bool cuda_graph_active() const override { return false; } + bool cuda_graph_captured() const override { return false; } + int32_t profile_idx() const override { return 0; } + std::vector input_info() const override { return {}; } + std::vector output_info() const override { return {}; } + bool has_input(const std::string& name) const override { + return name == "cross_attention_mask"; + } + bool has_output(const std::string& name) const override { return name == "logits"; } + trtmc::DType tensor_dtype(const std::string&) const override { return trtmc::DType::kFloat32; } + std::vector tensor_shape(const std::string&) const override { return {kVocabSize}; } + std::vector input_profile_shape(const std::string&, int32_t, + trtmc::ProfileShapeSelector) const override { + return {}; + } + int32_t optimization_profile_count() const override { return 1; } + void* device_ptr(const std::string&) override { return nullptr; } + void bind_external(const std::string&, void*) override {} + void bind_external(const std::string&, void*, const std::vector&) override {} + int32_t input_rank(const std::string&) const override { return 1; } + bool input_is_dynamic(const std::string&) const override { return false; } + void reset_execution_context() override {} + void set_timing_label(std::string) override {} + bool ok() const override { return is_ok; } + void keep_alive(std::shared_ptr) override {} + + bool is_ok{true}; + int calls{0}; + int stop_at_step{3}; + int32_t emit_token{15}; + std::vector logits_buffer_; +}; + +class MockBartInferenceState final : public trtmc::BartInferenceState { + public: + void reset() override { + ++reset_calls; + position_ = 0; + } + + void bind_to(trtmc::ITrtModule&) override { ++bind_calls; } + void prepare_step(trtmc::TensorMap&, int32_t) override { ++prepare_calls; } + + void advance(int32_t n_tokens = 1) override { + ++advance_calls; + position_ += n_tokens; + } + + int32_t position() const override { return position_; } + int32_t max_length() const override { return 128; } + int32_t num_layers() const override { return kDecoderLayers; } + bool needs_attention_mask() const override { return true; } + std::size_t device_memory_bytes() const override { return 0; } + const char* state_type() const override { return "mock_state"; } + bool ok() const override { return is_ok; } + + bool is_ok{true}; + int reset_calls{0}; + int bind_calls{0}; + int prepare_calls{0}; + int advance_calls{0}; + int32_t position_{0}; +}; + +std::unique_ptr make_pipeline(MockEncoderModule** out_enc = nullptr, + MockDecoderModule** out_dec = nullptr, + MockBartInferenceState** out_state = nullptr, + MockTokenizer** out_tok = nullptr) { + auto enc = std::make_unique(); + auto dec = std::make_unique(); + auto state = std::make_unique(); + auto tok = std::make_shared(); + + if (out_enc) + *out_enc = enc.get(); + if (out_dec) + *out_dec = dec.get(); + if (out_state) + *out_state = state.get(); + if (out_tok) + *out_tok = tok.get(); + + return std::make_unique(std::move(enc), std::move(dec), std::move(state), + kHiddenSize, kDecoderLayers, kMaxSourceLength, + kDecoderStartTokenId, kEosTokenId, kBosTokenId, + kPadTokenId, kVocabSize, nullptr, std::move(tok)); +} + +void require(bool condition, const char* message) { + if (!condition) + throw std::runtime_error(std::string("contract assertion failed: ") + message); +} + +template +void rejects(Fn&& fn, const char* message) { + try { + fn(); + } catch (const Exception&) { + return; + } + throw std::runtime_error(std::string("expected exception not thrown: ") + message); +} + +void test_task_bindings_and_primary_task() { + auto pipeline = make_pipeline(); + require(std::string(pipeline->task()) == "conditional_text_generation", + "primary task must be conditional_text_generation"); + + const auto bindings = pipeline->task_bindings(); + require(bindings.size() == 1, "pipeline must bind exactly 1 task"); + require(bindings[0].task == trtmc::internal::IConditionalTextGeneration::kTask, + "bound task must be conditional_text_generation"); + require(bindings[0].config_fields.size() == 3, "bound task must expose 3 config fields"); + require(bindings[0].config_fields[0].name == "max_new_tokens", + "field name must be max_new_tokens"); + require(bindings[0].config_fields[0].kind == trtmc::internal::ConfigKind::I64, + "field kind must be I64"); + require(bindings[0].config_fields[1].name == "temperature", "field name must be temperature"); + require(bindings[0].config_fields[1].kind == trtmc::internal::ConfigKind::F64, + "field kind must be F64"); + require(bindings[0].config_fields[2].name == "top_k", "field name must be top_k"); + require(bindings[0].config_fields[2].kind == trtmc::internal::ConfigKind::I64, + "field kind must be I64"); +} + +void test_conditional_text_generation_text_input() { + MockEncoderModule* enc = nullptr; + MockDecoderModule* dec = nullptr; + MockBartInferenceState* state = nullptr; + MockTokenizer* tok = nullptr; + + auto pipeline = make_pipeline(&enc, &dec, &state, &tok); + dec->stop_at_step = 3; + dec->emit_token = 25; + + const trtmc::internal::ConditionalTextGenerationRequest req{ + trtmc::internal::TextSource{"Translate this sentence"}}; + const auto result = pipeline->run(req, {}); + + require(enc->calls == 1, "encoder must be run once"); + require(dec->calls == 4, "decoder must run 4 steps (3 valid + 1 EOS)"); + require(state->reset_calls == 1, "state reset must be called"); + require(state->bind_calls == 1, "state bind_to must be called"); + require(tok->encode_calls == 1, "tokenizer encode must be called for text input"); + require(tok->decode_calls == 1, "tokenizer decode must be called for output tokens"); + require(result.token_ids.size() == 3, "must produce 3 tokens before EOS"); + require(result.token_ids[0] == 25 && result.token_ids[1] == 25 && result.token_ids[2] == 25, + "output token ids must match decoder emit"); + require(!result.text.empty(), "result text must be non-empty"); +} + +void test_conditional_text_generation_token_ids_input() { + MockEncoderModule* enc = nullptr; + MockDecoderModule* dec = nullptr; + MockTokenizer* tok = nullptr; + + auto pipeline = make_pipeline(&enc, &dec, nullptr, &tok); + dec->stop_at_step = 2; + dec->emit_token = 18; + + const std::vector ids = {10, 11, 12}; + const trtmc::internal::ConditionalTextGenerationRequest req{ + trtmc::internal::TextSource{trtmc::Span{ids.data(), ids.size()}}}; + const auto result = pipeline->run(req, {}); + + require(tok->encode_calls == 0, "token ids input must not invoke tokenizer encode"); + require(result.token_ids.size() == 2, "must generate 2 tokens"); + require(result.token_ids[0] == 18 && result.token_ids[1] == 18, + "generated token ids must match decoder emit"); +} + +void test_max_new_tokens_config() { + MockDecoderModule* dec = nullptr; + auto pipeline = make_pipeline(nullptr, &dec, nullptr, nullptr); + dec->stop_at_step = 100; // Does not emit EOS early + dec->emit_token = 12; + + const int64_t custom_max = 5; + const trtmc::internal::ConfigEntry entries[] = { + {"max_new_tokens", trtmc::internal::ConfigValue{custom_max}}}; + const trtmc::internal::ConfigView cfg{entries, 1}; + + const trtmc::internal::ConditionalTextGenerationRequest req{ + trtmc::internal::TextSource{"Test prompt"}}; + const auto result = pipeline->run(req, cfg); + + require(result.token_ids.size() == 5, "generation must stop at custom max_new_tokens limit"); +} + +void test_early_eos_stopping() { + MockDecoderModule* dec = nullptr; + auto pipeline = make_pipeline(nullptr, &dec, nullptr, nullptr); + dec->stop_at_step = 1; // Emits EOS on step 2 + dec->emit_token = 30; + + const trtmc::internal::ConditionalTextGenerationRequest req{ + trtmc::internal::TextSource{"Another prompt"}}; + const auto result = pipeline->run(req, {}); + + require(result.token_ids.size() == 1, "generation must stop immediately after EOS token"); + require(result.token_ids[0] == 30, "first token must match"); +} + +void test_invalid_inputs() { + auto pipeline = make_pipeline(); + + // Empty text + rejects( + [&] { + pipeline->run( + trtmc::internal::ConditionalTextGenerationRequest{trtmc::internal::TextSource{""}}, + {}); + }, + "empty text input must throw invalid_argument"); + + // Empty token ids span + const std::vector empty_ids; + rejects( + [&] { + pipeline->run( + trtmc::internal::ConditionalTextGenerationRequest{ + trtmc::Span{empty_ids.data(), empty_ids.size()}}, + {}); + }, + "empty token span must throw invalid_argument"); + + // Negative token id + const std::vector neg_ids = {5, -1, 10}; + rejects( + [&] { + pipeline->run( + trtmc::internal::ConditionalTextGenerationRequest{ + trtmc::Span{neg_ids.data(), neg_ids.size()}}, + {}); + }, + "negative token id must throw invalid_argument"); + + // Token id >= vocab_size + const std::vector oob_ids = {5, kVocabSize, 10}; + rejects( + [&] { + pipeline->run( + trtmc::internal::ConditionalTextGenerationRequest{ + trtmc::Span{oob_ids.data(), oob_ids.size()}}, + {}); + }, + "out of bounds token id must throw invalid_argument"); + + // Input sequence exceeding max_source_length + const std::vector overlong_ids(static_cast(kMaxSourceLength) + 1, 10); + rejects( + [&] { + pipeline->run( + trtmc::internal::ConditionalTextGenerationRequest{ + trtmc::Span{overlong_ids.data(), overlong_ids.size()}}, + {}); + }, + "input exceeding max_source_length must throw invalid_argument"); +} + +void test_invalid_config() { + auto pipeline = make_pipeline(); + + // Unknown config field + const trtmc::internal::ConfigEntry unknown[] = { + {"unknown_param", trtmc::internal::ConfigValue{int64_t{42}}}}; + rejects( + [&] { + pipeline->run( + trtmc::internal::ConditionalTextGenerationRequest{trtmc::internal::TextSource{"a"}}, + {unknown, 1}); + }, + "unknown config field must throw ConfigError"); + + // Config type mismatch (double instead of I64) + const trtmc::internal::ConfigEntry bad_type[] = { + {"max_new_tokens", trtmc::internal::ConfigValue{double{128.0}}}}; + rejects( + [&] { + pipeline->run( + trtmc::internal::ConditionalTextGenerationRequest{trtmc::internal::TextSource{"a"}}, + {bad_type, 1}); + }, + "config type mismatch must throw ConfigError"); + + // Non-positive max_new_tokens (0) + const trtmc::internal::ConfigEntry zero_tokens[] = { + {"max_new_tokens", trtmc::internal::ConfigValue{int64_t{0}}}}; + rejects( + [&] { + pipeline->run( + trtmc::internal::ConditionalTextGenerationRequest{trtmc::internal::TextSource{"a"}}, + {zero_tokens, 1}); + }, + "zero max_new_tokens must throw ConfigError"); + + // Negative max_new_tokens (-5) + const trtmc::internal::ConfigEntry neg_tokens[] = { + {"max_new_tokens", trtmc::internal::ConfigValue{int64_t{-5}}}}; + rejects( + [&] { + pipeline->run( + trtmc::internal::ConditionalTextGenerationRequest{trtmc::internal::TextSource{"a"}}, + {neg_tokens, 1}); + }, + "negative max_new_tokens must throw ConfigError"); +} + +void test_constructor_validates_dependencies() { + // Null encoder + rejects( + [&] { + trtmc::BartPipeline(nullptr, std::make_unique(), + std::make_unique(), kHiddenSize, + kDecoderLayers, kMaxSourceLength, kDecoderStartTokenId, kEosTokenId, + kBosTokenId, kPadTokenId, kVocabSize, nullptr, + std::make_shared()); + }, + "null encoder must throw runtime_error"); + + // Null decoder + rejects( + [&] { + trtmc::BartPipeline(std::make_unique(), nullptr, + std::make_unique(), kHiddenSize, + kDecoderLayers, kMaxSourceLength, kDecoderStartTokenId, kEosTokenId, + kBosTokenId, kPadTokenId, kVocabSize, nullptr, + std::make_shared()); + }, + "null decoder must throw runtime_error"); + + // Null state + rejects( + [&] { + trtmc::BartPipeline(std::make_unique(), + std::make_unique(), nullptr, kHiddenSize, + kDecoderLayers, kMaxSourceLength, kDecoderStartTokenId, kEosTokenId, + kBosTokenId, kPadTokenId, kVocabSize, nullptr, + std::make_shared()); + }, + "null state must throw runtime_error"); + + // Null tokenizer + rejects( + [&] { + trtmc::BartPipeline(std::make_unique(), + std::make_unique(), + std::make_unique(), kHiddenSize, + kDecoderLayers, kMaxSourceLength, kDecoderStartTokenId, kEosTokenId, + kBosTokenId, kPadTokenId, kVocabSize, nullptr, nullptr); + }, + "null tokenizer must throw runtime_error"); + + // Non-positive vocab size + rejects( + [&] { + trtmc::BartPipeline( + std::make_unique(), std::make_unique(), + std::make_unique(), kHiddenSize, kDecoderLayers, + kMaxSourceLength, kDecoderStartTokenId, kEosTokenId, kBosTokenId, kPadTokenId, 0, + nullptr, std::make_shared()); + }, + "zero vocab size must throw runtime_error"); +} + +void test_legacy_generate() { + auto pipeline = make_pipeline(); + trtmc::TextGenerationConfig cfg; + cfg.max_new_tokens = 4; + const auto res = pipeline->generate("Legacy prompt", cfg); + require(!res.token_ids.empty(), "legacy generate must produce tokens"); + require(!res.text.empty(), "legacy generate must decode text"); +} + +void test_result_lifetime() { + trtmc::internal::TextResult saved_result; + { + auto pipeline = make_pipeline(); + saved_result = pipeline->run( + trtmc::internal::ConditionalTextGenerationRequest{trtmc::internal::TextSource{"Test"}}, + {}); + } + // Result remains fully valid after pipeline is destroyed + require(!saved_result.text.empty(), "result text must persist after pipeline destruction"); + require(!saved_result.token_ids.empty(), + "result token_ids must persist after pipeline destruction"); +} + +} // namespace + +int main() { + try { + test_task_bindings_and_primary_task(); + test_conditional_text_generation_text_input(); + test_conditional_text_generation_token_ids_input(); + test_max_new_tokens_config(); + test_early_eos_stopping(); + test_invalid_inputs(); + test_invalid_config(); + test_constructor_validates_dependencies(); + test_legacy_generate(); + test_result_lifetime(); + std::cout << "All BART Task contract tests passed.\n"; + return 0; + } catch (const std::exception& error) { + std::cerr << "FAIL: " << error.what() << '\n'; + return 1; + } +} diff --git a/families/bart/tests/manifests/bart-base-tp4.json b/families/bart/tests/manifests/bart-base-tp4.json index 21ed5541b3..54ac8a3425 100644 --- a/families/bart/tests/manifests/bart-base-tp4.json +++ b/families/bart/tests/manifests/bart-base-tp4.json @@ -4,7 +4,7 @@ "hf_revision": "aadd2ab0ae0c8268c7c9693540e9904811f36177", "bundle": "bart-base-tp4.bundle", "family": "bart", - "task": "text_generation", + "task": "conditional_text_generation", "testcases": [ { "name": "bart-base-tp4", diff --git a/families/bart/tests/manifests/bart-base.json b/families/bart/tests/manifests/bart-base.json index 448a47e840..da16f3bf8f 100644 --- a/families/bart/tests/manifests/bart-base.json +++ b/families/bart/tests/manifests/bart-base.json @@ -4,7 +4,7 @@ "hf_revision": "aadd2ab0ae0c8268c7c9693540e9904811f36177", "bundle": "bart-base.bundle", "family": "bart", - "task": "text_generation", + "task": "conditional_text_generation", "precision": "fp16", "testcases": [ { diff --git a/families/bart/tests/sdk_consumer.c b/families/bart/tests/sdk_consumer.c new file mode 100644 index 0000000000..709977dce0 --- /dev/null +++ b/families/bart/tests/sdk_consumer.c @@ -0,0 +1,118 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include +#include + +static trtmc_string_view text(const char* value) { + const trtmc_string_view result = {value, (uint64_t)strlen(value)}; + return result; +} + +static void json_string(trtmc_string_view value) { + uint64_t i; + putchar('"'); + for (i = 0; i < value.size; ++i) { + const unsigned char byte = (unsigned char)value.data[i]; + if (byte == '"' || byte == '\\') { + putchar('\\'); + putchar(byte); + } else if (byte < 32) { + printf("\\u%04x", (unsigned)byte); + } else { + putchar(byte); + } + } + putchar('"'); +} + +int main(int argc, char** argv) { + const trtmc_core_api_v1* core = NULL; + const trtmc_api_header* header = NULL; + const trtmc_conditional_text_generation_api_v1* task = NULL; + trtmc_model* model = NULL; + trtmc_result* result = NULL; + trtmc_error* error = NULL; + trtmc_load_options_v1 options = {0}; + trtmc_conditional_text_generation_request_v1 request = {0}; + trtmc_text_result_view_v1 view = {0}; + trtmc_config_entry_v1 config_entries[1] = {0}; + trtmc_config_view_v1 config = {0}; + uint64_t i, field_count = 0; + int status = 1; + + if (argc != 4 && argc != 5) { + fprintf(stderr, "Usage: %s BUNDLE RUNTIME_ROOT TEXT [MAX_NEW_TOKENS]\n", argv[0]); + return 2; + } + if (trtmc_get_api(1, 0, &core) != TRTMC_OK || !core) { + fprintf(stderr, "unable to obtain the v1 C API\n"); + goto cleanup; + } + options.struct_size = sizeof(options); + options.runtime_root = text(argv[2]); + if (core->model_load(text(argv[1]), &options, &model, &error) != TRTMC_OK) + goto cleanup; + if (core->model_get_task_api(model, text(TRTMC_TASK_CONDITIONAL_TEXT_GENERATION), 1, 0, &header, + &error) != TRTMC_OK) + goto cleanup; + if (!header || header->byte_size < sizeof(*task)) { + fprintf(stderr, "incomplete conditional_text_generation C table\n"); + goto cleanup; + } + task = (const trtmc_conditional_text_generation_api_v1*)header; + if (core->config_field_count(model, text(TRTMC_TASK_CONDITIONAL_TEXT_GENERATION), 1, 0, + &field_count, &error) != TRTMC_OK) + goto cleanup; + if (field_count != 3) { + fprintf(stderr, "bart must expose exactly 3 config fields\n"); + goto cleanup; + } + + request.source.kind = TRTMC_TEXT_UTF8; + request.source.as.text = text(argv[3]); + if (argc >= 5) { + config_entries[0].name = text("max_new_tokens"); + config_entries[0].value.kind = TRTMC_CONFIG_I64; + config_entries[0].value.as.i64 = atoll(argv[4]); + config.entries = config_entries; + config.count = 1; + } + if (task->run(model, &request, argc >= 5 ? &config : NULL, &result, &error) != TRTMC_OK) + goto cleanup; + + core->model_release(model); + model = NULL; + + if (task->result_view(result, &view, &error) != TRTMC_OK) + goto cleanup; + + printf("{\"task\":\"conditional_text_generation\",\"text\":"); + json_string(view.text); + printf(",\"tokens\":["); + for (i = 0; i < view.token_count; ++i) { + if (i) + putchar(','); + printf("%" PRId32, view.tokens[i]); + } + printf("]}\n"); + status = 0; + +cleanup: + if (error && core) { + const trtmc_string_view message = core->error_message(error); + fprintf(stderr, "%.*s\n", (int)message.size, message.data); + core->error_release(error); + } + if (result && core) + core->result_release(result); + if (model && core) + core->model_release(model); + return status; +} diff --git a/families/bart/tests/sdk_consumer.cpp b/families/bart/tests/sdk_consumer.cpp new file mode 100644 index 0000000000..b97d0ba0c4 --- /dev/null +++ b/families/bart/tests/sdk_consumer.cpp @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include + +namespace { +void json_string(std::string_view value) { + std::cout << '"'; + for (const unsigned char byte : value) { + if (byte == '"' || byte == '\\') + std::cout << '\\' << static_cast(byte); + else if (byte < 32) + std::cout << "\\u" << std::hex << std::setw(4) << std::setfill('0') + << static_cast(byte) << std::dec << std::setfill(' '); + else + std::cout << static_cast(byte); + } + std::cout << '"'; +} +} // namespace + +int main(int argc, char** argv) { + if (argc != 4 && argc != 5) { + std::cerr << "Usage: " << argv[0] << " BUNDLE RUNTIME_ROOT TEXT [MAX_NEW_TOKENS]\n"; + return 2; + } + try { + auto result = [&] { + trtmc::LoadOptions options; + options.runtime_root = argv[2]; + auto model = trtmc::Model::load(argv[1], options); + const auto task = model.task(); + if (task.config_fields().size() != 3) + throw std::runtime_error("bart must expose exactly 3 config fields"); + trtmc::RunConfig config; + if (argc >= 5) { + config.set("max_new_tokens", std::int64_t(std::stoll(argv[4]))); + } + return task.run({trtmc::TextSource{std::string(argv[3])}}, config); + }(); + std::cout << "{\"task\":\"conditional_text_generation\",\"text\":"; + json_string(result.text()); + std::cout << ",\"tokens\":["; + const auto tokens = result.token_ids(); + for (std::size_t i = 0; i < tokens.size(); ++i) { + if (i) + std::cout << ','; + std::cout << tokens[i]; + } + std::cout << "]}\n"; + return 0; + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/families/bart/tests/test_e2e.py b/families/bart/tests/test_e2e.py index 1c51654ea3..2fe18b7912 100644 --- a/families/bart/tests/test_e2e.py +++ b/families/bart/tests/test_e2e.py @@ -15,7 +15,7 @@ from tensorrt_model_connect import BuildRequest, build FAMILY = "bart" -TASKS = frozenset({"text_generation"}) +TASKS = frozenset({"conditional_text_generation", "text_generation"}) TEST_ROOT = Path(__file__).resolve().parent MANIFEST_ROOT = TEST_ROOT / "manifests" THRESHOLD_ROOT = TEST_ROOT / "thresholds" @@ -308,6 +308,38 @@ def _assert_parity(actual, expected, manifest: dict, case: dict, thresholds: dic assert _edit_distance(actual_text, expected_text) <= limit +def _assert_sdk_consumers(runtime_root: Path, bundle: Path, case: dict, expected_text: str) -> None: + native_build = _required_path(os.environ.get("TRTMC_NATIVE_BUILD_DIR"), "TRTMC_NATIVE_BUILD_DIR") + env = os.environ.copy() + env["LD_LIBRARY_PATH"] = ":".join( + value for value in (str(runtime_root), env.get("LD_LIBRARY_PATH", "")) if value + ) + outputs = [] + for language in ("c", "cpp"): + consumer = native_build / f"test_bart_sdk_{language}" + assert consumer.is_file(), f"build the family-owned SDK consumer: {consumer.name}" + max_tokens = str(int(case.get("max_new_tokens", 20))) + completed = subprocess.run( + [str(consumer), str(bundle), str(runtime_root), _case_text(case), max_tokens], + check=True, + capture_output=True, + text=True, + env=env, + timeout=600, + ) + payloads = [line for line in completed.stdout.splitlines() if line.startswith("{")] + assert len(payloads) == 1, f"SDK {language} consumer must return one result" + actual = json.loads(payloads[0]) + record_evidence(f"sdk_{language}", actual) + assert actual["task"] == "conditional_text_generation" + assert "text" in actual + assert "tokens" in actual and len(actual["tokens"]) > 0 + if expected_text: + assert actual["text"] == expected_text + outputs.append(actual) + assert outputs[0] == outputs[1] + + def test_official_checkpoint_e2e(case_name: str, tmp_path: Path) -> None: _, manifest, case = CASES[case_name] record_evidence("inputs", {"manifest": manifest, "case": CASES[case_name][-1]}) @@ -325,3 +357,6 @@ def test_official_checkpoint_e2e(case_name: str, tmp_path: Path) -> None: record_evidence("reference", expected) with evidence_stage("compare"): _assert_parity(actual, expected, manifest, case, record_evidence("thresholds", _thresholds(case_name))) + if int(manifest["tensor_parallel_size"]) == 1: + with evidence_stage("sdk"): + _assert_sdk_consumers(runtime_root, bundle, case, str(actual.get("text", ""))) diff --git a/families/bart/tests/test_model.py b/families/bart/tests/test_model.py new file mode 100644 index 0000000000..073712e4ee --- /dev/null +++ b/families/bart/tests/test_model.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Focused contracts for the bart Task SDK migration.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +try: + import tensorrt # noqa: F401 +except ModuleNotFoundError: + sys.modules["tensorrt"] = SimpleNamespace() + +try: + import ml_dtypes # noqa: F401 +except ModuleNotFoundError: + sys.modules["ml_dtypes"] = SimpleNamespace(bfloat16=object()) + +try: + import safetensors # noqa: F401 +except ModuleNotFoundError: + sys.modules["safetensors"] = SimpleNamespace(safe_open=None) + +from families.bart import model # noqa: E402 +from tensorrt_model_connect import BuildRequest # noqa: E402 + + +def _checkpoint(tmp_path: Path) -> Path: + tmp_path.mkdir(parents=True, exist_ok=True) + (tmp_path / "config.json").write_text( + json.dumps( + { + "model_type": "bart", + "vocab_size": 100, + "hidden_size": 32, + "encoder_layers": 1, + "decoder_layers": 1, + "encoder_attention_heads": 2, + "decoder_attention_heads": 2, + "encoder_ffn_dim": 64, + "decoder_ffn_dim": 64, + "max_position_embeddings": 256, + "decoder_start_token_id": 2, + "bos_token_id": 0, + "eos_token_id": 2, + "pad_token_id": 1, + } + ), + encoding="utf-8", + ) + return tmp_path + + +class _Writer: + def __init__(self) -> None: + self.sections: dict[str, object] = {} + + def set_header(self, **header) -> None: + self.sections["header"] = header + + def add_bytes(self, name: str, value: bytes) -> None: + self.sections[name] = value + + def add_json(self, name: str, value: object) -> None: + self.sections[name] = value + + +def _patch_build(monkeypatch) -> None: + monkeypatch.setattr(model._BartModel, "load_weights", lambda self, *a, **k: {}) + monkeypatch.setattr(model._BartModel, "build_engine", lambda self, *a, **k: b"plan") + monkeypatch.setattr( + model._BartModel, "build_vision_engine", lambda self, *a, **k: b"encoder_plan" + ) + + +@pytest.mark.parametrize("task", ["conditional_text_generation", "text_generation"]) +def test_build_accepts_every_supported_task(tmp_path: Path, monkeypatch, task: str) -> None: + _patch_build(monkeypatch) + writer = _Writer() + + model.build( + BuildRequest( + model_dir=_checkpoint(tmp_path), + output_path=tmp_path / "unused.bundle", + family="bart", + task=task, + precision="fp32", + ), + writer, + ) + + assert writer.sections["header"]["task"] == task + assert writer.sections["engine.plan"] == b"plan" + assert writer.sections["encoder.plan"] == b"encoder_plan" + assert writer.sections["runtime.json"]["tensor_parallel_size"] == 1 + + +@pytest.mark.parametrize( + "task", + [ + "classification", + "text_to_audio", + "text_to_embedding", + "unconditional_text_generation", + ], +) +def test_build_rejects_unsupported_tasks(tmp_path: Path, monkeypatch, task: str) -> None: + _patch_build(monkeypatch) + + with pytest.raises(ValueError, match="bart supports only task="): + model.build( + BuildRequest( + model_dir=_checkpoint(tmp_path), + output_path=tmp_path / "unused.bundle", + family="bart", + task=task, + precision="fp32", + ), + _Writer(), + ) diff --git a/families/bart/tests/test_support.py b/families/bart/tests/test_support.py new file mode 100644 index 0000000000..fabc45c0cd --- /dev/null +++ b/families/bart/tests/test_support.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dependency-free semantic Task discovery for bart.""" + +from families.bart.support import describe +from tensorrt_model_connect.model_support import ModelMetadata + + +def test_primary_task_matches_the_semantic_runtime(): + support = describe(ModelMetadata(config={"model_type": "bart"}, model_index={})) + assert support is not None + assert support.tasks == ("conditional_text_generation", "text_generation") + assert support.default_task == "conditional_text_generation" + + +def test_mbart_model_type_supported(): + support = describe(ModelMetadata(config={"model_type": "mbart"}, model_index={})) + assert support is not None + assert support.tasks == ("conditional_text_generation", "text_generation") + assert support.default_task == "conditional_text_generation" + + +def test_unrelated_identity_is_not_claimed(): + assert describe(ModelMetadata(config={"model_type": "unrelated-model"}, model_index={})) is None +