From 92e7b8df719d87b0bd3ab5278bf395316e272854 Mon Sep 17 00:00:00 2001 From: Mohak Gupta Date: Thu, 10 Sep 2026 10:09:16 +0530 Subject: [PATCH 1/2] fix(marian): check cudaMalloc status and own the device buffers MarianPipeline discarded the result of every cudaMalloc. The constructor allocated a cross-attention key and value buffer per decoder layer, and generate() lazily allocated the encoder mask, none of them checked. A failed allocation left a null pointer that the pipeline then bound to the decoder as if it were real memory. Checking the status alone would not have been enough: the buffers lived in std::vector, destroying that vector does not free the device allocations its elements point at, and a constructor that throws never runs its own destructor - so any layer failing after the first would leak everything allocated before it. Each buffer, including the encoder mask, is now held in a family-local DeviceBuffer, so cleanup happens during unwinding. The explicit destructor is now redundant and removed. Adds test_marian_cross_kv_alloc, failing each of the eight allocations in turn against CPU CUDA stubs and asserting which buffer the error names. Verified it fails against the previous raw-pointer shape and passes here. Needs no GPU and links no cudart. Same treatment as bark #1201, bart #1222, whisper #1223, m2m_100 #1232 and t5, per the review on #1221. Signed-off-by: Mohak Gupta --- families/marian/runtime/CMakeLists.txt | 17 +++ families/marian/runtime/device_buffer.h | 51 +++++++ families/marian/runtime/plugin.cpp | 56 ++++---- .../tests/cpp/test_marian_cross_kv_alloc.cpp | 129 ++++++++++++++++++ 4 files changed, 224 insertions(+), 29 deletions(-) create mode 100644 families/marian/runtime/device_buffer.h create mode 100644 families/marian/tests/cpp/test_marian_cross_kv_alloc.cpp diff --git a/families/marian/runtime/CMakeLists.txt b/families/marian/runtime/CMakeLists.txt index c7db7b1fc3..42bb62f3b7 100644 --- a/families/marian/runtime/CMakeLists.txt +++ b/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() diff --git a/families/marian/runtime/device_buffer.h b/families/marian/runtime/device_buffer.h new file mode 100644 index 0000000000..dfdec74dee --- /dev/null +++ b/families/marian/runtime/device_buffer.h @@ -0,0 +1,51 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +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}; +}; + +} // namespace marian +} // namespace trtmc diff --git a/families/marian/runtime/plugin.cpp b/families/marian/runtime/plugin.cpp index e2b75c30e6..6f7b7bbd5f 100644 --- a/families/marian/runtime/plugin.cpp +++ b/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,25 +188,19 @@ class MarianPipeline final : public ITextGeneration { model_id_(std::move(model_id_str)) { cross_kv_bytes_ = static_cast(max_enc_seq_len_) * 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) { - 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); + const size_t layer = static_cast(i); + if (cross_k_ptrs_[layer].allocate(cross_kv_bytes_) != cudaSuccess) + throw std::runtime_error( + "MarianPipeline: unable to allocate cross-attention key buffer"); + if (cross_v_ptrs_[layer].allocate(cross_kv_bytes_) != cudaSuccess) + throw std::runtime_error( + "MarianPipeline: unable to allocate cross-attention value buffer"); } } - ~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_); - } - TextResult generate(const std::string& prompt, const TextGenerationConfig& cfg) override { if (!tokenizer_) throw std::runtime_error("MarianPipeline: no tokenizer configured"); @@ -267,10 +262,12 @@ 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(i)], encoder_output_host_.data(), - cross_kv_bytes_, cudaMemcpyHostToDevice, stream_); - cudaMemcpyAsync(cross_v_ptrs_[static_cast(i)], encoder_output_host_.data(), - cross_kv_bytes_, cudaMemcpyHostToDevice, stream_); + cudaMemcpyAsync(cross_k_ptrs_[static_cast(i)].get(), + encoder_output_host_.data(), cross_kv_bytes_, cudaMemcpyHostToDevice, + stream_); + cudaMemcpyAsync(cross_v_ptrs_[static_cast(i)].get(), + encoder_output_host_.data(), cross_kv_bytes_, cudaMemcpyHostToDevice, + stream_); } cudaStreamSynchronize(stream_); @@ -278,19 +275,20 @@ class MarianPipeline final : public ITextGeneration { for (int32_t i = 0; i < actual_enc_len_; ++i) enc_mask_host[static_cast(i)] = 0.0f; size_t mask_bytes = static_cast(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_); + if (enc_mask_device_.get() == nullptr && + enc_mask_device_.allocate(mask_bytes) != cudaSuccess) + throw std::runtime_error("MarianPipeline: unable to allocate encoder mask buffer"); + 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(i)]); - decoder_->bind_external(cv_name, cross_v_ptrs_[static_cast(i)]); + decoder_->bind_external(ck_name, cross_k_ptrs_[static_cast(i)].get()); + decoder_->bind_external(cv_name, cross_v_ptrs_[static_cast(i)].get()); } - decoder_->bind_external("encoder_mask", enc_mask_device_); + decoder_->bind_external("encoder_mask", enc_mask_device_.get()); } std::vector run_decoder(int32_t max_new_tokens, int32_t eos_id) { @@ -362,13 +360,13 @@ class MarianPipeline final : public ITextGeneration { std::shared_ptr tokenizer_; std::string model_id_; - std::vector cross_k_ptrs_; - std::vector cross_v_ptrs_; + std::vector cross_k_ptrs_; + std::vector cross_v_ptrs_; size_t cross_kv_bytes_{0}; std::vector 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) { diff --git a/families/marian/tests/cpp/test_marian_cross_kv_alloc.cpp b/families/marian/tests/cpp/test_marian_cross_kv_alloc.cpp new file mode 100644 index 0000000000..e2db156f86 --- /dev/null +++ b/families/marian/tests/cpp/test_marian_cross_kv_alloc.cpp @@ -0,0 +1,129 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Exercises the ownership the Marian pipeline uses for its cross-attention K/V +// buffers, against CPU CUDA stubs so each allocation can be failed in turn +// without a GPU. MarianPipeline itself needs live TensorRT modules to build, so +// this drives the same allocation loop over the same owning type. + +#include "families/marian/runtime/device_buffer.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +int g_fail_on_allocation = 0; +int g_allocation_count = 0; +std::set 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; + } +} + +// The allocation loop from MarianPipeline's constructor, over the same type. +void allocate_cross_kv(std::vector& keys, + std::vector& values, int32_t layers, + std::size_t bytes) { + keys.resize(static_cast(layers)); + values.resize(static_cast(layers)); + for (int32_t i = 0; i < layers; ++i) { + const std::size_t layer = static_cast(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"); + } +} + +} // 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(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 keys; + std::vector values; + 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 keys; + std::vector values; + allocate_cross_kv(keys, values, layers, bytes); + check(g_outstanding.size() == static_cast(2 * layers), + "every layer should hold a key and a value buffer"); + } + check(g_outstanding.empty(), "leaving scope should release every buffer"); + + 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; +} From f0befc97df4b87da57ec40cf1ef67696ec559560 Mon Sep 17 00:00:00 2001 From: Mohak Gupta Date: Fri, 11 Sep 2026 10:28:17 +0530 Subject: [PATCH 2/2] test(marian): drive the pipeline's own allocation paths, not a copy The test reproduced the constructor's allocation loop locally and only exercised the reproduction, so a regression in the real constructor or in the encoder mask allocation would have passed it. Moves both paths into family-local helpers - allocate_cross_kv and ensure_encoder_mask - that the pipeline calls and the test now calls directly, and adds coverage for the encoder mask: a failure throws and holds nothing, a success allocates once, a repeat call reuses it. Confirmed the test now sees production changes: dropping the status check inside allocate_cross_kv fails four checks. Signed-off-by: Mohak Gupta --- families/marian/runtime/device_buffer.h | 25 ++++++++ families/marian/runtime/plugin.cpp | 17 +----- .../tests/cpp/test_marian_cross_kv_alloc.cpp | 59 +++++++++++-------- 3 files changed, 64 insertions(+), 37 deletions(-) diff --git a/families/marian/runtime/device_buffer.h b/families/marian/runtime/device_buffer.h index dfdec74dee..3232dfadc6 100644 --- a/families/marian/runtime/device_buffer.h +++ b/families/marian/runtime/device_buffer.h @@ -6,7 +6,10 @@ #pragma once #include +#include #include +#include +#include namespace trtmc { namespace marian { @@ -47,5 +50,27 @@ class DeviceBuffer { 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& keys, std::vector& values, + int32_t layers, std::size_t bytes) { + keys.resize(static_cast(layers)); + values.resize(static_cast(layers)); + for (int32_t i = 0; i < layers; ++i) { + const std::size_t layer = static_cast(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 diff --git a/families/marian/runtime/plugin.cpp b/families/marian/runtime/plugin.cpp index 6f7b7bbd5f..704e6864b4 100644 --- a/families/marian/runtime/plugin.cpp +++ b/families/marian/runtime/plugin.cpp @@ -188,17 +188,8 @@ class MarianPipeline final : public ITextGeneration { model_id_(std::move(model_id_str)) { cross_kv_bytes_ = static_cast(max_enc_seq_len_) * 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 size_t layer = static_cast(i); - if (cross_k_ptrs_[layer].allocate(cross_kv_bytes_) != cudaSuccess) - throw std::runtime_error( - "MarianPipeline: unable to allocate cross-attention key buffer"); - if (cross_v_ptrs_[layer].allocate(cross_kv_bytes_) != cudaSuccess) - throw std::runtime_error( - "MarianPipeline: unable to allocate cross-attention value buffer"); - } + 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 { @@ -275,9 +266,7 @@ class MarianPipeline final : public ITextGeneration { for (int32_t i = 0; i < actual_enc_len_; ++i) enc_mask_host[static_cast(i)] = 0.0f; size_t mask_bytes = static_cast(max_enc_seq_len_) * sizeof(float); - if (enc_mask_device_.get() == nullptr && - enc_mask_device_.allocate(mask_bytes) != cudaSuccess) - throw std::runtime_error("MarianPipeline: unable to allocate encoder mask buffer"); + marian::ensure_encoder_mask(enc_mask_device_, mask_bytes); cudaMemcpyAsync(enc_mask_device_.get(), enc_mask_host.data(), mask_bytes, cudaMemcpyHostToDevice, stream_); cudaStreamSynchronize(stream_); diff --git a/families/marian/tests/cpp/test_marian_cross_kv_alloc.cpp b/families/marian/tests/cpp/test_marian_cross_kv_alloc.cpp index e2db156f86..81100325bc 100644 --- a/families/marian/tests/cpp/test_marian_cross_kv_alloc.cpp +++ b/families/marian/tests/cpp/test_marian_cross_kv_alloc.cpp @@ -3,10 +3,10 @@ * SPDX-License-Identifier: Apache-2.0 */ -// Exercises the ownership the Marian pipeline uses for its cross-attention K/V -// buffers, against CPU CUDA stubs so each allocation can be failed in turn -// without a GPU. MarianPipeline itself needs live TensorRT modules to build, so -// this drives the same allocation loop over the same owning type. +// 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" @@ -33,23 +33,6 @@ void check(bool condition, const char* what) { } } -// The allocation loop from MarianPipeline's constructor, over the same type. -void allocate_cross_kv(std::vector& keys, - std::vector& values, int32_t layers, - std::size_t bytes) { - keys.resize(static_cast(layers)); - values.resize(static_cast(layers)); - for (int32_t i = 0; i < layers; ++i) { - const std::size_t layer = static_cast(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"); - } -} - } // namespace extern "C" { @@ -90,7 +73,7 @@ int main() { try { std::vector keys; std::vector values; - allocate_cross_kv(keys, values, layers, bytes); + 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 @@ -114,12 +97,42 @@ int main() { { std::vector keys; std::vector values; - allocate_cross_kv(keys, values, layers, bytes); + trtmc::marian::allocate_cross_kv(keys, values, layers, bytes); check(g_outstanding.size() == static_cast(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;