diff --git a/examples/models/muse-glimmer/runtime/engine/sampling_cuda.cu b/examples/models/muse-glimmer/runtime/engine/sampling_cuda.cu index e776200ddaf..d96108de5a8 100644 --- a/examples/models/muse-glimmer/runtime/engine/sampling_cuda.cu +++ b/examples/models/muse-glimmer/runtime/engine/sampling_cuda.cu @@ -268,6 +268,33 @@ __global__ void accept_with_probability_kernel( uniform_from_uint32(curand(&local_rng)) < probabilities[index]; } +__global__ void exclude_tokens_kernel( + float* probabilities, + int64_t row_count, + int64_t row_size, + const uint64_t* excluded_tokens) { + const int64_t row = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + if (row < row_count && excluded_tokens[row] < row_size) { + probabilities[row * row_size + excluded_tokens[row]] = 0.0f; + } +} + +__global__ void normalize_probability_rows_kernel( + float* probabilities, + const double* cumulative, + int64_t total_size, + int64_t row_size) { + const int64_t offset = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + if (offset < total_size) { + const int64_t row = offset / row_size; + const float denominator = static_cast( + cumulative[row * row_size + row_size - 1]); + probabilities[offset] /= denominator; + } +} + } // namespace struct SamplingWorkspace::Impl { @@ -626,4 +653,70 @@ cudaError_t accept_with_probability( return cudaGetLastError(); } +cudaError_t sample_excluding_token_in_place( + float* probabilities, + int64_t row_count, + int64_t row_size, + const uint64_t* excluded_tokens, + uint64_t* sampled_tokens, + SamplingWorkspace& workspace, + cudaStream_t stream) { + if (probabilities == nullptr || excluded_tokens == nullptr || + sampled_tokens == nullptr || row_count <= 0 || row_size <= 0 || + row_count > std::numeric_limits::max() || + row_size > std::numeric_limits::max() / row_count) { + return cudaErrorInvalidValue; + } + cudaError_t error = workspace.reserve(row_count, row_size, stream); + if (error != cudaSuccess) { + return error; + } + auto& state = *workspace.impl_; + const int row_blocks = static_cast( + (row_count + kSamplingThreads - 1) / kSamplingThreads); + exclude_tokens_kernel<<>>( + probabilities, row_count, row_size, excluded_tokens); + error = cudaGetLastError(); + if (error != cudaSuccess) { + return error; + } + + const int64_t total_size = state.total_size; + const int item_blocks = static_cast( + (total_size + kSamplingThreads - 1) / kSamplingThreads); + probabilities_to_double_kernel<<< + item_blocks, kSamplingThreads, 0, stream>>>( + probabilities, total_size, state.weights); + error = cudaGetLastError(); + if (error != cudaSuccess) { + return error; + } + for (int64_t row = 0; row < row_count; ++row) { + error = cub::DeviceScan::InclusiveSum( + state.temporary_storage, + state.temporary_storage_bytes, + state.weights + row * row_size, + state.cumulative + row * row_size, + static_cast(row_size), + stream); + if (error != cudaSuccess) { + return error; + } + } + normalize_probability_rows_kernel<<< + item_blocks, kSamplingThreads, 0, stream>>>( + probabilities, state.cumulative, total_size, row_size); + error = cudaGetLastError(); + if (error != cudaSuccess) { + return error; + } + return categorical_sample( + probabilities, + row_count, + row_size, + sampled_tokens, + workspace, + stream); +} + } // namespace muse_glimmer::cuda diff --git a/examples/models/muse-glimmer/runtime/engine/sampling_cuda.h b/examples/models/muse-glimmer/runtime/engine/sampling_cuda.h index 6f61327bb88..ba93187add4 100644 --- a/examples/models/muse-glimmer/runtime/engine/sampling_cuda.h +++ b/examples/models/muse-glimmer/runtime/engine/sampling_cuda.h @@ -61,6 +61,14 @@ class SamplingWorkspace { uint8_t*, SamplingWorkspace&, cudaStream_t); + friend cudaError_t sample_excluding_token_in_place( + float*, + int64_t, + int64_t, + const uint64_t*, + uint64_t*, + SamplingWorkspace&, + cudaStream_t); }; // Computes one argmax per contiguous row of `values`. @@ -109,4 +117,15 @@ cudaError_t accept_with_probability( SamplingWorkspace& workspace, cudaStream_t stream); +// CUDA counterpart of muse_glimmer::sample_excluding_token_in_place for batched rows. +// Mutates each probability row by excluding and renormalizing before sampling. +cudaError_t sample_excluding_token_in_place( + float* probabilities, + int64_t row_count, + int64_t row_size, + const uint64_t* excluded_tokens, + uint64_t* sampled_tokens, + SamplingWorkspace& workspace, + cudaStream_t stream); + } // namespace muse_glimmer::cuda diff --git a/examples/models/muse-glimmer/tests/sampling_cuda_test.cpp b/examples/models/muse-glimmer/tests/sampling_cuda_test.cpp index bcdd2d51887..01a60438a90 100644 --- a/examples/models/muse-glimmer/tests/sampling_cuda_test.cpp +++ b/examples/models/muse-glimmer/tests/sampling_cuda_test.cpp @@ -318,4 +318,75 @@ TEST(CudaSamplingTest, AcceptanceMatchesHostSemantics) { ASSERT_CUDA_SUCCESS(cudaFree(device_probabilities)); } +TEST(CudaSamplingTest, ExcludingTokenMatchesHostSemantics) { + constexpr int64_t kRows = 12000; + constexpr int64_t kRowSize = 3; + const std::array distribution = {0.2f, 0.5f, 0.3f}; + std::vector probabilities(kRows * kRowSize); + for (int64_t row = 0; row < kRows; ++row) { + std::copy( + distribution.begin(), + distribution.end(), + probabilities.begin() + row * kRowSize); + } + std::vector excluded(kRows, 1); + + float* device_probabilities = nullptr; + uint64_t* device_excluded = nullptr; + uint64_t* device_tokens = nullptr; + ASSERT_CUDA_SUCCESS(cudaMalloc( + &device_probabilities, probabilities.size() * sizeof(float))); + ASSERT_CUDA_SUCCESS( + cudaMalloc(&device_excluded, excluded.size() * sizeof(uint64_t))); + ASSERT_CUDA_SUCCESS( + cudaMalloc(&device_tokens, kRows * sizeof(uint64_t))); + ASSERT_CUDA_SUCCESS(cudaMemcpy( + device_probabilities, + probabilities.data(), + probabilities.size() * sizeof(float), + cudaMemcpyHostToDevice)); + ASSERT_CUDA_SUCCESS(cudaMemcpy( + device_excluded, + excluded.data(), + excluded.size() * sizeof(uint64_t), + cudaMemcpyHostToDevice)); + + muse_glimmer::cuda::SamplingWorkspace workspace; + ASSERT_CUDA_SUCCESS(workspace.reserve(kRows, kRowSize, nullptr)); + ASSERT_CUDA_SUCCESS(workspace.set_seed(9012, nullptr)); + ASSERT_CUDA_SUCCESS(muse_glimmer::cuda::sample_excluding_token_in_place( + device_probabilities, + kRows, + kRowSize, + device_excluded, + device_tokens, + workspace, + nullptr)); + std::vector tokens(kRows); + ASSERT_CUDA_SUCCESS(cudaMemcpy( + tokens.data(), + device_tokens, + tokens.size() * sizeof(uint64_t), + cudaMemcpyDeviceToHost)); + ASSERT_CUDA_SUCCESS(cudaMemcpy( + probabilities.data(), + device_probabilities, + probabilities.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + + int64_t token_zero_count = 0; + for (int64_t row = 0; row < kRows; ++row) { + EXPECT_NEAR(probabilities[row * kRowSize], 0.4f, 1e-6f); + EXPECT_EQ(probabilities[row * kRowSize + 1], 0.0f); + EXPECT_NEAR(probabilities[row * kRowSize + 2], 0.6f, 1e-6f); + ASSERT_NE(tokens[row], 1); + token_zero_count += tokens[row] == 0; + } + EXPECT_NEAR(static_cast(token_zero_count) / kRows, 0.4, 0.02); + + ASSERT_CUDA_SUCCESS(cudaFree(device_tokens)); + ASSERT_CUDA_SUCCESS(cudaFree(device_excluded)); + ASSERT_CUDA_SUCCESS(cudaFree(device_probabilities)); +} + } // namespace