diff --git a/cmake/global_variables.cmake b/cmake/global_variables.cmake index 687b668a74..22282fb11d 100644 --- a/cmake/global_variables.cmake +++ b/cmake/global_variables.cmake @@ -81,6 +81,8 @@ file(GLOB generator_srcs CONFIGURE_DEPENDS "${GENERATORS_ROOT}/cuda/session_options.cpp" "${GENERATORS_ROOT}/nvtensorrtrtx/*.h" "${GENERATORS_ROOT}/nvtensorrtrtx/*.cpp" + "${GENERATORS_ROOT}/amdgpu/*.h" + "${GENERATORS_ROOT}/amdgpu/*.cpp" "${GENERATORS_ROOT}/vitisai/*.h" "${GENERATORS_ROOT}/vitisai/*.cpp" "${GENERATORS_ROOT}/dml/session_options.h" diff --git a/src/amdgpu/interface.cpp b/src/amdgpu/interface.cpp new file mode 100644 index 0000000000..bc5cba9141 --- /dev/null +++ b/src/amdgpu/interface.cpp @@ -0,0 +1,305 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// Modifications Copyright(C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// +// GPU-resident KV cache for the AMDGPU device. +// +// p_device_ is an opaque device-buffer handle, so no pointer arithmetic or +// dereference. All copies go through ORT's CopyTensors. Offset copies fall +// back to CPU staging via CopyThroughCpu. + +#include "../generators.h" +#include "../search.h" +#include "interface.h" + +#include + +namespace Generators { +namespace AMDGPU { + +const char* device_label = "amdgpu"; +const char* label_cpu = "cpu"; + +struct GpuMemory final : DeviceBuffer { + GpuMemory(size_t size, Ort::Allocator* allocator, const OrtMemoryInfo* memory_info) + : owned_{true}, ort_allocator_{allocator}, ort_memory_info_{memory_info} { + size_in_bytes_ = size; + p_device_ = static_cast(ort_allocator_->Alloc(size_in_bytes_)); + } + + GpuMemory(void* p, size_t size, Ort::Allocator* allocator, const OrtMemoryInfo* memory_info) + : owned_{false}, ort_allocator_{allocator}, ort_memory_info_{memory_info} { + size_in_bytes_ = size; + p_device_ = static_cast(p); + } + + ~GpuMemory() override { + if (owned_) + ort_allocator_->Free(p_device_); + if (p_cpu_) + free(p_cpu_); + } + + const char* GetType() const override { return device_label; } + + void AllocateCpu() override { + if (!p_cpu_) + p_cpu_ = static_cast(malloc(size_in_bytes_)); + } + + void CopyDeviceToCpu() override { + if (!ort_allocator_) + throw std::runtime_error("AMDGPU allocator not initialized"); + + AllocateCpu(); + + int64_t shape_val = static_cast(size_in_bytes_); + std::span shape{&shape_val, 1}; + auto src_tensor = OrtValue::CreateTensor(*ort_memory_info_, p_device_, size_in_bytes_, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8); + + auto cpu_mem_info = OrtMemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + auto dst_tensor = OrtValue::CreateTensor(*cpu_mem_info, p_cpu_, size_in_bytes_, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8); + + const std::vector src_ptrs = {src_tensor.get()}; + const std::vector dst_ptrs = {dst_tensor.get()}; + GetOrtEnv().CopyTensors(src_ptrs, dst_ptrs, nullptr); + } + + void CopyCpuToDevice() override { + if (!ort_allocator_) + throw std::runtime_error("AMDGPU allocator not initialized"); + assert(p_cpu_); + + int64_t shape_val = static_cast(size_in_bytes_); + std::span shape{&shape_val, 1}; + auto cpu_mem_info = OrtMemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + auto src_tensor = OrtValue::CreateTensor(*cpu_mem_info, p_cpu_, size_in_bytes_, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8); + + auto dst_tensor = OrtValue::CreateTensor(*ort_memory_info_, p_device_, size_in_bytes_, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8); + + const std::vector src_ptrs = {src_tensor.get()}; + const std::vector dst_ptrs = {dst_tensor.get()}; + GetOrtEnv().CopyTensors(src_ptrs, dst_ptrs, nullptr); + } + + void CopyFrom(size_t begin_dest, DeviceBuffer& source, size_t begin_source, size_t size_in_bytes) override { + if (!ort_allocator_) + throw std::runtime_error("AMDGPU allocator not initialized"); + + // Opaque handle, so wrap the whole buffer and let CopyTensors do the copy. + if (strcmp(source.GetType(), device_label) == 0 && begin_source == 0 && begin_dest == 0) { + // Full-buffer device-to-device copy. + int64_t shape_val = static_cast(size_in_bytes); + std::span shape{&shape_val, 1}; + auto src_tensor = OrtValue::CreateTensor(*ort_memory_info_, source.p_device_, size_in_bytes, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8); + auto dst_tensor = OrtValue::CreateTensor(*ort_memory_info_, p_device_, size_in_bytes, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8); + + const std::vector src_ptrs = {src_tensor.get()}; + const std::vector dst_ptrs = {dst_tensor.get()}; + GetOrtEnv().CopyTensors(src_ptrs, dst_ptrs, nullptr); + } else if (strcmp(source.GetType(), label_cpu) == 0 && begin_source == 0 && begin_dest == 0) { + // Full-buffer CPU-to-device copy. + int64_t shape_val = static_cast(size_in_bytes); + std::span shape{&shape_val, 1}; + auto cpu_mem_info = OrtMemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + auto src_tensor = OrtValue::CreateTensor(*cpu_mem_info, source.p_device_, size_in_bytes, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8); + auto dst_tensor = OrtValue::CreateTensor(*ort_memory_info_, p_device_, size_in_bytes, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8); + + const std::vector src_ptrs = {src_tensor.get()}; + const std::vector dst_ptrs = {dst_tensor.get()}; + GetOrtEnv().CopyTensors(src_ptrs, dst_ptrs, nullptr); + } else { + // Offset copies can't sub-buffer-view an opaque handle, so stage through CPU. + CopyThroughCpu(*this, begin_dest, source, begin_source, size_in_bytes); + } + } + + void Zero() override { + if (!ort_allocator_) + throw std::runtime_error("AMDGPU allocator not initialized"); + + // TODO: device-side zero to avoid the host staging buffer. Off the decode hot path for now. + std::vector zero_buffer(size_in_bytes_, 0); + + int64_t shape_val = static_cast(size_in_bytes_); + std::span shape{&shape_val, 1}; + auto cpu_mem_info = OrtMemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + auto src_tensor = OrtValue::CreateTensor(*cpu_mem_info, zero_buffer.data(), size_in_bytes_, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8); + + auto dst_tensor = OrtValue::CreateTensor(*ort_memory_info_, p_device_, size_in_bytes_, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8); + + const std::vector src_ptrs = {src_tensor.get()}; + const std::vector dst_ptrs = {dst_tensor.get()}; + GetOrtEnv().CopyTensors(src_ptrs, dst_ptrs, nullptr); + } + + bool owned_; // If we own the memory, we free it on destruction + Ort::Allocator* ort_allocator_; + const OrtMemoryInfo* ort_memory_info_; +}; + +const char* pinned_device_label = "amdgpu_pinned"; + +// DeviceBuffer over a host-accessible allocation: one mapped pointer is both CPU-writable +// and GPU-readable, so p_cpu_ aliases p_device_ and the copy methods are no-ops. Used for +// decode inputs so the CPU updates them in place with no roundtrip. +struct PinnedMemory final : DeviceBuffer { + PinnedMemory(size_t size, Ort::Allocator* allocator) : owned_{true}, ort_allocator_{allocator} { + size_in_bytes_ = size; + p_device_ = static_cast(ort_allocator_->Alloc(size_in_bytes_)); + p_cpu_ = p_device_; // mapped: one pointer for both + } + + PinnedMemory(void* p, size_t size, Ort::Allocator* allocator) : owned_{false}, ort_allocator_{allocator} { + size_in_bytes_ = size; + p_device_ = static_cast(p); + p_cpu_ = p_device_; + } + + ~PinnedMemory() override { + if (owned_) + ort_allocator_->Free(p_device_); + // p_cpu_ aliases p_device_, do not free it separately. + } + + const char* GetType() const override { return pinned_device_label; } + void AllocateCpu() override {} // p_cpu_ already valid (== p_device_) + void CopyDeviceToCpu() override {} // same memory: nothing to copy + void CopyCpuToDevice() override {} + + void CopyFrom(size_t begin_dest, DeviceBuffer& source, size_t begin_source, size_t size_in_bytes) override { + // Pinned memory is CPU-addressable; source may be on any device. + CopyThroughCpu(*this, begin_dest, source, begin_source, size_in_bytes); + } + + void Zero() override { memset(p_device_, 0, size_in_bytes_); } + + bool owned_; + Ort::Allocator* ort_allocator_; +}; + +struct InterfaceImpl : DeviceInterface { + DeviceType GetType() const override { return DeviceType::AMDGPU; } + + void InitOrt(const OrtApi& api, Ort::Allocator& allocator) override { + Ort::api = &api; + assert(!ort_allocator_); + ort_allocator_ = &allocator; + // Cache the memory info so tensors wrapping p_device_ carry the allocator's device attributes. + ort_memory_info_ = &ort_allocator_->GetInfo(); + } + + Ort::Allocator& GetAllocator() override { + return *ort_allocator_; + } + + void InitHostAccessible(Ort::Allocator& allocator) override { + ort_pinned_allocator_ = &allocator; + } + + Ort::Allocator* GetHostAccessibleAllocator() override { + return ort_pinned_allocator_; + } + + Ort::Allocator* PinnedAllocator() const { return ort_pinned_allocator_; } + + std::shared_ptr AllocateBase(size_t size) override { + return std::make_shared(size, ort_allocator_, ort_memory_info_); + } + + std::shared_ptr WrapMemoryBase(void* p, size_t size) override { + return std::make_shared(p, size, ort_allocator_, ort_memory_info_); + } + + std::unique_ptr CreateGreedy(const GeneratorParams& params) override { + return GetDeviceInterface(DeviceType::CPU)->CreateGreedy(params); + } + + std::unique_ptr CreateBeam(const GeneratorParams& params) override { + return GetDeviceInterface(DeviceType::CPU)->CreateBeam(params); + } + + void Synchronize() override {} + + // The umbrella EP selects its backend from provider options (notably "profile"). The trivial + // device-init session that EnsureDeviceOrtInit builds gets empty options by default, so forward + // the user's umbrella options here or the umbrella has no backend to create for that session. + void ShapeInitSessionProviderOptions(Config::ProviderOptions& init_options, + const Config::ProviderOptions* user_options) const override { + if (user_options) { + for (const auto& opt : user_options->options) { + init_options.options.emplace_back(opt); + } + } + } + + private: + Ort::Allocator* ort_allocator_{}; + const OrtMemoryInfo* ort_memory_info_{}; + // Host-accessible allocator, set by InitHostAccessible when one is available. + Ort::Allocator* ort_pinned_allocator_{}; +}; + +// Inputs-only interface: allocations come from the host-accessible allocator, everything else +// delegates to the base AMDGPU interface. Updates run on the CPU in place. +struct PinnedInputsImpl : DeviceInterface { + explicit PinnedInputsImpl(InterfaceImpl& base) : base_{base} {} + + DeviceType GetType() const override { return DeviceType::AMDGPU; } + void InitOrt(const OrtApi& api, Ort::Allocator& allocator) override { base_.InitOrt(api, allocator); } + Ort::Allocator& GetAllocator() override { return *base_.PinnedAllocator(); } + Ort::Allocator* GetHostAccessibleAllocator() override { return base_.PinnedAllocator(); } + + std::shared_ptr AllocateBase(size_t size) override { + return std::make_shared(size, base_.PinnedAllocator()); + } + std::shared_ptr WrapMemoryBase(void* p, size_t size) override { + return std::make_shared(p, size, base_.PinnedAllocator()); + } + + std::unique_ptr CreateGreedy(const GeneratorParams& params) override { return base_.CreateGreedy(params); } + std::unique_ptr CreateBeam(const GeneratorParams& params) override { return base_.CreateBeam(params); } + void Synchronize() override { base_.Synchronize(); } + + bool Cast(void* input, void* output, ONNXTensorElementDataType input_type, + ONNXTensorElementDataType output_type, size_t element_count) override { + return base_.Cast(input, output, input_type, output_type, element_count); + } + + // In-place CPU updates on the pinned buffer. Delegate to the CPU interface and report + // success so the caller skips its copy-back path. + bool UpdatePositionIds(void* position_ids, int batch_beam_size, int total_length, + int new_kv_length, ONNXTensorElementDataType type) override { + return GetDeviceInterface(DeviceType::CPU)->UpdatePositionIds(position_ids, batch_beam_size, total_length, new_kv_length, type); + } + bool UpdateAttentionMask(void* next_mask_data, void* mask_data, int batch_beam_size, int new_kv_length, + int total_length, int max_length, bool update_only, + ONNXTensorElementDataType type) override { + return GetDeviceInterface(DeviceType::CPU)->UpdateAttentionMask(next_mask_data, mask_data, batch_beam_size, new_kv_length, total_length, max_length, update_only, type); + } + + InterfaceImpl& base_; +}; + +} // namespace AMDGPU + +static std::unique_ptr g_amdgpu_device; +static std::unique_ptr g_amdgpu_pinned_inputs; + +DeviceInterface* GetAMDGPUInterface() { + if (!g_amdgpu_device) + g_amdgpu_device = std::make_unique(); + return g_amdgpu_device.get(); +} + +DeviceInterface* GetAMDGPUPinnedInputsInterface() { + auto* base = static_cast(GetAMDGPUInterface()); + if (!base->GetHostAccessibleAllocator()) + return nullptr; // no pinned allocator, caller falls back to the device interface + if (!g_amdgpu_pinned_inputs) + g_amdgpu_pinned_inputs = std::make_unique(*base); + return g_amdgpu_pinned_inputs.get(); +} + +} // namespace Generators diff --git a/src/amdgpu/interface.h b/src/amdgpu/interface.h new file mode 100644 index 0000000000..0957f9a1ff --- /dev/null +++ b/src/amdgpu/interface.h @@ -0,0 +1,17 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// Modifications Copyright(C) 2026 Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +namespace Generators { + +DeviceInterface* GetAMDGPUInterface(); + +// Inputs-only interface backed by the host-accessible allocator. Decode inputs allocated +// here are CPU-writable and GPU-readable, so per-step updates happen in place with no +// roundtrip. KV and scoring keep the default interface. Null if no host-accessible allocator. +DeviceInterface* GetAMDGPUPinnedInputsInterface(); + +} // namespace Generators diff --git a/src/amdgpu/session_options.cpp b/src/amdgpu/session_options.cpp new file mode 100644 index 0000000000..534ebbe31f --- /dev/null +++ b/src/amdgpu/session_options.cpp @@ -0,0 +1,149 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// Modifications Copyright(C) 2026 Advanced Micro Devices, Inc. All rights reserved. + +#include "session_options.h" + +#include "../generators.h" +#include "../models/session_options.h" +#include "interface.h" + +#include +#include + +#if defined(_WIN32) +#include +#endif + +namespace Generators::AMDGPUExecutionProvider { + +namespace { + +constexpr auto ep_registration_name_ = "AMDGPUExecutionProvider"; +#if defined(_WIN32) +constexpr auto ep_filename_ = "amdgpu-ep.dll"; +#else +constexpr auto ep_filename_ = "libamdgpu-ep.so"; +#endif + +// The AMDGPU umbrella is a plugin EP: unlike legacy in-proc EPs it must be registered on the +// OrtEnv via RegisterExecutionProviderLibrary before AppendExecutionProvider_V2 can find its +// OrtEpDevice. Mirror the RyzenAI interface's self-registration (resolve the DLL next to the +// genai/ort module or the executable) so the C model_benchmark -- which has no --ep_library flag +// -- still loads the umbrella. Registration is keyed per-OrtEnv; re-registration is benign. +void EnsureUmbrellaEpRegistered() { + namespace fs = std::filesystem; + std::error_code ec; + + fs::path ep_path; + +#if defined(_WIN32) + const auto hmod_of = [](LPCVOID func) -> HMODULE { + MEMORY_BASIC_INFORMATION mbi; + if (VirtualQuery(func, &mbi, sizeof(mbi)) && mbi.AllocationBase) { + return reinterpret_cast(mbi.AllocationBase); + } + return nullptr; + }; + + const auto find_next_to_module = [&](HMODULE hmod) -> fs::path { + wchar_t buffer[MAX_PATH + 1] = {0}; + if (GetModuleFileNameW(hmod, buffer, MAX_PATH + 1)) { + if (auto dir = fs::path{buffer}.remove_filename(); !dir.empty()) { + if (auto candidate = dir / ep_filename_; fs::exists(candidate, ec)) { + return candidate; + } + } + } + return {}; + }; + + if (ep_path.empty()) { + // next to onnxruntime-genai.dll (GetAMDGPUInterface is a genai symbol in that module) + if (const auto hmod = hmod_of(reinterpret_cast(&GetAMDGPUInterface))) { + ep_path = find_next_to_module(hmod); + } + } + if (ep_path.empty()) { + // next to onnxruntime.dll + if (const auto hmod = hmod_of(reinterpret_cast(Ort::api->RegisterExecutionProviderLibrary))) { + ep_path = find_next_to_module(hmod); + } + } + if (ep_path.empty()) { + // next to the current executable + if (const auto hmod = GetModuleHandleA(nullptr)) { + ep_path = find_next_to_module(hmod); + } + } +#endif // _WIN32 + + if (ep_path.empty()) { + ep_path = fs::current_path(ec) / ep_filename_; + } + + try { + Ort::RegisterExecutionProviderLibrary(&GetOrtEnv(), ep_registration_name_, ep_path.native().c_str()); + } catch (const Ort::Exception& e) { + if (std::string(e.what()).find("already registered") == std::string::npos) { + throw std::runtime_error("Failed to register AMDGPU execution provider library from '" + + ep_path.string() + "': " + e.what()); + } + } +} + +// Emit static-padding hints so the EP pads the prefill token axis to max_length and +// compiles it once, instead of recompiling per prompt length. +void SetStaticPaddingConfig(OrtSessionOptions& session_options, const Config& config) { + const auto& decoder = config.model.decoder; + const std::string seq_len = std::to_string(config.search.max_length); + const std::string pad_inputs = + decoder.inputs.input_ids + ":1," + decoder.inputs.position_ids + ":1"; + const std::string pad_outputs = decoder.outputs.logits + ":1"; + + session_options.AddConfigEntry("ep.migraphx.static_pad_seq", "1"); + session_options.AddConfigEntry("ep.migraphx.static_pad_seq_len", seq_len.c_str()); + session_options.AddConfigEntry("ep.migraphx.static_pad_inputs", pad_inputs.c_str()); + session_options.AddConfigEntry("ep.migraphx.static_pad_outputs", pad_outputs.c_str()); + + session_options.AddConfigEntry("ep.migraphx.hip_graph_enable", "1"); +} + +// The umbrella EP consumes its own provider options (e.g. "profile") from session-config +// entries prefixed "ep.amdgpuexecutionprovider." (see the umbrella's CreateEp), not from the +// AppendExecutionProvider_V2 ep_options channel -- those do not surface through +// GetSessionOptionsConfigEntries. Bridge the genai_config provider options into that prefixed +// form so the umbrella selects the requested backend (profile=hip -> hipgpu backend). +void ForwardUmbrellaProviderOptions(OrtSessionOptions& session_options, + const Config::ProviderOptions& provider_options) { + const std::string umbrella_prefix = "ep.amdgpuexecutionprovider."; + for (const auto& [key, value] : provider_options.options) { + session_options.AddConfigEntry((umbrella_prefix + key).c_str(), value.c_str()); + } +} + +} // namespace + +DeviceInterface* AppendExecutionProvider(OrtSessionOptions& session_options, + const Config::ProviderOptions& provider_options, + const Config& config, + bool /*disable_graph_capture*/) { + SetStaticPaddingConfig(session_options, config); + + // Umbrella-level hint: the model architecture drives the EP's backend routing. + session_options.AddConfigEntry("ep.amdgpuexecutionprovider.model_arch", config.model.type.c_str()); + + // DirectML backend: host-accessible decode inputs. + session_options.AddConfigEntry("ep.directml.enable_host_accessible", "1"); + + EnsureUmbrellaEpRegistered(); + ForwardUmbrellaProviderOptions(session_options, provider_options); + + AppendExecutionProviderV2(session_options, provider_options, + DeviceType::AMDGPU, ep_registration_name_); + + return GetAMDGPUInterface(); +} + +} // namespace Generators::AMDGPUExecutionProvider diff --git a/src/amdgpu/session_options.h b/src/amdgpu/session_options.h new file mode 100644 index 0000000000..9696bae245 --- /dev/null +++ b/src/amdgpu/session_options.h @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// Modifications Copyright(C) 2026 Advanced Micro Devices, Inc. All rights reserved. +#pragma once + +#include "../generators.h" + +namespace Generators::AMDGPUExecutionProvider { + +DeviceInterface* AppendExecutionProvider(OrtSessionOptions& session_options, + const Config::ProviderOptions& provider_options, + const Config& config, + bool disable_graph_capture = false); + +} // namespace Generators::AMDGPUExecutionProvider diff --git a/src/config.cpp b/src/config.cpp index fd4a740511..5b9bb92472 100644 --- a/src/config.cpp +++ b/src/config.cpp @@ -44,6 +44,10 @@ std::string_view NormalizeProviderName(std::string_view name) { return "RyzenAI"; } else if (lower_name == "nvtensorrtrtx") { return "NvTensorRtRtx"; + } else if (lower_name == "amdgpu" || + lower_name == "amdgpuexecutionprovider") { + // Accept canonical and catalog forms, all route to AMDGPU. + return "AMDGPU"; } return name; // Return name unchanged } @@ -1519,6 +1523,8 @@ bool IsGraphCaptureEnabled(const Config::SessionOptions& session_options) { } } return true; + } else if (provider_options->name == "AMDGPU") { + return true; } else if (provider_options->name == "WebGPU") { for (const auto& value : provider_options->options) { if (value.first == "enableGraphCapture" && value.second == "1") { diff --git a/src/generators.cpp b/src/generators.cpp index a287ab3711..0f11740fb8 100644 --- a/src/generators.cpp +++ b/src/generators.cpp @@ -23,6 +23,7 @@ #include "webgpu/interface.h" #include "openvino/interface.h" #include "ryzenai/interface.h" +#include "amdgpu/interface.h" #include "engine/engine.h" #if defined(_WIN32) @@ -335,6 +336,9 @@ DeviceInterface* OrtGlobals::GetDeviceInterface(DeviceType type) { owned_interfaces_.push_back(CreateRyzenAIInterface(*env_)); slot = owned_interfaces_.back().get(); break; + case DeviceType::AMDGPU: + slot = GetAMDGPUInterface(); // static singleton owned by amdgpu/interface.cpp + break; case DeviceType::CPU: default: owned_interfaces_.push_back(CreateCpuInterface()); @@ -363,6 +367,8 @@ std::string to_string(DeviceType device_type) { return "NvTensorRtRtx"; case DeviceType::RyzenAI: return "RyzenAI"; + case DeviceType::AMDGPU: + return "AMDGPU"; default: throw std::runtime_error("Unknown device type"); } @@ -774,6 +780,7 @@ void Generator::GenerateNextToken() { auto next_tokens = search_->GetNextTokens(); if (last_action_ == Action::rewound) search_->AppendTokens(next_tokens); + ComputeLogits(next_tokens); } if (guidance_logits_processor_) { diff --git a/src/generators.h b/src/generators.h index 2c0e60a057..79d3361641 100644 --- a/src/generators.h +++ b/src/generators.h @@ -172,6 +172,10 @@ struct OrtGlobals { // ~allocator_ runs first. std::unique_ptr session_; std::unique_ptr allocator_; + // Optional host-accessible allocator for decode inputs, owned by the OrtEnv (do not free). + // Null if unavailable, in which case inputs stay on the default device allocator. + Ort::Allocator* host_accessible_allocator_{}; + int device_id_{}; // Device this allocator is bound to (0 unless a specific device was selected). }; Allocator device_allocators_[static_cast(DeviceType::MAX)]; diff --git a/src/models/kv_cache.cpp b/src/models/kv_cache.cpp index 30c9ba23e1..e65bd83cc6 100644 --- a/src/models/kv_cache.cpp +++ b/src/models/kv_cache.cpp @@ -432,6 +432,7 @@ DefaultKeyValueCache::DefaultKeyValueCache(State& state) } presents_.push_back(OrtValue::CreateTensor(Allocator(), tensor_shape, type_)); + // WebGPU has no Zero() implementation; every other backend (incl. AMDGPU) zero-inits the KV. if (Device().GetType() != DeviceType::WEBGPU) { ByteWrapTensor(Device(), *presents_.back()).Zero(); } diff --git a/src/models/logits.cpp b/src/models/logits.cpp index 32d1992c43..772486b6e7 100644 --- a/src/models/logits.cpp +++ b/src/models/logits.cpp @@ -13,7 +13,7 @@ Logits::Logits(State& state) : state_{state}, shape_{static_cast(state_.params_->BatchBeamSize()), 0, model_.config_->model.vocab_size}, type_{model_.session_info_.GetOutputDataType(model_.config_->model.decoder.outputs.logits)} { - output_raw_ = std::make_unique(model_.p_device_inputs_, type_); + output_raw_ = std::make_unique(model_.p_device_logits_, type_); input_sequence_lengths.resize(state_.params_->search.batch_size); @@ -42,10 +42,10 @@ DeviceSpan Logits::Get() { const size_t num_beams = state_.params_->search.num_beams; // create new OrtValue for logits_of_last_token and use output_last_tokens_ to hold it - output_last_tokens_ = OrtValue::CreateTensor(model_.p_device_inputs_->GetAllocator(), shape_last, type_); + output_last_tokens_ = OrtValue::CreateTensor(model_.p_device_logits_->GetAllocator(), shape_last, type_); if (type_ == Ort::TypeToTensorType || type_ == Ort::TypeToTensorType) - logits_of_last_token_fp32_ = OrtValue::CreateTensor(model_.p_device_inputs_->GetAllocator(), shape_); + logits_of_last_token_fp32_ = OrtValue::CreateTensor(model_.p_device_logits_->GetAllocator(), shape_); logits_of_last_token = output_last_tokens_.get(); @@ -53,7 +53,7 @@ DeviceSpan Logits::Get() { size_t vocab_index = 0; // Simpler math to have this index go up by vocab_size for every logit chunk we process auto logits_raw = output_raw_->GetByteSpan(); - auto logits_last_tokens = ByteWrapTensor(*model_.p_device_inputs_, *logits_of_last_token); + auto logits_last_tokens = ByteWrapTensor(*model_.p_device_logits_, *logits_of_last_token); for (int batch_index = 0; batch_index < state_.params_->search.batch_size; batch_index++) { // Find the first non pad token from the end @@ -71,12 +71,12 @@ DeviceSpan Logits::Get() { // Convert from float16/bfloat16 to float32 if necessary if (type_ == Ort::TypeToTensorType || type_ == Ort::TypeToTensorType) { - Cast(*logits_of_last_token, logits_of_last_token_fp32_, *model_.p_device_inputs_, Ort::TypeToTensorType); + Cast(*logits_of_last_token, logits_of_last_token_fp32_, *model_.p_device_logits_, Ort::TypeToTensorType); logits_of_last_token = logits_of_last_token_fp32_.get(); } if (logits_.empty() || logits_of_last_token->GetTensorMutableRawData() != logits_.Span().data()) - logits_ = WrapTensor(*model_.p_device_inputs_, *logits_of_last_token); + logits_ = WrapTensor(*model_.p_device_logits_, *logits_of_last_token); return logits_; } diff --git a/src/models/model.cpp b/src/models/model.cpp index 5210115fe1..6ea17d9b88 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -32,6 +32,7 @@ #include "videochat_flash_processor.h" #include "mistral3_image_processor.h" #include "../dml/interface.h" +#include "../amdgpu/interface.h" #include "../openvino/interface.h" #include "../qnn/interface.h" #include "../ryzenai/interface.h" @@ -434,7 +435,7 @@ void EnsureDeviceOrtInit(DeviceInterface& device, const Config& config) { // This ensures memory allocated on-device for model inputs/outputs is valid for the lifetime of GenAI. // Names for the device types used by 'SetProviderSessionOptions' - static const char* device_type_names[] = {"CPU (Not used, see above)", "cuda", "DML", "WebGPU", "QNN", "QNN", "OpenVINO (Not used, see above)", "NvTensorRtRtx", "RyzenAI"}; + static const char* device_type_names[] = {"CPU (Not used, see above)", "cuda", "DML", "WebGPU", "QNN", "QNN", "OpenVINO (Not used, see above)", "NvTensorRtRtx", "RyzenAI", "AMDGPU"}; static_assert(std::size(device_type_names) == static_cast(DeviceType::MAX)); // Create an OrtSessionOptions and set the options to use the DeviceType we're using here @@ -466,15 +467,25 @@ void EnsureDeviceOrtInit(DeviceInterface& device, const Config& config) { const auto trivial_model = GetTrivialModel(); allocator.session_ = OrtSession::Create(GetOrtEnv(), trivial_model.data(), trivial_model.size(), session_options.get()); - // Names for the device memory types used by 'OrtMemoryInfo::Create' - static const char* device_memory_type_names[] = {"CPU (Not used, see above)", "Cuda", "DML", "WebGPU_Buf", "QnnHtpShared", "QnnHtpShared", "OpenVINO (Not used, see above)", "Cuda", "Cpu"}; + // Names for the device memory types used by 'OrtMemoryInfo::Create'. AMDGPU is "Hip". + static const char* device_memory_type_names[] = {"CPU (Not used, see above)", "Cuda", "DML", "WebGPU_Buf", "QnnHtpShared", "QnnHtpShared", "OpenVINO (Not used, see above)", "Cuda", "Cpu", "Hip"}; static_assert(std::size(device_memory_type_names) == static_cast(DeviceType::MAX)); // Get the allocator from the OrtSession for the DeviceType (it's called 'AllocatorCreate' but it's really 'AllocatorGet') auto name = device_memory_type_names[static_cast(type)]; + // AMDGPU: use the selected device's id rather than a hardcoded 0. Single-GPU resolves to 0. + if (type == DeviceType::AMDGPU) { + auto ep_devices = FindRegisteredEpDevices("AMDGPUExecutionProvider"); + if (user_provider_options) + ep_devices = ApplyDeviceFiltering(*user_provider_options, ep_devices); + if (!ep_devices.empty()) { + if (const OrtMemoryInfo* mi = Ort::api->EpDevice_MemoryInfo(ep_devices.front(), OrtDeviceMemoryType_DEFAULT)) + Ort::ThrowOnError(Ort::api->MemoryInfoGetId(mi, &allocator.device_id_)); + } + } try { auto memory_info = OrtMemoryInfo::Create(name, OrtAllocatorType::OrtDeviceAllocator, - 0, OrtMemType::OrtMemTypeDefault); + allocator.device_id_, OrtMemType::OrtMemTypeDefault); allocator.allocator_ = Ort::Allocator::Create(*allocator.session_, *memory_info); } catch (const Ort::Exception& e) { // WebGPU memory type name changed from "WebGPU_Buffer" to "WebGPU_Buf" in ORT 1.24.3. @@ -499,6 +510,33 @@ void EnsureDeviceOrtInit(DeviceInterface& device, const Config& config) { throw std::runtime_error("Unexpected failure to create device memory allocator for " + std::string(name)); } device.InitOrt(*Ort::api, *allocator.allocator_); + + // Host-accessible allocator for decode inputs (AMDGPU only). Request it via GetSharedAllocator; + // if unavailable, creation returns null and callers fall back to the default device inputs path. + if (!allocator.host_accessible_allocator_ && type == DeviceType::AMDGPU) { + try { + // Use the host-accessible memory-info on the same device as the compute allocator. + const OrtMemoryInfo* host_info = nullptr; + for (const OrtEpDevice* ep_device : FindRegisteredEpDevices("AMDGPUExecutionProvider")) { + const OrtMemoryInfo* mi = + Ort::api->EpDevice_MemoryInfo(ep_device, OrtDeviceMemoryType_HOST_ACCESSIBLE); + if (!mi) + continue; + int host_device_id = 0; + Ort::ThrowOnError(Ort::api->MemoryInfoGetId(mi, &host_device_id)); + if (host_device_id == allocator.device_id_) { + host_info = mi; + break; + } + } + if (host_info) + allocator.host_accessible_allocator_ = GetOrtEnv().GetSharedAllocator(*host_info); + } catch (const Ort::Exception&) { + allocator.host_accessible_allocator_ = nullptr; + } + if (allocator.host_accessible_allocator_) + device.InitHostAccessible(*allocator.host_accessible_allocator_); + } } void SessionInfo::Add(OrtSession& session) { @@ -592,6 +630,18 @@ Model::Model(std::unique_ptr config) : config_{std::move(config)} { else p_device_inputs_ = GetDeviceInterface(DeviceType::CPU); + // Host-accessible decode inputs (AMDGPU): route the small decode inputs through a + // host-accessible interface so the CPU updates them in place with no per-step roundtrip. + // KV cache and scoring stay on the default interface. Falls back if no allocator. + if (p_device_->GetType() == DeviceType::AMDGPU && + p_device_->GetHostAccessibleAllocator() != nullptr) { + if (auto* pinned = GetAMDGPUPinnedInputsInterface()) + p_device_inputs_ = pinned; + } + + // Logits are CPU-read; AMDGPU host-accessible inputs aren't CPU-read-coherent, so route to CPU. + p_device_logits_ = (p_device_->GetType() == DeviceType::AMDGPU) ? GetDeviceInterface(DeviceType::CPU) : p_device_inputs_; + // Search and sampling are performed on the CPU for all device types, // except for CUDA and NvTensorRtRtx, where this is performed on the device. if (p_device_->GetType() == DeviceType::CUDA || diff --git a/src/models/model.h b/src/models/model.h index 680197a1ef..72b9d73692 100644 --- a/src/models/model.h +++ b/src/models/model.h @@ -173,6 +173,7 @@ struct Model : std::enable_shared_from_this, LeakChecked, External DeviceInterface* p_device_{}; // The device we're running on (matches device_type_) used for things that work the same on all devices DeviceInterface* p_device_inputs_{}; // For some model inputs, the device might be the CPU device (all but KV cache currently for WebGPU and DML) + DeviceInterface* p_device_logits_{}; // Logits are CPU-read; on AMDGPU the host-accessible input allocator isn't CPU-read-coherent, so logits use CPU. Others match p_device_inputs_. DeviceInterface* p_device_scoring_{}; // Device for search/scoring (sequences, token allocation). DeviceInterface* p_device_kvcache_{}; // The kvcache is always allocated in device memory (TODO: Remove in favor of just p_device_?) diff --git a/src/models/onnxruntime_api.h b/src/models/onnxruntime_api.h index 17dd637c41..af5888c069 100644 --- a/src/models/onnxruntime_api.h +++ b/src/models/onnxruntime_api.h @@ -527,6 +527,11 @@ struct OrtEnv { OrtEnv& CreateAndRegisterAllocator(const OrtMemoryInfo& mem_info, const OrtArenaCfg& arena_cfg); ///< Wraps OrtApi::CreateAndRegisterAllocator + /// \brief Get an EP-advertised shared allocator matching mem_info (e.g. HOST_ACCESSIBLE), or + /// nullptr if none exists. Wraps OrtApi::GetSharedAllocator. The returned allocator is owned by + /// the OrtEnv — do NOT delete it. + Ort::Allocator* GetSharedAllocator(const OrtMemoryInfo& mem_info) const; + /// \brief Copy tensors between devices. Wraps OrtApi::CopyTensors /// \param src_tensors Array of source OrtValue tensors /// \param dst_tensors Array of destination OrtValue tensors (must be pre-allocated) @@ -827,6 +832,12 @@ struct OrtSession { struct OrtMemoryInfo { static std::unique_ptr CreateCpu(OrtAllocatorType type, OrtMemType mem_type1); static std::unique_ptr Create(const char* name, OrtAllocatorType type, int id, OrtMemType mem_type); + // Wraps CreateMemoryInfo_V2 — lets us request a specific OrtDeviceMemoryType (e.g. + // HOST_ACCESSIBLE) so an EP that registers both DEFAULT and HOST_ACCESSIBLE allocators + // (MIGraphX HipPinned, DML CUSTOM/L0) hands back the host-accessible one. + static std::unique_ptr CreateV2(const char* name, OrtMemoryInfoDeviceType device_type, + uint32_t vendor_id, int32_t device_id, + OrtDeviceMemoryType mem_type, OrtAllocatorType allocator_type); std::string GetAllocatorName() const; OrtAllocatorType GetAllocatorType() const; diff --git a/src/models/onnxruntime_inline.h b/src/models/onnxruntime_inline.h index eeb4ccff54..c9ec2751cd 100644 --- a/src/models/onnxruntime_inline.h +++ b/src/models/onnxruntime_inline.h @@ -256,6 +256,15 @@ inline std::unique_ptr OrtMemoryInfo::Create(const char* name, Or return std::unique_ptr{p}; } +inline std::unique_ptr OrtMemoryInfo::CreateV2(const char* name, OrtMemoryInfoDeviceType device_type, + uint32_t vendor_id, int32_t device_id, + OrtDeviceMemoryType mem_type, OrtAllocatorType allocator_type) { + OrtMemoryInfo* p; + Ort::ThrowOnError(Ort::api->CreateMemoryInfo_V2(name, device_type, vendor_id, device_id, mem_type, + /*alignment=*/0, allocator_type, &p)); + return std::unique_ptr{p}; +} + inline std::unique_ptr OrtSyncStream::Create(const OrtEpDevice* ep_device, const OrtKeyValuePairs* stream_options) { OrtSyncStream* p_stream = nullptr; Ort::ThrowOnError(Ort::api->CreateSyncStreamForEpDevice(ep_device, stream_options, &p_stream)); @@ -433,6 +442,12 @@ inline OrtEnv& OrtEnv::CreateAndRegisterAllocator(const OrtMemoryInfo& mem_info, return *this; } +inline Ort::Allocator* OrtEnv::GetSharedAllocator(const OrtMemoryInfo& mem_info) const { + OrtAllocator* p = nullptr; + Ort::ThrowOnError(Ort::api->GetSharedAllocator(const_cast(this), &mem_info, &p)); + return static_cast(p); // env-owned; may be nullptr if no match +} + inline void OrtEnv::CopyTensors(const std::vector& src_tensors, const std::vector& dst_tensors, OrtSyncStream* stream) const { diff --git a/src/models/session_options.cpp b/src/models/session_options.cpp index db3efb98dd..149f6c083e 100644 --- a/src/models/session_options.cpp +++ b/src/models/session_options.cpp @@ -13,6 +13,7 @@ #include "../openvino/session_options.h" #include "../qnn/session_options.h" #include "../ryzenai/session_options.h" +#include "../amdgpu/session_options.h" #include "../vitisai/session_options.h" #include "../webgpu/session_options.h" @@ -164,6 +165,7 @@ DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options, // Dispatch table: maps provider name (as it appears in genai_config.json) to // the corresponding provider-specific AppendExecutionProvider function. static const std::unordered_map append_execution_provider{ + {"AMDGPU", AMDGPUExecutionProvider::AppendExecutionProvider}, {"CPU", CPUAppendExecutionProvider}, {"cuda", CUDAExecutionProvider::AppendExecutionProvider}, {"DML", DMLExecutionProvider::AppendExecutionProvider}, diff --git a/src/smartptrs.h b/src/smartptrs.h index d8ad90f53a..7d6b3ad3cc 100644 --- a/src/smartptrs.h +++ b/src/smartptrs.h @@ -100,6 +100,7 @@ enum struct DeviceType { OpenVINO, NvTensorRtRtx, RyzenAI, + AMDGPU, MAX }; @@ -110,6 +111,12 @@ struct DeviceInterface { virtual void InitOrt(const OrtApi& api, Ort::Allocator& allocator) = 0; virtual Ort::Allocator& GetAllocator() = 0; + // Host-accessible (CPU-writable, GPU-readable) allocator for decode inputs, if the backend and + // machine support it (MIGraphX HipPinned, DML CUSTOM/L0). Null default -> callers keep the + // current device-memory path. Set via InitHostAccessible after the EP allocator is created. + virtual Ort::Allocator* GetHostAccessibleAllocator() { return nullptr; } + virtual void InitHostAccessible(Ort::Allocator& /*allocator*/) {} + template DeviceSpan Allocate(size_t count) { return DeviceSpan(AllocateBase(sizeof(T) * count)); } virtual std::shared_ptr AllocateBase(size_t size) = 0;