diff --git a/src/linked/tvm_ffi/nvidia/flashinfer_sampling.yaml b/src/linked/tvm_ffi/nvidia/flashinfer_sampling.yaml new file mode 100644 index 000000000..2d801e4c0 --- /dev/null +++ b/src/linked/tvm_ffi/nvidia/flashinfer_sampling.yaml @@ -0,0 +1,3 @@ +python_distribution_package: flashinfer-jit-cache +python_distribution_version: ">=0.6.7,<0.7" +library_glob: flashinfer_jit_cache/jit_cache/sampling/sampling.so diff --git a/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.cu b/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.cu new file mode 100644 index 000000000..45fa3c97c --- /dev/null +++ b/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.cu @@ -0,0 +1,776 @@ +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "data_type.h" +#include "dispatcher.h" +#include "linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.h" +#include "native/cpu/caster_.h" +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" + +extern "C" { +int __tvm_ffi_softmax(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); +int __tvm_ffi_top_k_mask_logits(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); +int __tvm_ffi_top_p_sampling_from_probs(void*, const TVMFFIAny*, int32_t, + TVMFFIAny*); +int __tvm_ffi_top_k_top_p_sampling_from_probs(void*, const TVMFFIAny*, int32_t, + TVMFFIAny*); +} + +namespace infini::ops { +namespace { + +using OptionalTensorView = tvm::ffi::Optional; + +constexpr std::size_t kScratchBytes = 1024 * 1024; +constexpr std::size_t kAlignment = 256; +constexpr unsigned int kThreads = 256; + +void Require(bool condition, const char* message) { + if (!condition) { + throw std::invalid_argument(message); + } +} + +void CheckCuda(cudaError_t status, const char* operation) { + if (status != cudaSuccess) { + throw std::runtime_error(std::string{operation} + ": " + + cudaGetErrorString(status)); + } +} + +void ReportCuda(cudaError_t status, const char* operation) noexcept { + if (status != cudaSuccess) { + std::fprintf(stderr, "[InfiniOps] %s: %s\n", operation, + cudaGetErrorString(status)); + } +} + +void CheckTvmFfi(int status, const char* operation) { + if (status != 0) { + throw std::runtime_error(std::string{operation} + " (status " + + std::to_string(status) + ")"); + } +} + +void ReportTvmFfi(int status, const char* operation) noexcept { + if (status != 0) { + std::fprintf(stderr, "[InfiniOps] %s (status %d)\n", operation, status); + } +} + +std::size_t Align(std::size_t value) { + return (value + kAlignment - 1) & ~(kAlignment - 1); +} + +std::size_t AddWorkspaceRegion(std::size_t* offset, std::size_t size) { + *offset = Align(*offset); + const auto result = *offset; + *offset += size; + return result; +} + +struct WorkspaceLayout { + explicit WorkspaceLayout(std::size_t matrix_elements, + std::size_t batch_size) { + matrix_a = AddWorkspaceRegion(&size, matrix_elements * sizeof(float)); + matrix_b = AddWorkspaceRegion(&size, matrix_elements * sizeof(float)); + top_k = AddWorkspaceRegion(&size, batch_size * sizeof(int64_t)); + top_p = AddWorkspaceRegion(&size, batch_size * sizeof(float)); + valid = AddWorkspaceRegion(&size, batch_size * sizeof(uint8_t)); + indices = AddWorkspaceRegion(&size, batch_size * sizeof(int64_t)); + scratch = AddWorkspaceRegion(&size, kScratchBytes); + size = Align(size); + } + + std::size_t matrix_a{0}; + std::size_t matrix_b{0}; + std::size_t top_k{0}; + std::size_t top_p{0}; + std::size_t valid{0}; + std::size_t indices{0}; + std::size_t scratch{0}; + std::size_t size{0}; +}; + +class DeviceGuard { + public: + explicit DeviceGuard(int device_index) { + auto status = cudaGetDevice(&previous_device_); + CheckCuda(status, + "FlashInferSampling failed to query the current CUDA device"); + if (previous_device_ != device_index) { + status = cudaSetDevice(device_index); + CheckCuda(status, + "FlashInferSampling failed to select the input CUDA device"); + restore_ = true; + } + } + + ~DeviceGuard() { + if (!restore_) return; + const auto status = cudaSetDevice(previous_device_); + ReportCuda(status, + "FlashInferSampling failed to restore the CUDA device"); + } + + private: + int previous_device_{0}; + bool restore_{false}; +}; + +class StreamGuard { + public: + StreamGuard(int device_index, cudaStream_t stream) + : device_index_{device_index} { + const auto status = + TVMFFIEnvSetStream(kDLCUDA, device_index, stream, &previous_stream_); + CheckTvmFfi(status, + "FlashInferSampling failed to set the TVM FFI CUDA stream"); + } + + ~StreamGuard() { + const auto status = + TVMFFIEnvSetStream(kDLCUDA, device_index_, previous_stream_, nullptr); + ReportTvmFfi(status, + "FlashInferSampling failed to restore the TVM FFI CUDA stream"); + } + + private: + int device_index_{0}; + TVMFFIStreamHandle previous_stream_{nullptr}; +}; + +class Workspace { + public: + Workspace(void* external, std::size_t available, std::size_t required, + cudaStream_t stream) + : data_{external}, stream_{stream} { + if (external != nullptr) { + Require(available >= required, + "FlashInferSampling received insufficient workspace"); + return; + } + + CheckCuda(cudaMallocAsync(&data_, required, stream_), + "FlashInferSampling failed to allocate async workspace"); + owned_ = true; + } + + ~Workspace() { + if (owned_) { + ReportCuda(cudaFreeAsync(data_, stream_), + "FlashInferSampling failed to free async workspace"); + } + } + + Workspace(const Workspace&) = delete; + Workspace& operator=(const Workspace&) = delete; + + void* data() const { return data_; } + + private: + void* data_{nullptr}; + cudaStream_t stream_{nullptr}; + bool owned_{false}; +}; + +class StagingRecorder { + public: + StagingRecorder(cudaEvent_t event, cudaStream_t stream, bool* recorded) + : event_{event}, stream_{stream}, recorded_{recorded} { + *recorded_ = false; + } + + ~StagingRecorder() { + if (!active_) return; + + const auto status = cudaEventRecord(event_, stream_); + if (status == cudaSuccess) { + *recorded_ = true; + return; + } + + ReportCuda(status, + "FlashInferSampling failed to record staging completion"); + ReportCuda(cudaStreamSynchronize(stream_), + "FlashInferSampling failed to await staging after an error"); + } + + StagingRecorder(const StagingRecorder&) = delete; + StagingRecorder& operator=(const StagingRecorder&) = delete; + + void Record() { + active_ = false; + const auto status = cudaEventRecord(event_, stream_); + if (status != cudaSuccess) { + CheckCuda( + cudaStreamSynchronize(stream_), + "FlashInferSampling failed to await staging after a record error"); + CheckCuda(status, + "FlashInferSampling failed to record staging completion"); + } + *recorded_ = true; + } + + private: + cudaEvent_t event_{nullptr}; + cudaStream_t stream_{nullptr}; + bool* recorded_{nullptr}; + bool active_{true}; +}; + +DLDataType Dtype(DataType dtype) { + switch (dtype) { + case DataType::kInt32: + return {kDLInt, 32, 1}; + case DataType::kInt64: + return {kDLInt, 64, 1}; + case DataType::kFloat32: + return {kDLFloat, 32, 1}; + default: + throw std::invalid_argument( + "FlashInferSampling received an unsupported dtype"); + } +} + +DLTensor MakeTensor(void* data, int device_index, int32_t ndim, int64_t* shape, + DLDataType dtype) { + return {data, {kDLCUDA, device_index}, ndim, dtype, shape, nullptr, 0}; +} + +template +__global__ void CastLogits(float* dst, const Src* src, std::size_t count) { + for (auto index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < count; + index += static_cast(blockDim.x) * gridDim.x) { + dst[index] = Caster::Cast(src[index]); + } +} + +template +__global__ void GatherCastLogits(float* dst, const Src* src, + const Index* indices, std::size_t rows, + std::size_t source_rows, + std::size_t vocab_size) { + const auto count = rows * vocab_size; + for (auto index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < count; + index += static_cast(blockDim.x) * gridDim.x) { + const auto row = index / vocab_size; + const auto column = index % vocab_size; + const auto source_index = indices[row]; + if (source_index < 0 || + static_cast(source_index) >= source_rows) { + asm("trap;"); + return; + } + const auto source_row = static_cast(source_index); + dst[index] = Caster::Cast( + src[source_row * vocab_size + column]); + } +} + +unsigned int Blocks(std::size_t count) { + const auto blocks = (count + kThreads - 1) / kThreads; + return static_cast(std::min(blocks, 65535)); +} + +void CallSoftmax(DLTensor* scratch, DLTensor* logits, DLTensor* output) { + tvm::ffi::Function::InvokeExternC( + nullptr, __tvm_ffi_softmax, tvm::ffi::TensorView(scratch), + tvm::ffi::TensorView(logits), tvm::ffi::TensorView(output), + OptionalTensorView{}, 1.0, false); +} + +void CallTopKMask(DLTensor* logits, DLTensor* output, DLTensor* top_k, + DLTensor* scratch) { + tvm::ffi::Function::InvokeExternC( + nullptr, __tvm_ffi_top_k_mask_logits, tvm::ffi::TensorView(logits), + tvm::ffi::TensorView(output), + OptionalTensorView{tvm::ffi::TensorView(top_k)}, int64_t{0}, + tvm::ffi::TensorView(scratch)); +} + +OptionalTensorView OptionalView(DLTensor* tensor) { + return tensor == nullptr ? OptionalTensorView{} + : OptionalTensorView{tvm::ffi::TensorView(tensor)}; +} + +void CallTopP(DLTensor* probs, DLTensor* output, DLTensor* valid, + DLTensor* indices, DLTensor* top_p, bool deterministic, + uint64_t seed, uint64_t offset) { + tvm::ffi::Function::InvokeExternC( + nullptr, __tvm_ffi_top_p_sampling_from_probs, tvm::ffi::TensorView(probs), + tvm::ffi::TensorView(output), tvm::ffi::TensorView(valid), + OptionalView(indices), OptionalTensorView{tvm::ffi::TensorView(top_p)}, + 1.0, deterministic, OptionalTensorView{}, seed, OptionalTensorView{}, + offset); +} + +void CallJoint(DLTensor* probs, DLTensor* output, DLTensor* valid, + DLTensor* indices, DLTensor* top_k, DLTensor* top_p, + bool deterministic, uint64_t seed, uint64_t offset) { + tvm::ffi::Function::InvokeExternC( + nullptr, __tvm_ffi_top_k_top_p_sampling_from_probs, + tvm::ffi::TensorView(probs), tvm::ffi::TensorView(output), + tvm::ffi::TensorView(valid), OptionalView(indices), + OptionalTensorView{tvm::ffi::TensorView(top_k)}, 0.0, + OptionalTensorView{tvm::ffi::TensorView(top_p)}, 1.0, deterministic, + OptionalTensorView{}, seed, OptionalTensorView{}, offset); +} + +int64_t ReadTopK(const Tensor tensor, Tensor::Size row) { + const auto offset = row * tensor.stride(0); + return tensor.dtype() == DataType::kInt32 + ? static_cast(tensor.data())[offset] + : static_cast(tensor.data())[offset]; +} + +double ReadTopP(const Tensor tensor, Tensor::Size row) { + const auto offset = row * tensor.stride(0); + switch (tensor.dtype()) { + case DataType::kFloat16: + return Caster::Cast( + static_cast(tensor.data())[offset]); + case DataType::kBFloat16: + return Caster::Cast( + static_cast(tensor.data())[offset]); + case DataType::kFloat32: + return static_cast(tensor.data())[offset]; + case DataType::kFloat64: + return static_cast(tensor.data())[offset]; + default: + throw std::invalid_argument( + "FlashInferSampling received an unsupported top-p dtype"); + } +} + +float NormalizeTopP(double value) { + if (!(value > 0.0 && value < 1.0)) return 1.0f; + + const auto converted = static_cast(value); + if (converted <= 0.0f) { + return std::numeric_limits::denorm_min(); + } + if (converted >= 1.0f) { + return std::nextafter(1.0f, 0.0f); + } + return converted; +} + +void Validate(const Tensor logits, const Tensor top_k, const Tensor top_p, + const std::optional& indices, + const std::string& filter_apply_order, bool check_nan, + Tensor out) { + const auto logits_dtype = logits.dtype(); + Require(logits_dtype == DataType::kFloat16 || + logits_dtype == DataType::kBFloat16 || + logits_dtype == DataType::kFloat32, + "FlashInferSampling supports float16, bfloat16, or float32 logits"); + Require(logits.device().type() == Device::Type::kNvidia && + out.device() == logits.device() && logits.IsContiguous() && + out.IsContiguous(), + "FlashInferSampling requires contiguous NVIDIA logits and output"); + Require((top_k.dtype() == DataType::kInt32 || + top_k.dtype() == DataType::kInt64) && + top_k.device().type() == Device::Type::kCpu, + "FlashInferSampling requires host int32 or int64 top-k"); + Require((top_p.dtype() == DataType::kFloat16 || + top_p.dtype() == DataType::kBFloat16 || + top_p.dtype() == DataType::kFloat32 || + top_p.dtype() == DataType::kFloat64) && + top_p.device().type() == Device::Type::kCpu, + "FlashInferSampling requires host floating-point top-p"); + Require(out.dtype() == DataType::kInt32 || + out.dtype() == DataType::kInt64, + "FlashInferSampling requires int32 or int64 output"); + Require(filter_apply_order == "top_k_first" || + filter_apply_order == "joint", + "FlashInferSampling requires top_k_first or joint filter order"); + Require(!check_nan, "FlashInferSampling does not support check_nan"); + if (indices) { + Require((indices->device() == logits.device() || + indices->device().type() == Device::Type::kCpu) && + indices->IsContiguous() && indices->dtype() == out.dtype(), + "FlashInferSampling requires contiguous CPU or NVIDIA indices " + "matching output"); + } else { + Require(logits.size(0) == out.size(0), + "FlashInferSampling requires output batch size to match logits " + "when indices are absent"); + Require(out.dtype() == DataType::kInt32, + "FlashInferSampling requires int32 output when indices are " + "absent"); + } +} + +} // 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), + workspace_size_{ + WorkspaceLayout(static_cast(out.size(0)) * + static_cast(logits.size(1)), + static_cast(out.size(0))) + .size}, + logits_batch_size_{logits.size(0)}, + device_index_{logits.device().index()}, + top_k_dtype_{top_k.dtype()}, + top_p_dtype_{top_p.dtype()}, + out_dtype_{out.dtype()}, + indices_dtype_{indices ? std::optional{indices->dtype()} : std::nullopt}, + indices_device_{indices ? std::optional{indices->device()} + : std::nullopt}, + filter_apply_order_{filter_apply_order}, + deterministic_{deterministic} { + Validate(logits, top_k, top_p, indices, filter_apply_order, check_nan, out); + Require(vocab_size_ > 0 && + vocab_size_ <= static_cast( + std::numeric_limits::max()), + "FlashInferSampling requires a nonempty int32-sized vocabulary"); + if (batch_size_ == 0) return; + + DeviceGuard guard{device_index_}; + try { + for (auto& slot : staging_slots_) { + CheckCuda( + cudaMallocHost( + &slot.top_p, + static_cast(batch_size_) * sizeof(float)), + "FlashInferSampling failed to allocate top-p staging"); + CheckCuda( + cudaMallocHost( + &slot.top_k, + static_cast(batch_size_) * sizeof(int64_t)), + "FlashInferSampling failed to allocate top-k staging"); + CheckCuda( + cudaMallocHost( + &slot.indices, + static_cast(batch_size_) * sizeof(int64_t)), + "FlashInferSampling failed to allocate indices staging"); + + cudaEvent_t event{nullptr}; + CheckCuda(cudaEventCreateWithFlags(&event, cudaEventDisableTiming), + "FlashInferSampling failed to create staging event"); + slot.event = event; + } + } catch (...) { + ReleaseStagingSlots(); + throw; + } +} + +void Operator::ReleaseStagingSlots() noexcept { + for (auto& slot : staging_slots_) { + if (slot.event_recorded && slot.event != nullptr) { + const auto status = + cudaEventSynchronize(static_cast(slot.event)); + if (status != cudaSuccess) { + ReportCuda(status, "FlashInferSampling failed to await staging"); + ReportCuda(cudaDeviceSynchronize(), + "FlashInferSampling failed to await device work"); + } + } + if (slot.event != nullptr) { + ReportCuda(cudaEventDestroy(static_cast(slot.event)), + "FlashInferSampling failed to destroy staging event"); + } + if (slot.indices != nullptr) { + ReportCuda(cudaFreeHost(slot.indices), + "FlashInferSampling failed to free indices staging"); + } + if (slot.top_k != nullptr) { + ReportCuda(cudaFreeHost(slot.top_k), + "FlashInferSampling failed to free top-k staging"); + } + if (slot.top_p != nullptr) { + ReportCuda(cudaFreeHost(slot.top_p), + "FlashInferSampling failed to free top-p staging"); + } + slot = {}; + } +} + +Operator::~Operator() { + if (batch_size_ == 0) return; + + int previous_device{0}; + auto status = cudaGetDevice(&previous_device); + if (status != cudaSuccess) { + ReportCuda(status, + "FlashInferSampling failed to query the current CUDA device"); + return; + } + + const bool restore_device = previous_device != device_index_; + if (restore_device) { + status = cudaSetDevice(device_index_); + if (status != cudaSuccess) { + ReportCuda(status, + "FlashInferSampling failed to select the input CUDA device"); + return; + } + } + + ReleaseStagingSlots(); + if (restore_device) { + ReportCuda(cudaSetDevice(previous_device), + "FlashInferSampling failed to restore the CUDA device"); + } +} + +std::size_t Operator::workspace_size_in_bytes() const { + return workspace_size_; +} + +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 { + Require(logits.ndim() == 2 && logits.size(0) == logits_batch_size_ && + logits.size(1) == vocab_size_ && logits.dtype() == dtype_ && + logits.device().type() == Device::Type::kNvidia && + logits.device().index() == device_index_ && top_k.ndim() == 1 && + top_k.size(0) == batch_size_ && + top_k.dtype() == top_k_dtype_ && + top_k.device().type() == Device::Type::kCpu && + top_p.ndim() == 1 && top_p.size(0) == batch_size_ && + top_p.dtype() == top_p_dtype_ && + top_p.device().type() == Device::Type::kCpu && out.ndim() == 1 && + out.size(0) == batch_size_ && out.dtype() == out_dtype_ && + out.device() == logits.device() && + indices.has_value() == indices_dtype_.has_value() && + filter_apply_order == filter_apply_order_ && + deterministic == deterministic_, + "FlashInferSampling call metadata changed after descriptor creation"); + if (indices) { + Require(indices->ndim() == 1 && indices->size(0) == batch_size_ && + indices->dtype() == *indices_dtype_ && + indices->device() == *indices_device_, + "FlashInferSampling indices metadata changed after descriptor " + "creation"); + } + Require(!offset || *offset >= 0, + "FlashInferSampling requires a nonnegative offset"); + Validate(logits, top_k, top_p, indices, filter_apply_order, check_nan, out); + if (batch_size_ == 0) return; + + DeviceGuard device_guard{device_index_}; + const auto stream = static_cast(stream_); + StreamGuard stream_guard{device_index_, stream}; + std::lock_guard lock{mutex_}; + + auto slot_index = next_staging_slot_; + auto* slot = &staging_slots_[slot_index]; + auto status = cudaSuccess; + if (slot->event_recorded) { + status = cudaEventQuery(static_cast(slot->event)); + if (status == cudaErrorNotReady) { + const auto other_index = (slot_index + 1) % staging_slots_.size(); + auto* other = &staging_slots_[other_index]; + auto other_status = + other->event_recorded + ? cudaEventQuery(static_cast(other->event)) + : cudaSuccess; + if (other_status == cudaSuccess) { + slot_index = other_index; + slot = other; + } else if (other_status == cudaErrorNotReady) { + CheckCuda(cudaEventSynchronize(static_cast(slot->event)), + "FlashInferSampling failed to await staging slot"); + } else { + CheckCuda(other_status, + "FlashInferSampling failed to query staging event"); + } + } else if (status != cudaSuccess) { + CheckCuda(status, "FlashInferSampling failed to query staging event"); + } + } + next_staging_slot_ = (slot_index + 1) % staging_slots_.size(); + + const auto matrix_elements = static_cast(batch_size_) * + static_cast(logits.size(1)); + const WorkspaceLayout layout{matrix_elements, + static_cast(batch_size_)}; + Workspace workspace_owner{workspace_, workspace_size_in_bytes_, layout.size, + stream}; + auto* workspace = static_cast(workspace_owner.data()); + auto* matrix_a = reinterpret_cast(workspace + layout.matrix_a); + auto* matrix_b = reinterpret_cast(workspace + layout.matrix_b); + auto* top_k_device = workspace + layout.top_k; + auto* top_p_device = reinterpret_cast(workspace + layout.top_p); + auto* valid_device = workspace + layout.valid; + auto* scratch_device = workspace + layout.scratch; + auto* indices_device = workspace + layout.indices; + + for (Tensor::Size row = 0; row < batch_size_; ++row) { + slot->top_p[static_cast(row)] = + NormalizeTopP(ReadTopP(top_p, row)); + } + + const bool top_k_is_int64 = + filter_apply_order == "joint" && out.dtype() == DataType::kInt64; + std::size_t top_k_bytes{0}; + if (top_k_is_int64) { + for (Tensor::Size row = 0; row < batch_size_; ++row) { + const auto value = ReadTopK(top_k, row); + slot->top_k[static_cast(row)] = + value > 0 && value <= static_cast(vocab_size_) + ? value + : static_cast(vocab_size_); + } + top_k_bytes = + static_cast(batch_size_) * sizeof(int64_t); + } else { + auto* top_k_int32 = reinterpret_cast(slot->top_k); + for (Tensor::Size row = 0; row < batch_size_; ++row) { + const auto value = ReadTopK(top_k, row); + top_k_int32[static_cast(row)] = + value > 0 && value <= static_cast(vocab_size_) + ? static_cast(value) + : static_cast(vocab_size_); + } + top_k_bytes = + static_cast(batch_size_) * sizeof(int32_t); + } + + const void* staged_indices_data = indices ? indices->data() : nullptr; + std::size_t indices_bytes{0}; + if (indices && indices->device().type() == Device::Type::kCpu) { + if (indices->dtype() == DataType::kInt32) { + for (Tensor::Size row = 0; row < batch_size_; ++row) { + const auto value = static_cast(indices->data())[row]; + Require(value >= 0 && value < logits_batch_size_, + "FlashInferSampling received an out-of-range host index"); + } + } else { + for (Tensor::Size row = 0; row < batch_size_; ++row) { + const auto value = static_cast(indices->data())[row]; + Require(value >= 0 && value < logits_batch_size_, + "FlashInferSampling received an out-of-range host index"); + } + } + indices_bytes = static_cast(batch_size_) * + kDataTypeToSize.at(indices->dtype()); + std::memcpy(slot->indices, indices->data(), indices_bytes); + staged_indices_data = indices_device; + } + + StagingRecorder staging_recorder{static_cast(slot->event), + stream, &slot->event_recorded}; + CheckCuda( + cudaMemcpyAsync(top_p_device, slot->top_p, + static_cast(batch_size_) * sizeof(float), + cudaMemcpyHostToDevice, stream), + "FlashInferSampling failed to stage top-p values"); + CheckCuda(cudaMemcpyAsync(top_k_device, slot->top_k, top_k_bytes, + cudaMemcpyHostToDevice, stream), + "FlashInferSampling failed to stage top-k values"); + if (indices_bytes != 0) { + CheckCuda(cudaMemcpyAsync(indices_device, slot->indices, indices_bytes, + cudaMemcpyHostToDevice, stream), + "FlashInferSampling failed to stage indices"); + } + staging_recorder.Record(); + + DispatchFunc( + logits.dtype(), + [&](auto tag) { + using T = typename decltype(tag)::type; + if (!indices) { + CastLogits<<>>( + matrix_a, static_cast(logits.data()), matrix_elements); + } else if (indices->dtype() == DataType::kInt32) { + GatherCastLogits<<>>( + matrix_a, static_cast(logits.data()), + static_cast(staged_indices_data), + static_cast(batch_size_), + static_cast(logits_batch_size_), + static_cast(vocab_size_)); + } else { + GatherCastLogits<<>>( + matrix_a, static_cast(logits.data()), + static_cast(staged_indices_data), + static_cast(batch_size_), + static_cast(logits_batch_size_), + static_cast(vocab_size_)); + } + }, + "`FlashInferSampling` logits cast"); + CheckCuda(cudaGetLastError(), + "FlashInferSampling logits preparation kernel launch failed"); + + int64_t matrix_shape[2]{static_cast(batch_size_), + static_cast(vocab_size_)}; + int64_t batch_shape[1]{static_cast(batch_size_)}; + int64_t scratch_shape[1]{static_cast(kScratchBytes)}; + auto matrix_a_tensor = MakeTensor(matrix_a, device_index_, 2, matrix_shape, + Dtype(DataType::kFloat32)); + auto matrix_b_tensor = MakeTensor(matrix_b, device_index_, 2, matrix_shape, + Dtype(DataType::kFloat32)); + auto top_k_tensor = + MakeTensor(top_k_device, device_index_, 1, batch_shape, + Dtype(top_k_is_int64 ? DataType::kInt64 : DataType::kInt32)); + auto top_p_tensor = MakeTensor(top_p_device, device_index_, 1, batch_shape, + Dtype(DataType::kFloat32)); + auto valid_tensor = + MakeTensor(valid_device, device_index_, 1, batch_shape, {kDLBool, 8, 1}); + auto scratch_tensor = MakeTensor(scratch_device, device_index_, 1, + scratch_shape, {kDLUInt, 8, 1}); + auto output_tensor = + MakeTensor(out.data(), device_index_, 1, batch_shape, Dtype(out.dtype())); + const auto actual_seed = static_cast( + seed.value_or(static_cast(std::random_device{}()))); + const auto actual_offset = static_cast(offset.value_or(0)); + if (filter_apply_order == "top_k_first") { + CheckCuda(cudaMemsetAsync(scratch_device, 0, kScratchBytes, stream), + "FlashInferSampling failed to initialize row-state workspace"); + CallTopKMask(&matrix_a_tensor, &matrix_b_tensor, &top_k_tensor, + &scratch_tensor); + CallSoftmax(&scratch_tensor, &matrix_b_tensor, &matrix_a_tensor); + CallTopP(&matrix_a_tensor, &output_tensor, &valid_tensor, nullptr, + &top_p_tensor, deterministic, actual_seed, actual_offset); + } else { + CallSoftmax(&scratch_tensor, &matrix_a_tensor, &matrix_b_tensor); + CallJoint(&matrix_b_tensor, &output_tensor, &valid_tensor, nullptr, + &top_k_tensor, &top_p_tensor, deterministic, actual_seed, + actual_offset); + } + + CheckCuda(cudaGetLastError(), + "FlashInferSampling CUDA kernel launch failed"); +} + +} // namespace infini::ops diff --git a/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.h b/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.h new file mode 100644 index 000000000..72de5595d --- /dev/null +++ b/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.h @@ -0,0 +1,77 @@ +#ifndef INFINI_OPS_LINKED_TVM_FFI_NVIDIA_OPS_TOP_K_TOP_P_SAMPLING_FROM_LOGITS_FLASHINFER_H_ +#define INFINI_OPS_LINKED_TVM_FFI_NVIDIA_OPS_TOP_K_TOP_P_SAMPLING_FROM_LOGITS_FLASHINFER_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: + struct StagingSlot { + float* top_p{nullptr}; + int64_t* top_k{nullptr}; + void* indices{nullptr}; + void* event{nullptr}; + bool event_recorded{false}; + }; + + std::size_t workspace_size_{0}; + + Tensor::Size logits_batch_size_{0}; + + int device_index_{0}; + + DataType top_k_dtype_; + + DataType top_p_dtype_; + + DataType out_dtype_; + + std::optional indices_dtype_; + + std::optional indices_device_; + + std::string filter_apply_order_; + + bool deterministic_{false}; + + mutable std::array staging_slots_; + + mutable std::size_t next_staging_slot_{0}; + + mutable std::mutex mutex_; + + void ReleaseStagingSlots() noexcept; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TVM_FFI_NVIDIA_OPS_TOP_K_TOP_P_SAMPLING_FROM_LOGITS_FLASHINFER_H_ diff --git a/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.yaml b/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.yaml new file mode 100644 index 000000000..e3a4f14c0 --- /dev/null +++ b/src/linked/tvm_ffi/nvidia/ops/top_k_top_p_sampling_from_logits/flashinfer.yaml @@ -0,0 +1,8 @@ +library: flashinfer_sampling +link_libraries: + - tvm_ffi +required_symbols: + - __tvm_ffi_softmax + - __tvm_ffi_top_k_mask_logits + - __tvm_ffi_top_p_sampling_from_probs + - __tvm_ffi_top_k_top_p_sampling_from_probs diff --git a/src/linked/tvm_ffi/nvidia/tvm_ffi.yaml b/src/linked/tvm_ffi/nvidia/tvm_ffi.yaml new file mode 100644 index 000000000..c386e37cf --- /dev/null +++ b/src/linked/tvm_ffi/nvidia/tvm_ffi.yaml @@ -0,0 +1,4 @@ +python_distribution_package: apache-tvm-ffi +python_distribution_version: "==0.1.10" +library_glob: tvm_ffi/lib/libtvm_ffi.so +include_glob: tvm_ffi/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..ebd71fc51 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,337 @@ def test_top_k_top_p_sampling_from_logits( assert torch.all(torch.isin(first, allowed_tensor)) +def test_flashinfer_sampling_joint_host_indices(device, implementation_index): + if implementation_index != 16: + pytest.skip("FlashInfer linked-provider coverage") + + logits = torch.tensor( + ( + (9.0, 1.0, 0.0, -1.0), + (0.0, 8.0, 1.0, -1.0), + (-1.0, 0.0, 1.0, 7.0), + ), + dtype=torch.float32, + device=device, + ) + indices = torch.tensor((2, 0, 2, 1, 0, 1, 2), dtype=torch.int64) + batch_size = indices.numel() + top_k = torch.ones(batch_size, dtype=torch.int64) + top_p = torch.ones(batch_size, dtype=torch.float32) + out = torch.empty(batch_size, dtype=torch.int64, device=device) + + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 9, + out, + implementation_index, + indices=indices, + filter_apply_order="joint", + ) + + expected = torch.tensor((3, 0, 3, 1, 0, 1, 3), dtype=torch.int64, device=device) + assert torch.equal(out, expected) + + +def test_flashinfer_sampling_top_k_first_cuda_indices(device, implementation_index): + if implementation_index != 16: + pytest.skip("FlashInfer linked-provider coverage") + + logits = torch.tensor( + ( + (9.0, 1.0, 0.0, -1.0), + (0.0, 8.0, 1.0, -1.0), + (-1.0, 0.0, 1.0, 7.0), + ), + dtype=torch.bfloat16, + device=device, + ) + indices = torch.tensor((2, 0, 2, 1, 0), dtype=torch.int32, device=device) + batch_size = indices.numel() + top_k = torch.ones(batch_size, dtype=torch.int32) + top_p = torch.ones(batch_size, dtype=torch.float64) + 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, + indices=indices, + filter_apply_order="top_k_first", + ) + + expected = torch.tensor((3, 0, 3, 1, 0), dtype=torch.int32, device=device) + assert torch.equal(out, expected) + + +def test_flashinfer_sampling_offset(device, implementation_index): + if implementation_index != 16: + pytest.skip("FlashInfer linked-provider coverage") + + batch_size = 256 + logits = torch.zeros((batch_size, 4), dtype=torch.float32, device=device) + top_k = torch.full((batch_size,), 4, dtype=torch.int32) + top_p = torch.ones(batch_size, dtype=torch.float32) + first = torch.empty(batch_size, dtype=torch.int32, device=device) + repeated = torch.empty_like(first) + different_offset = torch.empty_like(first) + + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 9, + first, + implementation_index, + filter_apply_order="joint", + ) + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 9, + repeated, + implementation_index, + filter_apply_order="joint", + ) + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 10, + different_offset, + implementation_index, + filter_apply_order="joint", + ) + + assert torch.equal(first, repeated) + assert not torch.equal(first, different_offset) + + +def test_flashinfer_sampling_uses_handle_stream(device, implementation_index): + if device != "cuda" or implementation_index != 16: + pytest.skip("FlashInfer linked-provider stream coverage") + + batch_size = 64 + logits = torch.zeros((batch_size, 4), dtype=torch.float32, device=device) + logits[:, 0] = 1.0 + top_k = torch.ones(batch_size, dtype=torch.int32) + top_p = torch.ones(batch_size, dtype=torch.float32) + out = torch.full((batch_size,), -1, dtype=torch.int32, device=device) + stream = torch.cuda.Stream() + + def call_sampling(): + infini.ops.top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + None, + "joint", + True, + False, + 1234, + 9, + out, + stream=stream.cuda_stream, + implementation_index=implementation_index, + ) + + try: + call_sampling() + stream.synchronize() + out.fill_(-1) + torch.cuda.synchronize() + + with torch.cuda.stream(stream): + torch.cuda._sleep(50_000_000) + call_sampling() + + default_stream = torch.cuda.default_stream() + with torch.cuda.stream(default_stream): + snapshot = out.clone() + default_stream.synchronize() + assert torch.all(snapshot == -1) + + stream.synchronize() + assert torch.all(out == 0) + finally: + torch.cuda.synchronize() + + +def test_flashinfer_sampling_preserves_float64_top_p_underflow( + device, implementation_index +): + if implementation_index != 16: + pytest.skip("FlashInfer linked-provider coverage") + + batch_size = 4096 + logits = torch.tensor((0.0, -1.0), dtype=torch.float32, device=device).repeat( + batch_size, 1 + ) + top_k = torch.full((batch_size,), 2, dtype=torch.int32) + top_p = torch.full((batch_size,), 1e-300, dtype=torch.float64) + 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, + ) + + assert torch.all(out == 0) + + +def test_flashinfer_sampling_rejects_float64_logits(device, implementation_index): + if implementation_index != 16: + pytest.skip("FlashInfer linked-provider coverage") + + logits = torch.tensor( + ((1.0, 1.0000000001),), + dtype=torch.float64, + device=device, + ) + top_k = torch.ones(1, dtype=torch.int32) + top_p = torch.ones(1, dtype=torch.float32) + out = torch.empty(1, dtype=torch.int32, device=device) + + with pytest.raises(ValueError, match="float16, bfloat16, or float32 logits"): + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 9, + out, + implementation_index, + ) + + +def test_flashinfer_sampling_rejects_noncontiguous_logits(device, implementation_index): + if implementation_index != 16: + pytest.skip("FlashInfer linked-provider coverage") + + logits = torch.tensor( + ((2.0, 0.0), (0.0, 2.0)), + dtype=torch.float32, + device=device, + ).T + assert not logits.is_contiguous() + top_k = torch.ones(2, dtype=torch.int32) + top_p = torch.ones(2, dtype=torch.float32) + out = torch.empty(2, dtype=torch.int32, device=device) + + with pytest.raises(ValueError, match="contiguous NVIDIA logits"): + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 9, + out, + implementation_index, + ) + + +def test_flashinfer_sampling_rejects_check_nan(device, implementation_index): + if implementation_index != 16: + pytest.skip("FlashInfer linked-provider coverage") + + logits = torch.tensor( + ((float("nan"), 0.0),), + dtype=torch.float32, + device=device, + ) + top_k = torch.ones(1, dtype=torch.int32) + top_p = torch.ones(1, dtype=torch.float32) + out = torch.empty(1, dtype=torch.int32, device=device) + + with pytest.raises(ValueError, match="does not support check_nan"): + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 9, + out, + implementation_index, + check_nan=True, + ) + + +def test_flashinfer_sampling_rejects_out_of_range_host_indices( + device, implementation_index +): + if implementation_index != 16: + pytest.skip("FlashInfer linked-provider coverage") + + logits = torch.tensor( + ((2.0, 0.0), (0.0, 2.0)), + dtype=torch.float32, + device=device, + ) + indices = torch.tensor((0, 2), dtype=torch.int32) + top_k = torch.ones(2, dtype=torch.int32) + top_p = torch.ones(2, dtype=torch.float32) + out = torch.empty(2, dtype=torch.int32, device=device) + + with pytest.raises(ValueError, match="out-of-range host index"): + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 9, + out, + implementation_index, + indices=indices, + ) + + +@pytest.mark.parametrize( + "logits_batch,out_dtype,error", + ( + (1, torch.int32, "output batch size to match logits"), + (2, torch.int64, "int32 output when indices are absent"), + ), +) +def test_flashinfer_sampling_rejects_invalid_output_without_indices( + logits_batch, out_dtype, error, device, implementation_index +): + if implementation_index != 16: + pytest.skip("FlashInfer linked-provider coverage") + + batch_size = 2 + logits = torch.zeros((logits_batch, 4), dtype=torch.float32, device=device) + top_k = torch.ones(batch_size, dtype=torch.int32) + top_p = torch.ones(batch_size, dtype=torch.float32) + out = torch.empty(batch_size, dtype=out_dtype, device=device) + + with pytest.raises(ValueError, match=error): + _top_k_top_p_sampling_from_logits( + logits, + top_k, + top_p, + 1234, + 9, + out, + implementation_index, + ) + + def _top_k_top_p_sampling_from_logits( logits, top_k, @@ -58,15 +389,19 @@ def _top_k_top_p_sampling_from_logits( offset, out, implementation_index, + *, + indices=None, + filter_apply_order="top_k_first", + check_nan=False, ): infini.ops.top_k_top_p_sampling_from_logits( logits, top_k, top_p, - None, - "top_k_first", + indices, + filter_apply_order, True, - False, + check_nan, seed, offset, out,