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..6e15718cae --- /dev/null +++ b/families/t5/runtime/device_buffer.h @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#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}; +}; + +// 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 8e3e617b4d..f6e7ee4361 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,23 +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); - 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); - } - } - - ~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_); + 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 { @@ -281,10 +266,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 +280,20 @@ 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_); + t5::ensure_encoder_mask(enc_mask_device_, mask_bytes); + 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 +366,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..cc9ad5bf4e --- /dev/null +++ b/families/t5/tests/cpp/test_t5_cross_kv_alloc.cpp @@ -0,0 +1,140 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// 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" + +#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; + } +} + +} // 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; + 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 + // 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; + 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; + } + std::printf("t5 cross-attention allocation-failure checks passed\n"); + return 0; +}