From 7457b52f60e113b4ce52de80c310f9e4a26f2d62 Mon Sep 17 00:00:00 2001 From: Moviw Date: Sun, 13 Sep 2026 03:21:40 +0900 Subject: [PATCH] fix(flux): check the return value of ensure_buf's cudaMalloc flux_gpu_matmul's scratch buffers were grown via a bare cudaMalloc call whose cudaError_t was discarded, matching the bug fixed in bark by #1201 and reported for six sibling families by #1221. On an allocation failure the pointer stayed null while `bytes` was already updated to the requested size, so the buffer looked correctly sized to the next call and the first real use (cudaMemcpyAsync/cublasSgemm into a null pointer) failed far from the actual cause. Introduces a small DeviceBuffer RAII wrapper (matching the pattern already merged for bart/whisper and open for m2m_100/marian/t5 in #1221's sibling PRs) so ensure_buf frees the stale buffer, throws on allocation failure, and leaves `bytes` unchanged so the next call retries the grow instead of treating a missing buffer as already sized. Fixes #1221 (flux only; the other five families are separate PRs per the repo's one-family-per-PR convention). Signed-off-by: Moviw --- families/flux/runtime/CMakeLists.txt | 18 +++ families/flux/runtime/device_buffer.h | 71 +++++++++++ families/flux/runtime/gpu_matmul.cpp | 49 +++----- .../cpp/test_flux_device_buffer_alloc.cpp | 115 ++++++++++++++++++ 4 files changed, 221 insertions(+), 32 deletions(-) create mode 100644 families/flux/runtime/device_buffer.h create mode 100644 families/flux/tests/cpp/test_flux_device_buffer_alloc.cpp diff --git a/families/flux/runtime/CMakeLists.txt b/families/flux/runtime/CMakeLists.txt index 44e856c375..30091e9e4b 100644 --- a/families/flux/runtime/CMakeLists.txt +++ b/families/flux/runtime/CMakeLists.txt @@ -58,4 +58,22 @@ if(TRTMC_BUILD_TESTS) target_compile_options(${test_name} PRIVATE -Wall -Wextra -Wpedantic) add_test(NAME ${test_name} COMMAND ${test_name}) endforeach() + + # CPU CUDA stubs live in the test, so allocation failures can be injected + # without a GPU. Not linked against trtmc_model_flux: that would pull in + # the real cudart/cublas symbols the stubs below need to shadow. + add_executable(test_flux_device_buffer_alloc + ${PROJECT_SOURCE_DIR}/families/flux/tests/cpp/test_flux_device_buffer_alloc.cpp + ) + target_include_directories(test_flux_device_buffer_alloc PRIVATE + ${PROJECT_SOURCE_DIR} + ${PROJECT_SOURCE_DIR}/core/runtime/include + ) + target_include_directories(test_flux_device_buffer_alloc SYSTEM PRIVATE + ${TRTMC_CUDA_INCLUDE_DIR} + ) + target_compile_options(test_flux_device_buffer_alloc PRIVATE + -Wall -Wextra -Wpedantic + ) + add_test(NAME flux_device_buffer_alloc COMMAND test_flux_device_buffer_alloc) endif() diff --git a/families/flux/runtime/device_buffer.h b/families/flux/runtime/device_buffer.h new file mode 100644 index 0000000000..9967f66bc8 --- /dev/null +++ b/families/flux/runtime/device_buffer.h @@ -0,0 +1,71 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace trtmc { +namespace flux { + +// Owns one device allocation. Growing it frees the previous allocation before +// attempting the new one and leaves the pointer null on failure, so a failed +// grow can never leave a stale or dangling pointer for the caller to reuse. +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}; +}; + +// A device buffer reused and grown across calls (unlike the per-request +// buffers the other families own): `bytes` tracks the capacity actually +// allocated, not the size of the most recent request. +struct GrowableBuffer { + DeviceBuffer buf; + std::size_t bytes = 0; +}; + +// Grows `gb` to at least `need` bytes. A no-op when it is already large +// enough. On allocation failure the previous (too-small) buffer is already +// gone -- `bytes` is left unchanged so the next call retries the grow rather +// than treating the missing buffer as already sized. +inline void ensure_buf(GrowableBuffer& gb, std::size_t need) { + if (gb.bytes >= need) + return; + if (gb.buf.allocate(need) != cudaSuccess) + throw std::runtime_error("flux_gpu_matmul: unable to allocate device buffer"); + gb.bytes = need; +} + +} // namespace flux +} // namespace trtmc diff --git a/families/flux/runtime/gpu_matmul.cpp b/families/flux/runtime/gpu_matmul.cpp index f5ab449b67..4be23fd4df 100644 --- a/families/flux/runtime/gpu_matmul.cpp +++ b/families/flux/runtime/gpu_matmul.cpp @@ -7,6 +7,8 @@ #include "families/flux/runtime/gpu_matmul.h" +#include "families/flux/runtime/device_buffer.h" + #include #include #include @@ -18,20 +20,7 @@ namespace { cublasHandle_t g_cublas = nullptr; cudaStream_t g_stream = nullptr; -struct DevBuf { - float* ptr = nullptr; - size_t bytes = 0; -}; -DevBuf g_dA, g_dB, g_dC; - -void ensure_buf(DevBuf& buf, size_t need) { - if (buf.bytes >= need) - return; - if (buf.ptr) - cudaFree(buf.ptr); - cudaMalloc(reinterpret_cast(&buf.ptr), need); - buf.bytes = need; -} +flux::GrowableBuffer g_dA, g_dB, g_dC; } // namespace @@ -44,16 +33,9 @@ void flux_gpu_matmul_init() { } void flux_gpu_matmul_shutdown() { - auto free_buf = [](DevBuf& b) { - if (b.ptr) { - cudaFree(b.ptr); - b.ptr = nullptr; - b.bytes = 0; - } - }; - free_buf(g_dA); - free_buf(g_dB); - free_buf(g_dC); + g_dA = flux::GrowableBuffer{}; + g_dB = flux::GrowableBuffer{}; + g_dC = flux::GrowableBuffer{}; if (g_stream) { cudaStreamDestroy(g_stream); g_stream = nullptr; @@ -70,18 +52,21 @@ void flux_gpu_matmul_bias(const float* A, const float* B, const float* bias, flo const size_t sB = size_t(K) * N * sizeof(float); const size_t sC = size_t(M) * N * sizeof(float); - ensure_buf(g_dA, sA); - ensure_buf(g_dB, sB); - ensure_buf(g_dC, sC); + flux::ensure_buf(g_dA, sA); + flux::ensure_buf(g_dB, sB); + flux::ensure_buf(g_dC, sC); + + auto* dA = static_cast(g_dA.buf.get()); + auto* dB = static_cast(g_dB.buf.get()); + auto* dC = static_cast(g_dC.buf.get()); - cudaMemcpyAsync(g_dA.ptr, A, sA, cudaMemcpyHostToDevice, g_stream); - cudaMemcpyAsync(g_dB.ptr, B, sB, cudaMemcpyHostToDevice, g_stream); + cudaMemcpyAsync(dA, A, sA, cudaMemcpyHostToDevice, g_stream); + cudaMemcpyAsync(dB, B, sB, cudaMemcpyHostToDevice, g_stream); const float alpha = 1.0f, beta = 0.0f; - cublasSgemm(g_cublas, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, g_dB.ptr, N, g_dA.ptr, K, - &beta, g_dC.ptr, N); + cublasSgemm(g_cublas, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, dB, N, dA, K, &beta, dC, N); - cudaMemcpyAsync(out, g_dC.ptr, sC, cudaMemcpyDeviceToHost, g_stream); + cudaMemcpyAsync(out, dC, sC, cudaMemcpyDeviceToHost, g_stream); cudaStreamSynchronize(g_stream); if (bias) { diff --git a/families/flux/tests/cpp/test_flux_device_buffer_alloc.cpp b/families/flux/tests/cpp/test_flux_device_buffer_alloc.cpp new file mode 100644 index 0000000000..c5eee18884 --- /dev/null +++ b/families/flux/tests/cpp/test_flux_device_buffer_alloc.cpp @@ -0,0 +1,115 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Exercises flux::ensure_buf's grow-and-retry contract against CPU CUDA +// stubs, so an allocation failure can be injected without a GPU. The matmul +// scratch buffers are process-lifetime globals reused across calls, unlike +// the per-request buffers the other families own, so what matters here is +// that a failed grow releases the stale buffer and leaves the tracked size +// unchanged, rather than the constructor-unwind behavior those test. + +#include "families/flux/runtime/device_buffer.h" + +#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() { + using trtmc::flux::ensure_buf; + using trtmc::flux::GrowableBuffer; + + // A no-op grow (need <= bytes) must not touch the allocator at all. + { + GrowableBuffer gb; + g_allocation_count = 0; + ensure_buf(gb, 256); + check(g_allocation_count == 1, "the first grow from empty should allocate once"); + const int32_t first_count = g_allocation_count; + ensure_buf(gb, 128); + check(g_allocation_count == first_count, + "shrinking the request below the current capacity should not reallocate"); + } + check(g_outstanding.empty(), "leaving scope should release the buffer"); + + // A failed grow must release the stale buffer, leave the pointer null, + // and leave `bytes` unchanged so the next call retries rather than + // treating the missing buffer as already sized. + { + GrowableBuffer gb; + ensure_buf(gb, 256); + check(gb.bytes == 256, "a successful grow should record the new size"); + + g_fail_on_allocation = g_allocation_count + 1; + bool threw = false; + try { + ensure_buf(gb, 1024); + } catch (const std::runtime_error&) { + threw = true; + } + check(threw, "a failed grow should throw"); + check(gb.buf.get() == nullptr, "a failed grow must leave the pointer null, not stale"); + check(gb.bytes == 256, "a failed grow must not update the tracked size"); + check(g_outstanding.empty(), + "a failed grow must release the old buffer instead of leaking it"); + + // Retrying with the allocator working again must succeed instead of + // treating `bytes` as already covering the request. + g_fail_on_allocation = 0; + ensure_buf(gb, 1024); + check(gb.bytes == 1024, "retrying after a transient failure should grow normally"); + } + check(g_outstanding.empty(), "leaving scope should release the buffer"); + + if (g_failures != 0) { + std::fprintf(stderr, "%d check(s) failed\n", g_failures); + return 1; + } + std::printf("flux device buffer allocation-failure checks passed\n"); + return 0; +}