From 0100227151c6a93546f89fe6c624e49465935279 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 11 Aug 2026 21:33:53 +0800 Subject: [PATCH] feat(nvidia): add top-k top-p sampling provider --- .../kernel.cu | 361 ++++++++++++++++ .../kernel.cuh | 178 ++++++++ .../top_k_top_p_sampling_from_logits/kernel.h | 81 ++++ tests/test_cpp_api.py | 147 +++++++ .../test_top_k_top_p_sampling_from_logits.py | 392 +++++++++++++++++- 5 files changed, 1156 insertions(+), 3 deletions(-) create mode 100644 src/native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.cuh create mode 100644 src/native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.h diff --git a/src/native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.cu b/src/native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.cu new file mode 100644 index 000000000..abbc07be6 --- /dev/null +++ b/src/native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.cu @@ -0,0 +1,361 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "data_type.h" +#include "dispatcher.h" +#include "native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.cuh" +#include "native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.h" + +namespace infini::ops { +namespace { + +constexpr uint64_t kCounterIncrement = 0x9e3779b97f4a7c15ULL; + +uint64_t MixCounter(uint64_t value) { + value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL; + value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL; + return value ^ (value >> 31); +} + +double CounterBasedUniform(uint64_t seed, uint64_t counter) { + const auto bits = MixCounter(seed + (counter + 1) * kCounterIncrement); + constexpr double kInverseTwoToThe53 = + 1.0 / static_cast(uint64_t{1} << 53); + return (static_cast(bits >> 11) + 0.5) * kInverseTwoToThe53; +} + +class DeviceGuard { + public: + explicit DeviceGuard(int device_index) { + auto status = cudaGetDevice(&previous_device_); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` failed to query the current CUDA " + "device"); + + if (previous_device_ != device_index) { + status = cudaSetDevice(device_index); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` failed to select the input CUDA " + "device"); + restore_ = true; + } + } + + ~DeviceGuard() { + if (restore_) { + const auto status = cudaSetDevice(previous_device_); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` failed to restore the CUDA device"); + } + } + + private: + int previous_device_{0}; + + bool restore_{false}; +}; + +} // namespace + +Operator::Operator( + const Tensor logits, const Tensor top_k, const Tensor top_p, + const std::optional indices, const std::string filter_apply_order, + const bool deterministic, const bool check_nan, + const std::optional seed, const std::optional offset, + Tensor out) + : TopKTopPSamplingFromLogits(logits, top_k, top_p, indices, + filter_apply_order, deterministic, check_nan, + seed, offset, out), + device_index_{logits.device().index()} { + ValidateSupportedOptions(filter_apply_order, deterministic, check_nan); + ValidateHostTensor(top_k); + ValidateHostTensor(top_p); + ValidateIndices(indices); + assert(logits.IsContiguous() && + "The NVIDIA `TopKTopPSamplingFromLogits` provider requires " + "contiguous logits"); + assert(out.IsContiguous() && + "The NVIDIA `TopKTopPSamplingFromLogits` provider requires " + "contiguous output"); + assert(out.device() == logits.device() && + "The NVIDIA `TopKTopPSamplingFromLogits` provider requires logits " + "and output on the same device"); + assert(vocab_size_ > 0 && + "The NVIDIA `TopKTopPSamplingFromLogits` provider requires a " + "nonempty vocabulary"); + assert(vocab_size_ <= std::numeric_limits::max() && + "The NVIDIA `TopKTopPSamplingFromLogits` provider requires the " + "vocabulary size to fit in `int`"); + + DeviceGuard device_guard{device_index_}; + workspace_size_ = DispatchWorkspaceSize(dtype_, vocab_size_); +} + +Operator::~Operator() { + if (std::all_of(default_workspace_slots_.begin(), + default_workspace_slots_.end(), + [](const auto& slot) { return slot.workspace == nullptr; })) { + return; + } + + DeviceGuard device_guard{device_index_}; + for (auto& slot : default_workspace_slots_) { + if (slot.workspace == nullptr) continue; + + if (slot.completion_recorded) { + const auto status = cudaEventSynchronize(slot.completion); + assert( + status == cudaSuccess && + "`TopKTopPSamplingFromLogits` failed to synchronize CUDA workspace"); + } + + auto status = cudaFree(slot.workspace); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` failed to free CUDA workspace"); + + status = cudaEventDestroy(slot.completion); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` failed to destroy CUDA workspace " + "event"); + } +} + +std::size_t Operator::workspace_size_in_bytes() const { + return workspace_size_; +} + +Operator::DefaultWorkspaceSlot* +Operator::AcquireDefaultWorkspaceSlot(cudaStream_t stream) const { + for (auto& slot : default_workspace_slots_) { + if (slot.workspace != nullptr && slot.stream == stream) return &slot; + } + + for (auto& slot : default_workspace_slots_) { + if (slot.workspace == nullptr) { + auto status = cudaMalloc(&slot.workspace, workspace_size_); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` failed to allocate CUDA workspace"); + + status = + cudaEventCreateWithFlags(&slot.completion, cudaEventDisableTiming); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` failed to create CUDA workspace " + "event"); + slot.stream = stream; + return &slot; + } + } + + auto& slot = default_workspace_slots_[next_default_workspace_slot_]; + next_default_workspace_slot_ = + (next_default_workspace_slot_ + 1) % kDefaultWorkspaceSlotCount; + if (slot.completion_recorded) { + const auto status = cudaEventSynchronize(slot.completion); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` failed while waiting to reuse CUDA " + "workspace"); + slot.completion_recorded = false; + } + slot.stream = stream; + return &slot; +} + +void Operator::RecordDefaultWorkspaceUse(DefaultWorkspaceSlot* slot, + cudaStream_t stream) { + const auto status = cudaEventRecord(slot->completion, stream); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` failed to record CUDA workspace use"); + slot->completion_recorded = true; +} + +void Operator::operator()( + const Tensor logits, const Tensor top_k, const Tensor top_p, + const std::optional indices, const std::string filter_apply_order, + const bool deterministic, const bool check_nan, + const std::optional seed, const std::optional offset, + Tensor out) const { + ValidateSupportedOptions(filter_apply_order, deterministic, check_nan); + ValidateHostTensor(top_k); + ValidateHostTensor(top_p); + ValidateIndices(indices); + assert(logits.IsContiguous() && out.IsContiguous() && + "The NVIDIA `TopKTopPSamplingFromLogits` provider requires " + "contiguous logits and output"); + + if (batch_size_ == 0) return; + + DeviceGuard device_guard{device_index_}; + const auto stream = static_cast(stream_ ? stream_ : nullptr); + DefaultWorkspaceSlot* default_workspace_slot = nullptr; + void* workspace = workspace_; + if (workspace == nullptr) { + default_workspace_slot = AcquireDefaultWorkspaceSlot(stream); + workspace = default_workspace_slot->workspace; + } + const auto workspace_size = + workspace_ ? workspace_size_in_bytes_ : workspace_size_; + assert(workspace != nullptr && workspace_size >= workspace_size_ && + "`TopKTopPSamplingFromLogits` received insufficient workspace"); + + const uint64_t actual_seed = + seed.has_value() ? static_cast(*seed) + : static_cast(std::random_device{}()); + const uint64_t actual_offset = static_cast(offset.value_or(0)); + const int vocab_size = static_cast(vocab_size_); + using OutputTypes = List; + + DispatchFunc( + {static_cast(logits.dtype()), static_cast(out.dtype())}, + [&](auto list_tag) { + using T = TypeMapType(list_tag)>; + using Tidx = TypeMapType(list_tag)>; + const auto* logits_ptr = static_cast(logits.data()); + auto* out_ptr = static_cast(out.data()); + + for (Tensor::Size row = 0; row < batch_size_; ++row) { + const int64_t logits_row = indices.has_value() + ? ReadIndex(*indices, row) + : static_cast(row); + assert(logits_row >= 0 && + static_cast(logits_row) < logits.size(0) && + "The NVIDIA `TopKTopPSamplingFromLogits` provider received " + "an out-of-range row index"); + const int64_t requested_top_k = ReadTopK(top_k, row); + const int normalized_top_k = requested_top_k <= 0 + ? vocab_size + : static_cast(std::min( + requested_top_k, vocab_size)); + const double requested_top_p = ReadTopP(top_p, row); + const double normalized_top_p = + requested_top_p > 0.0 && requested_top_p < 1.0 ? requested_top_p + : 1.0; + + top_k_top_p_sampling_from_logits_detail::SampleRow< + Device::Type::kNvidia>( + workspace, workspace_size, out_ptr + row, + logits_ptr + logits_row * logits.stride(0), vocab_size, + normalized_top_k, normalized_top_p, filter_apply_order == "joint", + CounterBasedUniform(actual_seed, + actual_offset + static_cast(row)), + stream); + } + }, + "Operator::operator()"); + + const auto status = cudaGetLastError(); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` CUDA kernel launch failed"); + if (default_workspace_slot != nullptr) { + RecordDefaultWorkspaceUse(default_workspace_slot, stream); + } +} + +void Operator::ValidateSupportedOptions(const std::string& + filter_apply_order, + bool deterministic, bool check_nan) { + assert( + (filter_apply_order == "top_k_first" || filter_apply_order == "joint") && + "The NVIDIA `TopKTopPSamplingFromLogits` provider supports only " + "`top_k_first` and `joint`"); + assert(deterministic && + "The NVIDIA `TopKTopPSamplingFromLogits` provider supports only the " + "deterministic path"); + assert(!check_nan && + "The NVIDIA `TopKTopPSamplingFromLogits` provider does not support " + "`check_nan`"); +} + +void Operator::ValidateHostTensor(const Tensor tensor) { + assert(tensor.device().type() == Device::Type::kCpu && + "The NVIDIA `TopKTopPSamplingFromLogits` provider requires host-side " + "`top_k` and `top_p` tensors"); + assert(tensor.IsContiguous() && + "The NVIDIA `TopKTopPSamplingFromLogits` provider requires " + "contiguous `top_k` and `top_p` tensors"); +} + +void Operator::ValidateIndices(const std::optional& indices) { + if (!indices.has_value()) return; + + assert(indices->device().type() == Device::Type::kCpu && + "The NVIDIA `TopKTopPSamplingFromLogits` provider requires " + "host-side `indices`"); + assert(indices->IsContiguous() && + "The NVIDIA `TopKTopPSamplingFromLogits` provider requires " + "contiguous `indices`"); +} + +int64_t Operator::ReadTopK(const Tensor top_k, Tensor::Size row) { + const auto element_offset = row * top_k.stride(0); + if (top_k.dtype() == DataType::kInt32) { + return static_cast(top_k.data())[element_offset]; + } + return static_cast(top_k.data())[element_offset]; +} + +double Operator::ReadTopP( + const Tensor top_p, Tensor::Size row) { + const auto element_offset = row * top_p.stride(0); + switch (top_p.dtype()) { + case DataType::kFloat16: + return static_cast(top_p.data())[element_offset] + .ToFloat(); + case DataType::kBFloat16: + return static_cast(top_p.data())[element_offset] + .ToFloat(); + case DataType::kFloat32: + return static_cast(top_p.data())[element_offset]; + case DataType::kFloat64: + return static_cast(top_p.data())[element_offset]; + default: + assert(false && + "`TopKTopPSamplingFromLogits` received unsupported `top_p` " + "dtype"); + return 1.0; + } +} + +int64_t Operator::ReadIndex(const Tensor indices, Tensor::Size row) { + const auto element_offset = row * indices.stride(0); + if (indices.dtype() == DataType::kInt32) { + return static_cast(indices.data())[element_offset]; + } + return static_cast(indices.data())[element_offset]; +} + +std::size_t Operator::DispatchWorkspaceSize(DataType dtype, + Tensor::Size vocab_size) { + std::size_t workspace_size = 0; + DispatchFunc( + dtype, + [&](auto tag) { + using T = typename decltype(tag)::type; + workspace_size = + top_k_top_p_sampling_from_logits_detail::WorkspaceSize( + static_cast(vocab_size)); + }, + "Operator::DispatchWorkspaceSize"); + return workspace_size; +} + +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.cuh b/src/native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.cuh new file mode 100644 index 000000000..182a600aa --- /dev/null +++ b/src/native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.cuh @@ -0,0 +1,178 @@ +#ifndef INFINI_OPS_NVIDIA_TOP_K_TOP_P_SAMPLING_FROM_LOGITS_KERNEL_CUH_ +#define INFINI_OPS_NVIDIA_TOP_K_TOP_P_SAMPLING_FROM_LOGITS_KERNEL_CUH_ + +#include +#include +#include +#include +#include +#include + +#include "data_type.h" +#include "native/cuda/nvidia/caster.cuh" + +namespace infini::ops::top_k_top_p_sampling_from_logits_detail { + +constexpr std::size_t Align256(std::size_t size) { + return (size + 255) & ~std::size_t{255}; +} + +template +struct Workspace { + int32_t* indices; + T* sorted_logits; + int32_t* sorted_indices; + double* cumulative_probabilities; + void* temporary_storage; + std::size_t temporary_storage_size; +}; + +template +std::size_t WorkspaceSize(int vocab_size) { + std::size_t sort_size = 0; + auto status = cub::DeviceRadixSort::SortPairsDescending( + nullptr, sort_size, static_cast(nullptr), + static_cast(nullptr), static_cast(nullptr), + static_cast(nullptr), vocab_size, 0, sizeof(T) * 8, nullptr); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` failed to query CUB sort workspace"); + + std::size_t scan_size = 0; + status = cub::DeviceScan::InclusiveSum( + nullptr, scan_size, static_cast(nullptr), + static_cast(nullptr), vocab_size, nullptr); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` failed to query CUB scan workspace"); + + return Align256(sizeof(int32_t) * static_cast(vocab_size)) + + Align256(sizeof(T) * static_cast(vocab_size)) + + Align256(sizeof(int32_t) * static_cast(vocab_size)) + + Align256(sizeof(double) * static_cast(vocab_size)) + + Align256(sort_size > scan_size ? sort_size : scan_size); +} + +template +Workspace PartitionWorkspace(void* workspace, std::size_t workspace_size, + int vocab_size) { + auto* cursor = static_cast(workspace); + const auto indices_size = + Align256(sizeof(int32_t) * static_cast(vocab_size)); + const auto sorted_logits_size = + Align256(sizeof(T) * static_cast(vocab_size)); + const auto sorted_indices_size = + Align256(sizeof(int32_t) * static_cast(vocab_size)); + const auto probabilities_size = + Align256(sizeof(double) * static_cast(vocab_size)); + + auto* indices = reinterpret_cast(cursor); + cursor += indices_size; + auto* sorted_logits = reinterpret_cast(cursor); + cursor += sorted_logits_size; + auto* sorted_indices = reinterpret_cast(cursor); + cursor += sorted_indices_size; + auto* cumulative_probabilities = reinterpret_cast(cursor); + cursor += probabilities_size; + + const auto fixed_size = indices_size + sorted_logits_size + + sorted_indices_size + probabilities_size; + assert(workspace_size >= fixed_size && + "`TopKTopPSamplingFromLogits` received insufficient workspace"); + + return {indices, sorted_logits, + sorted_indices, cumulative_probabilities, + cursor, workspace_size - fixed_size}; +} + +__global__ void FillIndicesKernel(int32_t* indices, int vocab_size) { + const int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index < vocab_size) indices[index] = index; +} + +template +__global__ void LogitsToProbabilitiesKernel(const T* sorted_logits, + double* probabilities, + int vocab_size) { + const int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index >= vocab_size) return; + + const double maximum = + Caster::template Cast(sorted_logits[0]); + const double value = + Caster::template Cast(sorted_logits[index]); + probabilities[index] = exp(value - maximum); +} + +template +__global__ void SampleKernel(Tidx* out, const double* cumulative_probabilities, + const int32_t* sorted_indices, int vocab_size, + int top_k, double top_p, bool joint_filter, + double random_value) { + if (blockIdx.x != 0 || threadIdx.x != 0) return; + + int keep_count = top_k; + double retained_probability_mass = cumulative_probabilities[keep_count - 1]; + if (joint_filter) { + const double top_p_probability_mass = + top_p * cumulative_probabilities[vocab_size - 1]; + if (top_p_probability_mass < retained_probability_mass) { + retained_probability_mass = top_p_probability_mass; + } + } else if (top_p > 0.0 && top_p < 1.0) { + const double threshold = top_p * cumulative_probabilities[keep_count - 1]; + for (int i = 0; i < keep_count; ++i) { + if (cumulative_probabilities[i] >= threshold) { + keep_count = i + 1; + break; + } + } + retained_probability_mass = cumulative_probabilities[keep_count - 1]; + } + + const double threshold = random_value * retained_probability_mass; + int selected = 0; + while (selected + 1 < keep_count && + cumulative_probabilities[selected] < threshold) { + ++selected; + } + *out = static_cast(sorted_indices[selected]); +} + +template +void SampleRow(void* workspace, std::size_t workspace_size, Tidx* out, + const T* logits, int vocab_size, int top_k, double top_p, + bool joint_filter, double random_value, cudaStream_t stream) { + auto partition = PartitionWorkspace(workspace, workspace_size, vocab_size); + constexpr int kBlockSize = 256; + const auto grid_size = static_cast( + vocab_size / kBlockSize + (vocab_size % kBlockSize != 0)); + + FillIndicesKernel<<>>(partition.indices, + vocab_size); + + auto temporary_storage_size = partition.temporary_storage_size; + auto status = cub::DeviceRadixSort::SortPairsDescending( + partition.temporary_storage, temporary_storage_size, logits, + partition.sorted_logits, partition.indices, partition.sorted_indices, + vocab_size, 0, sizeof(T) * 8, stream); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` CUB radix sort failed"); + + LogitsToProbabilitiesKernel<<>>( + partition.sorted_logits, partition.cumulative_probabilities, vocab_size); + + temporary_storage_size = partition.temporary_storage_size; + status = cub::DeviceScan::InclusiveSum( + partition.temporary_storage, temporary_storage_size, + partition.cumulative_probabilities, partition.cumulative_probabilities, + vocab_size, stream); + assert(status == cudaSuccess && + "`TopKTopPSamplingFromLogits` CUB inclusive scan failed"); + + SampleKernel<<<1, 1, 0, stream>>>( + out, partition.cumulative_probabilities, partition.sorted_indices, + vocab_size, top_k, top_p, joint_filter, random_value); +} + +} // namespace infini::ops::top_k_top_p_sampling_from_logits_detail + +#endif // INFINI_OPS_NVIDIA_TOP_K_TOP_P_SAMPLING_FROM_LOGITS_KERNEL_CUH_ diff --git a/src/native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.h b/src/native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.h new file mode 100644 index 000000000..708ab434e --- /dev/null +++ b/src/native/cuda/nvidia/ops/top_k_top_p_sampling_from_logits/kernel.h @@ -0,0 +1,81 @@ +#ifndef INFINI_OPS_NVIDIA_TOP_K_TOP_P_SAMPLING_FROM_LOGITS_KERNEL_H_ +#define INFINI_OPS_NVIDIA_TOP_K_TOP_P_SAMPLING_FROM_LOGITS_KERNEL_H_ + +#include + +#include +#include +#include +#include +#include + +#include "base/top_k_top_p_sampling_from_logits.h" + +namespace infini::ops { + +template <> +class Operator + : public TopKTopPSamplingFromLogits { + public: + Operator(const Tensor logits, const Tensor top_k, const Tensor top_p, + const std::optional indices, + const std::string filter_apply_order, const bool deterministic, + const bool check_nan, const std::optional seed, + const std::optional offset, Tensor out); + + ~Operator() override; + + std::size_t workspace_size_in_bytes() const override; + + void operator()(const Tensor logits, const Tensor top_k, const Tensor top_p, + const std::optional indices, + const std::string filter_apply_order, + const bool deterministic, const bool check_nan, + const std::optional seed, + const std::optional offset, + Tensor out) const override; + + private: + static void ValidateSupportedOptions(const std::string& filter_apply_order, + bool deterministic, bool check_nan); + + static void ValidateHostTensor(const Tensor tensor); + + static void ValidateIndices(const std::optional& indices); + + static int64_t ReadTopK(const Tensor top_k, Tensor::Size row); + + static double ReadTopP(const Tensor top_p, Tensor::Size row); + + static int64_t ReadIndex(const Tensor indices, Tensor::Size row); + + static std::size_t DispatchWorkspaceSize(DataType dtype, + Tensor::Size vocab_size); + + struct DefaultWorkspaceSlot { + void* workspace{nullptr}; + cudaStream_t stream{nullptr}; + cudaEvent_t completion{nullptr}; + bool completion_recorded{false}; + }; + + DefaultWorkspaceSlot* AcquireDefaultWorkspaceSlot(cudaStream_t stream) const; + + static void RecordDefaultWorkspaceUse(DefaultWorkspaceSlot* slot, + cudaStream_t stream); + + int device_index_{0}; + + std::size_t workspace_size_{0}; + + static constexpr std::size_t kDefaultWorkspaceSlotCount = 2; + + mutable std::array + default_workspace_slots_{}; + + mutable std::size_t next_default_workspace_slot_{0}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_TOP_K_TOP_P_SAMPLING_FROM_LOGITS_KERNEL_H_ diff --git a/tests/test_cpp_api.py b/tests/test_cpp_api.py index 28cce589a..fe901bcfd 100644 --- a/tests/test_cpp_api.py +++ b/tests/test_cpp_api.py @@ -110,6 +110,66 @@ def test_cpp_polymorphic_context_smoke(tmp_path): _run([str(binary)]) +def test_cpp_top_k_top_p_sampling_uses_caller_workspace(tmp_path): + import infini.ops + + op_class = getattr(infini.ops, "TopKTopPSamplingFromLogits", None) + + if op_class is None or not op_class.active_implementation_indices("nvidia"): + pytest.skip("NVIDIA `TopKTopPSamplingFromLogits` is not active") + + install_prefix = _install_prefix() + include_dir = install_prefix / "include" + library_dir = _library_dir(install_prefix) + source_include_dir = Path(__file__).resolve().parents[1] / "src" + cuda_home = Path(os.environ.get("CUDA_HOME", "/usr/local/cuda")) + cuda_include_dir = cuda_home / "include" + cuda_library_dir = next( + ( + path + for path in ( + cuda_home / "lib64", + cuda_home / "targets" / "x86_64-linux" / "lib", + ) + if path.exists() + ), + None, + ) + + if not (cuda_include_dir / "cuda_runtime_api.h").exists(): + pytest.skip("CUDA headers are not available") + + if cuda_library_dir is None: + pytest.skip("CUDA runtime libraries are not available") + + source = tmp_path / "top_k_top_p_sampling_workspace.cc" + binary = tmp_path / "top_k_top_p_sampling_workspace" + source.write_text(_TOP_K_TOP_P_SAMPLING_WORKSPACE_SOURCE) + + _run( + [ + _compiler("CXX", "c++"), + "-std=c++17", + "-Werror", + "-Wno-error=deprecated-declarations", + f"-I{source_include_dir}", + f"-I{include_dir}", + f"-I{cuda_include_dir}", + str(source), + f"-L{library_dir}", + f"-L{cuda_library_dir}", + "-linfiniops", + "-linfinirt", + "-lcudart", + f"-Wl,-rpath,{library_dir}", + f"-Wl,-rpath,{cuda_library_dir}", + "-o", + str(binary), + ] + ) + _run([str(binary)]) + + @pytest.mark.parametrize( "header", ( @@ -174,6 +234,93 @@ def _run(command): raise AssertionError(output) from error +_TOP_K_TOP_P_SAMPLING_WORKSPACE_SOURCE = textwrap.dedent( + r""" + #include + #include + #include + + #include + #include + #include + #include + + int main() { + using infini::ops::DataType; + using infini::ops::Device; + using infini::ops::Handle; + using infini::ops::Operator; + using infini::ops::Tensor; + using infini::ops::TopKTopPSamplingFromLogits; + + const std::array host_logits{ + 0.0f, 5.0f, 1.0f, 2.0f, 0.0f, 1.0f, 2.0f, 5.0f}; + std::array host_top_k{1, 1}; + std::array host_top_p{1.0, 1.0}; + std::array host_out{}; + float* device_logits = nullptr; + int32_t* device_out = nullptr; + + if (cudaMalloc(reinterpret_cast(&device_logits), + sizeof(host_logits)) != cudaSuccess) { + return 1; + } + if (cudaMalloc(reinterpret_cast(&device_out), sizeof(host_out)) != + cudaSuccess) { + return 2; + } + if (cudaMemcpy(device_logits, host_logits.data(), sizeof(host_logits), + cudaMemcpyHostToDevice) != cudaSuccess) { + return 3; + } + + const Device host{Device::Type::kCpu}; + const Device nvidia{Device::Type::kNvidia}; + const Tensor logits(device_logits, Tensor::Shape{2, 4}, + DataType::kFloat32, nvidia); + const Tensor top_k(host_top_k.data(), Tensor::Shape{2}, + DataType::kInt32, host); + const Tensor top_p(host_top_p.data(), Tensor::Shape{2}, + DataType::kFloat64, host); + const Tensor out(device_out, Tensor::Shape{2}, DataType::kInt32, nvidia); + const std::optional indices; + const std::string filter_apply_order{"top_k_first"}; + const std::optional seed{1234}; + const std::optional offset{0}; + + Operator op( + logits, top_k, top_p, indices, filter_apply_order, true, false, seed, + offset, out); + const auto workspace_size = op.workspace_size_in_bytes(); + void* workspace = nullptr; + if (workspace_size == 0 || cudaMalloc(&workspace, workspace_size) != + cudaSuccess) { + return 4; + } + + Handle handle; + handle.set_workspace(workspace); + handle.set_workspace_size_in_bytes(workspace_size); + auto& callable = static_cast&>(op); + callable(handle, logits, top_k, top_p, indices, filter_apply_order, true, + false, seed, offset, out); + + if (cudaDeviceSynchronize() != cudaSuccess) return 5; + if (cudaMemcpy(host_out.data(), device_out, sizeof(host_out), + cudaMemcpyDeviceToHost) != cudaSuccess) { + return 6; + } + + cudaFree(workspace); + cudaFree(device_out); + cudaFree(device_logits); + + return host_out == std::array{1, 3} ? 0 : 7; + } + """ +).lstrip() + + _ADD_SMOKE_SOURCE = textwrap.dedent( r""" #include diff --git a/tests/test_top_k_top_p_sampling_from_logits.py b/tests/test_top_k_top_p_sampling_from_logits.py index d3c9b79af..de6836a4e 100644 --- a/tests/test_top_k_top_p_sampling_from_logits.py +++ b/tests/test_top_k_top_p_sampling_from_logits.py @@ -50,6 +50,385 @@ def test_top_k_top_p_sampling_from_logits( assert torch.all(torch.isin(first, allowed_tensor)) +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_top_k_top_p_sampling_from_logits_per_request_filters( + dtype, + device, + implementation_index, +): + logits = torch.tensor( + ( + (5.0, 4.0, 3.0, -10.0), + (5.0, 4.0, 3.0, -10.0), + (5.0, 4.0, 3.0, -10.0), + (2.0, 1.9, 1.8, 1.7), + ), + dtype=dtype, + device=device, + ) + top_k = torch.tensor((1, 3, 2, 2), dtype=torch.int32) + top_p = torch.tensor((1.0, 0.8, 1.0, 0.4), dtype=torch.float64) + out = torch.empty((4,), dtype=torch.int32, device=device) + + _top_k_top_p_sampling_from_logits( + logits, top_k, top_p, 1234, 9, out, implementation_index + ) + + assert out[0].item() == 0 + assert out[1].item() in (0, 1) + assert out[2].item() in (0, 1) + assert out[3].item() == 0 + + +def test_top_k_top_p_sampling_from_logits_top_k_bounds_and_singleton( + device, + implementation_index, +): + if device != "cuda": + pytest.skip("NVIDIA edge-case coverage requires CUDA") + + batch_size = 64 + vocab_size = 32 + logits = torch.zeros((batch_size, vocab_size), dtype=torch.float32, device=device) + top_p = torch.ones((batch_size,), dtype=torch.float64) + outputs = [] + + for top_k_value in (0, -3, vocab_size + 7): + top_k = torch.full((batch_size,), top_k_value, dtype=torch.int64) + out = torch.empty((batch_size,), dtype=torch.int32, device=device) + _top_k_top_p_sampling_from_logits( + logits, top_k, top_p, 1234, 9, out, implementation_index + ) + outputs.append(out) + + assert torch.equal(outputs[0], outputs[1]) + assert torch.equal(outputs[0], outputs[2]) + + singleton_logits = torch.randn((batch_size, 1), dtype=torch.float32, device=device) + singleton_top_k = torch.zeros((batch_size,), dtype=torch.int32) + singleton_out = torch.empty((batch_size,), dtype=torch.int32, device=device) + _top_k_top_p_sampling_from_logits( + singleton_logits, + singleton_top_k, + top_p, + 1234, + 9, + singleton_out, + implementation_index, + ) + + assert torch.count_nonzero(singleton_out).item() == 0 + + +def test_top_k_top_p_sampling_from_logits_large_offset( + device, + implementation_index, +): + if device != "cuda": + pytest.skip("NVIDIA counter-based RNG coverage requires CUDA") + + batch_size = 64 + vocab_size = 16 + logits = torch.zeros((batch_size, vocab_size), dtype=torch.float32, device=device) + top_k = torch.full((batch_size,), vocab_size, dtype=torch.int32) + top_p = torch.ones((batch_size,), dtype=torch.float64) + first = torch.empty((batch_size,), dtype=torch.int32, device=device) + shifted = torch.empty_like(first) + repeated = torch.empty_like(first) + offset = 2**62 + + _top_k_top_p_sampling_from_logits( + logits, top_k, top_p, 1234, offset, first, implementation_index + ) + _top_k_top_p_sampling_from_logits( + logits, top_k, top_p, 1234, offset + 1, shifted, implementation_index + ) + _top_k_top_p_sampling_from_logits( + logits, top_k, top_p, 1234, offset, repeated, implementation_index + ) + + assert torch.equal(first, repeated) + assert torch.equal(first[1:], shifted[:-1]) + assert not torch.equal(first, shifted) + + +def test_top_k_top_p_sampling_from_logits_nvidia_flat_distribution(pytestconfig): + requested_devices = pytestconfig.getoption("--devices") or () + if "nvidia" not in requested_devices: + return + + assert torch.cuda.is_available() + assert 0 in infini.ops.TopKTopPSamplingFromLogits.active_implementation_indices( + "nvidia" + ) + + batch_size = 4096 + vocab_size = 8 + logits = torch.zeros((batch_size, vocab_size), dtype=torch.float32, device="cuda") + top_k = torch.full((batch_size,), vocab_size, dtype=torch.int32) + top_p = torch.ones((batch_size,), dtype=torch.float64) + out = torch.empty((batch_size,), dtype=torch.int32, device="cuda") + + _top_k_top_p_sampling_from_logits(logits, top_k, top_p, 20260811, 0, out, 0) + + counts = torch.bincount(out.to(torch.int64), minlength=vocab_size).cpu() + expected_count = batch_size / vocab_size + assert counts.numel() == vocab_size + assert torch.all(torch.abs(counts - expected_count) < expected_count * 0.15) + + +def test_top_k_top_p_sampling_from_logits_joint_filter_semantics( + device, + implementation_index, +): + if device != "cuda": + pytest.skip("NVIDIA joint-filter coverage requires CUDA") + + batch_size = 4096 + probabilities = torch.tensor( + (0.4, 0.3, 0.2, 0.1), dtype=torch.float32, device=device + ) + logits = probabilities.log().expand(batch_size, -1).contiguous() + top_k = torch.full((batch_size,), 2, dtype=torch.int32) + top_p = torch.full((batch_size,), 0.5, dtype=torch.float64) + joint = torch.empty((batch_size,), dtype=torch.int32, device=device) + repeated_joint = torch.empty_like(joint) + top_k_first = torch.empty_like(joint) + + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 20260811, + 0, + joint, + implementation_index, + filter_apply_order="joint", + ) + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 20260811, + 0, + repeated_joint, + implementation_index, + filter_apply_order="joint", + ) + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 20260811, + 0, + top_k_first, + implementation_index, + ) + + assert torch.equal(joint, repeated_joint) + assert torch.count_nonzero(top_k_first).item() == 0 + assert torch.all((joint == 0) | (joint == 1)) + second_token_count = torch.count_nonzero(joint == 1).item() + assert batch_size * 0.15 < second_token_count < batch_size * 0.25 + + +@pytest.mark.parametrize( + "top_p_value, second_logit", + ( + (1e-300, 0.0), + (0.999999975, -17.72753356339242), + ), +) +def test_top_k_top_p_sampling_from_logits_float64_top_p_boundaries( + top_p_value, + second_logit, + device, + implementation_index, +): + if device != "cuda": + pytest.skip("NVIDIA float64 `top_p` coverage requires CUDA") + + logits = torch.tensor(((0.0, second_logit),), dtype=torch.float32, device=device) + top_k = torch.tensor((2,), dtype=torch.int32) + top_p = torch.tensor((top_p_value,), dtype=torch.float64) + out = torch.empty((1,), dtype=torch.int32, device=device) + + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 5441626385717431455, + 0, + out, + implementation_index, + ) + + assert out.item() == 0 + + +@pytest.mark.parametrize("indices_dtype", (torch.int32, torch.int64)) +def test_top_k_top_p_sampling_from_logits_host_indices( + indices_dtype, + device, + implementation_index, +): + if device != "cuda": + pytest.skip("NVIDIA host `indices` coverage requires CUDA") + + logits = torch.tensor( + ( + (0.0, 5.0, 1.0, 2.0), + (0.0, 1.0, 2.0, 5.0), + (5.0, 1.0, 2.0, 0.0), + ), + dtype=torch.float32, + device=device, + ) + indices = torch.tensor((2, 0), dtype=indices_dtype) + top_k = torch.ones((2,), dtype=torch.int32) + top_p = torch.ones((2,), dtype=torch.float64) + out = torch.empty((2,), dtype=indices_dtype, device=device) + + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 0, + out, + implementation_index, + indices=indices, + ) + + expected = torch.tensor((0, 1), dtype=indices_dtype, device=device) + assert torch.equal(out, expected) + + +def test_top_k_top_p_sampling_from_logits_default_workspace_per_stream( + device, + implementation_index, +): + if device != "cuda": + pytest.skip("NVIDIA multi-stream coverage requires CUDA") + + batch_size = 8 + vocab_size = 8192 + logits_a = torch.randn((batch_size, vocab_size), dtype=torch.float32, device=device) + logits_b = -logits_a.flip(1).contiguous() + top_k = torch.full((batch_size,), 64, dtype=torch.int32) + top_p = torch.full((batch_size,), 0.95, dtype=torch.float64) + baseline_a = torch.empty((batch_size,), dtype=torch.int32, device=device) + baseline_b = torch.empty_like(baseline_a) + + _top_k_top_p_sampling_from_logits( + logits_a, top_k, top_p, 1234, 9, baseline_a, implementation_index + ) + _top_k_top_p_sampling_from_logits( + logits_b, top_k, top_p, 1234, 9, baseline_b, implementation_index + ) + torch.cuda.synchronize() + + stream_a = torch.cuda.Stream() + stream_b = torch.cuda.Stream() + stream_a.wait_stream(torch.cuda.current_stream()) + stream_b.wait_stream(torch.cuda.current_stream()) + warm_a = torch.empty_like(baseline_a) + warm_b = torch.empty_like(baseline_b) + _top_k_top_p_sampling_from_logits( + logits_a, + top_k, + top_p, + 1234, + 9, + warm_a, + implementation_index, + stream=stream_a.cuda_stream, + ) + stream_a.synchronize() + _top_k_top_p_sampling_from_logits( + logits_b, + top_k, + top_p, + 1234, + 9, + warm_b, + implementation_index, + stream=stream_b.cuda_stream, + ) + stream_b.synchronize() + assert torch.equal(warm_a, baseline_a) + assert torch.equal(warm_b, baseline_b) + + outputs_a = [torch.empty_like(baseline_a) for _ in range(4)] + outputs_b = [torch.empty_like(baseline_b) for _ in range(4)] + gate_stream = torch.cuda.Stream() + gate_stream.wait_stream(torch.cuda.current_stream()) + gate = torch.cuda.Event() + with torch.cuda.stream(gate_stream): + torch.cuda._sleep(100_000_000) + gate.record() + + stream_a.wait_event(gate) + stream_b.wait_event(gate) + + for out_a, out_b in zip(outputs_a, outputs_b): + _top_k_top_p_sampling_from_logits( + logits_a, + top_k, + top_p, + 1234, + 9, + out_a, + implementation_index, + stream=stream_a.cuda_stream, + ) + _top_k_top_p_sampling_from_logits( + logits_b, + top_k, + top_p, + 1234, + 9, + out_b, + implementation_index, + stream=stream_b.cuda_stream, + ) + + stream_c = torch.cuda.Stream() + stream_c.wait_stream(torch.cuda.current_stream()) + turnover_c = torch.empty_like(baseline_a) + turnover_a = torch.empty_like(baseline_b) + _top_k_top_p_sampling_from_logits( + logits_a, + top_k, + top_p, + 1234, + 9, + turnover_c, + implementation_index, + stream=stream_c.cuda_stream, + ) + _top_k_top_p_sampling_from_logits( + logits_b, + top_k, + top_p, + 1234, + 9, + turnover_a, + implementation_index, + stream=stream_a.cuda_stream, + ) + + stream_a.synchronize() + stream_b.synchronize() + stream_c.synchronize() + + for out_a, out_b in zip(outputs_a, outputs_b): + assert torch.equal(out_a, baseline_a) + assert torch.equal(out_b, baseline_b) + assert torch.equal(turnover_c, baseline_a) + assert torch.equal(turnover_a, baseline_b) + + def _top_k_top_p_sampling_from_logits( logits, top_k, @@ -58,18 +437,25 @@ def _top_k_top_p_sampling_from_logits( offset, out, implementation_index, + *, + indices=None, + stream=None, + filter_apply_order="top_k_first", ): + if stream is None: + stream = get_stream(logits.device) + infini.ops.top_k_top_p_sampling_from_logits( logits, top_k, top_p, - None, - "top_k_first", + indices, + filter_apply_order, True, False, seed, offset, out, - stream=get_stream(logits.device), + stream=stream, implementation_index=implementation_index, )