diff --git a/examples/models/muse-glimmer/runtime/engine/sampling_cuda.cu b/examples/models/muse-glimmer/runtime/engine/sampling_cuda.cu index d96108de5a8..45792231cbf 100644 --- a/examples/models/muse-glimmer/runtime/engine/sampling_cuda.cu +++ b/examples/models/muse-glimmer/runtime/engine/sampling_cuda.cu @@ -295,6 +295,39 @@ __global__ void normalize_probability_rows_kernel( } } +__global__ void compute_residual_kernel( + const float* target, + const float* draft, + int64_t total_size, + float* residual, + double* residual_double) { + const int64_t offset = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + if (offset < total_size) { + const float value = fmaxf(0.0f, target[offset] - draft[offset]); + residual[offset] = value; + residual_double[offset] = static_cast(value); + } +} + +__global__ void normalize_residual_rows_kernel( + float* target, + const float* residual, + 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]); + if (denominator > 0.0f) { + target[offset] = residual[offset] / denominator; + } + } +} + } // namespace struct SamplingWorkspace::Impl { @@ -719,4 +752,68 @@ cudaError_t sample_excluding_token_in_place( stream); } +cudaError_t sample_from_residual_in_place( + float* target_probabilities, + const float* draft_probabilities, + int64_t row_count, + int64_t row_size, + uint64_t* sampled_tokens, + SamplingWorkspace& workspace, + cudaStream_t stream) { + if (target_probabilities == nullptr || draft_probabilities == 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 int64_t total_size = state.total_size; + const int item_blocks = static_cast( + (total_size + kSamplingThreads - 1) / kSamplingThreads); + compute_residual_kernel<<>>( + target_probabilities, + draft_probabilities, + total_size, + state.sort_keys_in, + 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_residual_rows_kernel<<< + item_blocks, kSamplingThreads, 0, stream>>>( + target_probabilities, + state.sort_keys_in, + state.cumulative, + total_size, + row_size); + error = cudaGetLastError(); + if (error != cudaSuccess) { + return error; + } + return categorical_sample( + target_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 ba93187add4..f32ffac9feb 100644 --- a/examples/models/muse-glimmer/runtime/engine/sampling_cuda.h +++ b/examples/models/muse-glimmer/runtime/engine/sampling_cuda.h @@ -69,6 +69,14 @@ class SamplingWorkspace { uint64_t*, SamplingWorkspace&, cudaStream_t); + friend cudaError_t sample_from_residual_in_place( + float*, + const float*, + int64_t, + int64_t, + uint64_t*, + SamplingWorkspace&, + cudaStream_t); }; // Computes one argmax per contiguous row of `values`. @@ -128,4 +136,16 @@ cudaError_t sample_excluding_token_in_place( SamplingWorkspace& workspace, cudaStream_t stream); +// CUDA counterpart of muse_glimmer::sample_from_residual_in_place for batched rows. +// Mutates target probabilities to normalized max(p-q, 0) when residual mass +// exists, otherwise preserves p, then samples one token per row. +cudaError_t sample_from_residual_in_place( + float* target_probabilities, + const float* draft_probabilities, + int64_t row_count, + int64_t row_size, + 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 01a60438a90..3af72629dc1 100644 --- a/examples/models/muse-glimmer/tests/sampling_cuda_test.cpp +++ b/examples/models/muse-glimmer/tests/sampling_cuda_test.cpp @@ -389,4 +389,93 @@ TEST(CudaSamplingTest, ExcludingTokenMatchesHostSemantics) { ASSERT_CUDA_SUCCESS(cudaFree(device_probabilities)); } +TEST(CudaSamplingTest, ResidualCorrectionMatchesHostSemantics) { + constexpr int64_t kRowsPerCase = 8000; + constexpr int64_t kRows = 2 * kRowsPerCase; + constexpr int64_t kRowSize = 3; + const std::array target_a = {0.5f, 0.3f, 0.2f}; + const std::array draft_a = {0.2f, 0.4f, 0.1f}; + const std::array target_b = {0.1f, 0.6f, 0.3f}; + std::vector target(kRows * kRowSize); + std::vector draft(kRows * kRowSize); + for (int64_t row = 0; row < kRows; ++row) { + const auto& row_target = row < kRowsPerCase ? target_a : target_b; + const auto& row_draft = row < kRowsPerCase ? draft_a : target_b; + std::copy( + row_target.begin(), + row_target.end(), + target.begin() + row * kRowSize); + std::copy( + row_draft.begin(), + row_draft.end(), + draft.begin() + row * kRowSize); + } + + float* device_target = nullptr; + float* device_draft = nullptr; + uint64_t* device_tokens = nullptr; + const size_t probability_bytes = target.size() * sizeof(float); + ASSERT_CUDA_SUCCESS(cudaMalloc(&device_target, probability_bytes)); + ASSERT_CUDA_SUCCESS(cudaMalloc(&device_draft, probability_bytes)); + ASSERT_CUDA_SUCCESS(cudaMalloc(&device_tokens, kRows * sizeof(uint64_t))); + ASSERT_CUDA_SUCCESS(cudaMemcpy( + device_target, + target.data(), + probability_bytes, + cudaMemcpyHostToDevice)); + ASSERT_CUDA_SUCCESS(cudaMemcpy( + device_draft, + draft.data(), + probability_bytes, + cudaMemcpyHostToDevice)); + + muse_glimmer::cuda::SamplingWorkspace workspace; + ASSERT_CUDA_SUCCESS(workspace.reserve(kRows, kRowSize, nullptr)); + ASSERT_CUDA_SUCCESS(workspace.set_seed(3456, nullptr)); + ASSERT_CUDA_SUCCESS(muse_glimmer::cuda::sample_from_residual_in_place( + device_target, + device_draft, + kRows, + kRowSize, + 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( + target.data(), + device_target, + probability_bytes, + cudaMemcpyDeviceToHost)); + + std::array counts_a{}; + std::array counts_b{}; + for (int64_t row = 0; row < kRows; ++row) { + ASSERT_LT(tokens[row], kRowSize); + auto& counts = row < kRowsPerCase ? counts_a : counts_b; + ++counts[tokens[row]]; + } + const std::array expected_a = {0.75f, 0.0f, 0.25f}; + for (int64_t token = 0; token < kRowSize; ++token) { + EXPECT_NEAR(target[token], expected_a[token], 1e-6f); + EXPECT_NEAR( + static_cast(counts_a[token]) / kRowsPerCase, + expected_a[token], + 0.025); + EXPECT_NEAR(target[kRowsPerCase * kRowSize + token], target_b[token], 1e-6f); + EXPECT_NEAR( + static_cast(counts_b[token]) / kRowsPerCase, + target_b[token], + 0.025); + } + + ASSERT_CUDA_SUCCESS(cudaFree(device_tokens)); + ASSERT_CUDA_SUCCESS(cudaFree(device_draft)); + ASSERT_CUDA_SUCCESS(cudaFree(device_target)); +} + } // namespace