diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 6f6ea4d94b958..1e703c704ea4f 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -3480,11 +3480,11 @@ This version of the operator has been available since version 1 of the 'com.micr The weight tensor B has shape [N, K] with one FP32 scale per `block_size` consecutive K values (`b_scale` of shape [N, ceil(K / block_size)]). The scaled weight value is `B_scaled[n, k] = fp8_e4m3(B[n, k]) * b_scale[n, k / block_size]`. - + When the optional scalar `a_scale` is provided, the activation values used in the multiplication are `A_scaled = fp8_e4m3(A / a_scale) * a_scale` (W8A8). Otherwise, A retains its FP16/BF16 precision (weight-only W8A16). - + The operator multiplies the activation by the transpose of B_scaled and adds the optional bias. The output has shape [..., N] and the same element type as A. @@ -7742,4 +7742,3 @@ No versioning maintained for experimental ops.
T : tensor(float)
Constrain input and output types to float32 tensors.
- diff --git a/docs/cuda_host_pageable_gather.md b/docs/cuda_host_pageable_gather.md new file mode 100644 index 0000000000000..cc5e9b3260ed0 --- /dev/null +++ b/docs/cuda_host_pageable_gather.md @@ -0,0 +1,27 @@ +# CUDA host-pageable `GatherBlockQuantized` + +The CUDA Execution Provider option `enable_host_pageable_gather` enables direct access to CPU-resident FP8 +`com.microsoft::GatherBlockQuantized` input data. It is disabled by default. + +Direct access requires a CUDA device that reports both `cudaDevAttrPageableMemoryAccess` and +`cudaDevAttrPageableMemoryAccessUsesHostPageTables`. If either device capability is unavailable, ONNX Runtime emits a +warning and uses the standard CUDA input path. + +The option does not create or manage a file mapping. The model initializer must already be supplied as CPU memory, +such as a file-backed mapping, and that mapping remains live for the session. The direct path does not register, +prefetch, hash, scan, or copy the complete initializer, and the initializer is not accounted as CUDA-resident memory. +Fallback copies are allocated by the CUDA initializer allocator; they are not currently included in capacity-aware +partitioning estimates. + +CUDA Graph capture is supported for direct host-pageable access. The initializer mapping must remain alive at the same +virtual address until the graph executable is destroyed, and indices, scales, and outputs must retain their normal +CUDA Graph-stable addresses. The direct path performs no mapping, registration, allocation, copy, capability query, or +synchronization during capture. Persistent-copy fallbacks are prepared during prepacking or an uncaptured warmup run; +if lazy fallback initialization is still required when capture starts, the run fails instead of allocating or copying +during capture. + +Non-constant and non-FP8 input data retains the standard GPU-input contract. Multiple nodes that use the same constant +initializer can create separate fallback copies; direct host access does not duplicate the initializer. + +This mode primarily reduces GPU memory capacity requirements. Performance depends on storage latency and operating +system page-cache state, so cold prefill can be slower and less predictable than using resident GPU memory. diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc index a67514ac560b2..6240f542432df 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc @@ -71,7 +71,8 @@ REGISTER_GATHERBLOCKQUANTIZED(Float4E2M1x2, BFloat16, int64_t); #endif // !defined(DISABLE_FLOAT4_TYPES) template -GatherBlockQuantized::GatherBlockQuantized(const OpKernelInfo& info) : CudaKernel(info) { +GatherBlockQuantized::GatherBlockQuantized(const OpKernelInfo& info) + : CudaKernel(info), direct_host_data_(false), data_is_constant_(false) { if constexpr (IsFpQuantizedV) { bits_ = 0; // Not applicable for FP8/FP4 data. } else { @@ -82,6 +83,36 @@ GatherBlockQuantized::GatherBlockQuantized(const OpKernelInfo& inf gather_axis_ = info.GetAttrOrDefault("gather_axis", 0); quantize_axis_ = info.GetAttrOrDefault("quantize_axis", 1); + const Tensor* constant_data = nullptr; + data_is_constant_ = info.TryGetConstantInput(0, &constant_data); + + int pageable_memory_access = 0; + int uses_host_page_tables = 0; +#if defined(CUDA_VERSION) && CUDA_VERSION >= 10020 + const bool attributes_available = + cudaDeviceGetAttribute(&pageable_memory_access, cudaDevAttrPageableMemoryAccess, GetDeviceId()) == cudaSuccess && + cudaDeviceGetAttribute(&uses_host_page_tables, cudaDevAttrPageableMemoryAccessUsesHostPageTables, + GetDeviceId()) == cudaSuccess; + if (!attributes_available) { + pageable_memory_access = 0; + uses_host_page_tables = 0; + cudaGetLastError(); + } +#endif + + const bool option_enabled = EnableHostPageableGather(); + direct_host_data_ = + SelectGatherBlockQuantizedDataPolicy(option_enabled, pageable_memory_access != 0, + uses_host_page_tables != 0, IsFp8QuantizedV, + data_is_constant_) == + GatherBlockQuantizedDataPolicy::DirectHost; + if (option_enabled && IsFp8QuantizedV && !direct_host_data_) { + LOGS_DEFAULT(WARNING) + << "enable_host_pageable_gather was requested, but direct host-pageable GatherBlockQuantized " + "access is unavailable because input 0 is not a constant initializer or the CUDA device lacks pageable " + "memory access through host page tables. Using the standard CUDA input path."; + } + // If block size is set, it has to be no smaller than 16 and must be power of 2. // block_size_ & (block_size_ - 1) == 0 checks if block_size_ only has 1 bit set. // block_size_ == 0 is only valid for FP8/FP4 data, meaning the whole quantize_axis dimension @@ -93,6 +124,55 @@ GatherBlockQuantized::GatherBlockQuantized(const OpKernelInfo& inf } } +template +Status GatherBlockQuantized::CreateDeviceCopy(const Tensor& tensor, AllocatorPtr alloc) const { + ORT_RETURN_IF_NOT(tensor.Location().device.Type() == OrtDevice::CPU, + "GatherBlockQuantized input 0 must reside in CPU memory."); + + const size_t bytes = tensor.SizeInBytes(); + if (bytes == 0) { + data_shape_.assign(tensor.Shape().GetDims().begin(), tensor.Shape().GetDims().end()); + device_data_.reset(); + return Status::OK(); + } + + auto device_data = IAllocator::MakeUniquePtr(alloc, bytes); + ORT_RETURN_IF_NOT(device_data != nullptr, "Failed to allocate persistent CUDA storage for GatherBlockQuantized."); + CUDA_RETURN_IF_ERROR(cudaMemcpy(device_data.get(), tensor.DataRaw(), bytes, cudaMemcpyHostToDevice)); + data_shape_.assign(tensor.Shape().GetDims().begin(), tensor.Shape().GetDims().end()); + device_data_ = std::move(device_data); + return Status::OK(); +} + +template +Status GatherBlockQuantized::PrePack( + const Tensor& tensor, int input_idx, AllocatorPtr alloc, + bool& is_packed, PrePackedWeights* prepacked_weights) { + is_packed = false; + if (input_idx != 0) { + return Status::OK(); + } + + if (!direct_host_data_ && tensor.Location().device.Type() != OrtDevice::CPU) { + return Status::OK(); + } + + std::lock_guard lock(device_data_mutex_); + if (direct_host_data_) { + ORT_RETURN_IF_NOT(tensor.Location().device.Type() == OrtDevice::CPU, + "Direct host-pageable GatherBlockQuantized requires a CPU-resident initializer."); + direct_host_data_ptr_ = tensor.Data(); + data_shape_.assign(tensor.Shape().GetDims().begin(), tensor.Shape().GetDims().end()); + } else { + ORT_RETURN_IF_ERROR(CreateDeviceCopy(tensor, std::move(alloc))); + } + is_packed = true; + if (prepacked_weights != nullptr) { + prepacked_weights->has_kernel_owned_packed_weights_ = true; + } + return Status::OK(); +} + template Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) const { const Tensor* data = ctx->Input(0); @@ -100,8 +180,9 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) const Tensor* scales = ctx->Input(2); const Tensor* zero_points = ctx->Input(3); - auto data_shape = data->Shape().GetDims(); - int64_t data_rank = data->Shape().NumDimensions(); + const gsl::span data_shape = + data != nullptr ? data->Shape().GetDims() : gsl::span{data_shape_}; + int64_t data_rank = static_cast(data_shape.size()); const int64_t gather_axis = HandleNegativeAxis(gather_axis_, data_rank); const int64_t quantize_axis = HandleNegativeAxis(quantize_axis_, data_rank); @@ -155,7 +236,33 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) return Status::OK(); } - const auto* data_ptr = data->Data(); + const T1* data_ptr = nullptr; + if (direct_host_data_) { + data_ptr = direct_host_data_ptr_ != nullptr ? direct_host_data_ptr_ + : data == nullptr ? nullptr + : data->Data(); + } else if (data_is_constant_) { + { + std::lock_guard lock(device_data_mutex_); + if (device_data_ == nullptr && data != nullptr && + data->Location().device.Type() == OrtDevice::CPU && data->SizeInBytes() != 0) { + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + CUDA_RETURN_IF_ERROR(cudaStreamIsCapturing(Stream(ctx), &capture_status)); + ORT_RETURN_IF_NOT( + capture_status == cudaStreamCaptureStatusNone, + "GatherBlockQuantized cannot initialize its persistent CUDA fallback copy during CUDA Graph capture. " + "Enable prepacking or run an uncaptured warmup iteration before capture."); + ORT_RETURN_IF_ERROR(CreateDeviceCopy(*data, Info().GetAllocator(OrtMemTypeDefault))); + } + data_ptr = device_data_ != nullptr ? static_cast(device_data_.get()) + : data == nullptr ? nullptr + : data->Data(); + } + } else { + data_ptr = data->Data(); + } + ORT_RETURN_IF_NOT(N == 0 || data_ptr != nullptr, + "GatherBlockQuantized fallback has no device-resident input 0."); const auto* indices_ptr = indices->Data(); const T1* zero_points_ptr = nullptr; if (zero_points != nullptr) { diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cuh b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cuh index b240455622474..97ae4688eca11 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cuh @@ -43,6 +43,15 @@ struct IsFpQuantized : std::true_type {}; template inline constexpr bool IsFpQuantizedV = IsFpQuantized::value; +template +inline constexpr bool IsFp8QuantizedV = +#if !defined(DISABLE_FLOAT8_TYPES) + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v; +#else + false; +#endif + struct GatherBlockQuantizedParam { cudaStream_t stream; int64_t after_gather_dim; diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.h b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.h index 7718b6dd06765..b7a5a84c9c1f1 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.h +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.h @@ -3,9 +3,10 @@ #pragma once +#include "contrib_ops/cuda/quantization/gather_block_quantized_data_policy.h" #include "core/providers/cuda/cuda_kernel.h" -#include +#include using namespace onnxruntime::cuda; @@ -20,12 +21,22 @@ class GatherBlockQuantized final : public CudaKernel { public: GatherBlockQuantized(const OpKernelInfo& info); Status ComputeInternal(OpKernelContext* context) const override; + Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + bool& is_packed, PrePackedWeights* prepacked_weights) override; private: + Status CreateDeviceCopy(const Tensor& tensor, AllocatorPtr alloc) const; + int64_t bits_; int64_t block_size_; int64_t gather_axis_; int64_t quantize_axis_; + bool direct_host_data_; + bool data_is_constant_; + mutable std::mutex device_data_mutex_; + const T1* direct_host_data_ptr_{}; + mutable IAllocatorUniquePtr device_data_; + mutable TensorShapeVector data_shape_; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized_data_policy.h b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized_data_policy.h new file mode 100644 index 0000000000000..a48f5e2ed3576 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized_data_policy.h @@ -0,0 +1,22 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +namespace onnxruntime::contrib::cuda { + +enum class GatherBlockQuantizedDataPolicy { + DeviceCopy, + DirectHost, +}; + +constexpr GatherBlockQuantizedDataPolicy SelectGatherBlockQuantizedDataPolicy( + bool option_enabled, bool pageable_memory_access, bool uses_host_page_tables, + bool is_fp8, bool is_constant_initializer) { + return option_enabled && pageable_memory_access && uses_host_page_tables && + is_fp8 && is_constant_initializer + ? GatherBlockQuantizedDataPolicy::DirectHost + : GatherBlockQuantizedDataPolicy::DeviceCopy; +} + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.h b/onnxruntime/core/providers/cuda/cuda_execution_provider.h index 4d1f1125786fd..ce4fed9ca17f6 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.h +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.h @@ -90,6 +90,7 @@ class CUDAExecutionProvider : public IExecutionProvider { bool IsNHWCPreferred() const { return info_.prefer_nhwc; } bool IsFuseConvBias() const { return info_.fuse_conv_bias; } bool UseTF32() const { return info_.use_tf32; } + bool EnableHostPageableGather() const { return info_.enable_host_pageable_gather; } // Attention kernel options parsed from sdpa_kernel cuda provider option. const AttentionKernelOptions* GetAttentionKernelOptions() const { diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc index 14899aee419ec..6b311c0b1c50b 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc @@ -38,6 +38,7 @@ constexpr const char* kUseEPLevelUnifiedStream = "use_ep_level_unified_stream"; constexpr const char* kUseTF32 = "use_tf32"; constexpr const char* kFuseConvBias = "fuse_conv_bias"; constexpr const char* kSdpaKernel = "sdpa_kernel"; +constexpr const char* kEnableHostPageableGather = "enable_host_pageable_gather"; constexpr const char* kExternalDataLoaderReadingThreads = "external_data_loader_reading_threads"; } // namespace provider_option_names @@ -132,6 +133,8 @@ CUDAExecutionProviderInfo CUDAExecutionProviderInfo::FromProviderOptions(const P .AddAssignmentToReference(cuda::provider_option_names::kUseEPLevelUnifiedStream, info.use_ep_level_unified_stream) .AddAssignmentToReference(cuda::provider_option_names::kUseTF32, info.use_tf32) .AddAssignmentToReference(cuda::provider_option_names::kSdpaKernel, info.sdpa_kernel) + .AddAssignmentToReference(cuda::provider_option_names::kEnableHostPageableGather, + info.enable_host_pageable_gather) .AddAssignmentToReference(cuda::provider_option_names::kFuseConvBias, info.fuse_conv_bias) .AddValueParser( cuda::provider_option_names::kExternalDataLoaderReadingThreads, @@ -201,6 +204,8 @@ ProviderOptions CUDAExecutionProviderInfo::ToProviderOptions(const CUDAExecution {cuda::provider_option_names::kUseTF32, MakeStringWithClassicLocale(info.use_tf32)}, {cuda::provider_option_names::kSdpaKernel, MakeStringWithClassicLocale(info.sdpa_kernel)}, {cuda::provider_option_names::kFuseConvBias, MakeStringWithClassicLocale(info.fuse_conv_bias)}, + {cuda::provider_option_names::kEnableHostPageableGather, + MakeStringWithClassicLocale(info.enable_host_pageable_gather)}, {cuda::provider_option_names::kExternalDataLoaderReadingThreads, MakeStringWithClassicLocale(info.external_data_loader_reading_threads)}, }; diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider_info.h b/onnxruntime/core/providers/cuda/cuda_execution_provider_info.h index b4d4cb758d2f9..c3b5a770440d3 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider_info.h +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider_info.h @@ -83,6 +83,7 @@ struct CUDAExecutionProviderInfo { bool fuse_conv_bias{false}; int sdpa_kernel{0}; + bool enable_host_pageable_gather{false}; // 0 disables the custom external-data loader and retains the framework's existing path. // 1 uses the pinned-buffer loader with synchronous reads. 2..64 use that many parallel read tasks per block. @@ -121,6 +122,7 @@ struct std::hash<::onnxruntime::CUDAExecutionProviderInfo> { onnxruntime::HashCombine(info.tunable_op.max_tuning_duration_ms, value); onnxruntime::HashCombine(info.sdpa_kernel, value); onnxruntime::HashCombine(info.enable_cudnn, value); + onnxruntime::HashCombine(info.enable_host_pageable_gather, value); onnxruntime::HashCombine(info.external_data_loader_reading_threads, value); // Memory pointers diff --git a/onnxruntime/core/providers/cuda/cuda_kernel.h b/onnxruntime/core/providers/cuda/cuda_kernel.h index b9d48e4925334..b873cca995662 100644 --- a/onnxruntime/core/providers/cuda/cuda_kernel.h +++ b/onnxruntime/core/providers/cuda/cuda_kernel.h @@ -125,6 +125,8 @@ class CudaKernel : public OpKernel { bool GetCudnnConvUseMaxWorkspace() const { return provider_->GetCudnnConvUseMaxWorkspace(); } bool GetCudnnConv1dPadToNc1d() const { return provider_->GetCudnnConv1dPadToNc1d(); } bool IsFuseConvBias() const { return provider_->IsFuseConvBias(); } + bool EnableHostPageableGather() const { return provider_->EnableHostPageableGather(); } + bool IsCudaGraphEnabled() const { return provider_->IsGraphCaptureEnabled(); } // Compatibility helper used by kernels that need the underlying ORT stream object. inline onnxruntime::Stream* GetComputeStream(OpKernelContext* ctx) const { diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index 00519584950b8..92e9646581917 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -251,7 +251,6 @@ struct CUDA_Provider : Provider { info.use_tf32 = params->use_tf32 != 0; info.sdpa_kernel = params->sdpa_kernel; info.external_data_loader_reading_threads = params->external_data_loader_reading_threads; - return std::make_shared(info); } diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc b/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc index 4c946c872ba11..28a8dc1b80256 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep.cc @@ -212,6 +212,8 @@ CudaEp::CudaEp(CudaEpFactory& factory, const Config& config, const OrtLogger& lo adapter_config.cudnn_conv1d_pad_to_nc1d = config_.cudnn_conv1d_pad_to_nc1d; adapter_config.enable_cudnn = config_.enable_cudnn; adapter_config.fuse_conv_bias = config_.fuse_conv_bias; + adapter_config.enable_cuda_graph = config_.enable_cuda_graph; + adapter_config.enable_host_pageable_gather = config_.enable_host_pageable_gather; adapter_config.sdpa_kernel = config_.sdpa_kernel; adapter_config.device_id = config_.device_id; adapter_config.do_copy_in_default_stream = config_.do_copy_in_default_stream; diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_ep.h b/onnxruntime/core/providers/cuda/plugin/cuda_ep.h index fdfe3ff02f535..dafe90cc733ab 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep.h @@ -34,6 +34,7 @@ class CudaEp : public onnxruntime::ep::adapter::Ep { bool fuse_conv_bias = false; ///< Enable cuDNN frontend conv+bias fusion. int sdpa_kernel = 0; ///< Attention backend bitmask override. bool enable_cuda_graph = false; ///< Enable CUDA graph capture and replay. + bool enable_host_pageable_gather = false; ///< Enable direct host-pageable FP8 gather access. int min_num_runs_before_cuda_graph_capture = 2; ///< Warm-up runs before graph capture begins. bool has_user_compute_stream = false; ///< Whether user provided an external CUDA stream. void* user_compute_stream = nullptr; ///< User-provided CUDA stream (cudaStream_t cast to void*). diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc index e675836508be2..d5fd7f80224de 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc @@ -652,6 +652,7 @@ OrtStatus* ORT_API_CALL CudaEpFactory::CreateEpImpl( const std::string fuse_conv_bias_key = ep_options_prefix + "fuse_conv_bias"; const std::string sdpa_kernel_key = ep_options_prefix + "sdpa_kernel"; const std::string enable_cuda_graph_key = ep_options_prefix + "enable_cuda_graph"; + const std::string enable_host_pageable_gather_key = ep_options_prefix + "enable_host_pageable_gather"; const std::string min_runs_key = ep_options_prefix + "min_num_runs_before_cuda_graph_capture"; const std::string has_user_compute_stream_key = ep_options_prefix + "has_user_compute_stream"; const std::string user_compute_stream_key = ep_options_prefix + "user_compute_stream"; @@ -688,6 +689,9 @@ OrtStatus* ORT_API_CALL CudaEpFactory::CreateEpImpl( read_session_config_bool( {enable_cuda_graph_key, "ep.cuda.enable_cuda_graph", "enable_cuda_graph"}, config.enable_cuda_graph); + read_session_config_bool( + {enable_host_pageable_gather_key, "ep.cuda.enable_host_pageable_gather", "enable_host_pageable_gather"}, + config.enable_host_pageable_gather); read_session_config_non_negative_int( {min_runs_key, "ep.cuda.min_num_runs_before_cuda_graph_capture"}, config.min_num_runs_before_cuda_graph_capture); diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h index 4ccb919eb8bc9..09c523a8d3b28 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h @@ -534,6 +534,8 @@ struct CudaKernelAdapterRuntimeConfig { bool cudnn_conv1d_pad_to_nc1d = false; bool enable_cudnn = true; bool fuse_conv_bias = false; + bool enable_cuda_graph = false; + bool enable_host_pageable_gather = false; int sdpa_kernel = 0; int device_id = 0; bool do_copy_in_default_stream = true; @@ -1173,7 +1175,10 @@ class CudaKernel : public OpKernel { bool GetCudnnConv1dPadToNc1d() const { return runtime_config_->cudnn_conv1d_pad_to_nc1d; } bool UseTF32() const { return use_tf32_; } bool IsFuseConvBias() const { return runtime_config_->fuse_conv_bias; } + bool EnableHostPageableGather() const { return runtime_config_->enable_host_pageable_gather; } + bool IsCudaGraphEnabled() const { return runtime_config_->enable_cuda_graph; } bool IsArchAvailable(int arch) const { return GetDeviceProp().major >= arch; } + int GetDeviceId() const { return device_id_; } // Delegate to the base OpKernel::Info() which holds a safe copy of OpKernelInfo. // Do NOT store a reference to the constructor parameter — it becomes dangling. const OpKernelInfo& Info() const { return OpKernel::Info(); } diff --git a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc index 44e591d0326cf..4675f54f82076 100644 --- a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc +++ b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc @@ -3,6 +3,9 @@ #include #include +#include +#include +#include #include #include #include @@ -11,6 +14,10 @@ #include #include +#ifndef _WIN32 +#include +#endif + #include "core/common/common.h" #include "core/framework/execution_provider.h" #include "test/common/cuda_op_test_utils.h" @@ -18,6 +25,17 @@ #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" +#ifdef USE_CUDA +#include "contrib_ops/cuda/quantization/gather_block_quantized_data_policy.h" +#include "core/graph/model.h" +#include "core/graph/node_attr_utils.h" +#include "core/platform/env.h" +#include "core/session/IOBinding.h" +#include "core/session/inference_session.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "test/util/include/temp_dir.h" +#endif + namespace onnxruntime { namespace test { @@ -1386,6 +1404,295 @@ TEST(GatherBlockQuantizedOpTest, FpFloat16Output) { } #ifdef USE_CUDA +TEST(GatherBlockQuantizedOpTest, HostPageablePolicySelection) { + using contrib::cuda::GatherBlockQuantizedDataPolicy; + using contrib::cuda::SelectGatherBlockQuantizedDataPolicy; + + EXPECT_EQ(SelectGatherBlockQuantizedDataPolicy(false, true, true, true, true), + GatherBlockQuantizedDataPolicy::DeviceCopy); + EXPECT_EQ(SelectGatherBlockQuantizedDataPolicy(false, false, false, true, true), + GatherBlockQuantizedDataPolicy::DeviceCopy); + EXPECT_EQ(SelectGatherBlockQuantizedDataPolicy(true, true, true, true, true), + GatherBlockQuantizedDataPolicy::DirectHost); + EXPECT_EQ(SelectGatherBlockQuantizedDataPolicy(true, false, true, true, true), + GatherBlockQuantizedDataPolicy::DeviceCopy); + EXPECT_EQ(SelectGatherBlockQuantizedDataPolicy(true, true, false, true, true), + GatherBlockQuantizedDataPolicy::DeviceCopy); + EXPECT_EQ(SelectGatherBlockQuantizedDataPolicy(true, true, true, false, true), + GatherBlockQuantizedDataPolicy::DeviceCopy); + EXPECT_EQ(SelectGatherBlockQuantizedDataPolicy(true, true, true, true, false), + GatherBlockQuantizedDataPolicy::DeviceCopy); +} + +#if !defined(DISABLE_FLOAT8_TYPES) +TEST(GatherBlockQuantizedOpTest, FpFallbackWithPrepackingDisabledCuda) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "CUDA not available"; + } + + auto cuda_ep = CudaExecutionProviderWithOptions( + ProviderOptions{{"enable_host_pageable_gather", "0"}}); + if (cuda_ep == nullptr) { + GTEST_SKIP() << "CUDA EP not available"; + } + + OpTester test("GatherBlockQuantized", 1, kMSDomain); + test.AddAttribute("gather_axis", 0); + test.AddAttribute("quantize_axis", 1); + test.AddAttribute("block_size", 0); + test.AddInput("data", {2, 2}, + {Float8E4M3FN(1.0f), Float8E4M3FN(2.0f), + Float8E4M3FN(3.0f), Float8E4M3FN(4.0f)}, + true); + test.AddInput("indices", {2}, {1, 0}); + test.AddInput("scales", {2, 1}, {2.0f, 0.5f}); + test.AddOutput("output", {2, 2}, {1.5f, 2.0f, 2.0f, 4.0f}); + + SessionOptions session_options; + ASSERT_STATUS_OK( + session_options.config_options.AddConfigEntry(kOrtSessionOptionsConfigDisablePrepacking, "1")); + test.Config(session_options); + + std::vector> providers; + providers.push_back(std::move(cuda_ep)); + test.ConfigEps(std::move(providers)); + test.RunWithConfig(); +} + +TEST(GatherBlockQuantizedOpTest, FpDirectHostPageableCuda) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "CUDA not available"; + } + + int pageable_memory_access = 0; + int uses_host_page_tables = 0; +#if defined(CUDA_VERSION) && CUDA_VERSION >= 10020 + if (cudaDeviceGetAttribute(&pageable_memory_access, cudaDevAttrPageableMemoryAccess, 0) != cudaSuccess || + cudaDeviceGetAttribute(&uses_host_page_tables, cudaDevAttrPageableMemoryAccessUsesHostPageTables, 0) != + cudaSuccess) { + cudaGetLastError(); + GTEST_SKIP() << "CUDA pageable-memory attributes are unavailable"; + } +#endif + if (pageable_memory_access == 0 || uses_host_page_tables == 0) { + GTEST_SKIP() << "CUDA device does not use host page tables for pageable memory"; + } + + auto cuda_ep = CudaExecutionProviderWithOptions( + ProviderOptions{{"enable_host_pageable_gather", "1"}}); + if (cuda_ep == nullptr) { + GTEST_SKIP() << "CUDA EP not available"; + } + + OpTester test("GatherBlockQuantized", 1, kMSDomain); + test.AddAttribute("gather_axis", 0); + test.AddAttribute("quantize_axis", 1); + test.AddAttribute("block_size", 0); + test.AddInput("data", {2, 2}, + {Float8E4M3FN(1.0f), Float8E4M3FN(2.0f), + Float8E4M3FN(3.0f), Float8E4M3FN(4.0f)}, + true); + test.AddInput("indices", {3}, {-1, 0, 1}); + test.AddInput("scales", {2, 1}, {MLFloat16(2.0f), MLFloat16(0.5f)}); + test.AddOutput("output", {3, 2}, + {MLFloat16(1.5f), MLFloat16(2.0f), + MLFloat16(2.0f), MLFloat16(4.0f), + MLFloat16(1.5f), MLFloat16(2.0f)}); + + std::vector> providers; + providers.push_back(std::move(cuda_ep)); + test.ConfigEps(std::move(providers)); + test.RunWithConfig(); +} + +TEST(GatherBlockQuantizedOpTest, FpDirectHostPageableCudaGraph) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "CUDA not available"; + } + + int pageable_memory_access = 0; + int uses_host_page_tables = 0; +#if defined(CUDA_VERSION) && CUDA_VERSION >= 10020 + if (cudaDeviceGetAttribute(&pageable_memory_access, cudaDevAttrPageableMemoryAccess, 0) != cudaSuccess || + cudaDeviceGetAttribute(&uses_host_page_tables, cudaDevAttrPageableMemoryAccessUsesHostPageTables, 0) != + cudaSuccess) { + cudaGetLastError(); + GTEST_SKIP() << "CUDA pageable-memory attributes are unavailable"; + } +#endif + if (pageable_memory_access == 0 || uses_host_page_tables == 0) { + GTEST_SKIP() << "CUDA device does not use host page tables for pageable memory"; + } + + auto cuda_ep = CudaExecutionProviderWithOptions( + ProviderOptions{{"enable_cuda_graph", "1"}, + {"enable_host_pageable_gather", "1"}}); + if (cuda_ep == nullptr) { + GTEST_SKIP() << "CUDA EP not available"; + } + IExecutionProvider* cuda_ep_ptr = cuda_ep.get(); + + const std::vector data = { + Float8E4M3FN(1.0f), Float8E4M3FN(2.0f), + Float8E4M3FN(3.0f), Float8E4M3FN(4.0f), + Float8E4M3FN(5.0f), Float8E4M3FN(6.0f), + Float8E4M3FN(7.0f), Float8E4M3FN(8.0f)}; + const size_t data_bytes = data.size() * sizeof(data[0]); + const auto temp_dir_path = + std::filesystem::temp_directory_path() / + ("ort_gather_block_quantized_cuda_graph_" + + std::to_string(reinterpret_cast(&pageable_memory_access))); + TemporaryDirectory temp_dir(temp_dir_path.native()); + const auto data_path = temp_dir_path / "data.bin"; + { + std::ofstream data_file(data_path, std::ios::binary); + ASSERT_TRUE(data_file.good()); + data_file.write(reinterpret_cast(data.data()), static_cast(data_bytes)); + ASSERT_TRUE(data_file.good()); + } + + Env::MappedMemoryPtr mapped_memory; + ASSERT_STATUS_OK(Env::Default().MapFileIntoMemory(data_path.c_str(), 0, data_bytes, mapped_memory)); + OrtMemoryInfo cpu_memory_info{CPU, OrtDeviceAllocator}; + Tensor mapped_tensor(DataTypeImpl::GetType(), TensorShape({4, 2}), + mapped_memory.get(), cpu_memory_info); + OrtValue mapped_data_value; + Tensor::InitOrtValue(std::move(mapped_tensor), mapped_data_value); + + std::unordered_map domain_to_version = {{onnxruntime::kMSDomain, 1}}; + std::vector model_specific_functions; + auto model = std::make_unique( + "gather_block_quantized_cuda_graph", true, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, model_specific_functions, + DefaultLoggingManager().DefaultLogger(), ModelOptions(true, true)); + auto& graph = model->MainGraph(); + + std::vector tensor_types; + tensor_types.reserve(4); + auto add_tensor_type = [&](int elem_type, std::initializer_list dims) { + tensor_types.emplace_back(); + auto* type = &tensor_types.back(); + type->mutable_tensor_type()->set_elem_type(elem_type); + auto* shape = type->mutable_tensor_type()->mutable_shape(); + for (const int64_t dim : dims) { + shape->add_dim()->set_dim_value(dim); + } + return type; + }; + + auto& data_arg = graph.GetOrCreateNodeArg( + "data", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E4M3FN, {4, 2})); + auto& indices_arg = graph.GetOrCreateNodeArg( + "indices", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_INT64, {2})); + auto& scales_arg = graph.GetOrCreateNodeArg( + "scales", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, {4, 1})); + auto& output_arg = graph.GetOrCreateNodeArg( + "output", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, {2, 2})); + + ONNX_NAMESPACE::TensorProto data_initializer; + data_initializer.set_name("data"); + data_initializer.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E4M3FN); + data_initializer.add_dims(4); + data_initializer.add_dims(2); + data_initializer.mutable_raw_data()->assign(data_bytes, '\0'); + graph.AddInitializedTensor(data_initializer); + + NodeAttributes attributes = { + {"block_size", utils::MakeAttribute("block_size", int64_t{0})}, + {"gather_axis", utils::MakeAttribute("gather_axis", int64_t{0})}, + {"quantize_axis", utils::MakeAttribute("quantize_axis", int64_t{1})}, + }; + auto& node = graph.AddNode("gather_block_quantized", "GatherBlockQuantized", + "CUDA Graph direct host-pageable test", + {&data_arg, &indices_arg, &scales_arg}, {&output_arg}, + &attributes, onnxruntime::kMSDomain); + node.SetExecutionProviderType(cuda_ep_ptr->Type()); + ASSERT_STATUS_OK(graph.Resolve()); + + std::string model_string; + ASSERT_TRUE(model->ToProto().SerializeToString(&model_string)); + std::stringstream model_stream(model_string); + + SessionOptions session_options; + ASSERT_STATUS_OK(session_options.AddInitializer("data", &mapped_data_value)); + { + InferenceSession session(session_options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(cuda_ep))); + auto device_allocators = cuda_ep_ptr->CreatePreferredAllocators(); + const OrtMemoryInfo* device_memory_info = nullptr; + for (const auto& allocator : device_allocators) { + if (allocator->Info().device.Type() == OrtDevice::GPU && + allocator->Info().mem_type == OrtMemTypeDefault) { + device_memory_info = &allocator->Info(); + break; + } + } + ASSERT_NE(device_memory_info, nullptr); + ASSERT_STATUS_OK(session.Load(model_stream)); + ASSERT_STATUS_OK(session.Initialize()); + auto device_allocator = session.GetAllocator(*device_memory_info); + ASSERT_NE(device_allocator, nullptr); + + auto make_gpu_value = [&](const auto& values, const TensorShape& shape) { + using T = typename std::decay_t::value_type; + Tensor cpu_tensor(DataTypeImpl::GetType(), shape, const_cast(values.data()), cpu_memory_info); + Tensor gpu_tensor(DataTypeImpl::GetType(), shape, device_allocator); + ORT_THROW_IF_ERROR(cuda_ep_ptr->GetDataTransfer()->CopyTensor(cpu_tensor, gpu_tensor)); + OrtValue value; + Tensor::InitOrtValue(std::move(gpu_tensor), value); + return value; + }; + + std::vector indices = {0, 2}; + const std::vector scales = {1.0f, 0.5f, 2.0f, 0.25f}; + auto indices_value = make_gpu_value(indices, TensorShape({2})); + auto scales_value = make_gpu_value(scales, TensorShape({4, 1})); + auto output_value = make_gpu_value(std::vector(4), TensorShape({2, 2})); + + std::unique_ptr io_binding; + ASSERT_STATUS_OK(session.NewIOBinding(&io_binding)); + ASSERT_STATUS_OK(io_binding->BindInput("indices", indices_value)); + ASSERT_STATUS_OK(io_binding->BindInput("scales", scales_value)); + ASSERT_STATUS_OK(io_binding->BindOutput("output", output_value)); + + RunOptions run_options; + ASSERT_STATUS_OK(run_options.config_options.AddConfigEntry("gpu_graph_id", "1")); + for (int i = 0; i < 3; ++i) { + ASSERT_STATUS_OK(session.Run(run_options, *io_binding)); + } + ASSERT_TRUE(cuda_ep_ptr->IsGraphCaptured(1)); + + auto verify_output = [&](std::initializer_list expected) { + std::vector actual(expected.size()); + Tensor cpu_output(DataTypeImpl::GetType(), TensorShape({2, 2}), actual.data(), cpu_memory_info); + ASSERT_STATUS_OK(cuda_ep_ptr->GetDataTransfer()->CopyTensor(output_value.Get(), cpu_output)); + EXPECT_EQ(actual, std::vector(expected.begin(), expected.end())); + }; + verify_output({1.0f, 2.0f, 10.0f, 12.0f}); + indices = {3, 1}; + indices = {3, 1}; + Tensor cpu_indices(DataTypeImpl::GetType(), TensorShape({2}), indices.data(), cpu_memory_info); + ASSERT_STATUS_OK(cuda_ep_ptr->GetDataTransfer()->CopyTensor(cpu_indices, *indices_value.GetMutable())); + ASSERT_STATUS_OK(session.Run(run_options, *io_binding)); + verify_output({1.75f, 2.0f, 1.5f, 2.0f}); + +#ifndef _WIN32 + auto* mapped_data = reinterpret_cast(mapped_memory.get()); + mapped_data[2] = Float8E4M3FN(6.0f); + mapped_data[3] = Float8E4M3FN(8.0f); + mapped_data[6] = Float8E4M3FN(2.0f); + mapped_data[7] = Float8E4M3FN(4.0f); + ASSERT_STATUS_OK(session.Run(run_options, *io_binding)); + verify_output({0.5f, 1.0f, 3.0f, 4.0f}); + + ASSERT_EQ(0, madvise(mapped_memory.get(), data_bytes, MADV_DONTNEED)); + ASSERT_STATUS_OK(session.Run(run_options, *io_binding)); + verify_output({1.75f, 2.0f, 1.5f, 2.0f}); +#endif + } +} +#endif + TEST(GatherBlockQuantizedOpTest, FpBFloat16OutputCuda) { if (!HasCudaEnvironment(0)) { GTEST_SKIP() << "CUDA not available"; diff --git a/onnxruntime/test/providers/cuda/test_cases/cuda_test_provider.cc b/onnxruntime/test/providers/cuda/test_cases/cuda_test_provider.cc index 01c7573b9de14..54b6557343d05 100644 --- a/onnxruntime/test/providers/cuda/test_cases/cuda_test_provider.cc +++ b/onnxruntime/test/providers/cuda/test_cases/cuda_test_provider.cc @@ -34,6 +34,26 @@ namespace onnxruntime { void InitializeRegistry(); void DeleteRegistry(); +TEST(CUDAProviderOptionsTest, HostPageableGatherRoundTripAndHash) { + const CUDAExecutionProviderInfo default_info = + CUDAExecutionProviderInfo::FromProviderOptions({}); + EXPECT_FALSE(default_info.enable_host_pageable_gather); + + const CUDAExecutionProviderInfo disabled_info = + CUDAExecutionProviderInfo::FromProviderOptions({{"enable_host_pageable_gather", "0"}}); + EXPECT_FALSE(disabled_info.enable_host_pageable_gather); + + const CUDAExecutionProviderInfo enabled_info = + CUDAExecutionProviderInfo::FromProviderOptions({{"enable_host_pageable_gather", "1"}}); + EXPECT_TRUE(enabled_info.enable_host_pageable_gather); + const ProviderOptions serialized = CUDAExecutionProviderInfo::ToProviderOptions(enabled_info); + ASSERT_EQ(serialized.count("enable_host_pageable_gather"), 1u); + EXPECT_EQ(serialized.at("enable_host_pageable_gather"), "1"); + EXPECT_TRUE(CUDAExecutionProviderInfo::FromProviderOptions(serialized).enable_host_pageable_gather); + EXPECT_NE(std::hash{}(disabled_info), + std::hash{}(enabled_info)); +} + struct ProviderInfo_CUDA_TestImpl : ProviderInfo_CUDA { OrtStatus* SetCurrentGpuDeviceId(_In_ int) override { return nullptr; diff --git a/onnxruntime/test/util/default_providers.cc b/onnxruntime/test/util/default_providers.cc index 26a65a44c74c5..5a07194ba75be 100644 --- a/onnxruntime/test/util/default_providers.cc +++ b/onnxruntime/test/util/default_providers.cc @@ -201,6 +201,17 @@ std::unique_ptr CudaExecutionProviderWithOptions(const OrtCU #endif } +std::unique_ptr CudaExecutionProviderWithOptions(const ProviderOptions& provider_options) { +#ifdef USE_CUDA + Ort::CUDAProviderOptions cuda_options; + cuda_options.Update(provider_options); + return CudaExecutionProviderWithOptions(&*cuda_options); +#else + ORT_UNUSED_PARAMETER(provider_options); + return nullptr; +#endif +} + std::unique_ptr DefaultDnnlExecutionProvider() { #ifdef USE_DNNL OrtDnnlProviderOptions dnnl_options; diff --git a/onnxruntime/test/util/include/default_providers.h b/onnxruntime/test/util/include/default_providers.h index 306397d8745a2..345e9d4de529e 100644 --- a/onnxruntime/test/util/include/default_providers.h +++ b/onnxruntime/test/util/include/default_providers.h @@ -39,6 +39,7 @@ std::unique_ptr DefaultCudaExecutionProvider(); std::unique_ptr DefaultCudaNHWCExecutionProvider(); #endif std::unique_ptr CudaExecutionProviderWithOptions(const OrtCUDAProviderOptionsV2* provider_options); +std::unique_ptr CudaExecutionProviderWithOptions(const ProviderOptions& provider_options); std::unique_ptr DefaultDnnlExecutionProvider(); std::unique_ptr DnnlExecutionProviderWithOptions(const OrtDnnlProviderOptions* provider_options); // std::unique_ptr DefaultTvmExecutionProvider();