fix(marian): check cudaMalloc status for cross-attention and encoder mask buffers

families/marian/runtime/CMakeLists.txt

@@ -47,4 +47,21 @@ if(TRTMC_BUILD_TESTS)
   target_link_libraries(test_marian_runtime_config PRIVATE trtmc_model_marian)
   target_compile_options(test_marian_runtime_config PRIVATE -Wall -Wextra -Wpedantic)
   add_test(NAME marian_runtime_config COMMAND test_marian_runtime_config)
+
+  # CPU CUDA stubs live in the test, so allocation failures can be injected
+  # without a GPU or cudart.
+  add_executable(test_marian_cross_kv_alloc
+    ${PROJECT_SOURCE_DIR}/families/marian/tests/cpp/test_marian_cross_kv_alloc.cpp
+  )
+  target_include_directories(test_marian_cross_kv_alloc PRIVATE
+    ${PROJECT_SOURCE_DIR}
+    ${PROJECT_SOURCE_DIR}/core/runtime/include
+  )
+  target_include_directories(test_marian_cross_kv_alloc SYSTEM PRIVATE
+    ${TRTMC_CUDA_INCLUDE_DIR}
+  )
+  target_compile_options(test_marian_cross_kv_alloc PRIVATE
+    -Wall -Wextra -Wpedantic
+  )
+  add_test(NAME marian_cross_kv_alloc COMMAND test_marian_cross_kv_alloc)
 endif()

families/marian/runtime/device_buffer.h

@@ -0,0 +1,76 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <cuda_runtime_api.h>
+#include <stdexcept>
+#include <vector>
+
+namespace trtmc {
+namespace marian {
+
+// Owns one device allocation. Holding the cross-attention buffers in these
+// rather than in raw pointers means a constructor that throws part way through
+// still releases what it already allocated: members that finished constructing
+// are destroyed during unwinding, and a raw pointer in a vector is not.
+class DeviceBuffer {
+  public:
+    DeviceBuffer() = default;
+
+    ~DeviceBuffer() { cudaFree(ptr_); }
+
+    DeviceBuffer(const DeviceBuffer&) = delete;
+    DeviceBuffer& operator=(const DeviceBuffer&) = delete;
+
+    DeviceBuffer(DeviceBuffer&& other) noexcept : ptr_(other.ptr_) { other.ptr_ = nullptr; }
+
+    DeviceBuffer& operator=(DeviceBuffer&& other) noexcept {
+        if (this != &other) {
+            cudaFree(ptr_);
+            ptr_ = other.ptr_;
+            other.ptr_ = nullptr;
+        }
+        return *this;
+    }
+
+    cudaError_t allocate(std::size_t bytes) {
+        cudaFree(ptr_);
+        ptr_ = nullptr;
+        return cudaMalloc(&ptr_, bytes);
+    }
+
+    void* get() const { return ptr_; }
+
+  private:
+    void* ptr_{nullptr};
+};
+
+// The pipeline's own allocation paths, kept here so the test can drive the
+// same code the constructor and setup_cross_attention run.
+inline void allocate_cross_kv(std::vector<DeviceBuffer>& keys, std::vector<DeviceBuffer>& values,
+                              int32_t layers, std::size_t bytes) {
+    keys.resize(static_cast<std::size_t>(layers));
+    values.resize(static_cast<std::size_t>(layers));
+    for (int32_t i = 0; i < layers; ++i) {
+        const std::size_t layer = static_cast<std::size_t>(i);
+        if (keys[layer].allocate(bytes) != cudaSuccess)
+            throw std::runtime_error(
+                "MarianPipeline: unable to allocate cross-attention key buffer");
+        if (values[layer].allocate(bytes) != cudaSuccess)
+            throw std::runtime_error(
+                "MarianPipeline: unable to allocate cross-attention value buffer");
+    }
+}
+
+inline void ensure_encoder_mask(DeviceBuffer& mask, std::size_t bytes) {
+    if (mask.get() == nullptr && mask.allocate(bytes) != cudaSuccess)
+        throw std::runtime_error("MarianPipeline: unable to allocate encoder mask buffer");
+}
+
+} // namespace marian
+} // namespace trtmc

families/marian/runtime/plugin.cpp

@@ -12,6 +12,7 @@
 //   3. Run decoder autoregressively with cross-attention to encoder output
 //   4. Detokenize output
 
+#include "families/marian/runtime/device_buffer.h"
 #include "families/marian/runtime/distributed_runtime.h"
 #include "families/marian/runtime/kv_cache.h"
 #include "families/marian/runtime/plugin_helpers.h"
@@ -187,23 +188,8 @@ class MarianPipeline final : public ITextGeneration {
           model_id_(std::move(model_id_str)) {
         cross_kv_bytes_ = static_cast<size_t>(max_enc_seq_len_) *
                           static_cast<size_t>(hidden_size_) * sizeof(float);
-        for (int32_t i = 0; i < num_decoder_layers_; ++i) {
-            void* dk = nullptr;
-            void* dv = nullptr;
-            cudaMalloc(&dk, cross_kv_bytes_);
-            cudaMalloc(&dv, cross_kv_bytes_);
-            cross_k_ptrs_.push_back(dk);
-            cross_v_ptrs_.push_back(dv);
-        }
-    }
-
-    ~MarianPipeline() override {
-        for (auto* p : cross_k_ptrs_)
-            cudaFree(p);
-        for (auto* p : cross_v_ptrs_)
-            cudaFree(p);
-        if (enc_mask_device_)
-            cudaFree(enc_mask_device_);
+        marian::allocate_cross_kv(cross_k_ptrs_, cross_v_ptrs_, num_decoder_layers_,
+                                  cross_kv_bytes_);
     }
 
     TextResult generate(const std::string& prompt, const TextGenerationConfig& cfg) override {
@@ -267,30 +253,31 @@ class MarianPipeline final : public ITextGeneration {
 
     void setup_cross_attention() {
         for (int32_t i = 0; i < num_decoder_layers_; ++i) {
-            cudaMemcpyAsync(cross_k_ptrs_[static_cast<size_t>(i)], encoder_output_host_.data(),
-                            cross_kv_bytes_, cudaMemcpyHostToDevice, stream_);
-            cudaMemcpyAsync(cross_v_ptrs_[static_cast<size_t>(i)], encoder_output_host_.data(),
-                            cross_kv_bytes_, cudaMemcpyHostToDevice, stream_);
+            cudaMemcpyAsync(cross_k_ptrs_[static_cast<size_t>(i)].get(),
+                            encoder_output_host_.data(), cross_kv_bytes_, cudaMemcpyHostToDevice,
+                            stream_);
+            cudaMemcpyAsync(cross_v_ptrs_[static_cast<size_t>(i)].get(),
+                            encoder_output_host_.data(), cross_kv_bytes_, cudaMemcpyHostToDevice,
+                            stream_);
         }
         cudaStreamSynchronize(stream_);
 
         std::vector<float> enc_mask_host(static_cast<size_t>(max_enc_seq_len_), -1e9f);
         for (int32_t i = 0; i < actual_enc_len_; ++i)
             enc_mask_host[static_cast<size_t>(i)] = 0.0f;
         size_t mask_bytes = static_cast<size_t>(max_enc_seq_len_) * sizeof(float);
-        if (!enc_mask_device_)
-            cudaMalloc(&enc_mask_device_, mask_bytes);
-        cudaMemcpyAsync(enc_mask_device_, enc_mask_host.data(), mask_bytes, cudaMemcpyHostToDevice,
-                        stream_);
+        marian::ensure_encoder_mask(enc_mask_device_, mask_bytes);
+        cudaMemcpyAsync(enc_mask_device_.get(), enc_mask_host.data(), mask_bytes,
+                        cudaMemcpyHostToDevice, stream_);
         cudaStreamSynchronize(stream_);
 
         for (int32_t i = 0; i < num_decoder_layers_; ++i) {
             std::string ck_name = "cross_k_" + std::to_string(i);
             std::string cv_name = "cross_v_" + std::to_string(i);
-            decoder_->bind_external(ck_name, cross_k_ptrs_[static_cast<size_t>(i)]);
-            decoder_->bind_external(cv_name, cross_v_ptrs_[static_cast<size_t>(i)]);
+            decoder_->bind_external(ck_name, cross_k_ptrs_[static_cast<size_t>(i)].get());
+            decoder_->bind_external(cv_name, cross_v_ptrs_[static_cast<size_t>(i)].get());
         }
-        decoder_->bind_external("encoder_mask", enc_mask_device_);
+        decoder_->bind_external("encoder_mask", enc_mask_device_.get());
     }
 
     std::vector<int32_t> run_decoder(int32_t max_new_tokens, int32_t eos_id) {
@@ -362,13 +349,13 @@ class MarianPipeline final : public ITextGeneration {
     std::shared_ptr<ITokenizer> tokenizer_;
     std::string model_id_;
 
-    std::vector<void*> cross_k_ptrs_;
-    std::vector<void*> cross_v_ptrs_;
+    std::vector<marian::DeviceBuffer> cross_k_ptrs_;
+    std::vector<marian::DeviceBuffer> cross_v_ptrs_;
     size_t cross_kv_bytes_{0};
 
     std::vector<float> encoder_output_host_;
     int32_t actual_enc_len_{0};
-    void* enc_mask_device_{nullptr};
+    marian::DeviceBuffer enc_mask_device_;
 };
 
 ITask* create_marian(const FamilyContext& context) {

families/marian/tests/cpp/test_marian_cross_kv_alloc.cpp

@@ -0,0 +1,142 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Drives the Marian pipeline's own allocation paths - allocate_cross_kv, which
+// the constructor runs, and ensure_encoder_mask, which setup_cross_attention
+// runs - against CPU CUDA stubs so each allocation can be failed in turn
+// without a GPU. MarianPipeline itself needs live TensorRT modules to build.
+
+#include "families/marian/runtime/device_buffer.h"
+
+#include <cstdint>
+#include <cstdio>
+#include <cuda_runtime.h>
+#include <set>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace {
+
+int g_fail_on_allocation = 0;
+int g_allocation_count = 0;
+std::set<void*> g_outstanding;
+std::uintptr_t g_next_address = 0x1000;
+int g_failures = 0;
+
+void check(bool condition, const char* what) {
+    if (!condition) {
+        std::fprintf(stderr, "FAIL: %s\n", what);
+        ++g_failures;
+    }
+}
+
+} // namespace
+
+extern "C" {
+
+cudaError_t cudaMalloc(void** devPtr, size_t size) {
+    (void)size;
+    ++g_allocation_count;
+    if (g_fail_on_allocation != 0 && g_allocation_count == g_fail_on_allocation) {
+        *devPtr = nullptr;
+        return cudaErrorMemoryAllocation;
+    }
+    void* address = reinterpret_cast<void*>(g_next_address);
+    g_next_address += 0x1000;
+    g_outstanding.insert(address);
+    *devPtr = address;
+    return cudaSuccess;
+}
+
+cudaError_t cudaFree(void* devPtr) {
+    if (devPtr != nullptr) {
+        g_outstanding.erase(devPtr);
+    }
+    return cudaSuccess;
+}
+
+} // extern "C"
+
+int main() {
+    const int32_t layers = 4;
+    const std::size_t bytes = 1024;
+
+    // Four layers means eight allocations. Fail each in turn; whatever the loop
+    // acquired before the failure must be released as the exception unwinds.
+    for (int failing = 1; failing <= 2 * layers; ++failing) {
+        g_fail_on_allocation = failing;
+        g_allocation_count = 0;
+        bool threw = false;
+        try {
+            std::vector<trtmc::marian::DeviceBuffer> keys;
+            std::vector<trtmc::marian::DeviceBuffer> values;
+            trtmc::marian::allocate_cross_kv(keys, values, layers, bytes);
+        } catch (const std::runtime_error& error) {
+            threw = true;
+            // Allocations alternate key, value, so an odd failure point is a
+            // key buffer and an even one is a value buffer.
+            const char* expected =
+                (failing % 2) == 1
+                    ? "MarianPipeline: unable to allocate cross-attention key buffer"
+                    : "MarianPipeline: unable to allocate cross-attention value buffer";
+            check(std::string(error.what()) == expected,
+                  "the error should name the buffer that actually failed");
+        }
+        check(threw, "a failed cross-attention allocation should throw");
+        check(g_outstanding.empty(),
+              "a failed allocation must release everything already acquired");
+        g_outstanding.clear();
+    }
+
+    // With no injected failure the buffers are held, then released on scope exit.
+    g_fail_on_allocation = 0;
+    g_allocation_count = 0;
+    {
+        std::vector<trtmc::marian::DeviceBuffer> keys;
+        std::vector<trtmc::marian::DeviceBuffer> values;
+        trtmc::marian::allocate_cross_kv(keys, values, layers, bytes);
+        check(g_outstanding.size() == static_cast<std::size_t>(2 * layers),
+              "every layer should hold a key and a value buffer");
+    }
+    check(g_outstanding.empty(), "leaving scope should release every buffer");
+
+    // The encoder mask is allocated lazily by setup_cross_attention. A failure
+    // must throw and hold nothing; a success must allocate exactly once, and a
+    // repeat call must reuse it rather than allocate again.
+    {
+        trtmc::marian::DeviceBuffer mask;
+        g_fail_on_allocation = 1;
+        g_allocation_count = 0;
+        bool mask_threw = false;
+        try {
+            trtmc::marian::ensure_encoder_mask(mask, bytes);
+        } catch (const std::runtime_error& error) {
+            mask_threw = true;
+            check(std::string(error.what()) ==
+                      "MarianPipeline: unable to allocate encoder mask buffer",
+                  "the encoder mask error should name the encoder mask");
+        }
+        check(mask_threw, "a failed encoder mask allocation should throw");
+        check(mask.get() == nullptr && g_outstanding.empty(),
+              "a failed encoder mask allocation must hold nothing");
+
+        g_fail_on_allocation = 0;
+        g_allocation_count = 0;
+        trtmc::marian::ensure_encoder_mask(mask, bytes);
+        check(mask.get() != nullptr && g_outstanding.size() == 1,
+              "the encoder mask should allocate once");
+        trtmc::marian::ensure_encoder_mask(mask, bytes);
+        check(g_allocation_count == 1, "a second call must reuse the mask, not reallocate");
+    }
+    check(g_outstanding.empty(), "leaving scope should release the encoder mask");
+
+    if (g_failures != 0) {
+        std::fprintf(stderr, "%d check(s) failed\n", g_failures);
+        return 1;
+    }
+    std::printf("marian cross-attention allocation-failure checks passed\n");
+    return 0;
+}