diff --git a/src/migraphx/mgx_ep.cc b/src/migraphx/mgx_ep.cc index b0df633..d0dcf4f 100644 --- a/src/migraphx/mgx_ep.cc +++ b/src/migraphx/mgx_ep.cc @@ -456,6 +456,8 @@ ExecutionProvider::ExecutionProvider(const ProviderFactory& factory, std::string PARSE_ENV_VAR(env_var::kDumpSubgraphs, dump_subgraphs_); PARSE_ENV_VAR(env_var::kDumpEpContextModel, context_enable_); PARSE_ENV_VAR(env_var::kExhaustiveTune, exhaustive_tune_); + PARSE_ENV_VAR(env_var::kCoresidentPrograms, coresident_programs_); + PARSE_ENV_VAR(env_var::kMaxResidentPrograms, max_resident_programs_); PARSE_ENV_VAR(env_var::kMlssUseSpecificOps, mlss_use_specific_ops_); platform::SetEnvironmentVar("MIGRAPHX_MLSS_USE_SPECIFIC_OPS", mlss_use_specific_ops_); @@ -798,6 +800,11 @@ Ort::Status ExecutionProvider::CreateNodeComputeInfoFromGraph(const Ort::ConstGr t_, onnx_options, program, + coresident_programs_, + max_resident_programs_, + {}, // resident_programs + {}, // lru_order + {}, // active_shape_key enable_fp16_, enable_bf16_, enable_fp8_, @@ -1017,6 +1024,50 @@ Ort::Status NodeComputeInfo::Compute(ComputeState& compute_state, const Ort::Ker } } + // Registers a program in the resident cache as the most-recently-used entry + // and evicts least-recently-used entries until the bound is met. The key just + // registered is never the eviction victim. + const auto register_resident = [&compute_state](const std::string& key, const migraphx::program& prog) { + compute_state.active_shape_key = key; + compute_state.resident_programs[key] = prog; + compute_state.lru_order.remove(key); + compute_state.lru_order.push_back(key); + while (compute_state.resident_programs.size() > compute_state.max_resident_programs && + !compute_state.lru_order.empty()) { + const auto victim{compute_state.lru_order.front()}; + compute_state.lru_order.pop_front(); + if (victim != key) { + compute_state.resident_programs.erase(victim); + } + } + }; + + // On a shape switch, stash the active program in the resident cache under its + // (previous) shape key, then check whether the incoming shape is already + // resident. On a hit, restore it and skip the reload entirely (a lookup that + // replaces a multi-second reload and finalize). Default off: falls through to + // the legacy single-slot overwrite path below. + if (!input_shapes_match && compute_state.coresident_programs) { + const std::string shape_key{hash::ToHex(input_shapes_hash)}; + // Save the currently-active program under the previously-active key + // (if we have one and it isn't already stored). + if (!compute_state.active_shape_key.empty() && + compute_state.resident_programs.find(compute_state.active_shape_key) == + compute_state.resident_programs.end()) { + compute_state.resident_programs.emplace(compute_state.active_shape_key, program); + compute_state.lru_order.push_back(compute_state.active_shape_key); + } + if (const auto it{compute_state.resident_programs.find(shape_key)}; + it != compute_state.resident_programs.end()) { + program = it->second; // shared_ptr handle copy, no reload/finalize + param_shapes = program.get_parameter_shapes(); + input_shapes_match = true; + // Mark the restored program most-recently-used and enforce the bound + // (the stash above may have pushed the cache one past max). + register_resident(shape_key, program); + } + } + // If the input shapes are different (e.g., LLMs), the EP needs to reparse and recompile the program if (!input_shapes_match) { migraphx::program_parameters compile_params{}; @@ -1061,6 +1112,12 @@ Ort::Status NodeComputeInfo::Compute(ComputeState& compute_state, const Ort::Ker } } param_shapes = program.get_parameter_shapes(); + + // The freshly loaded/compiled program is now the active one; register it + // in the resident cache and enforce the LRU bound. + if (compute_state.coresident_programs) { + register_resident(hash::ToHex(input_shapes_hash), program); + } } migraphx::program_parameters compute_params; diff --git a/src/migraphx/mgx_ep.h b/src/migraphx/mgx_ep.h index b0b3cd2..ddc5400 100644 --- a/src/migraphx/mgx_ep.h +++ b/src/migraphx/mgx_ep.h @@ -4,8 +4,10 @@ #pragma once #include +#include #include #include +#include #include #include @@ -34,6 +36,12 @@ constexpr auto kComputeMode = "ORT_MIGRAPHX_COMPUTE_MODE"sv; constexpr auto kINT8UseNativeCalibrationTable = "ORT_MIGRAPHX_INT8_USE_NATIVE_CALIBRATION_TABLE"sv; constexpr auto kExhaustiveTune = "ORT_MIGRAPHX_EXHAUSTIVE_TUNE"sv; constexpr auto kMlssUseSpecificOps = "ORT_MIGRAPHX_MLSS_USE_SPECIFIC_OPS"sv; +// Keep multiple shape-specialized programs (for example an LLM's prefill and +// decode programs) resident at once so a shape switch is a map lookup instead of +// a full reload and finalize. Default off (single-slot, legacy overwrite +// behavior). When on, the cache is bounded by kMaxResidentPrograms. +constexpr auto kCoresidentPrograms = "ORT_MIGRAPHX_CORESIDENT_PROGRAMS"sv; +constexpr auto kMaxResidentPrograms = "ORT_MIGRAPHX_MAX_RESIDENT_PROGRAMS"sv; } // namespace env_vars struct ComputeState { @@ -41,7 +49,18 @@ struct ComputeState { int device_id; const migraphx::target& t; migraphx::onnx_options onnx_options; - migraphx::program program; + migraphx::program program; // currently-active program (last shape run) + // shape-hash -> resident program cache. When coresident_programs is on, + // programs for distinct input shapes are kept here so switching is a lookup, + // not a reload and finalize. lru_order is most-recent-last; on overflow past + // max_resident_programs the front (least-recently-used) entry is evicted. The + // program run every step stays at the back and so is never the victim. + // active_shape_key is the key of the program currently in the `program` slot. + bool coresident_programs{}; + size_t max_resident_programs{4}; + std::unordered_map resident_programs; + std::list lru_order; + std::string active_shape_key; bool enable_fp16{}; bool enable_bf16{}; bool enable_fp8{}; @@ -146,6 +165,8 @@ struct ExecutionProvider : OrtEp, ApiPtrs { bool enable_fp8_{}; bool enable_int8_{}; bool exhaustive_tune_{}; + bool coresident_programs_{}; + size_t max_resident_programs_{4}; std::string mlss_use_specific_ops_{}; bool int8_calibration_cache_available_{}; bool int8_use_native_calibration_table_{};