diff --git a/families/bark/runtime/CMakeLists.txt b/families/bark/runtime/CMakeLists.txt index c094a0f62..810a87349 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 5d8c8c3f4..11a91a568 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 20146a4ef..8c6ae75c4 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 000000000..86ac37d31 --- /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 000000000..3b954b781 --- /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; +}