From 5d6aa5a75703cd9f6cf7be907503c880323375aa Mon Sep 17 00:00:00 2001 From: Mohak Gupta Date: Thu, 10 Sep 2026 10:08:13 +0530 Subject: [PATCH 1/2] fix(t5): check cudaMalloc status and own the device buffers T5Pipeline 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_t5_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 and m2m_100 Signed-off-by: Mohak Gupta #1232, per the review on #1221. --- families/t5/runtime/CMakeLists.txt | 17 +++ families/t5/runtime/device_buffer.h | 51 +++++++ families/t5/runtime/plugin.cpp | 56 ++++---- .../t5/tests/cpp/test_t5_cross_kv_alloc.cpp | 126 ++++++++++++++++++ 4 files changed, 221 insertions(+), 29 deletions(-) create mode 100644 families/t5/runtime/device_buffer.h create mode 100644 families/t5/tests/cpp/test_t5_cross_kv_alloc.cpp diff --git a/families/t5/runtime/CMakeLists.txt b/families/t5/runtime/CMakeLists.txt index 23ba08440c..49ae234c0a 100644 --- a/families/t5/runtime/CMakeLists.txt +++ b/families/t5/runtime/CMakeLists.txt @@ -47,4 +47,21 @@ if(TRTMC_BUILD_TESTS) target_link_libraries(test_t5_runtime_config PRIVATE trtmc_model_t5) target_compile_options(test_t5_runtime_config PRIVATE -Wall -Wextra -Wpedantic) add_test(NAME t5_runtime_config COMMAND test_t5_runtime_config) + + # CPU CUDA stubs live in the test, so allocation failures can be injected + # without a GPU or cudart. + add_executable(test_t5_cross_kv_alloc + ${PROJECT_SOURCE_DIR}/families/t5/tests/cpp/test_t5_cross_kv_alloc.cpp + ) + target_include_directories(test_t5_cross_kv_alloc PRIVATE + ${PROJECT_SOURCE_DIR} + ${PROJECT_SOURCE_DIR}/core/runtime/include + ) + target_include_directories(test_t5_cross_kv_alloc SYSTEM PRIVATE + ${TRTMC_CUDA_INCLUDE_DIR} + ) + target_compile_options(test_t5_cross_kv_alloc PRIVATE + -Wall -Wextra -Wpedantic + ) + add_test(NAME t5_cross_kv_alloc COMMAND test_t5_cross_kv_alloc) endif() diff --git a/families/t5/runtime/device_buffer.h b/families/t5/runtime/device_buffer.h new file mode 100644 index 0000000000..264cda5afd --- /dev/null +++ b/families/t5/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 t5 { + +// 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 t5 +} // namespace trtmc diff --git a/families/t5/runtime/plugin.cpp b/families/t5/runtime/plugin.cpp index 8e3e617b4d..956bfecea8 100644 --- a/families/t5/runtime/plugin.cpp +++ b/families/t5/runtime/plugin.cpp @@ -12,6 +12,7 @@ // 3. Run decoder autoregressively with cross-attention to encoder output // 4. Detokenize output +#include "families/t5/runtime/device_buffer.h" #include "families/t5/runtime/distributed_runtime.h" #include "families/t5/runtime/kv_cache.h" #include "families/t5/runtime/plugin_helpers.h" @@ -186,25 +187,19 @@ class T5Pipeline final : public ITextGeneration { // Allocate cross-attention device buffers (one per decoder layer) 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( + "T5Pipeline: unable to allocate cross-attention key buffer"); + if (cross_v_ptrs_[layer].allocate(cross_kv_bytes_) != cudaSuccess) + throw std::runtime_error( + "T5Pipeline: unable to allocate cross-attention value buffer"); } } - ~T5Pipeline() 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("T5Pipeline: no tokenizer configured"); @@ -281,10 +276,12 @@ class T5Pipeline final : public ITextGeneration { // Copy raw encoder output to all decoder layer cross_k/cross_v inputs. // The per-layer K/V projections are baked into the decoder TRT graph. 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_); @@ -293,21 +290,22 @@ class T5Pipeline 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("T5Pipeline: unable to allocate encoder mask buffer"); + cudaMemcpyAsync(enc_mask_device_.get(), enc_mask_host.data(), mask_bytes, + cudaMemcpyHostToDevice, stream_); cudaStreamSynchronize(stream_); // Bind cross-attention buffers to decoder 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()); } // Bind encoder mask - 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) { @@ -380,15 +378,15 @@ class T5Pipeline final : public ITextGeneration { std::string model_id_; // Cross-attention device buffers - 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}; // Encoder output (host copy) std::vector encoder_output_host_; int32_t actual_enc_len_{0}; // Encoder attention mask (device) - void* enc_mask_device_{nullptr}; + t5::DeviceBuffer enc_mask_device_; }; ITask* create_t5(const FamilyContext& context) { diff --git a/families/t5/tests/cpp/test_t5_cross_kv_alloc.cpp b/families/t5/tests/cpp/test_t5_cross_kv_alloc.cpp new file mode 100644 index 0000000000..b1a6940c28 --- /dev/null +++ b/families/t5/tests/cpp/test_t5_cross_kv_alloc.cpp @@ -0,0 +1,126 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Exercises the ownership the T5 pipeline uses for its cross-attention K/V +// buffers, against CPU CUDA stubs so each allocation can be failed in turn +// without a GPU. T5Pipeline itself needs live TensorRT modules to build, so +// this drives the same allocation loop over the same owning type. + +#include "families/t5/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 T5Pipeline'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("T5Pipeline: unable to allocate cross-attention key buffer"); + if (values[layer].allocate(bytes) != cudaSuccess) + throw std::runtime_error("T5Pipeline: 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 ? "T5Pipeline: unable to allocate cross-attention key buffer" + : "T5Pipeline: 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("t5 cross-attention allocation-failure checks passed\n"); + return 0; +} From 8cf3f72d66aa63881eb1c3a0972c47ec0f2b0554 Mon Sep 17 00:00:00 2001 From: Mohak Gupta Date: Fri, 11 Sep 2026 10:30:22 +0530 Subject: [PATCH 2/2] test(t5): 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/t5/runtime/device_buffer.h | 23 ++++++++ families/t5/runtime/plugin.cpp | 16 +----- .../t5/tests/cpp/test_t5_cross_kv_alloc.cpp | 56 ++++++++++++------- 3 files changed, 60 insertions(+), 35 deletions(-) diff --git a/families/t5/runtime/device_buffer.h b/families/t5/runtime/device_buffer.h index 264cda5afd..6e15718cae 100644 --- a/families/t5/runtime/device_buffer.h +++ b/families/t5/runtime/device_buffer.h @@ -6,7 +6,10 @@ #pragma once #include +#include #include +#include +#include namespace trtmc { namespace t5 { @@ -47,5 +50,25 @@ 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("T5Pipeline: unable to allocate cross-attention key buffer"); + if (values[layer].allocate(bytes) != cudaSuccess) + throw std::runtime_error("T5Pipeline: 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("T5Pipeline: unable to allocate encoder mask buffer"); +} + } // namespace t5 } // namespace trtmc diff --git a/families/t5/runtime/plugin.cpp b/families/t5/runtime/plugin.cpp index 956bfecea8..f6e7ee4361 100644 --- a/families/t5/runtime/plugin.cpp +++ b/families/t5/runtime/plugin.cpp @@ -187,17 +187,7 @@ class T5Pipeline final : public ITextGeneration { // Allocate cross-attention device buffers (one per decoder layer) 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( - "T5Pipeline: unable to allocate cross-attention key buffer"); - if (cross_v_ptrs_[layer].allocate(cross_kv_bytes_) != cudaSuccess) - throw std::runtime_error( - "T5Pipeline: unable to allocate cross-attention value buffer"); - } + t5::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 { @@ -290,9 +280,7 @@ class T5Pipeline 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("T5Pipeline: unable to allocate encoder mask buffer"); + t5::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/t5/tests/cpp/test_t5_cross_kv_alloc.cpp b/families/t5/tests/cpp/test_t5_cross_kv_alloc.cpp index b1a6940c28..cc9ad5bf4e 100644 --- a/families/t5/tests/cpp/test_t5_cross_kv_alloc.cpp +++ b/families/t5/tests/cpp/test_t5_cross_kv_alloc.cpp @@ -3,10 +3,10 @@ * SPDX-License-Identifier: Apache-2.0 */ -// Exercises the ownership the T5 pipeline uses for its cross-attention K/V -// buffers, against CPU CUDA stubs so each allocation can be failed in turn -// without a GPU. T5Pipeline itself needs live TensorRT modules to build, so -// this drives the same allocation loop over the same owning type. +// Drives the T5 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. T5Pipeline itself needs live TensorRT modules to build. #include "families/t5/runtime/device_buffer.h" @@ -33,21 +33,6 @@ void check(bool condition, const char* what) { } } -// The allocation loop from T5Pipeline'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("T5Pipeline: unable to allocate cross-attention key buffer"); - if (values[layer].allocate(bytes) != cudaSuccess) - throw std::runtime_error("T5Pipeline: unable to allocate cross-attention value buffer"); - } -} - } // namespace extern "C" { @@ -88,7 +73,7 @@ int main() { try { std::vector keys; std::vector values; - allocate_cross_kv(keys, values, layers, bytes); + trtmc::t5::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 @@ -111,12 +96,41 @@ int main() { { std::vector keys; std::vector values; - allocate_cross_kv(keys, values, layers, bytes); + trtmc::t5::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::t5::DeviceBuffer mask; + g_fail_on_allocation = 1; + g_allocation_count = 0; + bool mask_threw = false; + try { + trtmc::t5::ensure_encoder_mask(mask, bytes); + } catch (const std::runtime_error& error) { + mask_threw = true; + check(std::string(error.what()) == "T5Pipeline: 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::t5::ensure_encoder_mask(mask, bytes); + check(mask.get() != nullptr && g_outstanding.size() == 1, + "the encoder mask should allocate once"); + trtmc::t5::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;