From 7198d7a579ef6c2609b2ec076f887b75418aec6d Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Tue, 21 Jul 2026 11:55:29 -0700 Subject: [PATCH 01/21] 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 fcd1f90c83..735e7ec304 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 } @@ -1543,6 +1547,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 f5f0da3779..f5676d8dd1 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) @@ -367,6 +368,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()); @@ -395,6 +399,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"); } @@ -855,6 +861,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 0e380ad556..bd58dfc6df 100644 --- a/src/models/kv_cache.cpp +++ b/src/models/kv_cache.cpp @@ -517,7 +517,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 f18e0fe623..154e9600c4 100644 --- a/src/models/onnxruntime_api.h +++ b/src/models/onnxruntime_api.h @@ -572,6 +572,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) @@ -871,6 +876,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 0e3ba831af..4a276fe4b0 100644 --- a/src/smartptrs.h +++ b/src/smartptrs.h @@ -126,6 +126,7 @@ enum struct DeviceType { OpenVINO, NvTensorRtRtx, RyzenAI, + AMDGPU, MAX }; @@ -137,6 +138,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 0a7138ac65b80bea89e35a4640f3942841167f5f Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Tue, 21 Jul 2026 11:56:10 -0700 Subject: [PATCH 02/21] 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 a7d61a8554dd4d76de03bca3a547f1571794a9ba Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Tue, 21 Jul 2026 11:56:37 -0700 Subject: [PATCH 03/21] 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 67877250e0d9e423a0ec5ff1bbe81a005207841f Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Tue, 21 Jul 2026 21:52:36 -0700 Subject: [PATCH 04/21] 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 da43039cd01df7b42d85cedd6bbad2a93d2e8926 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Tue, 21 Jul 2026 22:06:37 -0700 Subject: [PATCH 05/21] 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 6b1c0557950ee59e702ccc258bca8a306e0a95ed Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Wed, 22 Jul 2026 02:33:39 -0700 Subject: [PATCH 06/21] 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 bd58dfc6df..5b39c35029 100644 --- a/src/models/kv_cache.cpp +++ b/src/models/kv_cache.cpp @@ -517,10 +517,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 f498c698a8b33339206090c04b1c0756f2a1d845 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Wed, 22 Jul 2026 02:33:40 -0700 Subject: [PATCH 07/21] 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 6d3a782714057e5b167749c284143b2ca228f820 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Wed, 22 Jul 2026 13:59:24 -0700 Subject: [PATCH 08/21] 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 093ac7427c82df6ba25e8f865969de45297b4764 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Wed, 22 Jul 2026 18:30:30 -0700 Subject: [PATCH 09/21] 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 e8ce62c8a482931a8ed879219abb532b25376ba4 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Wed, 22 Jul 2026 18:58:48 -0700 Subject: [PATCH 10/21] 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 099c277de9672a5939b27713563382f2a9f53b34 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Fri, 24 Jul 2026 17:21:24 -0700 Subject: [PATCH 11/21] 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"); From 03e49734cfce156ad53280906e0aeb646192607b Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Thu, 6 Aug 2026 16:15:56 -0700 Subject: [PATCH 12/21] Drop the AMDGPU mention from the KV zero-init comment (review #2165) Applies the reviewer's suggested wording: the exclusion is about WebGPU, so naming the backends that do zero-init adds nothing. --- src/models/kv_cache.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/kv_cache.cpp b/src/models/kv_cache.cpp index 5b39c35029..1da26da820 100644 --- a/src/models/kv_cache.cpp +++ b/src/models/kv_cache.cpp @@ -517,7 +517,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. + // WebGPU has no Zero() implementation; every other backend zero-inits the KV. if (Device().GetType() != DeviceType::WEBGPU) { ByteWrapTensor(Device(), *presents_.back()).Zero(); } From 86b64ff8a0c7db0671f28f886178e2141d2dc1dd Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Thu, 6 Aug 2026 16:16:58 -0700 Subject: [PATCH 13/21] Fold the host-accessible inputs case into the p_device_inputs_ chain (review #2165) The host-accessible interface was selected in a separate if that reassigned p_device_inputs_ after the chain had already set it. It is now an else if in that chain, so the field is assigned once. The interface getter already returns null when no host-accessible allocator exists, so the extra guard on it went away with the restructure. --- src/models/model.cpp | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/src/models/model.cpp b/src/models/model.cpp index 526aedc704..c1d9b2352a 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -625,25 +625,23 @@ Model::Model(std::unique_ptr config) : config_{std::move(config)} { CreateSessionOptions(); EnsureDeviceOrtInit(*p_device_, *config_); + // Inputs-only interface backed by a host-accessible allocation, so the CPU updates the small + // decode inputs in place with no per-step roundtrip. Null if the device offers no such allocator. + DeviceInterface* p_host_accessible_inputs = + p_device_->GetType() == DeviceType::AMDGPU ? GetAMDGPUPinnedInputsInterface() : nullptr; + // Only CUDA, TRT-RTX, RyzenAI and DML does every input on the device // For WebGPU, use device memory only if graph capture is enabled, otherwise use CPU if (p_device_->GetType() == DeviceType::CUDA || p_device_->GetType() == DeviceType::DML || p_device_->GetType() == DeviceType::NvTensorRtRtx || p_device_->GetType() == DeviceType::RyzenAI || (p_device_->GetType() == DeviceType::WEBGPU && IsGraphCaptureEnabled(config_->model.decoder.session_options))) p_device_inputs_ = p_device_; + else if (p_host_accessible_inputs) + p_device_inputs_ = p_host_accessible_inputs; 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. + // Logits are read back on the CPU every step, which is slow from a host-accessible allocation. 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, From 8e06f0173c76ba47c64d072a0a82a77f9b381f69 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Thu, 6 Aug 2026 16:17:26 -0700 Subject: [PATCH 14/21] Move the AMDGPU allocator setup into the device interface (review #2165) EnsureDeviceOrtInit carried the AMD-specific device-id lookup and the host-accessible allocator acquisition inline, including the provider name as a literal. Both now sit behind DeviceInterface virtuals with no-op defaults, following the existing ShapeInitSessionProviderOptions pattern, so the shared path no longer branches on the device: GetDeviceId - id the allocators bind to, 0 unless the device resolves one from EP metadata InitDeviceAllocators - lets a device set up any additional allocators it offers once the device allocator exists The AMDGPU implementations move to src/amdgpu/interface.cpp, which also drops the SetAMDGPUDeviceId free function that existed only to reach back into the interface from the shared path. --- src/amdgpu/interface.cpp | 50 ++++++++++++++++++++++++++++++++-------- src/amdgpu/interface.h | 4 ---- src/models/model.cpp | 46 ++++++------------------------------ src/smartptrs.h | 15 ++++++++---- 4 files changed, 59 insertions(+), 56 deletions(-) diff --git a/src/amdgpu/interface.cpp b/src/amdgpu/interface.cpp index cc167b8405..6d34ed9a5b 100644 --- a/src/amdgpu/interface.cpp +++ b/src/amdgpu/interface.cpp @@ -11,6 +11,7 @@ #include "../generators.h" #include "../search.h" +#include "../models/session_options.h" #include "interface.h" #include @@ -21,6 +22,9 @@ namespace AMDGPU { const char* device_label = "amdgpu"; const char* label_cpu = "cpu"; +// Registration name the EP is discovered under, used to look up its advertised memory-info. +constexpr const char* kExecutionProviderName = "AMDGPUExecutionProvider"; + 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} { @@ -194,16 +198,48 @@ struct InterfaceImpl : DeviceInterface { return *ort_allocator_; } - // The MIGraphX backend registers its device memory-info under "Hip", keyed on the device id. + // The 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; } + // Read the id of the EP device the model will run on, so the allocators bind to it rather than + // assuming device 0. Resolves to 0 on a single-GPU machine. + int GetDeviceId(const ProviderOptions* user_options) override { + auto ep_devices = FindRegisteredEpDevices(kExecutionProviderName); + if (user_options) + ep_devices = ApplyDeviceFiltering(*user_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, &device_id_)); + } + return device_id_; + } - void InitHostAccessible(Ort::Allocator& allocator) override { - ort_pinned_allocator_ = &allocator; + // Pick up the host-accessible allocator advertised on the same device as the compute allocator, + // so the pinned decode inputs live on the device that runs the model. If the EP advertises none, + // this leaves the allocator null and callers keep the default device-memory path. + void InitDeviceAllocators(const ProviderOptions* /*user_options*/, int device_id) override { + if (ort_pinned_allocator_) + return; + try { + for (const OrtEpDevice* ep_device : FindRegisteredEpDevices(kExecutionProviderName)) { + 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 == device_id) { + ort_pinned_allocator_ = GetOrtEnv().GetSharedAllocator(*mi); + break; + } + } + } catch (const Ort::Exception&) { + ort_pinned_allocator_ = nullptr; + } } Ort::Allocator* GetHostAccessibleAllocator() override { @@ -233,7 +269,7 @@ struct InterfaceImpl : DeviceInterface { private: Ort::Allocator* ort_allocator_{}; const OrtMemoryInfo* ort_memory_info_{}; - // Host-accessible allocator, set by InitHostAccessible when one is available. + // Host-accessible allocator, set by InitDeviceAllocators when the EP advertises one. Ort::Allocator* ort_pinned_allocator_{}; int device_id_{}; }; @@ -291,10 +327,6 @@ 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 bc3f33d76c..0957f9a1ff 100644 --- a/src/amdgpu/interface.h +++ b/src/amdgpu/interface.h @@ -9,10 +9,6 @@ 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/models/model.cpp b/src/models/model.cpp index c1d9b2352a..0a4a014640 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -492,17 +492,8 @@ 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_); - } + // Bind the allocator to the selected device rather than assuming device 0. + allocator.device_id_ = device.GetDeviceId(user_provider_options); try { auto memory_info = device.GetMemoryInfo(); allocator.allocator_ = Ort::Allocator::Create(*allocator.session_, *memory_info); @@ -515,32 +506,10 @@ void EnsureDeviceOrtInit(DeviceInterface& device, const Config& config) { } 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_); - } + // Let the device set up any additional allocators it offers (e.g. host-accessible memory for + // decode inputs). Devices that offer none leave the defaults in place. + device.InitDeviceAllocators(user_provider_options, allocator.device_id_); + allocator.host_accessible_allocator_ = device.GetHostAccessibleAllocator(); } void SessionInfo::Add(OrtSession& session) { @@ -627,8 +596,7 @@ Model::Model(std::unique_ptr config) : config_{std::move(config)} { // Inputs-only interface backed by a host-accessible allocation, so the CPU updates the small // decode inputs in place with no per-step roundtrip. Null if the device offers no such allocator. - DeviceInterface* p_host_accessible_inputs = - p_device_->GetType() == DeviceType::AMDGPU ? GetAMDGPUPinnedInputsInterface() : nullptr; + DeviceInterface* p_host_accessible_inputs = GetAMDGPUPinnedInputsInterface(); // Only CUDA, TRT-RTX, RyzenAI and DML does every input on the device // For WebGPU, use device memory only if graph capture is enabled, otherwise use CPU diff --git a/src/smartptrs.h b/src/smartptrs.h index 4a276fe4b0..e6819dc781 100644 --- a/src/smartptrs.h +++ b/src/smartptrs.h @@ -138,11 +138,18 @@ 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. + // Host-accessible (CPU-writable, GPU-readable) allocator for decode inputs, if the device + // supports it. Null default -> callers keep the current device-memory path. virtual Ort::Allocator* GetHostAccessibleAllocator() { return nullptr; } - virtual void InitHostAccessible(Ort::Allocator& /*allocator*/) {} + + // Called once after the device allocator is created, so a device that offers additional + // allocators (e.g. host-accessible memory) can set them up. The default sets up nothing. + // `device_id` is the id the device allocator was created on. + virtual void InitDeviceAllocators(const ProviderOptions* /*user_options*/, int /*device_id*/) {} + + // Id of the EP device this interface's allocators should bind to. 0 unless the device resolves a + // specific one from EP metadata. + virtual int GetDeviceId(const ProviderOptions* /*user_options*/) { return 0; } template DeviceSpan Allocate(size_t count) { return DeviceSpan(AllocateBase(sizeof(T) * count)); } From 07d1b30c98c73d4eed1ae3ada21bef5eb3b91734 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Thu, 6 Aug 2026 16:17:34 -0700 Subject: [PATCH 15/21] Drop vendor-specific detail from shared-code comments (review #2165) Comments in model.cpp, model.h, smartptrs.h, onnxruntime_api.h and generators.cpp named AMD backends while describing device-agnostic code. The behaviour they document is not vendor-specific, so the names are dropped and the wording shortened. --- src/generators.cpp | 2 +- src/models/model.h | 2 +- src/models/onnxruntime_api.h | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/generators.cpp b/src/generators.cpp index f5676d8dd1..db611979e3 100644 --- a/src/generators.cpp +++ b/src/generators.cpp @@ -369,7 +369,7 @@ DeviceInterface* OrtGlobals::GetDeviceInterface(DeviceType type) { slot = owned_interfaces_.back().get(); break; case DeviceType::AMDGPU: - slot = GetAMDGPUInterface(); // static singleton owned by amdgpu/interface.cpp + slot = GetAMDGPUInterface(); break; case DeviceType::CPU: default: diff --git a/src/models/model.h b/src/models/model.h index 06b8264dcd..0dc648a9a4 100644 --- a/src/models/model.h +++ b/src/models/model.h @@ -188,7 +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_logits_{}; // Logits are read back on the CPU every step 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 154e9600c4..93129409b2 100644 --- a/src/models/onnxruntime_api.h +++ b/src/models/onnxruntime_api.h @@ -878,7 +878,7 @@ struct OrtMemoryInfo { 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. + // 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); From d84dcdbb33debce4470437c21db2c5a60e8deca6 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Thu, 6 Aug 2026 17:10:52 -0700 Subject: [PATCH 16/21] Remove the unused CreateMemoryInfo_V2 wrapper (review #2165) The reviewer asked for this to be an overload of Create rather than a separately named CreateV2. It turned out to have no callers left: the allocator setup now reads the memory-info the EP advertises instead of building one, so the wrapper is dropped rather than renamed. --- src/models/onnxruntime_api.h | 6 ------ src/models/onnxruntime_inline.h | 9 --------- 2 files changed, 15 deletions(-) diff --git a/src/models/onnxruntime_api.h b/src/models/onnxruntime_api.h index 93129409b2..45216d1c70 100644 --- a/src/models/onnxruntime_api.h +++ b/src/models/onnxruntime_api.h @@ -876,12 +876,6 @@ 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 - // 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 b46336ec03..0d63e5f00f 100644 --- a/src/models/onnxruntime_inline.h +++ b/src/models/onnxruntime_inline.h @@ -260,15 +260,6 @@ 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)); From 05297cc8338ecab987b8323681261e3304d3ba17 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Thu, 6 Aug 2026 17:10:52 -0700 Subject: [PATCH 17/21] Drop the memory-type name comment from the AMDGPU interface (review #2165) The name sits next to the interface that uses it, so the comment restates what the code already shows. --- src/amdgpu/interface.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/amdgpu/interface.cpp b/src/amdgpu/interface.cpp index 6d34ed9a5b..af24d2fb0e 100644 --- a/src/amdgpu/interface.cpp +++ b/src/amdgpu/interface.cpp @@ -198,7 +198,6 @@ struct InterfaceImpl : DeviceInterface { return *ort_allocator_; } - // The 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); From 1c7dc11e7a03b7339a8570d958c567161f01e36d Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Thu, 6 Aug 2026 17:45:42 -0700 Subject: [PATCH 18/21] AMDGPU: self-register the plugin EP library when the host has not A plugin EP is only discoverable once its library is registered on the OrtEnv, and the C model_benchmark has no option to pass a path. Resolve it the way the RyzenAI interface does and register it, skipping entirely when the EP is already registered or the library is not found. Also accept AMDGPU in the benchmark's execution-provider list. --- benchmark/c/options.cpp | 7 ++-- src/amdgpu/interface.cpp | 7 +--- src/amdgpu/interface.h | 3 ++ src/amdgpu/session_options.cpp | 77 +++++++++++++++++++++++++++++++++- 4 files changed, 85 insertions(+), 9 deletions(-) diff --git a/benchmark/c/options.cpp b/benchmark/c/options.cpp index f23e5228e9..b6319087fb 100644 --- a/benchmark/c/options.cpp +++ b/benchmark/c/options.cpp @@ -30,7 +30,7 @@ namespace { << " -i,--input_folder \n" << " Path to the ONNX model directory to benchmark, compatible with onnxruntime-genai.\n" << " -e,--execution_provider \n" - << " Execution provider to use. Valid values are: cpu, cuda, dml, NvTensorRtRtx. Default: " << defaults.execution_provider << "\n" + << " Execution provider to use. Valid values are: cpu, cuda, dml, NvTensorRtRtx, AMDGPU. Default: " << defaults.execution_provider << "\n" << " -b,--batch_size \n" << " Number of sequences to generate in parallel. Default: " << defaults.batch_size << "\n" << " Prompt options:\n" @@ -98,8 +98,9 @@ std::string ReadFileContent(std::string_view file_path) { } void ValidateExecutionProvider(const std::string& provider) { - if (provider != "cpu" && provider != "cuda" && provider != "dml" && provider != "NvTensorRtRtx") { - throw std::runtime_error("Invalid execution provider: " + provider + ". Valid values are: cpu, cuda, dml, NvTensorRtRtx"); + if (provider != "cpu" && provider != "cuda" && provider != "dml" && provider != "NvTensorRtRtx" && + provider != "AMDGPU") { + throw std::runtime_error("Invalid execution provider: " + provider + ". Valid values are: cpu, cuda, dml, NvTensorRtRtx, AMDGPU"); } } diff --git a/src/amdgpu/interface.cpp b/src/amdgpu/interface.cpp index af24d2fb0e..916e47c189 100644 --- a/src/amdgpu/interface.cpp +++ b/src/amdgpu/interface.cpp @@ -22,9 +22,6 @@ namespace AMDGPU { const char* device_label = "amdgpu"; const char* label_cpu = "cpu"; -// Registration name the EP is discovered under, used to look up its advertised memory-info. -constexpr const char* kExecutionProviderName = "AMDGPUExecutionProvider"; - 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} { @@ -206,7 +203,7 @@ struct InterfaceImpl : DeviceInterface { // Read the id of the EP device the model will run on, so the allocators bind to it rather than // assuming device 0. Resolves to 0 on a single-GPU machine. int GetDeviceId(const ProviderOptions* user_options) override { - auto ep_devices = FindRegisteredEpDevices(kExecutionProviderName); + auto ep_devices = FindRegisteredEpDevices(kAMDGPUExecutionProviderName); if (user_options) ep_devices = ApplyDeviceFiltering(*user_options, ep_devices); if (!ep_devices.empty()) { @@ -224,7 +221,7 @@ struct InterfaceImpl : DeviceInterface { if (ort_pinned_allocator_) return; try { - for (const OrtEpDevice* ep_device : FindRegisteredEpDevices(kExecutionProviderName)) { + for (const OrtEpDevice* ep_device : FindRegisteredEpDevices(kAMDGPUExecutionProviderName)) { const OrtMemoryInfo* mi = Ort::api->EpDevice_MemoryInfo(ep_device, OrtDeviceMemoryType_HOST_ACCESSIBLE); if (!mi) diff --git a/src/amdgpu/interface.h b/src/amdgpu/interface.h index 0957f9a1ff..974ac77c4c 100644 --- a/src/amdgpu/interface.h +++ b/src/amdgpu/interface.h @@ -7,6 +7,9 @@ namespace Generators { +// Name the EP library is registered under, and that its OrtEpDevice is discovered by. +constexpr const char* kAMDGPUExecutionProviderName = "AMDGPUExecutionProvider"; + DeviceInterface* GetAMDGPUInterface(); // Inputs-only interface backed by the host-accessible allocator. Decode inputs allocated diff --git a/src/amdgpu/session_options.cpp b/src/amdgpu/session_options.cpp index 6e844094bb..74becf912a 100644 --- a/src/amdgpu/session_options.cpp +++ b/src/amdgpu/session_options.cpp @@ -5,13 +5,86 @@ #include "session_options.h" +#include +#include + +#include "../models/env_utils.h" #include "../models/session_options.h" #include "interface.h" +#if defined(_WIN32) +#include +#endif + namespace Generators::AMDGPUExecutionProvider { namespace { +constexpr const char* kEpPathEnvKey = "AMDGPU_EP_PATH"; +#if defined(_WIN32) +constexpr const char* kEpFilename = "amdgpu-ep.dll"; +#else +constexpr const char* kEpFilename = "libamdgpu-ep.so"; +#endif + +// A plugin EP is only discoverable once its library is registered on the OrtEnv. Hosts that can +// supply a path do that themselves; resolve it here for the ones that cannot. No-op if the EP is +// already registered or the library is not found, so an explicit registration always wins. +void EnsureUmbrellaEpRegistered() { + if (!FindRegisteredEpDevices(kAMDGPUExecutionProviderName).empty()) + return; + + std::error_code ec; + std::filesystem::path ep_path = GetEnv(kEpPathEnvKey); + +#if defined(_WIN32) + const auto module_of = [](const void* address) -> HMODULE { + MEMORY_BASIC_INFORMATION mbi; + if (VirtualQuery(address, &mbi, sizeof(mbi)) && mbi.AllocationBase) + return reinterpret_cast(mbi.AllocationBase); + return nullptr; + }; + + const auto find_next_to_module = [&](HMODULE module) -> std::filesystem::path { + wchar_t buffer[MAX_PATH + 1] = {0}; + if (GetModuleFileNameW(module, buffer, MAX_PATH + 1)) + if (const auto dir = std::filesystem::path{buffer}.remove_filename(); !dir.empty()) + if (auto path = dir / kEpFilename; std::filesystem::exists(path, ec)) + return path; + return {}; + }; + + if (ep_path.empty()) + // next to onnxruntime-genai, using a symbol in that module as the address marker + if (const auto module = module_of(reinterpret_cast(&GetAMDGPUInterface))) + ep_path = find_next_to_module(module); + + if (ep_path.empty()) + // next to onnxruntime + if (const auto module = module_of(reinterpret_cast(Ort::api->RegisterExecutionProviderLibrary))) + ep_path = find_next_to_module(module); + + if (ep_path.empty()) + // next to the current executable + if (const auto module = GetModuleHandleA(nullptr)) + ep_path = find_next_to_module(module); +#endif + + if (ep_path.empty()) + ep_path = std::filesystem::current_path(ec) / kEpFilename; + + if (!std::filesystem::exists(ep_path, ec)) + return; + + try { + Ort::RegisterExecutionProviderLibrary(&GetOrtEnv(), kAMDGPUExecutionProviderName, ep_path.native().c_str()); + } catch (const Ort::Exception& e) { + // Registered but advertising no device: the check above cannot see that, ORT reports it here. + if (std::string(e.what()).find("already registered") == std::string::npos) + throw; + } +} + // 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) { @@ -35,6 +108,8 @@ DeviceInterface* AppendExecutionProvider(OrtSessionOptions& session_options, const Config::ProviderOptions& provider_options, const Config& config, bool /*disable_graph_capture*/) { + EnsureUmbrellaEpRegistered(); + SetStaticPaddingConfig(session_options, config); // Umbrella-level hint: the model architecture drives the EP's backend routing. @@ -44,7 +119,7 @@ DeviceInterface* AppendExecutionProvider(OrtSessionOptions& session_options, session_options.AddConfigEntry("ep.directml.enable_host_accessible", "1"); AppendExecutionProviderV2(session_options, provider_options, - DeviceType::AMDGPU, "AMDGPUExecutionProvider"); + DeviceType::AMDGPU, kAMDGPUExecutionProviderName); return GetAMDGPUInterface(); } From 94515ef14b56cf65d6a8e50e40928eaf52dffc00 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Mon, 10 Aug 2026 12:21:47 -0700 Subject: [PATCH 19/21] Select the host-accessible inputs interface through the device (review #2165) model.cpp called an AMDGPU-specific function to get it, which is the last EP-specialized call in that file. It is now a DeviceInterface virtual that defaults to null, so the shared path asks the device instead of naming one. --- src/amdgpu/interface.cpp | 16 +++++++++++----- src/amdgpu/interface.h | 5 ----- src/models/model.cpp | 3 +-- src/smartptrs.h | 4 ++++ 4 files changed, 16 insertions(+), 12 deletions(-) diff --git a/src/amdgpu/interface.cpp b/src/amdgpu/interface.cpp index 916e47c189..a218f4b3e9 100644 --- a/src/amdgpu/interface.cpp +++ b/src/amdgpu/interface.cpp @@ -242,6 +242,9 @@ struct InterfaceImpl : DeviceInterface { return ort_pinned_allocator_; } + // Defined below, once PinnedInputsImpl is complete. + DeviceInterface* GetHostAccessibleDevice() override; + Ort::Allocator* PinnedAllocator() const { return ort_pinned_allocator_; } std::shared_ptr AllocateBase(size_t size) override { @@ -323,13 +326,16 @@ DeviceInterface* GetAMDGPUInterface() { 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 +namespace AMDGPU { + +DeviceInterface* InterfaceImpl::GetHostAccessibleDevice() { + if (!ort_pinned_allocator_) + return nullptr; if (!g_amdgpu_pinned_inputs) - g_amdgpu_pinned_inputs = std::make_unique(*base); + g_amdgpu_pinned_inputs = std::make_unique(*this); return g_amdgpu_pinned_inputs.get(); } +} // namespace AMDGPU + } // namespace Generators diff --git a/src/amdgpu/interface.h b/src/amdgpu/interface.h index 974ac77c4c..9ad8fd65b0 100644 --- a/src/amdgpu/interface.h +++ b/src/amdgpu/interface.h @@ -12,9 +12,4 @@ constexpr const char* kAMDGPUExecutionProviderName = "AMDGPUExecutionProvider"; 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/models/model.cpp b/src/models/model.cpp index 0a4a014640..2b4145c61e 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -35,7 +35,6 @@ #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" @@ -596,7 +595,7 @@ Model::Model(std::unique_ptr config) : config_{std::move(config)} { // Inputs-only interface backed by a host-accessible allocation, so the CPU updates the small // decode inputs in place with no per-step roundtrip. Null if the device offers no such allocator. - DeviceInterface* p_host_accessible_inputs = GetAMDGPUPinnedInputsInterface(); + DeviceInterface* p_host_accessible_inputs = p_device_->GetHostAccessibleDevice(); // Only CUDA, TRT-RTX, RyzenAI and DML does every input on the device // For WebGPU, use device memory only if graph capture is enabled, otherwise use CPU diff --git a/src/smartptrs.h b/src/smartptrs.h index e6819dc781..7ef55b8e53 100644 --- a/src/smartptrs.h +++ b/src/smartptrs.h @@ -142,6 +142,10 @@ struct DeviceInterface { // supports it. Null default -> callers keep the current device-memory path. virtual Ort::Allocator* GetHostAccessibleAllocator() { return nullptr; } + // Inputs-only interface backed by that allocator, so the decode inputs are updated in place. + // Null default -> callers keep the current device-memory path. + virtual DeviceInterface* GetHostAccessibleDevice() { return nullptr; } + // Called once after the device allocator is created, so a device that offers additional // allocators (e.g. host-accessible memory) can set them up. The default sets up nothing. // `device_id` is the id the device allocator was created on. From 9fb38da5cd8168f66f7f3f7f7083715024a709a1 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Mon, 10 Aug 2026 12:21:47 -0700 Subject: [PATCH 20/21] Correct the p_device_logits_ comment (review #2165) It read as though logits are always read back on the CPU, which is not true for the EPs that keep them on the device. --- src/models/model.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/model.h b/src/models/model.h index 0dc648a9a4..671b4e9369 100644 --- a/src/models/model.h +++ b/src/models/model.h @@ -188,7 +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 read back on the CPU every step + DeviceInterface* p_device_logits_{}; // Matches p_device_inputs_ unless the device reads logits back on the CPU 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 7a2cd43f69428c3f60a567f3db156a8186e263fb Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Mon, 10 Aug 2026 12:53:22 -0700 Subject: [PATCH 21/21] Drop a stray blank line in GenerateNextToken It was accidental whitespace with no effect, and it sat in a region main has since rewritten, so it was the only thing conflicting with upstream. --- src/generators.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/generators.cpp b/src/generators.cpp index db611979e3..8c470c0ec4 100644 --- a/src/generators.cpp +++ b/src/generators.cpp @@ -861,7 +861,6 @@ void Generator::GenerateNextToken() { auto next_tokens = search_->GetNextTokens(); if (last_action_ == Action::rewound) search_->AppendTokens(next_tokens); - ComputeLogits(next_tokens); } if (guidance_logits_processor_) {