From 85ff8fd24a8035fd5499323f6d61daf2ae001e47 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Tue, 21 Jul 2026 11:55:29 -0700 Subject: [PATCH 01/11] Add AMDGPU execution provider with GPU-resident KV cache Add the AMDGPU execution provider to ONNX Runtime GenAI. The AMDGPU EP resolves a profile to a concrete backend (MIGraphX or DML) at runtime; OGA drives it as a single DeviceType::AMDGPU. Provider naming: exposed as "amdgpu"; OGA also accepts "AMDGPUExecutionProvider" (the catalog form used by the AMD-shipped Windows ML EP MSIX) so test harnesses that match config strings against WinML-discovered names work without bypass hacks. Both normalize to "AMDGPU". GPU-resident KV cache: a GPU-resident DeviceInterface keeps the KV cache on the device (no per-token CPU-to-GPU roundtrip), with backend-agnostic opaque DeviceBuffer copies that dispatch to the active backend. Static-shape prefill: emit ep.migraphx.static_pad_* and hip_graph_enable session-config entries so the EP pads the prefill token axis and reuses a captured graph. DML ignores the migraphx-namespaced keys. Known limitations: - Beam search not supported (needs past_present_share_buffer=true, which requires num_beams=1) --- cmake/global_variables.cmake | 2 + src/amdgpu/interface.cpp | 294 ++++++++++++++++++++++++++++++++ src/amdgpu/interface.h | 17 ++ src/amdgpu/session_options.cpp | 46 +++++ src/amdgpu/session_options.h | 16 ++ src/config.cpp | 6 + src/generators.cpp | 7 + src/models/kv_cache.cpp | 5 +- src/models/model.cpp | 3 +- src/models/onnxruntime_api.h | 11 ++ src/models/onnxruntime_inline.h | 15 ++ src/models/session_options.cpp | 2 + src/smartptrs.h | 7 + 13 files changed, 429 insertions(+), 2 deletions(-) create mode 100644 src/amdgpu/interface.cpp create mode 100644 src/amdgpu/interface.h create mode 100644 src/amdgpu/session_options.cpp create mode 100644 src/amdgpu/session_options.h diff --git a/cmake/global_variables.cmake b/cmake/global_variables.cmake index 4c9ee32be1..8cd033163f 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..93ac5166f1 --- /dev/null +++ b/src/amdgpu/interface.cpp @@ -0,0 +1,294 @@ +// 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 {} + + 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..d6a8bc9e48 --- /dev/null +++ b/src/amdgpu/session_options.cpp @@ -0,0 +1,46 @@ +// 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 "../models/session_options.h" +#include "interface.h" + +namespace Generators::AMDGPUExecutionProvider { + +namespace { + +// 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"); +} + +} // namespace + +DeviceInterface* AppendExecutionProvider(OrtSessionOptions& session_options, + const Config::ProviderOptions& provider_options, + const Config& config, + bool /*disable_graph_capture*/) { + SetStaticPaddingConfig(session_options, config); + + AppendExecutionProviderV2(session_options, provider_options, + DeviceType::AMDGPU, "AMDGPUExecutionProvider"); + + 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 502ad529d1..f3b8bdad47 100644 --- a/src/config.cpp +++ b/src/config.cpp @@ -45,6 +45,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 } @@ -1523,6 +1527,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 8a8fe3f586..4056ccec7b 100644 --- a/src/generators.cpp +++ b/src/generators.cpp @@ -25,6 +25,7 @@ #include "webgpu/interface.h" #include "openvino/interface.h" #include "ryzenai/interface.h" +#include "amdgpu/interface.h" #include "engine/engine.h" #if defined(_WIN32) @@ -360,6 +361,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()); @@ -388,6 +392,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"); } @@ -848,6 +854,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/models/kv_cache.cpp b/src/models/kv_cache.cpp index 3bea7b6ace..455ebebbf5 100644 --- a/src/models/kv_cache.cpp +++ b/src/models/kv_cache.cpp @@ -508,7 +508,10 @@ DefaultKeyValueCache::DefaultKeyValueCache(State& state) } presents_.push_back(OrtValue::CreateTensor(Allocator(), tensor_shape, type_)); - if (Device().GetType() != DeviceType::WEBGPU) { + // WebGPU has no Zero() implementation; AMDGPU skips it for the same reason (Stage A + // DeviceInterface throws on Zero — KV is fresh per-Generator so zeroing is optional). + if (Device().GetType() != DeviceType::WEBGPU && + Device().GetType() != DeviceType::AMDGPU) { ByteWrapTensor(Device(), *presents_.back()).Zero(); } } diff --git a/src/models/model.cpp b/src/models/model.cpp index 65d1dd1e3e..70966baf46 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -35,6 +35,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" @@ -459,7 +460,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 diff --git a/src/models/onnxruntime_api.h b/src/models/onnxruntime_api.h index ec356249db..5ef8037b05 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) @@ -826,6 +831,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 2d68126f62..a873152008 100644 --- a/src/models/onnxruntime_inline.h +++ b/src/models/onnxruntime_inline.h @@ -260,6 +260,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)); @@ -437,6 +446,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..613a939a03 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" @@ -171,6 +172,7 @@ DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options, {"OpenVINO", OpenVINOExecutionProvider::AppendExecutionProvider}, {"RyzenAI", RyzenAIExecutionProvider::AppendExecutionProvider}, {"QNN", QNNExecutionProvider::AppendExecutionProvider}, + {"AMDGPU", AMDGPUExecutionProvider::AppendExecutionProvider}, {"VitisAI", VitisAIExecutionProvider::AppendExecutionProvider}, {"WebGPU", WebGPUExecutionProvider::AppendExecutionProvider}, }; diff --git a/src/smartptrs.h b/src/smartptrs.h index d9dc63338f..e28127dbe5 100644 --- a/src/smartptrs.h +++ b/src/smartptrs.h @@ -100,6 +100,7 @@ enum struct DeviceType { OpenVINO, NvTensorRtRtx, RyzenAI, + AMDGPU, MAX }; @@ -111,6 +112,12 @@ struct DeviceInterface { virtual Ort::Allocator& GetAllocator() = 0; virtual std::unique_ptr GetMemoryInfo() const = 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; From 976aacfd6e9a62f186446d67304afeb92072c476 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Tue, 21 Jul 2026 11:56:10 -0700 Subject: [PATCH 02/11] AMDGPU EP: host-accessible decode inputs Route the small decode inputs (input_ids/position_ids/attention_mask) through a host-accessible (CPU-writable, GPU-readable) allocator so the CPU updates them in place with no per-step copy. Resolved via GetSharedAllocator; KV cache and scoring stay on the default device interface, and the path falls back to default inputs if no host-accessible allocator is available. Single-GPU only for now (device_id 0). --- src/generators.h | 3 +++ src/models/model.cpp | 27 +++++++++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/src/generators.h b/src/generators.h index d606420152..b90df69db5 100644 --- a/src/generators.h +++ b/src/generators.h @@ -183,6 +183,9 @@ 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_{}; }; Allocator device_allocators_[static_cast(DeviceType::MAX)]; diff --git a/src/models/model.cpp b/src/models/model.cpp index 70966baf46..aa2ad7bdeb 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -503,6 +503,24 @@ void EnsureDeviceOrtInit(DeviceInterface& device, const Config& config) { throw std::runtime_error("Unexpected failure to create device memory allocator for " + to_string(type)); } 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 { + // device_id is hardcoded 0, so this is correct for single-GPU only. On multi-GPU the + // device_id differs and GetSharedAllocator returns null (falls back, perf loss not error). + // TODO(multi-gpu): query EpDevice::MemoryInfo(HOST_ACCESSIBLE) instead of reconstructing. + auto host_info = OrtMemoryInfo::CreateV2("pinned", OrtMemoryInfoDeviceType_GPU, + /*vendor_id=*/0x1002, /*device_id=*/0, + OrtDeviceMemoryType_HOST_ACCESSIBLE, OrtDeviceAllocator); + 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) { @@ -596,6 +614,15 @@ 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; + } + // 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 || From 01f882f44f37eb9542373de18742c34c28401eae Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Tue, 21 Jul 2026 11:56:37 -0700 Subject: [PATCH 03/11] AMDGPU EP: route logits off the host-accessible allocator Logits is GPU-written and CPU-read (the sampler), the opposite of the pinned decode inputs. On AMDGPU the inputs use a host-accessible allocator whose heap is not CPU-read-coherent, so reading logits from it returns stale data. Route logits to the CPU interface instead via a new p_logits_ member; only the decode inputs stay on the host-accessible allocator. --- src/models/logits.cpp | 18 ++++++++++++------ src/models/logits.h | 6 ++++++ 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/src/models/logits.cpp b/src/models/logits.cpp index 32d1992c43..e2d5c9a1d5 100644 --- a/src/models/logits.cpp +++ b/src/models/logits.cpp @@ -13,7 +13,13 @@ 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_); + // Logits is GPU-written / CPU-read (the sampler). For AMDGPU, p_device_inputs_ points at the + // host-accessible allocator (DML WRITE_COMBINE is not CPU-read-coherent), so route logits to the + // CPU interface instead; other backends keep logits on their device-inputs interface. + p_logits_ = (model_.p_device_->GetType() == DeviceType::AMDGPU) + ? GetDeviceInterface(DeviceType::CPU) + : model_.p_device_inputs_; + output_raw_ = std::make_unique(p_logits_, type_); input_sequence_lengths.resize(state_.params_->search.batch_size); @@ -42,10 +48,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(p_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(p_logits_->GetAllocator(), shape_); logits_of_last_token = output_last_tokens_.get(); @@ -53,7 +59,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(*p_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 +77,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_, *p_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(*p_logits_, *logits_of_last_token); return logits_; } diff --git a/src/models/logits.h b/src/models/logits.h index 82c433180a..4c6a4b20b4 100644 --- a/src/models/logits.h +++ b/src/models/logits.h @@ -31,6 +31,12 @@ struct Logits { std::unique_ptr output_raw_; // Raw logits output from model + // Device interface for the logits output + last-token/fp32 scratch. Normally p_device_inputs_, but + // logits is GPU-written / CPU-read (opposite of the pinned decode inputs). For AMDGPU the inputs use + // a host-accessible allocator whose heap is NOT CPU-read-coherent (DML CUSTOM/L0/WRITE_COMBINE), so + // logits go on the CPU interface instead; only the 3 decode inputs use the pinned allocator. + DeviceInterface* p_logits_{}; + std::vector input_sequence_lengths; // OrtValue wrapped in a DeviceMemory object to make it universal DeviceSpan logits_; From d621a68ed22a6589e91f81b5f66bbe7f4915b559 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Tue, 21 Jul 2026 21:52:36 -0700 Subject: [PATCH 04/11] AMDGPU: alphabetize provider dispatch entry (review #2165) --- src/models/session_options.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/session_options.cpp b/src/models/session_options.cpp index 613a939a03..149f6c083e 100644 --- a/src/models/session_options.cpp +++ b/src/models/session_options.cpp @@ -165,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}, @@ -172,7 +173,6 @@ DeviceInterface* SetProviderSessionOptions(OrtSessionOptions& session_options, {"OpenVINO", OpenVINOExecutionProvider::AppendExecutionProvider}, {"RyzenAI", RyzenAIExecutionProvider::AppendExecutionProvider}, {"QNN", QNNExecutionProvider::AppendExecutionProvider}, - {"AMDGPU", AMDGPUExecutionProvider::AppendExecutionProvider}, {"VitisAI", VitisAIExecutionProvider::AppendExecutionProvider}, {"WebGPU", WebGPUExecutionProvider::AppendExecutionProvider}, }; From 9a2337ff2d018635191e9df302727a1a6b6a4607 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Tue, 21 Jul 2026 22:06:37 -0700 Subject: [PATCH 05/11] AMDGPU: move p_device_logits_ into Model device-family (review #2165) --- src/models/logits.cpp | 18 ++++++------------ src/models/logits.h | 6 ------ src/models/model.cpp | 3 +++ src/models/model.h | 1 + 4 files changed, 10 insertions(+), 18 deletions(-) diff --git a/src/models/logits.cpp b/src/models/logits.cpp index e2d5c9a1d5..772486b6e7 100644 --- a/src/models/logits.cpp +++ b/src/models/logits.cpp @@ -13,13 +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)} { - // Logits is GPU-written / CPU-read (the sampler). For AMDGPU, p_device_inputs_ points at the - // host-accessible allocator (DML WRITE_COMBINE is not CPU-read-coherent), so route logits to the - // CPU interface instead; other backends keep logits on their device-inputs interface. - p_logits_ = (model_.p_device_->GetType() == DeviceType::AMDGPU) - ? GetDeviceInterface(DeviceType::CPU) - : model_.p_device_inputs_; - output_raw_ = std::make_unique(p_logits_, type_); + output_raw_ = std::make_unique(model_.p_device_logits_, type_); input_sequence_lengths.resize(state_.params_->search.batch_size); @@ -48,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(p_logits_->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(p_logits_->GetAllocator(), shape_); + logits_of_last_token_fp32_ = OrtValue::CreateTensor(model_.p_device_logits_->GetAllocator(), shape_); logits_of_last_token = output_last_tokens_.get(); @@ -59,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(*p_logits_, *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 @@ -77,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_, *p_logits_, 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(*p_logits_, *logits_of_last_token); + logits_ = WrapTensor(*model_.p_device_logits_, *logits_of_last_token); return logits_; } diff --git a/src/models/logits.h b/src/models/logits.h index 4c6a4b20b4..82c433180a 100644 --- a/src/models/logits.h +++ b/src/models/logits.h @@ -31,12 +31,6 @@ struct Logits { std::unique_ptr output_raw_; // Raw logits output from model - // Device interface for the logits output + last-token/fp32 scratch. Normally p_device_inputs_, but - // logits is GPU-written / CPU-read (opposite of the pinned decode inputs). For AMDGPU the inputs use - // a host-accessible allocator whose heap is NOT CPU-read-coherent (DML CUSTOM/L0/WRITE_COMBINE), so - // logits go on the CPU interface instead; only the 3 decode inputs use the pinned allocator. - DeviceInterface* p_logits_{}; - std::vector input_sequence_lengths; // OrtValue wrapped in a DeviceMemory object to make it universal DeviceSpan logits_; diff --git a/src/models/model.cpp b/src/models/model.cpp index aa2ad7bdeb..aa32b0adf8 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -623,6 +623,9 @@ Model::Model(std::unique_ptr config) : config_{std::move(config)} { 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 24ddca224e..06b8264dcd 100644 --- a/src/models/model.h +++ b/src/models/model.h @@ -188,6 +188,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_?) From 4598f2df1f8f4f2b6da524db0b4673aeabf8aa1b Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Wed, 22 Jul 2026 02:33:39 -0700 Subject: [PATCH 06/11] AMDGPU: zero-init KV cache (drop stale skip) (review #2165) --- src/models/kv_cache.cpp | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/models/kv_cache.cpp b/src/models/kv_cache.cpp index 455ebebbf5..b382e50ee2 100644 --- a/src/models/kv_cache.cpp +++ b/src/models/kv_cache.cpp @@ -508,10 +508,8 @@ DefaultKeyValueCache::DefaultKeyValueCache(State& state) } presents_.push_back(OrtValue::CreateTensor(Allocator(), tensor_shape, type_)); - // WebGPU has no Zero() implementation; AMDGPU skips it for the same reason (Stage A - // DeviceInterface throws on Zero — KV is fresh per-Generator so zeroing is optional). - if (Device().GetType() != DeviceType::WEBGPU && - Device().GetType() != DeviceType::AMDGPU) { + // WebGPU has no Zero() implementation; every other backend (incl. AMDGPU) zero-inits the KV. + if (Device().GetType() != DeviceType::WEBGPU) { ByteWrapTensor(Device(), *presents_.back()).Zero(); } } From a82b78021e52724322d5f35f37a146be62c50fe4 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Wed, 22 Jul 2026 02:33:40 -0700 Subject: [PATCH 07/11] AMDGPU: source host-accessible mem-info from EpDevice metadata (review #2165) --- src/models/model.cpp | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/src/models/model.cpp b/src/models/model.cpp index aa32b0adf8..0455da9499 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -508,13 +508,19 @@ void EnsureDeviceOrtInit(DeviceInterface& device, const Config& config) { // if unavailable, creation returns null and callers fall back to the default device inputs path. if (!allocator.host_accessible_allocator_ && type == DeviceType::AMDGPU) { try { - // device_id is hardcoded 0, so this is correct for single-GPU only. On multi-GPU the - // device_id differs and GetSharedAllocator returns null (falls back, perf loss not error). - // TODO(multi-gpu): query EpDevice::MemoryInfo(HOST_ACCESSIBLE) instead of reconstructing. - auto host_info = OrtMemoryInfo::CreateV2("pinned", OrtMemoryInfoDeviceType_GPU, - /*vendor_id=*/0x1002, /*device_id=*/0, - OrtDeviceMemoryType_HOST_ACCESSIBLE, OrtDeviceAllocator); - allocator.host_accessible_allocator_ = GetOrtEnv().GetSharedAllocator(*host_info); + // Query the AMDGPU EP's advertised HOST_ACCESSIBLE memory-info instead of reconstructing it with + // hardcoded vendor/device ids, so it carries the real ids for this machine (multi-GPU safe). If + // none is advertised, leave it unset -> callers fall back to the default device inputs path. + const OrtMemoryInfo* host_info = nullptr; + for (const OrtEpDevice* ep_device : FindRegisteredEpDevices("AMDGPUExecutionProvider")) { + if (const OrtMemoryInfo* mi = + Ort::api->EpDevice_MemoryInfo(ep_device, OrtDeviceMemoryType_HOST_ACCESSIBLE)) { + host_info = mi; + break; + } + } + if (host_info) + allocator.host_accessible_allocator_ = GetOrtEnv().GetSharedAllocator(*host_info); } catch (const Ort::Exception&) { allocator.host_accessible_allocator_ = nullptr; } From 8d742c1fba4078139238a6c6256e9d308309b054 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Wed, 22 Jul 2026 13:59:24 -0700 Subject: [PATCH 08/11] AMDGPU: bind compute + host-accessible allocators to the selected device (review #2165) Resolve the AMDGPU device id from the filtered EP device instead of hardcoding 0. The EP keys its allocator on this id, so a hardcoded 0 pinned compute to device 0 regardless of the selected device. Correlate the host-accessible pool to the same id so pinned decode inputs live on the device that runs the model. Single-GPU resolves to id 0, unchanged. --- src/amdgpu/interface.cpp | 14 ++++++++++++++ src/amdgpu/interface.h | 4 ++++ src/generators.h | 1 + src/models/model.cpp | 24 +++++++++++++++++++----- 4 files changed, 38 insertions(+), 5 deletions(-) diff --git a/src/amdgpu/interface.cpp b/src/amdgpu/interface.cpp index 93ac5166f1..e310970f58 100644 --- a/src/amdgpu/interface.cpp +++ b/src/amdgpu/interface.cpp @@ -194,6 +194,14 @@ struct InterfaceImpl : DeviceInterface { return *ort_allocator_; } + // The MIGraphX backend registers its device memory-info under "Hip", keyed on the device id. + std::unique_ptr GetMemoryInfo() const override { + return OrtMemoryInfo::Create("Hip", OrtAllocatorType::OrtDeviceAllocator, + device_id_, OrtMemType::OrtMemTypeDefault); + } + + void SetDeviceId(int device_id) { device_id_ = device_id; } + void InitHostAccessible(Ort::Allocator& allocator) override { ort_pinned_allocator_ = &allocator; } @@ -227,6 +235,7 @@ struct InterfaceImpl : DeviceInterface { const OrtMemoryInfo* ort_memory_info_{}; // Host-accessible allocator, set by InitHostAccessible when one is available. Ort::Allocator* ort_pinned_allocator_{}; + int device_id_{}; }; // Inputs-only interface: allocations come from the host-accessible allocator, everything else @@ -238,6 +247,7 @@ struct PinnedInputsImpl : DeviceInterface { 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::unique_ptr GetMemoryInfo() const override { return base_.GetMemoryInfo(); } std::shared_ptr AllocateBase(size_t size) override { return std::make_shared(size, base_.PinnedAllocator()); @@ -282,6 +292,10 @@ DeviceInterface* GetAMDGPUInterface() { return g_amdgpu_device.get(); } +void SetAMDGPUDeviceId(int device_id) { + static_cast(GetAMDGPUInterface())->SetDeviceId(device_id); +} + DeviceInterface* GetAMDGPUPinnedInputsInterface() { auto* base = static_cast(GetAMDGPUInterface()); if (!base->GetHostAccessibleAllocator()) diff --git a/src/amdgpu/interface.h b/src/amdgpu/interface.h index 0957f9a1ff..bc3f33d76c 100644 --- a/src/amdgpu/interface.h +++ b/src/amdgpu/interface.h @@ -9,6 +9,10 @@ namespace Generators { DeviceInterface* GetAMDGPUInterface(); +// Device this interface's allocators bind to. Resolved from the selected EP device before the +// allocator is created; defaults to 0. +void SetAMDGPUDeviceId(int device_id); + // 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. diff --git a/src/generators.h b/src/generators.h index b90df69db5..275b903531 100644 --- a/src/generators.h +++ b/src/generators.h @@ -186,6 +186,7 @@ struct OrtGlobals { // 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/model.cpp b/src/models/model.cpp index 0455da9499..526aedc704 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -492,6 +492,17 @@ 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()); + // 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_)); + } + SetAMDGPUDeviceId(allocator.device_id_); + } try { auto memory_info = device.GetMemoryInfo(); allocator.allocator_ = Ort::Allocator::Create(*allocator.session_, *memory_info); @@ -508,13 +519,16 @@ void EnsureDeviceOrtInit(DeviceInterface& device, const Config& config) { // if unavailable, creation returns null and callers fall back to the default device inputs path. if (!allocator.host_accessible_allocator_ && type == DeviceType::AMDGPU) { try { - // Query the AMDGPU EP's advertised HOST_ACCESSIBLE memory-info instead of reconstructing it with - // hardcoded vendor/device ids, so it carries the real ids for this machine (multi-GPU safe). If - // none is advertised, leave it unset -> callers fall back to the default device inputs path. + // 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")) { - if (const OrtMemoryInfo* mi = - Ort::api->EpDevice_MemoryInfo(ep_device, OrtDeviceMemoryType_HOST_ACCESSIBLE)) { + 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; } From a1a5be7fdd491852e3a7a1a55bdff454e89a7e00 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Wed, 22 Jul 2026 18:30:30 -0700 Subject: [PATCH 09/11] AMDGPU: forward model_arch to the umbrella EP for backend routing The umbrella EP routes backends by model architecture but OGA never sent it, so every model routed as non-LLM. Emit config.model.type as the ep.amdgpuexecutionprovider.model_arch provider option alongside the existing static-padding hints. --- src/amdgpu/session_options.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/amdgpu/session_options.cpp b/src/amdgpu/session_options.cpp index d6a8bc9e48..9bcfaa9800 100644 --- a/src/amdgpu/session_options.cpp +++ b/src/amdgpu/session_options.cpp @@ -37,6 +37,9 @@ DeviceInterface* AppendExecutionProvider(OrtSessionOptions& session_options, 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()); + AppendExecutionProviderV2(session_options, provider_options, DeviceType::AMDGPU, "AMDGPUExecutionProvider"); From 85cfef232697e93aade26abb283afd3cf24c4a06 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Wed, 22 Jul 2026 18:58:48 -0700 Subject: [PATCH 10/11] AMDGPU: fix clang-format violations Trailing-comment spacing and argument-continuation alignment flagged by the lint-cpp CI check (clang-format 20.1.0). Formatting only, no behavior change. --- src/amdgpu/interface.cpp | 7 +++---- src/models/onnxruntime_inline.h | 4 ++-- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/src/amdgpu/interface.cpp b/src/amdgpu/interface.cpp index e310970f58..cc167b8405 100644 --- a/src/amdgpu/interface.cpp +++ b/src/amdgpu/interface.cpp @@ -164,8 +164,8 @@ struct PinnedMemory final : DeviceBuffer { } 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 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 { @@ -274,8 +274,7 @@ struct PinnedInputsImpl : DeviceInterface { 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); + 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_; diff --git a/src/models/onnxruntime_inline.h b/src/models/onnxruntime_inline.h index a873152008..b46336ec03 100644 --- a/src/models/onnxruntime_inline.h +++ b/src/models/onnxruntime_inline.h @@ -261,8 +261,8 @@ inline std::unique_ptr OrtMemoryInfo::Create(const char* name, Or } 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) { + 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)); From 91c1fd5218e9bc2f8c8daec85796de8158e94dc8 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Fri, 24 Jul 2026 17:21:24 -0700 Subject: [PATCH 11/11] AMDGPU: enable DirectML host-accessible decode inputs Emit the ep.directml.enable_host_accessible provider option so the DirectML backend uses host-accessible decode inputs. Sits alongside the existing static-padding and model_arch config entries. --- src/amdgpu/session_options.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/amdgpu/session_options.cpp b/src/amdgpu/session_options.cpp index 9bcfaa9800..6e844094bb 100644 --- a/src/amdgpu/session_options.cpp +++ b/src/amdgpu/session_options.cpp @@ -40,6 +40,9 @@ DeviceInterface* AppendExecutionProvider(OrtSessionOptions& session_options, // 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"); + AppendExecutionProviderV2(session_options, provider_options, DeviceType::AMDGPU, "AMDGPUExecutionProvider");