feat(bart): migrate to semantic Task SDK

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)

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()

families/bart/runtime/pipeline.cpp

@@ -0,0 +1,246 @@
+/*
+ * 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 <algorithm>
+#include <cstring>
+#include <limits>
+#include <stdexcept>
+
+namespace trtmc {
+
+BartPipeline::BartPipeline(std::unique_ptr<ITrtModule> encoder, std::unique_ptr<ITrtModule> decoder,
+                           std::unique_ptr<BartInferenceState> 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<ITokenizer> 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<std::size_t>(max_source_length_) *
+                      static_cast<std::size_t>(hidden_size_) * sizeof(float);
+    cross_k_ptrs_.resize(static_cast<std::size_t>(num_decoder_layers_));
+    cross_v_ptrs_.resize(static_cast<std::size_t>(num_decoder_layers_));
+    for (int32_t i = 0; i < num_decoder_layers_; ++i) {
+        const std::size_t layer = static_cast<std::size_t>(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<internal::TaskInstance> BartPipeline::task_bindings() {
+    return {internal::bind<internal::IConditionalTextGeneration>(*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<std::int64_t>(config, bart_config_fields, "max_new_tokens")) {
+        if (*val <= 0 || *val > std::numeric_limits<int32_t>::max())
+            throw internal::ConfigError(
+                "BartPipeline: max_new_tokens must be positive and fit in int32");
+        max_tokens = static_cast<int32_t>(*val);
+    }
+
+    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<std::vector<int32_t>, int32_t>
+BartPipeline::prepare_encoder_input_semantic(const internal::TextSource& source) {
+    std::vector<int32_t> ids;
+    if (const auto* view = std::get_if<std::string_view>(&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<Span<const std::int32_t>>(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<int32_t>(ids.size()) > max_source_length_)
+        throw std::invalid_argument("BartPipeline: input exceeds maximum sequence length");
+
+    int32_t copy_len = static_cast<int32_t>(ids.size());
+    std::vector<int32_t> padded(static_cast<std::size_t>(max_source_length_), pad_token_id_);
+    std::copy_n(ids.begin(), ids.size(), padded.begin());
+    return {std::move(padded), copy_len};
+}
+
+std::pair<std::vector<int32_t>, int32_t>
+BartPipeline::prepare_encoder_input(const std::string& prompt) {
+    std::vector<int32_t> 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<int32_t>(ids.size()), max_source_length_);
+    std::vector<int32_t> padded(static_cast<std::size_t>(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<int32_t>& 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<std::size_t>(max_source_length_), -1e9f);
+    for (int32_t i = 0; i < actual_len; ++i)
+        encoder_attention_mask_[static_cast<std::size_t>(i)] = 0.0f;
+
+    TensorMap inputs;
+    Tensor ids_tensor;
+    ids_tensor.data = const_cast<int32_t*>(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<std::size_t>(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<std::size_t>(i)].get());
+        decoder_->bind_external("cross_v" + s, cross_v_ptrs_[static_cast<std::size_t>(i)].get());
+    }
+}
+
+std::vector<int32_t> BartPipeline::run_decoder(int32_t max_new_tokens) {
+    state_->reset();
+    state_->bind_to(*decoder_);
+    std::vector<int32_t> output_ids;
+    std::vector<float> 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<float>& 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<std::size_t>(num));
+    std::memcpy(logits.data(), it->second.data, num * sizeof(float));
+    state_->advance();
+}
+
+} // namespace trtmc

families/bart/runtime/pipeline.h

@@ -0,0 +1,86 @@
+/*
+ * 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/task.h"
+#include "trtmc/runtime/trt_module.h"
+
+#include <cstdint>
+#include <cuda_runtime_api.h>
+#include <memory>
+#include <string>
+#include <string_view>
+#include <utility>
+#include <vector>
+
+namespace trtmc {
+
+inline constexpr internal::ConfigField bart_config_fields[] = {
+    {"max_new_tokens", internal::ConfigKind::I64, std::int64_t{128}, "Generation token cap"},
+};
+
+class BartPipeline final : public internal::IModel,
+                           public internal::IConditionalTextGeneration,
+                           public ITextGeneration {
+  public:
+    BartPipeline(std::unique_ptr<ITrtModule> encoder, std::unique_ptr<ITrtModule> decoder,
+                 std::unique_ptr<BartInferenceState> 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<ITokenizer> 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<internal::TaskInstance> 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<std::vector<int32_t>, int32_t> prepare_encoder_input(const std::string& prompt);
+    std::pair<std::vector<int32_t>, int32_t>
+    prepare_encoder_input_semantic(const internal::TextSource& source);
+    void run_encoder(const std::vector<int32_t>& padded_ids, int32_t actual_len);
+    void setup_cross_attention();
+    std::vector<int32_t> run_decoder(int32_t max_new_tokens);
+    void run_decoder_step(int32_t token_id, std::vector<float>& logits);
+
+    std::unique_ptr<ITrtModule> encoder_;
+    std::unique_ptr<ITrtModule> decoder_;
+    std::unique_ptr<BartInferenceState> 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<ITokenizer> tokenizer_;
+    std::string task_;
+    std::string model_id_;
+    std::vector<bart::DeviceBuffer> cross_k_ptrs_;
+    std::vector<bart::DeviceBuffer> cross_v_ptrs_;
+    std::vector<float> encoder_attention_mask_;
+    std::size_t cross_kv_bytes_{0};
+    int32_t actual_enc_len_{0};
+};
+
+} // namespace trtmc