[executorch][muse-glimmer] Add CUDA residual correction sampling

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<int64_t>(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<double>(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<int64_t>(blockIdx.x) * blockDim.x +
+      threadIdx.x;
+  if (offset < total_size) {
+    const int64_t row = offset / row_size;
+    const float denominator = static_cast<float>(
+        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<int>::max() ||
+      row_size > std::numeric_limits<int>::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<int>(
+      (total_size + kSamplingThreads - 1) / kSamplingThreads);
+  compute_residual_kernel<<<item_blocks, kSamplingThreads, 0, stream>>>(
+      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<int>(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

examples/models/muse-glimmer/runtime/engine/sampling_cuda.h

@@ -1,3 +1,3 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
@@ -69,6 +69,14 @@
       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 @@
     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

examples/models/muse-glimmer/tests/sampling_cuda_test.cpp

@@ -1,3 +1,3 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
@@ -389,4 +389,93 @@
   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<float, kRowSize> target_a = {0.5f, 0.3f, 0.2f};
+  const std::array<float, kRowSize> draft_a = {0.2f, 0.4f, 0.1f};
+  const std::array<float, kRowSize> target_b = {0.1f, 0.6f, 0.3f};
+  std::vector<float> target(kRows * kRowSize);
+  std::vector<float> 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<uint64_t> 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<int64_t, kRowSize> counts_a{};
+  std::array<int64_t, kRowSize> 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<float, kRowSize> 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<double>(counts_a[token]) / kRowsPerCase,
+        expected_a[token],
+        0.025);
+    EXPECT_NEAR(target[kRowsPerCase * kRowSize + token], target_b[token], 1e-6f);
+    EXPECT_NEAR(
+        static_cast<double>(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