diff --git a/examples/models/muse-glimmer/BUCK b/examples/models/muse-glimmer/BUCK new file mode 100644 index 00000000000..8d2cadeb964 --- /dev/null +++ b/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"), + ], +) diff --git a/examples/models/muse-glimmer/runtime/engine/sampling_cuda.cu b/examples/models/muse-glimmer/runtime/engine/sampling_cuda.cu new file mode 100644 index 00000000000..5f8f6171f0d --- /dev/null +++ b/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 + +#include +#include + +#include + +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(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(row_count), kArgmaxThreads, 0, stream>>>( + values, row_size, indices); + return cudaGetLastError(); +} + +} // 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 new file mode 100644 index 00000000000..904381d58e3 --- /dev/null +++ b/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 + +#include + +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 diff --git a/examples/models/muse-glimmer/tests/sampling_cuda_test.cpp b/examples/models/muse-glimmer/tests/sampling_cuda_test.cpp new file mode 100644 index 00000000000..59622b2d8e3 --- /dev/null +++ b/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 +#include + +#include +#include + +#include +#include + +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 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 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