fix(flux): check the return value of ensure_buf's cudaMalloc

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()

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 <cstddef>
+#include <cuda_runtime_api.h>
+#include <stdexcept>
+
+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

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 <cstdlib>
 #include <cublas_v2.h>
 #include <cuda_runtime_api.h>
@@ -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<void**>(&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<float*>(g_dA.buf.get());
+    auto* dB = static_cast<float*>(g_dB.buf.get());
+    auto* dC = static_cast<float*>(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) {

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 <cstdint>
+#include <cstdio>
+#include <cuda_runtime.h>
+#include <set>
+#include <stdexcept>
+
+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() {
+    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;
+}