From d2c66ac95ff8ec74989c51314b3a74de298735af Mon Sep 17 00:00:00 2001 From: Mohak Gupta Date: Mon, 21 Sep 2026 14:33:10 +0530 Subject: [PATCH] fix(bark): check the multinomial device query before sizing the launch bark_compute_torch_multinomial_execution_policy discarded the return value of both cudaGetDevice and cudaGetDeviceProperties. When either call fails the property block stays zero initialized, so the resident block count is zero, the grid is empty, and counter_offset divides by a zero thread count. On x86 that is a SIGFPE inside the policy, before the sampler can report anything. Check both calls and throw the CUDA error the way the rest of the bark runtime reports a failed query, and reject a zero thread count before the offset arithmetic. Thread counts and offsets on a healthy device are unchanged. The computation is host code, so move it out of the CUDA translation unit and cover it with a CPU test that links CUDA stubs and injects each failure; the GPU lane is not what gates a pull request. The curand offset stride the kernel and the policy both read now lives in the feature header so the two cannot drift. Signed-off-by: Mohak Gupta --- families/bark/runtime/CMakeLists.txt | 19 +++ .../bark/runtime/sparse_multinomial_kernel.cu | 31 ---- .../bark/runtime/sparse_multinomial_kernel.h | 4 + .../runtime/sparse_multinomial_policy.cpp | 64 ++++++++ .../cpp/test_bark_multinomial_policy.cpp | 141 ++++++++++++++++++ 5 files changed, 228 insertions(+), 31 deletions(-) create mode 100644 families/bark/runtime/sparse_multinomial_policy.cpp create mode 100644 families/bark/tests/cpp/test_bark_multinomial_policy.cpp diff --git a/families/bark/runtime/CMakeLists.txt b/families/bark/runtime/CMakeLists.txt index c094a0f62a..810a873490 100644 --- a/families/bark/runtime/CMakeLists.txt +++ b/families/bark/runtime/CMakeLists.txt @@ -10,6 +10,7 @@ add_library(trtmc_model_bark SHARED plugin.cpp plugin_helpers.cpp sampler.cpp + sparse_multinomial_policy.cpp sparse_multinomial_kernel.cu ) @@ -83,4 +84,22 @@ if(TRTMC_BUILD_TESTS) -Wall -Wextra -Wpedantic ) add_test(NAME bark_sampler_alloc COMMAND test_bark_sampler_alloc) + + # Compiles the host multinomial policy against CPU CUDA stubs defined in the + # test, so a failed device query can be injected without a GPU or cudart. + add_executable(test_bark_multinomial_policy + ${PROJECT_SOURCE_DIR}/families/bark/tests/cpp/test_bark_multinomial_policy.cpp + sparse_multinomial_policy.cpp + ) + target_include_directories(test_bark_multinomial_policy PRIVATE + ${PROJECT_SOURCE_DIR} + ${PROJECT_SOURCE_DIR}/core/runtime/include + ) + target_include_directories(test_bark_multinomial_policy SYSTEM PRIVATE + ${TRTMC_CUDA_INCLUDE_DIR} + ) + target_compile_options(test_bark_multinomial_policy PRIVATE + -Wall -Wextra -Wpedantic + ) + add_test(NAME bark_multinomial_policy COMMAND test_bark_multinomial_policy) endif() diff --git a/families/bark/runtime/sparse_multinomial_kernel.cu b/families/bark/runtime/sparse_multinomial_kernel.cu index 5d8c8c3f43..11a91a568f 100644 --- a/families/bark/runtime/sparse_multinomial_kernel.cu +++ b/families/bark/runtime/sparse_multinomial_kernel.cu @@ -5,7 +5,6 @@ #include "families/bark/runtime/sparse_multinomial_kernel.h" -#include #include #include #include @@ -14,9 +13,7 @@ namespace trtmc { namespace { -constexpr int kDistributionBlockSize = 256; constexpr int kSamplerBlockSize = 128; -constexpr uint64_t kGeneratorOffsetsPerCurandCall = 4; __device__ float torch_exponential_from_uniform(float value) { const float log_value = value >= 1.0F - FLT_EPSILON / 2.0F ? -FLT_EPSILON / 2.0F : logf(value); @@ -85,34 +82,6 @@ __global__ void sparse_multinomial_exact_kernel(const int32_t* __restrict__ indi } // namespace -BarkTorchMultinomialExecutionPolicy bark_compute_torch_multinomial_execution_policy(int32_t numel) { - if (numel <= 0) { - return {}; - } - - int device = 0; - cudaGetDevice(&device); - cudaDeviceProp properties{}; - cudaGetDeviceProperties(&properties, device); - - const uint32_t blocks_per_sm = - static_cast(properties.maxThreadsPerMultiProcessor / kDistributionBlockSize); - const uint32_t grid = - std::min(static_cast(properties.multiProcessorCount) * blocks_per_sm, - static_cast((static_cast(numel) + kDistributionBlockSize - 1) / - kDistributionBlockSize)); - const uint64_t total_threads = static_cast(grid) * kDistributionBlockSize; - const uint64_t counter_offset = - ((static_cast(numel) - 1) / (total_threads * kGeneratorOffsetsPerCurandCall) + - 1) * - kGeneratorOffsetsPerCurandCall; - - BarkTorchMultinomialExecutionPolicy policy; - policy.total_threads = static_cast(total_threads); - policy.counter_offset = counter_offset; - return policy; -} - void bark_gpu_sparse_torch_multinomial_exact(const int32_t* d_indices, const float* d_probs, int32_t rows, int32_t vocab_size, int32_t keep, uint64_t seed, uint64_t base_offset, diff --git a/families/bark/runtime/sparse_multinomial_kernel.h b/families/bark/runtime/sparse_multinomial_kernel.h index 20146a4ef4..8c6ae75c42 100644 --- a/families/bark/runtime/sparse_multinomial_kernel.h +++ b/families/bark/runtime/sparse_multinomial_kernel.h @@ -10,6 +10,10 @@ namespace trtmc { +// The kernel strides the curand offset by this and the policy sizes its offset +// with it, so both sides read one value. +inline constexpr uint64_t kGeneratorOffsetsPerCurandCall = 4; + struct BarkTorchMultinomialExecutionPolicy { int32_t total_threads{0}; uint64_t counter_offset{0}; diff --git a/families/bark/runtime/sparse_multinomial_policy.cpp b/families/bark/runtime/sparse_multinomial_policy.cpp new file mode 100644 index 0000000000..86ac37d313 --- /dev/null +++ b/families/bark/runtime/sparse_multinomial_policy.cpp @@ -0,0 +1,64 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "families/bark/runtime/sparse_multinomial_kernel.h" + +#include +#include +#include +#include + +namespace trtmc { + +namespace { + +constexpr int kDistributionBlockSize = 256; + +} // namespace + +BarkTorchMultinomialExecutionPolicy bark_compute_torch_multinomial_execution_policy(int32_t numel) { + if (numel <= 0) { + return {}; + } + + int device = 0; + const cudaError_t device_status = cudaGetDevice(&device); + if (device_status != cudaSuccess) { + throw std::runtime_error("cudaGetDevice failed for the bark multinomial launch: " + + std::string(cudaGetErrorString(device_status))); + } + + cudaDeviceProp properties{}; + const cudaError_t properties_status = cudaGetDeviceProperties(&properties, device); + if (properties_status != cudaSuccess) { + throw std::runtime_error( + "cudaGetDeviceProperties failed for the bark multinomial launch: " + + std::string(cudaGetErrorString(properties_status))); + } + + const uint32_t blocks_per_sm = + static_cast(properties.maxThreadsPerMultiProcessor / kDistributionBlockSize); + const uint32_t grid = + std::min(static_cast(properties.multiProcessorCount) * blocks_per_sm, + static_cast((static_cast(numel) + kDistributionBlockSize - 1) / + kDistributionBlockSize)); + const uint64_t total_threads = static_cast(grid) * kDistributionBlockSize; + // A query can succeed and still report no usable occupancy. + if (total_threads == 0) { + throw std::runtime_error("bark multinomial launch policy computed no threads"); + } + + const uint64_t counter_offset = + ((static_cast(numel) - 1) / (total_threads * kGeneratorOffsetsPerCurandCall) + + 1) * + kGeneratorOffsetsPerCurandCall; + + BarkTorchMultinomialExecutionPolicy policy; + policy.total_threads = static_cast(total_threads); + policy.counter_offset = counter_offset; + return policy; +} + +} // namespace trtmc diff --git a/families/bark/tests/cpp/test_bark_multinomial_policy.cpp b/families/bark/tests/cpp/test_bark_multinomial_policy.cpp new file mode 100644 index 0000000000..3b954b781c --- /dev/null +++ b/families/bark/tests/cpp/test_bark_multinomial_policy.cpp @@ -0,0 +1,141 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Compiles the host multinomial policy against CPU CUDA stubs, so a failed +// device query can be injected without a GPU or cudart. Verifies that a failed +// query throws instead of sizing a launch from an empty property block, and +// that a healthy device keeps the thread count and generator offset it had. + +#include "families/bark/runtime/sparse_multinomial_kernel.h" + +#include +#include +#include +#include +#include + +namespace { + +cudaError_t g_device_status = cudaSuccess; +cudaError_t g_properties_status = cudaSuccess; +cudaDeviceProp g_properties{}; +int g_device_queries = 0; +int g_property_queries = 0; +int g_failures = 0; + +void check(bool condition, const char* what) { + if (!condition) { + std::fprintf(stderr, "FAIL: %s\n", what); + ++g_failures; + } +} + +void reset_stubs() { + g_device_status = cudaSuccess; + g_properties_status = cudaSuccess; + g_properties = cudaDeviceProp{}; + g_device_queries = 0; + g_property_queries = 0; +} + +// Returns true when the call threw, and checks that the message names the +// operation that actually failed. +bool throws_naming(const char* expected, const char* what) { + try { + trtmc::bark_compute_torch_multinomial_execution_policy(1024); + } catch (const std::runtime_error& error) { + check(std::string(error.what()).find(expected) != std::string::npos, what); + return true; + } + check(false, what); + return false; +} + +} // namespace + +extern "C" { + +cudaError_t cudaGetDevice(int* device) { + ++g_device_queries; + if (g_device_status != cudaSuccess) { + return g_device_status; + } + *device = 0; + return cudaSuccess; +} + +cudaError_t cudaGetDeviceProperties(cudaDeviceProp* properties, int device) { + (void)device; + ++g_property_queries; + if (g_properties_status != cudaSuccess) { + return g_properties_status; + } + *properties = g_properties; + return cudaSuccess; +} + +const char* cudaGetErrorString(cudaError_t error) { + return error == cudaErrorInsufficientDriver ? "insufficient driver" : "invalid device ordinal"; +} + +} // extern "C" + +int main() { + // A healthy device keeps the values the sampler divides by. The grid is + // capped by resident blocks (108 * 6 = 648) rather than by the request. + reset_stubs(); + g_properties.multiProcessorCount = 108; + g_properties.maxThreadsPerMultiProcessor = 1536; + const trtmc::BarkTorchMultinomialExecutionPolicy policy = + trtmc::bark_compute_torch_multinomial_execution_policy(1000000); + check(policy.total_threads == 165888, "a 108-SM device should keep its resident thread count"); + check(policy.counter_offset == 8, "a 108-SM device should keep its generator offset"); + + // A request smaller than one resident grid is capped by the request itself. + reset_stubs(); + g_properties.multiProcessorCount = 108; + g_properties.maxThreadsPerMultiProcessor = 1536; + const trtmc::BarkTorchMultinomialExecutionPolicy small = + trtmc::bark_compute_torch_multinomial_execution_policy(256); + check(small.total_threads == 256, "a one-block request should keep one block of threads"); + check(small.counter_offset == 4, "a one-block request should keep its generator offset"); + + // An empty request needs no device query at all. + reset_stubs(); + const trtmc::BarkTorchMultinomialExecutionPolicy empty = + trtmc::bark_compute_torch_multinomial_execution_policy(0); + check(empty.total_threads == 0, "an empty request should have no threads"); + check(empty.counter_offset == 0, "an empty request should have no offset"); + check(g_device_queries == 0, "an empty request should not query the device"); + + // A device lookup that fails must throw, and must not go on to query the + // properties of a device it never resolved. + reset_stubs(); + g_device_status = cudaErrorInsufficientDriver; + check(throws_naming("cudaGetDevice failed", "a failed device lookup should throw naming it"), + "a failed device lookup should throw"); + check(g_property_queries == 0, "a failed device lookup should not query the device properties"); + + // A property lookup that fails must throw naming its own error string. + reset_stubs(); + g_properties_status = cudaErrorInvalidDevice; + check(throws_naming("invalid device ordinal", + "a failed property lookup should report the CUDA error"), + "a failed property lookup should throw"); + + // The reported fault: a query that succeeds but reports no usable occupancy + // used to divide by a zero thread count. + reset_stubs(); + check(throws_naming("computed no threads", + "a zeroed property block should throw instead of dividing by zero"), + "a zeroed property block should throw"); + + if (g_failures != 0) { + std::fprintf(stderr, "%d check(s) failed\n", g_failures); + return 1; + } + std::printf("bark multinomial policy device-query checks passed\n"); + return 0; +}