[executorch][muse-glimmer] Add CUDA argmax sampling primitive

examples/models/muse-glimmer/BUCK

@@ -0,0 +1,43 @@
+load("@fbcode//tools/build/buck:nvcc_flags.bzl", "get_nvcc_arch_args")
+load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target")
+load("@fbcode_macros//build_defs/lib:re_test_utils.bzl", "re_test_utils")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+oncall("executorch")
+
+# GPU-side counterparts of the host sampling primitives in
+# runtime/engine/sampling.h, used by the CUDA DFlash decode loop.
+fbcode_target(_kind = runtime.cxx_library,
+    name = "sampling_cuda",
+    srcs = ["runtime/engine/sampling_cuda.cu"],
+    headers = [
+        "runtime/engine/sampling.h",
+        "runtime/engine/sampling_cuda.h",
+    ],
+    nvcc_flags = get_nvcc_arch_args() + [
+        "-_NVCC_HOST_COMPILER_FLAG_",
+        "gcc",
+    ],
+    external_deps = [
+        ("cuda", None, "cuda-lazy"),
+    ],
+    visibility = ["PUBLIC"],
+)
+
+fbcode_target(_kind = runtime.cxx_test,
+    name = "sampling_cuda_test",
+    srcs = ["tests/sampling_cuda_test.cpp"],
+    keep_gpu_sections = True,
+    remote_execution = re_test_utils.remote_execution(
+        platform = "gpu-remote-execution",
+        subplatform = "A100-exclusive",
+    ),
+    deps = [
+        ":sampling_cuda",
+        "//executorch/extension/tensor:tensor",
+    ],
+    external_deps = [
+        "gtest",
+        ("cuda", None, "cuda-lazy"),
+    ],
+)

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

@@ -0,0 +1,94 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/examples/models/muse-glimmer/runtime/engine/sampling_cuda.h>
+
+#include <cuda_runtime.h>
+#include <math_constants.h>
+
+#include <cstdint>
+
+namespace muse_glimmer::cuda {
+namespace {
+
+constexpr int kArgmaxThreads = 256;
+
+struct ArgmaxCandidate {
+  float value;
+  uint64_t index;
+};
+
+__device__ ArgmaxCandidate better_candidate(
+    ArgmaxCandidate lhs,
+    ArgmaxCandidate rhs) {
+  if (rhs.value > lhs.value ||
+      (rhs.value == lhs.value && rhs.index < lhs.index)) {
+    return rhs;
+  }
+  return lhs;
+}
+
+__global__ void argmax_index_kernel(
+    const float* __restrict__ values,
+    int64_t row_size,
+    uint64_t* __restrict__ indices) {
+  const int64_t row = blockIdx.x;
+  const float* row_values = values + row * row_size;
+
+  ArgmaxCandidate candidate{-CUDART_INF_F, uint64_t{0}};
+  for (int64_t token = threadIdx.x; token < row_size;
+       token += blockDim.x) {
+    candidate = better_candidate(
+        candidate,
+        ArgmaxCandidate{row_values[token], static_cast<uint64_t>(token)});
+  }
+
+  __shared__ float shared_values[kArgmaxThreads];
+  __shared__ uint64_t shared_indices[kArgmaxThreads];
+  shared_values[threadIdx.x] = candidate.value;
+  shared_indices[threadIdx.x] = candidate.index;
+  __syncthreads();
+
+  for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+    if (threadIdx.x < stride) {
+      const ArgmaxCandidate reduced = better_candidate(
+          ArgmaxCandidate{
+              shared_values[threadIdx.x], shared_indices[threadIdx.x]},
+          ArgmaxCandidate{
+              shared_values[threadIdx.x + stride],
+              shared_indices[threadIdx.x + stride]});
+      shared_values[threadIdx.x] = reduced.value;
+      shared_indices[threadIdx.x] = reduced.index;
+    }
+    __syncthreads();
+  }
+
+  if (threadIdx.x == 0) {
+    indices[row] = shared_indices[0];
+  }
+}
+
+} // namespace
+
+cudaError_t argmax_index(
+    const float* values,
+    int64_t row_count,
+    int64_t row_size,
+    uint64_t* indices,
+    cudaStream_t stream) {
+  if (values == nullptr || indices == nullptr || row_count <= 0 ||
+      row_size <= 0) {
+    return cudaErrorInvalidValue;
+  }
+  argmax_index_kernel<<<
+      static_cast<unsigned int>(row_count), kArgmaxThreads, 0, stream>>>(
+      values, row_size, indices);
+  return cudaGetLastError();
+}
+
+} // namespace muse_glimmer::cuda

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

@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// CUDA counterparts of the host sampling primitives in sampling.h.
+
+#pragma once
+
+#include <cuda_runtime_api.h>
+
+#include <cstdint>
+
+namespace muse_glimmer::cuda {
+
+// Computes one argmax per contiguous row of `values`.
+//
+// `values` and `indices` must point to CUDA memory. Equal maxima select the
+// lowest token index, matching muse_glimmer::argmax_index. The launch is asynchronous
+// with respect to `stream`.
+cudaError_t argmax_index(
+    const float* values,
+    int64_t row_count,
+    int64_t row_size,
+    uint64_t* indices,
+    cudaStream_t stream);
+
+} // namespace muse_glimmer::cuda

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

@@ -0,0 +1,75 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/examples/models/muse-glimmer/runtime/engine/sampling.h>
+#include <executorch/examples/models/muse-glimmer/runtime/engine/sampling_cuda.h>
+
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+
+#include <cstdint>
+#include <vector>
+
+namespace {
+
+#define ASSERT_CUDA_SUCCESS(expression)               \
+  do {                                                \
+    const cudaError_t error = (expression);           \
+    ASSERT_EQ(error, cudaSuccess)                     \
+        << cudaGetErrorString(error);                 \
+  } while (false)
+
+TEST(CudaSamplingTest, ArgmaxMatchesHostForBatchedRowsAndTies) {
+  constexpr int64_t kRows = 3;
+  constexpr int64_t kRowSize = 513;
+  std::vector<float> host_values(kRows * kRowSize, -1000.0f);
+
+  host_values[0 * kRowSize + 1] = 8.0f;
+  host_values[0 * kRowSize + 257] = 8.0f;
+  host_values[1 * kRowSize + 512] = -2.0f;
+  host_values[1 * kRowSize + 17] = -3.0f;
+  host_values[2 * kRowSize + 256] = 4.0f;
+  host_values[2 * kRowSize + 511] = 4.0f;
+
+  float* device_values = nullptr;
+  uint64_t* device_indices = nullptr;
+  ASSERT_CUDA_SUCCESS(
+      cudaMalloc(&device_values, host_values.size() * sizeof(float)));
+  ASSERT_CUDA_SUCCESS(cudaMalloc(&device_indices, kRows * sizeof(uint64_t)));
+  ASSERT_CUDA_SUCCESS(cudaMemcpy(
+      device_values,
+      host_values.data(),
+      host_values.size() * sizeof(float),
+      cudaMemcpyHostToDevice));
+
+  ASSERT_CUDA_SUCCESS(muse_glimmer::cuda::argmax_index(
+      device_values, kRows, kRowSize, device_indices, nullptr));
+  std::vector<uint64_t> actual(kRows);
+  ASSERT_CUDA_SUCCESS(cudaMemcpy(
+      actual.data(),
+      device_indices,
+      actual.size() * sizeof(uint64_t),
+      cudaMemcpyDeviceToHost));
+
+  for (int64_t row = 0; row < kRows; ++row) {
+    const uint64_t expected = muse_glimmer::argmax_index(
+        host_values.data() + row * kRowSize, kRowSize);
+    EXPECT_EQ(actual[row], expected) << "row " << row;
+  }
+
+  ASSERT_CUDA_SUCCESS(cudaFree(device_indices));
+  ASSERT_CUDA_SUCCESS(cudaFree(device_values));
+}
+
+TEST(CudaSamplingTest, ArgmaxRejectsInvalidArguments) {
+  EXPECT_EQ(
+      muse_glimmer::cuda::argmax_index(nullptr, 1, 1, nullptr, nullptr),
+      cudaErrorInvalidValue);
+}
+
+} // namespace