diff --git a/examples/python/common.py b/examples/python/common.py index a9002a3522..1ebd1fee82 100644 --- a/examples/python/common.py +++ b/examples/python/common.py @@ -50,13 +50,21 @@ def register_ep(ep: str, ep_path: str, use_winml: bool) -> None: print(f"Registered {ep} from {ep_path} successfully!") -def get_config(path: str, ep: str, ep_options: dict[str, str] = {}, search_options: dict[str, int] = {}) -> og.Config: +def get_config( + path: str, + ep: str, + ep_path: str | None, + ep_options: dict[str, str] = {}, + search_options: dict[str, int] = {}, +) -> og.Config: """ Get og.Config object and set EP-specific and search-specific options inside it Args: path (str): Path to model folder containing GenAI config ep (str): Name of execution provider to set + ep_path (str | None): Path to an external execution provider library. If set, the + registered library is used and providers from the GenAI config are preserved. ep_options (dict[str, str]): Map of EP-specific option names and their values search_options (dict[str, int]): Map of search-specific option names and their values Returns: @@ -66,7 +74,7 @@ def get_config(path: str, ep: str, ep_options: dict[str, str] = {}, search_optio # - If follow_config, then use the default EP stored inside the GenAI config. # - Otherwise, override the stored EP by clearing all providers and appending the desired one. config = og.Config(path) - if ep != "follow_config": + if not ep_path and ep != "follow_config": config.clear_providers() if ep != "cpu": print(f"Setting model to {ep}") diff --git a/examples/python/model-chat.py b/examples/python/model-chat.py index 0edad43198..edc6ded89f 100644 --- a/examples/python/model-chat.py +++ b/examples/python/model-chat.py @@ -19,7 +19,7 @@ def main(args): print("Loading model...") # Create model - config = get_config(args.model_path, args.execution_provider) + config = get_config(args.model_path, args.execution_provider, args.ep_path) model = og.Model(config) if args.verbose: print("Model loaded") diff --git a/examples/python/model-generate.py b/examples/python/model-generate.py index 67fb1a5795..9d22564916 100644 --- a/examples/python/model-generate.py +++ b/examples/python/model-generate.py @@ -32,7 +32,9 @@ def main(args): prompts = [text] setattr(args, "batch_size", len(prompts)) search_config = {"batch_size": args.batch_size, "chunk_size": args.chunk_size, "num_beams": args.num_beams} - config = get_config(args.model_path, args.execution_provider, ep_options={}, search_options=search_config) + config = get_config( + args.model_path, args.execution_provider, args.ep_path, ep_options={}, search_options=search_config + ) model = og.Model(config) if args.verbose: diff --git a/examples/python/model-mm.py b/examples/python/model-mm.py index 5a938cd670..d69d6195ca 100644 --- a/examples/python/model-mm.py +++ b/examples/python/model-mm.py @@ -32,7 +32,7 @@ def main(args): print("Loading model...") # Create model - config = get_config(args.model_path, args.execution_provider) + config = get_config(args.model_path, args.execution_provider, args.ep_path) model = og.Model(config) if args.verbose: print("Model loaded") diff --git a/examples/python/model-qa.py b/examples/python/model-qa.py index 733140bae9..aed1b1f23c 100644 --- a/examples/python/model-qa.py +++ b/examples/python/model-qa.py @@ -29,7 +29,7 @@ def main(args): print("Loading model...") # Create model - config = get_config(args.model_path, args.execution_provider) + config = get_config(args.model_path, args.execution_provider, args.ep_path) model = og.Model(config) if args.verbose: print("Model loaded") diff --git a/examples/python/nemotron_speech.py b/examples/python/nemotron_speech.py index 54b2150c55..2061cd1528 100644 --- a/examples/python/nemotron_speech.py +++ b/examples/python/nemotron_speech.py @@ -96,7 +96,7 @@ def simulate_microphone(model_path, audio_path, execution_provider, use_vad=None audio = load_audio(audio_path, sample_rate) duration = len(audio) / sample_rate - config = get_config(model_path, execution_provider) + config = get_config(model_path, execution_provider, None) selected_lang = None if language is not None: if language not in LANG_TO_ID: diff --git a/src/models/recurrent_state.cpp b/src/models/recurrent_state.cpp index 7d2f4104c2..8be92368a3 100644 --- a/src/models/recurrent_state.cpp +++ b/src/models/recurrent_state.cpp @@ -88,73 +88,44 @@ RecurrentState::RecurrentState(State& state) const int num_layers = static_cast(layer_indices_.size()); - pasts_.resize(num_layers * 2); + const bool past_present_share_buffer = state_.params_->IsPastPresentShareBufferEnabled(model_.config_->model.type); + if (!past_present_share_buffer) { + throw std::runtime_error( + "RecurrentState requires past_present_share_buffer=true. " + "Set past_present_share_buffer to true in genai_config.json."); + } + presents_.reserve(num_layers * 2); auto& allocator = model_.p_device_kvcache_->GetAllocator(); + // Qwen3.5 linear-attention state is a compressed recurrent state, not a + // token-indexed KV cache. For graph replay, bind each state tensor as both + // past input and present output so ORT/TRT-RTX sees stable addresses. + // The EP/plugin kernels must read the previous contents before writing the + // updated state back to the same buffer. for (int i = 0; i < num_layers; ++i) { - pasts_[i * 2] = OrtValue::CreateTensor(allocator, conv_shape_, conv_type_); - pasts_[i * 2 + 1] = OrtValue::CreateTensor(allocator, recurrent_shape_, recurrent_type_); - presents_.push_back(OrtValue::CreateTensor(allocator, conv_shape_, conv_type_)); presents_.push_back(OrtValue::CreateTensor(allocator, recurrent_shape_, recurrent_type_)); } - // Zero-initialize past and present states - ZeroStates(pasts_); ZeroStates(presents_); } void RecurrentState::Add() { if (layer_indices_.empty()) return; - input_index_ = state_.inputs_.size(); - output_index_ = state_.outputs_.size(); - const int num_layers = static_cast(layer_indices_.size()); for (int i = 0; i < num_layers * 2; ++i) { - state_.inputs_.push_back(pasts_[i].get()); + auto* past = presents_[i].get(); + state_.inputs_.push_back(past); state_.input_names_.push_back(input_name_strings_[i].c_str()); state_.outputs_.push_back(presents_[i].get()); state_.output_names_.push_back(output_name_strings_[i].c_str()); } - - // Cache byte spans for the graph-capture copy path. These tensors are - // fixed-shape and never reallocated, so the spans remain valid for the - // entire generation lifetime. - if (state_.params_->use_graph_capture) { - auto& device = *model_.p_device_kvcache_; - past_byte_spans_.reserve(num_layers * 2); - present_byte_spans_.reserve(num_layers * 2); - for (int i = 0; i < num_layers * 2; ++i) { - past_byte_spans_.push_back(ByteWrapTensor(device, *pasts_[i])); - present_byte_spans_.push_back(ByteWrapTensor(device, *presents_[i])); - } - } } void RecurrentState::Update() { - if (layer_indices_.empty()) return; - - const int num_layers = static_cast(layer_indices_.size()); - - if (state_.params_->use_graph_capture) { - // When graph capture is enabled, we must not swap pointers because the - // graph has captured the original memory addresses. Instead, copy - // present→past in-place so the pointers remain stable. Uses cached byte - // spans to avoid recomputing tensor metadata each step. - for (int i = 0; i < num_layers * 2; ++i) { - past_byte_spans_[i].CopyFrom(present_byte_spans_[i]); - } - // No need to rebind state_.inputs_/outputs_ — pointers are unchanged. - } else { - for (int i = 0; i < num_layers * 2; ++i) { - std::swap(pasts_[i], presents_[i]); - state_.inputs_[input_index_ + i] = pasts_[i].get(); - state_.outputs_[output_index_ + i] = presents_[i].get(); - } - } } void RecurrentState::RewindTo(size_t index) { @@ -168,19 +139,10 @@ void RecurrentState::RewindTo(size_t index) { ") is a no-op. Recurrent states cannot be partially rewound."); return; } - - const int num_layers = static_cast(layer_indices_.size()); - - // Zero existing buffers in-place instead of reallocating, to preserve - // device pointers and avoid invalidating captured graphs. - ZeroStates(pasts_); + // Shared recurrent states keep stable input/output pointers for graph replay. + // Reset the state contents in place without rebinding. ZeroStates(presents_); - - // Re-bind state pointers (swap may have changed which OrtValue is past vs present) - for (int i = 0; i < num_layers * 2; ++i) { - state_.inputs_[input_index_ + i] = pasts_[i].get(); - state_.outputs_[output_index_ + i] = presents_[i].get(); - } + return; } void RecurrentState::ZeroStates(std::vector>& states) { diff --git a/src/models/recurrent_state.h b/src/models/recurrent_state.h index dc7d716a4c..e7c3c42734 100644 --- a/src/models/recurrent_state.h +++ b/src/models/recurrent_state.h @@ -27,21 +27,12 @@ struct RecurrentState { std::vector layer_indices_; // Interleaved as [conv_0, recurrent_0, conv_1, recurrent_1, ...] - std::vector> pasts_; std::vector> presents_; - // Cached byte spans for graph-capture copy path (avoids recomputing - // tensor metadata on every decode step for fixed-shape tensors). - std::vector> past_byte_spans_; - std::vector> present_byte_spans_; - // Kept alive for state_ const char* pointers std::vector input_name_strings_; std::vector output_name_strings_; - size_t input_index_{~0U}; - size_t output_index_{~0U}; - ONNXTensorElementDataType conv_type_{}; ONNXTensorElementDataType recurrent_type_{}; diff --git a/test/test_models/qwen35-hybrid-preprocessing/genai_config.json b/test/test_models/qwen35-hybrid-preprocessing/genai_config.json index d03eacee34..9ae14816c4 100644 --- a/test/test_models/qwen35-hybrid-preprocessing/genai_config.json +++ b/test/test_models/qwen35-hybrid-preprocessing/genai_config.json @@ -67,7 +67,7 @@ "no_repeat_ngram_size": 0, "num_beams": 1, "num_return_sequences": 1, - "past_present_share_buffer": false, + "past_present_share_buffer": true, "repetition_penalty": 1.0, "temperature": 1.0, "top_k": 1,